Repository: databricks/koalas

Issue:
Support bool list-like object for column selection in `loc` indexer.
We should support bool list-like object for `loc` indexer.

```py
>>> df
     a    b
0  1.0  NaN
1  2.0  1.0
2  NaN  NaN
>>> df.columns.str.contains("a", regex=True)
array([ True, False])
>>> df.loc[:, df.columns.str.contains("a", regex=True)]
Traceback (most recent call last):
...
KeyError: "['True'] not in index"
```

whereas:

```py
>>> pdf.loc[:, pdf.columns.str.contains("a", regex=True)]
     a
0  1.0
1  2.0
2  NaN
```

See also: https://github.com/databricks/koalas/issues/2043#issuecomment-774278358


Discussion:

Resolving patch:
diff --git a/databricks/koalas/indexing.py b/databricks/koalas/indexing.py
--- a/databricks/koalas/indexing.py
+++ b/databricks/koalas/indexing.py
@@ -20,8 +20,9 @@
 from abc import ABCMeta, abstractmethod
 from collections.abc import Iterable
 from functools import reduce
-from typing import Any, Optional, List, Tuple, TYPE_CHECKING, Union
+from typing import Any, Optional, List, Tuple, TYPE_CHECKING, Union, cast, Sized
 
+import pandas as pd
 from pandas.api.types import is_list_like
 from pyspark import sql as spark
 from pyspark.sql import functions as F
@@ -756,9 +757,15 @@ class LocIndexer(LocIndexerLike):
 
     - A conditional boolean Series derived from the DataFrame or Series
 
+    - A boolean array of the same length as the column axis being sliced,
+      e.g. ``[True, False, True]``.
+
+    - An alignable boolean pandas Series to the column axis being sliced.
+      The index of the key will be aligned before masking.
+
     Not allowed inputs which pandas allows are:
 
-    - A boolean array of the same length as the axis being sliced,
+    - A boolean array of the same length as the row axis being sliced,
       e.g. ``[True, False, True]``.
     - A ``callable`` function with one argument (the calling Series, DataFrame
       or Panel) and that returns valid output for indexing (one of the above)
@@ -851,6 +858,22 @@ class LocIndexer(LocIndexerLike):
                 max_speed
     sidewinder          7
 
+    A boolean array of the same length as the column axis being sliced.
+
+    >>> df.loc[:, [False, True]]
+                shield
+    cobra            2
+    viper            5
+    sidewinder       8
+
+    An alignable boolean Series to the column axis being sliced.
+
+    >>> df.loc[:, pd.Series([False, True], index=['max_speed', 'shield'])]
+                shield
+    cobra            2
+    viper            5
+    sidewinder       8
+
     **Setting values**
 
     Setting value for all items matching the list of labels.
@@ -1201,6 +1224,43 @@ def _select_cols_by_iterable(
             ]
             data_spark_columns = list(cols_sel)
             data_dtypes = None
+        elif all(isinstance(key, bool) for key in cols_sel) or all(
+            isinstance(key, np.bool_) for key in cols_sel
+        ):
+            if len(cast(Sized, cols_sel)) != len(self._internal.column_labels):
+                raise IndexError(
+                    "Boolean index has wrong length: %s instead of %s"
+                    % (len(cast(Sized, cols_sel)), len(self._internal.column_labels))
+                )
+            if isinstance(cols_sel, pd.Series):
+                if not cols_sel.index.sort_values().equals(self._kdf.columns.sort_values()):
+                    raise SparkPandasIndexingError(
+                        "Unalignable boolean Series provided as indexer "
+                        "(index of the boolean Series and of the indexed object do not match)"
+                    )
+                else:
+                    column_labels = [
+                        column_label
+                        for column_label in self._internal.column_labels
+                        if cols_sel[column_label if len(column_label) > 1 else column_label[0]]
+                    ]
+                    data_spark_columns = [
+                        self._internal.spark_column_for(column_label)
+                        for column_label in column_labels
+                    ]
+                    data_dtypes = [
+                        self._internal.dtype_for(column_label) for column_label in column_labels
+                    ]
+            else:
+                column_labels = [
+                    self._internal.column_labels[i] for i, col in enumerate(cols_sel) if col
+                ]
+                data_spark_columns = [
+                    self._internal.data_spark_columns[i] for i, col in enumerate(cols_sel) if col
+                ]
+                data_dtypes = [
+                    self._internal.data_dtypes[i] for i, col in enumerate(cols_sel) if col
+                ]
         elif any(isinstance(key, tuple) for key in cols_sel) and any(
             not is_name_like_tuple(key) for key in cols_sel
         ):


Test patch:
diff --git a/databricks/koalas/tests/test_indexing.py b/databricks/koalas/tests/test_indexing.py
--- a/databricks/koalas/tests/test_indexing.py
+++ b/databricks/koalas/tests/test_indexing.py
@@ -593,6 +593,20 @@ def test_loc2d(self):
         self.assert_eq(kdf.loc[:, "a":"d"], pdf.loc[:, "a":"d"])
         self.assert_eq(kdf.loc[:, "c":"d"], pdf.loc[:, "c":"d"])
 
+        # bool list-like column select
+        bool_list = [True, False]
+        self.assert_eq(kdf.loc[:, bool_list], pdf.loc[:, bool_list])
+        self.assert_eq(kdf.loc[:, np.array(bool_list)], pdf.loc[:, np.array(bool_list)])
+
+        pser = pd.Series(bool_list, index=pdf.columns)
+        self.assert_eq(kdf.loc[:, pser], pdf.loc[:, pser])
+        pser = pd.Series(list(reversed(bool_list)), index=list(reversed(pdf.columns)))
+        self.assert_eq(kdf.loc[:, pser], pdf.loc[:, pser])
+
+        self.assertRaises(IndexError, lambda: kdf.loc[:, bool_list[:-1]])
+        self.assertRaises(IndexError, lambda: kdf.loc[:, np.array(bool_list + [True])])
+        self.assertRaises(SparkPandasIndexingError, lambda: kdf.loc[:, pd.Series(bool_list)])
+
         # non-string column names
         kdf = self.kdf2
         pdf = self.pdf2
@@ -642,6 +656,17 @@ def test_loc2d_multiindex_columns(self):
         self.assertRaises(KeyError, lambda: kdf.loc[:, "bar":("baz", "one")])
         self.assertRaises(KeyError, lambda: kdf.loc[:, ("bar", "two"):"bar"])
 
+        # bool list-like column select
+        bool_list = [True, False, True, False]
+        self.assert_eq(kdf.loc[:, bool_list], pdf.loc[:, bool_list])
+        self.assert_eq(kdf.loc[:, np.array(bool_list)], pdf.loc[:, np.array(bool_list)])
+
+        pser = pd.Series(bool_list, index=pdf.columns)
+        self.assert_eq(kdf.loc[:, pser], pdf.loc[:, pser])
+
+        pser = pd.Series(list(reversed(bool_list)), index=list(reversed(pdf.columns)))
+        self.assert_eq(kdf.loc[:, pser], pdf.loc[:, pser])
+
         # non-string column names
         arrays = [np.array([0, 0, 1, 1]), np.array([1, 2, 1, 2])]