add coverage for expanding exprs and regex selection

pola-rs · Jan 28, 2023 · d6451d5 · d6451d5
1 parent 65e7888
commit d6451d5
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 17 deletions.
diff --git a/py-polars/tests/unit/test_df.py b/py-polars/tests/unit/test_df.py
@@ -2482,27 +2482,51 @@ def test_with_columns() -> None:
     assert_frame_equal(dx, expected)
 
     # automatically upconvert multi-output expressions to struct
-    df = pl.DataFrame({"x1": [1, 2, 6], "x2": [1, 2, 3]}).with_columns(
-        pct_change=pl.col(["x1", "x2"]).pct_change()
-    )
-    # ┌─────┬─────┬────────────────┐
-    # │ x1  ┆ x2  ┆ pct_change     │
-    # │ --- ┆ --- ┆ ---            │
-    # │ i64 ┆ i64 ┆ struct[2]      │
-    # ╞═════╪═════╪════════════════╡
-    # │ 1   ┆ 1   ┆ {null,null}    │
-    # │ 2   ┆ 2   ┆ {1.0,1.0}      │
-    # │ 6   ┆ 3   ┆ {2.0,0.5}      │
-    # └─────┴─────┴────────────────┘
-    assert df.to_dicts() == [
-        {"x1": 1, "x2": 1, "pct_change": {"x1": None, "x2": None}},
-        {"x1": 2, "x2": 2, "pct_change": {"x1": 1.0, "x2": 1.0}},
-        {"x1": 6, "x2": 3, "pct_change": {"x1": 2.0, "x2": 0.5}},
+    ldf = (
+        pl.DataFrame({"x1": [1, 2, 6], "x2": [1, 2, 3]})
+        .lazy()
+        .with_columns(
+            pct_change=pl.col(["x1", "x2"]).pct_change(),
+            maxes=pl.all().max().suffix("_max"),
+            xcols=pl.col("^x.*$"),
+        )
+    )
+    # ┌─────┬─────┬─────────────┬───────────┬───────────┐
+    # │ x1  ┆ x2  ┆ pct_change  ┆ maxes     ┆ xcols     │
+    # │ --- ┆ --- ┆ ---         ┆ ---       ┆ ---       │
+    # │ i64 ┆ i64 ┆ struct[2]   ┆ struct[2] ┆ struct[2] │
+    # ╞═════╪═════╪═════════════╪═══════════╪═══════════╡
+    # │ 1   ┆ 1   ┆ {null,null} ┆ {6,3}     ┆ {1,1}     │
+    # │ 2   ┆ 2   ┆ {1.0,1.0}   ┆ {6,3}     ┆ {2,2}     │
+    # │ 6   ┆ 3   ┆ {2.0,0.5}   ┆ {6,3}     ┆ {6,3}     │
+    # └─────┴─────┴─────────────┴───────────┴───────────┘
+    assert ldf.collect().to_dicts() == [
+        {
+            "x1": 1,
+            "x2": 1,
+            "pct_change": {"x1": None, "x2": None},
+            "maxes": {"x1_max": 6, "x2_max": 3},
+            "xcols": {"x1": 1, "x2": 1},
+        },
+        {
+            "x1": 2,
+            "x2": 2,
+            "pct_change": {"x1": 1.0, "x2": 1.0},
+            "maxes": {"x1_max": 6, "x2_max": 3},
+            "xcols": {"x1": 2, "x2": 2},
+        },
+        {
+            "x1": 6,
+            "x2": 3,
+            "pct_change": {"x1": 2.0, "x2": 0.5},
+            "maxes": {"x1_max": 6, "x2_max": 3},
+            "xcols": {"x1": 6, "x2": 3},
+        },
     ]
 
     # require at least one of exprs / **named_exprs
     with pytest.raises(ValueError):
-        _ = df.with_columns()
+        _ = ldf.with_columns()
 
 
 def test_len_compute(df: pl.DataFrame) -> None:

diff --git a/py-polars/tests/unit/test_meta.py b/py-polars/tests/unit/test_meta.py
@@ -58,3 +58,4 @@ def test_meta_has_multiple_outputs() -> None:
 def test_meta_is_regex_projection() -> None:
     e = pl.col("^.*$").alias("bar")
     assert e.meta.is_regex_projection()
+    assert e.meta.has_multiple_outputs()