Topic multi-partition writer (write-by-key)

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Add a topic multi-partition writer (`topic_client.multiwriter(...)`) that routes messages across partitions by their `key`, with Kafka-hash and key-range partition choosers, and transparently resends in-flight messages to child partitions on an auto-partition split (no loss, no duplicates); expose partition `key_range` on `describe_topic` results
+
 ## 3.30.0 ##
 * Query session attach stream handles `NodeShutdown` and `SessionShutdown` session hints: on `NodeShutdown` the session's node connection is pessimized and the session is retired, on `SessionShutdown` the session is retired without touching the node
 * Honor the `max_bytes` parameter in topic reader `receive_batch`/`receive_batch_with_tx` (previously accepted but silently ignored); add `max_bytes` to the async reader as well

docs/topic.rst

@@ -259,6 +259,66 @@ For high-throughput pipelines, buffer writes and gather futures:
             raise f.exception()
 
 
+Writing by Key (Multiple Partitions)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+A regular writer targets a single partition for its whole lifetime. To spread load across
+all partitions of a topic from one logical writer — while keeping every message with the same
+key on the same partition (so per-key ordering is preserved) — use ``topic_client.multiwriter()``.
+
+Each message carries a ``key``; the writer hashes it and routes the message to the owning
+partition, maintaining a separate underlying writer per partition. This is the client-side
+companion to auto-partitioning.
+
+**Synchronous:**
+
+.. code-block:: python
+
+    with driver.topic_client.multiwriter("/local/my-topic") as writer:
+        writer.write(ydb.TopicWriterMessage(data="event", key="user-42"))
+        writer.write(ydb.TopicWriterMessage(data="event", key="user-7"))
+
+**Asynchronous:**
+
+.. code-block:: python
+
+    async with driver.topic_client.multiwriter("/local/my-topic") as writer:
+        await writer.write(ydb.TopicWriterMessage(data="event", key="user-42"))
+
+**Partition choosers** decide how a key maps to a partition. By default the writer picks one
+automatically after describing the topic — the key-range chooser for auto-partitioned topics,
+the Kafka-hash chooser otherwise — so ``multiwriter(topic)`` works out of the box. You can also
+set one explicitly:
+
+* :class:`~ydb.TopicWriterPartitionByKeyBound` — hashes the key and selects the partition whose
+  server-side key range owns it. Mirrors YDB auto-partitioning, so a key lands where the server
+  expects it. Used automatically for auto-partitioned topics.
+* :class:`~ydb.TopicWriterPartitionByKeyKafka` — ``murmur2(key) % partitions_count``,
+  Kafka-compatible. Best for topics with a fixed partition count.
+
+.. code-block:: python
+
+    writer = driver.topic_client.multiwriter(
+        "/local/my-topic",
+        partition_chooser=ydb.TopicWriterPartitionByKeyBound(),
+        producer_id_prefix="my-app",  # each partition writer uses "<prefix>-<partition_id>"
+    )
+
+The multi-writer accepts the same ``codec``, ``encoders``, ``auto_seqno``, ``auto_created_at``
+and buffer-limit parameters as :meth:`writer`, and exposes ``write``, ``write_with_ack``,
+``flush`` and ``close`` with the same semantics. ``wait_init()`` differs: it waits until the
+topic has been described and the partition set is known, and (unlike the single-partition writer)
+returns nothing, because the multi-writer manages a stream per partition rather than one stream.
+
+.. note::
+
+   When an auto-partitioned partition is split (one into two) or merged (two into one), the
+   multi-writer re-describes the topic, routes subsequent keys to the new child partition(s), and
+   transparently resends the messages that were still in flight to the retired partition(s).
+   Messages already persisted before the change are not resent (so no duplicates), and per-key
+   ordering is preserved.
+
+
 Writer Backpressure
 ^^^^^^^^^^^^^^^^^^^
 

examples/topic/multiwriter_example.py

@@ -0,0 +1,59 @@
+"""Examples for the topic multi-partition writer (write-by-key).
+
+The multi-writer routes each message to a partition by its ``key``, keeping every
+message with the same key on the same partition (per-key ordering). By default it
+picks a partition chooser automatically: the key-range chooser for
+auto-partitioned topics, the Kafka-hash chooser otherwise.
+"""
+
+import asyncio
+
+import ydb
+
+
+def write_by_key_sync(db: ydb.Driver, topic_path: str):
+    with db.topic_client.multiwriter(topic_path, producer_id_prefix="my-app") as writer:
+        writer.write(ydb.TopicWriterMessage(data="event-a", key="user-42"))
+        writer.write(ydb.TopicWriterMessage(data="event-b", key="user-7"))
+        # messages with the same key always land on the same partition, in order
+        writer.write(ydb.TopicWriterMessage(data="event-c", key="user-42"))
+        writer.flush()
+
+
+async def write_by_key_async(db: ydb.aio.Driver, topic_path: str):
+    async with db.topic_client.multiwriter(topic_path, producer_id_prefix="my-app") as writer:
+        await writer.write(ydb.TopicWriterMessage(data="event-a", key="user-42"))
+        # wait for the server ack of a specific message
+        await writer.write_with_ack(ydb.TopicWriterMessage(data="event-b", key="user-7"))
+
+
+def write_by_key_with_explicit_chooser(db: ydb.Driver, topic_path: str):
+    # Force the Kafka-compatible hash chooser (murmur2(key) % partitions_count).
+    with db.topic_client.multiwriter(
+        topic_path,
+        partition_chooser=ydb.TopicWriterPartitionByKeyKafka(),
+    ) as writer:
+        writer.write(ydb.TopicWriterMessage(data="event", key="user-42"))
+
+
+def run_sync():
+    with ydb.Driver(
+        connection_string="grpc://localhost:2135?database=/local",
+        credentials=ydb.credentials.AnonymousCredentials(),
+    ) as db:
+        db.wait(timeout=5, fail_fast=True)
+        write_by_key_sync(db, "/local/topic")
+
+
+async def run_async():
+    async with ydb.aio.Driver(
+        connection_string="grpc://localhost:2135?database=/local",
+        credentials=ydb.credentials.AnonymousCredentials(),
+    ) as db:
+        await db.wait(timeout=5, fail_fast=True)
+        await write_by_key_async(db, "/local/topic")
+
+
+if __name__ == "__main__":
+    run_sync()
+    asyncio.run(run_async())

tests/topics/test_topic_writer.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import datetime
 from typing import List  # noqa: F401
 
 import pytest
@@ -324,3 +325,182 @@ class TestException(Exception):
                     writer.write_with_ack("123")
 
                 raise TestException()
+
+
+@pytest.mark.asyncio
+class TestTopicMultiWriterAsyncIO:
+    async def _recreate(self, driver, path, consumer, **kwargs):
+        try:
+            await driver.topic_client.drop_topic(path)
+        except ydb.SchemeError:
+            pass
+        await driver.topic_client.create_topic(path=path, consumers=[consumer], **kwargs)
+
+    def _auto_partitioning(self):
+        return ydb.TopicAutoPartitioningSettings(
+            strategy=ydb.TopicAutoPartitioningStrategy.SCALE_UP,
+            up_utilization_percent=1,
+            down_utilization_percent=1,
+            stabilization_window=datetime.timedelta(seconds=1),
+        )
+
+    async def test_key_range_exposed_for_autopartitioned_topic(self, driver, database, topic_consumer):
+        path = database + "/mw-keyrange"
+        await self._recreate(
+            driver,
+            path,
+            topic_consumer,
+            min_active_partitions=2,
+            max_active_partitions=50,
+            auto_partitioning_settings=self._auto_partitioning(),
+        )
+        desc = await driver.topic_client.describe_topic(path)
+        assert any(p.key_range is not None for p in desc.partitions)
+
+    async def test_write_by_key_preserves_per_key_order(self, driver, database, topic_consumer):
+        path = database + "/mw-plain"
+        await self._recreate(driver, path, topic_consumer, min_active_partitions=3)
+
+        keys = ["user-1", "user-2", "user-3", "user-4", "user-5"]
+        per_key = 8
+        async with driver.topic_client.multiwriter(path, producer_id_prefix="mw") as writer:
+            await writer.wait_init()
+            assert isinstance(writer._chooser, ydb.TopicWriterPartitionByKeyKafka)
+            for i in range(per_key):
+                for key in keys:
+                    await writer.write(ydb.TopicWriterMessage(data=("%s:%d" % (key, i)).encode(), key=key))
+            await writer.flush()
+
+        total = per_key * len(keys)
+        received = {key: [] for key in keys}
+        async with driver.topic_client.reader(path, consumer=topic_consumer) as reader:
+            for _ in range(total):
+                message = await asyncio.wait_for(reader.receive_message(), timeout=30)
+                key, index = message.data.decode().split(":")
+                received[key].append(int(index))
+                reader.commit(message)
+
+        for key in keys:
+            assert received[key] == list(range(per_key)), (key, received[key])
+
+    async def test_write_by_key_on_autopartitioned_topic(self, driver, database, topic_consumer):
+        path = database + "/mw-auto"
+        await self._recreate(
+            driver,
+            path,
+            topic_consumer,
+            min_active_partitions=2,
+            max_active_partitions=50,
+            auto_partitioning_settings=self._auto_partitioning(),
+        )
+
+        keys = ["alpha", "beta", "gamma", "delta"]
+        per_key = 5
+        async with driver.topic_client.multiwriter(path, producer_id_prefix="mw") as writer:
+            await writer.wait_init()
+            # auto-partitioned topics report key ranges -> adaptive default picks the bound chooser
+            assert isinstance(writer._chooser, ydb.TopicWriterPartitionByKeyBound)
+            for i in range(per_key):
+                for key in keys:
+                    # write_with_ack verifies the server accepts bound-routed writes
+                    await writer.write_with_ack(ydb.TopicWriterMessage(data=("%s:%d" % (key, i)).encode(), key=key))
+
+        total = per_key * len(keys)
+        seen = 0
+        async with driver.topic_client.reader(path, consumer=topic_consumer) as reader:
+            for _ in range(total):
+                message = await asyncio.wait_for(reader.receive_message(), timeout=30)
+                seen += 1
+                reader.commit(message)
+        assert seen == total
+
+    async def test_write_by_key_survives_partition_split(self, driver, database, topic_consumer):
+        # Aggressive auto-partitioning + a low write-speed limit force the topic to split
+        # under load, exercising the resend path. The test keeps writing until a split is
+        # observed, then asserts exactly-once delivery (no loss, no duplicates).
+        path = database + "/mw-split"
+        await self._recreate(
+            driver,
+            path,
+            topic_consumer,
+            min_active_partitions=1,
+            max_active_partitions=100,
+            partition_write_speed_bytes_per_second=1024,
+            auto_partitioning_settings=self._auto_partitioning(),
+        )
+
+        partitions_before = len((await driver.topic_client.describe_topic(path)).partitions)
+        payload = b"x" * 512
+        written = 0
+        max_batches = 15
+        async with driver.topic_client.multiwriter(path, producer_id_prefix="mw-split") as writer:
+            for _ in range(max_batches):
+                for _ in range(100):
+                    await writer.write(
+                        ydb.TopicWriterMessage(data=b"%d:%s" % (written, payload), key="k%d" % (written % 32))
+                    )
+                    written += 1
+                await writer.flush()
+                # give the auto-partitioning actuator time to measure and split
+                await asyncio.sleep(1.5)
+                if len((await driver.topic_client.describe_topic(path)).partitions) > partitions_before:
+                    break
+            await writer.flush()
+
+        total = written
+        partitions_after = len((await driver.topic_client.describe_topic(path)).partitions)
+        if partitions_after <= partitions_before:
+            # Auto-partitioning did not split the topic in this environment (single-node
+            # clusters often don't actuate). The split/resend path is covered deterministically
+            # by the unit tests; here we only assert exactly-once when a split actually happened.
+            pytest.skip("topic did not split under load; resend path covered by unit tests")
+
+        seen = set()
+        duplicates = 0
+        async with driver.topic_client.reader(path, consumer=topic_consumer) as reader:
+            while len(seen) < total:
+                try:
+                    message = await asyncio.wait_for(reader.receive_message(), timeout=30)
+                except asyncio.TimeoutError:
+                    break
+                index = int(message.data.split(b":", 1)[0])
+                if index in seen:
+                    duplicates += 1
+                seen.add(index)
+                reader.commit(message)
+
+        assert duplicates == 0, "resend produced duplicate messages"
+        assert seen == set(range(total)), "some messages were lost (partitions %d->%d)" % (
+            partitions_before,
+            partitions_after,
+        )
+
+
+class TestTopicMultiWriterSync:
+    def test_write_by_key_preserves_per_key_order(self, driver_sync, database, topic_consumer):
+        path = database + "/mw-sync"
+        try:
+            driver_sync.topic_client.drop_topic(path)
+        except ydb.SchemeError:
+            pass
+        driver_sync.topic_client.create_topic(path=path, consumers=[topic_consumer], min_active_partitions=3)
+
+        keys = ["a", "b", "c"]
+        per_key = 6
+        with driver_sync.topic_client.multiwriter(path, producer_id_prefix="mw-sync") as writer:
+            for i in range(per_key):
+                for key in keys:
+                    writer.write(ydb.TopicWriterMessage(data=("%s:%d" % (key, i)).encode(), key=key))
+            writer.flush()
+
+        total = per_key * len(keys)
+        received = {key: [] for key in keys}
+        with driver_sync.topic_client.reader(path, consumer=topic_consumer) as reader:
+            for _ in range(total):
+                message = reader.receive_message(timeout=30)
+                key, index = message.data.decode().split(":")
+                received[key].append(int(index))
+                reader.commit(message)
+
+        for key in keys:
+            assert received[key] == list(range(per_key)), (key, received[key])