From cceac0e4fa34cbb727b0eb5bf22cb07666e1ccc4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=BB=A7=E7=9B=9B?= Date: Fri, 15 Oct 2021 15:37:16 +0800 Subject: [PATCH] Fix df.loc when providing empty list --- mars/dataframe/indexing/index_lib.py | 14 +++++++++++--- .../indexing/tests/test_indexing_execution.py | 11 +++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/mars/dataframe/indexing/index_lib.py b/mars/dataframe/indexing/index_lib.py index b5ad90aa9d..91b0483d68 100644 --- a/mars/dataframe/indexing/index_lib.py +++ b/mars/dataframe/indexing/index_lib.py @@ -756,7 +756,10 @@ def preprocess(self, for chunk_index, pos in chunk_index_to_pos.items(): # chunk_index and pos are all list with 1 element abs_pos = pos[0] + cum_nsplit[chunk_index[0]] - chunk_labels = to_numpy(pd_index[abs_pos]) + if isinstance(pd_index, pd.RangeIndex) and len(abs_pos) == 0: + chunk_labels = np.array([], dtype=pd_index.dtype) + else: + chunk_labels = to_numpy(pd_index[abs_pos]) chunk_index_to_labels[chunk_index[0]] = chunk_labels index_info.is_label_asc_sorted = is_asc_sorted @@ -778,6 +781,8 @@ def process(self, tileable = context.tileable input_axis = index_info.input_axis chunk_index_to_labels = index_info.chunk_index_to_labels + full_label_size = sum(labels.sum() for labels + in chunk_index_to_labels.values()) other_index_to_iter = dict() chunk_index_to_info = context.chunk_index_to_info.copy() @@ -786,8 +791,11 @@ def process(self, chunk_labels = chunk_index_to_labels[i] size = chunk_labels.size - if size == 0 and tileable.shape[0] > 0: - # not effected when tileable not empty and no index chosen + if size == 0 and full_label_size > 0 and tileable.shape[0] > 0: + # not effected when + # 1) tileable not empty + # 2) full index not empty + # 3) no index chosen for this chunk 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 index 6e069a2656..d4a2978ff9 100644 --- a/mars/dataframe/indexing/tests/test_indexing_execution.py +++ b/mars/dataframe/indexing/tests/test_indexing_execution.py @@ -301,6 +301,17 @@ def test_loc_getitem(setup): expected = raw1.loc['a3', 'b'] assert result == expected + # test empty list + df = df1.loc[[]] + result = df.execute().fetch() + expected = raw1.loc[[]] + pd.testing.assert_frame_equal(result, expected) + + df = df2.loc[[]] + result = df.execute().fetch() + expected = raw2.loc[[]] + pd.testing.assert_frame_equal(result, expected) + df = df2.loc[1:4, 'b':'d'] result = df.execute().fetch() expected = raw2.loc[1:4, 'b': 'd']