This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5989-e9b208ada375a435b7f346851290a07d6ca5e51a in repository https://gitbox.apache.org/repos/asf/texera.git
commit f5c77295a9db77420ecb2dd62eec1f7a05b82e46 Author: Matthew B. <[email protected]> AuthorDate: Tue Jul 7 16:23:22 2026 -0700 refactor(pyamber): share receiver-batch construction in partitioners (#5989) ### What changes were proposed in this PR? - Added `Partitioner.build_receiver_batches(channels)` to the base class, the single home for the ordered `(receiver, batch)` list built via `dict.fromkeys` (which preserves channel order where a set literal would not). - Replaced the three identical inline copies of that construction in `RoundRobinPartitioner`, `HashBasedShufflePartitioner`, and `RangeBasedShufflePartitioner` with a call to the shared helper, dropping the drifted copy-pasted comments. - No behavior change: each partitioner still builds the same ordered, deduplicated receiver list. ### Any related issues, documentation, discussions? Closes: #5988 ### How was this PR tested? - Run `cd amber && PYTHONPATH=src/main/python python -m pytest src/test/python/core/architecture/sendsemantics/test_partitioners.py -q`, expect all 35 tests to pass. - The order/dedup behavior owned by the new helper is covered by the existing `TestRoundRobinPartitioner::test_init_preserves_channel_order` and `test_init_dedupes_duplicate_channels_preserving_first_seen_order`; confirm they still pass against the refactored construction. ### Was this PR authored or co-authored using generative AI tooling? Co-authored with Claude Opus 4.8 in compliance with ASF Co-authored-by: Xinyuan Lin <[email protected]> --- .../sendsemantics/hash_based_shuffle_partitioner.py | 14 ++------------ .../python/core/architecture/sendsemantics/partitioner.py | 12 ++++++++++++ .../sendsemantics/range_based_shuffle_partitioner.py | 14 ++------------ .../architecture/sendsemantics/round_robin_partitioner.py | 15 ++------------- 4 files changed, 18 insertions(+), 37 deletions(-) diff --git a/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py index de018e5a3c..ceee764c1b 100644 --- a/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py +++ b/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py @@ -37,18 +37,8 @@ class HashBasedShufflePartitioner(Partitioner): super().__init__(set_one_of(Partitioning, partitioning)) logger.debug(f"got {partitioning}") self.batch_size = partitioning.batch_size - # Partitioning contains an ordered list of downstream worker ids. - # Currently we are using the index of such an order to choose - # a downstream worker to send tuples to. - # Must use dict.fromkeys to ensure the order of receiver workers - # from partitioning is preserved (using `{}` to create a set - # does not preserve order and will not work correctly.) - self.receivers = [ - (rid, []) - for rid in dict.fromkeys( - channel.to_worker_id for channel in partitioning.channels - ) - ] + # Indexed by hash_code to choose the downstream worker to send to. + self.receivers = self.build_receiver_batches(partitioning.channels) self.hash_attribute_names = partitioning.hash_attribute_names @overrides diff --git a/amber/src/main/python/core/architecture/sendsemantics/partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/partitioner.py index c4aac57cbe..ed2332d5cb 100644 --- a/amber/src/main/python/core/architecture/sendsemantics/partitioner.py +++ b/amber/src/main/python/core/architecture/sendsemantics/partitioner.py @@ -32,6 +32,18 @@ class Partitioner(ABC): def __init__(self, partitioning: Message): self.partitioning: Partitioning = get_one_of(partitioning) + @staticmethod + def build_receiver_batches( + channels, + ) -> typing.List[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]: + # An ordered (receiver, batch) pair per distinct downstream worker. + # dict.fromkeys preserves the channel order; a set literal would not, + # which breaks input-port materialization reader threads. + return [ + (rid, []) + for rid in dict.fromkeys(channel.to_worker_id for channel in channels) + ] + def add_tuple_to_batch( self, tuple_: Tuple ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]: diff --git a/amber/src/main/python/core/architecture/sendsemantics/range_based_shuffle_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/range_based_shuffle_partitioner.py index 28aff35935..3f8f83c504 100644 --- a/amber/src/main/python/core/architecture/sendsemantics/range_based_shuffle_partitioner.py +++ b/amber/src/main/python/core/architecture/sendsemantics/range_based_shuffle_partitioner.py @@ -37,18 +37,8 @@ class RangeBasedShufflePartitioner(Partitioner): super().__init__(set_one_of(Partitioning, partitioning)) logger.info(f"got {partitioning}") self.batch_size = partitioning.batch_size - # Partitioning contains an ordered list of downstream worker ids. - # Currently we are using the index of such an order to choose - # a downstream worker to send tuples to. - # Must use dict.fromkeys to ensure the order of receiver workers - # from partitioning is preserved (using `{}` to create a set - # does not preserve order and will not work correctly.) - self.receivers = [ - (rid, []) - for rid in dict.fromkeys( - channel.to_worker_id for channel in partitioning.channels - ) - ] + # Indexed by get_receiver_index to choose the downstream worker to send to. + self.receivers = self.build_receiver_batches(partitioning.channels) self.range_attribute_names = partitioning.range_attribute_names self.range_min = partitioning.range_min self.range_max = partitioning.range_max diff --git a/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py index 87c3fee87d..21c1a3ca0f 100644 --- a/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py +++ b/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py @@ -35,19 +35,8 @@ class RoundRobinPartitioner(Partitioner): def __init__(self, partitioning: RoundRobinPartitioning): super().__init__(set_one_of(Partitioning, partitioning)) self.batch_size = partitioning.batch_size - # Partitioning contains an ordered list of downstream worker ids. - # Currently we are using the index of such an order to choose - # a downstream worker to send tuples to. - # Must use dict.fromkeys to ensure the order of receiver workers - # from partitioning is preserved (using `{}` to create a set - # does not preserve order and will not work with input-port - # materialization reader threads.) - self.receivers = [ - (rid, []) - for rid in dict.fromkeys( - channel.to_worker_id for channel in partitioning.channels - ) - ] + # Indexed by round_robin_index to choose the downstream worker to send to. + self.receivers = self.build_receiver_batches(partitioning.channels) self.round_robin_index = 0 @overrides
