Joinery Model
Example
Joinery Models are experimental
Dataclasses and Pydantic models can define a schema, but transformations still refer to column names as strings. For example, take this simple transformation:
from dataclasses import dataclass
from typing import Annotated
import polars as pl
from data_joinery import Schema, Strict, transform
@dataclass
class Order:
"""Input columns for an order."""
price: float
quantity: int
@dataclass
class OrderTotal:
"""Output column for the order total."""
total_price: float
@transform
def calculate_total(
orders: Annotated[pl.LazyFrame, Strict(Order)],
) -> Annotated[pl.LazyFrame, Strict(OrderTotal)]:
"""Multiply each order's price by its quantity."""
return orders.select((pl.col("price") * pl.col("quantity")).alias("total_price"))
orders = Schema(Order).create_dataframe(
[Order(price=12.5, quantity=2)],
pl.LazyFrame,
)
print(calculate_total(orders).collect())
Outputs:
from dataclasses import dataclass
from typing import Annotated
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql import functions as F
from data_joinery import Schema, Strict, transform
@dataclass
class Order:
"""Input columns for an order."""
price: float
quantity: int
@dataclass
class OrderTotal:
"""Output column for the order total."""
total_price: float
@transform
def calculate_total(
orders: Annotated[DataFrame, Strict(Order)],
) -> Annotated[DataFrame, Strict(OrderTotal)]:
"""Multiply each order's price by its quantity."""
return orders.select((F.col("price") * F.col("quantity")).alias("total_price"))
spark = SparkSession.builder.getOrCreate()
orders = Schema(Order).create_dataframe(
[Order(price=12.5, quantity=2)],
DataFrame,
session=spark,
)
calculate_total(orders).show()
Outputs:
This is fine but we don't get any type safety or autocomplete on the column names for our inputs and outputs. It means we have to keep switching between the model definitions and our transformations to help us remember what each column is called.
I think this is a common complaint when working with dataframes. This is solved by an extension
of dataclasses called a JoineryModel. These models behave similarly to standard dataclasses but
they retain fields as class attributes with string values:
from dataclasses import dataclass
from data_joinery.schemas.joinery_model import Column, JoineryModel, column
@dataclass
class DataclassOrder:
"""An order defined with a normal dataclass."""
price: float
quantity: int = 1
class JoineryOrder(JoineryModel):
"""An order with column names available on the class."""
price: Column[float]
quantity: Column[int] = column(default=1)
# A required dataclass field has no class attribute.
print("DataclassOrder has a price attribute:", hasattr(DataclassOrder, "price"))
print("DataclassOrder.quantity:", DataclassOrder.quantity)
# JoineryModel class attributes return column names, even for defaulted fields.
print("JoineryOrder.price:", JoineryOrder.price)
print("JoineryOrder.quantity:", JoineryOrder.quantity)
# Both instances return the stored values.
plain_order = DataclassOrder(price=12.5)
joinery_order = JoineryOrder(price=12.5)
print("Dataclass instance:", plain_order.price, plain_order.quantity)
print("JoineryModel instance:", joinery_order.price, joinery_order.quantity)
Outputs:
DataclassOrder has a price attribute: False
DataclassOrder.quantity: 1
JoineryOrder.price: price
JoineryOrder.quantity: quantity
Dataclass instance: 12.5 1
JoineryModel instance: 12.5 1
We can use this in our transformations. Our earlier example can now use Order.price
and Order.quantity. Another advantage of JoineryModel is that we can override the column
name with the column function. This is especially useful when we want to use a
column name that is not a valid Python identifier:
from typing import Annotated
import polars as pl
from data_joinery import Schema, Strict, transform
from data_joinery.schemas.joinery_model import Column, JoineryModel, column
class Order(JoineryModel):
"""Input columns for an order."""
price: Column[float]
quantity: Column[int]
class OrderTotal(JoineryModel):
"""Output column for the order total."""
total_price: Column[float] = column(name="Total Price")
@transform
def calculate_total(
orders: Annotated[pl.LazyFrame, Strict(Order)],
) -> Annotated[pl.LazyFrame, Strict(OrderTotal)]:
"""Multiply each order's price by its quantity."""
return orders.select(
(pl.col(Order.price) * pl.col(Order.quantity)).alias(OrderTotal.total_price)
)
orders = Schema(Order).create_dataframe(
[Order(price=12.5, quantity=2)],
pl.LazyFrame,
)
print(calculate_total(orders).collect())
Outputs:
from typing import Annotated
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql import functions as F
from data_joinery import Schema, Strict, transform
from data_joinery.schemas.joinery_model import Column, JoineryModel, column
class Order(JoineryModel):
"""Input columns for an order."""
price: Column[float]
quantity: Column[int]
class OrderTotal(JoineryModel):
"""Output column for the order total."""
total_price: Column[float] = column(name="Total Price")
@transform
def calculate_total(
orders: Annotated[DataFrame, Strict(Order)],
) -> Annotated[DataFrame, Strict(OrderTotal)]:
"""Multiply each order's price by its quantity."""
return orders.select(
(F.col(Order.price) * F.col(Order.quantity)).alias(OrderTotal.total_price)
)
spark = SparkSession.builder.getOrCreate()
orders = Schema(Order).create_dataframe(
[Order(price=12.5, quantity=2)],
DataFrame,
session=spark,
)
calculate_total(orders).show()
Outputs: