Repository: databricks/koalas

Issue:
.loc behavior for scalar values
Assuming a simple DataFrame `kdf = ks.DataFrame({'A': [1,2,3]})`, running `kdf.loc[1]` results in the following exception:
```SparkPandasNotImplementedError: Cannot use a scalar value for row selection with Spark. You are trying to use pandas function .loc[..., ...], use spark function select, where```

However, `kdf.loc[1:1]` behaves the same as in pandas and returns the wanted column but as a DataFrame while `.loc[1]` would return a Series.

The `LocIndexer` docs state the following:
```
Allowed inputs are:
    - A single label, e.g. ``5`` or ``'a'``, (note that ``5`` is
      interpreted as a *label* of the index, and **never** as an
      integer position along the index) for column selection.
```

Can someone point out the rationale behind this behavior? My only guess would be that the exception tries to keep users from using the label as an integer position, which is not even the point of `.loc` but rather `.iloc`, right?


Discussion:
The notion of order in Spark's DataFrame doesn't exist. If we read a DataFrame from the disk, we perform some transformations and then we save it. We won't find the same order as the first. 
I think this due to the fact that Spark divides the output into chunks.


What you say is true for regular Spark DataFrames. However, Koalas stores pandas indices as a dedicated *\_\_index_level_n\_\_* column. Given an example DataFrame `kdf = ks.DataFrame({'A': list('abc')})`, it is for example easily possible to get the second column (containing `b`) using `kdf._sdf.where(F.col('__index_level_0__') == 1)`. This what the `.loc` method is made for. I'm aware that getting the same row with `.iloc[1]` is not feasible with PySpark for the reasons you've explained.
Here is the previous discussion: https://github.com/databricks/koalas/pull/324#discussion_r283658407

If we want to strictly follow pandas' behavior, we need to count the matched row to change the return value if the matched row is only one or not.
I was not sure whether we want to run the extra `count` job for this.
Also I was not sure what the return value should be, Koalas DataFrame/Series or pandas' ones for the case?
The implementation itself should be easy if we want anyway.
We can certainly add scalar addressing, but I want to hear from more users first, because it is easily an antipattern when using spark:

```py
[process_group(df.loc[x]) for x in ...]
``` 


Resolving patch:
diff --git a/databricks/koalas/indexing.py b/databricks/koalas/indexing.py
--- a/databricks/koalas/indexing.py
+++ b/databricks/koalas/indexing.py
@@ -25,7 +25,7 @@
 from pyspark.sql.types import BooleanType, StringType, LongType
 from pyspark.sql.utils import AnalysisException
 
-from databricks.koalas.internal import _InternalFrame, HIDDEN_COLUMNS, NATURAL_ORDER_COLUMN_NAME
+from databricks.koalas.internal import _InternalFrame, NATURAL_ORDER_COLUMN_NAME
 from databricks.koalas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError
 from databricks.koalas.utils import column_index_level, name_like_string
 
@@ -99,7 +99,7 @@ def __getitem__(self, key):
             row_sel, col_sel = key
         else:
             assert self._is_series, type(self._kdf_or_kser)
-            if not isinstance(key, str) and len(key) != 1:
+            if isinstance(key, tuple) and len(key) != 1:
                 raise TypeError("Use Series.at like .at[row_index]")
             row_sel = key
             col_sel = self._internal.column_index[0]
@@ -138,17 +138,12 @@ def __getitem__(self, key):
         from databricks.koalas.series import Series
 
         if self._is_series:
-            if isinstance(key, tuple):
-                if len(key) > 1:
-                    raise SparkPandasIndexingError('Too many indexers')
-                key = key[0]
-
             if isinstance(key, Series) and key._kdf is not self._kdf_or_kser._kdf:
                 kdf = self._kdf_or_kser.to_frame()
                 kdf['__temp_col__'] = key
                 return type(self)(kdf[self._kdf_or_kser.name])[kdf['__temp_col__']]
 
-            cond, limit = self._select_rows(key)
+            cond, limit, remaining_index = self._select_rows(key)
             if cond is None and limit is None:
                 return self._kdf_or_kser
 
@@ -171,7 +166,7 @@ def __getitem__(self, key):
                 return type(self)(kdf)[kdf['__temp_col__'],
                                        cols_sel][list(self._kdf_or_kser.columns)]
 
-            cond, limit = self._select_rows(rows_sel)
+            cond, limit, remaining_index = self._select_rows(rows_sel)
             column_index, column_scols, returns_series = self._select_cols(cols_sel)
 
             if cond is None and limit is None and returns_series:
@@ -179,6 +174,20 @@ def __getitem__(self, key):
                                                   column_index=[column_index[0]]),
                               anchor=self._kdf_or_kser)
 
+        if remaining_index is not None:
+            index_scols = self._internal.index_scols[-remaining_index:]
+            index_map = self._internal.index_map[-remaining_index:]
+        else:
+            index_scols = self._internal.index_scols
+            index_map = self._internal.index_map
+
+        if self._internal.column_index_names is None:
+            column_index_names = None
+        else:
+            # Manage column index names
+            level = column_index_level(column_index)
+            column_index_names = self._internal.column_index_names[-level:]
+
         try:
             sdf = self._internal._sdf
             if cond is not None:
@@ -189,29 +198,34 @@ def __getitem__(self, key):
                 else:
                     sdf = sdf.limit(sdf.count() + limit)
 
-            sdf = sdf.select(self._internal.index_scols + column_scols)
-
-            if self._internal.column_index_names is None:
-                column_index_names = None
-            else:
-                # Manage column index names
-                level = column_index_level(column_index)
-                column_index_names = self._internal.column_index_names[-level:]
-
-            internal = _InternalFrame(sdf=sdf,
-                                      index_map=self._internal.index_map,
-                                      column_index=column_index,
-                                      column_index_names=column_index_names)
-            kdf = DataFrame(internal)
+            sdf = sdf.select(index_scols + column_scols)
         except AnalysisException:
             raise KeyError('[{}] don\'t exist in columns'
                            .format([col._jc.toString() for col in column_scols]))
 
+        internal = _InternalFrame(sdf=sdf,
+                                  index_map=index_map,
+                                  column_index=column_index,
+                                  column_index_names=column_index_names)
+        kdf = DataFrame(internal)
+
         if returns_series:
-            return Series(kdf._internal.copy(scol=kdf._internal.column_scols[0]),
-                          anchor=kdf)
+            kdf_or_kser = Series(kdf._internal.copy(scol=kdf._internal.column_scols[0]),
+                                 anchor=kdf)
+        else:
+            kdf_or_kser = kdf
+
+        if remaining_index is not None and remaining_index == 0:
+            pdf_or_pser = kdf_or_kser.head(2).to_pandas()
+            length = len(pdf_or_pser)
+            if length == 0:
+                raise KeyError(name_like_string(key))
+            elif length == 1:
+                return pdf_or_pser.iloc[0]
+            else:
+                return kdf_or_kser
         else:
-            return kdf
+            return kdf_or_kser
 
 
 class LocIndexer(_LocIndexerLike):
@@ -265,12 +279,12 @@ class LocIndexer(_LocIndexerLike):
     viper               4       5
     sidewinder          7       8
 
-    A single label for row selection is not allowed.
+    Single label. Note this returns the row as a Series.
 
     >>> df.loc['viper']
-    Traceback (most recent call last):
-     ...
-    databricks.koalas.exceptions.SparkPandasNotImplementedError: ...
+    max_speed    4
+    shield       5
+    Name: viper, dtype: int64
 
     List of labels. Note using ``[[]]`` returns a DataFrame.
     Also note that Koalas behaves just a filter without reordering by the labels.
@@ -285,13 +299,24 @@ class LocIndexer(_LocIndexerLike):
     viper               4       5
     sidewinder          7       8
 
-    Single label for column
+    Single label for column.
+
+    >>> df.loc['cobra', 'shield']
+    2
+
+    List of labels for row.
 
     >>> df.loc[['cobra'], 'shield']
     cobra    2
     Name: shield, dtype: int64
 
-    List of labels for column. Note using list returns a DataFrame.
+    List of labels for column.
+
+    >>> df.loc['cobra', ['shield']]
+    shield    2
+    Name: cobra, dtype: int64
+
+    List of labels for both row and column.
 
     >>> df.loc[['cobra'], ['shield']]
            shield
@@ -397,14 +422,14 @@ def _select_rows(self, rows_sel):
 
         if isinstance(rows_sel, Series):
             assert isinstance(rows_sel.spark_type, BooleanType), rows_sel.spark_type
-            return rows_sel._scol, None
+            return rows_sel._scol, None, None
         elif isinstance(rows_sel, slice):
             assert len(self._internal.index_columns) > 0
             if rows_sel.step is not None:
                 LocIndexer._raiseNotImplemented("Cannot use step with Spark.")
             if rows_sel == slice(None):
                 # If slice is None - select everything, so nothing to do
-                return None, None
+                return None, None, None
             elif len(self._internal.index_columns) == 1:
                 sdf = self._kdf_or_kser._internal.sdf
                 index = self._kdf_or_kser.index
@@ -451,7 +476,8 @@ def _select_rows(self, rows_sel):
 
                 # if start and stop are same, just get all start(or stop) values
                 if start == stop:
-                    return index_column._scol == F.lit(rows_sel.start).cast(index_data_type), None
+                    return (index_column._scol == F.lit(rows_sel.start).cast(index_data_type),
+                            None, None)
 
                 cond = []
                 if start is not None:
@@ -460,30 +486,34 @@ def _select_rows(self, rows_sel):
                     cond.append(stop_order_column <= F.lit(stop).cast(stop_order_column_type))
 
                 if len(cond) > 0:
-                    return reduce(lambda x, y: x & y, cond), None
+                    return reduce(lambda x, y: x & y, cond), None, None
             else:
                 LocIndexer._raiseNotImplemented("Cannot use slice for MultiIndex with Spark.")
-        elif isinstance(rows_sel, str):
-            LocIndexer._raiseNotImplemented(
-                "Cannot use a scalar value for row selection with Spark.")
-        else:
-            try:
-                rows_sel = list(rows_sel)
-            except TypeError:
-                LocIndexer._raiseNotImplemented(
-                    "Cannot use a scalar value for row selection with Spark.")
+        elif is_list_like(rows_sel) and not isinstance(rows_sel, tuple):
+            rows_sel = list(rows_sel)
             if len(rows_sel) == 0:
-                return F.lit(False), None
+                return F.lit(False), None, None
             elif len(self._internal.index_columns) == 1:
                 index_column = self._kdf_or_kser.index.to_series()
                 index_data_type = index_column.spark_type
                 if len(rows_sel) == 1:
-                    return index_column._scol == F.lit(rows_sel[0]).cast(index_data_type), None
+                    return (index_column._scol == F.lit(rows_sel[0]).cast(index_data_type),
+                            None, None)
                 else:
-                    return index_column._scol.isin(
-                        [F.lit(r).cast(index_data_type) for r in rows_sel]), None
+                    return (index_column._scol.isin([F.lit(r).cast(index_data_type)
+                                                     for r in rows_sel]),
+                            None, None)
             else:
                 LocIndexer._raiseNotImplemented("Cannot select with MultiIndex with Spark.")
+        else:
+            if not isinstance(rows_sel, tuple):
+                rows_sel = (rows_sel,)
+            if len(rows_sel) > len(self._internal.index_map):
+                raise SparkPandasIndexingError('Too many indexers')
+
+            rows = [scol == value for scol, value in zip(self._internal.index_scols, rows_sel)]
+            return (reduce(lambda x, y: x & y, rows),
+                    None, len(self._internal.index_map) - len(rows_sel))
 
     def _get_from_multiindex_column(self, key, indexes=None):
         """ Select columns from multi-index columns.
@@ -752,20 +782,22 @@ def _raiseNotImplemented(description):
     def _select_rows(self, rows_sel):
         from databricks.koalas.indexes import Index
 
-        if isinstance(rows_sel, Index):
+        if isinstance(rows_sel, tuple) and len(rows_sel) > 1:
+            raise SparkPandasIndexingError('Too many indexers')
+        elif isinstance(rows_sel, Index):
             assert isinstance(rows_sel.spark_type, BooleanType), rows_sel.spark_type
-            return rows_sel._scol, None
+            return rows_sel._scol, None, None
         elif isinstance(rows_sel, slice):
             if rows_sel == slice(None):
                 # If slice is None - select everything, so nothing to do
-                return None, None
+                return None, None, None
             elif (rows_sel.start is not None) or (rows_sel.step is not None):
                 ILocIndexer._raiseNotImplemented("Cannot use start or step with Spark.")
             elif not isinstance(rows_sel.stop, int):
                 raise TypeError("cannot do slice indexing with these indexers [{}] of {}"
                                 .format(rows_sel.stop, type(rows_sel.stop)))
             else:
-                return None, rows_sel.stop
+                return None, rows_sel.stop, None
         else:
             ILocIndexer._raiseNotImplemented(".iloc requires numeric slice or conditional "
                                              "boolean Index, got {}".format(type(rows_sel)))
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
--- a/databricks/koalas/series.py
+++ b/databricks/koalas/series.py
@@ -37,6 +37,7 @@
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
 from databricks.koalas.config import get_option
 from databricks.koalas.base import IndexOpsMixin
+from databricks.koalas.exceptions import SparkPandasIndexingError
 from databricks.koalas.frame import DataFrame
 from databricks.koalas.generic import _Frame
 from databricks.koalas.internal import (_InternalFrame, NATURAL_ORDER_COLUMN_NAME,
@@ -4264,43 +4265,12 @@ def __len__(self):
         return len(self.to_dataframe())
 
     def __getitem__(self, key):
-        if isinstance(key, Series) and isinstance(key.spark_type, BooleanType):
+        try:
             return self.loc[key]
-
-        if not isinstance(key, tuple):
-            key = (key,)
-        if len(self._internal._index_map) < len(key):
+        except SparkPandasIndexingError:
             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 == 0:
-                raise KeyError(name_like_string(key))
-            if length == 1:
-                return pdf[self._internal.data_columns[0]].iloc[0]
-
-            key_string = name_like_string(key)
-            sdf = sdf.withColumn(SPARK_INDEX_NAME_FORMAT(0), F.lit(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("__"):
             raise AttributeError(item)


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
@@ -223,6 +223,14 @@ def test_loc(self):
         self.assert_eq(kdf.loc[1000:], pdf.loc[1000:])
         self.assert_eq(kdf.loc[-2000:-1000], pdf.loc[-2000:-1000])
 
+        self.assert_eq(kdf.loc[5], pdf.loc[5])
+        self.assert_eq(kdf.loc[9], pdf.loc[9])
+        self.assert_eq(kdf.a.loc[5], pdf.a.loc[5])
+        self.assert_eq(kdf.a.loc[9], pdf.a.loc[9])
+
+        self.assertRaises(KeyError, lambda: kdf.loc[10])
+        self.assertRaises(KeyError, lambda: kdf.a.loc[10])
+
         # duplicated index test
         pdf = pd.DataFrame(
             [1, 2, 3, 4, 5, 6, 7, 8, 9],
@@ -284,11 +292,18 @@ def test_loc_multiindex(self):
         pdf = self.pdf
         pdf = pdf.set_index('b', append=True)
 
-        self.assert_eq(kdf[['a']], pdf[['a']])
-
         self.assert_eq(kdf.loc[:], pdf.loc[:])
         self.assertRaises(NotImplementedError, lambda: kdf.loc[5:5])
 
+        self.assert_eq(kdf.loc[5], pdf.loc[5])
+        self.assert_eq(kdf.loc[9], pdf.loc[9])
+        # TODO: self.assert_eq(kdf.loc[(5, 3)], pdf.loc[(5, 3)])
+        # TODO: self.assert_eq(kdf.loc[(9, 0)], pdf.loc[(9, 0)])
+        self.assert_eq(kdf.a.loc[5], pdf.a.loc[5])
+        self.assert_eq(kdf.a.loc[9], pdf.a.loc[9])
+        self.assertTrue((kdf.a.loc[(5, 3)] == pdf.a.loc[(5, 3)]).all())
+        self.assert_eq(kdf.a.loc[(9, 0)], pdf.a.loc[(9, 0)])
+
     def test_loc2d_multiindex(self):
         kdf = self.kdf
         kdf = kdf.set_index('b', append=True)
@@ -325,6 +340,11 @@ 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])
 
+        self.assert_eq(kdf.loc[5, 'a'], pdf.loc[5, 'a'])
+        self.assert_eq(kdf.loc[9, 'a'], pdf.loc[9, 'a'])
+        self.assert_eq(kdf.loc[5, ['a']], pdf.loc[5, ['a']])
+        self.assert_eq(kdf.loc[9, ['a']], pdf.loc[9, ['a']])
+
     def test_loc2d_multiindex_columns(self):
         arrays = [np.array(['bar', 'bar', 'baz', 'baz']),
                   np.array(['one', 'two', 'one', 'two'])]