dagster-delta

Dagster deltalake implementation for Pyarrow & Polars. Originally forked from dagster-deltalake with customizations.

The IO Managers support partition mapping, custom write modes, special metadata configuration for advanced use cases.

The supported write modes:

• error
• append
• overwrite
• ignore
• merge
• create_or_replace

Merge

dagster-delta supports MERGE execution with a couple pre-defined MERGE types (dagster_delta.config.MergeType):

• deduplicate_insert <- Deduplicates on write
• update_only <- updates only the matches records
• upsert <- updates existing matches and inserts non matched records
• replace_and_delete_unmatched <- updates existing matches and deletes unmatched
• custom <- custom Merge with MergeOperationsConfig

Example:

from dagster_delta import DeltaLakePolarsIOManager, WriteMode, MergeConfig, MergeType
from dagster_delta_polars import DeltaLakePolarsIOManager

@asset(
    key_prefix=["my_schema"]  # will be used as the schema (parent folder) in Delta Lake
)
def my_table() -> pl.DataFrame:  # the name of the asset will be the table name
    ...

defs = Definitions(
    assets=[my_table],
    resources={"io_manager": DeltaLakePolarsIOManager(
        root_uri="s3://bucket",
        mode=WriteMode.merge, # or just "merge"
        merge_config=MergeConfig(
            merge_type=MergeType.upsert,
            predicate="s.a = t.a",
            source_alias="s",
            target_alias="t",
        )
    )}
)

Custom merge (gives full control)

from dagster_delta import DeltaLakePolarsIOManager, WriteMode, MergeConfig, MergeType, MergeOperationsConfig
from dagster_delta_polars import DeltaLakePolarsIOManager

@asset(
    key_prefix=["my_schema"]  # will be used as the schema (parent folder) in Delta Lake
)
def my_table() -> pl.DataFrame:  # the name of the asset will be the table name
    ...

defs = Definitions(
    assets=[my_table],
    resources={"io_manager": DeltaLakePolarsIOManager(
        root_uri="s3://bucket",
        mode=WriteMode.merge, # or just "merge"
        merge_config=MergeConfig(
            merge_type=MergeType.custom,
            predicate="s.a = t.a",
            source_alias="s",
            target_alias="t",
            merge_operations_config=MergeOperationsConfig(
                when_not_matched_insert_all=[WhenNotMatchedInsertAll(predicate="s.price > 600")],
                when_matched_update_all=[WhenMatchedUpdateAll()],
            ),
        )
    )}
)

Special metadata configurations

Add additional table_configuration

Specify additional table configurations for configuration in write_deltalake.

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    metadata={"table_configuration": {
        "delta.enableChangeDataFeed": "true"
    }},
)
def my_asset() -> pl.DataFrame:
    ...

Overwrite the write mode

Override the write mode to be used in write_deltalake.

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    metadata={"mode": "append"},
)
def my_asset() -> pl.DataFrame:
    ...

Overwrite the custom_metadata

Override the custom_metadata to be used in write_deltalake.

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    metadata={"custom_metadata": {"owner":"John Doe"}},
)
def my_asset() -> pl.DataFrame:
    ...

Overwrite the write schema_mode

Override the schema_mode to be used in write_deltalake.

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    metadata={"schema_mode": "merge"},
)
def my_asset() -> pl.DataFrame:
    ...

Overwrite the writer_properties

Override the writer_properties to be used in write_deltalake.

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    metadata={"writer_properties": {
        "compression": "SNAPPY",
    }},
)
def my_asset() -> pl.DataFrame:
    ...

Overwrite the merge_predicate

Override the merge_predicate to be used with merge execution.

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    metadata={"merge_predicate": "s.foo = t.foo AND s.bar = t.bar"},
)
def my_asset() -> pl.DataFrame:
    ...

Overwrite the schema

Override the schema of where the table will be saved

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    metadata={"schema": "custom_db_schema"},
)
def my_asset() -> pl.DataFrame:
    ...

Set the columns that need to be read

Override the columns to only load these columns in

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    ins = {
        "upstream_asset": dg.AssetIn(metadata={"columns":["foo","bar"]})
    }
)
def my_asset(upstream_asset) -> pl.DataFrame:
    ...

Override table name using root_name

Instead of using the asset_name for the table name it's possible to set a custom table name using the root_name in the asset defintion metadata.

This is useful where you have two or multiple assets who have the same table structure, but each asset is a subset of the full table partition_definition, and it wasn't possible to combine this into a single asset due to requiring different underlying Op logic and/or upstream assets:

import polars as pl
import dagster as dg

@dg.asset(
    io_manager_key = "deltalake_io_manager",
    partitions_def=dg.StaticPartitionsDefinition(["a", "b"]),
    metadata={
        "partition_expr": "foo",
        "root_name": "asset_partitioned",
        },
)
def asset_partitioned_1(upstream_1: pl.DataFrame, upstream_2: pl.DataFrame) -> pl.DataFrame:
    ...

@dg.asset(
    partitions_def=dg.StaticPartitionsDefinition(["c", "d"]),
    metadata={
        "partition_expr": "foo",
        "root_name": "asset_partitioned",
        },
)
def asset_partitioned_2(upstream_3: pl.DataFrame, upstream_4: pl.DataFrame) -> pl.DataFrame:
    ...

Effectively this would be the flow:

                 {static_partition_def: [a,b]}
┌───────────┐
│upstream 1 ├─┐ ┌────────────────────────┐
└───────────┘ │ │                        │            write to storage on partition (a,b)
┌───────────┐ └─►   asset_partitioned_1  ├──────────────────────┐
│upstream 2 ├───►                        │                      │
└───────────┘   └────────────────────────┘       ┌──────────────▼──────────────────┐
                                                 │                     partitions  │
                                                 │  asset_partitioned:             │
                                                 │                     [a,b,c,d]   │
┌───────────┐   ┌────────────────────────┐       └──────────────▲──────────────────┘
│upstream 3 ├──┐│                        │                      │
└───────────┘  └►   asset_partitioned_2  │                      │
┌───────────┐ ┌─►                        ├──────────────────────┘
│upstream 4 ├─┘ └────────────────────────┘            write to storage on partition (c,d)
└───────────┘
                 {static_partition_def: [c,d]}