more tests (#4163)

pola-rs · Jul 26, 2022 · d51d4a8 · d51d4a8
1 parent ba400f4
commit d51d4a8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -1982,3 +1982,14 @@ def test_with_columns() -> None:
     # at least one of exprs/**named_exprs required
     with pytest.raises(ValueError):
         _ = df.with_columns()
+
+
+def test_len_compute(df: pl.DataFrame) -> None:
+    df = df.with_column(pl.struct(["list_bool", "cat"]).alias("struct"))
+    filtered = df.filter(pl.col("bools"))
+    for col in filtered.columns:
+        assert len(filtered[col]) == 1
+
+    taken = df[[1, 2], :]
+    for col in taken.columns:
+        assert len(taken[col]) == 2
diff --git a/py-polars/tests/test_pivot.py b/py-polars/tests/test_pivot.py
@@ -75,6 +75,10 @@ def test_pivot_categorical_index() -> None:
         columns=[("A", pl.Categorical), ("B", pl.Categorical)],
     )
 
+    assert df.pivot(values="B", index=["A"], columns="B", aggregate_fn="count").to_dict(
+        False
+    ) == {"A": ["Fire", "Water"], "Car": [1, 2], "Ship": [1, None]}
+
     df = pl.DataFrame(
         {
             "A": ["Fire", "Water", "Water", "Fire"],