Creating dataframes

Creating Dataframes

You can easily create example DataFrames from lists of model instances. This is particularly useful for building test fixtures for a transformation. Pass the desired frame class to select its backend. Spark also requires its session through the session keyword:

PolarsPySpark

from dataclasses import dataclass

import polars as pl

from data_joinery import Schema


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


rows = [
    Customer("1", "Alice", True),
    Customer("2", "Bob", False),
    Customer("3", "Charlie", True),
]

schema = Schema(Customer)
df = schema.create_dataframe(rows, pl.LazyFrame)
print(df.collect())

Outputs:

shape: (3, 3)
┌─────────────┬─────────┬───────────┐
│ customer_id ┆ name    ┆ is_active │
│ ---         ┆ ---     ┆ ---       │
│ str         ┆ str     ┆ bool      │
╞═════════════╪═════════╪═══════════╡
│ 1           ┆ Alice   ┆ true      │
│ 2           ┆ Bob     ┆ false     │
│ 3           ┆ Charlie ┆ true      │
└─────────────┴─────────┴───────────┘

from dataclasses import dataclass

from pyspark.sql import DataFrame, SparkSession

from data_joinery import Schema


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


spark = SparkSession.builder.getOrCreate()

rows = [
    Customer("1", "Alice", True),
    Customer("2", "Bob", False),
    Customer("3", "Charlie", True),
]

schema = Schema(Customer)
df = schema.create_dataframe(rows, DataFrame, session=spark)
df.show()

Outputs:

+-----------+-------+---------+
|customer_id|   name|is_active|
+-----------+-------+---------+
|          1|  Alice|     true|
|          2|    Bob|    false|
|          3|Charlie|     true|
+-----------+-------+---------+