test: ensure expected data has every column of theresult data (#2296)

Author: EdAbati

CONTRIBUTING.md

@@ -109,7 +109,7 @@ If you want to run PySpark-related tests, you'll need to have Java installed. Re
 
 2. Install Narwhals: `uv pip install -e . --group local-dev"`. This will include fast-ish core libraries and dev dependencies.
    If you also want to test other libraries like Dask , PySpark, and Modin, you can install them too with
-   `uv pip install -e ".[dask, pyspark, modin]" --group local-dev"`.
+   `uv pip install -e ".[dask, pyspark, modin]" --group local-dev`.
 
 You should also install pre-commit:
 

tests/expr_and_series/cum_count_test.py

@@ -75,14 +75,19 @@ def test_lazy_cum_count_grouped(
         constructor(
             {
                 "arg entina": [None, 2, 3],
-                "ban gkock": [1, 0, 2],
+                "ban gkok": [1, 0, 2],
                 "i ran": [0, 1, 2],
                 "g": [1, 1, 1],
             }
         )
     )
     result = df.with_columns(
-        nw.col("arg entina").cum_count(reverse=reverse).over("g", order_by="ban gkock")
+        nw.col("arg entina").cum_count(reverse=reverse).over("g", order_by="ban gkok")
     ).sort("i ran")
-    expected = {"arg entina": expected_a, "ban gkock": [1, 0, 2], "i ran": [0, 1, 2]}
+    expected = {
+        "arg entina": expected_a,
+        "ban gkok": [1, 0, 2],
+        "i ran": [0, 1, 2],
+        "g": [1, 1, 1],
+    }
     assert_equal_data(result, expected)

tests/expr_and_series/cum_max_test.py

@@ -83,7 +83,7 @@ def test_lazy_cum_max_grouped(
     result = df.with_columns(
         nw.col("a").cum_max(reverse=reverse).over("g", order_by="b")
     ).sort("i")
-    expected = {"a": expected_a, "b": [1, 0, 2], "i": [0, 1, 2]}
+    expected = {"a": expected_a, "b": [1, 0, 2], "i": [0, 1, 2], "g": [1, 1, 1]}
     assert_equal_data(result, expected)
 
 
@@ -134,6 +134,7 @@ def test_lazy_cum_max_ordered_by_nulls(
         "a": expected_a,
         "b": [1, -1, 3, 2, 5, 0, None],
         "i": [0, 1, 2, 3, 4, 5, 6],
+        "g": [1, 1, 1, 1, 1, 1, 1],
     }
     assert_equal_data(result, expected)
 

tests/expr_and_series/cum_min_test.py

@@ -83,7 +83,7 @@ def test_lazy_cum_min_grouped(
     result = df.with_columns(
         nw.col("a").cum_min(reverse=reverse).over("g", order_by="b")
     ).sort("i")
-    expected = {"a": expected_a, "b": [1, 0, 2], "i": [0, 1, 2]}
+    expected = {"a": expected_a, "b": [1, 0, 2], "i": [0, 1, 2], "g": [1, 1, 1]}
     assert_equal_data(result, expected)
 
 
@@ -134,6 +134,7 @@ def test_lazy_cum_min_ordered_by_nulls(
         "a": expected_a,
         "b": [1, -1, 3, 2, 5, 0, None],
         "i": [0, 1, 2, 3, 4, 5, 6],
+        "g": [1, 1, 1, 1, 1, 1, 1],
     }
     assert_equal_data(result, expected)
 

tests/expr_and_series/cum_prod_test.py

@@ -88,22 +88,24 @@ def test_lazy_cum_prod_grouped(
     if "cudf" in str(constructor):
         # https://github.com/rapidsai/cudf/issues/18159
         request.applymarker(pytest.mark.xfail)
-    if "sqlframe" in str(constructor):
-        # https://github.com/eakmanrq/sqlframe/issues/348
-        request.applymarker(pytest.mark.xfail)
 
     df = nw.from_native(
         constructor(
             {
                 "arg entina": [1, 2, 3],
-                "ban gkock": [1, 0, 2],
+                "ban gkok": [1, 0, 2],
                 "i ran": [0, 1, 2],
                 "g": [1, 1, 1],
             }
         )
     )
     result = df.with_columns(
-        nw.col("arg entina").cum_prod(reverse=reverse).over("g", order_by="ban gkock")
+        nw.col("arg entina").cum_prod(reverse=reverse).over("g", order_by="ban gkok")
     ).sort("i ran")
-    expected = {"arg entina": expected_a, "ban gkock": [1, 0, 2], "i ran": [0, 1, 2]}
+    expected = {
+        "arg entina": expected_a,
+        "ban gkok": [1, 0, 2],
+        "i ran": [0, 1, 2],
+        "g": [1, 1, 1],
+    }
     assert_equal_data(result, expected)

tests/expr_and_series/cum_sum_test.py

@@ -62,16 +62,21 @@ def test_lazy_cum_sum_grouped(
         constructor(
             {
                 "arg entina": [1, 2, 3],
-                "ban gkock": [1, 0, 2],
+                "ban gkok": [1, 0, 2],
                 "i ran": [0, 1, 2],
                 "g": [1, 1, 1],
             }
         )
     )
     result = df.with_columns(
-        nw.col("arg entina").cum_sum(reverse=reverse).over("g", order_by="ban gkock")
+        nw.col("arg entina").cum_sum(reverse=reverse).over("g", order_by="ban gkok")
     ).sort("i ran")
-    expected = {"arg entina": expected_a, "ban gkock": [1, 0, 2], "i ran": [0, 1, 2]}
+    expected = {
+        "arg entina": expected_a,
+        "ban gkok": [1, 0, 2],
+        "i ran": [0, 1, 2],
+        "g": [1, 1, 1],
+    }
     assert_equal_data(result, expected)
 
 
@@ -110,19 +115,20 @@ def test_lazy_cum_sum_ordered_by_nulls(
         constructor(
             {
                 "arg entina": [1, 2, 3, 1, 2, 3, 4],
-                "ban gkock": [1, -1, 3, 2, 5, 0, None],
+                "ban gkok": [1, -1, 3, 2, 5, 0, None],
                 "i ran": [0, 1, 2, 3, 4, 5, 6],
                 "g": [1, 1, 1, 1, 1, 1, 1],
             }
         )
     )
     result = df.with_columns(
-        nw.col("arg entina").cum_sum(reverse=reverse).over("g", order_by="ban gkock")
+        nw.col("arg entina").cum_sum(reverse=reverse).over("g", order_by="ban gkok")
     ).sort("i ran")
     expected = {
         "arg entina": expected_a,
-        "ban gkock": [1, -1, 3, 2, 5, 0, None],
+        "ban gkok": [1, -1, 3, 2, 5, 0, None],
         "i ran": [0, 1, 2, 3, 4, 5, 6],
+        "g": [1, 1, 1, 1, 1, 1, 1],
     }
     assert_equal_data(result, expected)
 
@@ -156,15 +162,15 @@ def test_lazy_cum_sum_ungrouped(
         constructor(
             {
                 "arg entina": [2, 3, 1],
-                "ban gkock": [0, 2, 1],
+                "ban gkok": [0, 2, 1],
                 "i ran": [1, 2, 0],
             }
         )
     ).sort("i ran")
     result = df.with_columns(
-        nw.col("arg entina").cum_sum(reverse=reverse).over(order_by="ban gkock")
+        nw.col("arg entina").cum_sum(reverse=reverse).over(order_by="ban gkok")
     ).sort("i ran")
-    expected = {"arg entina": expected_a, "ban gkock": [1, 0, 2], "i ran": [0, 1, 2]}
+    expected = {"arg entina": expected_a, "ban gkok": [1, 0, 2], "i ran": [0, 1, 2]}
     assert_equal_data(result, expected)
 
 
@@ -197,17 +203,17 @@ def test_lazy_cum_sum_ungrouped_ordered_by_nulls(
         constructor(
             {
                 "arg entina": [1, 2, 3, 1, 2, 3, 4],
-                "ban gkock": [1, -1, 3, 2, 5, 0, None],
+                "ban gkok": [1, -1, 3, 2, 5, 0, None],
                 "i ran": [0, 1, 2, 3, 4, 5, 6],
             }
         )
     ).sort("i ran")
     result = df.with_columns(
-        nw.col("arg entina").cum_sum(reverse=reverse).over(order_by="ban gkock")
+        nw.col("arg entina").cum_sum(reverse=reverse).over(order_by="ban gkok")
     ).sort("i ran")
     expected = {
         "arg entina": expected_a,
-        "ban gkock": [1, -1, 3, 2, 5, 0, None],
+        "ban gkok": [1, -1, 3, 2, 5, 0, None],
         "i ran": [0, 1, 2, 3, 4, 5, 6],
     }
     assert_equal_data(result, expected)

tests/expr_and_series/over_test.py

@@ -88,17 +88,12 @@ def test_over_multiple(constructor: Constructor) -> None:
     if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3):
         pytest.skip()
     df = nw.from_native(constructor(data))
-    expected = {
-        "a": ["a", "a", "b", "b", "b"],
-        "b": [1, 2, 3, 3, 5],
-        "c": [5, 4, 3, 1, 2],
-        "i": list(range(5)),
-        "c_min": [5, 4, 1, 1, 2],
-    }
     expected = {
         "a": ["a", "a", "b", "b", "b"],
         "b": [1, 2, 3, 5, 3],
         "c": [5, 4, 3, 2, 1],
+        "i": list(range(5)),
+        "c_min": [5, 4, 1, 2, 1],
     }
 
     result = df.with_columns(c_min=nw.col("c").min().over("a", "b")).sort("i")

tests/utils.py

@@ -48,7 +48,7 @@ def get_module_version_as_tuple(module_name: str) -> tuple[int, ...]:
 
 def zip_strict(left: Sequence[Any], right: Sequence[Any]) -> Iterator[Any]:
     if len(left) != len(right):
-        msg = f"left {len(left)=} != right {len(right)=}"  # pragma: no cover
+        msg = f"{len(left)=} != {len(right)=}\nLeft: {left}\nRight: {right}"  # pragma: no cover
         raise ValueError(msg)  # pragma: no cover
     return zip(left, right)
 
@@ -70,21 +70,6 @@ def _to_comparable_list(column_values: Any) -> Any:
     return list(column_values)
 
 
-def _sort_dict_by_key(
-    data_dict: Mapping[str, list[Any]], key: str
-) -> dict[str, list[Any]]:  # pragma: no cover
-    sort_list = data_dict[key]
-    sorted_indices = sorted(
-        range(len(sort_list)),
-        key=lambda i: (
-            (sort_list[i] is None)
-            or (isinstance(sort_list[i], float) and math.isnan(sort_list[i])),
-            sort_list[i],
-        ),
-    )
-    return {key: [value[i] for i in sorted_indices] for key, value in data_dict.items()}
-
-
 def assert_equal_data(result: Any, expected: Mapping[str, Any]) -> None:
     is_duckdb = (
         hasattr(result, "_compliant_frame")
@@ -103,7 +88,9 @@ def assert_equal_data(result: Any, expected: Mapping[str, Any]) -> None:
         result = result.collect(**kwargs.get(result.implementation, {}))
 
     if hasattr(result, "columns"):
-        for idx, (col, key) in enumerate(zip(result.columns, expected.keys())):
+        for idx, (col, key) in enumerate(
+            zip_strict(result.columns, list(expected.keys()))
+        ):
             assert col == key, f"Expected column name {key} at index {idx}, found {col}"
     result = {key: _to_comparable_list(result[key]) for key in expected}
     assert list(result.keys()) == list(expected.keys()), (