[OPIK-8253] [SDK] [DOCS] feat: parallel chunked dataset item reads

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

@@ -415,6 +415,63 @@ You can download a dataset from Opik using the `get_dataset` method:
   ```
 </CodeBlocks>
 
+### 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:
+
+```python title="Python" language="python"
+from opik import Opik
+
+client = Opik()
+dataset = client.get_dataset(name="My dataset")
+
+# The whole dataset as one list, downloaded over 8 threads
+items = dataset.get_items(num_threads=8)
+```
+
+The thread count never changes the result, only how quickly it arrives.
+
+`get_items()` returns the whole dataset as a single list, so the call does not return until every
+item has been downloaded and the full result is held in memory. `stream_items()` reads the same
+items in chunks instead, yielding each chunk as soon as it arrives. Use it when you want to start
+processing before the download finishes, or when the dataset is too large to hold in memory all at
+once:
+
+```python title="Python" language="python"
+from opik import Opik
+
+client = Opik()
+dataset = client.get_dataset(name="My dataset")
+
+# One chunk at a time, instead of the whole dataset at once
+for chunk in dataset.stream_items(chunk_size=5000, num_threads=8):
+    process(chunk)  # a list of dicts, exactly as get_items() returns them
+```
+
+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 read only the first N items:
+
+```python title="Python" language="python"
+for chunk in dataset.stream_items(
+    num_threads=8,
+    filter_string='data.category = "geography"',
+    nb_samples=10_000,
+):
+    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. 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
 
 You can filter dataset items using the `filter_string` parameter on the `get_items()` method or when

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

@@ -21,6 +21,19 @@
 
 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.
+DATASET_ITEMS_READ_MAX_THREADS = 32
+
 # Parallel dataset insert requires a backend that serializes concurrent dataset
 # version writes. On backends older than this version, concurrent batches
 # sharing one batch_group_id raced and could 500 or silently drop rows; 2.2.8 is

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

@@ -73,6 +73,28 @@ def __internal_api__stream_items_as_dataclasses__(
         """
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def __internal_api__stream_item_chunks__(
+        self,
+        chunk_size: int,
+        num_threads: int,
+        nb_samples: Optional[int],
+        filter_string: Optional[str],
+    ) -> Iterator[List[Dict[str, Any]]]:
+        """
+        Stream dataset items as chunks of raw dictionaries.
+
+        Args:
+            chunk_size: Number of items per chunk.
+            num_threads: Number of chunks fetched concurrently.
+            nb_samples: Maximum number of items to retrieve.
+            filter_string: Optional OQL filter string to filter dataset items.
+
+        Yields:
+            Lists of dictionaries representing the dataset items.
+        """
+        raise NotImplementedError
+
     def to_pandas(self) -> "pd.DataFrame":
         """
         Convert the dataset items to a pandas DataFrame.
@@ -99,12 +121,27 @@ def get_items(
         self,
         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.
 
         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
+                more load on the backend. Capped at
+                ``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.
 
@@ -122,14 +159,116 @@ def get_items(
 
         Returns:
             A list of dictionaries representing the dataset items.
+
+        Raises:
+            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.
         """
-        dataset_items_as_dicts = [
-            {"id": item.id, **item.get_content()}
-            for item in self.__internal_api__stream_items_as_dataclasses__(
-                nb_samples=nb_samples, filter_string=filter_string
+        return [
+            item
+            for chunk in self.stream_items(
+                chunk_size=chunk_size,
+                filter_string=filter_string,
+                nb_samples=nb_samples,
+                num_threads=num_threads,
             )
+            for item in chunk
         ]
-        return dataset_items_as_dicts
+
+    def stream_items(
+        self,
+        chunk_size: int = constants.DATASET_STREAM_BATCH_SIZE,
+        num_threads: int = constants.DATASET_ITEMS_READ_NUM_THREADS,
+        filter_string: Optional[str] = None,
+        nb_samples: Optional[int] = None,
+    ) -> Iterator[List[Dict[str, Any]]]:
+        """
+        Read dataset items in chunks, fetching the chunks concurrently.
+
+        The chunked counterpart to :meth:`get_items`, which is itself built on
+        this method: chunks are fetched in parallel and are handed back as
+        plain dictionaries without going through the typed REST layer. Prefer
+        it over :meth:`get_items` when you want to start processing before the
+        whole dataset has been downloaded, or when the dataset is too large to
+        hold in memory all at once.
+
+        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
+                (``constants.DATASET_STREAM_BATCH_SIZE``). Fetching a chunk
+                costs a fixed overhead whatever its size, so lowering this
+                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. 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 ``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):
+            ...     process(chunk)
+
+        Note:
+            ``nb_samples`` items are read starting from the beginning of the
+            dataset, so the same call reads the same items whatever the thread
+            count.
+        """
+        if isinstance(chunk_size, bool) or not isinstance(chunk_size, int):
+            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:
+            raise ValueError("num_threads must be a positive integer")
+        if nb_samples is not None and (
+            isinstance(nb_samples, bool)
+            or not isinstance(nb_samples, int)
+            or nb_samples < 1
+        ):
+            raise ValueError("nb_samples must be a positive integer")
+
+        return self.__internal_api__stream_item_chunks__(
+            chunk_size=chunk_size,
+            num_threads=min(num_threads, constants.DATASET_ITEMS_READ_MAX_THREADS),
+            nb_samples=nb_samples,
+            filter_string=filter_string,
+        )
 
     @abc.abstractmethod
     def get_version_info(
@@ -283,6 +422,24 @@ def __internal_api__stream_items_as_dataclasses__(
             dataset_version=self._version_info.version_hash,
         )
 
+    @override
+    def __internal_api__stream_item_chunks__(
+        self,
+        chunk_size: int,
+        num_threads: int,
+        nb_samples: Optional[int],
+        filter_string: Optional[str],
+    ) -> Iterator[List[Dict[str, Any]]]:
+        return rest_operations.stream_dataset_item_chunks(
+            rest_client=self._rest_client,
+            dataset_id=self._dataset_id,
+            chunk_size=chunk_size,
+            num_threads=num_threads,
+            nb_samples=nb_samples,
+            filter_string=filter_string,
+            dataset_version=self._version_info.version_hash,
+        )
+
     @override
     def get_version_info(
         self,
@@ -391,6 +548,11 @@ def from_public(
         # Backend may already hold items we haven't seen; lazy-sync on first
         # insert so content-hash dedup still works without paying a sync now.
         dataset_.__internal_api__hashes_synced__ = False
+        # The response already carries the id, so seed the cached_property
+        # rather than paying a get-dataset-by-name round trip the first time
+        # something (a read, an item delete) needs it.
+        if dataset_fern.id is not None:
+            dataset_.__dict__["id"] = dataset_fern.id
         return dataset_
 
     @functools.cached_property
@@ -973,6 +1135,55 @@ def __internal_api__stream_items_as_dataclasses__(
             dataset_version=None,
         )
 
+    @override
+    def __internal_api__stream_item_chunks__(
+        self,
+        chunk_size: int,
+        num_threads: int,
+        nb_samples: Optional[int],
+        filter_string: Optional[str],
+    ) -> Iterator[List[Dict[str, Any]]]:
+        # 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,
+            num_threads=num_threads,
+            nb_samples=nb_samples,
+            filter_string=filter_string,
+            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,