Skip to content

Transformations

Data Joinery provides a decorator for defining transformations on DataFrames. This allows you to annotate input and output schemas making it much clearer what the transformation does:

For transforms that write or print data, see Side effects.

from dataclasses import dataclass
from typing import Annotated

import polars as pl

from data_joinery import Project, Strict, transform
from data_joinery.schemas import Schema


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@transform
def filter_active_customers(
    customers: Annotated[pl.LazyFrame, Project(Customer)],
) -> Annotated[pl.LazyFrame, Strict(Customer)]:
    return customers.filter(pl.col("is_active"))


customers = Schema(Customer).create_dataframe(
    [
        Customer(customer_id="1", name="Alice", is_active=True),
        Customer(customer_id="2", name="Bob", is_active=False),
        Customer(customer_id="3", name="Charlie", is_active=True),
    ],
    pl.LazyFrame,
)

active_customers = filter_active_customers(customers)
print(active_customers.collect())

Outputs:

shape: (2, 3)
┌─────────────┬─────────┬───────────┐
│ customer_id ┆ name    ┆ is_active │
│ ---         ┆ ---     ┆ ---       │
│ str         ┆ str     ┆ bool      │
╞═════════════╪═════════╪═══════════╡
│ 1           ┆ Alice   ┆ true      │
│ 3           ┆ Charlie ┆ true      │
└─────────────┴─────────┴───────────┘

from dataclasses import dataclass
from typing import Annotated

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Project, Strict, transform
from data_joinery.schemas import Schema


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@transform
def filter_active_customers(
    customers: Annotated[DataFrame, Project(Customer)],
) -> Annotated[DataFrame, Strict(Customer)]:
    return customers.filter(customers.is_active)


spark = SparkSession.builder.getOrCreate()

customers = Schema(Customer).create_dataframe(
    [
        Customer(customer_id="1", name="Alice", is_active=True),
        Customer(customer_id="2", name="Bob", is_active=False),
        Customer(customer_id="3", name="Charlie", is_active=True),
    ],
    DataFrame,
    session=spark,
)

active_customers = filter_active_customers(customers)
active_customers.show()

Outputs:

+-----------+-------+---------+
|customer_id|   name|is_active|
+-----------+-------+---------+
|          1|  Alice|     true|
|          3|Charlie|     true|
+-----------+-------+---------+

Defining a contract

You can define an input contract by annotating any input parameters with a valid coercion mode and schema. The example above accepts a projected input and checks its output strictly. The contract lets the transformation accept input with extra fields while keeping its output precise.

For read steps of wide/nested tables you might want to use Project as the output coercion mode instead of writing out the full schema explicitly. This can be especially useful if you only want to select a few fields of highly nested data:

@transform
def read_nested_event_data(
    spark: Annotated[SparkSession, Context()],
    path: str
) -> Annotated[DataFrame, Project(Event)]:
    return spark.read.parquet(path)

Running Transformations

Call a transformation on a DataFrame. You can chain transformations in either backend:

from dataclasses import dataclass
from typing import Annotated

import polars as pl

from data_joinery import Schema, Strict, transform


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@dataclass
class CustomerCount:
    customer_count: int


@transform
def filter_active_customers(
    customers: Annotated[pl.LazyFrame, Strict(Customer)],
) -> Annotated[pl.LazyFrame, Strict(Customer)]:
    return customers.filter(pl.col("is_active"))


@transform
def get_customer_count(
    customers: Annotated[pl.LazyFrame, Strict(Customer)],
) -> Annotated[pl.LazyFrame, Strict(CustomerCount)]:
    return customers.select(pl.len().cast(pl.Int64).alias("customer_count"))


customer_schema = Schema(Customer)

customers = customer_schema.create_dataframe(
    [
        Customer("1", "Alice", True),
        Customer("2", "Bob", False),
        Customer("3", "Charlie", True),
    ],
    pl.LazyFrame,
)

active_customer_count = get_customer_count(filter_active_customers(customers))
print(active_customer_count.collect())

Outputs:

shape: (1, 1)
┌────────────────┐
│ customer_count │
│ ---            │
│ i64            │
╞════════════════╡
│ 2              │
└────────────────┘

from dataclasses import dataclass
from typing import Annotated

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema, Strict, transform


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@dataclass
class CustomerCount:
    customer_count: int


@transform
def filter_active_customers(
    customers: Annotated[DataFrame, Strict(Customer)],
) -> Annotated[DataFrame, Strict(Customer)]:
    return customers.filter(customers.is_active)


@transform
def get_customer_count(
    customers: Annotated[DataFrame, Strict(Customer)],
) -> Annotated[DataFrame, Strict(CustomerCount)]:
    return customers.groupBy().count().withColumnRenamed("count", "customer_count")


spark = SparkSession.builder.getOrCreate()

customer_schema = Schema(Customer)

customers = customer_schema.create_dataframe(
    [
        Customer("1", "Alice", True),
        Customer("2", "Bob", False),
        Customer("3", "Charlie", True),
    ],
    DataFrame,
    session=spark,
)

active_customer_count = customers.transform(filter_active_customers).transform(
    get_customer_count
)
active_customer_count.show()

Outputs:

+--------------+
|customer_count|
+--------------+
|             2|
+--------------+

A schema mismatch in a transformation chain raises an error:

import traceback
from dataclasses import dataclass
from typing import Annotated

import polars as pl

from data_joinery import Schema, SchemaCoercionError, Strict, transform


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@dataclass
class CustomerCount:
    customer_count: int


@transform
def filter_active_customers(
    customers: Annotated[pl.LazyFrame, Strict(Customer)],
) -> Annotated[pl.LazyFrame, Strict(Customer)]:
    return customers.filter(pl.col("is_active"))


@transform
def get_customer_count(
    customers: Annotated[pl.LazyFrame, Strict(Customer)],
) -> Annotated[pl.LazyFrame, Strict(CustomerCount)]:
    return customers.select(pl.len().alias("count"))


customer_schema = Schema(Customer)

customers = customer_schema.create_dataframe(
    [
        Customer("1", "Alice", True),
        Customer("2", "Bob", False),
        Customer("3", "Charlie", True),
    ],
    pl.LazyFrame,
)
try:
    get_customer_count(filter_active_customers(customers))
except SchemaCoercionError:
    print(traceback.format_exc(limit=1))

Outputs:

Traceback (most recent call last):
  File "/home/runner/work/data-joinery/data-joinery/docs_src/learn/transform/index/usage_example_bad_polars.py", line 47, in <module>
    get_customer_count(filter_active_customers(customers))
data_joinery.schemas.coercion_error.SchemaCoercionError: Transform 'get_customer_count' return value, Schema 'CustomerCount': Cannot coerce DataFrame (mode='strict'):
  - Missing field 'customer_count': expected Int64. Add the field or update the model.
  - Additional field 'count': got UInt32. Remove it or add it to the model.

import traceback
from dataclasses import dataclass
from typing import Annotated

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema, SchemaCoercionError, Strict, transform


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@dataclass
class CustomerCount:
    customer_count: int


@transform
def filter_active_customers(
    customers: Annotated[DataFrame, Strict(Customer)],
) -> Annotated[DataFrame, Strict(Customer)]:
    return customers.filter(customers.is_active)


@transform
def get_customer_count(
    customers: Annotated[DataFrame, Strict(Customer)],
) -> Annotated[DataFrame, Strict(CustomerCount)]:
    return customers.groupBy().count()


spark = SparkSession.builder.getOrCreate()

customer_schema = Schema(Customer)

customers = customer_schema.create_dataframe(
    [
        Customer("1", "Alice", True),
        Customer("2", "Bob", False),
        Customer("3", "Charlie", True),
    ],
    DataFrame,
    session=spark,
)
try:
    customers.transform(filter_active_customers).transform(get_customer_count)
except SchemaCoercionError:
    print(traceback.format_exc(limit=1))

Outputs:

Traceback (most recent call last):
  File "/home/runner/work/data-joinery/data-joinery/docs_src/learn/transform/index/usage_example_bad_spark.py", line 50, in <module>
    customers.transform(filter_active_customers).transform(get_customer_count)
data_joinery.schemas.coercion_error.SchemaCoercionError: Transform 'get_customer_count' return value, Schema 'CustomerCount': Cannot coerce DataFrame (mode='strict'):
  - Missing field 'customer_count': expected LongType(). Add the field or update the model.
  - Additional field 'count': got LongType(). Remove it or add it to the model.

Warning

A PySpark transformation that performs an action such as collect() may process data before a later schema check raises SchemaCoercionError.

Lazy vs Eager

Polars supports two modes of execution, lazy and eager. You can use either mode with Data Joinery but lazy mode is strongly preferred as you will benefit from the performance optimisations Polars can make. Because both Polars and PySpark use lazy evaluation, the idea that you are 'running' transformations is a bit misleading. In reality you are just chaining up transformations and the intermediate results are not materialised.