[Data] Add Join & Aggregations on external shuffle v2 - #65897
Conversation
Signed-off-by: You-Cheng Lin <mses010108@gmail.com>
There was a problem hiding this comment.
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.
| 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),) |
There was a problem hiding this comment.
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.
| 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),) |
| 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)) |
There was a problem hiding this comment.
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>
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
Additional information