Avoid iterative tiling for df.loc[:, fields] (#2685)

Author: Xuye (Chris) Qin

mars/dataframe/indexing/index_lib.py

@@ -242,17 +242,12 @@ def preprocess(self, index_info: IndexInfo, context: IndexHandlerContext) -> Non
             index_value = [tileable.index_value, tileable.columns_value][input_axis]
 
         # check if chunks have unknown shape
-        check = False
-        if index_value.has_value():
-            # index_value has value,
-            check = True
-        elif self._slice_all(index_info.raw_index):
-            # if slice on all data
-            check = True
-
-        if check:
-            if any(np.isnan(ns) for ns in tileable.nsplits[input_axis]):
-                yield []
+        if (
+            not self._slice_all(index_info.raw_index)
+            and index_value.has_value()
+            and any(np.isnan(ns) for ns in tileable.nsplits[input_axis])
+        ):  # pragma: no cover
+            yield []
 
     def set_chunk_index_info(
         cls,
@@ -297,6 +292,27 @@ def set_chunk_index_info(
 
         chunk_index_info.set(ChunkIndexAxisInfo(**kw))
 
+    def _process_slice_all_index(
+        self,
+        tileable: Tileable,
+        index_info: IndexInfo,
+        input_axis: int,
+        context: IndexHandlerContext,
+    ) -> None:
+        index_to_info = context.chunk_index_to_info.copy()
+        for chunk_index, chunk_index_info in index_to_info.items():
+            i = chunk_index[input_axis]
+            size = tileable.nsplits[input_axis][i]
+            self.set_chunk_index_info(
+                context,
+                index_info,
+                chunk_index,
+                chunk_index_info,
+                i,
+                slice(None),
+                size,
+            )
+
     def _process_has_value_index(
         self,
         tileable: Tileable,
@@ -306,17 +322,14 @@ def _process_has_value_index(
         context: IndexHandlerContext,
     ) -> None:
         pd_index = index_value.to_pandas()
-        if self._slice_all(index_info.raw_index):
-            slc = slice(None)
-        else:
-            # turn label-based slice into position-based slice
-            start, end = pd_index.slice_locs(
-                index_info.raw_index.start,
-                index_info.raw_index.stop,
-                index_info.raw_index.step,
-                kind="loc",
-            )
-            slc = slice(start, end, index_info.raw_index.step)
+        # turn label-based slice into position-based slice
+        start, end = pd_index.slice_locs(
+            index_info.raw_index.start,
+            index_info.raw_index.stop,
+            index_info.raw_index.step,
+            kind="loc",
+        )
+        slc = slice(start, end, index_info.raw_index.step)
 
         cum_nsplit = [0] + np.cumsum(tileable.nsplits[index_info.input_axis]).tolist()
         # split position-based slice into chunk slices
@@ -379,7 +392,9 @@ def process(self, index_info: IndexInfo, context: IndexHandlerContext) -> None:
         else:
             index_value = [tileable.index_value, tileable.columns_value][input_axis]
 
-        if index_value.has_value() or self._slice_all(index_info.raw_index):
+        if self._slice_all(index_info.raw_index):
+            self._process_slice_all_index(tileable, index_info, input_axis, context)
+        elif index_value.has_value():
             self._process_has_value_index(
                 tileable, index_info, index_value, input_axis, context
             )
@@ -829,10 +844,13 @@ def parse(self, raw_index, context: IndexHandlerContext) -> IndexInfo:
     def preprocess(self, index_info: IndexInfo, context: IndexHandlerContext) -> None:
         tileable = context.tileable
         op = context.op
-        if has_unknown_shape(tileable):
-            yield []
 
         input_axis = index_info.input_axis
+
+        # check unknown shape
+        if any(np.isnan(s) for s in tileable.nsplits[input_axis]):
+            yield []
+
         if tileable.ndim == 2:
             index_value = [tileable.index_value, tileable.columns_value][input_axis]
         else:

mars/dataframe/indexing/tests/test_indexing.py

@@ -645,6 +645,12 @@ def test_dataframe_loc():
     for loc_chunk, chunk in zip(tiled_loc_df.chunks, tiled_df.chunks):
         assert loc_chunk.index_value.key == chunk.index_value.key
 
+    # test loc on filtered df
+    df2 = df[df["x"] < 1]
+    loc_df = df2.loc[:, ["y", "x"]]
+    tiled_loc_df = tile(loc_df)
+    assert len(tiled_loc_df.chunks) == 3
+
 
 def test_loc_use_iloc():
     raw = pd.DataFrame([[1, 3, 3], [4, 2, 6], [7, 8, 9]], columns=["x", "y", "z"])

mars/dataframe/indexing/tests/test_indexing_execution.py

@@ -312,6 +312,10 @@ def test_loc_getitem(setup):
     result = df.execute().fetch()
     expected = raw2.loc[:, "b"]
     pd.testing.assert_series_equal(result, expected)
+    df = df2.loc[:, ["b", "a"]]
+    result = df.execute().fetch()
+    expected = raw2.loc[:, ["b", "a"]]
+    pd.testing.assert_frame_equal(result, expected)
 
     # 'b' is non-unique
     df = df3.loc[:, "b"]
@@ -336,6 +340,11 @@ def test_loc_getitem(setup):
     result = df.execute().fetch()
     expected = raw2.loc[[3, 0, 1], ["c", "a", "d"]]
     pd.testing.assert_frame_equal(result, expected)
+    df = df2[df2["a"] < 10]
+    df = df.loc[[3, 0, 1], ["c", "a", "d"]]
+    result = df.execute().fetch()
+    expected = raw2.loc[[3, 0, 1], ["c", "a", "d"]]
+    pd.testing.assert_frame_equal(result, expected)
 
     # label-based fancy index, asc sorted
     df = df2.loc[[0, 1, 3], ["a", "c", "d"]]