Repository: modin-project/modin

Issue:
PERF: don't call `compute_sliced_len.remote` when `row_labels/col_labels` == `slice(None)`
This should speedup `partition.mask` operation.


Discussion:

Resolving patch:
diff --git a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py
--- a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py
+++ b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py
@@ -187,13 +187,21 @@ def mask(self, row_labels, col_labels):
         """
         new_obj = super().mask(row_labels, col_labels)
         if isinstance(row_labels, slice) and isinstance(self._length_cache, Future):
-            new_obj._length_cache = DaskWrapper.deploy(
-                compute_sliced_len, row_labels, self._length_cache
-            )
+            if row_labels == slice(None):
+                # fast path - full axis take
+                new_obj._length_cache = self._length_cache
+            else:
+                new_obj._length_cache = DaskWrapper.deploy(
+                    compute_sliced_len, row_labels, self._length_cache
+                )
         if isinstance(col_labels, slice) and isinstance(self._width_cache, Future):
-            new_obj._width_cache = DaskWrapper.deploy(
-                compute_sliced_len, col_labels, self._width_cache
-            )
+            if col_labels == slice(None):
+                # fast path - full axis take
+                new_obj._width_cache = self._width_cache
+            else:
+                new_obj._width_cache = DaskWrapper.deploy(
+                    compute_sliced_len, col_labels, self._width_cache
+                )
         return new_obj
 
     def __copy__(self):
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
@@ -228,15 +228,23 @@ def mask(self, row_labels, col_labels):
         if isinstance(row_labels, slice) and isinstance(
             self._length_cache, ObjectIDType
         ):
-            new_obj._length_cache = compute_sliced_len.remote(
-                row_labels, self._length_cache
-            )
+            if row_labels == slice(None):
+                # fast path - full axis take
+                new_obj._length_cache = self._length_cache
+            else:
+                new_obj._length_cache = compute_sliced_len.remote(
+                    row_labels, self._length_cache
+                )
         if isinstance(col_labels, slice) and isinstance(
             self._width_cache, ObjectIDType
         ):
-            new_obj._width_cache = compute_sliced_len.remote(
-                col_labels, self._width_cache
-            )
+            if col_labels == slice(None):
+                # fast path - full axis take
+                new_obj._width_cache = self._width_cache
+            else:
+                new_obj._width_cache = compute_sliced_len.remote(
+                    col_labels, self._width_cache
+                )
         logger.debug(f"EXIT::Partition.mask::{self._identity}")
         return new_obj
 


Test patch: