Fix Series.__getitem__ to work properly

databricks/koalas/series.py

@@ -42,7 +42,8 @@
 from databricks.koalas.internal import IndexMap, _InternalFrame, SPARK_INDEX_NAME_FORMAT
 from databricks.koalas.missing.series import _MissingPandasLikeSeries
 from databricks.koalas.plot import KoalasSeriesPlotMethods
-from databricks.koalas.utils import validate_arguments_and_invoke_function, scol_for
+from databricks.koalas.utils import (validate_arguments_and_invoke_function, scol_for,
+                                     tuple_like_strings)
 from databricks.koalas.datetimes import DatetimeMethods
 from databricks.koalas.strings import StringMethods
 
@@ -3560,7 +3561,37 @@ def __len__(self):
         return len(self.to_dataframe())
 
     def __getitem__(self, key):
-        return self._with_new_scol(self._scol.__getitem__(key))
+        if not isinstance(key, tuple):
+            key = (key,)
+        if len(self._internal._index_map) < len(key):
+            raise KeyError("Key length ({}) exceeds index depth ({})"
+                           .format(len(key), len(self._internal.index_map)))
+
+        cols = (self._internal.index_scols[len(key):] +
+                [self._internal.scol_for(self._internal.column_index[0])])
+        rows = [self._internal.scols[level] == index
+                for level, index in enumerate(key)]
+        sdf = self._internal.sdf \
+            .select(cols) \
+            .where(reduce(lambda x, y: x & y, rows))
+
+        if len(self._internal._index_map) == len(key):
+            # if sdf has one column and one data, return data only without frame
+            pdf = sdf.limit(2).toPandas()
+            length = len(pdf)
+            if length == 1:
+                return pdf[self.name].iloc[0]
+
+            key_string = tuple_like_strings(key) if len(key) > 1 else key[0]
+            sdf = sdf.withColumn(SPARK_INDEX_NAME_FORMAT(0), F.lit(str(key_string)))
+            internal = _InternalFrame(sdf=sdf, index_map=[(SPARK_INDEX_NAME_FORMAT(0), None)])
+            return _col(DataFrame(internal))
+
+        internal = self._internal.copy(
+            sdf=sdf,
+            index_map=self._internal._index_map[len(key):])
+
+        return _col(DataFrame(internal))
 
     def __getattr__(self, item: str_type) -> Any:
         if item.startswith("__") or item.startswith("_pandas_") or item.startswith("_spark_"):

databricks/koalas/tests/test_series.py

@@ -729,6 +729,31 @@ def test_duplicates(self):
         self.assert_eq(pser.drop_duplicates().sort_values(),
                        kser.drop_duplicates().sort_values())
 
+    def test_getitem(self):
+        pser = pd.Series([10, 20, 15, 30, 45], ['A', 'A', 'B', 'C', 'D'])
+        kser = ks.Series(pser)
+
+        self.assert_eq(kser['A'], pser['A'])
+        self.assert_eq(kser['B'], pser['B'])
+
+        # for MultiIndex
+        midx = pd.MultiIndex([['a', 'b', 'c'],
+                              ['lama', 'cow', 'falcon'],
+                              ['speed', 'weight', 'length']],
+                             [[0, 0, 0, 0, 0, 0, 1, 1, 1],
+                              [0, 0, 0, 1, 1, 1, 2, 2, 2],
+                              [0, 0, 0, 0, 1, 2, 0, 1, 2]])
+        pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3],
+                         name='0', index=midx)
+        kser = ks.Series(pser)
+
+        self.assert_eq(kser['a'], pser['a'])
+        self.assert_eq(kser['a', 'lama'], pser['a', 'lama'])
+
+        msg = r"'Key length \(4\) exceeds index depth \(3\)'"
+        with self.assertRaisesRegex(KeyError, msg):
+            kser[('a', 'lama', 'speed', 'x')]
+
     def test_keys(self):
         midx = pd.MultiIndex([['lama', 'cow', 'falcon'],
                               ['speed', 'weight', 'length']],

databricks/koalas/utils.py

@@ -358,3 +358,16 @@ def column_index_level(column_index: List[Tuple[str, ...]]) -> int:
         levels = set(0 if idx is None else len(idx) for idx in column_index)
         assert len(levels) == 1, levels
         return list(levels)[0]
+
+
+def tuple_like_strings(items):
+    """
+    Return the tuple-like strings from items
+
+    Examples
+    --------
+    >>> items = ('a', 'b', 'c')
+    >>> tuple_like_strings(items)
+    '(a, b, c)'
+    """
+    return '(%s)' % ', '.join(items)