refactor cutting new batch

Author: vgvoleg

ydb/_topic_reader/datatypes.py

@@ -179,3 +179,19 @@ def _extend(self, batch: PublicBatch) -> None:
     def _pop(self) -> Tuple[List[PublicMessage], bool]:
         msgs_left = True if len(self.messages) > 1 else False
         return self.messages.pop(0), msgs_left
+
+    def _pop_batch(self, size: int) -> PublicBatch:
+        initial_length = len(self.messages)
+        one_message_size = self._bytes_size // initial_length
+
+        new_batch = PublicBatch(
+            messages=self.messages[:size],
+            _partition_session=self._partition_session,
+            _bytes_size=one_message_size * size,
+            _codec=self._codec,
+        )
+
+        self.messages = self.messages[size:]
+        self._bytes_size = one_message_size * (initial_length - size)
+
+        return new_batch

ydb/_topic_reader/topic_reader_asyncio.py

@@ -421,9 +421,9 @@ def receive_batch_nowait(self, max_messages: Optional[int] = None):
             self._buffer_release_bytes(batch._bytes_size)
             return batch
 
-        cutted_batch, remaining_batch = self._cut_batch_by_max_messages(batch, max_messages)
+        cutted_batch = batch._pop_batch(size=max_messages)
 
-        self._message_batches[part_sess_id] = remaining_batch
+        self._message_batches[part_sess_id] = batch
         self._buffer_release_bytes(cutted_batch._bytes_size)
 
         return cutted_batch