fix(rust, python): block projection pushdown on unnest (pola-rs#5093)

zundertj · Jan 7, 2023 · fbbe4d1 · fbbe4d1
1 parent 6cf54e4
commit fbbe4d1
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 2 deletions.
diff --git a/polars/polars-lazy/polars-plan/Cargo.toml b/polars/polars-lazy/polars-plan/Cargo.toml
@@ -28,7 +28,7 @@ parquet = ["polars-core/parquet", "polars-io/parquet"]
 ipc = ["polars-io/ipc"]
 json = ["polars-io/json"]
 csv-file = ["polars-io/csv-file"]
-temporal = ["polars-core/temporal", "dtype-date", "dtype-datetime"]
+temporal = ["polars-core/temporal", "dtype-date", "dtype-datetime", "dtype-time"]
 # debugging purposes
 fmt = ["polars-core/fmt"]
 strings = ["polars-core/strings", "polars-ops/strings"]
@@ -40,6 +40,7 @@ dtype-i16 = ["polars-core/dtype-i16"]
 dtype-date = ["polars-core/dtype-date", "polars-time/dtype-date", "temporal"]
 dtype-datetime = ["polars-core/dtype-datetime", "polars-time/dtype-datetime", "temporal"]
 dtype-duration = ["polars-core/dtype-duration", "polars-time/dtype-duration", "temporal"]
+dtype-time = ["polars-core/dtype-time", "polars-time/dtype-time"]
 dtype-categorical = ["polars-core/dtype-categorical"]
 dtype-struct = ["polars-core/dtype-struct"]
 object = ["polars-core/object"]

diff --git a/polars/polars-lazy/src/frame/mod.rs b/polars/polars-lazy/src/frame/mod.rs
@@ -1187,7 +1187,17 @@ impl LazyFrame {
         };
         self.map(
             move |df| df.unnest(&cols2),
-            Some(AllowedOptimizations::default()),
+            Some(AllowedOptimizations {
+                projection_pushdown: false,
+                predicate_pushdown: true,
+                type_coercion: true,
+                simplify_expr: true,
+                file_caching: true,
+                aggregate_pushdown: true,
+                slice_pushdown: true,
+                #[cfg(feature = "cse")]
+                common_subplan_elimination: true,
+            }),
             Some(Arc::new(udf_schema)),
             Some("unnest"),
         )

diff --git a/py-polars/tests/unit/test_projections.py b/py-polars/tests/unit/test_projections.py
@@ -39,3 +39,26 @@ def test_double_projection_pushdown() -> None:
             .select(["c0", "c1"])
         ).describe_optimized_plan()
     )
+
+
+def test_unnest_projection_pushdown() -> None:
+    lf = pl.DataFrame({"x|y|z": [1, 2], "a|b|c": [2, 3]}).lazy()
+
+    mlf = (
+        lf.melt()
+        .with_column(pl.col("variable").str.split_exact("|", 2))
+        .unnest("variable")
+    )
+    mlf = mlf.select(
+        [
+            pl.col("field_1").cast(pl.Categorical).alias("row"),
+            pl.col("field_2").cast(pl.Categorical).alias("col"),
+            pl.col("value"),
+        ]
+    )
+    out = mlf.collect().to_dict(False)
+    assert out == {
+        "row": ["y", "y", "b", "b"],
+        "col": ["z", "z", "c", "c"],
+        "value": [1, 2, 2, 3],
+    }