register_spark_dataframe_namespace

laktory.api.register_spark_dataframe_namespace(name)

Decorator for registering a custom namespace on PySpark's DataFrame class.

The decorated class receives the native PySpark DataFrame as its first __init__ argument, so its methods are written in pure PySpark without any Narwhals knowledge.

Use with DATAFRAME_API=NATIVE in pipeline YAML.

PARAMETER	DESCRIPTION
name	Namespace name, used as func_name: "<name>.<method>" in YAML. TYPE: str

Examples:

import pyspark.sql.functions as F
import laktory as lk


@lk.api.register_spark_dataframe_namespace("custom")
class CustomOps:
    def __init__(self, _df):
        self._df = _df

    def with_x2(self):
        return self._df.withColumn("x2", F.col("x1") * 2)

transformer:
  nodes:
    - func_name: custom.with_x2
      dataframe_api: NATIVE

References

• Spark Extension

Source code in laktory/api/namespace.py

def register_spark_dataframe_namespace(name: str):
    """
    Decorator for registering a custom namespace on PySpark's ``DataFrame`` class.

    The decorated class receives the **native PySpark DataFrame** as its first
    ``__init__`` argument, so its methods are written in pure PySpark without
    any Narwhals knowledge.

    Use with ``DATAFRAME_API=NATIVE`` in pipeline YAML.

    Parameters
    ----------
    name:
        Namespace name, used as ``func_name: "<name>.<method>"`` in YAML.

    Examples
    --------
    ```py
    import pyspark.sql.functions as F
    import laktory as lk


    @lk.api.register_spark_dataframe_namespace("custom")
    class CustomOps:
        def __init__(self, _df):
            self._df = _df

        def with_x2(self):
            return self._df.withColumn("x2", F.col("x1") * 2)
    ```

    ```yaml
    transformer:
      nodes:
        - func_name: custom.with_x2
          dataframe_api: NATIVE
    ```

    References
    ----------
    * [Spark Extension](https://www.laktory.ai/concepts/extension_custom/)
    """

    def wrapper(ns_cls: type):
        from pyspark.sql import DataFrame as SparkDataFrame

        setattr(SparkDataFrame, name, SparkNameSpace(name, ns_cls))
        try:
            from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame

            setattr(ConnectDataFrame, name, SparkNameSpace(name, ns_cls))
        except ImportError:
            pass
        return ns_cls

    return wrapper