This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/main by this push:
new 988c655597 refactor(pyamber): extract shared base for indexed shuffle
partitioners (#8920)
988c655597 is described below
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)])
