Repository: databricks/koalas

Issue:
Can't select multiple columns when a dataframe has multi-index columns
```
import databricks.koalas as ks
import numpy as np

arrays = [np.array(['bar', 'bar', 'baz', 'baz']),
          np.array(['one', 'two', 'one', 'two'])]

kdf = ks.DataFrame(np.random.randn(3, 4), index=['A', 'B', 'C'], columns=arrays)
kdf[['bar', 'baz']]
```

This code gave:

```
---------------------------------------------------------------------------
Py4JJavaError                             Traceback (most recent call last)
~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/pyspark/sql/utils.py in deco(*a, **kw)
     62         try:
---> 63             return f(*a, **kw)
     64         except py4j.protocol.Py4JJavaError as e:

~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/py4j/protocol.py in get_return_value(answer, gateway_client, target_id, name)
    327                     "An error occurred while calling {0}{1}{2}.\n".
--> 328                     format(target_id, ".", name), value)
    329             else:

Py4JJavaError: An error occurred while calling o715.select.
: org.apache.spark.sql.AnalysisException: cannot resolve '`bar`' given input columns: [('bar', 'one'), __index_level_0__, ('bar', 'two'), ('baz', 'one'), ('baz', 'two')];;
'Project [__index_level_0__#147, 'bar, 'baz]
+- LogicalRDD [__index_level_0__#147, ('bar', 'one')#148, ('bar', 'two')#149, ('baz', 'one')#150, ('baz', 'two')#151], false

	at org.apache.spark.sql.catalyst.analysis.package$AnalysisErrorAt.failAnalysis(package.scala:42)
	at org.apache.spark.sql.catalyst.analysis.CheckAnalysis$$anonfun$checkAnalysis$1$$anonfun$apply$3.applyOrElse(CheckAnalysis.scala:110)
	at org.apache.spark.sql.catalyst.analysis.CheckAnalysis$$anonfun$checkAnalysis$1$$anonfun$apply$3.applyOrElse(CheckAnalysis.scala:107)
	at org.apache.spark.sql.catalyst.trees.TreeNode$$anonfun$transformUp$1.apply(TreeNode.scala:278)
	at org.apache.spark.sql.catalyst.trees.TreeNode$$anonfun$transformUp$1.apply(TreeNode.scala:278)
	at org.apache.spark.sql.catalyst.trees.CurrentOrigin$.withOrigin(TreeNode.scala:70)
	at org.apache.spark.sql.catalyst.trees.TreeNode.transformUp(TreeNode.scala:277)
	at org.apache.spark.sql.catalyst.plans.QueryPlan$$anonfun$transformExpressionsUp$1.apply(QueryPlan.scala:93)
	at org.apache.spark.sql.catalyst.plans.QueryPlan$$anonfun$transformExpressionsUp$1.apply(QueryPlan.scala:93)
	at org.apache.spark.sql.catalyst.plans.QueryPlan$$anonfun$1.apply(QueryPlan.scala:105)
	at org.apache.spark.sql.catalyst.plans.QueryPlan$$anonfun$1.apply(QueryPlan.scala:105)
	at org.apache.spark.sql.catalyst.trees.CurrentOrigin$.withOrigin(TreeNode.scala:70)
	at org.apache.spark.sql.catalyst.plans.QueryPlan.transformExpression$1(QueryPlan.scala:104)
	at org.apache.spark.sql.catalyst.plans.QueryPlan.org$apache$spark$sql$catalyst$plans$QueryPlan$$recursiveTransform$1(QueryPlan.scala:116)
	at org.apache.spark.sql.catalyst.plans.QueryPlan$$anonfun$org$apache$spark$sql$catalyst$plans$QueryPlan$$recursiveTransform$1$2.apply(QueryPlan.scala:121)
	at scala.collection.TraversableLike$$anonfun$map$1.apply(TraversableLike.scala:234)
	at scala.collection.TraversableLike$$anonfun$map$1.apply(TraversableLike.scala:234)
	at scala.collection.mutable.ResizableArray$class.foreach(ResizableArray.scala:59)
	at scala.collection.mutable.ArrayBuffer.foreach(ArrayBuffer.scala:48)
	at scala.collection.TraversableLike$class.map(TraversableLike.scala:234)
	at scala.collection.AbstractTraversable.map(Traversable.scala:104)
	at org.apache.spark.sql.catalyst.plans.QueryPlan.org$apache$spark$sql$catalyst$plans$QueryPlan$$recursiveTransform$1(QueryPlan.scala:121)
	at org.apache.spark.sql.catalyst.plans.QueryPlan$$anonfun$2.apply(QueryPlan.scala:126)
	at org.apache.spark.sql.catalyst.trees.TreeNode.mapProductIterator(TreeNode.scala:187)
	at org.apache.spark.sql.catalyst.plans.QueryPlan.mapExpressions(QueryPlan.scala:126)
	at org.apache.spark.sql.catalyst.plans.QueryPlan.transformExpressionsUp(QueryPlan.scala:93)
	at org.apache.spark.sql.catalyst.analysis.CheckAnalysis$$anonfun$checkAnalysis$1.apply(CheckAnalysis.scala:107)
	at org.apache.spark.sql.catalyst.analysis.CheckAnalysis$$anonfun$checkAnalysis$1.apply(CheckAnalysis.scala:85)
	at org.apache.spark.sql.catalyst.trees.TreeNode.foreachUp(TreeNode.scala:127)
	at org.apache.spark.sql.catalyst.analysis.CheckAnalysis$class.checkAnalysis(CheckAnalysis.scala:85)
	at org.apache.spark.sql.catalyst.analysis.Analyzer.checkAnalysis(Analyzer.scala:95)
	at org.apache.spark.sql.catalyst.analysis.Analyzer$$anonfun$executeAndCheck$1.apply(Analyzer.scala:108)
	at org.apache.spark.sql.catalyst.analysis.Analyzer$$anonfun$executeAndCheck$1.apply(Analyzer.scala:105)
	at org.apache.spark.sql.catalyst.plans.logical.AnalysisHelper$.markInAnalyzer(AnalysisHelper.scala:201)
	at org.apache.spark.sql.catalyst.analysis.Analyzer.executeAndCheck(Analyzer.scala:105)
	at org.apache.spark.sql.execution.QueryExecution.analyzed$lzycompute(QueryExecution.scala:57)
	at org.apache.spark.sql.execution.QueryExecution.analyzed(QueryExecution.scala:55)
	at org.apache.spark.sql.execution.QueryExecution.assertAnalyzed(QueryExecution.scala:47)
	at org.apache.spark.sql.Dataset$.ofRows(Dataset.scala:78)
	at org.apache.spark.sql.Dataset.org$apache$spark$sql$Dataset$$withPlan(Dataset.scala:3406)
	at org.apache.spark.sql.Dataset.select(Dataset.scala:1334)
	at sun.reflect.GeneratedMethodAccessor62.invoke(Unknown Source)
	at sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
	at java.lang.reflect.Method.invoke(Method.java:498)
	at py4j.reflection.MethodInvoker.invoke(MethodInvoker.java:244)
	at py4j.reflection.ReflectionEngine.invoke(ReflectionEngine.java:357)
	at py4j.Gateway.invoke(Gateway.java:282)
	at py4j.commands.AbstractCommand.invokeMethod(AbstractCommand.java:132)
	at py4j.commands.CallCommand.execute(CallCommand.java:79)
	at py4j.GatewayConnection.run(GatewayConnection.java:238)
	at java.lang.Thread.run(Thread.java:748)


During handling of the above exception, another exception occurred:

AnalysisException                         Traceback (most recent call last)
~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/databricks/koalas/indexing.py in __getitem__(self, key)
    419         try:
--> 420             kdf = DataFrame(sdf.select(self._kdf._internal.index_scols + columns))
    421         except AnalysisException:

~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/pyspark/sql/dataframe.py in select(self, *cols)
   1319         """
-> 1320         jdf = self._jdf.select(self._jcols(*cols))
   1321         return DataFrame(jdf, self.sql_ctx)

~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/py4j/java_gateway.py in __call__(self, *args)
   1256         return_value = get_return_value(
-> 1257             answer, self.gateway_client, self.target_id, self.name)
   1258 

~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/pyspark/sql/utils.py in deco(*a, **kw)
     68             if s.startswith('org.apache.spark.sql.AnalysisException: '):
---> 69                 raise AnalysisException(s.split(': ', 1)[1], stackTrace)
     70             if s.startswith('org.apache.spark.sql.catalyst.analysis'):

AnalysisException: "cannot resolve '`bar`' given input columns: [('bar', 'one'), __index_level_0__, ('bar', 'two'), ('baz', 'one'), ('baz', 'two')];;\n'Project [__index_level_0__#147, 'bar, 'baz]\n+- LogicalRDD [__index_level_0__#147, ('bar', 'one')#148, ('bar', 'two')#149, ('baz', 'one')#150, ('baz', 'two')#151], false\n"

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
<ipython-input-20-e01dbeecebf3> in <module>
      6 
      7 kdf = ks.DataFrame(np.random.randn(3, 4), index=['A', 'B', 'C'], columns=arrays)
----> 8 kdf[['bar', 'baz']]

~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/databricks/koalas/frame.py in __getitem__(self, key)
   6158 
   6159     def __getitem__(self, key):
-> 6160         return self._pd_getitem(key)
   6161 
   6162     def __setitem__(self, key, value):

~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/databricks/koalas/frame.py in _pd_getitem(self, key)
   6115             raise NotImplementedError(key)
   6116         if isinstance(key, list):
-> 6117             return self.loc[:, key]
   6118         if isinstance(key, DataFrame):
   6119             # TODO Should not implement alignment, too dangerous?

~/.pyenv/versions/anaconda3-5.3.1/envs/koalas/lib/python3.6/site-packages/databricks/koalas/indexing.py in __getitem__(self, key)
    421         except AnalysisException:
    422             raise KeyError('[{}] don\'t exist in columns'
--> 423                            .format([col._jc.toString() for col in columns]))
    424         kdf._internal = kdf._internal.copy(
    425             data_columns=kdf._internal.data_columns[-len(columns):],

KeyError: "[['bar', 'baz']] don't exist in columns"
```


Discussion:
I'll take this.

Resolving patch:
diff --git a/databricks/koalas/indexing.py b/databricks/koalas/indexing.py
--- a/databricks/koalas/indexing.py
+++ b/databricks/koalas/indexing.py
@@ -19,6 +19,7 @@
 """
 from functools import reduce
 
+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
@@ -143,18 +144,28 @@ def __getitem__(self, key):
         if len(self._kdf._internal.index_columns) != 1:
             raise ValueError("'.at' only supports indices with level 1 right now")
 
-        column = key[1] if self._ks is None else self._ks.name
+        if self._ks is None:
+            if self._kdf._internal.column_index is not None:
+                column = dict(zip(self._kdf._internal.column_index,
+                                  self._kdf._internal.data_columns)).get(key[1], None)
+                if column is None:
+                    raise KeyError(key[1])
+            else:
+                column = key[1]
+        else:
+            column = self._ks.name
+
         if column is not None and column not in self._kdf._internal.data_columns:
-            raise KeyError("%s" % column)
-        series = self._ks if self._ks is not None else self._kdf[column]
+            raise KeyError(column)
+        sdf = self._ks._kdf._sdf if self._ks is not None else self._kdf._sdf
 
         row = key[0] if self._ks is None else key
-        pdf = (series._kdf._sdf
+        pdf = (sdf
                .where(self._kdf._internal.index_scols[0] == row)
                .select(_make_col(column))
                .toPandas())
         if len(pdf) < 1:
-            raise KeyError("%s" % row)
+            raise KeyError(row)
 
         values = pdf.iloc[:, 0].values
         return values[0] if len(values) == 1 else values
@@ -403,7 +414,13 @@ def raiseNotImplemented(description):
                 raiseNotImplemented("Cannot select with MultiIndex with Spark.")
 
         # make cols_sel a 1-tuple of string if a single string
-        if isinstance(cols_sel, (str, Series)):
+        column_index = self._kdf._internal.column_index
+        if isinstance(cols_sel, str):
+            if column_index is not None:
+                return self[rows_sel, [cols_sel]]._get_from_multiindex_column(cols_sel)
+            else:
+                cols_sel = _make_col(cols_sel)
+        elif isinstance(cols_sel, Series):
             cols_sel = _make_col(cols_sel)
         elif isinstance(cols_sel, slice) and cols_sel != slice(None):
             raise raiseNotImplemented("Can only select columns either by name or reference or all")
@@ -415,15 +432,27 @@ def raiseNotImplemented(description):
         elif isinstance(cols_sel, spark.Column):
             columns = [cols_sel]
         else:
-            columns = [_make_col(c) for c in cols_sel]
+            if column_index is not None:
+                column_to_index = list(zip(self._kdf._internal.data_columns,
+                                           self._kdf._internal.column_index))
+                columns, column_index = zip(*[(_make_col(column), idx)
+                                              for key in cols_sel
+                                              for column, idx in column_to_index
+                                              if idx[0] == key])
+                columns, column_index = list(columns), list(column_index)
+            else:
+                columns = [_make_col(c) for c in cols_sel]
+
         try:
             kdf = DataFrame(sdf.select(self._kdf._internal.index_scols + columns))
         except AnalysisException:
             raise KeyError('[{}] don\'t exist in columns'
                            .format([col._jc.toString() for col in columns]))
+
         kdf._internal = kdf._internal.copy(
             data_columns=kdf._internal.data_columns[-len(columns):],
-            index_map=self._kdf._internal.index_map)
+            index_map=self._kdf._internal.index_map,
+            column_index=column_index)
         if cols_sel is not None and isinstance(cols_sel, spark.Column):
             from databricks.koalas.series import _col
             return _col(kdf)
@@ -655,9 +684,18 @@ def raiseNotImplemented(description):
         except AnalysisException:
             raise KeyError('[{}] don\'t exist in columns'
                            .format([col._jc.toString() for col in columns]))
+
+        column_index = self._kdf._internal.column_index
+        if column_index is not None:
+            if cols_sel is not None and isinstance(cols_sel, (Series, int)):
+                column_index = None
+            else:
+                column_index = pd.MultiIndex.from_tuples(column_index)[cols_sel].tolist()
+
         kdf._internal = kdf._internal.copy(
             data_columns=kdf._internal.data_columns[-len(columns):],
-            index_map=self._kdf._internal.index_map)
+            index_map=self._kdf._internal.index_map,
+            column_index=column_index)
         if cols_sel is not None and isinstance(cols_sel, (Series, int)):
             from databricks.koalas.series import _col
             return _col(kdf)


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
--- a/databricks/koalas/tests/test_dataframe.py
+++ b/databricks/koalas/tests/test_dataframe.py
@@ -151,6 +151,9 @@ def test_multiindex_column_access(self):
         self.assert_eq(kdf['']['h'], pdf['']['h'])
         self.assert_eq(kdf['i'], pdf['i'])
 
+        self.assert_eq(kdf[['a', 'e']], pdf[['a', 'e']])
+        self.assert_eq(kdf[['e', 'a']], pdf[['e', 'a']])
+
     def test_repr_cache_invalidation(self):
         # If there is any cache, inplace operations should invalidate it.
         df = ks.range(10)
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
@@ -165,6 +165,15 @@ def test_at(self):
         with self.assertRaises(TypeError):
             kdf.at[3, 'b'] = 10
 
+    def test_at_multiindex_columns(self):
+        arrays = [np.array(['bar', 'bar', 'baz', 'baz']),
+                  np.array(['one', 'two', 'one', 'two'])]
+
+        pdf = pd.DataFrame(np.random.randn(3, 4), index=['A', 'B', 'C'], columns=arrays)
+        kdf = ks.from_pandas(pdf)
+
+        self.assert_eq(kdf.at['B', ('bar', 'one')], pdf.at['B', ('bar', 'one')])
+
     def test_loc(self):
         kdf = self.kdf
         pdf = self.pdf
@@ -271,6 +280,16 @@ def test_loc2d(self):
         self.assertRaises(SparkPandasIndexingError, lambda: kdf.a.loc[3:, 3])
         self.assertRaises(SparkPandasIndexingError, lambda: kdf.a.loc[kdf.a % 2 == 0, 3])
 
+    def test_loc2d_multiindex_columns(self):
+        arrays = [np.array(['bar', 'bar', 'baz', 'baz']),
+                  np.array(['one', 'two', 'one', 'two'])]
+
+        pdf = pd.DataFrame(np.random.randn(3, 4), index=['A', 'B', 'C'], columns=arrays)
+        kdf = ks.from_pandas(pdf)
+
+        self.assert_eq(kdf.loc['B':'B', 'bar'], pdf.loc['B':'B', 'bar'])
+        self.assert_eq(kdf.loc['B':'B', ['bar']], pdf.loc['B':'B', ['bar']])
+
     def test_loc2d_with_known_divisions(self):
         pdf = pd.DataFrame(np.random.randn(20, 5),
                            index=list('abcdefghijklmnopqrst'),
@@ -465,6 +484,24 @@ def test_iloc(self):
             self.assert_eq(kdf.iloc[:-1, indexer], pdf.iloc[:-1, indexer])
             self.assert_eq(kdf.iloc[kdf.index == 2, indexer], pdf.iloc[pdf.index == 2, indexer])
 
+    def test_iloc_multiindex_columns(self):
+        arrays = [np.array(['bar', 'bar', 'baz', 'baz']),
+                  np.array(['one', 'two', 'one', 'two'])]
+
+        pdf = pd.DataFrame(np.random.randn(3, 4), index=['A', 'B', 'C'], columns=arrays)
+        kdf = ks.from_pandas(pdf)
+
+        for indexer in [0,
+                        [0],
+                        [0, 1],
+                        [1, 0],
+                        [False, True, True, True],
+                        slice(0, 1)]:
+            self.assert_eq(kdf.iloc[:, indexer], pdf.iloc[:, indexer])
+            self.assert_eq(kdf.iloc[:1, indexer], pdf.iloc[:1, indexer])
+            self.assert_eq(kdf.iloc[:-1, indexer], pdf.iloc[:-1, indexer])
+            self.assert_eq(kdf.iloc[kdf.index == 'B', indexer], pdf.iloc[pdf.index == 'B', indexer])
+
     def test_iloc_series(self):
         pseries = pd.Series([1, 2, 3])
         kseries = ks.from_pandas(pseries)