fix(dataset): pin reads to one version and stop blocking on early exit

Pages are addressed by offset and the backend sorts newest id first, so an
item inserted mid-read landed at offset 0 and shifted every unfetched page:
the read returned one item twice and skipped another. The cursor stream got
snapshot semantics for free from its `id < last_retrieved_id` seek, so this
was a regression against it, and it landed on callers that match items by
exact equality. Every page of a read now carries one version hash, resolved
once. Backends without dataset versioning fall back to a live read, which is
documented on stream_items().

The thread pool also moves out of a `with` block: the yields happen inside
its scope, so closing the generator ran shutdown(wait=True) and blocked the
caller until every outstanding page returned -- at GC time, on an unrelated
line. It now shuts down with cancel_futures=True and wait=False.

Adds get_items(chunk_size=...) so the page size is tunable from the simple
API as well.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: alexkuzmik

sdks/python/src/opik/api_objects/dataset/dataset.py

@@ -122,6 +122,7 @@ def get_items(
         nb_samples: Optional[int] = None,
         filter_string: Optional[str] = None,
         num_threads: int = constants.DATASET_ITEMS_READ_NUM_THREADS,
+        chunk_size: int = constants.DATASET_STREAM_BATCH_SIZE,
     ) -> List[Dict[str, Any]]:
         """
         Retrieve dataset items as a list of dictionaries.
@@ -137,6 +138,10 @@ def get_items(
                 ``constants.DATASET_ITEMS_READ_MAX_THREADS``. Use
                 :meth:`stream_items` instead when the dataset is too large to
                 hold in memory all at once.
+            chunk_size: Number of items fetched per request. See
+                :meth:`stream_items` for how to pick it; the whole result is
+                materialized either way, so this only trades request count
+                against per-request size.
             filter_string: Optional OQL filter string to filter dataset items.
                 Supports filtering by tags, data fields, metadata, etc.
 
@@ -156,12 +161,15 @@ def get_items(
             A list of dictionaries representing the dataset items.
 
         Raises:
-            ValueError: If ``num_threads`` is not a positive integer, or
+            ValueError: If ``num_threads`` is not a positive integer, if
+                ``chunk_size`` is not a positive integer or exceeds
+                ``constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE``, or if
                 ``nb_samples`` is not a positive integer.
         """
         return [
             item
             for chunk in self.stream_items(
+                chunk_size=chunk_size,
                 filter_string=filter_string,
                 nb_samples=nb_samples,
                 num_threads=num_threads,
@@ -189,6 +197,13 @@ def stream_items(
         Items have exactly the shape :meth:`get_items` returns: the item's
         data plus its ``id``.
 
+        The read is pinned to a single dataset version, so items inserted or
+        deleted while it is in progress do not affect it. On backends where
+        dataset versioning is unavailable there is no version to pin to and the
+        live state is read instead; a concurrent insert can then shift the
+        remaining pages, returning one item twice and skipping another. Read a
+        :class:`DatasetVersion` explicitly if you need that guarantee there.
+
         Args:
             chunk_size: Number of items per chunk, defaulting to and capped at
                 the same batch size the typed item stream reads with
@@ -1138,9 +1153,37 @@ def __internal_api__stream_item_chunks__(
             num_threads=num_threads,
             nb_samples=nb_samples,
             filter_string=filter_string,
-            dataset_version=None,
+            dataset_version=self._resolve_read_version(),
         )
 
+    def _resolve_read_version(self) -> Optional[str]:
+        """The version hash every page of one read is pinned to, if there is one.
+
+        Pages are addressed by offset, and the backend sorts newest id first, so
+        an item inserted mid-read lands at offset 0 and shifts every page that
+        has not been fetched yet -- returning one item twice and skipping
+        another. Reading a single version instead makes the whole read a
+        snapshot, which is what the cursor-based stream got for free from its
+        ``id < last_retrieved_id`` seek.
+
+        Returns None when the backend has no version to pin to (versioning
+        disabled, or a dataset with no versions yet); the read then falls back
+        to the live state and stays vulnerable to that shift, which is called
+        out on :meth:`stream_items`.
+        """
+        version_info = self.get_version_info()
+        version_hash = version_info.version_hash if version_info else None
+
+        if version_hash is None:
+            LOGGER.debug(
+                "No dataset version to pin the read of dataset %s to; reading "
+                "the live state, which may return an item twice or skip one if "
+                "items are inserted or deleted while the read is in progress.",
+                self._name,
+            )
+
+        return version_hash
+
     def insert_from_json(
         self,
         json_array: str,

sdks/python/src/opik/api_objects/dataset/parallel_items_reader.py

@@ -117,13 +117,19 @@ def _read_pages(
         worker_count,
     )
 
-    with futures.ThreadPoolExecutor(
+    # Deliberately not a `with` block: its __exit__ is shutdown(wait=True), and
+    # because the yields below happen inside it, closing the generator -- an
+    # early `break`, or garbage collection -- would block the caller until every
+    # outstanding page came back. At a 2000-item page that is a visible stall on
+    # what looks like a plain `break`, and at GC time it would surface on an
+    # unrelated line.
+    pool = futures.ThreadPoolExecutor(
         max_workers=worker_count, thread_name_prefix="opik_dataset_items_read"
-    ) as pool:
+    )
+    try:
         # Two pages in flight per worker: enough that a worker always has the
         # next page waiting, while keeping the reader's memory bounded by the
-        # look-ahead rather than by the size of the dataset. A consumer that
-        # stops iterating leaves at most this many requests to drain.
+        # look-ahead rather than by the size of the dataset.
         in_flight: Deque[futures.Future] = collections.deque(
             pool.submit(fetch_page, page)
             for page in itertools.islice(remaining_pages, worker_count * 2)
@@ -137,6 +143,11 @@ def _read_pages(
                 in_flight.append(pool.submit(fetch_page, next_page))
 
             yield _page_items(page)
+    finally:
+        # Drops the pages that haven't started and doesn't join the ones that
+        # have, so abandoning the iterator returns immediately. On the normal
+        # path there is nothing left in flight, so this is a no-op.
+        pool.shutdown(wait=False, cancel_futures=True)
 
 
 def _build_page_fetcher(

sdks/python/tests/unit/api_objects/dataset/test_stream_item_chunks.py

@@ -2,6 +2,7 @@
 
 import json
 import threading
+import time
 from typing import Any, Dict, List, Optional
 from unittest.mock import Mock, patch
 
@@ -81,17 +82,36 @@ def _rest_items(count: int) -> List[Dict[str, Any]]:
     ]
 
 
-def _build_dataset(endpoint: Optional[FakeItemsEndpoint]) -> Dataset:
+def _mock_rest_client(endpoint=None, version_hash: Optional[str] = None) -> Mock:
+    """A rest client whose dataset-id and version lookups are both controlled.
+
+    ``version_hash=None`` models a backend with no version to pin a read to,
+    which is the default so the assertions about the ``version`` query
+    parameter stay meaningful.
+    """
     mock_rest_client = Mock()
     mock_rest_client.datasets.get_dataset_by_identifier.return_value.id = DATASET_ID
+
+    versions_page = Mock()
+    versions_page.content = (
+        [Mock(version_hash=version_hash)] if version_hash is not None else []
+    )
+    mock_rest_client.datasets.list_dataset_versions.return_value = versions_page
+
     if endpoint is not None:
         mock_rest_client._client_wrapper.httpx_client.request.side_effect = endpoint
 
+    return mock_rest_client
+
+
+def _build_dataset(
+    endpoint: Optional[FakeItemsEndpoint], version_hash: Optional[str] = None
+) -> Dataset:
     return Dataset(
         name="test-dataset",
         description=None,
         project_name=None,
-        rest_client=mock_rest_client,
+        rest_client=_mock_rest_client(endpoint, version_hash),
     )
 
 
@@ -466,33 +486,25 @@ def test_stream_items__malformed_total__raises_instead_of_truncating(total):
     if total is not None:
         body["total"] = total
 
-    mock_rest_client = Mock()
-    mock_rest_client.datasets.get_dataset_by_identifier.return_value.id = DATASET_ID
-    mock_rest_client._client_wrapper.httpx_client.request.side_effect = (
-        _endpoint_returning(body)
-    )
     dataset = Dataset(
         name="test-dataset",
         description=None,
         project_name=None,
-        rest_client=mock_rest_client,
+        rest_client=_mock_rest_client(_endpoint_returning(body)),
     )
 
     with pytest.raises(exceptions.OpikException, match="Malformed response"):
         list(dataset.stream_items())
 
 
 def test_stream_items__non_list_content__raises():
-    mock_rest_client = Mock()
-    mock_rest_client.datasets.get_dataset_by_identifier.return_value.id = DATASET_ID
-    mock_rest_client._client_wrapper.httpx_client.request.side_effect = (
-        _endpoint_returning({"total": 2, "content": {"not": "a list"}})
-    )
     dataset = Dataset(
         name="test-dataset",
         description=None,
         project_name=None,
-        rest_client=mock_rest_client,
+        rest_client=_mock_rest_client(
+            _endpoint_returning({"total": 2, "content": {"not": "a list"}})
+        ),
     )
 
     with pytest.raises(exceptions.OpikException, match="Malformed response"):
@@ -501,16 +513,11 @@ def test_stream_items__non_list_content__raises():
 
 def test_stream_items__total_zero_with_empty_content__reads_nothing():
     """A genuinely empty dataset is not malformed."""
-    mock_rest_client = Mock()
-    mock_rest_client.datasets.get_dataset_by_identifier.return_value.id = DATASET_ID
-    mock_rest_client._client_wrapper.httpx_client.request.side_effect = (
-        _endpoint_returning({"total": 0, "content": []})
-    )
     dataset = Dataset(
         name="test-dataset",
         description=None,
         project_name=None,
-        rest_client=mock_rest_client,
+        rest_client=_mock_rest_client(_endpoint_returning({"total": 0, "content": []})),
     )
 
     assert list(dataset.stream_items()) == []
@@ -561,3 +568,127 @@ def test_get_items__chunk_size_cap_applies_through_get_items():
         call["params"]["size"] <= constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE
         for call in endpoint.calls
     )
+
+
+def test_stream_items__version_available__every_page_pinned_to_it():
+    """Pages are addressed by offset, so they must all read one version --
+    otherwise an insert landing at offset 0 shifts the unfetched pages."""
+    endpoint = FakeItemsEndpoint(_rest_items(25))
+    dataset = _build_dataset(endpoint, version_hash="v-hash-abc")
+
+    list(dataset.stream_items(chunk_size=10, num_threads=4))
+
+    assert len(endpoint.calls) == 3
+    assert {call["params"]["version"] for call in endpoint.calls} == {"v-hash-abc"}
+
+
+def test_stream_items__no_version_available__reads_the_live_state():
+    endpoint = FakeItemsEndpoint(_rest_items(25))
+    dataset = _build_dataset(endpoint, version_hash=None)
+
+    list(dataset.stream_items(chunk_size=10, num_threads=4))
+
+    assert {call["params"]["version"] for call in endpoint.calls} == {None}
+
+
+def test_stream_items__version_resolved_once__not_per_page():
+    endpoint = FakeItemsEndpoint(_rest_items(50))
+    dataset = _build_dataset(endpoint, version_hash="v-hash-abc")
+
+    list(dataset.stream_items(chunk_size=10, num_threads=4))
+
+    assert dataset._rest_client.datasets.list_dataset_versions.call_count == 1
+
+
+def test_stream_items__version_lookup_deferred_until_iteration():
+    endpoint = FakeItemsEndpoint(_rest_items(10))
+    dataset = _build_dataset(endpoint, version_hash="v-hash-abc")
+    versions = dataset._rest_client.datasets.list_dataset_versions
+
+    stream = dataset.stream_items()
+
+    assert versions.call_count == 0
+    next(iter(stream))
+    assert versions.call_count == 1
+
+
+def test_get_items__pins_the_read_to_a_version_too():
+    endpoint = FakeItemsEndpoint(_rest_items(25))
+    dataset = _build_dataset(endpoint, version_hash="v-hash-abc")
+
+    dataset.get_items(chunk_size=10)
+
+    assert {call["params"]["version"] for call in endpoint.calls} == {"v-hash-abc"}
+
+
+class ShiftingItemsEndpoint(FakeItemsEndpoint):
+    """Simulates an insert landing between page 1 and page 2.
+
+    Ids sort newest-first, so a new item takes offset 0 and pushes every later
+    item one slot further down -- the exact shift that makes an offset-paged
+    read of a live dataset return one item twice and skip another.
+    """
+
+    def __call__(self, path, *, method, params):
+        response = super().__call__(path, method=method, params=params)
+        if params["page"] == 1 and params.get("version") is None:
+            self._items.insert(0, {"id": "i-new", "data": {"question": "inserted"}})
+        return response
+
+
+def test_stream_items__unversioned_read_with_a_concurrent_insert__shifts():
+    """Documents the failure mode the version pin exists to prevent, so a
+    regression in the pinning shows up as this test starting to pass."""
+    endpoint = ShiftingItemsEndpoint(
+        [{"id": f"i{i}", "data": {"question": f"q{i}"}} for i in range(4)]
+    )
+    dataset = _build_dataset(endpoint, version_hash=None)
+
+    ids = [
+        item["id"]
+        for chunk in dataset.stream_items(chunk_size=2, num_threads=1)
+        for item in chunk
+    ]
+
+    # i1 is returned twice and i3 never arrives.
+    assert len(ids) != len(set(ids)), (
+        "expected the unversioned read to duplicate an item once the dataset shifted"
+    )
+
+
+def test_stream_items__version_pinned_read_is_unaffected_by_a_concurrent_insert():
+    endpoint = ShiftingItemsEndpoint(
+        [{"id": f"i{i}", "data": {"question": f"q{i}"}} for i in range(4)]
+    )
+    dataset = _build_dataset(endpoint, version_hash="v-hash-abc")
+
+    ids = [
+        item["id"]
+        for chunk in dataset.stream_items(chunk_size=2, num_threads=1)
+        for item in chunk
+    ]
+
+    assert ids == ["i0", "i1", "i2", "i3"]
+    assert len(ids) == len(set(ids))
+
+
+def test_stream_items__abandoned_early__does_not_block_on_in_flight_pages():
+    """Closing the generator must not join the outstanding requests: the yields
+    happen inside the pool's scope, so a `with` block would turn a plain
+    `break` into a wait for every page still in flight."""
+    page_delay = 0.5
+    endpoint = FakeItemsEndpoint(_rest_items(200), delay_seconds=page_delay)
+    dataset = _build_dataset(endpoint)
+
+    stream = dataset.stream_items(chunk_size=10, num_threads=4)
+    for _ in stream:
+        break  # abandon after the first chunk, with pages still in flight
+
+    started = time.perf_counter()
+    stream.close()
+    elapsed = time.perf_counter() - started
+
+    assert elapsed < page_delay, (
+        f"closing the stream blocked for {elapsed:.2f}s; it must not wait for "
+        "in-flight pages"
+    )