yggdrasil.spark.extensions

extensions

Spark DataFrame extension helpers for aliases and resampling.

getAliases

getAliases(
    obj: Union[DataFrame, Column, str, Iterable[Union[DataFrame, Column, str]]],
    full: bool = True,
) -> list[str]

Return aliases for Spark columns/dataframes or collections.

Parameters:

Name	Type	Description	Default
obj	Union[DataFrame, Column, str, Iterable[Union[DataFrame, Column, str]]]	Spark DataFrame/Column, string, or iterable of these.	required
full	bool	Whether to return full qualified names.	True

Returns:

Type	Description
list[str]	List of alias strings.

getAlias

getAlias(obj: Union[DataFrame, Column, str], full: bool = True) -> str

Parse a column name out of a PySpark Column repr string.

latest

latest(
    df: DataFrame,
    partitionBy: List[Union[str, Column]],
    orderBy: List[Union[str, Column]],
) -> SparkSQL.DataFrame

Return the latest rows per partition based on ordering.

Parameters:

Name	Type	Description	Default
df	DataFrame	Spark DataFrame.	required
partitionBy	List[Union[str, Column]]	Columns to partition by.	required
orderBy	List[Union[str, Column]]	Columns to order by.	required

Returns:

Type	Description
DataFrame	Spark DataFrame with latest rows per partition.

checkJoin

checkJoin(
    df: DataFrame,
    other: DataFrame,
    on: Optional[Union[str, List[str], Column, List[Column]]] = None,
    *args: Any,
    **kwargs: Any
) -> SparkSQL.DataFrame

Join two DataFrames with schema-aware column casting.

Parameters:

Name	Type	Description	Default
df	DataFrame	Left Spark DataFrame.	required
other	DataFrame	Right Spark DataFrame.	required
on	Optional[Union[str, List[str], Column, List[Column]]]	Join keys or mapping.	None
*args	Any	Positional args passed to join.	()
**kwargs	Any	Keyword args passed to join.	{}

Returns:

Type	Description
DataFrame	Joined Spark DataFrame.

checkMapInArrow

checkMapInArrow(
    df: DataFrame,
    func: Callable[[Iterable[RecordBatch]], Iterable[RecordBatch]],
    schema: Union[StructType, str],
    *args: Any,
    **kwargs: Any
) -> SparkSQL.DataFrame

Wrap mapInArrow to enforce output schema conversion.

Parameters:

Name	Type	Description	Default
df	DataFrame	Spark DataFrame.	required
func	Callable[[Iterable[RecordBatch]], Iterable[RecordBatch]]	Generator function yielding RecordBatches.	required
schema	Union[StructType, str]	Output schema (Spark StructType or DDL string).	required
*args	Any	Positional args passed to mapInArrow.	()
**kwargs	Any	Keyword args passed to mapInArrow.	{}

Returns:

Type	Description
DataFrame	Spark DataFrame with enforced schema.

checkMapInPandas

checkMapInPandas(
    df: DataFrame,
    func: Callable[[Iterable[DataFrame]], Iterable[DataFrame]],
    schema: Union[StructType, str],
    *args: Any,
    **kwargs: Any
) -> SparkSQL.DataFrame

Wrap mapInPandas to enforce output schema conversion.

Parameters:

Name	Type	Description	Default
df	DataFrame	Spark DataFrame.	required
func	Callable[[Iterable[DataFrame]], Iterable[DataFrame]]	Generator function yielding pandas DataFrames.	required
schema	Union[StructType, str]	Output schema (Spark StructType or DDL string).	required
*args	Any	Positional args passed to mapInPandas.	()
**kwargs	Any	Keyword args passed to mapInPandas.	{}

Returns:

Type	Description
DataFrame	Spark DataFrame with enforced schema.