dataframe: fix row indexing when result is empty

Signed-off-by: mgqa34 <mgq3374541@163.com>
FederatedAI · Jul 24, 2023 · d845c7a · d845c7a
1 parent cea0e11
commit d845c7a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git 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
@@ -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)