[PERF] Remove calls to remote_len_partition

daft/daft.pyi

@@ -934,7 +934,7 @@ class PhysicalPlanScheduler:
     def num_partitions(self) -> int: ...
     def to_partition_tasks(
         self, psets: dict[str, list[PartitionT]], is_ray_runner: bool
-    ) -> physical_plan.MaterializedPhysicalPlan: ...
+    ) -> physical_plan.InProgressPhysicalPlan: ...
 
 class LogicalPlanBuilder:
     """

daft/dataframe/dataframe.py

@@ -198,8 +198,9 @@ def iter_partitions(self) -> Iterator[Union[Table, "RayObjectRef"]]:
         else:
             # Execute the dataframe in a streaming fashion.
             context = get_context()
-            partitions_iter = context.runner().run_iter(self._builder)
-            yield from partitions_iter
+            results_iter = context.runner().run_iter(self._builder)
+            for result in results_iter:
+                yield result.partition()
 
     @DataframePublicAPI
     def __repr__(self) -> str:

daft/execution/execution_step.py

@@ -4,7 +4,7 @@
 import pathlib
 import sys
 from dataclasses import dataclass, field
-from typing import Generic, TypeVar
+from typing import Generic
 
 if sys.version_info < (3, 8):
     from typing_extensions import Protocol
@@ -26,14 +26,15 @@
 from daft.logical.map_partition_ops import MapPartitionOp
 from daft.logical.schema import Schema
 from daft.runners.partitioning import (
+    MaterializedResult,
     PartialPartitionMetadata,
     PartitionMetadata,
+    PartitionT,
     TableParseCSVOptions,
     TableReadOptions,
 )
 from daft.table import Table, table_io
 
-PartitionT = TypeVar("PartitionT")
 ID_GEN = itertools.count()
 
 
@@ -251,35 +252,6 @@ def __repr__(self) -> str:
         return super().__str__()
 
 
-class MaterializedResult(Protocol[PartitionT]):
-    """A protocol for accessing the result partition of a PartitionTask.
-
-    Different Runners can fill in their own implementation here.
-    """
-
-    def partition(self) -> PartitionT:
-        """Get the partition of this result."""
-        ...
-
-    def vpartition(self) -> Table:
-        """Get the vPartition of this result."""
-        ...
-
-    def metadata(self) -> PartitionMetadata:
-        """Get the metadata of the partition in this result."""
-        ...
-
-    def cancel(self) -> None:
-        """If possible, cancel execution of this PartitionTask."""
-        ...
-
-    def _noop(self, _: PartitionT) -> None:
-        """Implement this as a no-op.
-        https://peps.python.org/pep-0544/#overriding-inferred-variance-of-protocol-classes
-        """
-        ...
-
-
 class Instruction(Protocol):
     """An instruction is a function to run over a list of partitions.
 

daft/execution/physical_plan.py

@@ -38,19 +38,22 @@
 )
 from daft.expressions import ExpressionsProjection
 from daft.logical.schema import Schema
-from daft.runners.partitioning import PartialPartitionMetadata
+from daft.runners.partitioning import (
+    MaterializedResult,
+    PartialPartitionMetadata,
+    PartitionT,
+)
 
 logger = logging.getLogger(__name__)
 
-PartitionT = TypeVar("PartitionT")
 T = TypeVar("T")
 
 
 # A PhysicalPlan that is still being built - may yield both PartitionTaskBuilders and PartitionTasks.
 InProgressPhysicalPlan = Iterator[Union[None, PartitionTask[PartitionT], PartitionTaskBuilder[PartitionT]]]
 
 # A PhysicalPlan that is complete and will only yield PartitionTasks or final PartitionTs.
-MaterializedPhysicalPlan = Iterator[Union[None, PartitionTask[PartitionT], PartitionT]]
+MaterializedPhysicalPlan = Iterator[Union[None, PartitionTask[PartitionT], MaterializedResult[PartitionT]]]
 
 
 def _stage_id_counter():
@@ -738,7 +741,7 @@ def materialize(
         # Check if any inputs finished executing.
         while len(materializations) > 0 and materializations[0].done():
             done_task = materializations.popleft()
-            yield done_task.partition()
+            yield done_task._result
 
         # Materialize a single dependency.
         try:

daft/execution/rust_physical_plan_shim.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Iterator, TypeVar, cast
+from typing import Iterator, cast
 
 from daft.daft import (
     FileFormat,
@@ -19,11 +19,9 @@
 from daft.expressions import Expression, ExpressionsProjection
 from daft.logical.map_partition_ops import MapPartitionOp
 from daft.logical.schema import Schema
-from daft.runners.partitioning import PartialPartitionMetadata
+from daft.runners.partitioning import PartialPartitionMetadata, PartitionT
 from daft.table import Table
 
-PartitionT = TypeVar("PartitionT")
-
 
 def scan_with_tasks(
     scan_tasks: list[ScanTask],

daft/runners/partitioning.py

@@ -92,6 +92,40 @@
 PartitionT = TypeVar("PartitionT")
 
 
+class MaterializedResult(Generic[PartitionT]):
+    """A protocol for accessing the result partition of a PartitionTask.
+
+    Different Runners can fill in their own implementation here.
+    """
+
+    @abstractmethod
+    def partition(self) -> PartitionT:
+        """Get the partition of this result."""
+        ...
+
+    @abstractmethod
+    def vpartition(self) -> Table:
+        """Get the vPartition of this result."""
+        ...
+
+    @abstractmethod
+    def metadata(self) -> PartitionMetadata:
+        """Get the metadata of the partition in this result."""
+        ...
+
+    @abstractmethod
+    def cancel(self) -> None:
+        """If possible, cancel execution of this PartitionTask."""
+        ...
+
+    @abstractmethod
+    def _noop(self, _: PartitionT) -> None:
+        """Implement this as a no-op.
+        https://peps.python.org/pep-0544/#overriding-inferred-variance-of-protocol-classes
+        """
+        ...
+
+
 class PartitionSet(Generic[PartitionT]):
     def _get_merged_vpartition(self) -> Table:
         raise NotImplementedError()
@@ -126,7 +160,7 @@
         raise NotImplementedError()
 
     @abstractmethod
-    def set_partition(self, idx: PartID, part: PartitionT) -> None:
+    def set_partition(self, idx: PartID, part: MaterializedResult[PartitionT]) -> None:
         raise NotImplementedError()
 
     @abstractmethod

daft/runners/pyrunner.py

@@ -16,13 +16,14 @@
     StorageConfig,
 )
 from daft.execution import physical_plan
-from daft.execution.execution_step import Instruction, MaterializedResult, PartitionTask
+from daft.execution.execution_step import Instruction, PartitionTask
 from daft.filesystem import glob_path_with_stats
 from daft.internal.gpu import cuda_device_count
 from daft.logical.builder import LogicalPlanBuilder
 from daft.logical.schema import Schema
 from daft.runners import runner_io
 from daft.runners.partitioning import (
+    MaterializedResult,
     PartID,
     PartitionCacheEntry,
     PartitionMetadata,
@@ -52,8 +53,8 @@ def _get_merged_vpartition(self) -> Table:
     def get_partition(self, idx: PartID) -> Table:
         return self._partitions[idx]
 
-    def set_partition(self, idx: PartID, part: Table) -> None:
-        self._partitions[idx] = part
+    def set_partition(self, idx: PartID, part: MaterializedResult[Table]) -> None:
+        self._partitions[idx] = part.partition()
 
     def delete_partition(self, idx: PartID) -> None:
         del self._partitions[idx]
@@ -119,11 +120,11 @@ def runner_io(self) -> PyRunnerIO:
         return PyRunnerIO()
 
     def run(self, builder: LogicalPlanBuilder) -> PartitionCacheEntry:
-        partitions = list(self.run_iter(builder))
+        results = list(self.run_iter(builder))
 
         result_pset = LocalPartitionSet({})
-        for i, partition in enumerate(partitions):
-            result_pset.set_partition(i, partition)
+        for i, result in enumerate(results):
+            result_pset.set_partition(i, result)
 
         pset_entry = self.put_partition_set_into_cache(result_pset)
         return pset_entry
@@ -133,7 +134,7 @@ def run_iter(
         builder: LogicalPlanBuilder,
         # NOTE: PyRunner does not run any async execution, so it ignores `results_buffer_size` which is essentially 0
         results_buffer_size: int | None = None,
-    ) -> Iterator[Table]:
+    ) -> Iterator[PyMaterializedResult]:
         # Optimize the logical plan.
         builder = builder.optimize()
         # Finalize the logical plan and get a physical plan scheduler for translating the
@@ -147,13 +148,16 @@ def run_iter(
         # Get executable tasks from planner.
         tasks = plan_scheduler.to_partition_tasks(psets, is_ray_runner=False)
         with profiler("profile_PyRunner.run_{datetime.now().isoformat()}.json"):
-            partitions_gen = self._physical_plan_to_partitions(tasks)
-            yield from partitions_gen
+            results_gen = self._physical_plan_to_partitions(tasks)
+            yield from results_gen
 
     def run_iter_tables(self, builder: LogicalPlanBuilder, results_buffer_size: int | None = None) -> Iterator[Table]:
-        return self.run_iter(builder, results_buffer_size=results_buffer_size)
+        for result in self.run_iter(builder, results_buffer_size=results_buffer_size):
+            yield result.partition()
 
-    def _physical_plan_to_partitions(self, plan: physical_plan.MaterializedPhysicalPlan) -> Iterator[Table]:
+    def _physical_plan_to_partitions(
+        self, plan: physical_plan.MaterializedPhysicalPlan[Table]
+    ) -> Iterator[PyMaterializedResult]:
         inflight_tasks: dict[str, PartitionTask] = dict()
         inflight_tasks_resources: dict[str, ResourceRequest] = dict()
         future_to_task: dict[futures.Future, str] = dict()
@@ -171,7 +175,9 @@ def _physical_plan_to_partitions(self, plan: physical_plan.MaterializedPhysicalP
                             # Blocked on already dispatched tasks; await some tasks.
                             break
 
-                        elif isinstance(next_step, Table):
+                        elif isinstance(next_step, MaterializedResult):
+                            assert isinstance(next_step, PyMaterializedResult)
+
                             # A final result.
                             yield next_step
                             next_step = next(plan)