From d845c7ab71e51b1c20aaba3e6344451eb247e513 Mon Sep 17 00:00:00 2001 From: mgqa34 Date: Mon, 24 Jul 2023 15:33:15 +0800 Subject: [PATCH] dataframe: fix row indexing when result is empty Signed-off-by: mgqa34 --- python/fate/arch/dataframe/_dataframe.py | 11 +++++++---- python/fate/arch/dataframe/ops/_dimension_scaling.py | 3 +++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/python/fate/arch/dataframe/_dataframe.py b/python/fate/arch/dataframe/_dataframe.py index c9d2dd6062..8fc8812e6f 100644 --- a/python/fate/arch/dataframe/_dataframe.py +++ b/python/fate/arch/dataframe/_dataframe.py @@ -101,15 +101,18 @@ def weight(self): @property def shape(self) -> "tuple": - if not self.__count: + if self.__count is None: if self._sample_id_indexer: items = self._sample_id_indexer.count() elif self._match_id_indexer: items = self._match_id_indexer.count() else: - items = self._block_table.mapValues(lambda block: 0 if block is None else len(block[0])).reduce( - lambda size1, size2: size1 + size2 - ) + if self._block_table.count() == 0: + items = 0 + else: + items = self._block_table.mapValues(lambda block: 0 if block is None else len(block[0])).reduce( + lambda size1, size2: size1 + size2 + ) self.__count = items return self.__count, len(self._data_manager.schema.columns) diff --git a/python/fate/arch/dataframe/ops/_dimension_scaling.py b/python/fate/arch/dataframe/ops/_dimension_scaling.py index e52685567d..1164ba45b9 100644 --- a/python/fate/arch/dataframe/ops/_dimension_scaling.py +++ b/python/fate/arch/dataframe/ops/_dimension_scaling.py @@ -219,6 +219,9 @@ def _retrieval(blocks, t: torch.Tensor): _flatten_func = functools.partial(_flatten_partition, block_num=df.data_manager.block_num) retrieval_raw_table = retrieval_block_table.mapPartitions(_flatten_func, use_previous_behavior=False) + if retrieval_raw_table.count() == 0: + return df.empty_frame() + partition_order_mappings = get_partition_order_by_raw_table(retrieval_raw_table) to_blocks_func = functools.partial(to_blocks, dm=df.data_manager, partition_mappings=partition_order_mappings)