feat: Implement shuffle map side spill support (not flight shuffle)

Author: cancai

daft/context.py

@@ -186,6 +186,7 @@ def set_execution_config(
     maintain_order: bool | None = None,
     enable_dynamic_batching: bool | None = None,
     dynamic_batching_strategy: str | None = None,
+    shuffle_spill_threshold: int | None = None,
 ) -> DaftContext:
     """Globally sets various configuration parameters which control various aspects of Daft execution.
 
@@ -229,6 +230,7 @@ def set_execution_config(
         maintain_order: Whether to maintain order during execution. Defaults to True. Some blocking sink operators (e.g. write_parquet) won't respect this flag and will always keep maintain_order as false, and propagate to child operators. It's useful to set this to False for running df.collect() when no ordering is required.
         enable_dynamic_batching: Whether to enable dynamic batching. Defaults to False.
         dynamic_batching_strategy: The strategy to use for dynamic batching. Defaults to 'auto'.
+        shuffle_spill_threshold: Memory threshold in bytes for shuffle spill. Defaults to None (no spill).
     """
     # Replace values in the DaftExecutionConfig with user-specified overrides
     ctx = get_context()
@@ -265,6 +267,7 @@ def set_execution_config(
             maintain_order=maintain_order,
             enable_dynamic_batching=enable_dynamic_batching,
             dynamic_batching_strategy=dynamic_batching_strategy,
+            shuffle_spill_threshold=shuffle_spill_threshold,
         )
 
         ctx._ctx._daft_execution_config = new_daft_execution_config

daft/daft/__init__.pyi

@@ -2032,18 +2032,20 @@ class LocalPhysicalPlan:
     def from_logical_plan_builder(builder: LogicalPlanBuilder) -> LocalPhysicalPlan: ...
 
 class RayPartitionRef:
-    object_ref: ray.ObjectRef
+    object_refs: list[ray.ObjectRef]
     num_rows: int
     size_bytes: int
 
-    def __init__(self, object_ref: ray.ObjectRef, num_rows: int, size_bytes: int): ...
+    def __init__(self, object_refs: list[ray.ObjectRef], num_rows: int, size_bytes: int): ...
 
 class RaySwordfishTask:
     def name(self) -> str: ...
+    def num_partitions(self) -> int: ...
     def plan(self) -> LocalPhysicalPlan: ...
     def psets(self) -> dict[str, list[RayPartitionRef]]: ...
     def config(self) -> PyDaftExecutionConfig: ...
     def context(self) -> dict[str, str]: ...
+    def is_into_batches(self) -> bool: ...
 
 class RayTaskResult:
     @staticmethod
@@ -2136,6 +2138,7 @@ class PyDaftExecutionConfig:
         maintain_order: bool | None = None,
         enable_dynamic_batching: bool | None = None,
         dynamic_batching_strategy: str | None = None,
+        shuffle_spill_threshold: int | None = None,
     ) -> PyDaftExecutionConfig: ...
     @property
     def enable_scan_task_split_and_merge(self) -> bool: ...

daft/runners/flotilla.py

@@ -75,8 +75,9 @@ async def run_plan(
         from daft.daft import PyDaftContext
 
         with profile():
-            psets = {k: await asyncio.gather(*v) for k, v in psets.items()}
-            psets_mp = {k: [v._micropartition for v in v] for k, v in psets.items()}
+            psets_mp: dict[str, list[PyMicroPartition]] = {
+                k: [ray.get(r)._micropartition for r in refs] for k, refs in psets.items()
+            }
 
             metas = []
             native_executor = NativeExecutor()
@@ -131,6 +132,8 @@ class RaySwordfishTaskHandle:
 
     result_handle: ray.ObjectRef
     actor_handle: ray.actor.ActorHandle
+    num_partitions: int
+    is_into_batches: bool
     task: asyncio.Task[RayTaskResult] | None = None
 
     async def _get_result(self) -> RayTaskResult:
@@ -142,11 +145,34 @@ async def _get_result(self) -> RayTaskResult:
             task_metadata: SwordfishTaskMetadata = await metadata_ref
             assert len(results) == len(task_metadata.partition_metadatas)
 
+            # Pack the results into partitions
+            num_partitions = self.num_partitions
+            partition_refs = []
+
+            # We rely on the task metadata for now because IntoBatches is the only operator
+            # that dynamically generates partitions without a fixed mapping.
+            is_into_batches = self.is_into_batches
+            if is_into_batches:
+                for res, meta in zip(results, task_metadata.partition_metadatas):
+                    partition_refs.append(RayPartitionRef([res], meta.num_rows, meta.size_bytes or 0))
+            else:
+                packed_results: list[list[ray.ObjectRef]] = [[] for _ in range(num_partitions)]
+                packed_metadatas: list[list[PartitionMetadata]] = [[] for _ in range(num_partitions)]
+
+                for i, (res, meta) in enumerate(zip(results, task_metadata.partition_metadatas)):
+                    part_idx = i % num_partitions
+                    packed_results[part_idx].append(res)
+                    packed_metadatas[part_idx].append(meta)
+
+                for i in range(num_partitions):
+                    chunks = packed_results[i]
+                    metas = packed_metadatas[i]
+                    total_rows = sum(m.num_rows for m in metas)
+                    total_bytes = sum(m.size_bytes or 0 for m in metas)
+                    partition_refs.append(RayPartitionRef(chunks, total_rows, total_bytes))
+
             return RayTaskResult.success(
-                [
-                    RayPartitionRef(result, metadata.num_rows, metadata.size_bytes or 0)
-                    for result, metadata in zip(results, task_metadata.partition_metadatas)
-                ],
+                partition_refs,
                 task_metadata.stats,
             )
         except (ray.exceptions.ActorDiedError, ray.exceptions.ActorUnschedulableError):
@@ -179,13 +205,18 @@ def __init__(
         self.actor_handle = actor_handle
 
     def submit_task(self, task: RaySwordfishTask) -> RaySwordfishTaskHandle:
-        psets = {k: [v.object_ref for v in v] for k, v in task.psets().items()}
-        result_handle = self.actor_handle.run_plan.options(name=task.name()).remote(
-            task.plan(), task.config(), psets, task.context()
+        psets = {k: [obj_ref for p in v for obj_ref in p.object_refs] for k, v in task.psets().items()}
+        result_handle = self.actor_handle.run_plan.remote(
+            task.plan(),
+            task.config(),
+            psets,
+            task.context(),
         )
         return RaySwordfishTaskHandle(
-            result_handle,
-            self.actor_handle,
+            result_handle=result_handle,
+            actor_handle=self.actor_handle,
+            num_partitions=task.num_partitions(),
+            is_into_batches=task.is_into_batches(),
         )
 
     def shutdown(self) -> None:
@@ -259,12 +290,17 @@ def __init__(self, dashboard_url: str | None = None) -> None:
     def run_plan(
         self,
         plan: DistributedPhysicalPlan,
-        partition_sets: dict[str, PartitionSet[ray.ObjectRef]],
+        partition_sets: dict[str, PartitionSet[list[ray.ObjectRef]]],
     ) -> None:
-        psets = {
-            k: [RayPartitionRef(v.partition(), v.metadata().num_rows, v.metadata().size_bytes or 0) for v in v.values()]
-            for k, v in partition_sets.items()
-        }
+        psets = {}
+        for k, v in partition_sets.items():
+            partition_refs = []
+            for val in v.values():
+                partition_refs.append(
+                    RayPartitionRef(val.partition(), val.metadata().num_rows, val.metadata().size_bytes or 0)
+                )
+            psets[k] = partition_refs
+
         self.curr_plans[plan.idx()] = plan
         self.curr_result_gens[plan.idx()] = self.plan_runner.run_plan(plan, psets)
 
@@ -289,7 +325,7 @@ async def get_next_partition(self, plan_id: str) -> RayMaterializedResult | Reco
             [PartitionMetadata(next_partition_ref.num_rows, next_partition_ref.size_bytes)]
         )
         materialized_result = RayMaterializedResult(
-            partition=next_partition_ref.object_ref,
+            partition=next_partition_ref.object_refs,
             metadatas=metadata_accessor,
             metadata_idx=0,
         )
@@ -373,7 +409,7 @@ def __init__(self) -> None:
     def stream_plan(
         self,
         plan: DistributedPhysicalPlan,
-        partition_sets: dict[str, PartitionSet[RayMaterializedResult]],
+        partition_sets: dict[str, PartitionSet[list[ray.ObjectRef]]],
     ) -> Generator[RayMaterializedResult, None, RecordBatch]:
         plan_id = plan.idx()
         ray.get(self.runner.run_plan.remote(plan, partition_sets))

daft/runners/ray_runner.py

@@ -282,14 +282,14 @@ def _to_pandas_ref(df: pd.DataFrame | ray.ObjectRef) -> ray.ObjectRef:
         raise ValueError(f"Expected a Ray object ref or a Pandas DataFrame, got {type(df)}")
 
 
-class RayPartitionSet(PartitionSet[ray.ObjectRef]):
+class RayPartitionSet(PartitionSet[list[ray.ObjectRef]]):
     _results: dict[PartID, RayMaterializedResult]
 
     def __init__(self) -> None:
         super().__init__()
         self._results = {}
 
-    def items(self) -> list[tuple[PartID, MaterializedResult[ray.ObjectRef]]]:
+    def items(self) -> list[tuple[PartID, MaterializedResult[list[ray.ObjectRef]]]]:
         return [(pid, result) for pid, result in sorted(self._results.items())]
 
     def _get_merged_micropartition(self, schema: Schema) -> MicroPartition:
@@ -298,22 +298,30 @@ def _get_merged_micropartition(self, schema: Schema) -> MicroPartition:
             assert ids_and_partitions[0][0] == 0
             assert ids_and_partitions[-1][0] + 1 == len(ids_and_partitions)
 
-        all_partitions = ray.get([part.partition() for id, part in ids_and_partitions])
-        return MicroPartition.concat_or_empty(all_partitions, schema)
+        all_refs = []
+        for _, part in ids_and_partitions:
+            all_refs.extend(part.partition())
+
+        all_micropartitions = ray.get(all_refs)
+        return MicroPartition.concat_or_empty(all_micropartitions, schema)
 
     def _get_preview_micropartitions(self, num_rows: int) -> list[MicroPartition]:
         ids_and_partitions = self.items()
         preview_parts = []
         for _, mat_result in ids_and_partitions:
-            ref: ray.ObjectRef = mat_result.partition()
-            part: MicroPartition = ray.get(ref)
-            part_len = len(part)
-            if part_len >= num_rows:  # if this part has enough rows, take what we need and break
-                preview_parts.append(part.slice(0, num_rows))
+            refs: list[ray.ObjectRef] = mat_result.partition()
+            parts: list[MicroPartition] = ray.get(refs)
+            for part in parts:
+                part_len = len(part)
+                if part_len >= num_rows:
+                    preview_parts.append(part.slice(0, num_rows))
+                    num_rows = 0
+                    break
+                else:
+                    num_rows -= part_len
+                    preview_parts.append(part)
+            if num_rows == 0:
                 break
-            else:  # otherwise, take the whole part and keep going
-                num_rows -= part_len
-                preview_parts.append(part)
         return preview_parts
 
     def to_ray_dataset(self) -> RayDataset:
@@ -350,9 +358,9 @@ def _make_dask_dataframe_partition_from_micropartition(partition: MicroPartition
         return cast("dd.DataFrame", dd.from_delayed(ddf_parts, meta=meta))
 
     def get_partition(self, idx: PartID) -> RayMaterializedResult:
-        return self._results[idx].partition()
+        return self._results[idx]
 
-    def set_partition(self, idx: PartID, result: MaterializedResult[ray.ObjectRef]) -> None:
+    def set_partition(self, idx: PartID, result: MaterializedResult[list[ray.ObjectRef]]) -> None:
         assert isinstance(result, RayMaterializedResult)
         self._results[idx] = result
 
@@ -377,8 +385,11 @@ def num_partitions(self) -> int:
         return len(self._results)
 
     def wait(self) -> None:
-        deduped_object_refs = {r.partition() for r in self._results.values()}
-        ray.wait(list(deduped_object_refs), fetch_local=False, num_returns=len(deduped_object_refs))
+        all_refs = []
+        for r in self._results.values():
+            all_refs.extend(r.partition())
+        deduped_object_refs = list(set(all_refs))
+        ray.wait(deduped_object_refs, fetch_local=False, num_returns=len(deduped_object_refs))
 
 
 def _from_arrow_type_with_ray_data_extensions(arrow_type: pa.DataType) -> DataType:
@@ -444,7 +455,7 @@ def partition_set_from_ray_dataset(
         pset = RayPartitionSet()
 
         for i, obj in enumerate(daft_micropartitions):
-            pset.set_partition(i, RayMaterializedResult(obj))
+            pset.set_partition(i, RayMaterializedResult([obj]))
         return (
             pset,
             daft_schema,
@@ -476,7 +487,7 @@ def partition_set_from_dask_dataframe(
         pset = RayPartitionSet()
 
         for i, obj in enumerate(daft_micropartitions):
-            pset.set_partition(i, RayMaterializedResult(obj))
+            pset.set_partition(i, RayMaterializedResult([obj]))
         return (
             pset,
             schemas[0],
@@ -487,7 +498,7 @@ def partition_set_from_dask_dataframe(
 
 
 @ray.remote  # type: ignore[untyped-decorator]
-def get_metas(*partitions: MicroPartition) -> list[PartitionMetadata]:
+def get_metas(partitions: list[MicroPartition]) -> list[PartitionMetadata]:
     return [PartitionMetadata.from_table(partition) for partition in partitions]
 
 
@@ -665,7 +676,7 @@ def run_iter_tables(
         self, builder: LogicalPlanBuilder, results_buffer_size: int | None = None
     ) -> Iterator[MicroPartition]:
         for result in self.run_iter(builder, results_buffer_size=results_buffer_size):
-            yield ray.get(result.partition())
+            yield result.micropartition()
 
     def _collect_into_cache(
         self, results_iter: Generator[RayMaterializedResult, None, RecordBatch]
@@ -689,48 +700,62 @@ def run(self, builder: LogicalPlanBuilder) -> tuple[PartitionCacheEntry, RecordB
         results_iter = self.run_iter(builder)
         return self._collect_into_cache(results_iter)
 
-    def put_partition_set_into_cache(self, pset: PartitionSet[ray.ObjectRef]) -> PartitionCacheEntry:
+    def put_partition_set_into_cache(self, pset: PartitionSet[list[ray.ObjectRef]]) -> PartitionCacheEntry:
         if isinstance(pset, LocalPartitionSet):
             new_pset = RayPartitionSet()
             metadata_accessor = PartitionMetadataAccessor.from_metadata_list([v.metadata() for v in pset.values()])
             for i, (pid, py_mat_result) in enumerate(pset.items()):
-                new_pset.set_partition(
-                    pid, RayMaterializedResult(ray.put(py_mat_result.partition()), metadata_accessor, i)
-                )
+                part = py_mat_result.partition()
+                new_pset.set_partition(pid, RayMaterializedResult([ray.put(part)], metadata_accessor, i))
             pset = new_pset
         return self._part_set_cache.put_partition_set(pset=pset)
 
     def runner_io(self) -> RayRunnerIO:
         return RayRunnerIO()
 
 
-class RayMaterializedResult(MaterializedResult[ray.ObjectRef]):
+class RayMaterializedResult(MaterializedResult[list[ray.ObjectRef]]):
     def __init__(
         self,
-        partition: ray.ObjectRef[Any],
+        partition: list[ray.ObjectRef[Any]],
         metadatas: PartitionMetadataAccessor | None = None,
         metadata_idx: int | None = None,
     ):
+        assert isinstance(partition, list)
         self._partition = partition
         if metadatas is None:
             assert metadata_idx is None
             metadatas = PartitionMetadataAccessor(get_metas.remote(self._partition))
-            metadata_idx = 0
+
         self._metadatas = metadatas
         self._metadata_idx = metadata_idx
 
-    def partition(self) -> ray.ObjectRef:
+    def partition(self) -> list[ray.ObjectRef]:
         return self._partition
 
     def micropartition(self) -> MicroPartition:
-        return ray.get(self._partition)
+        parts = ray.get(self._partition)
+        return MicroPartition.concat(parts)
 
     def metadata(self) -> PartitionMetadata:
-        assert self._metadata_idx is not None
-        return self._metadatas.get_index(self._metadata_idx)
+        all_metas = self._metadatas._get_metadatas()
+
+        if self._metadata_idx is not None:
+            return all_metas[self._metadata_idx]
+
+        total_rows = sum(m.num_rows for m in all_metas)
+        total_bytes = 0
+        for m in all_metas:
+            if m.size_bytes is not None:
+                total_bytes += m.size_bytes
+        return PartitionMetadata(
+            num_rows=total_rows,
+            size_bytes=total_bytes if total_bytes > 0 else None,
+        )
 
     def cancel(self) -> None:
-        return ray.cancel(self._partition)
+        for p in self._partition:
+            ray.cancel(p)
 
     def _noop(self, _: ray.ObjectRef) -> None:
         return None