Context
Context allows you to inject additional dependencies or configuration into your transformations. You
can do this using the Context() marker in your transformation definition:
A PySpark pipeline that only requires a Spark session can use the built-in SparkContext:
Define a custom dataclass when the pipeline has additional dependencies:
from dataclasses import dataclass
from pathlib import Path
from typing import Annotated
import polars as pl
from data_joinery import Context, Pipeline, Strict, transform
@dataclass
class A:
a: str
@dataclass
class B:
b: str
@dataclass
class PathConfig:
input_path: Path
output_path: Path
@dataclass(frozen=True)
class PipelineDependencies:
paths: PathConfig
@transform
def read_data(
path_config: Annotated[PathConfig, Context()],
) -> Annotated[pl.LazyFrame, Strict(A)]:
return pl.scan_csv(path_config.input_path, schema_overrides={"a": pl.String})
@transform
def transform_data(
a: Annotated[pl.LazyFrame, Strict(A)],
) -> Annotated[pl.LazyFrame, Strict(B)]:
return a.rename({"a": "b"})
@transform
def write_data(
b: Annotated[pl.LazyFrame, Strict(B)],
path_config: Annotated[PathConfig, Context()],
) -> None:
print("Write to ", path_config.output_path.name)
print(b.collect())
order_metrics = Pipeline(PipelineDependencies)
read_data_step = order_metrics.add_step(read_data)
transform_data_step = order_metrics.add_step(transform_data)
write_data_step = order_metrics.add_step(write_data)
order_metrics.connect(read_data_step, transform_data_step)
order_metrics.connect(transform_data_step, write_data_step)
input_path = Path(__file__).parent / "example.csv"
output_path = Path(__file__).parent / "output.parquet"
path_config = PathConfig(input_path=input_path, output_path=output_path)
context = PipelineDependencies(path_config)
order_metrics.run(context)
Outputs:
from dataclasses import dataclass
from pathlib import Path
from typing import Annotated
from pyspark.sql import DataFrame, SparkSession
from data_joinery import Context, Pipeline, Strict, transform
@dataclass
class A:
a: str
@dataclass
class B:
b: str
@dataclass
class PathConfig:
input_path: Path
output_path: Path
@dataclass(frozen=True)
class PipelineDependencies:
spark: SparkSession
paths: PathConfig
@transform
def read_data(
spark: Annotated[SparkSession, Context()],
path_config: Annotated[PathConfig, Context()],
) -> Annotated[DataFrame, Strict(A)]:
return spark.read.csv(str(path_config.input_path), header=True)
@transform
def transform_data(
a: Annotated[DataFrame, Strict(A)],
) -> Annotated[DataFrame, Strict(B)]:
return a.withColumnRenamed("a", "b")
@transform
def write_data(
b: Annotated[DataFrame, Strict(B)],
path_config: Annotated[PathConfig, Context()],
) -> None:
print("Write to ", path_config.output_path.name)
b.show()
order_metrics = Pipeline(PipelineDependencies)
read_data_step = order_metrics.add_step(read_data)
transform_data_step = order_metrics.add_step(transform_data)
write_data_step = order_metrics.add_step(write_data)
order_metrics.connect(read_data_step, transform_data_step)
order_metrics.connect(transform_data_step, write_data_step)
input_path = Path(__file__).parent / "example.csv"
output_path = Path(__file__).parent / "output.parquet"
path_config = PathConfig(input_path=input_path, output_path=output_path)
context = PipelineDependencies(SparkSession.builder.getOrCreate(), path_config)
order_metrics.run(context)
Outputs:
Context is automatically injected into your transformation from the context dataclass declared by
the pipeline. Each field type identifies one dependency, so a context cannot contain multiple fields
of the same type. Use your own class definitions instead of built-in types like str or int when
their meaning matters. One common example is a run date for reading day-partitioned data. It is best
practice to define a distinct date subclass so the dependency is explicit:
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Annotated
import polars as pl
from data_joinery import Context, Pipeline, Strict, transform
@dataclass
class A:
snapshot_date: str
class RunDate(date):
pass
@dataclass(frozen=True)
class PipelineDependencies:
run_date: RunDate
@transform
def read_data(
run_date: Annotated[RunDate, Context()],
) -> Annotated[pl.LazyFrame, Strict(A)]:
expected_path = Path(__file__).parent / f"example_{run_date.strftime('%Y%m%d')}.csv"
return pl.scan_csv(expected_path, schema_overrides={"snapshot_date": pl.String})
@transform
def write_data(b: Annotated[pl.LazyFrame, Strict(A)]) -> None:
print(b.collect())
order_metrics = Pipeline(PipelineDependencies)
read_data_step = order_metrics.add_step(read_data)
write_data_step = order_metrics.add_step(write_data)
order_metrics.connect(read_data_step, write_data_step)
context = PipelineDependencies(RunDate(2026, 1, 1))
order_metrics.run(context)
Outputs:
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Annotated
from pyspark.sql import DataFrame, SparkSession
from data_joinery import Context, Pipeline, Strict, transform
@dataclass
class A:
snapshot_date: str
class RunDate(date):
pass
@dataclass(frozen=True)
class PipelineDependencies:
spark: SparkSession
run_date: RunDate
@transform
def read_data(
spark: Annotated[SparkSession, Context()],
run_date: Annotated[RunDate, Context()],
) -> Annotated[DataFrame, Strict(A)]:
expected_path = Path(__file__).parent / f"example_{run_date.strftime('%Y%m%d')}.csv"
return spark.read.csv(str(expected_path), header=True)
@transform
def write_data(b: Annotated[DataFrame, Strict(A)]) -> None:
b.show()
order_metrics = Pipeline(PipelineDependencies)
read_data_step = order_metrics.add_step(read_data)
write_data_step = order_metrics.add_step(write_data)
order_metrics.connect(read_data_step, write_data_step)
context = PipelineDependencies(SparkSession.builder.getOrCreate(), RunDate(2026, 1, 1))
order_metrics.run(context)
Outputs:
Validating Context
The pipeline validates every transformation against its declared context when the step is added. If a transformation requires a dependency that the context dataclass does not provide, construction fails immediately:
import traceback
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Annotated
import polars as pl
from data_joinery import Context, Pipeline, Strict, transform
@dataclass
class A:
snapshot_date: str
class RunDate(date):
pass
@dataclass(frozen=True)
class PipelineDependencies:
run_date: date
@transform
def read_data(
run_date: Annotated[RunDate, Context()],
) -> Annotated[pl.LazyFrame, Strict(A)]:
expected_path = Path(__file__).parent / f"example_{run_date.strftime('%Y%m%d')}.csv"
return pl.scan_csv(expected_path)
@transform
def write_data(b: Annotated[pl.LazyFrame, Strict(A)]) -> None:
print(b.collect())
try:
order_metrics = Pipeline(PipelineDependencies)
order_metrics.add_step(read_data)
except TypeError:
print(traceback.format_exc(limit=1))
Outputs:
Traceback (most recent call last):
File "/home/runner/work/data-joinery/data-joinery/docs_src/learn/pipelines/context/context_example_unsatisifed_polars.py", line 41, in <module>
order_metrics.add_step(read_data)
TypeError: Transform 'read_data' parameter 'run_date' requires context dependency RunDate, but PipelineDependencies does not provide it. Add a field of that type to the context dataclass or change the Context annotation.
import traceback
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Annotated
from pyspark.sql import DataFrame, SparkSession
from data_joinery import Context, Pipeline, Strict, transform
@dataclass
class A:
snapshot_date: str
class RunDate(date):
pass
@dataclass(frozen=True)
class PipelineDependencies:
spark: SparkSession
run_date: date
@transform
def read_data(
spark: Annotated[SparkSession, Context()],
run_date: Annotated[RunDate, Context()],
) -> Annotated[DataFrame, Strict(A)]:
expected_path = Path(__file__).parent / f"example_{run_date.strftime('%Y%m%d')}.csv"
return spark.read.csv(str(expected_path), header=True)
@transform
def write_data(b: Annotated[DataFrame, Strict(A)]) -> None:
b.show()
try:
order_metrics = Pipeline(PipelineDependencies)
order_metrics.add_step(read_data)
except TypeError:
print(traceback.format_exc(limit=1))
Outputs:
Traceback (most recent call last):
File "/home/runner/work/data-joinery/data-joinery/docs_src/learn/pipelines/context/context_example_unsatisifed_spark.py", line 43, in <module>
order_metrics.add_step(read_data)
TypeError: Transform 'read_data' parameter 'run_date' requires context dependency RunDate, but PipelineDependencies does not provide it. Add a field of that type to the context dataclass or change the Context annotation.