Test

modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py

@@ -94,10 +94,8 @@ def apply(self, func: Callable, *args, **kwargs):
 
         Note: the function is not applied immediately but is added to the execution tree.
         """
-        self.drain_call_queue()
         de = DeferredExecution(self._data, func, args, kwargs)
-        de.ref_count(1)
-        return self.__constructor__(*de.exec())
+        return self.__constructor__(de)
 
     @_inherit_docstrings(PandasDataframePartition.add_to_apply_calls)
     def add_to_apply_calls(self, func, *args, length=None, width=None, **kwargs):
@@ -234,8 +232,8 @@ def length(self, materialize=True):
         int or ray.ObjectRef
             The length of the object.
         """
-        self.drain_call_queue()
         if self._length_cache is None:
+            self.drain_call_queue()
             self._length_cache, self._width_cache = _get_index_and_columns.remote(
                 self._data
             )
@@ -259,8 +257,8 @@ def width(self, materialize=True):
         int or ray.ObjectRef
             The width of the object.
         """
-        self.drain_call_queue()
         if self._width_cache is None:
+            self.drain_call_queue()
             self._length_cache, self._width_cache = _get_index_and_columns.remote(
                 self._data
             )
@@ -284,8 +282,8 @@ def ip(self, materialize=True):
         str
             IP address of the node that holds the data.
         """
-        self.drain_call_queue()
         if self._ip_cache is None:
+            self.drain_call_queue()
             self._ip_cache = _get_ip.remote()
         if materialize and isinstance(self._ip_cache, ObjectIDType):
             self._ip_cache = RayWrapper.materialize(self._ip_cache)