Honor max_bytes in topic reader receive_batch

The parameter was accepted by the sync reader but silently ignored, and was
missing from the async reader. Cap the returned batch by both max_messages and
max_bytes (approximate, from the batch's aggregate size), always returning at
least one message. Threaded through the sync and async readers.

Author: vgvoleg

CHANGELOG.md

@@ -1,3 +1,4 @@
+* 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
 * Support the timezone-carrying types `TzDate`, `TzDatetime` and `TzTimestamp` in query results and parameters (reading them previously raised `AttributeError`)
 * Drop support for Python 3.8 and 3.9; the minimum supported version is now Python 3.10
 * Accept native `datetime.datetime` values for `Datetime` and `Datetime64` query parameters (previously only integer seconds since the epoch were accepted)

docs/topic.rst

@@ -456,8 +456,8 @@ Receiving Messages
         process(message)
     reader.commit(batch)   # commit the whole batch at once
 
-    # With size limits
-    batch = reader.receive_batch(max_messages=100)
+    # With size limits (either or both; at least one message is always returned)
+    batch = reader.receive_batch(max_messages=100, max_bytes=1024 * 1024)
 
     # Asynchronous
     batch = await reader.receive_batch()

ydb/_topic_reader/datatypes.py

@@ -197,13 +197,34 @@ def _pop(self) -> Tuple[PublicMessage, bool]:
         msgs_left = True if len(self.messages) > 1 else False
         return self.messages.pop(0), msgs_left
 
-    def _pop_batch(self, message_count: int) -> PublicBatch:
+    def _pop_batch(self, max_messages: Optional[int] = None, max_bytes: Optional[int] = None) -> PublicBatch:
+        """Split off and return a prefix of the batch, capped by max_messages and/or
+        max_bytes. The remainder stays in self (empty if the whole batch was taken).
+        At least one message is always taken so a caller whose max_bytes is smaller
+        than a single message still makes progress."""
         initial_length = len(self.messages)
+        one_message_size = self._bytes_size // initial_length
 
-        if message_count >= initial_length:
-            raise ValueError("Pop batch with size >= actual size is not supported.")
+        message_count = initial_length
+        if max_messages is not None:
+            message_count = min(message_count, max_messages)
+        if max_bytes is not None:
+            # Only an aggregate byte size is known, so the per-message size (and
+            # hence the cut) is approximate.
+            message_count = min(message_count, max(1, max_bytes // max(1, one_message_size)))
 
-        one_message_size = self._bytes_size // initial_length
+        if message_count >= initial_length:
+            # Take the whole batch, keeping its exact byte size (the proportional
+            # split below would drop the integer-division remainder).
+            new_batch = PublicBatch(
+                messages=self.messages,
+                _partition_session=self._partition_session,
+                _bytes_size=self._bytes_size,
+                _codec=self._codec,
+            )
+            self.messages = []
+            self._bytes_size = 0
+            return new_batch
 
         new_batch = PublicBatch(
             messages=self.messages[:message_count],

ydb/_topic_reader/topic_reader_asyncio.py

@@ -117,35 +117,47 @@ async def wait_message(self):
     async def receive_batch(
         self,
         max_messages: typing.Union[int, None] = None,
+        max_bytes: typing.Union[int, None] = None,
     ) -> typing.Union[datatypes.PublicBatch, None]:
         """
         Get one messages batch from reader.
         All messages in a batch from same partition.
 
+        The batch is capped by max_messages and/or max_bytes when set; at least
+        one message is always returned. max_bytes uses the batch's server-reported
+        size, so the cut is approximate.
+
         use asyncio.wait_for for wait with timeout.
         """
-        logger.debug("receive_batch max_messages=%s", max_messages)
+        logger.debug("receive_batch max_messages=%s max_bytes=%s", max_messages, max_bytes)
         await self._reconnector.wait_message()
         return self._reconnector.receive_batch_nowait(
             max_messages=max_messages,
+            max_bytes=max_bytes,
         )
 
     async def receive_batch_with_tx(
         self,
         tx: "BaseQueryTxContext",
         max_messages: typing.Union[int, None] = None,
+        max_bytes: typing.Union[int, None] = None,
     ) -> typing.Union[datatypes.PublicBatch, None]:
         """
         Get one messages batch with tx from reader.
         All messages in a batch from same partition.
 
+        The batch is capped by max_messages and/or max_bytes when set; at least
+        one message is always returned. max_bytes uses the batch's server-reported
+        size, so the cut is approximate.
+
         use asyncio.wait_for for wait with timeout.
         """
-        logger.debug("receive_batch_with_tx tx=%s max_messages=%s", tx, max_messages)
+        logger.debug("receive_batch_with_tx tx=%s max_messages=%s max_bytes=%s", tx, max_messages, max_bytes)
         await self._reconnector.wait_message()
         return self._reconnector.receive_batch_with_tx_nowait(
             tx=tx,
             max_messages=max_messages,
+            max_bytes=max_bytes,
         )
 
     async def receive_message(self) -> typing.Optional[datatypes.PublicMessage]:
@@ -293,18 +305,22 @@ async def wait_message(self):
             await self._state_changed.wait()
             self._state_changed.clear()
 
-    def receive_batch_nowait(self, max_messages: Optional[int] = None):
+    def receive_batch_nowait(self, max_messages: Optional[int] = None, max_bytes: Optional[int] = None):
         if self._stream_reader is None:
             return None
         return self._stream_reader.receive_batch_nowait(
             max_messages=max_messages,
+            max_bytes=max_bytes,
         )
 
-    def receive_batch_with_tx_nowait(self, tx: "BaseQueryTxContext", max_messages: Optional[int] = None):
+    def receive_batch_with_tx_nowait(
+        self, tx: "BaseQueryTxContext", max_messages: Optional[int] = None, max_bytes: Optional[int] = None
+    ):
         if self._stream_reader is None:
             return None
         batch = self._stream_reader.receive_batch_nowait(
             max_messages=max_messages,
+            max_bytes=max_bytes,
         )
 
         self._init_tx(tx)
@@ -651,7 +667,7 @@ def _return_batch_to_queue(self, part_sess_id: int, batch: datatypes.PublicBatch
         if part_sess_id in self._partition_sessions and self._partition_sessions[part_sess_id].ended:
             self._message_batches.move_to_end(part_sess_id, last=False)
 
-    def receive_batch_nowait(self, max_messages: Optional[int] = None):
+    def receive_batch_nowait(self, max_messages: Optional[int] = None, max_bytes: Optional[int] = None):
         first_error = self._get_first_error()
         if first_error is not None:
             raise first_error
@@ -661,13 +677,10 @@ def receive_batch_nowait(self, max_messages: Optional[int] = None):
 
         part_sess_id, batch = self._get_first_batch()
 
-        if max_messages is None or len(batch.messages) <= max_messages:
-            self._buffer_release_bytes(batch._bytes_size)
-            return batch
-
-        cutted_batch = batch._pop_batch(message_count=max_messages)
+        cutted_batch = batch._pop_batch(max_messages=max_messages, max_bytes=max_bytes)
 
-        self._return_batch_to_queue(part_sess_id, batch)
+        if not batch.empty():
+            self._return_batch_to_queue(part_sess_id, batch)
 
         self._buffer_release_bytes(cutted_batch._bytes_size)
 

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -1367,6 +1367,79 @@ async def test_read_batch_max_messages(
         assert len(batch.messages) == actual_messages
         assert stream_reader._message_batches == OrderedDict(batches_after)
 
+    @pytest.mark.parametrize(
+        "max_messages,max_bytes,actual_messages",
+        [
+            (None, None, 4),  # no limits -> whole batch
+            (None, 100, 4),  # max_bytes above total -> whole batch
+            (None, 25, 2),  # batch is 40 bytes / 4 msgs = 10 each -> ~2 fit
+            (None, 10, 1),  # exactly one message
+            (None, 5, 1),  # below one message -> still return at least one
+            (3, 25, 2),  # both limits -> smaller (bytes) wins
+            (1, 100, 1),  # both limits -> smaller (messages) wins
+        ],
+    )
+    async def test_read_batch_max_bytes(
+        self,
+        stream_reader,
+        max_messages: typing.Optional[int],
+        max_bytes: typing.Optional[int],
+        actual_messages: int,
+    ):
+        stream_reader._message_batches = OrderedDict(
+            {
+                0: PublicBatch(
+                    messages=[stub_message(i) for i in range(1, 5)],
+                    _partition_session=stub_partition_session(),
+                    _bytes_size=40,
+                    _codec=Codec.CODEC_RAW,
+                )
+            }
+        )
+        batch = stream_reader.receive_batch_nowait(max_messages=max_messages, max_bytes=max_bytes)
+        assert len(batch.messages) == actual_messages
+
+    async def test_receive_whole_batch_releases_exact_bytes(self, stream_reader, monkeypatch):
+        # Taking the whole batch must release its exact byte size back to the buffer,
+        # even when the size is not divisible by the message count (no rounding loss).
+        released = []
+        monkeypatch.setattr(stream_reader, "_buffer_release_bytes", released.append)
+        stream_reader._message_batches = OrderedDict(
+            {
+                0: PublicBatch(
+                    messages=[stub_message(1), stub_message(2), stub_message(3)],
+                    _partition_session=stub_partition_session(),
+                    _bytes_size=100,  # not divisible by 3
+                    _codec=Codec.CODEC_RAW,
+                )
+            }
+        )
+        batch = stream_reader.receive_batch_nowait()
+        assert len(batch.messages) == 3
+        assert released == [100]  # full size released, not (100 // 3) * 3 == 99
+        assert stream_reader._message_batches == OrderedDict()  # nothing left buffered
+
+    async def test_receive_partial_batch_releases_only_cut_bytes(self, stream_reader, monkeypatch):
+        # A partial read releases only the taken bytes; the remainder stays buffered.
+        released = []
+        monkeypatch.setattr(stream_reader, "_buffer_release_bytes", released.append)
+        stream_reader._message_batches = OrderedDict(
+            {
+                0: PublicBatch(
+                    messages=[stub_message(1), stub_message(2), stub_message(3)],
+                    _partition_session=stub_partition_session(),
+                    _bytes_size=90,  # 30 per message
+                    _codec=Codec.CODEC_RAW,
+                )
+            }
+        )
+        batch = stream_reader.receive_batch_nowait(max_messages=1)
+        assert len(batch.messages) == 1
+        assert released == [30]  # only the cut is released
+        remainder = stream_reader._message_batches[0]
+        assert len(remainder.messages) == 2
+        assert remainder._bytes_size == 60  # remainder stays buffered
+
     async def test_receive_batch_nowait(self, stream, stream_reader, partition_session):
         assert stream_reader.receive_batch_nowait() is None
 

ydb/_topic_reader/topic_reader_sync.py

@@ -120,6 +120,7 @@ def receive_batch(
         return self._caller.safe_call_with_result(
             self._async_reader.receive_batch(
                 max_messages=max_messages,
+                max_bytes=max_bytes,
             ),
             timeout,
         )
@@ -145,6 +146,7 @@ def receive_batch_with_tx(
             self._async_reader.receive_batch_with_tx(
                 tx=tx,
                 max_messages=max_messages,
+                max_bytes=max_bytes,
             ),
             timeout,
         )