Skip to content

Coercing Schemas

Coercion in Data Joinery checks or converts an input DataFrame against a specified schema. There are several modes available to control how strict/relaxed the coercion process should be for your transformations.

Use Schema(Model).coerce_dataframe() to take an input DataFrame and return either a new DataFrame that matches the model's schema or raise a validation error if coercion fails based on the mode's strictness.

Let's take the following schema and walk through how each mode works:

from dataclasses import dataclass

from data_joinery import Schema


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)

Coercion Modes

Depending on your workflow and needs, you can choose from several coercion modes.

Strict (Default)

In strict mode, the input DataFrame must have exactly the same named fields and types as the specified model. Any extra columns, missing columns or type mismatches will result in an error. Field order and nullability are ignored, and the input field order is preserved. This DataFrame has an extra column so will fail the strict validation check:

import traceback
from dataclasses import dataclass

import polars as pl

from data_joinery import Schema, SchemaCoercionError


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class ModelWithExtraColumn:
    field1: str
    field2: str
    sub_schema: SubModel
    extra_field: str


input_schema = Schema(ModelWithExtraColumn)


input_df = input_schema.create_dataframe(
    [
        ModelWithExtraColumn(
            field1="value1",
            field2="value2",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
            extra_field="extra_value",
        )
    ],
    pl.LazyFrame,
)

try:
    schema.coerce_dataframe(input_df, mode="strict")
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/coercion/dataclass_coercion_strict_fail_extra_column_polars.py", line 49, in <module>
    schema.coerce_dataframe(input_df, mode="strict")
data_joinery.schemas.coercion_error.SchemaCoercionError: Schema 'Model': Cannot coerce DataFrame (mode='strict'):
  - Additional field 'extra_field': got String. Remove it or add it to the model.

import traceback
from dataclasses import dataclass

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema, SchemaCoercionError


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class ModelWithExtraColumn:
    field1: str
    field2: str
    sub_schema: SubModel
    extra_field: str


input_schema = Schema(ModelWithExtraColumn)


spark = SparkSession.builder.getOrCreate()

input_df = input_schema.create_dataframe(
    [
        ModelWithExtraColumn(
            field1="value1",
            field2="value2",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
            extra_field="extra_value",
        )
    ],
    DataFrame,
    session=spark,
)

try:
    schema.coerce_dataframe(input_df, mode="strict")
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/coercion/dataclass_coercion_strict_fail_extra_column_spark.py", line 52, in <module>
    schema.coerce_dataframe(input_df, mode="strict")
data_joinery.schemas.coercion_error.SchemaCoercionError: Schema 'Model': Cannot coerce DataFrame (mode='strict'):
  - Additional field 'extra_field': got StringType(). Remove it or add it to the model.

However if the Dataframe is equivalent it will pass and return the same DataFrame:

from dataclasses import dataclass

import polars as pl

from data_joinery import Schema


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class SubModel2:
    field2: str
    field1: str
    sub_schema: SubModel


input_schema = Schema(SubModel2)


input_df = input_schema.create_dataframe(
    [
        SubModel2(
            field1="value1",
            field2="value2",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
        )
    ],
    pl.LazyFrame,
)

coerced_df = schema.coerce_dataframe(input_df, mode="strict")
print(coerced_df.collect())

Outputs:

shape: (1, 3)
┌────────┬────────┬───────────────────────────┐
│ field2 ┆ field1 ┆ sub_schema                │
│ ---    ┆ ---    ┆ ---                       │
│ str    ┆ str    ┆ struct[2]                 │
╞════════╪════════╪═══════════════════════════╡
│ value2 ┆ value1 ┆ {"subvalue1","subvalue2"} │
└────────┴────────┴───────────────────────────┘

from dataclasses import dataclass

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class SubModel2:
    field2: str
    field1: str
    sub_schema: SubModel


input_schema = Schema(SubModel2)


spark = SparkSession.builder.getOrCreate()

input_df = input_schema.create_dataframe(
    [
        SubModel2(
            field1="value1",
            field2="value2",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
        )
    ],
    DataFrame,
    session=spark,
)

coerced_df = schema.coerce_dataframe(input_df, mode="strict")
coerced_df.show()

Outputs:

+------+------+--------------------+
|field2|field1|          sub_schema|
+------+------+--------------------+
|value2|value1|{subvalue1, subva...|
+------+------+--------------------+

Strict does not care about ordering of the fields.

Project Top Level

project_top_level projects only the model's top-level columns. If there is a difference in struct fields it will fail:

import traceback
from dataclasses import dataclass

import polars as pl

from data_joinery import Schema, SchemaCoercionError


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class SubModelWithExtraField:
    subfield1: str
    subfield2: str
    extra_field: str


@dataclass
class ModelWithExtraStructField:
    field1: str
    field2: str
    sub_schema: SubModelWithExtraField


input_schema = Schema(ModelWithExtraStructField)


input_df = input_schema.create_dataframe(
    [
        ModelWithExtraStructField(
            field1="value1",
            field2="value2",
            sub_schema=SubModelWithExtraField(
                subfield1="subvalue1", subfield2="subvalue2", extra_field="extra_value"
            ),
        )
    ],
    pl.LazyFrame,
)

try:
    schema.coerce_dataframe(input_df, mode="project_top_level")
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/coercion/dataclass_coercion_project_top_level_fail_extra_struct_field_polars.py", line 56, in <module>
    schema.coerce_dataframe(input_df, mode="project_top_level")
data_joinery.schemas.coercion_error.SchemaCoercionError: Schema 'Model': Cannot coerce DataFrame (mode='project_top_level'):
  - Additional field 'sub_schema.extra_field': got String. Remove it or use 'project' to drop nested extra fields.

import traceback
from dataclasses import dataclass

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema, SchemaCoercionError


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class SubModelWithExtraField:
    subfield1: str
    subfield2: str
    extra_field: str


@dataclass
class ModelWithExtraStructField:
    field1: str
    field2: str
    sub_schema: SubModelWithExtraField


input_schema = Schema(ModelWithExtraStructField)


spark = SparkSession.builder.getOrCreate()

input_df = input_schema.create_dataframe(
    [
        ModelWithExtraStructField(
            field1="value1",
            field2="value2",
            sub_schema=SubModelWithExtraField(
                subfield1="subvalue1", subfield2="subvalue2", extra_field="extra_value"
            ),
        )
    ],
    DataFrame,
    session=spark,
)

try:
    schema.coerce_dataframe(input_df, mode="project_top_level")
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/coercion/dataclass_coercion_project_top_level_fail_extra_struct_field_spark.py", line 59, in <module>
    schema.coerce_dataframe(input_df, mode="project_top_level")
data_joinery.schemas.coercion_error.SchemaCoercionError: Schema 'Model': Cannot coerce DataFrame (mode='project_top_level'):
  - Additional field 'sub_schema.extra_field': got StringType(). Remove it or use 'project' to drop nested extra fields.

Project

The mode project projects nested struct fields, removing all fields that are not in the model:

from dataclasses import dataclass

import polars as pl

from data_joinery import Schema


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class SubModelWithExtraField:
    subfield1: str
    subfield2: str
    extra_field: str


@dataclass
class ModelWithExtraStructField:
    field1: str
    field2: str
    sub_schema: SubModelWithExtraField


input_schema = Schema(ModelWithExtraStructField)


input_df = input_schema.create_dataframe(
    [
        ModelWithExtraStructField(
            field1="value1",
            field2="value2",
            sub_schema=SubModelWithExtraField(
                subfield1="subvalue1", subfield2="subvalue2", extra_field="extra_value"
            ),
        )
    ],
    pl.LazyFrame,
)

coerced_df = schema.coerce_dataframe(input_df, mode="project")
print(coerced_df.collect())

Outputs:

shape: (1, 3)
┌────────┬────────┬───────────────────────────┐
│ field1 ┆ field2 ┆ sub_schema                │
│ ---    ┆ ---    ┆ ---                       │
│ str    ┆ str    ┆ struct[2]                 │
╞════════╪════════╪═══════════════════════════╡
│ value1 ┆ value2 ┆ {"subvalue1","subvalue2"} │
└────────┴────────┴───────────────────────────┘

from dataclasses import dataclass

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class SubModelWithExtraField:
    subfield1: str
    subfield2: str
    extra_field: str


@dataclass
class ModelWithExtraStructField:
    field1: str
    field2: str
    sub_schema: SubModelWithExtraField


input_schema = Schema(ModelWithExtraStructField)


spark = SparkSession.builder.getOrCreate()

input_df = input_schema.create_dataframe(
    [
        ModelWithExtraStructField(
            field1="value1",
            field2="value2",
            sub_schema=SubModelWithExtraField(
                subfield1="subvalue1", subfield2="subvalue2", extra_field="extra_value"
            ),
        )
    ],
    DataFrame,
    session=spark,
)

coerced_df = schema.coerce_dataframe(input_df, mode="project")
coerced_df.show(truncate=50)

Outputs:

+------+------+----------------------+
|field1|field2|            sub_schema|
+------+------+----------------------+
|value1|value2|{subvalue1, subvalue2}|
+------+------+----------------------+

Project will still fail if there are missing columns:

import traceback
from dataclasses import dataclass

import polars as pl

from data_joinery import Schema, SchemaCoercionError


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class ModelWithMissingColumn:
    field1: str
    sub_schema: SubModel


input_schema = Schema(ModelWithMissingColumn)


input_df = input_schema.create_dataframe(
    [
        ModelWithMissingColumn(
            field1="value1",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
        )
    ],
    pl.LazyFrame,
)

try:
    schema.coerce_dataframe(input_df, mode="project")
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/coercion/dataclass_coercion_project_fail_missing_columns_polars.py", line 45, in <module>
    schema.coerce_dataframe(input_df, mode="project")
data_joinery.schemas.coercion_error.SchemaCoercionError: Schema 'Model': Cannot coerce DataFrame (mode='project'):
  - Missing field 'field2': expected String. Add the field or update the model.

import traceback
from dataclasses import dataclass

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema, SchemaCoercionError


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class ModelWithMissingColumn:
    field1: str
    sub_schema: SubModel


input_schema = Schema(ModelWithMissingColumn)


spark = SparkSession.builder.getOrCreate()

input_df = input_schema.create_dataframe(
    [
        ModelWithMissingColumn(
            field1="value1",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
        )
    ],
    DataFrame,
    session=spark,
)

try:
    schema.coerce_dataframe(input_df, mode="project")
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/coercion/dataclass_coercion_project_fail_missing_columns_spark.py", line 48, in <module>
    schema.coerce_dataframe(input_df, mode="project")
data_joinery.schemas.coercion_error.SchemaCoercionError: Schema 'Model': Cannot coerce DataFrame (mode='project'):
  - Missing field 'field2': expected StringType(). Add the field or update the model.

Project Cast

The most relaxed mode, project_cast, recursively projects fields and casts values to the model's field types. This will still fail if the backend cannot cast a type safely, but this can be useful if reading external data and you want to easily align your DataFrame with the model's schema. It will also still fail if there are missing columns:

from dataclasses import dataclass

import polars as pl

from data_joinery import Schema


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class ModelWithExtraFieldAndDifferentType:
    field1: int
    field2: str
    sub_schema: SubModel


input_schema = Schema(ModelWithExtraFieldAndDifferentType)


input_df = input_schema.create_dataframe(
    [
        ModelWithExtraFieldAndDifferentType(
            field1=123,
            field2="value2",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
        )
    ],
    pl.LazyFrame,
)

coerced_df = schema.coerce_dataframe(input_df, mode="project_cast")
print(coerced_df.collect())

Outputs:

shape: (1, 3)
┌────────┬────────┬───────────────────────────┐
│ field1 ┆ field2 ┆ sub_schema                │
│ ---    ┆ ---    ┆ ---                       │
│ str    ┆ str    ┆ struct[2]                 │
╞════════╪════════╪═══════════════════════════╡
│ 123    ┆ value2 ┆ {"subvalue1","subvalue2"} │
└────────┴────────┴───────────────────────────┘

from dataclasses import dataclass

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema


@dataclass
class SubModel:
    subfield1: str
    subfield2: str


@dataclass
class Model:
    field1: str
    field2: str
    sub_schema: SubModel


schema = Schema(Model)


@dataclass
class ModelWithExtraFieldAndDifferentType:
    field1: int
    field2: str
    sub_schema: SubModel


input_schema = Schema(ModelWithExtraFieldAndDifferentType)


spark = SparkSession.builder.getOrCreate()

input_df = input_schema.create_dataframe(
    [
        ModelWithExtraFieldAndDifferentType(
            field1=123,
            field2="value2",
            sub_schema=SubModel(subfield1="subvalue1", subfield2="subvalue2"),
        )
    ],
    DataFrame,
    session=spark,
)

coerced_df = schema.coerce_dataframe(input_df, mode="project_cast")
coerced_df.show(truncate=50)

Outputs:

+------+------+----------------------+
|field1|field2|            sub_schema|
+------+------+----------------------+
|   123|value2|{subvalue1, subvalue2}|
+------+------+----------------------+

Type casting in Data Joinery follows the backend's casting rules. If the cast is permitted then a conversion attempt will happen to try to change the data type. This could still raise a runtime error if the cast is not possible for a specific value in your dataframe.

Schema mismatches such as missing fields fail during coercion. In project_cast mode, a value that cannot be cast may fail later, when Polars runs the query.