Repository: databricks/koalas

Issue:
Add complete .iloc support
cc https://github.com/databricks/koalas/issues/247

Not sure what https://github.com/databricks/koalas/pull/364 intended but may we please have real `iloc` support?  Given the current error message, it sounds like the case below could simply be translated to a spark SQL / DataFrame query.  Since the input to `iloc` necessarily fits in memory in the driver (else the user program OOMs), such translation is unlikely to cause an OOM w/in `koalas`.  Also the error message "use spark function select, where" isn't very clear... does it mean use Spark SQL?  Like I need to register my DataFrame as a temp table and do a SQL query and then create a `koalas` dataframe from that?

```
>>> import findspark
>>> findspark.init()
>>> import pandas as pd
>>> import databricks.koalas as ks
>>> ks.__version__
'0.12.0'
>>> pd.__version__
'0.24.2'
>>> import pyspark
>>> pyspark.__version__
'2.4.3'
>>> import numpy as np
>>> np.__version__
'1.15.2'

>>> pdf = pd.DataFrame({'a': [1, 2, 3]})
>>> kdf = ks.DataFrame({'a': [1, 2, 3]})
... ( spark logging removed ) ...

>>> pdf.iloc[np.array([0, 1])]
   a
0  1
1  2
>>> pdf.iloc[np.array([1, 2])]
   a
1  2
2  3

>>> pdf.iloc[np.array([0, 1])]
   a
0  1
1  2
>>> kdf.iloc[np.array([0, 1])]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.6/dist-packages/databricks/koalas/indexing.py", line 626, in __getitem__
    "got {}".format(rows_sel))
  File "/usr/local/lib/python3.6/dist-packages/databricks/koalas/indexing.py", line 601, in raiseNotImplemented
    spark_target_function="select, where")
databricks.koalas.exceptions.SparkPandasNotImplementedError: .iloc requires numeric slice or conditional boolean Index, got [0 1] You are trying to use pandas function .iloc[..., ...], use spark function select, where
```


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
@@ -17,7 +17,7 @@
 """
 A loc indexer for Koalas DataFrame/Series.
 """
-from collections import OrderedDict
+from collections import OrderedDict, Iterable
 from functools import reduce
 
 from pandas.api.types import is_list_like
@@ -25,6 +25,7 @@
 from pyspark.sql import functions as F
 from pyspark.sql.types import BooleanType, LongType
 from pyspark.sql.utils import AnalysisException
+import numpy as np
 
 from databricks.koalas.internal import _InternalFrame, NATURAL_ORDER_COLUMN_NAME
 from databricks.koalas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError
@@ -797,15 +798,16 @@ class iLocIndexer(_LocIndexerLike):
     Allowed inputs are:
 
     - An integer for column selection, e.g. ``5``.
+    - A list or array of integers for row selection with distinct index values,
+      e.g. ``[3, 4, 0]``
     - A list or array of integers for column selection, e.g. ``[4, 3, 0]``.
     - A boolean array for column selection.
-    - A slice object with ints for column selection, e.g. ``1:7``.
-    - A slice object with ints without start and step for row selection, e.g. ``:7``.
+    - A slice object with ints for row and column selection, e.g. ``1:7``.
 
     Not allowed inputs which pandas allows are:
 
-    - An integer for row selection, e.g. ``5``.
-    - A list or array of integers for row selection, e.g. ``[4, 3, 0]``.
+    - A list or array of integers for row selection with duplicated indexes,
+      e.g. ``[4, 4, 0]``.
     - A boolean array for row selection.
     - A ``callable`` function with one argument (the calling Series, DataFrame
       or Panel) and that returns valid output for indexing (one of the above).
@@ -846,12 +848,9 @@ class iLocIndexer(_LocIndexerLike):
     d    4
     Name: 0, dtype: int64
 
-    A list of integers for row selection is not allowed.
-
     >>> df.iloc[[0]]
-    Traceback (most recent call last):
-     ...
-    databricks.koalas.exceptions.SparkPandasNotImplementedError: ...
+       a  b  c  d
+    0  1  2  3  4
 
     With a `slice` object.
 
@@ -991,10 +990,47 @@ def verify_type(i):
         elif isinstance(rows_sel, int):
             sdf = self._internal.spark_frame
             return (sdf[self._sequence_col] == rows_sel), None, 0
+        elif isinstance(rows_sel, Iterable):
+            sdf = self._internal.spark_frame
+
+            if any(
+                isinstance(key, (int, np.int, np.int64, np.int32)) and key < 0 for key in rows_sel
+            ):
+                offset = sdf.count()
+            else:
+                offset = 0
+
+            new_rows_sel = []
+            for key in list(rows_sel):
+                if not isinstance(key, (int, np.int, np.int64, np.int32)):
+                    raise TypeError(
+                        "cannot do positional indexing with these indexers [{}] of {}".format(
+                            key, type(key)
+                        )
+                    )
+                if key < 0:
+                    key = key + offset
+                new_rows_sel.append(key)
+
+            if len(new_rows_sel) != len(set(new_rows_sel)):
+                raise NotImplementedError(
+                    "Duplicated row selection is not currently supported; "
+                    "however, normalised index was [%s]" % new_rows_sel
+                )
+
+            sequence_scol = sdf[self._sequence_col]
+            cond = []
+            for key in new_rows_sel:
+                cond.append(sequence_scol == F.lit(int(key)).cast(LongType()))
+
+            if len(cond) == 0:
+                cond = [F.lit(False)]
+            return reduce(lambda x, y: x | y, cond), None, None
         else:
             iLocIndexer._raiseNotImplemented(
-                ".iloc requires numeric slice or conditional "
-                "boolean Index, got {}".format(type(rows_sel))
+                ".iloc requires numeric slice, conditional "
+                "boolean Index or a sequence of positions as int, "
+                "got {}".format(type(rows_sel))
             )
 
     def _select_cols(self, 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
@@ -695,6 +695,33 @@ def test_iloc_slice_rows_sel(self):
                 self.assert_eq(kdf.iloc[rows_sel].sort_index(), pdf.iloc[rows_sel].sort_index())
                 self.assert_eq(kdf.A.iloc[rows_sel].sort_index(), pdf.A.iloc[rows_sel].sort_index())
 
+    def test_iloc_iterable_rows_sel(self):
+        pdf = pd.DataFrame({"A": [1, 2] * 5, "B": [3, 4] * 5, "C": [5, 6] * 5})
+        kdf = ks.from_pandas(pdf)
+
+        for rows_sel in [
+            [],
+            np.array([0, 1]),
+            [1, 2],
+            np.array([-3]),
+            [3],
+            np.array([-2]),
+            [8, 3, -5],
+        ]:
+            with self.subTest(rows_sel=rows_sel):
+                self.assert_eq(kdf.iloc[rows_sel].sort_index(), pdf.iloc[rows_sel].sort_index())
+                self.assert_eq(kdf.A.iloc[rows_sel].sort_index(), pdf.A.iloc[rows_sel].sort_index())
+
+            with self.subTest(rows_sel=rows_sel):
+                self.assert_eq(
+                    kdf.iloc[rows_sel, :].sort_index(), pdf.iloc[rows_sel, :].sort_index()
+                )
+
+            with self.subTest(rows_sel=rows_sel):
+                self.assert_eq(
+                    kdf.iloc[rows_sel, :1].sort_index(), pdf.iloc[rows_sel, :1].sort_index()
+                )
+
     def test_setitem(self):
         pdf = pd.DataFrame(
             [[1, 2], [4, 5], [7, 8]],
@@ -736,18 +763,6 @@ def test_iloc_raises(self):
         pdf = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]})
         kdf = ks.from_pandas(pdf)
 
-        with self.assertRaisesRegex(
-            SparkPandasNotImplementedError,
-            ".iloc requires numeric slice or conditional boolean Index",
-        ):
-            kdf.iloc[[0, 1], :]
-
-        with self.assertRaisesRegex(
-            SparkPandasNotImplementedError,
-            ".iloc requires numeric slice or conditional boolean Index",
-        ):
-            kdf.A.iloc[[0, 1]]
-
         with self.assertRaisesRegex(SparkPandasIndexingError, "Only accepts pairs of candidates"):
             kdf.iloc[[0, 1], [0, 1], [1, 2]]
 
@@ -824,3 +839,7 @@ def test_index_operator_int(self):
         # Index loc search
         self.assert_eq(kdf.A[5], pdf.A[5])
         self.assert_eq(kdf.A[3], pdf.A[3])
+        with self.assertRaisesRegex(
+            NotImplementedError, "Duplicated row selection is not currently supported"
+        ):
+            kdf.iloc[[1, 1]]