diff --git a/src/msfabricutils/etl/load/delta_table.py b/src/msfabricutils/etl/load/delta_table.py index e8f7b8c..4c4ca66 100644 --- a/src/msfabricutils/etl/load/delta_table.py +++ b/src/msfabricutils/etl/load/delta_table.py @@ -1,3 +1,4 @@ +from typing import Any import polars as pl from deltalake import DeltaTable, write_deltalake @@ -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. @@ -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. @@ -37,8 +40,6 @@ def upsert( from msfabricutils.etl import upsert import polars as pl - - config = get_default_config() data = pl.DataFrame({...}) upsert( @@ -56,7 +57,7 @@ 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 = [] @@ -64,6 +65,9 @@ def upsert( 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() @@ -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 @@ -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 @@ -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: @@ -140,6 +227,7 @@ def overwrite(table_or_uri: str | DeltaTable, df: PolarsFrame) -> None: table_or_uri=dt, data=df, mode="overwrite", + **delta_write_options, ) @@ -147,6 +235,7 @@ 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. @@ -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 @@ -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) @@ -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 @@ -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) @@ -229,10 +325,16 @@ 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. @@ -240,6 +342,7 @@ def incremental(table_or_uri: str | DeltaTable, df: PolarsFrame, incremental_col 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 @@ -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) @@ -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, ) - -