### Install and Activate Development Environment

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/contributing.md

Installs Python 3.11, creates a virtual environment named 'typedspark', activates it, and installs project dependencies.

```bash
pyenv install 3.11
pyenv virtualenv 3.11 typedspark
pyenv activate typedspark
pip install -r requirements.txt
pip install -r requirements-dev.txt
```

--------------------------------

### Install Typedspark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/index.md

Install the typedspark library using pip. This command installs the base package.

```bash
pip install typedspark
```

--------------------------------

### Install Pre-commit Hook

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/contributing.md

Installs the pre-commit hook to enforce code quality checks before committing.

```bash
pre-commit install
```

--------------------------------

### Initialize SparkSession

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/autocomplete_in_notebooks.ipynb

Sets up a SparkSession with progress display disabled and log level set to ERROR. This is a common setup for Spark applications.

```python
from pyspark.sql import SparkSession

spark = SparkSession.Builder().config("spark.ui.showConsoleProgress", "false").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")
```

--------------------------------

### Install Typedspark with PySpark Dependency

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/index.md

Install typedspark along with its PySpark dependency. Use this if your environment does not have PySpark pre-installed.

```bash
pip install "typedspark[pyspark]"
```

--------------------------------

### Unit Test Example with TypedSpark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_empty_datasets.ipynb

An example of a full unit test using `typedspark` for DataFrame transformations. It utilizes `create_partially_filled_dataset` to set up test data and `assert_df_equality` from `chispa` for validation.

```python
from pyspark.sql import SparkSession
from pyspark.sql.types import LongType, StringType
from typedspark import Column, DataSet, Schema, create_partially_filled_dataset, transform_to_schema
from chispa.dataframe_comparer import assert_df_equality


class Person(Schema):
    name: Column[StringType]
    age: Column[LongType]


def birthday(df: DataSet[Person]) -> DataSet[Person]:
    return transform_to_schema(df, Person, {Person.age: Person.age + 1})

def test_birthday(spark: SparkSession):
    df = create_partially_filled_dataset(
        spark,
        Person,
        {
            Person.name: ["Alice", "Bob"],
            Person.age: [20, 30],
        },
    )

    observed = birthday(df)
    expected = create_partially_filled_dataset(
        spark,
        Person,
        {
            Person.name: ["Alice", "Bob"],
            Person.age: [21, 31],
        },
    )

    assert_df_equality(observed, expected, ignore_row_order=True, ignore_nullable=True)
```

--------------------------------

### Linting with Multiple Schemas and Type Checking

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/type_checking.ipynb

Shows how linting can catch type mismatches between expected and actual DataSet types in function arguments. Includes examples of correct and incorrect type usage.

```python
class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]
    age: Column[LongType]


class Address(Schema):
    street: Column[StringType]
    number: Column[LongType]


def birthday(df: DataSet[Person]) -> DataSet[Person]:
    return DataSet[Person](
        df.withColumn(Person.age.str, Person.age + 1),
    )


df_1 = DataSet[Person](
    spark.createDataFrame(
        pd.DataFrame(
            dict(
                id=[1, 2, 3],
                name=["John", "Jane", "Jack"],
                age=[20, 30, 40],
            )
        )
    )
)
# no linting error
birthday(df_1)

df_2 = DataSet[Address](
    spark.createDataFrame(
        pd.DataFrame(
            dict(
                street=["Lynton Walk", "Canada Square", "Chapelside Avenue"],
                number=[1, 2, 3],
            )
        )
    )
)
try:
    # linting error: expected DataSet[Person], observed DataSet[Address]
    birthday(df_2)
except:
    pass

df_3 = spark.createDataFrame(
    pd.DataFrame(
        dict(
            id=[1, 2, 3],
            name=["John", "Jane", "Jack"],
            age=[20, 30, 40],
        )
    )
)
try:
    # linting error: expected DataSet[Person], observed DataFrame
    birthday(df_3)
except:
    pass
```

--------------------------------

### Handle Ambiguous Columns in transform_to_schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/column_ambiguity.ipynb

Illustrates an error encountered when using `transform_to_schema` with ambiguous columns from a join. The error message guides on how to explicitly map columns.

```python
from typedspark import transform_to_schema


class PersonWithJob(Schema):
    id: Column[IntegerType]
    name: Column[StringType]
    age: Column[IntegerType]
    salary: Column[IntegerType]


try:
    (
        transform_to_schema(
            df_a.join(
                df_b,
                person.id == job.id,
            ),
            PersonWithJob,
        ).show()
    )
except ValueError as e:
    print(e)
```

--------------------------------

### Load DataSet and Schema with underscore table name

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Loads a DataSet and Schema from a table whose name starts with an underscore. TypedSpark renames the attribute to 'u_person' to enable autocompletion.

```python
db = Database()

persons, Person = db.u_person()
```

--------------------------------

### Save DataFrame to a table with underscore prefix

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Saves the DataFrame 'df' as a table named '_person' in the default database. Table names starting with an underscore can cause issues with autocompletion.

```python
df.write.saveAsTable("default._person")
```

--------------------------------

### Initialize SparkSession and create databases

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Initializes a SparkSession and creates necessary databases if they do not already exist.

```python
spark = SparkSession.builder.getOrCreate()
spark.sql("CREATE DATABASE IF NOT EXISTS vet")
spark.sql("CREATE DATABASE IF NOT EXISTS library")
spark.sql("CREATE DATABASE IF NOT EXISTS store")
```

--------------------------------

### Build Documentation Locally

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/contributing.md

Cleans existing documentation and builds the HTML documentation locally within the 'docs/' directory.

```bash
cd docs/; make clean; make html; cd ..
```

--------------------------------

### Create and Overwrite Tables from Empty Datasets

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Creates Parquet tables for 'owners', 'pets', and 'appointments' in the 'vet' database, overwriting existing ones.

```python
create_empty_dataset(spark, Owners).write.saveAsTable(
    "vet.owners", format="parquet", mode="overwrite"
)
create_empty_dataset(spark, Pets).write.saveAsTable("vet.pets", format="parquet", mode="overwrite")
create_empty_dataset(spark, Appointments).write.saveAsTable(
    "vet.appointments", format="parquet", mode="overwrite"
)
```

--------------------------------

### Create Partially Filled Dataset with Schema A

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/advanced_linting_support.ipynb

Demonstrates creating a TypedSpark DataSet with a defined schema 'A' and populating it with data.

```python
from typedspark import Column, Schema, DataSet, create_partially_filled_dataset
from pyspark.sql.types import StringType


class A(Schema):
    a: Column[StringType]


df = create_partially_filled_dataset(
    spark,
    A,
    {
        A.a: ["a", "b", "c"],
    },
)
res = df.filter(A.a == "a")
```

--------------------------------

### Run Documentation Notebooks

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/contributing.md

Executes all documentation notebooks in the 'docs/' directory and strips metadata for manageable PR diffs.

```bash
sh docs/run_notebooks.sh
```

--------------------------------

### Load DataSet and Schema using Databases

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Loads a DataSet and its corresponding Schema from the 'person_table' in the default database using the Databases instance.

```python
persons, Person = db.default.person_table()
```

--------------------------------

### Import Catalogs and create SparkSession

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Imports the Catalogs class and initializes a SparkSession, which is necessary for catalog operations.

```python
from typedspark import Catalogs, create_schema
from pyspark.sql import SparkSession
import pyspark.sql.functions as F

spark = SparkSession.builder.getOrCreate()
```

--------------------------------

### Instantiate Catalogs

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Creates an instance of the Catalogs class, which provides access to Spark catalog metadata.

```python
db = Catalogs(spark)
```

--------------------------------

### Load DataSet and Schema using Catalogs

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Loads a DataSet and its corresponding Schema from the 'person_table' in the default catalog using the Catalogs instance. Autocompletion for table and column names is available.

```python
persons, Person = db.spark_catalog.default.person_table()
```

--------------------------------

### Create Partially Filled Dataset (Column-wise)

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_empty_datasets.ipynb

Shows how to create a DataFrame with some columns pre-filled and others left as NULL, using a column-wise definition. This is useful when you have many columns but few rows.

```python
from typedspark import Column, Schema, create_partially_filled_dataset
from pyspark.sql.types import LongType, StringType


class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]
    age: Column[LongType]

df_partially_filled = create_partially_filled_dataset(
    spark,
    Person,
    {
        Person.id: [1, 2, 3],
        Person.name: ["John", "Jane", "Jack"],
    },
)
df_partially_filled.show()
```

--------------------------------

### Create Empty Dataset from Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_empty_datasets.ipynb

Demonstrates creating an entirely empty DataFrame based on a defined `Person` schema using `create_empty_dataset`. All columns will contain NULL values.

```python
from typedspark import Column, Schema, create_empty_dataset
from pyspark.sql.types import LongType, StringType


class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]
    age: Column[LongType]


df_empty = create_empty_dataset(spark, Person)
df_empty.show()
```

--------------------------------

### Load DataSet and Schema using Database

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Loads a DataSet and its corresponding Schema from the 'person_table' in the default database using the Database instance. The .load() method is used here.

```python
person, Person = db.person_table.load()
```

--------------------------------

### Generate Schema with documentation

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Prints the schema of the 'Person' class, including documentation and metadata. This is useful for migrating existing tables to TypedSpark and generating documentation.

```python
Person.print_schema(include_documentation=True)
```

--------------------------------

### Create and Populate Vaccinations Table

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Creates the 'vet.vaccinations' table using a partially filled dataset with vaccination details, overwriting if it exists.

```python
date = datetime(2023, 10, 2)

create_partially_filled_dataset(
    spark,
    Vaccinations,
    {
        Vaccinations.vaccination_id: [1, 2, 3, 4, 5, 6, 7],
        Vaccinations.pet_id: [1, 2, 3, 1, 3, 2, 3],
        Vaccinations.vaccine_name: [
            "rabies",
            "rabies",
            "rabies",
            "lyme",
            "lyme",
            "influenza",
            "influenza",
        ],
        Vaccinations.next_due_date: [
            date + timedelta(days=32),
            date + timedelta(days=6),
            date + timedelta(days=12),
            date + timedelta(days=15),
            date + timedelta(days=2),
            date + timedelta(days=1),
            date + timedelta(days=3),
        ],
        Vaccinations.vaccine_date: [
            date + timedelta(days=32) - timedelta(days=365),
            date + timedelta(days=6) - timedelta(days=365),
            date + timedelta(days=12) - timedelta(days=365),
            date + timedelta(days=15) - timedelta(days=365),
            date + timedelta(days=2) - timedelta(days=365),
            date + timedelta(days=1) - timedelta(days=365),
            date + timedelta(days=3) - timedelta(days=365),
        ],
    },
).write.saveAsTable("vet.vaccinations", format="parquet", mode="overwrite")
```

--------------------------------

### Runtime Type Checking with Valid Data

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/type_checking.ipynb

Demonstrates successful initialization of a DataSet with data that conforms to the defined schema. No errors are raised.

```python
import pandas as pd
from typedspark import Column, DataSet, Schema
from pyspark.sql.types import LongType, StringType


class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]
    age: Column[LongType]


df = spark.createDataFrame(
    pd.DataFrame(
        dict(
            id=[1, 2, 3],
            name=["John", "Jane", "Jack"],
            age=[20, 30, 40],
        )
    )
)
# no errors raised
df = DataSet[Person](df)
df.show()
```

--------------------------------

### Initialize TypedSpark Databases

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Initializes the Databases class from TypedSpark, optionally specifying a catalog name. This is an alternative to Catalogs for loading DataSets and Schemas.

```python
from typedspark import Databases

db = Databases()
```

--------------------------------

### Create and Load Dataset with StructType

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/structtypes_in_notebooks.ipynb

Creates a partially filled dataset with a nested StructType column and registers it as a temporary view. Then, it loads the table using TypedSpark for type-safe access.

```python
create_partially_filled_dataset(
    spark,
    Container,
    {
        Container.values: create_partially_filled_dataset(
            spark,
            Values,
            {Values.a: [1, 2, 3]},
        ).collect(),
    },
).createOrReplaceTempView("structtype_table")

container, ContainerSchema = load_table(spark, "structtype_table", "Container")
```

--------------------------------

### Create Partially Filled Dataset (Row-wise)

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_empty_datasets.ipynb

Demonstrates creating a DataFrame with specified values for certain columns in a row-wise fashion. This is beneficial when you have few columns but many rows.

```python
from typedspark import Column, Schema, create_partially_filled_dataset
from pyspark.sql.types import LongType, StringType


class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]
    age: Column[LongType]

create_partially_filled_dataset(
    spark,
    Person,
    [
        {Person.name: "Alice", Person.age: 20},
        {Person.name: "Bob", Person.age: 30},
        {Person.name: "Charlie", Person.age: 40},
        {Person.name: "Dave", Person.age: 50},
        {Person.name: "Eve", Person.age: 60},
        {Person.name: "Frank", Person.age: 70},
        {Person.name: "Grace", Person.age: 80},
    ],
).show()
```

--------------------------------

### Basic DataFrame Transformation to Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/transforming_datasets.ipynb

Demonstrates a standard pattern for transforming a DataFrame to a new schema using explicit joins, filters, and column selections.

```python
from pyspark.sql.types import IntegerType, StringType
from typedspark import Column, Schema, DataSet


class Person(Schema):
    name: Column[StringType]
    job_id: Column[IntegerType]


class Job(Schema):
    id: Column[IntegerType]
    function: Column[StringType]
    hourly_rate: Column[IntegerType]


class PersonWithJob(Person, Job):
    id: Column[IntegerType]
    name: Column[StringType]
    job_name: Column[StringType]
    rate: Column[IntegerType]


def get_plumbers(persons: DataSet[Person], jobs: DataSet[Job]) -> DataSet[PersonWithJob]:
    return DataSet[PersonWithJob](
        jobs.filter(Job.function == "plumber")
        .join(persons, Job.id == Person.job_id)
        .withColumn(PersonWithJob.job_name.str, Job.function)
        .withColumn(PersonWithJob.rate.str, Job.hourly_rate)
        .select(*PersonWithJob.all_column_names())
    )
```

--------------------------------

### Applying Birthday Transformation with DataSetImplements

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/dataset_implements.ipynb

Demonstrates applying the generic 'birthday' function (defined using DataSetImplements) to 'person' and 'pet' datasets. It also shows how a type error is raised when attempting to apply it to a 'fruit' dataset, which lacks the 'age' column.

```python
# returns a DataSet[Person]
happy_person = birthday(person)

# returns a DataSet[Pet]
happy_pet = birthday(pet)

try:
    # Raises a linting error:
    # Argument of type "DataSet[Fruit]" cannot be assigned to
    # parameter "df" of type "DataSetImplements[Age, T@birthday]"
    birthday(fruit)
except Exception as e:
    pass
```

--------------------------------

### Creating a Partially Filled Dataset

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/transforming_datasets.ipynb

Creates a DataFrame with some columns pre-filled and others left null, using `create_partially_filled_dataset`.

```python
from typedspark import create_partially_filled_dataset

df = create_partially_filled_dataset(spark, Job, {Job.hourly_rate: [10, 20, 30]})
```

--------------------------------

### Define Schemas and Create Empty Datasets

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/dataset_implements.ipynb

Defines several Pydantic-like schemas (Person, Pet, Fruit) using TypedSpark's Schema and Column types, and then creates empty datasets for each schema. This sets up the data structures for subsequent transformations.

```python
from pyspark.sql.types import LongType, StringType
from typedspark import (
    Schema,
    Column,
    create_empty_dataset,
)


class Person(Schema):
    name: Column[StringType]
    age: Column[LongType]
    job: Column[StringType]


class Pet(Schema):
    name: Column[StringType]
    age: Column[LongType]
    type: Column[StringType]


class Fruit(Schema):
    type: Column[StringType]


person = create_empty_dataset(spark, Person)
pet = create_empty_dataset(spark, Pet)
fruit = create_empty_dataset(spark, Fruit)
```

--------------------------------

### Register Schema with Alias for Self-Joins

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/column_ambiguity.ipynb

Demonstrates how to handle self-joins by registering the same schema to a DataFrame multiple times with different aliases using `register_schema_to_dataset_with_alias`. This prevents ambiguity when joining a DataFrame with itself.

```python
from typedspark import register_schema_to_dataset_with_alias

df = create_partially_filled_dataset(
    spark,
    Person,
    {
        Person.id: [1, 2, 3],
        Person.name: ["Alice", "Bob", "Charlie"],
        Person.age: [20, 30, 40],
    },
)

df_a, person_a = register_schema_to_dataset_with_alias(df, Person, alias="a")
df_b, person_b = register_schema_to_dataset_with_alias(df, Person, alias="b")

df_a.join(df_b, person_a.id == person_b.id).show()
```

--------------------------------

### Extending a Schema with New Columns

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/subclassing_schemas.ipynb

Demonstrates how to subclass an existing schema (Person) to add new columns (age) for pipeline stages. Use when a function adds specific columns to an existing dataset.

```python
from typedspark import Column, Schema, DataSet
from pyspark.sql.types import LongType, StringType
from pyspark.sql.functions import lit


class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]


class PersonWithAge(Person):
    age: Column[LongType]


def foo(df: DataSet[Person]) -> DataSet[PersonWithAge]:
    return DataSet[PersonWithAge](
        df.withColumn(PersonWithAge.age, lit(42)),
    )
```

--------------------------------

### Delta Live Table Definition with TypedSpark Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/documentation.ipynb

Integrate TypedSpark schemas with Databricks Delta Live Tables to automatically generate table documentation in the UI.

```python
@dlt.table(**Person.get_dlt_kwargs())
def table_definition() -> DataSet[Person]:
    # your table definition here
```

--------------------------------

### Define and Create Sample DataFrame

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_schema_in_notebook.ipynb

Defines a `Vaccinations` schema using TypedSpark and creates a sample DataFrame. This DataFrame is then used to demonstrate schema generation.

```python
from datetime import timedelta, datetime
from pyspark.sql.functions import first
from pyspark.sql.types import LongType, StringType, DateType
from typedspark import Column, Schema, create_partially_filled_dataset, create_schema

date = datetime(2023, 10, 2)


class Vaccinations(Schema):
    vaccination_id: Column[LongType]
    pet_id: Column[LongType]
    vaccine_name: Column[StringType]
    vaccine_date: Column[DateType]
    next_due_date: Column[DateType]


vaccinations = create_partially_filled_dataset(
    spark,
    Vaccinations,
    {
        Vaccinations.vaccination_id: [1, 2, 3, 4, 5, 6, 7],
        Vaccinations.pet_id: [1, 2, 3, 1, 3, 2, 3],
        Vaccinations.vaccine_name: [
            "rabies",
            "rabies",
            "rabies",
            "lyme",
            "lyme",
            "influenza",
            "influenza",
        ],
        Vaccinations.next_due_date: [
            date + timedelta(days=32),
            date + timedelta(days=6),
            date + timedelta(days=12),
            date + timedelta(days=15),
            date + timedelta(days=2),
            date + timedelta(days=1),
            date + timedelta(days=3),
        ],
        Vaccinations.vaccine_date: [
            date + timedelta(days=32) - timedelta(days=365),
            date + timedelta(days=6) - timedelta(days=365),
            date + timedelta(days=12) - timedelta(days=365),
            date + timedelta(days=15) - timedelta(days=365),
            date + timedelta(days=2) - timedelta(days=365),
            date + timedelta(days=1) - timedelta(days=365),
            date + timedelta(days=3) - timedelta(days=365),
        ],
    },
)
vaccinations.show()
```

--------------------------------

### Initialize TypedSpark Catalogs

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Initializes the Catalogs class from TypedSpark. This class is used to load DataSets and Schemas from Spark's catalog.

```python
from typedspark import Catalogs

db = Catalogs()
```

--------------------------------

### Handling Duplicate Keys in transform_to_schema Transformations

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/transforming_datasets.ipynb

Demonstrates how attempting to use duplicate keys in the transformations dictionary for `transform_to_schema` raises a ValueError. This highlights the requirement for unique column names in the transformations.

```python
try:
    transform_to_schema(
        df,
        Job,
        {
            Job.hourly_rate: Job.hourly_rate + 3,
            Job.hourly_rate: Job.hourly_rate * 2,
        },
    )
except ValueError as e:
    print(e)
```

--------------------------------

### Initialize TypedSpark Database

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Initializes the Database class from TypedSpark, optionally specifying a database name. This class is used for loading tables from a single database.

```python
from typedspark import Database

db = Database()
```

--------------------------------

### Create DataFrame and Register Temp View

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/autocomplete_in_notebooks.ipynb

Creates a Spark DataFrame from a Pandas DataFrame and registers it as a temporary view named 'person_table'. This prepares data for subsequent operations.

```python
import pandas as pd

( \
    spark.createDataFrame(
        pd.DataFrame(
            dict(
                name=["Jack", "John", "Jane"],
                age=[20, 30, 40],
            )
        )
    ).createOrReplaceTempView("person_table")
)
```

--------------------------------

### Load Single Dataset with Default Schema Name

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Loads a single dataset using `load_table` with the default schema naming convention. This is useful when you only need to load one table and its schema can be inferred or doesn't require a specific name.

```python
from typedspark import load_table

person, Person = load_table(spark, "person_table")
Person
```

```python
from pyspark.sql.types import LongType, StringType

from typedspark import Column, Schema


class DynamicallyLoadedSchema(Schema):
    name: Column[StringType]
    age: Column[LongType]
```

--------------------------------

### Transform DataFrame to a New Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/advanced_linting_support.ipynb

Demonstrates using the transform function with transform_to_schema to apply a transformation and change the DataFrame's schema from A to B.

```python
from typedspark import transform_to_schema
from pyspark.sql.functions import lit


class B(A):
    b: Column[StringType]


def foo(df: DataSet[A]) -> DataSet[A]:
    return transform_to_schema(
        df,
        B,
        {
            B.b: lit("hi"),
        },
    )


res = create_partially_filled_dataset(
    spark,
    A,
    {
        A.a: ["a", "b", "c"],
    },
).transform(foo)
```

--------------------------------

### Generate DataSet with StructType Column using create_structtype_row

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/structtype_columns.ipynb

Generates a `DataSet` with `StructType` columns by explicitly creating each nested row using `create_structtype_row`. This provides fine-grained control over the structure of each row.

```python
from typedspark import create_structtype_row

create_partially_filled_dataset(
    spark,
    Actions,
    [
        {
            Actions.consequeces: create_structtype_row(
                Values, {Values.name: "a", Values.severity: 1}
            ),
        },
        {
            Actions.consequeces: create_structtype_row(
                Values, {Values.name: "b", Values.severity: 2}
            ),
        },
        {
            Actions.consequeces: create_structtype_row(
                Values, {Values.name: "c", Values.severity: 3}
            ),
        },
    ],
).show()
```

--------------------------------

### Define Schemas for Vet Database Entities

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Defines Pydantic-like schemas for 'Appointments', 'Pets', 'Vaccinations', and 'Owners' using TypedSpark's Column and Schema.

```python
class Appointments(Schema):
    appointment_id: Column[LongType]
    pet_id: Column[LongType]
    appointment_date: Column[DateType]
    appointment_reason: Column[StringType]
    veterinarian_name: Column[StringType]
    notes: Column[StringType]


class Pets(Schema):
    pet_id: Column[LongType]
    owner_id: Column[LongType]
    pet_name: Column[StringType]
    species: Column[StringType]
    breed: Column[StringType]
    age: Column[LongType]
    birthdate: Column[DateType]
    gender: Column[StringType]


class Vaccinations(Schema):
    vaccination_id: Column[LongType]
    pet_id: Column[LongType]
    vaccine_name: Column[StringType]
    vaccine_date: Column[DateType]
    next_due_date: Column[DateType]


class Owners(Schema):
    owner_id: Column[LongType]
    first_name: Column[StringType]
    last_name: Column[StringType]
    email: Column[StringType]
    phone_number: Column[StringType]
    address: Column[StringType]
```

--------------------------------

### Create Dynamic Schema with TypedSpark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Creates a TypedSpark schema dynamically from a pivoted DataFrame. This allows for type-safe access to DataFrame columns.

```python
pivot, Pivot = create_schema(pivot)
```

```python
Pivot
```

--------------------------------

### Define Spark Schema with Column Comments

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/documentation.ipynb

Define a Spark schema with detailed column metadata and comments using TypedSpark's Column and ColumnMeta.

```python
from typing import Annotated
from pyspark.sql.types import DateType, StringType
from typedspark import Column, ColumnMeta, Schema


class Person(Schema):
    """Dimension table that contains information about a person."""

    person_id: Annotated[
        Column[StringType],
        ColumnMeta(comment="Unique person id"),
    ]
    gender: Annotated[Column[StringType], ColumnMeta(comment="Gender of the person")]
    birthdate: Annotated[Column[DateType], ColumnMeta(comment="Date of birth of the person")]
    job_id: Annotated[Column[StringType], ColumnMeta(comment="Id of the job")]
```

--------------------------------

### Show DataSet contents

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Displays the contents of the loaded 'persons' DataSet. This is a standard Spark DataFrame operation.

```python
persons.show()
```

--------------------------------

### Corrected transform_to_schema with Combined Transformations

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/transforming_datasets.ipynb

Shows the correct way to apply multiple transformations to a single column within `transform_to_schema` by combining them into a single expression.

```python
transform_to_schema(
    df,
    Job,
    {
        Job.hourly_rate: (Job.hourly_rate + 3) * 2,
    },
).show()
```

--------------------------------

### Condensed DataFrame Transformation with transform_to_schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/transforming_datasets.ipynb

Shows a more concise way to transform a DataFrame to a new schema using the `transform_to_schema` function, simplifying the process.

```python
from typedspark import transform_to_schema

def get_plumbers(persons: DataSet[Person], jobs: DataSet[Job]) -> DataSet[PersonWithJob]:
    return transform_to_schema(
        jobs.filter(
            Job.function == "plumber",
        ).join(
            persons,
            Job.id == Person.job_id,
        ),
        PersonWithJob,
        {
            PersonWithJob.job_name: Job.function,
            PersonWithJob.rate: Job.hourly_rate,
        },
    )
```

--------------------------------

### Define Custom ColumnMeta and Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/subclass_column_meta.ipynb

Define a custom ColumnMeta class inheriting from TypedSpark's ColumnMeta and use it to add custom attributes like 'primary_key' to schema columns.

```python
from dataclasses import dataclass
from typing import Annotated
from pyspark.sql.types import LongType, StringType
from typedspark import ColumnMeta, Schema
from typedspark._core.column import Column


@dataclass
class MyColumnMeta(ColumnMeta):
    primary_key: bool = False


class Persons(Schema):
    id: Annotated[
        Column[LongType],
        MyColumnMeta(
            comment="Identifies the person",
            primary_key=True,
        ),
    ]
    name: Column[StringType]
    age: Column[LongType]


Persons.get_metadata()
```

--------------------------------

### Transform Data to StructType Column

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/structtype_columns.ipynb

Use `transform_to_schema` and `structtype_column` to create a StructType column from existing columns. The `transformations` dictionary requires unique keys for each column.

```python
from typedspark import create_partially_filled_dataset, transform_to_schema, structtype_column


class Input(Schema):
    a: Column[StringType]
    b: Column[IntegerType]


df = create_partially_filled_dataset(
    spark,
    Input,
    {
        Input.a: ["a", "b", "c"],
        Input.b: [1, 2, 3],
    },
)

transform_to_schema(
    df,
    Actions,
    {
        Actions.consequeces: structtype_column(
            Actions.consequeces.dtype.schema,
            {
                Actions.consequeces.dtype.schema.name: Input.a,
                Actions.consequeces.dtype.schema.severity: Input.b,
            },
        )
    },
).show()
```

--------------------------------

### Load Single Dataset with Custom Schema Name

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Loads a single dataset using `load_table` and explicitly sets the schema name using the `schema_name` argument. This is beneficial for ensuring clarity and proper naming of the loaded schema, especially in complex projects.

```python
person, Person = load_table(spark, "person_table", schema_name="Person")
Person
```

```python
from pyspark.sql.types import LongType, StringType

from typedspark import Column, Schema


class Person(Schema):
    name: Column[StringType]
    age: Column[LongType]
```

--------------------------------

### Generic Birthday Transformation using DataSetImplements

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/dataset_implements.ipynb

Defines a generic 'birthday' function using DataSetImplements, which allows the function to operate on any dataset that implements the 'Age' protocol (i.e., has an 'age' column). This approach is more flexible as it doesn't require explicit schema listing.

```python
from typing import Protocol

from typedspark import DataSetImplements


class Age(Schema, Protocol):
    age: Column[LongType]


T = TypeVar("T", bound=Schema)


def birthday(df: DataSetImplements[Age, T]) -> DataSet[T]:
    return transform_to_schema(
        df,
        df.typedspark_schema,
        {Age.age: Age.age + 1},
    )
```

--------------------------------

### Test Filtering Dogs with TypedSpark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/ide.ipynb

Unit test for the `get_dogs` function using chispa for DataFrame comparison. It creates a partially filled dataset and asserts the filtered output.

```python
from chispa.dataframe_comparer import assert_df_equality
from pyspark.sql import SparkSession
from typedspark import create_partially_filled_dataset


def test_get_dogs(spark: SparkSession):
    pets = create_partially_filled_dataset(
        spark,
        Pets,
        {
            Pets.pet_id: [1, 2, 3],
            Pets.species: ["dog", "cat", "dog"],
        },
    )

    observed = get_dogs(pets)
    expected = create_partially_filled_dataset(
        spark,
        Pets,
        {
            Pets.pet_id: [1, 3],
            Pets.species: ["dog", "dog"],
        },
    )

    assert_df_equality(
        observed,
        expected,
        ignore_row_order=True,
        ignore_nullable=True,
    )
```

--------------------------------

### Dynamically Generated Schema Definition

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_schema_in_notebook.ipynb

Shows the structure of the schema generated by `create_schema()` after a pivot operation. This schema can be directly imported and used.

```python
from pyspark.sql.types import DateType, LongType

from typedspark import Column, Schema


class DynamicallyLoadedSchema(Schema):
    pet_id: Column[LongType]
    influenza: Column[DateType]
    lyme: Column[DateType]
    rabies: Column[DateType]
```

--------------------------------

### Import necessary libraries for TypedSpark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Imports required modules from PySpark and TypedSpark for schema definition and data manipulation.

```python
from pyspark.sql import SparkSession
from pyspark.sql.types import DateType, LongType, StringType
from typedspark import Column, Schema, create_empty_dataset, create_partially_filled_dataset
from datetime import datetime, timedelta
```

--------------------------------

### Pivot DataFrame using Monkeypatch

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_schema_in_notebook.ipynb

This snippet shows how to pivot a DataFrame and aggregate data using a monkeypatched `to_typedspark()` method directly within a function chain. It requires prior import of TypedSpark components.

```python
pivot, Pivot = (
    vaccinations.groupby(Vaccinations.pet_id)
    .pivot(Vaccinations.vaccine_name.str)
    .agg(first(Vaccinations.next_due_date))
    .to_typedspark()
)
pivot.show()
```

--------------------------------

### Load Table with TypedSpark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/autocomplete_in_notebooks.ipynb

Loads a table into a DataFrame and infers its schema using TypedSpark's load_table function. The inferred schema is available via the second return value.

```python
from typedspark import load_table

df, Person = load_table(spark, "person_table")
```

--------------------------------

### Runtime Type Checking for Missing Columns

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/type_checking.ipynb

Illustrates the error raised when the data is missing columns that are defined in the schema. A TypeError is caught and printed.

```python
df = spark.createDataFrame(
    pd.DataFrame(
        dict(
            id=[1, 2, 3],
            name=["John", "Jane", "Jack"],
        )
    )
)
try:
    DataSet[Person](df)
except TypeError as e:
    print(e)
```

--------------------------------

### Runtime Type Checking for Extra Columns in Data

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/type_checking.ipynb

Demonstrates the error raised when the data contains columns not defined in the schema. The error message suggests adding the missing columns to the schema.

```python
df = spark.createDataFrame(
    pd.DataFrame(
        dict(
            id=[1, 2, 3],
            name=["John", "Jane", "Jack"],
            age=[20, 30, 40],
            gender=["male", "female", "male"],
        )
    )
)
try:
    DataSet[Person](df)
except TypeError as e:
    print(e)
```

--------------------------------

### UnionByName Two Datasets of Same Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/advanced_linting_support.ipynb

Shows how to combine two DataFrames with the same schema using unionByName, maintaining the schema type for the result.

```python
df_a = create_partially_filled_dataset(spark, A, {A.a: ["a", "b", "c"]})
df_b = create_partially_filled_dataset(spark, A, {A.a: ["d", "e", "f"]})

res = df_a.unionByName(df_b)
```

--------------------------------

### Generate DataSet with StructType Column from DataFrame

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/structtype_columns.ipynb

Creates a `DataSet` with a `StructType` column by collecting rows from another `DataSet`. This is useful when the nested data is already structured as rows.

```python
from typedspark import create_partially_filled_dataset

values = create_partially_filled_dataset(
    spark,
    Values,
    {
        Values.severity: [1, 2, 3],
    },
)

actions = create_partially_filled_dataset(
    spark,
    Actions,
    {
        Actions.consequeces: values.collect(),
    },
)
actions.show()
```

--------------------------------

### Drop Databases

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Drops multiple databases ('vet', 'library', 'store') using Spark SQL. This is typically used for cleanup after testing or development.

```python
spark.sql("DROP DATABASE vet CASCADE")
spark.sql("DROP DATABASE library CASCADE")
spark.sql("DROP DATABASE store CASCADE")
```

--------------------------------

### Access a specific table using Catalogs

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Retrieves the 'vaccinations' DataFrame and its corresponding TypedSpark schema from the 'vet' database.

```python
vaccinations, Vaccinations = db.spark_catalog.vet.vaccinations()
```

--------------------------------

### Pivot DataFrame and Generate Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_schema_in_notebook.ipynb

Pivots the DataFrame and uses `create_schema()` to dynamically generate a schema from the pivoted result. The new schema `Pivot` can be used for type-safe filtering.

```python
pivot = (
    vaccinations.groupby(Vaccinations.pet_id)
    .pivot(Vaccinations.vaccine_name.str)
    .agg(first(Vaccinations.next_due_date))
)

pivot, Pivot = create_schema(pivot)
pivot.show()
```

```python
pivot.filter(Pivot.influenza.isNotNull()).show()
```

--------------------------------

### Display inferred Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Displays the inferred Schema class for the 'persons' DataSet. This shows the structure and types of the data.

```python
Person
```

--------------------------------

### Filter DataSet using Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/loading_datasets_in_notebooks.ipynb

Filters the 'persons' DataSet to show individuals older than 25, using the inferred 'Person' Schema for column access. This demonstrates type-safe filtering.

```python
persons.filter(Person.age > 25).show()
```

--------------------------------

### Accessing Monkeypatched Schema Class

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/create_schema_in_notebook.ipynb

After using `to_typedspark()`, the corresponding schema class (e.g., `Pivot`) becomes available. This snippet shows how to access and inspect this dynamically generated schema class.

```python
Pivot
```

--------------------------------

### Display Schema with TypedSpark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/structtypes_in_notebooks.ipynb

Displays the schema of the loaded table using TypedSpark's schema object. This provides a structured view of the DataFrame's schema.

```python
ContainerSchema
```

--------------------------------

### Merging Schemas for Joins

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/subclassing_schemas.ipynb

Illustrates subclassing to merge multiple schemas (PersonA, PersonB) into a combined schema (PersonAB) for operations like joins. Use when combining data from different sources with overlapping or distinct columns.

```python
class PersonA(Schema):
    id: Column[LongType]
    name: Column[StringType]


class PersonB(Schema):
    id: Column[LongType]
    age: Column


class PersonAB(PersonA, PersonB):
    pass


def foo(df_a: DataSet[PersonA], df_b: DataSet[PersonB]) -> DataSet[PersonAB]:
    return DataSet[PersonAB](
        df_a.join(df_b, PersonAB.id),
    )
```

--------------------------------

### Define StructType Columns

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/structtype_columns.ipynb

Define nested schemas using TypedSpark's Schema and Column types. This allows for type-safe access to sub-columns.

```python
from pyspark.sql.types import IntegerType, StringType
from typedspark import DataSet, StructType, Schema, Column


class Values(Schema):
    name: Column[StringType]
    severity: Column[IntegerType]


class Actions(Schema):
    consequeces: Column[StructType[Values]]
```

--------------------------------

### Filter Pets Dataset for Dogs

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/ide.ipynb

A function to filter a DataSet of Pets to include only those with the species 'dog'. It utilizes TypedSpark's schema-aware filtering.

```python
def get_dogs(pets: DataSet[Pets]) -> DataSet[Pets]:
    return pets.filter(Pets.species == "dog")
```

--------------------------------

### Display all vaccinations

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Shows all records in the 'vaccinations' DataFrame.

```python
vaccinations.show()
```

--------------------------------

### Generic Birthday Transformation using DataSet

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/dataset_implements.ipynb

Defines a generic 'birthday' function that increments the 'age' column for datasets typed as Person or Pet. This approach requires explicit type hinting for schemas that support the transformation.

```python
from typing import TypeVar, Union

from typedspark import DataSet, transform_to_schema

T = TypeVar("T", bound=Union[Person, Pet])


def birthday(df: DataSet[T]) -> DataSet[T]:
    return transform_to_schema(
        df,
        df.typedspark_schema,
        {Person.age: Person.age + 1},
    )
```

--------------------------------

### Define Typed Function for Filtering

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/advanced_linting_support.ipynb

Illustrates a Python function that takes and returns a DataSet of a specific schema, preserving schema information through filtering.

```python
def foo(df: DataSet[A]) -> DataSet[A]:
    return df.filter(A.a == "a")
```

--------------------------------

### Define Schemas for Pets, Vaccinations, and Owners

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/ide.ipynb

Define Pydantic-like schemas for Spark DataFrames using TypedSpark's Schema and Column classes. This allows for type-safe access to DataFrame columns.

```python
from pyspark.sql.types import DateType, LongType, StringType
from typedspark import DataSet, Column, Schema


class Pets(Schema):
    pet_id: Column[LongType]
    owner_id: Column[LongType]
    pet_name: Column[StringType]
    species: Column[StringType]
    breed: Column[StringType]
    age: Column[LongType]
    birthdate: Column[DateType]
    gender: Column[StringType]


class Vaccinations(Schema):
    vaccination_id: Column[LongType]
    pet_id: Column[LongType]
    vaccine_name: Column[StringType]
    vaccine_date: Column[DateType]
    next_due_date: Column[DateType]


class Owners(Schema):
    owner_id: Column[LongType]
    first_name: Column[StringType]
    last_name: Column[StringType]
    email: Column[StringType]
    phone_number: Column[StringType]
    address: Column[StringType]
```

--------------------------------

### Vaccinations Schema Definition

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Displays the Python code for the TypedSpark schema definition of the Vaccinations table.

```python
from pyspark.sql.types import DateType, LongType, StringType

from typedspark import Column, Schema


class Vaccinations(Schema):
    vaccination_id: Column[LongType]
    pet_id: Column[LongType]
    vaccine_name: Column[StringType]
    vaccine_date: Column[DateType]
    next_due_date: Column[DateType]
```

--------------------------------

### Pivot Data with TypedSpark

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/videos/notebook.ipynb

Pivots vaccination data by pet ID, aggregating the next due date for each vaccine type. This is useful for transforming long-format data into a wide format for easier analysis.

```python
pivot = (
    vaccinations.groupby(Vaccinations.pet_id)
    .pivot(Vaccinations.vaccine_name.str)
    .agg(F.first(Vaccinations.next_due_date))
)
```

```python
pivot.show()
```

--------------------------------

### Implementing a Birthday Function with TypedSpark Schema Attributes

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/schema_attributes.ipynb

This function increments the 'age' column using TypedSpark schema attributes (`Person.age.str` and `Person.age`). This approach provides auto-completion and facilitates refactoring.

```python
def birthday(df: DataSet[Person]) -> DataSet[Person]:
    return DataSet[Person](
        df.withColumn(Person.age.str, Person.age + 1),
    )
```

--------------------------------

### Runtime Type Checking Ignoring Extra Columns

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/type_checking.ipynb

Shows that columns prefixed with '__' are ignored during schema validation, preventing errors when extra columns are present in the data.

```python
df = spark.createDataFrame(
    pd.DataFrame(
        dict(
            id=[1, 2, 3],
            name=["John", "Jane", "Jack"],
            age=[20, 30, 40],
            __extra_column=[1, 2, 3],
        )
    )
)
# no errors raised because __extra_column is ignored during the check
df = DataSet[Person](df)
df.show()
```

--------------------------------

### Resolve Ambiguous Columns with Schema Registration

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/column_ambiguity.ipynb

Resolves ambiguous column references in a join operation by registering schemas to their respective datasets using `register_schema_to_dataset`. This allows Spark to correctly identify columns.

```python
from typedspark import register_schema_to_dataset

person = register_schema_to_dataset(df_a, Person)
job = register_schema_to_dataset(df_b, Job)
( 
    df_a.join(
        df_b,
        person.id == job.id,
    ).show()
)
```

--------------------------------

### Typedspark DataFrame Function with Schema

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/index.md

A Python function using Typedspark's DataSet and Schema for explicit column-level type annotations. This provides enhanced type safety and code clarity.

```python
from typedspark import Column, DataSet, Schema
from pyspark.sql.types import LongType, StringType

class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]
    age: Column[LongType]

def foo(df: DataSet[Person]) -> DataSet[Person]:
    # do stuff
    return df
```

--------------------------------

### Join with Ambiguous Columns

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/column_ambiguity.ipynb

Demonstrates a join operation between two DataFrames with overlapping column names, leading to an ambiguous reference error. This highlights the need for schema registration.

```python
from typedspark import Column, Schema, create_partially_filled_dataset
from pyspark.sql.types import IntegerType, StringType


class Person(Schema):
    id: Column[IntegerType]
    name: Column[StringType]
    age: Column[IntegerType]


class Job(Schema):
    id: Column[IntegerType]
    salary: Column[IntegerType]


df_a = create_partially_filled_dataset(spark, Person, {Person.id: [1, 2, 3]})
df_b = create_partially_filled_dataset(spark, Job, {Job.id: [1, 2, 3]})

try:
    df_a.join(df_b, Person.id == Job.id)
except Exception as e:
    print(e)
```

--------------------------------

### Generate Dataset with Complex Data Types

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/complex_datatypes.ipynb

Create a partially filled Spark DataFrame using a TypedSpark schema that includes complex types like arrays, maps, decimals, intervals, dates, and timestamps. This function requires a SparkSession, a schema class, and a dictionary mapping schema fields to data.

```python
from datetime import date, datetime, timedelta
from decimal import Decimal
from pyspark.sql.types import DateType, TimestampType
from typedspark._utils.create_dataset import create_partially_filled_dataset


class MoreValues(Values):
    date: Column[DateType]
    timestamp: Column[TimestampType]


create_partially_filled_dataset(
    spark,
    MoreValues,
    {
        MoreValues.array: [["a", "b", "c"]],
        MoreValues.map: [{"a": "b"}],
        MoreValues.decimal: [Decimal(32)],
        MoreValues.interval: [timedelta(days=1, hours=2, minutes=3, seconds=4)],
        MoreValues.date: [date(2020, 1, 1)],
        MoreValues.timestamp: [datetime(2020, 1, 1, 10, 15)],
    },
).show()
```

--------------------------------

### Defining a TypedSpark Schema for a Person Dataset

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/schema_attributes.ipynb

This snippet defines a `Person` schema using TypedSpark, specifying column names and their corresponding PySpark data types. It serves as a foundation for type-safe DataFrame operations.

```python
from typedspark import Column, DataSet, Schema
from pyspark.sql.types import LongType, StringType
from pyspark.sql.functions import col


class Person(Schema):
    id: Column[LongType]
    name: Column[StringType]
    age: Column[LongType]
```

--------------------------------

### Define Nested StructType Schemas

Source: https://github.com/kaiko-ai/typedspark/blob/main/docs/source/structtypes_in_notebooks.ipynb

Defines nested schemas using TypedSpark's Schema and Column for StructType. This allows for type-safe access to nested fields.

```python
from typedspark import Schema, Column, StructType, create_partially_filled_dataset, load_table
from pyspark.sql.types import IntegerType


class Values(Schema):
    a: Column[IntegerType]
    b: Column[IntegerType]


class Container(Schema):
    values: Column[StructType[Values]]
```