Overriding types

Types can be overridden for a specific backend by annotating the fields in your models. If you need to override the default type for multiple backends you can just annotate more than once. Take this example schema where we override the default types for a python int. Only if the annotation is present for the backend does the type get overridden. The tabs show backend-specific overrides for the same Python field types:

from dataclasses import dataclass
from typing import Annotated

import polars as pl

from data_joinery import Schema


@dataclass
class AnnotatedSchema:
    number: int
    short_number: Annotated[int, pl.Int8]
    another_short_number: Annotated[int, pl.Int16]


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

Schema({'number': Int64, 'short_number': Int8, 'another_short_number': Int16})

from dataclasses import dataclass
from typing import Annotated

from pyspark.sql.types import ShortType, StructType

from data_joinery import Schema


@dataclass
class AnnotatedSchema:
    number: int
    short_number: Annotated[int, ShortType()]
    another_short_number: Annotated[int, ShortType()]


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

root
 |-- number: long (nullable = true)
 |-- short_number: short (nullable = true)
 |-- another_short_number: short (nullable = true)

Backend specific type behaviour

It isn't a perfect mapping between Python types and backend-specific types, and there may not be equivalent types between Polars and PySpark as well. The backends offer different integer widths, so converting a frame between them may change the available precision or raise an overflow error.

Spark

By default Data Joinery will use the documented conversions to map Python types to Spark types. However, this might not be suitable for all use cases for example if you have a DecimalType or VarcharType. These types cannot be expressed in normal Python so you have to use Annotated to override the type in the schema.

There are also some Python types that can map to multiple spark types and this module has made a choice. For example int -> LongType and float -> DoubleType. For a full list of default mappings see default type mappings.

The below example shows the effect of annotating your fields has on the resulting schema:

from dataclasses import dataclass
from typing import Annotated

from pyspark.sql.types import (
    ByteType,
    DoubleType,
    FloatType,
    IntegerType,
    LongType,
    ShortType,
    StructType,
)

from data_joinery import Schema


@dataclass
class AnnotatedSchema:
    my_byte: Annotated[int, ByteType()]
    my_short: Annotated[int, ShortType()]
    my_int: Annotated[int, IntegerType()]
    my_long: Annotated[int, LongType()]
    my_float: Annotated[float, FloatType()]
    my_double: Annotated[float, DoubleType()]


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

Outputs:

root
 |-- my_byte: byte (nullable = true)
 |-- my_short: short (nullable = true)
 |-- my_int: integer (nullable = true)
 |-- my_long: long (nullable = true)
 |-- my_float: float (nullable = true)
 |-- my_double: double (nullable = true)

Polars

Polars type mapping can be found in the default type mappings reference..