Repository: modin-project/modin

Issue:
REFACTOR: remove code duplication for `get_indices` function



Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -774,6 +774,28 @@ def from_arrow(cls, at, return_dims=False):
         """
         return cls.from_pandas(at.to_pandas(), return_dims=return_dims)
 
+    @classmethod
+    def get_objects_from_partitions(cls, partitions):
+        """
+        Get the objects wrapped by `partitions`.
+
+        Parameters
+        ----------
+        partitions : np.ndarray
+            NumPy array with ``PandasDataframePartition``-s.
+
+        Returns
+        -------
+        list
+            The objects wrapped by `partitions`.
+
+        Notes
+        -----
+        This method should be implemented in a more efficient way for engines that support
+        getting objects in parallel.
+        """
+        return [partition.get() for partition in partitions]
+
     @classmethod
     def get_indices(cls, axis, partitions, index_func=None):
         """
@@ -803,16 +825,15 @@ def get_indices(cls, axis, partitions, index_func=None):
         func = cls.preprocess_func(index_func)
         if axis == 0:
             new_idx = (
-                [idx.apply(func).get() for idx in partitions.T[0]]
+                [idx.apply(func) for idx in partitions.T[0]]
                 if len(partitions.T)
                 else []
             )
         else:
             new_idx = (
-                [idx.apply(func).get() for idx in partitions[0]]
-                if len(partitions)
-                else []
+                [idx.apply(func) for idx in partitions[0]] if len(partitions) else []
             )
+        new_idx = cls.get_objects_from_partitions(new_idx)
         # TODO FIX INFORMATION LEAK!!!!1!!1!!
         return new_idx[0].append(new_idx[1:]) if len(new_idx) else new_idx
 
diff --git a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py
--- a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py
+++ b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py
@@ -23,7 +23,6 @@
     PandasOnDaskDataframeRowPartition,
 )
 from .partition import PandasOnDaskDataframePartition
-from modin.error_message import ErrorMessage
 
 
 class PandasOnDaskDataframePartitionManager(PandasDataframePartitionManager):
@@ -35,45 +34,19 @@ class PandasOnDaskDataframePartitionManager(PandasDataframePartitionManager):
     _row_partition_class = PandasOnDaskDataframeRowPartition
 
     @classmethod
-    def get_indices(cls, axis, partitions, index_func):
+    def get_objects_from_partitions(cls, partitions):
         """
-        Get the internal indices stored in the partitions.
+        Get the objects wrapped by `partitions` in parallel.
 
         Parameters
         ----------
-        axis : {0, 1}
-            Axis to extract the labels over.
         partitions : np.ndarray
-            The array of partitions from which need to extract the labels.
-        index_func : callable
-            The function to be used to extract the indices.
+            NumPy array with ``PandasDataframePartition``-s.
 
         Returns
         -------
-        pandas.Index
-            A pandas Index object.
-
-        Notes
-        -----
-        These are the global indices of the object. This is mostly useful
-        when you have deleted rows/columns internally, but do not know
-        which ones were deleted.
+        list
+            The objects wrapped by `partitions`.
         """
         client = default_client()
-        ErrorMessage.catch_bugs_and_request_email(not callable(index_func))
-        func = cls.preprocess_func(index_func)
-        if axis == 0:
-            # We grab the first column of blocks and extract the indices
-            new_idx = (
-                [idx.apply(func).future for idx in partitions.T[0]]
-                if len(partitions.T)
-                else []
-            )
-        else:
-            new_idx = (
-                [idx.apply(func).future for idx in partitions[0]]
-                if len(partitions)
-                else []
-            )
-        new_idx = client.gather(new_idx)
-        return new_idx[0].append(new_idx[1:]) if len(new_idx) else new_idx
+        return client.gather([partition.future for partition in partitions])
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
@@ -30,7 +30,6 @@
 from .partition import PandasOnRayDataframePartition
 from modin.core.execution.ray.generic.modin_aqp import call_progress_bar
 from modin.core.storage_formats.pandas.utils import compute_chunksize
-from modin.error_message import ErrorMessage
 
 
 def progress_bar_wrapper(f):
@@ -95,47 +94,21 @@ class PandasOnRayDataframePartitionManager(GenericRayDataframePartitionManager):
     _row_partition_class = PandasOnRayDataframeRowPartition
 
     @classmethod
-    def get_indices(cls, axis, partitions, index_func=None):
+    def get_objects_from_partitions(cls, partitions):
         """
-        Get the internal indices stored in the partitions.
+        Get the objects wrapped by `partitions` in parallel.
 
         Parameters
         ----------
-        axis : {0, 1}
-            Axis to extract the labels over.
         partitions : np.ndarray
             NumPy array with ``PandasDataframePartition``-s.
-        index_func : callable, default: None
-            The function to be used to extract the indices.
 
         Returns
         -------
-        pandas.Index
-            A ``pandas.Index`` object.
-
-        Notes
-        -----
-        These are the global indices of the object. This is mostly useful
-        when you have deleted rows/columns internally, but do not know
-        which ones were deleted.
+        list
+            The objects wrapped by `partitions`.
         """
-        ErrorMessage.catch_bugs_and_request_email(not callable(index_func))
-        func = cls.preprocess_func(index_func)
-        if axis == 0:
-            # We grab the first column of blocks and extract the indices
-            new_idx = (
-                [idx.apply(func).oid for idx in partitions.T[0]]
-                if len(partitions.T)
-                else []
-            )
-        else:
-            new_idx = (
-                [idx.apply(func).oid for idx in partitions[0]]
-                if len(partitions)
-                else []
-            )
-        new_idx = ray.get(new_idx)
-        return new_idx[0].append(new_idx[1:]) if len(new_idx) else new_idx
+        return ray.get([partition.oid for partition in partitions])
 
     @classmethod
     def concat(cls, axis, left_parts, right_parts):


Test patch: