dagster-polars 0.0.2

Dagster integration library for Polars

dagster-polars

Polars integration library for Dagster.

Features

• All IOManagers log various metadata about the DataFrame - size, schema, sample, stats, ...
• For all IOManagers the "columns" input metadata key can be used to select a subset of columns to load
• BasePolarsUPathIOManager is a base class for IO managers that work with Polars DataFrames. Shouldn't be used directly unless you want to implement your own IOManager.
  • returns the correct type (polars.DataFrame or polars.LazyFrame) based on the type annotation
  • inherits all the features of the UPathIOManager - works with local and remote filesystems (like S3), supports loading multiple partitions (use dict[str, pl.DataFrame] type annotation), ...
  • Implemented serialization formats:
    • PolarsParquetIOManager - for reading and writing files in Apache Parquet format. Supports reading partitioned Parquet datasets (for example, often produced by Spark).
• BigQueryPolarsIOManager - for reading and writing data from/to BigQuery. Supports writing partitioned tables ("partition_expr" input metadata key must be specified).

Quickstart

Installation

pip install dagster-polars

To use the BigQueryPolarsIOManager you need to install the gcp extra:

pip install 'dagster-polars[gcp]'

Usage

import polars as pl
from dagster import asset, Definitions
from dagster_polars import PolarsParquetIOManager


@asset(io_manager_key="polars_parquet_io_manager")
def upstream() -> pl.DataFrame:
    df: pl.DataFrame = ...
    return df


@asset(io_manager_key="polars_parquet_io_manager")
def downstream(upstream: pl.LazyFrame) -> pl.DataFrame:
    df = ...  # some lazy operations with `upstream`
    return df.collect()


definitions = Definitions(
    assets=[upstream, downstream],
    resources={
        "polars_parquet_io_manager": PolarsParquetIOManager(base_dir="/remote/or/local/path")
    }
)

Development

Installation

poetry install
poetry run pre-commit install

Testing

poetry run pytest

TODO

• Add PolarsDeltaIOManager
• Data validation like in dagster-pandas
• Maybe use DagsterTypeLoader ?