Skip to content

Defining a Schema

Schemas wrap around Pydantic or Python dataclasses. They allow you to define the structure of your data in a backend agnostic way. The same schema can be used to define a schema for any backend you choose.

Dataclasses

Take the below example for a simple dataclass schema:

from dataclasses import dataclass

import polars as pl

from data_joinery import Schema


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


schema = Schema(Customer)
print(schema.native_schema(pl.Schema))

Outputs:

Schema({'customer_id': String, 'name': String, 'is_active': Boolean, 'discount_rate': Float64})

from dataclasses import dataclass

from pyspark.sql.types import StructType

from data_joinery import Schema


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


schema = Schema(Customer)
print(schema.native_schema(StructType).treeString())

Outputs:

root
 |-- customer_id: string (nullable = true)
 |-- name: string (nullable = true)
 |-- is_active: boolean (nullable = true)
 |-- discount_rate: double (nullable = true)

The method native_schema can be used to generate the schema for a specific backend, such as PySpark or Polars.

It doesn't matter what dataclass configuration you choose (frozen, slots, etc.) However if you decide to use dataclass inheritance it is best practice to use kw_only=True to avoid issues with field ordering:

from dataclasses import dataclass

import polars as pl

from data_joinery import Schema


@dataclass(kw_only=True)
class AddressMixin:
    number: str
    street: str
    postal_code: str


@dataclass(kw_only=True)
class Customer(AddressMixin):
    customer_id: str
    name: str


schema = Schema(Customer)
print(schema.native_schema(pl.Schema))

Outputs:

Schema({'number': String, 'street': String, 'postal_code': String, 'customer_id': String, 'name': String})

from dataclasses import dataclass

from pyspark.sql.types import StructType

from data_joinery import Schema


@dataclass(kw_only=True)
class AddressMixin:
    number: str
    street: str
    postal_code: str


@dataclass(kw_only=True)
class Customer(AddressMixin):
    customer_id: str
    name: str


schema = Schema(Customer)
print(schema.native_schema(StructType).treeString())

Outputs:

root
 |-- number: string (nullable = true)
 |-- street: string (nullable = true)
 |-- postal_code: string (nullable = true)
 |-- customer_id: string (nullable = true)
 |-- name: string (nullable = true)

The field order will follow standard dataclass rules

Pydantic models

The same example schema can be defined with Pydantic as well:

import polars as pl
from pydantic import BaseModel, Field

from data_joinery import Schema


class Customer(BaseModel):
    customer_id: str
    name: str
    is_active: bool
    discount_rate: float = Field(default=0.0, ge=0.0, le=1.0)


schema = Schema(Customer)
print(schema.native_schema(pl.Schema))

Outputs:

Schema({'customer_id': String, 'name': String, 'is_active': Boolean, 'discount_rate': Float64})

from pydantic import BaseModel, Field
from pyspark.sql.types import StructType

from data_joinery import Schema


class Customer(BaseModel):
    customer_id: str
    name: str
    is_active: bool
    discount_rate: float = Field(default=0.0, ge=0.0, le=1.0)


schema = Schema(Customer)
print(schema.native_schema(StructType).treeString())

Outputs:

root
 |-- customer_id: string (nullable = true)
 |-- name: string (nullable = true)
 |-- is_active: boolean (nullable = true)
 |-- discount_rate: double (nullable = true)
Pydantic field validation is not enforced

Data Joinery does not enforce Pydantic validation rules. It can be useful to include them in your models for documentation and for generating valid test fixtures, but only the data type is enforced.

Nested schemas

If your data contains nested structs or arrays, describe them with nested dataclasses or Pydantic models:

from dataclasses import dataclass

import polars as pl

from data_joinery import Schema


@dataclass
class Location:
    latitude: float
    longitude: float


@dataclass
class Employee:
    name: str
    phone_number: str


@dataclass
class Customer:
    customer_id: str
    name: str
    employees: list[Employee]
    location: Location


schema = Schema(Customer)
print(schema.native_schema(pl.Schema))

Outputs:

Schema({'customer_id': String, 'name': String, 'employees': List(Struct({'name': String, 'phone_number': String})), 'location': Struct({'latitude': Float64, 'longitude': Float64})})

from dataclasses import dataclass

from pyspark.sql.types import StructType

from data_joinery import Schema


@dataclass
class Location:
    latitude: float
    longitude: float


@dataclass
class Employee:
    name: str
    phone_number: str


@dataclass
class Customer:
    customer_id: str
    name: str
    employees: list[Employee]
    location: Location


schema = Schema(Customer)
print(schema.native_schema(StructType).treeString())

Outputs:

root
 |-- customer_id: string (nullable = true)
 |-- name: string (nullable = true)
 |-- employees: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- name: string (nullable = true)
 |    |    |-- phone_number: string (nullable = true)
 |-- location: struct (nullable = true)
 |    |-- latitude: double (nullable = true)
 |    |-- longitude: double (nullable = true)