This is an automated email from the ASF dual-hosted git repository.
spetz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iggy.git
The following commit(s) were added to refs/heads/master by this push:
new be3663805 feat(python): support message partitioning strategies (#3927)
be3663805 is described below
commit be3663805829492a41c971b3056f2da8680e2035
Author: Gunther Xing <[email protected]>
AuthorDate: Tue Sep 8 02:23:55 2026 +0800
feat(python): support message partitioning strategies (#3927)
Closes #3896
---
examples/python/README.md | 2 +-
examples/python/basic/producer.py | 8 +-
foreign/python/apache_iggy.pyi | 66 ++++++++++--
foreign/python/src/client.rs | 37 +++++--
foreign/python/src/lib.rs | 3 +
foreign/python/src/options.rs | 1 +
foreign/python/src/partitioning.rs | 105 ++++++++++++++++++
foreign/python/tests/test_message_operations.py | 137 ++++++++++++++++++++++++
8 files changed, 341 insertions(+), 18 deletions(-)
diff --git a/examples/python/README.md b/examples/python/README.md
index 4e5c0c740..e10810592 100644
--- a/examples/python/README.md
+++ b/examples/python/README.md
@@ -12,7 +12,7 @@ docker run --rm -p 8080:8080 -p 3000:3000 -p 8090:8090 \
-e IGGY_NODE_ADVERTISED_ADDRESS=localhost apache/iggy:latest
# Or build from source (recommended for development)
-cd ../../ && cargo run --bin iggy-server
+cd ../../ && cargo run --bin iggy-server -- --with-default-root-credentials
--fresh
```
For server configuration options and help:
diff --git a/examples/python/basic/producer.py
b/examples/python/basic/producer.py
index cf3e4b84a..13a4ae34b 100644
--- a/examples/python/basic/producer.py
+++ b/examples/python/basic/producer.py
@@ -19,7 +19,7 @@ import argparse
import asyncio
from typing import NamedTuple
-from apache_iggy import IggyClient, StreamDetails, TopicDetails
+from apache_iggy import IggyClient, Partitioning, StreamDetails, TopicDetails
from apache_iggy import SendMessage as Message
from loguru import logger
@@ -106,7 +106,11 @@ async def produce_messages(client: IggyClient):
await client.send_messages(
stream=STREAM_NAME,
topic=TOPIC_NAME,
- partitioning=PARTITION_ID,
+ # A fixed strategy sends the whole batch to this partition.
+ # For topics with multiple partitions, Partitioning.balanced()
+ # distributes batches round-robin, while
+ # Partitioning.messages_key(key) keeps equal keys on the same
partition.
+ partitioning=Partitioning.partition_id(PARTITION_ID),
messages=messages,
)
n_sent_batches += 1
diff --git a/foreign/python/apache_iggy.pyi b/foreign/python/apache_iggy.pyi
index 0400a1a51..ffed2857d 100644
--- a/foreign/python/apache_iggy.pyi
+++ b/foreign/python/apache_iggy.pyi
@@ -46,6 +46,7 @@ __all__ = [
"MaxTopicSize",
"OptionSpec",
"Partition",
+ "Partitioning",
"Permissions",
"PollingStrategy",
"QuicConfig",
@@ -1570,15 +1571,32 @@ class IggyClient:
self,
stream: builtins.str | builtins.int,
topic: builtins.str | builtins.int,
- partitioning: builtins.int,
+ partitioning: Partitioning | builtins.int,
messages: list[SendMessage],
) -> collections.abc.Awaitable[SendMessagesResponse]:
r"""
- Sends a list of messages to the specified topic.
- Returns a SendMessagesResponse carrying the per-partition commit
- confirmations, or a PyRuntimeError on failure. The confirmation list is
- empty when the server reports no offsets, and the legacy server never
- reports any.
+ Sends a batch of messages to a topic using the selected partitioning
strategy.
+
+ Args:
+ stream: Stream identifier as `str | int`.
+ topic: Topic identifier as `str | int`.
+ partitioning: A `Partitioning` strategy or an integer partition ID.
+ Use `Partitioning.balanced()`,
`Partitioning.partition_id(id)`, or
+ `Partitioning.messages_key(key)`. An integer is shorthand for
+ `Partitioning.partition_id(id)`.
+ messages: Messages to send as `list[SendMessage]`.
+
+ Returns:
+ An awaitable that resolves to `SendMessagesResponse`. Its
confirmations
+ report the committed partition and batch base offset. The list is
empty
+ when the server reports no offsets, including on the legacy server.
+
+ Raises:
+ ValueError: If a string stream or topic identifier is invalid.
+ TypeError: If `partitioning` or `messages` has an unsupported type.
+ OverflowError: If a numeric stream, topic, or partition ID is
outside
+ the supported unsigned 32-bit range.
+ RuntimeError: If the request fails.
"""
def poll_messages(
self,
@@ -1889,6 +1907,42 @@ class Partition:
The number of messages in the partition.
"""
[email protected]
+class Partitioning:
+ r"""
+ Defines how a batch of messages is assigned to a topic partition.
+ """
+ @staticmethod
+ def balanced() -> Partitioning:
+ r"""
+ Routes the batch to one partition selected by round-robin.
+ """
+ @staticmethod
+ def partition_id(partition_id: builtins.int) -> Partitioning:
+ r"""
+ Routes the batch to the specified partition.
+
+ `partition_id` must be between 0 and `2**32 - 1`. The topic must
contain
+ that partition when the batch is sent.
+
+ Raises:
+ TypeError: If `partition_id` is not an integer.
+ OverflowError: If `partition_id` is outside the supported unsigned
+ 32-bit range.
+ """
+ @staticmethod
+ def messages_key(key: builtins.str | bytes) -> Partitioning:
+ r"""
+ Routes the batch to one partition selected by hashing `key`.
+
+ `key` may be `str` or `bytes`. Strings are encoded as UTF-8; the
encoded
+ key must contain between 1 and 255 bytes.
+
+ Raises:
+ ValueError: If the encoded key is empty or exceeds 255 bytes.
+ TypeError: If `key` is not `str` or `bytes`.
+ """
+
@typing.final
class Permissions:
r"""
diff --git a/foreign/python/src/client.rs b/foreign/python/src/client.rs
index 7d19de995..168ef2f6d 100644
--- a/foreign/python/src/client.rs
+++ b/foreign/python/src/client.rs
@@ -19,7 +19,7 @@ use bytes::Bytes;
use iggy::prelude::{
AutoCommit as RustAutoCommit, Consumer as RustConsumer, IggyClient as
RustIggyClient,
IggyExpiry as RustIggyExpiry, IggyMessage as RustMessage, MaxTopicSize as
RustMaxTopicSize,
- PollingStrategy as RustPollingStrategy, *,
+ Partitioning as RustPartitioning, PollingStrategy as RustPollingStrategy,
*,
};
use pyo3::PyRef;
use pyo3::prelude::*;
@@ -40,6 +40,7 @@ use crate::consumer::{
use crate::duration::{py_delta_to_iggy_duration, reject_zero};
use crate::identifier::PyIdentifier;
use crate::options::OptionSpec as PyOptionSpec;
+use crate::partitioning::PyPartitioning;
use crate::permissions::Permissions as PyPermissions;
use crate::receive_message::{PollingStrategy, ReceiveMessage};
use crate::send_message::{SendMessage, SendMessagesResponse as
PySendMessagesResponse};
@@ -1190,18 +1191,36 @@ impl IggyClient {
})
}
- /// Sends a list of messages to the specified topic.
- /// Returns a SendMessagesResponse carrying the per-partition commit
- /// confirmations, or a PyRuntimeError on failure. The confirmation list is
- /// empty when the server reports no offsets, and the legacy server never
- /// reports any.
+ /// Sends a batch of messages to a topic using the selected partitioning
strategy.
+ ///
+ /// Args:
+ /// stream: Stream identifier as `str | int`.
+ /// topic: Topic identifier as `str | int`.
+ /// partitioning: A `Partitioning` strategy or an integer partition ID.
+ /// Use `Partitioning.balanced()`,
`Partitioning.partition_id(id)`, or
+ /// `Partitioning.messages_key(key)`. An integer is shorthand for
+ /// `Partitioning.partition_id(id)`.
+ /// messages: Messages to send as `list[SendMessage]`.
+ ///
+ /// Returns:
+ /// An awaitable that resolves to `SendMessagesResponse`. Its
confirmations
+ /// report the committed partition and batch base offset. The list is
empty
+ /// when the server reports no offsets, including on the legacy server.
+ ///
+ /// Raises:
+ /// ValueError: If a string stream or topic identifier is invalid.
+ /// TypeError: If `partitioning` or `messages` has an unsupported type.
+ /// OverflowError: If a numeric stream, topic, or partition ID is
outside
+ /// the supported unsigned 32-bit range.
+ /// RuntimeError: If the request fails.
#[gen_stub(override_return_type(type_repr="collections.abc.Awaitable[SendMessagesResponse]",
imports=("collections.abc")))]
fn send_messages<'a>(
&self,
py: Python<'a>,
stream: PyIdentifier,
topic: PyIdentifier,
- partitioning: u32,
+ #[gen_stub(override_type(type_repr = "Partitioning | builtins.int"))]
+ partitioning: PyPartitioning,
#[gen_stub(override_type(type_repr = "list[SendMessage]"))] messages:
&Bound<'_, PyList>,
) -> PyResult<Bound<'a, PyAny>> {
let messages: Vec<SendMessage> = messages
@@ -1218,7 +1237,7 @@ impl IggyClient {
let stream = Identifier::try_from(stream)?;
let topic = Identifier::try_from(topic)?;
- let partitioning = Partitioning::partition_id(partitioning);
+ let partitioning = RustPartitioning::from(partitioning);
let inner = self.inner.clone();
future_into_py(py, async move {
@@ -1246,7 +1265,7 @@ impl IggyClient {
polling_strategy: &PollingStrategy,
count: u32,
auto_commit: bool,
- partition_id: Option<u32>,
+ #[gen_stub(override_type(type_repr = "builtins.int | None"))]
partition_id: Option<u32>,
) -> PyResult<Bound<'a, PyAny>> {
let consumer = RustConsumer::try_from(consumer)?;
let stream = Identifier::try_from(stream)?;
diff --git a/foreign/python/src/lib.rs b/foreign/python/src/lib.rs
index 80a123ee0..78b777918 100644
--- a/foreign/python/src/lib.rs
+++ b/foreign/python/src/lib.rs
@@ -21,6 +21,7 @@ mod consumer;
mod duration;
mod identifier;
mod options;
+mod partitioning;
mod permissions;
mod receive_message;
mod send_message;
@@ -39,6 +40,7 @@ use consumer::{
ConsumerGroupMember, IggyConsumer, ReceiveMessageIterator,
};
use options::OptionSpec;
+use partitioning::Partitioning;
use permissions::{GlobalPermissions, Permissions, StreamPermissions,
TopicPermissions};
use pyo3::prelude::*;
use receive_message::{PollingStrategy, ReceiveMessage};
@@ -73,6 +75,7 @@ fn apache_iggy(_py: Python, m: &Bound<'_, PyModule>) ->
PyResult<()> {
m.add_class::<IggyExpiry>()?;
m.add_class::<MaxTopicSize>()?;
m.add_class::<OptionSpec>()?;
+ m.add_class::<Partitioning>()?;
m.add_class::<Partition>()?;
m.add_class::<Consumer>()?;
m.add_class::<ConsumerGroup>()?;
diff --git a/foreign/python/src/options.rs b/foreign/python/src/options.rs
index fa489c550..a0d4fa809 100644
--- a/foreign/python/src/options.rs
+++ b/foreign/python/src/options.rs
@@ -64,6 +64,7 @@ impl OptionSpec {
/// options ride that codec.
#[gen_stub(override_return_type(type_repr = "HeaderValue | None"))]
#[getter]
+ #[gen_stub(override_return_type(type_repr = "HeaderValue | None"))]
pub fn default_value<'a>(&self, py: Python<'a>) ->
PyResult<Option<Bound<'a, HeaderValue>>> {
if self.inner.default_value.is_empty() {
return Ok(None);
diff --git a/foreign/python/src/partitioning.rs
b/foreign/python/src/partitioning.rs
new file mode 100644
index 000000000..83d9715d8
--- /dev/null
+++ b/foreign/python/src/partitioning.rs
@@ -0,0 +1,105 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use iggy::prelude::Partitioning as RustPartitioning;
+use pyo3::{exceptions::PyValueError, prelude::*, types::PyBytes};
+use pyo3_stub_gen::{
+ derive::{gen_stub_pyclass, gen_stub_pymethods},
+ impl_stub_type,
+};
+
+/// Defines how a batch of messages is assigned to a topic partition.
+#[derive(Clone)]
+#[pyclass(from_py_object)]
+#[gen_stub_pyclass]
+pub struct Partitioning {
+ pub(crate) inner: RustPartitioning,
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl Partitioning {
+ /// Routes the batch to one partition selected by round-robin.
+ #[staticmethod]
+ pub fn balanced() -> Self {
+ Self {
+ inner: RustPartitioning::balanced(),
+ }
+ }
+
+ /// Routes the batch to the specified partition.
+ ///
+ /// `partition_id` must be between 0 and `2**32 - 1`. The topic must
contain
+ /// that partition when the batch is sent.
+ ///
+ /// Raises:
+ /// TypeError: If `partition_id` is not an integer.
+ /// OverflowError: If `partition_id` is outside the supported unsigned
+ /// 32-bit range.
+ #[staticmethod]
+ pub fn partition_id(partition_id: u32) -> Self {
+ Self {
+ inner: RustPartitioning::partition_id(partition_id),
+ }
+ }
+
+ /// Routes the batch to one partition selected by hashing `key`.
+ ///
+ /// `key` may be `str` or `bytes`. Strings are encoded as UTF-8; the
encoded
+ /// key must contain between 1 and 255 bytes.
+ ///
+ /// Raises:
+ /// ValueError: If the encoded key is empty or exceeds 255 bytes.
+ /// TypeError: If `key` is not `str` or `bytes`.
+ #[staticmethod]
+ pub fn messages_key(py: Python<'_>, key: PyMessagesKey) -> PyResult<Self> {
+ let key = match key {
+ PyMessagesKey::String(key) => key.into_bytes(),
+ PyMessagesKey::Bytes(key) => key.extract::<Vec<u8>>(py)?,
+ };
+ let inner = RustPartitioning::messages_key(&key)
+ .map_err(|error| PyValueError::new_err(error.to_string()))?;
+ Ok(Self { inner })
+ }
+}
+
+#[derive(FromPyObject)]
+pub enum PyMessagesKey {
+ #[pyo3(transparent, annotation = "str")]
+ String(String),
+ #[pyo3(transparent, annotation = "bytes")]
+ Bytes(Py<PyBytes>),
+}
+impl_stub_type!(PyMessagesKey = String | PyBytes);
+
+#[derive(FromPyObject)]
+pub(crate) enum PyPartitioning {
+ #[pyo3(transparent, annotation = "Partitioning")]
+ Strategy(Partitioning),
+ #[pyo3(transparent, annotation = "int")]
+ PartitionId(u32),
+}
+impl_stub_type!(PyPartitioning = Partitioning | isize);
+
+impl From<PyPartitioning> for RustPartitioning {
+ fn from(partitioning: PyPartitioning) -> Self {
+ match partitioning {
+ PyPartitioning::Strategy(partitioning) => partitioning.inner,
+ PyPartitioning::PartitionId(partition_id) =>
Self::partition_id(partition_id),
+ }
+ }
+}
diff --git a/foreign/python/tests/test_message_operations.py
b/foreign/python/tests/test_message_operations.py
index bbf5b1d0e..af5924c4d 100644
--- a/foreign/python/tests/test_message_operations.py
+++ b/foreign/python/tests/test_message_operations.py
@@ -25,12 +25,54 @@ from apache_iggy import (
HeaderKey,
HeaderValue,
IggyClient,
+ Partitioning,
PollingStrategy,
UserHeaders,
)
from apache_iggy import SendMessage as Message
+class TestPartitioning:
+ """Test message partitioning strategy construction."""
+
+ @pytest.mark.unit
+ def test_balanced_and_partition_id_strategies(self):
+ assert isinstance(Partitioning.balanced(), Partitioning)
+ assert isinstance(Partitioning.partition_id(1), Partitioning)
+ assert isinstance(Partitioning.partition_id(2**32 - 1), Partitioning)
+
+ @pytest.mark.unit
+ @pytest.mark.parametrize("key", [b"customer-42", "customer-42", "客户-42"])
+ def test_messages_key_accepts_bytes_and_strings(self, key):
+ assert isinstance(Partitioning.messages_key(key), Partitioning)
+
+ @pytest.mark.unit
+ @pytest.mark.parametrize("key", [b"a" * 255, "a" * 255, "界" * 85])
+ def test_messages_key_accepts_255_bytes(self, key):
+ assert isinstance(Partitioning.messages_key(key), Partitioning)
+
+ @pytest.mark.unit
+ @pytest.mark.parametrize("key", [b"", "", b"a" * 256, "a" * 256, "界" * 86])
+ def test_messages_key_rejects_invalid_encoded_length(self, key):
+ with pytest.raises(ValueError):
+ Partitioning.messages_key(key)
+
+ @pytest.mark.unit
+ @pytest.mark.parametrize("partition_id", [-1, 2**32])
+ def test_partition_id_rejects_values_outside_u32(self, partition_id):
+ with pytest.raises(OverflowError):
+ Partitioning.partition_id(partition_id)
+
+ @pytest.mark.unit
+ def test_partitioning_rejects_invalid_types(self):
+ with pytest.raises(TypeError):
+ # pyrefly: ignore # bad-argument-type
+ Partitioning.partition_id("0")
+ with pytest.raises(TypeError):
+ # pyrefly: ignore # bad-argument-type
+ Partitioning.messages_key(1)
+
+
class TestMessageOperations:
"""Test message sending, polling, and processing."""
@@ -112,6 +154,101 @@ class TestMessageOperations:
assert confirmation.partition_id == partition_id
assert confirmation.base_offset == polled_messages[0].offset()
+ @pytest.mark.asyncio
+ async def test_send_messages_with_partition_id_strategy(
+ self, iggy_client: IggyClient, unique_name
+ ):
+ stream_name = unique_name()
+ topic_name = unique_name()
+ partition_id = 2
+
+ await iggy_client.create_stream(stream_name)
+ await iggy_client.create_topic(
+ stream=stream_name, name=topic_name, partitions_count=3
+ )
+
+ response = await iggy_client.send_messages(
+ stream=stream_name,
+ topic=topic_name,
+ partitioning=Partitioning.partition_id(partition_id),
+ messages=[Message("fixed partition 1"), Message("fixed partition
2")],
+ )
+
+ assert len(response.confirmations) == 1
+ assert response.confirmations[0].partition_id == partition_id
+
+ with pytest.raises(RuntimeError):
+ await iggy_client.send_messages(
+ stream=stream_name,
+ topic=topic_name,
+ partitioning=Partitioning.partition_id(3),
+ messages=[Message("missing partition")],
+ )
+
+ @pytest.mark.asyncio
+ async def test_send_messages_with_balanced_strategy(
+ self, iggy_client: IggyClient, unique_name
+ ):
+ stream_name = unique_name()
+ topic_name = unique_name()
+ partitions_count = 3
+
+ await iggy_client.create_stream(stream_name)
+ await iggy_client.create_topic(
+ stream=stream_name,
+ name=topic_name,
+ partitions_count=partitions_count,
+ )
+
+ responses = [
+ await iggy_client.send_messages(
+ stream=stream_name,
+ topic=topic_name,
+ partitioning=Partitioning.balanced(),
+ messages=[Message(f"balanced {index}")],
+ )
+ for index in range(partitions_count)
+ ]
+
+ assert all(len(response.confirmations) == 1 for response in responses)
+ assert (
+ len({response.confirmations[0].partition_id for response in
responses})
+ == partitions_count
+ )
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("key", [b"customer-42", "customer-42"])
+ async def test_send_messages_with_same_key_uses_same_partition(
+ self, iggy_client: IggyClient, unique_name, key
+ ):
+ stream_name = unique_name()
+ topic_name = unique_name()
+
+ await iggy_client.create_stream(stream_name)
+ await iggy_client.create_topic(
+ stream=stream_name, name=topic_name, partitions_count=3
+ )
+ partitioning = Partitioning.messages_key(key)
+
+ first = await iggy_client.send_messages(
+ stream=stream_name,
+ topic=topic_name,
+ partitioning=partitioning,
+ messages=[Message("first")],
+ )
+ second = await iggy_client.send_messages(
+ stream=stream_name,
+ topic=topic_name,
+ partitioning=partitioning,
+ messages=[Message("second")],
+ )
+
+ assert len(first.confirmations) == 1
+ assert len(second.confirmations) == 1
+ assert (
+ first.confirmations[0].partition_id ==
second.confirmations[0].partition_id
+ )
+
@pytest.mark.asyncio
async def test_send_and_poll_messages_as_bytes(
self, iggy_client: IggyClient, unique_name
