From fd895edac5faf9e3a45024ae220957eb8de24e5c Mon Sep 17 00:00:00 2001 From: Ye Cao Date: Fri, 15 Nov 2024 10:49:50 +0800 Subject: [PATCH] Support to read chunk id in stream object. Signed-off-by: Ye Cao --- python/vineyard/io/stream.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/vineyard/io/stream.py b/python/vineyard/io/stream.py index 047cbcc9e..530453f5f 100644 --- a/python/vineyard/io/stream.py +++ b/python/vineyard/io/stream.py @@ -264,6 +264,12 @@ def __init__(self, client, stream: ObjectID, resolver=None): self._resolver = resolver self._client.open_stream(stream, 'r') + def next_chunk_id(self) -> ObjectID: + try: + return self._client.next_chunk_id(self._stream) + except StreamDrainedException as e: + raise StopIteration('No more chunks') from e + def next(self) -> object: try: chunk = self._client.next_chunk(self._stream)