Fix df.loc when df is empty (#2524)

mars-project · Oct 15, 2021 · 53ac012 · 53ac012
1 parent b552835
commit 53ac012
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/mars/dataframe/indexing/index_lib.py b/mars/dataframe/indexing/index_lib.py
@@ -786,8 +786,8 @@ def process(self,
             chunk_labels = chunk_index_to_labels[i]
             size = chunk_labels.size
 
-            if size == 0:
-                # not effected
+            if size == 0 and tileable.shape[0] > 0:
+                # not effected when tileable not empty and no index chosen
                 del context.chunk_index_to_info[chunk_index]
                 continue
 

diff --git a/mars/dataframe/indexing/tests/test_indexing_execution.py b/mars/dataframe/indexing/tests/test_indexing_execution.py
@@ -282,12 +282,14 @@ def test_loc_getitem(setup):
     # index that is timestamp
     raw5 = raw1.copy()
     raw5.index = pd.date_range('2020-1-1', periods=5)
+    raw6 = raw1[:0]
 
     df1 = md.DataFrame(raw1, chunk_size=2)
     df2 = md.DataFrame(raw2, chunk_size=2)
     df3 = md.DataFrame(raw3, chunk_size=2)
     df4 = md.DataFrame(raw4, chunk_size=2)
     df5 = md.DataFrame(raw5, chunk_size=2)
+    df6 = md.DataFrame(raw6)
 
     df = df2.loc[3, 'b']
     result = df.execute().fetch()
@@ -382,6 +384,12 @@ def test_loc_getitem(setup):
     expected = raw5.loc['2020-1-1', 'c']
     assert result == expected
 
+    # test empty df
+    df = df6.loc[[]]
+    result = df.execute().fetch()
+    expected = raw6.loc[[]]
+    pd.testing.assert_frame_equal(result, expected)
+
 
 def test_dataframe_getitem(setup):
     data = pd.DataFrame(np.random.rand(10, 5), columns=['c1', 'c2', 'c3', 'c4', 'c5'])