fix(dataset): address review findings on the chunked reader

- Validate the page response instead of defaulting `total` to zero. A
  missing, non-numeric or negative `total` capped the read at page 1 and
  returned part of the dataset as if it were all of it; it now raises.
- Cap `chunk_size` at DATASET_STREAM_BATCH_SIZE. It was forwarded straight to
  the endpoint's `size`, so oversized values could materialize an unbounded
  page and values above int32 failed only after the request went out. Rejects
  rather than clamps: quietly returning smaller pages than asked for would
  look like the argument had no effect.
- Defer the dataset-id lookup in `stream_items()` to first iteration by making
  the Dataset hook a generator. The laziness test only checked httpx and
  passed despite the eager by-name lookup; it now asserts the lookup too.
- Document the `nb_samples` valid range and its ValueError.
- Cover the `Dataset.from_public` id seeding, including the no-id fallback.

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

Author: alexkuzmik

apps/opik-documentation/documentation/fern/docs-v2/evaluation/advanced/manage_datasets.mdx

@@ -418,8 +418,7 @@ You can download a dataset from Opik using the `get_dataset` method:
 ### Downloading large datasets faster
 
 Dataset items are fetched a page at a time, and those pages are downloaded concurrently. Raise
-`num_threads` to speed up a large read — the default is 4, and 8 is usually worth trying on
-datasets of tens of thousands of items:
+`num_threads` to speed up a large read — the default is 4:
 
 ```python title="Python" language="python"
 from opik import Opik
@@ -451,7 +450,7 @@ for chunk in dataset.stream_items(chunk_size=5000, num_threads=8):
 ```
 
 Chunks arrive in dataset order; only the last one may be shorter than `chunk_size`. Both methods
-accept the same `filter_string`, and `nb_samples` to stop after a given number of items:
+accept the same `filter_string`, and `nb_samples` to read only the first N items:
 
 ```python title="Python" language="python"
 for chunk in dataset.stream_items(
@@ -462,11 +461,15 @@ for chunk in dataset.stream_items(
     process(chunk)
 ```
 
+`nb_samples` must be a positive integer — omit it or pass `None` to read everything. Passing `0`
+or a negative value raises `ValueError` rather than being treated as a limit.
+
 <Note>
-  `chunk_size` controls how many items each request fetches (default 2000). Fetching a chunk costs
-  a fixed overhead whatever its size, so lowering it makes the whole read slower — raise it to
-  trade memory for speed. Lower it only when individual items are large, since up to
-  `2 * num_threads` chunks are held in memory at once.
+  `chunk_size` controls how many items each request fetches. It defaults to 2000, which is also
+  the maximum — a larger value raises `ValueError`, so that peak memory stays bounded. Fetching a
+  chunk costs a fixed overhead whatever its size, so lowering it makes the whole read slower;
+  lower it only when individual items are large, since up to `2 * num_threads` chunks are held in
+  memory at once.
 </Note>
 
 ## Filtering datasets programmatically

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

@@ -22,6 +22,13 @@
 DATASET_STREAM_BATCH_SIZE = 2000
 
 DATASET_ITEMS_READ_NUM_THREADS = 4
+# Page-size ceiling for reads, deliberately the same as the batch size above: a
+# read should never ask the backend for a bigger page than the SDK's own read
+# batch, so peak memory stays bounded the way it was before pages were fetched
+# in parallel. Unlike the thread ceiling this one rejects rather than clamps --
+# silently handing back smaller pages than asked for would look like the
+# argument had no effect.
+DATASET_ITEMS_READ_MAX_CHUNK_SIZE = DATASET_STREAM_BATCH_SIZE
 # Ceiling on dataset read threads. The SDK's httpx client pools 100
 # connections, so a caller passing an arbitrarily large num_threads would
 # otherwise queue pages behind the pool instead of speeding anything up.

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

@@ -127,7 +127,9 @@ def get_items(
         Retrieve dataset items as a list of dictionaries.
 
         Args:
-            nb_samples: Maximum number of items to retrieve. If not set, all items are returned.
+            nb_samples: Maximum number of items to retrieve. Must be a positive
+                integer; omit it or pass ``None`` to return all items. Zero and
+                negative values raise rather than being treated as a limit.
             num_threads: Number of item pages fetched concurrently. Must be a
                 positive integer, defaults to 4; pass ``1`` to fetch
                 sequentially. Raising it speeds up large reads at the cost of
@@ -188,30 +190,34 @@ def stream_items(
         data plus its ``id``.
 
         Args:
-            chunk_size: Number of items per chunk, defaulting to the same
-                batch size the typed item stream reads with
+            chunk_size: Number of items per chunk, defaulting to and capped at
+                the same batch size the typed item stream reads with
                 (``constants.DATASET_STREAM_BATCH_SIZE``). Fetching a chunk
                 costs a fixed overhead whatever its size, so lowering this
-                makes the whole read slower; raise it to trade memory for
-                speed, bearing in mind that up to ``2 * num_threads`` chunks
-                are held at once.
+                makes the whole read slower; lower it when the items are
+                individually large, bearing in mind that up to
+                ``2 * num_threads`` chunks are held in memory at once.
             num_threads: Number of chunks fetched concurrently. Must be a
                 positive integer, defaults to 4; pass ``1`` to fetch
                 sequentially. Capped at
                 ``constants.DATASET_ITEMS_READ_MAX_THREADS``.
             filter_string: Optional OQL filter string to filter dataset items.
                 Accepts the same expressions as :meth:`get_items`.
-            nb_samples: Maximum number of items to read. If not set, the whole
-                dataset is read.
+            nb_samples: Maximum number of items to read. Must be a positive
+                integer; omit it or pass ``None`` to read the whole dataset.
+                Zero and negative values raise rather than being treated as a
+                limit.
 
         Yields:
             Lists of dictionaries representing the dataset items, in dataset
             order. The last chunk may be shorter than ``chunk_size``; empty
             chunks are never yielded.
 
         Raises:
-            ValueError: If ``chunk_size`` or ``num_threads`` is not a positive
-                integer, or ``nb_samples`` is not a positive integer.
+            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.
 
         Example:
             >>> for chunk in dataset.stream_items(chunk_size=2000, num_threads=8):
@@ -226,6 +232,11 @@ def stream_items(
             raise ValueError("chunk_size must be a positive integer")
         if chunk_size < 1:
             raise ValueError("chunk_size must be a positive integer")
+        if chunk_size > constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE:
+            raise ValueError(
+                "chunk_size must not exceed "
+                f"{constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE}, got {chunk_size}"
+            )
         if isinstance(num_threads, bool) or not isinstance(num_threads, int):
             raise ValueError("num_threads must be a positive integer")
         if num_threads < 1:
@@ -1117,7 +1128,10 @@ def __internal_api__stream_item_chunks__(
         nb_samples: Optional[int],
         filter_string: Optional[str],
     ) -> Iterator[List[Dict[str, Any]]]:
-        return rest_operations.stream_dataset_item_chunks(
+        # A generator, so `self.id` -- which resolves the dataset by name over
+        # REST when it hasn't been seeded -- is not touched until the caller
+        # actually starts iterating. Keeps `stream_items()` free of I/O.
+        yield from rest_operations.stream_dataset_item_chunks(
             rest_client=self._rest_client,
             dataset_id=self.id,
             chunk_size=chunk_size,

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

@@ -20,12 +20,19 @@
 from typing import Any, Callable, Deque, Dict, Iterator, List, Optional
 import collections
 
+import opik.exceptions as exceptions
 from opik.rest_api import client as rest_api_client
 from opik.rest_api.core import api_error as rest_api_error
 from opik.rest_client_configurator import retry_decorator
 
 LOGGER = logging.getLogger(__name__)
 
+_MALFORMED_PAGE = (
+    "Malformed response from the dataset items endpoint. "
+    "The page count for the rest of the read is derived from this response, so "
+    "continuing would silently return only part of the dataset."
+)
+
 
 def stream_item_chunks(
     rest_client: rest_api_client.OpikApi,
@@ -91,9 +98,9 @@ def _read_pages(
     how many pages there are to fan out over.
     """
     first_page = fetch_page(1)
+    total = _page_total(first_page)
     yield _page_items(first_page)
 
-    total = first_page.get("total") or 0
     wanted = total if max_items is None else min(total, max_items)
     last_page = math.ceil(wanted / chunk_size)
     if last_page <= 1:
@@ -172,5 +179,38 @@ def fetch_page(page: int) -> Dict[str, Any]:
 
 
 def _page_items(page: Dict[str, Any]) -> List[Dict[str, Any]]:
-    content: List[Dict[str, Any]] = page.get("content") or []
+    content = page.get("content")
+    if content is None:
+        content = []
+    if not isinstance(content, list):
+        raise exceptions.OpikException(
+            f"{_MALFORMED_PAGE} Expected 'content' to be a list, got "
+            f"{type(content).__name__}."
+        )
     return content
+
+
+def _page_total(page: Dict[str, Any]) -> int:
+    """Read the item count the rest of the read is planned against.
+
+    Validated rather than defaulted: the page count is derived from ``total``,
+    so a missing or non-numeric value would silently cap the read at the first
+    page and hand back part of the dataset as if it were all of it.
+    """
+    if not isinstance(page, dict):
+        raise exceptions.OpikException(
+            f"{_MALFORMED_PAGE} Expected a JSON object, got {type(page).__name__}."
+        )
+
+    total = page.get("total")
+    # bool is an int subclass, and True would silently read as 1 item.
+    if isinstance(total, bool) or not isinstance(total, int):
+        raise exceptions.OpikException(
+            f"{_MALFORMED_PAGE} Expected an integer 'total', got {total!r}."
+        )
+    if total < 0:
+        raise exceptions.OpikException(
+            f"{_MALFORMED_PAGE} Expected a non-negative 'total', got {total}."
+        )
+
+    return total

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

@@ -176,3 +176,38 @@ def test_backend_returns_none_count__property_returns_none():
 
     assert count is None
     mock_rest_client.datasets.get_dataset_by_id.assert_called_once()
+
+
+def test_from_public__response_carries_id__id_returned_without_a_lookup():
+    """The get-dataset response already holds the id, so reads must not pay a
+    second by-name lookup for it."""
+    mock_rest_client = Mock()
+    mock_dataset_public = DatasetPublic(
+        id="01a06bd3-f379-7231-a8ff-842808c8ba38",
+        name="test_dataset",
+        dataset_items_count=7,
+    )
+
+    dataset = Dataset.from_public(
+        dataset_fern=mock_dataset_public,
+        project_name="Test project",
+        rest_client=mock_rest_client,
+    )
+
+    assert dataset.id == "01a06bd3-f379-7231-a8ff-842808c8ba38"
+    mock_rest_client.datasets.get_dataset_by_identifier.assert_not_called()
+
+
+def test_from_public__response_without_id__falls_back_to_the_lookup():
+    mock_rest_client = Mock()
+    mock_rest_client.datasets.get_dataset_by_identifier.return_value.id = "looked-up"
+    mock_dataset_public = DatasetPublic(name="test_dataset")
+
+    dataset = Dataset.from_public(
+        dataset_fern=mock_dataset_public,
+        project_name="Test project",
+        rest_client=mock_rest_client,
+    )
+
+    assert dataset.id == "looked-up"
+    mock_rest_client.datasets.get_dataset_by_identifier.assert_called_once()

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

@@ -7,6 +7,7 @@
 
 import pytest
 
+import opik.exceptions as exceptions
 from opik.api_objects import constants
 from opik.api_objects.dataset import rest_operations
 from opik.api_objects.dataset.dataset import Dataset
@@ -173,14 +174,21 @@ def test_stream_items__nb_samples_above_total__reads_everything():
 
 
 def test_stream_items__lazy__no_request_until_iterated():
+    """No I/O at all before the first chunk is pulled -- including the
+    by-name dataset-id lookup, which is a REST call of its own."""
     endpoint = FakeItemsEndpoint(_rest_items(10))
     dataset = _build_dataset(endpoint)
+    lookup = dataset._rest_client.datasets.get_dataset_by_identifier
 
     stream = dataset.stream_items()
 
     assert endpoint.calls == []
+    assert lookup.call_count == 0
+
     next(iter(stream))
+
     assert endpoint.requested_pages == [1]
+    assert lookup.call_count == 1
 
 
 def test_stream_items__data_with_id_key__real_item_id_wins():
@@ -430,3 +438,126 @@ def test_get_items__positional_args_unchanged__nb_samples_still_first():
     dataset = _build_dataset(endpoint)
 
     assert len(dataset.get_items(10)) == 10
+
+
+def _endpoint_returning(body: Dict[str, Any]) -> Mock:
+    """An items endpoint that returns one fixed, possibly malformed, page."""
+
+    def call(path: str, *, method: str, params: Dict[str, Any]) -> Mock:
+        response = Mock()
+        response.status_code = 200
+        response.headers = {}
+        response.text = ""
+        response.json.return_value = body
+        return response
+
+    return call
+
+
+@pytest.mark.parametrize(
+    "total",
+    [None, "abc", 1.5, True, -1],
+    ids=["missing", "string", "float", "bool", "negative"],
+)
+def test_stream_items__malformed_total__raises_instead_of_truncating(total):
+    """A bad `total` would cap the read at page 1 and look like a short
+    dataset, so it has to fail loudly rather than silently."""
+    body: Dict[str, Any] = {"content": _rest_items(2)}
+    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,
+    )
+
+    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,
+    )
+
+    with pytest.raises(exceptions.OpikException, match="Malformed response"):
+        list(dataset.stream_items())
+
+
+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,
+    )
+
+    assert list(dataset.stream_items()) == []
+
+
+def test_stream_items__chunk_size_at_the_cap__accepted():
+    endpoint = FakeItemsEndpoint(_rest_items(10))
+    dataset = _build_dataset(endpoint)
+
+    chunks = list(
+        dataset.stream_items(chunk_size=constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE)
+    )
+
+    assert sum(len(chunk) for chunk in chunks) == 10
+
+
+@pytest.mark.parametrize(
+    "chunk_size",
+    [
+        constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE + 1,
+        10**9,
+        2**31,
+    ],
+)
+def test_stream_items__chunk_size_above_the_cap__raises_before_any_request(chunk_size):
+    """Oversized pages are rejected client-side; forwarding them would let the
+    backend materialize an unbounded response, and values beyond int32 would
+    only fail after the request went out."""
+    endpoint = FakeItemsEndpoint(_rest_items(10))
+    dataset = _build_dataset(endpoint)
+
+    with pytest.raises(ValueError, match="chunk_size must not exceed"):
+        dataset.stream_items(chunk_size=chunk_size)
+
+    assert endpoint.calls == []
+
+
+def test_get_items__chunk_size_cap_applies_through_get_items():
+    """get_items() never exceeds the cap, since it uses the default."""
+    endpoint = FakeItemsEndpoint(
+        _rest_items(constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE + 5)
+    )
+    dataset = _build_dataset(endpoint)
+
+    dataset.get_items()
+
+    assert all(
+        call["params"]["size"] <= constants.DATASET_ITEMS_READ_MAX_CHUNK_SIZE
+        for call in endpoint.calls
+    )