spark_io_decorator

spark_io_decorator is a Python decorator utility for Apache Spark that enables seamless persistence of input and output DataFrames in DAG-based (Directed Acyclic Graph) workflows.

---

Installation

 pip install -i https://test.pypi.org/simple/ spark-io-decorator==0.0.1

Basic Example:

test/test_run.py has this example, don't forget to replace the base_path variable base_path = "<data_path>"

DAG root node (no input DataFrames)

from decorator import persist_df
from pyspark.sql import SparkSession

@persist_df(output_path=f"{base_path}/output/raw.parquet")
def load_initial_data():
    spark = SparkSession.getActiveSession()
    return spark.read.csv(f"{base_path}/source.csv", header=True, inferSchema=True)

Function with input DataFrames

@persist_df(
    input_paths=[f"{base_path}/input/df1.parquet", f"{base_path}/input/df2.parquet"],
    output_path=f"{base_path}/output/joined.parquet"
)
def join_by_id(df1, df2):
    """Performs an inner join on two dataframes by their "id" column"""
    return df1.join(df2, on="id")