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-8920-65a748ab664ea2e479fd291385514a8273c0c645 in repository https://gitbox.apache.org/repos/asf/texera.git
commit 988c655597f30e89cd2887dc8c7f27499ad0fef8 Author: Matthew B. <[email protected]> AuthorDate: Thu Oct 15 02:37:09 2026 +0000 refactor(pyamber): extract shared base for indexed shuffle partitioners (#8920) ### What changes were proposed in this PR? - Add `IndexedShufflePartitioner` (new `indexed_shuffle_partitioner.py`), a base that owns the shared shuffle machinery: the `__init__` (batch size + per-worker receiver slots), `add_tuple_to_batch`, `flush`, and `flush_state`. - Reduce `RoundRobinPartitioner`, `HashBasedShufflePartitioner`, and `RangeBasedShufflePartitioner` to just their routing: an abstract `_route(tuple_) -> int`, an `_advance()` hook (round-robin bumps its index), and a `_clear_batch_on_flush` flag (round-robin drains a flushed slot; hash/range do not). Each subclass shrinks by roughly 50 lines. - The two round-robin-only quirks (index advance, clearing a flushed slot) are made explicit via the hook and flag rather than changed, so behavior is preserved. `OneToOnePartitioner` and `BroadcastPartitioner` keep their distinct single-batch structure and are untouched. ### Any related issues, documentation, discussions? Closes: #8919 ### How was this PR tested? - Run the partitioner and output-manager suites from `amber/`: `python -m pytest src/test/python/core/architecture/sendsemantics/test_partitioners.py src/test/python/core/architecture/packaging/test_output_manager.py`. Expect all green: the 74 existing specs (which assert exact emitted tuples, residual slot state, `round_robin_index`, and `get_receiver_index` outputs) pass unchanged, serving as the behavior-preservation anchor. - Two tests were added to pin the one genuine divergence: round-robin clears a flushed slot's batch, while hash/range leave it intact. ### Was this PR authored or co-authored using generative AI tooling? Co-authored with Claude Opus 4.8 in compliance with ASF --- .../hash_based_shuffle_partitioner.py | 49 +++----------------- ...rtitioner.py => indexed_shuffle_partitioner.py} | 43 +++++++++++------ .../range_based_shuffle_partitioner.py | 49 +++----------------- .../sendsemantics/round_robin_partitioner.py | 54 +++++----------------- .../sendsemantics/test_partitioners.py | 17 +++++++ 5 files changed, 73 insertions(+), 139 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 ceee764c1b..3cf64256a1 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 @@ -15,65 +15,30 @@ # specific language governing permissions and limitations # under the License. -import typing from loguru import logger from overrides import overrides -from typing import Iterator -from core.architecture.sendsemantics.partitioner import Partitioner +from core.architecture.sendsemantics.indexed_shuffle_partitioner import ( + IndexedShufflePartitioner, +) from core.models import Tuple -from core.models.state import State -from core.util import set_one_of -from proto.org.apache.texera.amber.core import ActorVirtualIdentity -from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage from proto.org.apache.texera.amber.engine.architecture.sendsemantics import ( HashBasedShufflePartitioning, - Partitioning, ) -class HashBasedShufflePartitioner(Partitioner): +class HashBasedShufflePartitioner(IndexedShufflePartitioner): def __init__(self, partitioning: HashBasedShufflePartitioning): - super().__init__(set_one_of(Partitioning, partitioning)) + super().__init__(partitioning) logger.debug(f"got {partitioning}") - self.batch_size = partitioning.batch_size # 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 - def add_tuple_to_batch( - self, tuple_: Tuple - ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]: + def _route(self, tuple_: Tuple) -> int: partial_tuple = ( tuple_ if not self.hash_attribute_names else tuple_.get_partial_tuple(self.hash_attribute_names) ) - hash_code = hash(partial_tuple) % len(self.receivers) - receiver, batch = self.receivers[hash_code] - batch.append(tuple_) - if len(batch) == self.batch_size: - yield receiver, batch - self.receivers[hash_code] = (receiver, list()) - - @overrides - def flush( - self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage - ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]: - for receiver, batch in self.receivers: - if receiver == to: - if len(batch) > 0: - yield batch - yield ecm - - @overrides - def flush_state( - self, state: State - ) -> Iterator[ - typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]] - ]: - for receiver, batch in self.receivers: - if len(batch) > 0: - yield receiver, batch - yield receiver, state + return hash(partial_tuple) % len(self.receivers) diff --git a/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py b/amber/src/main/python/core/architecture/sendsemantics/indexed_shuffle_partitioner.py similarity index 66% copy from amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py copy to amber/src/main/python/core/architecture/sendsemantics/indexed_shuffle_partitioner.py index 21c1a3ca0f..bd12214307 100644 --- a/amber/src/main/python/core/architecture/sendsemantics/round_robin_partitioner.py +++ b/amber/src/main/python/core/architecture/sendsemantics/indexed_shuffle_partitioner.py @@ -16,6 +16,7 @@ # under the License. import typing +from abc import abstractmethod from overrides import overrides from typing import Iterator @@ -25,30 +26,44 @@ from core.models.state import State from core.util import set_one_of from proto.org.apache.texera.amber.core import ActorVirtualIdentity from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage -from proto.org.apache.texera.amber.engine.architecture.sendsemantics import ( - Partitioning, - RoundRobinPartitioning, -) +from proto.org.apache.texera.amber.engine.architecture.sendsemantics import Partitioning -class RoundRobinPartitioner(Partitioner): - def __init__(self, partitioning: RoundRobinPartitioning): +class IndexedShufflePartitioner(Partitioner): + """Base for partitioners that keep one (receiver, batch) slot per downstream + worker and route each tuple to a slot by a computed index. + + Subclasses provide the routing via ``_route`` and, if needed, advance any + routing state via ``_advance``. ``_clear_batch_on_flush`` controls whether a + yielded batch is cleared in place during ``flush``/``flush_state``. + """ + + # Whether flush/flush_state clear a yielded batch in place after emitting it. + _clear_batch_on_flush = False + + def __init__(self, partitioning): super().__init__(set_one_of(Partitioning, partitioning)) self.batch_size = partitioning.batch_size - # 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 + + @abstractmethod + def _route(self, tuple_: Tuple) -> int: + """Return the index of the receiver slot this tuple is routed to.""" + + def _advance(self) -> None: + """Advance any routing state after a tuple is placed. No-op by default.""" @overrides def add_tuple_to_batch( self, tuple_: Tuple ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]: - receiver, batch = self.receivers[self.round_robin_index] + index = self._route(tuple_) + receiver, batch = self.receivers[index] batch.append(tuple_) if len(batch) == self.batch_size: yield receiver, batch - self.receivers[self.round_robin_index] = (receiver, list()) - self.round_robin_index = (self.round_robin_index + 1) % len(self.receivers) + self.receivers[index] = (receiver, list()) + self._advance() @overrides def flush( @@ -58,7 +73,8 @@ class RoundRobinPartitioner(Partitioner): if receiver == to: if len(batch) > 0: yield batch - batch.clear() + if self._clear_batch_on_flush: + batch.clear() yield ecm @overrides @@ -70,5 +86,6 @@ class RoundRobinPartitioner(Partitioner): for receiver, batch in self.receivers: if len(batch) > 0: yield receiver, batch - batch.clear() + if self._clear_batch_on_flush: + batch.clear() yield receiver, state 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 3f8f83c504..3059e8db41 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 @@ -15,30 +15,23 @@ # specific language governing permissions and limitations # under the License. -import typing from loguru import logger from overrides import overrides -from typing import Iterator -from core.architecture.sendsemantics.partitioner import Partitioner +from core.architecture.sendsemantics.indexed_shuffle_partitioner import ( + IndexedShufflePartitioner, +) from core.models import Tuple -from core.models.state import State -from core.util import set_one_of -from proto.org.apache.texera.amber.core import ActorVirtualIdentity -from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage from proto.org.apache.texera.amber.engine.architecture.sendsemantics import ( RangeBasedShufflePartitioning, - Partitioning, ) -class RangeBasedShufflePartitioner(Partitioner): +class RangeBasedShufflePartitioner(IndexedShufflePartitioner): def __init__(self, partitioning: RangeBasedShufflePartitioning): - super().__init__(set_one_of(Partitioning, partitioning)) + super().__init__(partitioning) logger.info(f"got {partitioning}") - self.batch_size = partitioning.batch_size # 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 @@ -59,34 +52,6 @@ class RangeBasedShufflePartitioner(Partitioner): return int((column_val - self.range_min) // self.keys_per_receiver) @overrides - def add_tuple_to_batch( - self, tuple_: Tuple - ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]: + def _route(self, tuple_: Tuple) -> int: column_val = tuple_[self.range_attribute_names[0]] - receiver_index = self.get_receiver_index(column_val) - receiver, batch = self.receivers[receiver_index] - batch.append(tuple_) - if len(batch) == self.batch_size: - yield receiver, batch - self.receivers[receiver_index] = (receiver, list()) - - @overrides - def flush( - self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage - ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]: - for receiver, batch in self.receivers: - if receiver == to: - if len(batch) > 0: - yield batch - yield ecm - - @overrides - def flush_state( - self, state: State - ) -> Iterator[ - typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]] - ]: - for receiver, batch in self.receivers: - if len(batch) > 0: - yield receiver, batch - yield receiver, state + return self.get_receiver_index(column_val) 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 21c1a3ca0f..c220ee3dd0 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 @@ -15,60 +15,30 @@ # specific language governing permissions and limitations # under the License. -import typing from overrides import overrides -from typing import Iterator -from core.architecture.sendsemantics.partitioner import Partitioner +from core.architecture.sendsemantics.indexed_shuffle_partitioner import ( + IndexedShufflePartitioner, +) from core.models import Tuple -from core.models.state import State -from core.util import set_one_of -from proto.org.apache.texera.amber.core import ActorVirtualIdentity -from proto.org.apache.texera.amber.engine.architecture.rpc import EmbeddedControlMessage from proto.org.apache.texera.amber.engine.architecture.sendsemantics import ( - Partitioning, RoundRobinPartitioning, ) -class RoundRobinPartitioner(Partitioner): +class RoundRobinPartitioner(IndexedShufflePartitioner): + # Unlike the hash/range shuffles, round-robin drains a slot's batch on flush. + _clear_batch_on_flush = True + def __init__(self, partitioning: RoundRobinPartitioning): - super().__init__(set_one_of(Partitioning, partitioning)) - self.batch_size = partitioning.batch_size + super().__init__(partitioning) # 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 - def add_tuple_to_batch( - self, tuple_: Tuple - ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]: - receiver, batch = self.receivers[self.round_robin_index] - batch.append(tuple_) - if len(batch) == self.batch_size: - yield receiver, batch - self.receivers[self.round_robin_index] = (receiver, list()) - self.round_robin_index = (self.round_robin_index + 1) % len(self.receivers) - - @overrides - def flush( - self, to: ActorVirtualIdentity, ecm: EmbeddedControlMessage - ) -> Iterator[typing.Union[EmbeddedControlMessage, typing.List[Tuple]]]: - for receiver, batch in self.receivers: - if receiver == to: - if len(batch) > 0: - yield batch - batch.clear() - yield ecm + def _route(self, tuple_: Tuple) -> int: + return self.round_robin_index @overrides - def flush_state( - self, state: State - ) -> Iterator[ - typing.Tuple[ActorVirtualIdentity, typing.Union[State, typing.List[Tuple]]] - ]: - for receiver, batch in self.receivers: - if len(batch) > 0: - yield receiver, batch - batch.clear() - yield receiver, state + def _advance(self) -> None: + self.round_robin_index = (self.round_robin_index + 1) % len(self.receivers) diff --git a/amber/src/test/python/core/architecture/sendsemantics/test_partitioners.py b/amber/src/test/python/core/architecture/sendsemantics/test_partitioners.py index 2d40c53a41..7d70eb37de 100644 --- a/amber/src/test/python/core/architecture/sendsemantics/test_partitioners.py +++ b/amber/src/test/python/core/architecture/sendsemantics/test_partitioners.py @@ -210,6 +210,14 @@ class TestRoundRobinPartitioner: # A's batch is now drained, B's pending batch remains untouched assert partitioner.receivers[1][1] == [_tuple(k=2)] + def test_flush_clears_the_flushed_slot(self, partitioner): + # Round-robin drains a flushed slot's batch in place + # (_clear_batch_on_flush=True), unlike the hash/range shuffles. + list(partitioner.add_tuple_to_batch(_tuple(k=1))) # → A (index 0) + ecm = EmbeddedControlMessage() + list(partitioner.flush(_worker("A"), ecm)) + assert partitioner.receivers[0][1] == [] + def test_flush_to_unknown_receiver_emits_nothing(self, partitioner): ecm = EmbeddedControlMessage() assert list(partitioner.flush(_worker("Z"), ecm)) == [] @@ -282,6 +290,15 @@ class TestHashBasedShufflePartitioner: a_out = _snapshot(p.flush(p.receivers[0][0], ecm)) assert a_out == [[_hashable_tuple(k=1)], ecm] + def test_flush_does_not_clear_the_flushed_slot(self): + # Hash/range shuffles leave a flushed slot's batch intact + # (_clear_batch_on_flush=False), unlike round-robin. + p = self._partitioner(batch_size=10) + p.receivers[0] = (p.receivers[0][0], [_hashable_tuple(k=1)]) + ecm = EmbeddedControlMessage() + list(p.flush(p.receivers[0][0], ecm)) + assert p.receivers[0][1] == [_hashable_tuple(k=1)] + def test_flush_state_emits_pending_batches_and_state(self): p = self._partitioner(batch_size=10) p.receivers[0] = (p.receivers[0][0], [_hashable_tuple(k=1)])
