fix tests

pola-rs · Feb 10, 2025 · 8ec1326 · 8ec1326
1 parent 5438c4a
commit 8ec1326
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/py-polars/tests/unit/lazyframe/test_explain.py b/py-polars/tests/unit/lazyframe/test_explain.py
@@ -10,23 +10,24 @@ def test_lf_explain_format_tree() -> None:
     plan = lf.select("a").select(pl.col("a").sum() + pl.len())
 
     result = plan.explain(format="tree")
+    print(result)
 
     expected = """\
-                0                        1
-   ┌─────────────────────────────────────────────────
+                0                            1
+   ┌─────────────────────────────────────────────────────────
    │
    │        ╭────────╮
  0 │        │ SELECT │
    │        ╰───┬┬───╯
    │            ││
-   │            │╰───────────────────────╮
-   │            │                        │
-   │  ╭─────────┴──────────╮             │
-   │  │ expression:        │  ╭──────────┴──────────╮
-   │  │ [(col("a")         │  │ FROM:               │
- 1 │  │   .sum()) + (len() │  │ DF ["a", "b"]       │
-   │  │   .cast(Int64))]   │  │ PROJECT 1/2 COLUMNS │
-   │  ╰────────────────────╯  ╰─────────────────────╯
+   │            │╰───────────────────────────╮
+   │            │                            │
+   │  ╭─────────┴──────────╮                 │
+   │  │ expression:        │  ╭──────────────┴──────────────╮
+   │  │ [(col("a")         │  │ FROM:                       │
+ 1 │  │   .sum()) + (len() │  │ DF ["a", "b"]               │
+   │  │   .cast(Int64))]   │  │ PROJECT: ["a"]; 1/2 COLUMNS │
+   │  ╰────────────────────╯  ╰─────────────────────────────╯
 \
 """
     assert result == expected

diff --git a/py-polars/tests/unit/test_projections.py b/py-polars/tests/unit/test_projections.py
@@ -37,7 +37,7 @@ def test_unpivot_projection_pd_block_4997() -> None:
 
 def test_double_projection_pushdown() -> None:
     assert (
-        "PROJECT 2/3 COLUMNS"
+        "2/3 COLUMNS"
         in (
             pl.DataFrame({"c0": [], "c1": [], "c2": []})
             .lazy()
@@ -49,7 +49,7 @@ def test_double_projection_pushdown() -> None:
 
 def test_group_by_projection_pushdown() -> None:
     assert (
-        "PROJECT 2/3 COLUMNS"
+        "2/3 COLUMNS"
         in (
             pl.DataFrame({"c0": [], "c1": [], "c2": []})
             .lazy()
@@ -100,7 +100,7 @@ def test_hconcat_projection_pushdown() -> None:
     query = pl.concat([lf1, lf2], how="horizontal").select(["a", "d"])
 
     explanation = query.explain()
-    assert explanation.count("PROJECT 1/2 COLUMNS") == 2
+    assert explanation.count("1/2 COLUMNS") == 2
 
     out = query.collect()
     expected = pl.DataFrame({"a": [0, 1, 2], "d": [9, 10, 11]})
@@ -115,7 +115,7 @@ def test_hconcat_projection_pushdown_length_maintained() -> None:
     query = pl.concat([lf1, lf2], how="horizontal").select(["a"])
 
     explanation = query.explain()
-    assert "PROJECT 1/2 COLUMNS" in explanation
+    assert "1/2 COLUMNS" in explanation
 
     out = query.collect()
     expected = pl.DataFrame({"a": [0, 1, None, None]})
@@ -406,7 +406,7 @@ def test_rolling_key_projected_13617() -> None:
     df = pl.DataFrame({"idx": [1, 2], "value": ["a", "b"]}).set_sorted("idx")
     ldf = df.lazy().select(pl.col("value").rolling("idx", period="1i"))
     plan = ldf.explain(projection_pushdown=True)
-    assert r'DF ["idx", "value"]; PROJECT 2/2 COLUMNS' in plan
+    assert r"2/2 COLUMNS" in plan
     out = ldf.collect(projection_pushdown=True)
     assert out.to_dict(as_series=False) == {"value": [["a"], ["b"]]}
 
@@ -526,7 +526,7 @@ def test_projection_pushdown_semi_anti_no_selection(
 
     q_b = pl.LazyFrame({"b": [1, 2, 3], "c": [1, 2, 3]})
 
-    assert "PROJECT 1/2" in (
+    assert "1/2 COLUMNS" in (
         q_a.join(q_b, left_on="a", right_on="b", how=how).explain()
     )
 
@@ -536,7 +536,7 @@ def test_projection_empty_frame_len_16904() -> None:
 
     q = df.select(pl.len())
 
-    assert "PROJECT 0/0" in q.explain()
+    assert "0/0 COLUMNS" in q.explain()
 
     expect = pl.DataFrame({"len": [0]}, schema_overrides={"len": pl.UInt32()})
     assert_frame_equal(q.collect(), expect)
@@ -612,7 +612,7 @@ def test_with_columns_projection_pushdown() -> None:
     assert plan.startswith("WITH_COLUMNS:")
     # [dyn int: 1.alias("x"), dyn int: 1.alias("y")]
     # Csv SCAN [20 in-mem bytes]
-    assert plan.endswith("PROJECT 1/6 COLUMNS")
+    assert plan.endswith("1/6 COLUMNS")
 
 
 def test_projection_pushdown_height_20221() -> None:
@@ -641,7 +641,7 @@ def test_filter_count_projection_20902() -> None:
         }
     )
     assert (
-        "PROJECT 1/3"
+        "1/3 COLUMNS"
         in lineitem_ldf.filter(pl.col("l_partkey").is_between(10, 20))
         .select(pl.len())
         .explain()