python eager: allow masking of columns in indexing

pola-rs · Oct 29, 2021 · 1f1a949 · 1f1a949
1 parent 57658c0
commit 1f1a949
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/py-polars/polars/eager/frame.py b/py-polars/polars/eager/frame.py
@@ -941,6 +941,21 @@ def __getitem__(self, item: Any) -> Any:
                     df = self.__getitem__(self.columns[col_selection])
                     return df[row_selection]
 
+                # slice and boolean mask
+                # df[:2, [True, False, True]]
+                if isinstance(col_selection, (Sequence, pl.Series)):
+                    if (
+                        isinstance(col_selection[0], bool)
+                        or isinstance(col_selection, pl.Series)
+                        and col_selection.dtype() == pl.Boolean
+                    ):
+                        df = self.__getitem__(row_selection)
+                        select = []
+                        for col, valid in zip(df.columns, col_selection):
+                            if valid:
+                                select.append(col)
+                        return df.select(select)
+
                 # single slice
                 # df[:, unknown]
                 series = self.__getitem__(col_selection)

diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -212,6 +212,10 @@ def test_selection():
     # get_column by name
     assert df.get_column("a").to_list() == [1, 2, 3]
 
+    # select columns by mask
+    assert df[:2, [True, False, False]].shape == (2, 1)
+    assert df[:2, pl.Series([True, False, False])].shape == (2, 1)
+
     # column selection by string(s) in first dimension
     assert df["a"].to_list() == [1, 2, 3]
     assert df["b"].to_list() == [1.0, 2.0, 3.0]