Skip to content

[Data] Add Join & Aggregations on external shuffle v2 - #65897

Open
owenowenisme wants to merge 4 commits into
ray-project:masterfrom
owenowenisme:owenowenisme/Implement-join-and-aggregation-planner-on-external-shuffle
Open

[Data] Add Join & Aggregations on external shuffle v2#65897
owenowenisme wants to merge 4 commits into
ray-project:masterfrom
owenowenisme:owenowenisme/Implement-join-and-aggregation-planner-on-external-shuffle

Conversation

@owenowenisme

@owenowenisme owenowenisme commented Sep 3, 2026

Copy link
Copy Markdown
Member

Description

As title, this pr add join & aggregations on external shuffle v2, most code is mirrored from shuffle v2, we can refactor and unify those different operator later.

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.

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>
Signed-off-by: You-Cheng Lin <mses010108@gmail.com>
@owenowenisme owenowenisme added data Ray Data-related issues go add ONLY when ready to merge, run all tests labels Sep 3, 2026
@owenowenisme
owenowenisme marked this pull request as ready for review September 3, 2026 16:05
@owenowenisme
owenowenisme requested a review from a team as a code owner September 3, 2026 16:05

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request extends the external (on-disk, file-transport) hash shuffle implementation to support multi-input operations such as aggregations and joins. Key changes include updating ExternalHashShuffleReduceOp to accept multiple upstream map operators, pairing per-partition wrappers across inputs, and modifying the task bodies to handle multiple input handle lists. Feedback on these changes highlights a potential IndexError if blocks is empty during block transformation, and suggests restoring a table concatenation optimization in _decode_region to prevent performance degradation from processing numerous small tables.

Comment on lines +107 to +118
if block_transformer is not None:
arrow_inputs = [
TableBlockAccessor.try_convert_block_type(block, block_type=BlockType.ARROW)
for block in blocks
if BlockAccessor.for_block(block).num_rows() > 0
] or [
TableBlockAccessor.try_convert_block_type(
blocks[0], block_type=BlockType.ARROW
)
]
combined_input = transform_pyarrow.concat(arrow_inputs, promote_types=True)
blocks = (block_transformer(combined_input),)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If blocks is empty, accessing blocks[0] will raise an IndexError. We should add a guard to ensure blocks is non-empty before attempting to convert and concatenate the blocks.

Suggested change
if block_transformer is not None:
arrow_inputs = [
TableBlockAccessor.try_convert_block_type(block, block_type=BlockType.ARROW)
for block in blocks
if BlockAccessor.for_block(block).num_rows() > 0
] or [
TableBlockAccessor.try_convert_block_type(
blocks[0], block_type=BlockType.ARROW
)
]
combined_input = transform_pyarrow.concat(arrow_inputs, promote_types=True)
blocks = (block_transformer(combined_input),)
if block_transformer is not None and blocks:
arrow_inputs = [
TableBlockAccessor.try_convert_block_type(block, block_type=BlockType.ARROW)
for block in blocks
if BlockAccessor.for_block(block).num_rows() > 0
] or [
TableBlockAccessor.try_convert_block_type(
blocks[0], block_type=BlockType.ARROW
)
]
combined_input = transform_pyarrow.concat(arrow_inputs, promote_types=True)
blocks = (block_transformer(combined_input),)

Comment on lines +429 to +437
def _decode_region(base: int, size: int, accum: List[pa.Table]):
"""Decode one staging-file region's IPC frames into ``accum_tables``."""
pos = base
end = base + size
region_tables: List[pa.Table] = []
while pos < end:
length = struct.unpack(">Q", os.pread(fd, 8, pos))[0]
ipc_buf = os.pread(fd, length, pos + 8)
pos += 8 + length
region_tables.append(_read_ipc(ipc_buf, _compression))
if len(region_tables) == 1:
accum_tables.append(region_tables[0])
elif region_tables:
accum_tables.append(
transform_pyarrow.combine_chunks(
pa.concat_tables(region_tables)
)
)
accum.append(_read_ipc(ipc_buf, _compression))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The original optimization of concatenating all region tables into a single table (using pa.concat_tables and combine_chunks) was removed. Appending each individual IPC frame directly to accum can result in many small tables being passed to reduce_fn, which increases overhead and degrades performance. We should restore this concatenation optimization.

                    def _decode_region(base: int, size: int, accum: List[pa.Table]):
                        """Decode one staging-file region's IPC frames into ``accum_tables``."""
                        pos = base
                        end = base + size
                        region_tables: List[pa.Table] = []
                        while pos < end:
                            length = struct.unpack(">Q", os.pread(fd, 8, pos))[0]
                            ipc_buf = os.pread(fd, length, pos + 8)
                            pos += 8 + length
                            region_tables.append(_read_ipc(ipc_buf, _compression))
                        if len(region_tables) == 1:
                            accum.append(region_tables[0])
                        elif region_tables:
                            accum.append(
                                transform_pyarrow.combine_chunks(
                                    pa.concat_tables(region_tables)
                                )
                            )

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>
Signed-off-by: You-Cheng Lin <mses010108@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

data Ray Data-related issues go add ONLY when ready to merge, run all tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Ray fails to serialize self-reference objects

1 participant