Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 116 additions & 12 deletions src/msfabricutils/etl/load/delta_table.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from typing import Any

import polars as pl
from deltalake import DeltaTable, write_deltalake
Expand All @@ -18,6 +19,7 @@ def upsert(
primary_key_columns: str | list[str],
update_exclusion_columns: str | list[str] | None = None,
predicate_exclusion_columns: str | list[str] | None = None,
delta_merge_options: dict[str, Any] | None = None,
) -> dict[str, str]:
"""
Upserts dataframe into a Delta table using the provided primary key columns.
Expand All @@ -28,6 +30,7 @@ def upsert(
primary_key_columns (str | list[str]): Primary key column(s) for the upsert.
update_exclusion_columns (str | list[str] | None): Columns which will not be updated in the merge.
predicate_exclusion_columns (str | list[str] | None): Columns to exclude from the upsert. Difference between source and target of these columns will not trigger an update, however if there is a difference in the other columns, the row will be *also* be updated with the `excluded_columns`.
delta_merge_options (dict[str, Any] | None): Additional keyword arguments while merging to the Delta Table

Returns:
Result of the merge operation.
Expand All @@ -37,8 +40,6 @@ def upsert(
from msfabricutils.etl import upsert
import polars as pl


config = get_default_config()
data = pl.DataFrame({...})

upsert(
Expand All @@ -56,14 +57,17 @@ def upsert(
update_exclusion_columns = []

if isinstance(update_exclusion_columns, str):
update_exclusion_columns = [update_exclusion_columns]
update_exclusion_columns = [update_exclusion_columns]

if predicate_exclusion_columns is None:
predicate_exclusion_columns = []

if isinstance(predicate_exclusion_columns, str):
predicate_exclusion_columns = [predicate_exclusion_columns]

if delta_merge_options is None:
delta_merge_options = {}

if isinstance(df, pl.LazyFrame):
df = df.collect()

Expand Down Expand Up @@ -97,6 +101,7 @@ def upsert(
source_alias="source",
target_alias="target",
predicate=merge_predicate,
**delta_merge_options,
)
.when_matched_update(
predicate=when_matched_update_predicates, updates=when_matched_update_columns
Expand All @@ -107,13 +112,92 @@ def upsert(
return table_merger.execute()


def overwrite(table_or_uri: str | DeltaTable, df: PolarsFrame) -> None:
def upsert_cdf(
table_or_uri: str,
df: PolarsFrame,
primary_key_columns: list[str],
change_type_column: str = "_change_type",
change_type_insert: str = "insert",
change_type_update: str = "update",
change_type_delete: str = "delete",
delta_merge_options: dict[str, Any] | None = None,
):
"""
Upserts dataframe into a Delta table using the provided primary key columns using Change Data Feed (CDF).

Args:
table_or_uri (str): The URI of the target Delta table.
df (PolarsFrame): The dataframe to upsert.
primary_key_columns (str | list[str]): Primary key column(s) for the upsert.
change_type_column (str): The column in the dataframe containing the change type, i.e. `insert`, `update` and `delete`
change_type_insert (str): The value of the `insert` change type in the dataframe.
change_type_update (str): The value of the `update` change type in the dataframe.
change_type_delete (str): The value of the `delete` change type in the dataframe.
delta_write_options (dict[str, Any] | None): Additional keyword arguments while writing a Delta lake Table.

Returns:
Result of the merge operation.

Example:
```python
from msfabricutils.etl import upsert_cdf
import polars as pl

data = pl.DataFrame({...})

upsert_cdf(
"path/to/delta_table",
data,
primary_key_columns=["id"],
)
```
"""
if isinstance(primary_key_columns, str):
primary_key_columns = [primary_key_columns]

if delta_merge_options is None:
delta_merge_options = {}

data = df.lazy().collect()

if isinstance(table_or_uri, str):
dt = get_or_create_delta_table(table_or_uri, df.schema)
else:
dt = table_or_uri

merge_predicate = build_merge_predicate(primary_key_columns)


table_merger = (
dt.merge(
data,
source_alias="source",
target_alias="target",
predicate=merge_predicate,
**delta_merge_options,
)
.when_not_matched_insert_all(
predicate=f"source.{change_type_column} = '{change_type_insert}'"
)
.when_matched_update_all(predicate=f"source.{change_type_column} = '{change_type_update}'")
.when_matched_delete(predicate=f"source.{change_type_column} = '{change_type_delete}'")
)

return table_merger.execute()


def overwrite(
table_or_uri: str | DeltaTable,
df: PolarsFrame,
delta_write_options: dict[str, Any] | None = None,
) -> None:
"""
Overwrites the entire Delta table with the provided dataframe.

Args:
table_or_uri (str): The URI of the target Delta table.
df (PolarsFrame): The dataframe to write to the Delta table.
delta_write_options (dict[str, Any] | None): Additional keyword arguments while writing a Delta lake Table.

Example:
```python
Expand All @@ -126,6 +210,9 @@ def overwrite(table_or_uri: str | DeltaTable, df: PolarsFrame) -> None:
```
"""

if delta_write_options is None:
delta_write_options = {}

if isinstance(table_or_uri, str):
dt = get_or_create_delta_table(table_or_uri, df.schema)
else:
Expand All @@ -140,13 +227,15 @@ def overwrite(table_or_uri: str | DeltaTable, df: PolarsFrame) -> None:
table_or_uri=dt,
data=df,
mode="overwrite",
**delta_write_options,
)


def replace_range(
table_or_uri: str | DeltaTable,
df: PolarsFrame,
range_column: str,
delta_write_options: dict[str, Any] | None = None,
) -> None:
"""
Replaces a range of data in the Delta table based on a specified column.
Expand All @@ -155,6 +244,7 @@ def replace_range(
table_or_uri (str): The URI of the target Delta table.
df (PolarsFrame): The dataframe to write to the Delta table.
range_column (str): The column used to determine the range of data to replace. This replaces the data in the range of the `range_column` in the delta table based off the min and max values of the `range_column` in the dataframe.
delta_write_options (dict[str, Any] | None): Additional keyword arguments while writing a Delta lake Table.

Example:
```python
Expand All @@ -166,6 +256,8 @@ def replace_range(
replace_range("path/to/delta_table", data, range_column="date")
```
"""
if delta_write_options is None:
delta_write_options = {}

if isinstance(table_or_uri, str):
dt = get_or_create_delta_table(table_or_uri, df.schema)
Expand All @@ -189,20 +281,22 @@ def replace_range(
df = df.to_arrow()

write_deltalake(
table_or_uri=dt,
data=df,
mode="overwrite",
predicate=predicate,
table_or_uri=dt, data=df, mode="overwrite", predicate=predicate, **delta_write_options
)


def append(table_or_uri: str | DeltaTable, df: PolarsFrame) -> None:
def append(
table_or_uri: str | DeltaTable,
df: PolarsFrame,
delta_write_options: dict[str, Any] | None = None,
) -> None:
"""
Appends the provided dataframe to the Delta table.

Args:
table_or_uri (str): The URI of the target Delta table.
df (PolarsFrame): The dataframe to append to the Delta table.
delta_write_options (dict[str, Any] | None): Additional keyword arguments while writing a Delta lake Table.

Example:
```python
Expand All @@ -214,6 +308,8 @@ def append(table_or_uri: str | DeltaTable, df: PolarsFrame) -> None:
append("path/to/delta_table", data)
```
"""
if delta_write_options is None:
delta_write_options = {}

if isinstance(table_or_uri, str):
dt = get_or_create_delta_table(table_or_uri, df.schema)
Expand All @@ -229,17 +325,24 @@ def append(table_or_uri: str | DeltaTable, df: PolarsFrame) -> None:
table_or_uri=dt,
data=df,
mode="append",
**delta_write_options,
)


def incremental(table_or_uri: str | DeltaTable, df: PolarsFrame, incremental_column: str) -> None:
def incremental(
table_or_uri: str | DeltaTable,
df: PolarsFrame,
incremental_column: str,
delta_write_options: dict[str, Any] | None = None,
) -> None:
"""
Appends new data to the Delta table based on an incremental column.

Args:
table_or_uri (str): The URI of the target Delta table.
df (PolarsFrame): The dataframe to append to the Delta table.
incremental_column (str): The column used to determine new data to append. The source dataframe will only append rows where the value of the `incremental_column` is greater than the max value of the `incremental_column` in the delta table.
delta_write_options (dict[str, Any] | None): Additional keyword arguments while writing a Delta lake Table.

Example:
```python
Expand All @@ -251,6 +354,8 @@ def incremental(table_or_uri: str | DeltaTable, df: PolarsFrame, incremental_col
incremental("path/to/delta_table", data, incremental_column="timestamp")
```
"""
if delta_write_options is None:
delta_write_options = {}

if isinstance(table_or_uri, str):
dt = get_or_create_delta_table(table_or_uri, df.schema)
Expand All @@ -270,6 +375,5 @@ def incremental(table_or_uri: str | DeltaTable, df: PolarsFrame, incremental_col
table_or_uri=dt,
data=df,
mode="append",
**delta_write_options,
)