Skip to content

yggdrasil.plan.func_registry

func_registry

SQL function registry with Arrow-native UDF execution.

:class:FunctionRegistry maps SQL function names to Arrow compute kernels. Built-in functions (UPPER, LOWER, ABS, …) are pre-wired to pyarrow.compute kernels. User-defined functions register a Python callable that operates on pa.Array arguments and returns a pa.Array.

The same kernels auto-register in Spark via spark.udf.register when :meth:FunctionRegistry.register_in_spark is called, wrapping the Arrow callable in a pandas_udf so data stays columnar.

FunctionRegistry

FunctionRegistry()

register_udf

register_udf(
    name: str,
    kernel: ArrowKernel,
    *,
    min_args: int = 1,
    max_args: int | None = None
) -> FunctionMeta

Register a user-defined Arrow-array function.

apply_arrow

apply_arrow(name: str, *arrays: Array) -> pa.Array | None

Execute function on Arrow arrays. Returns None if no kernel.

apply_table

apply_table(name: str, table: Table, col_names: list[str]) -> pa.Array | None

Convenience: extract named columns and apply kernel.

register_in_spark

register_in_spark(spark_session: Any) -> int

Register all kerneled functions as Spark SQL UDFs.

Uses pandas_udf so data stays columnar (Arrow-backed). Returns the number of functions registered.

explode_table

explode_table(
    table: Table, list_col: str, out_col: str | None = None
) -> pa.Table

Explode a list column — equivalent to SQL LATERAL VIEW EXPLODE.

Each row with a list of N elements produces N rows; scalar columns are repeated. Uses Arrow-native list_flatten + list_parent_indices for zero-copy where possible.

posexplode_table

posexplode_table(
    table: Table,
    list_col: str,
    pos_col: str = "pos",
    out_col: str | None = None,
) -> pa.Table

Explode with position — equivalent to POSEXPLODE.