datazone-sdk 6.0.1.dev3__tar.gz → 6.0.1.dev5__tar.gz

{datazone_sdk-6.0.1.dev3 → datazone_sdk-6.0.1.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datazone-sdk
-Version: 6.0.1.dev3
+Version: 6.0.1.dev5
 Summary: Database and Delta storage client library for working with Delta Lake tables
 Author: Team Enigma
 Author-email: enigma@energinet.dk

{datazone_sdk-6.0.1.dev3 → datazone_sdk-6.0.1.dev5}/datazone/__init__.py

@@ -1,6 +1,13 @@
 from datazone.backtesting import backtest
 from datazone.caching import ParquetCache
 from datazone.db import DatabaseClient, SnapshotDatabaseClient
-from datazone.deltastorage import Field, HyperSlice, Schema, Store, Table
+from datazone.deltastorage import (
+    DeltaOverwriteFallbackWarning,
+    Field,
+    HyperSlice,
+    Schema,
+    Store,
+    Table,
+)
 
 from . import testing

{datazone_sdk-6.0.1.dev3 → datazone_sdk-6.0.1.dev5}/datazone/deltastorage/__init__.py

@@ -1,4 +1,4 @@
 from .schema import Field, Schema
 from .slicing import HyperSlice
 from .store import Store
-from .table import Table
+from .table import DeltaOverwriteFallbackWarning, Table

{datazone_sdk-6.0.1.dev3 → datazone_sdk-6.0.1.dev5}/datazone/deltastorage/table.py

@@ -1,4 +1,5 @@
 import datetime as dt
+import warnings
 from typing import Any, Optional
 
 import deltalake as dl
@@ -51,6 +52,35 @@ def _is_arrow_cast_commit_failure(error: CommitFailedError) -> bool:
     return "arrow_cast should have been simplified to cast" in str(error)
 
 
+def _is_safe_partition_filter(filter_: tuple) -> bool:
+    return filter_[1] in ["=", "in"]
+
+
+def _empty_to_none(filters: list[tuple]) -> list[tuple] | None:
+    return filters if len(filters) > 0 else None
+
+
+def _filter_to_polars_expr(filter_: tuple) -> pl.Expr:
+    col, op, val = filter_
+    if op == "=":
+        return pl.col(col) == val
+    if op == "in":
+        return pl.col(col).is_in(val)
+    if op == ">=":
+        return pl.col(col) >= val
+    if op == "<=":
+        return pl.col(col) <= val
+    if op == ">":
+        return pl.col(col) > val
+    if op == "<":
+        return pl.col(col) < val
+    raise ValueError(f"Unsupported operation: {op}")
+
+
+class DeltaOverwriteFallbackWarning(RuntimeWarning):
+    pass
+
+
 class Table:
     def __init__(
         self,
@@ -110,8 +140,10 @@ class Table:
         if hyper_slice is None:
             hyper_slice = []
 
-        # add generated filters to hyperslice
-        hyper_slice = self.schema().add_generated_filters(hyper_slice)
+        # Generated filters are an optimization for partition pruning. Keep them out
+        # of row filters because partition values may be represented as strings by
+        # PyArrow even when the Delta schema has a richer logical type.
+        partition_hyper_slice = self.schema().add_generated_filters(hyper_slice)
 
         delta_table = self.delta_table
         partition_cols = delta_table.metadata().partition_columns
@@ -119,9 +151,23 @@ class Table:
         if len(hyper_slice) == 0:
             file_filters = None
             partition_filters = None
+            post_filters = []
         else:
-            file_filters = hyper_slice
-            partition_filters = [f for f in hyper_slice if f[0] in partition_cols]
+            file_filters = _empty_to_none(
+                [f for f in hyper_slice if f[0] not in partition_cols]
+            )
+            partition_filters = _empty_to_none(
+                [
+                    f
+                    for f in partition_hyper_slice
+                    if f[0] in partition_cols and _is_safe_partition_filter(f)
+                ]
+            )
+            post_filters = [
+                f
+                for f in hyper_slice
+                if f[0] in partition_cols and not _is_safe_partition_filter(f)
+            ]
 
         pyarrow_table_existing_data = delta_table.to_pyarrow_table(
             columns=columns,
@@ -129,7 +175,10 @@ class Table:
             filters=file_filters,
         )
 
-        return pl.from_arrow(pyarrow_table_existing_data)
+        df = pl.from_arrow(pyarrow_table_existing_data)
+        for filter_ in post_filters:
+            df = df.filter(_filter_to_polars_expr(filter_))
+        return df
 
     def _to_writable_pyarrow_table(self, df: pl.DataFrame, schema: Schema) -> pa.Table:
         """Convert Polars dataframe to pyarrow table with casted schema.
@@ -183,6 +232,8 @@ class Table:
         schema = self.schema()
         data = self._to_writable_pyarrow_table(df=df, schema=schema)
 
+        hyper_slice = schema.add_generated_filters(hyper_slice)
+
         if len(hyper_slice) == 0:
             predicate = None
         else:
@@ -194,6 +245,13 @@ class Table:
             if predicate is None or not _is_arrow_cast_commit_failure(error):
                 raise
 
+            warnings.warn(
+                "Delta predicate overwrite failed with the DataFusion arrow_cast bug. "
+                "Falling back to delete followed by append; this is not atomic as a "
+                "single overwrite commit.",
+                DeltaOverwriteFallbackWarning,
+                stacklevel=2,
+            )
             self.delta_table.delete(predicate)
             self._write_deltalake(data=data, mode="append", predicate=None)
 

{datazone_sdk-6.0.1.dev3 → datazone_sdk-6.0.1.dev5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datazone-sdk"
-version = "6.0.1.dev3"
+version = "6.0.1.dev5"
 description = "Database and Delta storage client library for working with Delta Lake tables"
 authors = [{ name = "Team Enigma", email = "enigma@energinet.dk" }]
 requires-python = ">=3.10"