Mixing Backends
It is possible to write a transformation that takes a dataframe in one backend and outputs it in another. A common use case is to move an aggregated Spark frame to Polars on the driver node after an aggregation.
This example moves the records from Polars to Spark and then to Polars; it does not change
their values or columns. record_id uses Annotated metadata to give the same Python int a
different explicit representation in every backend: Polars UInt32, Spark LongType, and Polars
UInt16.
from dataclasses import dataclass
from typing import Annotated, cast
import polars as pl
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql.types import LongType, StructType
from data_joinery import Schema, Strict, transform
from data_joinery.contract import ProjectCast
@dataclass
class Record:
record_id: Annotated[int, pl.UInt32, LongType()]
label: str
record_schema = Schema(Record)
@transform
def polars_to_spark(
records: Annotated[pl.LazyFrame, Strict(Record)],
spark: SparkSession,
) -> Annotated[DataFrame, Strict(Record)]:
return spark.createDataFrame(
records.collect().to_dicts(), schema=record_schema.native_schema(StructType)
)
@transform
def spark_to_polars(
records: Annotated[DataFrame, Strict(Record)],
) -> Annotated[pl.LazyFrame, ProjectCast(Record)]:
return cast(pl.DataFrame, pl.from_arrow(records.toArrow())).lazy()
spark = SparkSession.builder.getOrCreate()
polars_records = record_schema.create_dataframe(
[Record(1, "first"), Record(2, "second")], pl.LazyFrame
)
spark_records = polars_to_spark(polars_records, spark)
polars_records_back = spark_to_polars(spark_records)
print("Polars:\n", polars_records.schema)
print("Spark:\n", spark_records.schema)
print("Polars Back:\n", polars_records_back.schema)
Outputs: