[Data][DNM] new Join op based on hash shuffle v2

[owenowenisme]

Signed-off-by: You-Cheng Lin mses010108@gmail.com> Thank you for contributing to Ray! 🚀

Please review the Ray Contribution Guide before opening a pull request.

⚠️ Remove these instructions before submitting your PR.

💡 Tip: Mark as draft if you want early feedback, or ready for review when it's complete.

Description

Briefly describe what this PR accomplishes and why it's needed.

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request introduces a task-based V2 hash shuffle and join implementation in Ray Data, adding dedicated ShuffleMapOp and ShuffleReduceOp physical operators along with planning logic and configuration options. The code review feedback highlights a critical correctness bug in outer joins where empty inputs cause the join to silently yield zero rows, suggesting that schemas should be passed to the reduce function to properly construct empty tables. Additionally, the feedback recommends improving resource cleanup during operator shutdown by explicitly destroying owned bundles in the queues, and ensuring type consistency by returning a list instead of a tuple for partition sentinels.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Critical Correctness Bug in Outer Joins with Empty Inputs

When one of the join inputs is completely empty (0 blocks), the corresponding ShuffleMapOp will not emit any bundles. In ShuffleReduceOp.all_inputs_done(), the missing input is filled with an empty placeholder RefBundle((), schema=None). This results in left_tables or right_tables being empty in _reduce.

Because of the check:

if not left_tables or not right_tables:
    return

the reduce function returns immediately and yields nothing. This is a critical correctness bug for outer joins (e.g., LEFT_OUTER, RIGHT_OUTER, FULL_OUTER). For example, a LEFT_OUTER join where the right side is empty should preserve all rows from the left side, but here it will silently return 0 rows!

Suggested Fix

Pass the left_schema and right_schema to _make_join_reduce_fn at planning time. If either side is empty, construct an empty table using the schema (e.g., left_schema.empty_table()) so that join_tables can correctly perform the outer join and preserve the non-empty side's rows.

def _make_join_reduce_fn(
    *,
    join_type: JoinType,
    left_key_col_names: Tuple[str, ...],
    right_key_col_names: Tuple[str, ...],
    left_schema: "pa.Schema",
    right_schema: "pa.Schema",
    left_columns_suffix: Optional[str] = None,
    right_columns_suffix: Optional[str] = None,
) -> "ReduceFn":
    """Build a V2-shuffle reduce fn that joins two co-partitioned inputs.

    The returned fn is the binary counterpart of ``JoiningAggregation.finalize``:
    given one partition's shards from both the left (input 0) and right
    (input 1) ``ShuffleMapOp``s, it concatenates each side and applies
    ``join_tables``.  Runs in blocking mode -- the whole partition of both
    inputs must be present before joining.
    """

    def _reduce(
        partition_id: int, tables_by_input: List[List["pa.Table"]]
    ) -> Iterator[Block]:
        assert (
            len(tables_by_input) == 2
        ), f"Join reduce expects exactly two inputs (got {len(tables_by_input)})"
        left_tables, right_tables = tables_by_input[0], tables_by_input[1]
        if not left_tables and not right_tables:
            return
        left_table = _combine(left_tables) if left_tables else left_schema.empty_table()
        right_table = _combine(right_tables) if right_tables else right_schema.empty_table()
        yield join_tables(
            left_table,
            right_table,
            join_type=join_type,
            left_key_col_names=left_key_col_names,
            right_key_col_names=right_key_col_names,
            left_columns_suffix=left_columns_suffix,
            right_columns_suffix=right_columns_suffix,
        )

    return _reduce

[gemini-code-assist]

Pass the left and right schemas to _make_join_reduce_fn to support correct outer join behavior when one of the inputs is empty.

    reduce_fn = _make_join_reduce_fn(
        join_type=logical_op.join_type,
        left_key_col_names=tuple(left_keys),
        right_key_col_names=tuple(right_keys),
        left_schema=physical_children[0].output_schema(),
        right_schema=physical_children[1].output_schema(),
        left_columns_suffix=logical_op.left_columns_suffix,
        right_columns_suffix=logical_op.right_columns_suffix,
    )

[gemini-code-assist]

BlockMetadata.input_files is typically expected to be a List[str]. Returning a List[str] instead of a Tuple[str, ...] ensures type consistency and avoids potential issues if downstream code expects a list.

Suggested change

	def make_partition_sentinel(partition_id: int) -> Tuple[str, ...]:
	return (f"{_PARTITION_ID_SENTINEL}{partition_id}",)
	def make_partition_sentinel(partition_id: int) -> List[str]:
	return [f"{_PARTITION_ID_SENTINEL}{partition_id}"]

[gemini-code-assist]

Explicitly Destroy Bundles in Queues on Shutdown

In _do_shutdown, self._partition_staging and self._output_queue are cleared, but destroy_if_owned() is not called on the bundles they contain. If these bundles own blocks, their object store leases/references might be held until Python's garbage collector runs. Explicitly calling destroy_if_owned() ensures immediate resource cleanup.

Suggested change

	def _do_shutdown(self, force: bool = False) -> None:
	super()._do_shutdown(force)
	self._shuffle_map_tasks.clear()
	self._merge_buffer_refs_by_node.clear()
	for bundles in self._merge_buffer_bundles_by_node.values():
	for bundle in bundles:
	bundle.destroy_if_owned()
	self._merge_buffer_bundles_by_node.clear()
	self._merge_buffer_bytes_by_node.clear()
	for queue in self._partition_staging.values():
	queue.clear()
	self._output_queue.clear()
	def _do_shutdown(self, force: bool = False) -> None:
	super()._do_shutdown(force)
	self._shuffle_map_tasks.clear()
	self._merge_buffer_refs_by_node.clear()
	for bundles in self._merge_buffer_bundles_by_node.values():
	for bundle in bundles:
	bundle.destroy_if_owned()
	self._merge_buffer_bundles_by_node.clear()
	self._merge_buffer_bytes_by_node.clear()
	for queue in self._partition_staging.values():
	while queue.has_next():
	queue.get_next().destroy_if_owned()
	queue.clear()
	while self._output_queue.has_next():
	self._output_queue.get_next().destroy_if_owned()
	self._output_queue.clear()

[gemini-code-assist]

Explicitly Destroy Bundles in Output Queue on Shutdown

In _do_shutdown, self._output_queue is cleared, but destroy_if_owned() is not called on the bundles it contains. Explicitly calling destroy_if_owned() ensures immediate resource cleanup.

Suggested change

	def _do_shutdown(self, force: bool = False) -> None:
	super()._do_shutdown(force)
	self._shuffle_reduce_tasks.clear()
	self._output_queue.clear()
	for pending in self._pending_inputs.values():
	for bundle in pending.values():
	bundle.destroy_if_owned()
	self._pending_inputs.clear()
	def _do_shutdown(self, force: bool = False) -> None:
	super()._do_shutdown(force)
	self._shuffle_reduce_tasks.clear()
	while self._output_queue:
	self._output_queue.popleft().destroy_if_owned()
	self._output_queue.clear()
	for pending in self._pending_inputs.values():
	for bundle in pending.values():
	bundle.destroy_if_owned()
	self._pending_inputs.clear()

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit 4a4cafc. Configure here.}

[cursor]

Outer join drops missing-side rows

High Severity

When V2 join reduce sees an empty shard list for either input, it returns without calling join_tables. After all_inputs_done, a blockless side is represented by a schema-less placeholder bundle, so shard gathering yields [] even when the other side has rows. Left/right/full outer joins then omit preserved-side rows (e.g. left outer with an empty right dataset).

Additional Locations (1)

• python/ray/data/_internal/execution/operators/shuffle_operators/shuffle_reduce_operator.py#L180-L187

 

^{Reviewed by Cursor Bugbot for commit 4a4cafc. Configure here.}

[cursor]

Map emits fewer partition bundles

Medium Severity

After all map work finishes, _maybe_emit_partition_bundles skips any partition whose staging queue is empty instead of emitting a schema-carrying bundle. If no shuffle map tasks ran because every upstream bundle had no block refs, no partition outputs are queued despite the documented “exactly num_partitions bundles” contract, so hash repartition can return far fewer blocks than requested.

 

^{Reviewed by Cursor Bugbot for commit 4a4cafc. Configure here.}