Additional testing

Author: dehume

libraries/dagster-polars/dagster_polars/io_managers/delta.py

@@ -271,9 +271,7 @@ def write_df_to_path(
 
                 if use_legacy_deltalake:
                     if engine == "rust":
-                        delta_write_options["predicate"] = self.get_predicate(
-                            context, df
-                        )
+                        delta_write_options["predicate"] = self.get_predicate(context)
 
                     elif engine == "pyarrow":
                         delta_write_options["partition_filters"] = (
@@ -283,7 +281,7 @@ def write_df_to_path(
                     else:
                         raise NotImplementedError(f"Invalid engine: {engine}")
                 else:
-                    delta_write_options["predicate"] = self.get_predicate(context, df)
+                    delta_write_options["predicate"] = self.get_predicate(context)
 
         if delta_write_options:
             context.log.debug(
@@ -472,18 +470,19 @@ def get_partition_filters(
     @staticmethod
     def get_predicate(
         context: InputContext | OutputContext,
-        df: pl.DataFrame | None = None,
     ) -> str | None:
         """Create a predicate for `deltalake` to select which partitions are overwritten.
 
         Returns `None` if the entire table is overwritten.
 
-        Partition keys are quoted as string literals by default. When the target
-        column in the written DataFrame is a `Date` type, the key is emitted as a
-        `DATE '...'` literal so the predicate compares against the native column
-        type. The column dtype is used as the signal rather than the
-        `PartitionsDefinition` type, since a `TimeWindowPartitionsDefinition` may
-        map onto a non-temporal column (e.g. a string `year` derived from `%Y`).
+        Partition keys are always compared with `=` (multiple keys are joined
+        with `OR`) using plain string literals. This works across column types
+        (string, `Date`, `Datetime`, integer, ...) because `deltalake`'s
+        DataFusion predicate parser implicitly coerces the string literal to the
+        column's type for `=`. The `IN (...)` operator is deliberately avoided:
+        it requires the literal type to match the column type exactly, so a
+        string literal fails against a `Date`/`Datetime` column, and there is no
+        literal form that both parses and matches for temporal columns (see #330).
 
         See documentation here:
         https://delta-io.github.io/delta-rs/usage/writing/#overwriting-part-of-the-table-data-using-a-predicate
@@ -499,12 +498,8 @@ def get_predicate(
                 f"Invalid context type: {type(context)}"
             )
 
-        schema = df.schema if df is not None else None
-
-        def key_to_predicate(key: str, column: str) -> str:
-            if schema is not None and schema.get(column) == pl.Date:
-                return f"DATE '{key}'"
-            return f"'{key}'"
+        def keys_to_predicate(column: str, keys: list[str]) -> str:
+            return " OR ".join(f"{column} = '{key}'" for key in keys)
 
         if partition_by is None or not context.has_asset_partitions:
             return
@@ -520,26 +515,18 @@ def key_to_predicate(key: str, column: str) -> str:
                     all_keys_by_dim[dim].append(key)
 
             predicate = " AND ".join(
-                [
-                    f"{partition_by[dim]} in ({', '.join(key_to_predicate(key, partition_by[dim]) for key in keys)})"
-                    for dim, keys in all_keys_by_dim.items()
-                ]
+                f"({keys_to_predicate(partition_by[dim], keys)})"
+                for dim, keys in all_keys_by_dim.items()
             )
 
         elif isinstance(partition_by, str):
             assert not isinstance(context.asset_partition_keys[0], MultiPartitionKey), (
                 f"Received string `partition_by` metadata value `{partition_by}`, "
                 f"but the partition_key is not a `MultiPartitionKey`: {context.asset_partition_keys[0]}"
             )
-
-            if len(context.asset_partition_keys) == 1:
-                predicate = f"{partition_by} = {key_to_predicate(context.asset_partition_keys[0], partition_by)}"
-            else:
-                keys = ", ".join(
-                    key_to_predicate(key, partition_by)
-                    for key in context.asset_partition_keys
-                )
-                predicate = f"{partition_by} in ({keys})"
+            predicate = keys_to_predicate(
+                partition_by, list(context.asset_partition_keys)
+            )
 
         else:
             raise NotImplementedError(

libraries/dagster-polars/dagster_polars_tests/test_polars_delta.py

@@ -386,6 +386,54 @@ def downstream_partitioned(
         assert DeltaTable(saved_path).metadata().partition_columns == ["year"]
 
 
+def test_polars_delta_native_partitioning_datetime_column_predicate(
+    polars_delta_io_manager: PolarsDeltaIOManager,
+    df_for_delta: pl.DataFrame,
+):
+    """The overwrite predicate must work for a ``Datetime`` partition column.
+
+    In ``0.27.12`` the predicate emitted a ``DATE '...'`` literal, which fails
+    to compare against a ``Datetime`` column (`Invalid comparison operation:
+    Timestamp <= ...`). A plain string literal is coerced by DataFusion to the
+    column's timestamp type, so overwriting one partition leaves the others
+    intact.
+
+    Note: this exercises the write/overwrite predicate path only. Reading back a
+    *single* partition of a ``Datetime``-partitioned table is a separate,
+    pre-existing delta-rs limitation (the partition value cannot be parsed as
+    ``timestamp_ntz``), so the table is read back in full here rather than via a
+    partitioned downstream asset.
+    """
+    manager = polars_delta_io_manager
+    df = df_for_delta
+
+    partitions_def = DailyPartitionsDefinition(start_date=datetime(2024, 1, 1))
+
+    @asset(
+        io_manager_def=manager,
+        partitions_def=partitions_def,
+        metadata={"partition_by": "ts"},
+    )
+    def upstream_partitioned(context: OpExecutionContext) -> pl.DataFrame:
+        return df.with_columns(
+            pl.lit(context.partition_key)
+            .str.strptime(pl.Datetime, "%Y-%m-%d")
+            .alias("ts")
+        )
+
+    saved_path = None
+    for partition_key in ["2024-01-01", "2024-01-02"]:
+        result = materialize([upstream_partitioned], partition_key=partition_key)
+        saved_path = get_saved_path(result, "upstream_partitioned")
+
+    assert saved_path is not None
+    assert DeltaTable(saved_path).metadata().partition_columns == ["ts"]
+
+    # Both partition writes succeeded and neither predicate clobbered the other.
+    written = pl.read_delta(saved_path)["ts"].unique().sort().to_list()
+    assert written == [datetime(2024, 1, 1), datetime(2024, 1, 2)]
+
+
 def test_polars_delta_native_multi_partitions(
     polars_delta_io_manager: PolarsDeltaIOManager,
     df_for_delta: pl.DataFrame,
@@ -541,51 +589,28 @@ def asset_schema_2(context: OpExecutionContext) -> pl.DataFrame:
 
 
 @pytest.mark.parametrize(
-    "partition_by, partition_keys, expected_filters, expected_predicate, df",
+    "partition_by, partition_keys, expected_filters, expected_predicate",
     [
-        ("col_name", ["a"], [("col_name", "in", ["a"])], "col_name = 'a'", None),
+        ("col_name", ["a"], [("col_name", "in", ["a"])], "col_name = 'a'"),
         (
             "col_name",
             ["a", "b"],
             [("col_name", "in", ["a", "b"])],
-            "col_name in ('a', 'b')",
-            None,
+            "col_name = 'a' OR col_name = 'b'",
         ),
         (
             {"col_name": "mapped_col"},
             [{"col_name": "a"}],
             [("mapped_col", "in", ["a"])],
-            "mapped_col in ('a')",
-            None,
+            "(mapped_col = 'a')",
         ),
         (
             {"col_name": "mapped_col"},
             [{"col_name": "a"}, {"col_name": "b"}],
             [("mapped_col", "in", ["a", "b"])],
-            "mapped_col in ('a', 'b')",
-            None,
-        ),
-        (None, [], [], None, None),
-        # Regression (#330): a time-window partition mapped onto a non-temporal
-        # (string) column must stay quoted as a string, not wrapped in `DATE`.
-        (
-            "year",
-            ["2025"],
-            [("year", "in", ["2025"])],
-            "year = '2025'",
-            pl.DataFrame({"year": ["2025"]}, schema={"year": pl.String}),
-        ),
-        # A genuine `Date` column is emitted as a `DATE '...'` literal so the
-        # predicate compares against the native column type.
-        (
-            "date",
-            ["2025-01-01"],
-            [("date", "in", ["2025-01-01"])],
-            "date = DATE '2025-01-01'",
-            pl.DataFrame(
-                {"date": [datetime(2025, 1, 1).date()]}, schema={"date": pl.Date}
-            ),
+            "(mapped_col = 'a' OR mapped_col = 'b')",
         ),
+        (None, [], [], None),
     ],
 )
 @pytest.mark.parametrize("context", [InputContext, OutputContext])
@@ -595,7 +620,6 @@ def test_partition_filters_predicate(
     partition_keys: list[str] | list[dict[str, str]],
     expected_filters: list[tuple[str, str, list[str]]],
     expected_predicate: str,
-    df: pl.DataFrame | None,
     context: type[InputContext | OutputContext],
 ):
     """Test that the partition filters and predicate are generated correctly."""
@@ -633,4 +657,4 @@ def test_partition_filters_predicate(
         )
 
     assert PolarsDeltaIOManager.get_partition_filters(context) == expected_filters
-    assert PolarsDeltaIOManager.get_predicate(context, df) == expected_predicate
+    assert PolarsDeltaIOManager.get_predicate(context) == expected_predicate