add pagination to examples metric

src/pyproject.toml

@@ -33,7 +33,7 @@ docs = [
 ]
 test = ["pytest", "coverage", "pre-commit"]
 benchmark = ["requests"]
-dev = ["valor-lite[nlp, openai, mistral, benchmark, test, docs]"]
+dev = ["valor-lite[nlp, openai, mistral, benchmark, test, docs]", "pyarrow-stubs"]
 
 [tool.black]
 line-length = 79

src/valor_lite/cache/compute.py

@@ -1,9 +1,10 @@
 import heapq
 import tempfile
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Generator
 
 import pyarrow as pa
+import pyarrow.compute as pc
 
 from valor_lite.cache.ephemeral import MemoryCacheReader, MemoryCacheWriter
 from valor_lite.cache.persistent import FileCacheReader, FileCacheWriter
@@ -50,7 +51,9 @@ def create_sort_key(
     batch_iterators = []
     batches = []
     for batch_idx, batch_iter in enumerate(
-        intermediate_source.iterate_fragments(batch_size=batch_size)
+        intermediate_source.iterate_fragment_batch_iterators(
+            batch_size=batch_size
+        )
     ):
         batch_iterators.append(batch_iter)
         batches.append(next(batch_iterators[batch_idx], None))
@@ -152,3 +155,57 @@ def sort(
             columns=columns,
             table_sort_override=table_sort_override,
         )
+
+
+def paginate_index(
+    source: MemoryCacheReader | FileCacheReader,
+    column_key: str,
+    modifier: pc.Expression | None = None,
+    limit: int | None = None,
+    offset: int = 0,
+) -> Generator[pa.Table, None, None]:
+    """
+    Iterate through a paginated cache reader.
+
+    Note this function expects unqiue keys to be fragment-aligned.
+    """
+    total = source.count_rows()
+    limit = limit if limit else total
+
+    # pagination broader than data scope
+    if offset == 0 and limit >= total:
+        for tbl in source.iterate_tables(filter=modifier):
+            yield tbl
+        return
+    elif offset >= total:
+        return
+
+    curr_idx = 0
+    for tbl in source.iterate_tables(filter=modifier):
+        if tbl.num_rows == 0:
+            continue
+
+        # sort the unique keys as they may be out of order
+        unique_values = pc.unique(tbl[column_key]).sort()  # type: ignore[reportAttributeAccessIssue]
+        n_unique = len(unique_values)
+        prev_idx = curr_idx
+        curr_idx += n_unique
+
+        # check for page overlap
+        if curr_idx <= offset:
+            continue
+        elif prev_idx >= (offset + limit):
+            return
+
+        # apply any pagination conditions
+        condition = pc.scalar(True)
+        if prev_idx < offset and curr_idx > offset:
+            condition &= (
+                pc.field(column_key) >= unique_values[offset - prev_idx]
+            )
+        if prev_idx < (offset + limit) and curr_idx > (offset + limit):
+            condition &= (
+                pc.field(column_key) < unique_values[offset + limit - prev_idx]
+            )
+
+        yield tbl.filter(condition)

src/valor_lite/cache/ephemeral.py

@@ -118,7 +118,7 @@ def iterate_tables_with_arrays(
                 [tbl[col].to_numpy() for col in table_columns]
             )
 
-    def iterate_fragments(
+    def iterate_fragment_batch_iterators(
         self, batch_size: int
     ) -> Iterator[Iterator[pa.RecordBatch]]:
         """

src/valor_lite/cache/persistent.py

@@ -154,6 +154,23 @@ def _retrieve(config: dict, key: str):
             compression=compression,
         )
 
+    def iterate_ordered_fragments(
+        self, filter: pc.Expression | None = None
+    ) -> Iterator[ds.Fragment]:
+        """
+        Iterate through ordered cache fragments.
+        """
+        dataset = ds.dataset(
+            source=self._path,
+            schema=self._schema,
+            format="parquet",
+        )
+        fragments = list(dataset.get_fragments(filter=filter))
+        for fragment in sorted(
+            fragments, key=lambda x: int(Path(x.path).stem)
+        ):
+            yield fragment
+
     def iterate_tables(
         self,
         columns: list[str] | None = None,
@@ -173,12 +190,7 @@ def iterate_tables(
         -------
         Iterator[pa.Table]
         """
-        dataset = ds.dataset(
-            source=self._path,
-            schema=self._schema,
-            format="parquet",
-        )
-        for fragment in dataset.get_fragments():
+        for fragment in self.iterate_ordered_fragments(filter=filter):
             yield fragment.to_table(columns=columns, filter=filter)
 
     def iterate_arrays(
@@ -243,7 +255,7 @@ def iterate_tables_with_arrays(
                 [tbl[col].to_numpy() for col in table_columns]
             )
 
-    def iterate_fragments(
+    def iterate_fragment_batch_iterators(
         self, batch_size: int
     ) -> Iterator[Iterator[pa.RecordBatch]]:
         """
@@ -258,12 +270,7 @@ def iterate_fragments(
         -------
         Iterator[Iterator[pa.RecordBatch]]
         """
-        dataset = ds.dataset(
-            source=self._path,
-            schema=self._schema,
-            format="parquet",
-        )
-        for fragment in dataset.get_fragments():
+        for fragment in self.iterate_ordered_fragments():
             yield fragment.to_batches(batch_size=batch_size)
 
 

src/valor_lite/classification/evaluator.py

@@ -528,23 +528,6 @@ def iterate_values(self, datums: pc.Expression | None = None):
             matches = tbl["match"].to_numpy()
             yield ids, scores, winners, matches
 
-    def iterate_values_with_tables(self, datums: pc.Expression | None = None):
-        for tbl in self._reader.iterate_tables(filter=datums):
-            ids = np.column_stack(
-                [
-                    tbl[col].to_numpy()
-                    for col in [
-                        "datum_id",
-                        "gt_label_id",
-                        "pd_label_id",
-                    ]
-                ]
-            )
-            scores = tbl["pd_score"].to_numpy()
-            winners = tbl["pd_winner"].to_numpy()
-            matches = tbl["match"].to_numpy()
-            yield ids, scores, winners, matches, tbl
-
     def compute_rocauc(self) -> dict[MetricType, list[Metric]]:
         """
         Compute ROCAUC.
@@ -723,6 +706,8 @@ def compute_examples(
         score_thresholds: list[float] = [0.0],
         hardmax: bool = True,
         datums: pc.Expression | None = None,
+        limit: int | None = None,
+        offset: int = 0,
     ) -> list[Metric]:
         """
         Compute examples per datum.
@@ -737,6 +722,10 @@ def compute_examples(
             Toggles whether a hardmax is applied to predictions.
         datums : pyarrow.compute.Expression, optional
             Option to filter datums by an expression.
+        limit : int, optional
+            Option to set a limit to the number of returned datum examples.
+        offset : int, default=0
+            Option to offset where examples are being created in the datum index.
 
         Returns
         -------
@@ -747,16 +736,29 @@ def compute_examples(
             raise ValueError("At least one score threshold must be passed.")
 
         metrics = []
-        for (
-            ids,
-            scores,
-            winners,
-            _,
-            tbl,
-        ) in self.iterate_values_with_tables(datums=datums):
-            if ids.size == 0:
+        for tbl in compute.paginate_index(
+            source=self._reader,
+            column_key="datum_id",
+            modifier=datums,
+            limit=limit,
+            offset=offset,
+        ):
+            if tbl.num_rows == 0:
                 continue
 
+            ids = np.column_stack(
+                [
+                    tbl[col].to_numpy()
+                    for col in [
+                        "datum_id",
+                        "gt_label_id",
+                        "pd_label_id",
+                    ]
+                ]
+            )
+            scores = tbl["pd_score"].to_numpy()
+            winners = tbl["pd_winner"].to_numpy()
+
             # extract external identifiers
             index_to_datum_id = create_mapping(
                 tbl, ids, 0, "datum_id", "datum_uid"
@@ -829,16 +831,23 @@ def compute_confusion_matrix_with_examples(
             )
             for score_idx, score_thresh in enumerate(score_thresholds)
         }
-        for (
-            ids,
-            scores,
-            winners,
-            _,
-            tbl,
-        ) in self.iterate_values_with_tables(datums=datums):
-            if ids.size == 0:
+        for tbl in self._reader.iterate_tables(filter=datums):
+            if tbl.num_rows == 0:
                 continue
 
+            ids = np.column_stack(
+                [
+                    tbl[col].to_numpy()
+                    for col in [
+                        "datum_id",
+                        "gt_label_id",
+                        "pd_label_id",
+                    ]
+                ]
+            )
+            scores = tbl["pd_score"].to_numpy()
+            winners = tbl["pd_winner"].to_numpy()
+
             # extract external identifiers
             index_to_datum_id = create_mapping(
                 tbl, ids, 0, "datum_id", "datum_uid"

src/valor_lite/object_detection/evaluator.py

@@ -600,6 +600,8 @@ def compute_examples(
         iou_thresholds: list[float],
         score_thresholds: list[float],
         datums: pc.Expression | None = None,
+        limit: int | None = None,
+        offset: int = 0,
     ) -> list[Metric]:
         """
         Computes examples at various thresholds.
@@ -614,6 +616,10 @@ def compute_examples(
             A list of score thresholds to compute metrics over.
         datums : pyarrow.compute.Expression, optional
             Option to filter datums by an expression.
+        limit : int, optional
+            Option to set a limit to the number of returned datum examples.
+        offset : int, default=0
+            Option to offset where examples are being created in the datum index.
 
         Returns
         -------
@@ -626,11 +632,6 @@ def compute_examples(
             raise ValueError("At least one score threshold must be passed.")
 
         metrics = []
-        tbl_columns = [
-            "datum_uid",
-            "gt_uid",
-            "pd_uid",
-        ]
         numeric_columns = [
             "datum_id",
             "gt_id",
@@ -640,14 +641,20 @@ def compute_examples(
             "iou",
             "pd_score",
         ]
-        for tbl, pairs in self._detailed_reader.iterate_tables_with_arrays(
-            columns=tbl_columns + numeric_columns,
-            numeric_columns=numeric_columns,
-            filter=datums,
+        for tbl in compute.paginate_index(
+            source=self._detailed_reader,
+            column_key="datum_id",
+            modifier=datums,
+            limit=limit,
+            offset=offset,
         ):
-            if pairs.size == 0:
+            if tbl.num_rows == 0:
                 continue
 
+            pairs = np.column_stack(
+                [tbl[col].to_numpy() for col in numeric_columns]
+            )
+
             index_to_datum_id = {}
             index_to_groundtruth_id = {}
             index_to_prediction_id = {}

tests/classification/test_examples.py

@@ -681,3 +681,65 @@ def test_examples_without_hardmax_animal_example(
         assert m in expected_metrics
     for m in expected_metrics:
         assert m in actual_metrics
+
+
+def test_examples_with_color_example_paginated(
+    loader: Loader,
+    classifications_color_example: list[Classification],
+):
+
+    loader.add_data(classifications_color_example)
+    evaluator = loader.finalize()
+
+    actual_metrics = evaluator.compute_examples(
+        score_thresholds=[0.5],
+        limit=3,
+        offset=1,
+    )
+
+    actual_metrics = [m.to_dict() for m in actual_metrics]
+    expected_metrics = [
+        {
+            "type": "Examples",
+            "value": {
+                "datum_id": "uid1",
+                "true_positives": [],
+                "false_positives": ["blue"],
+                "false_negatives": ["white"],
+            },
+            "parameters": {
+                "score_threshold": 0.5,
+                "hardmax": True,
+            },
+        },
+        {
+            "type": "Examples",
+            "value": {
+                "datum_id": "uid2",
+                "true_positives": [],
+                "false_positives": [],
+                "false_negatives": ["red"],
+            },
+            "parameters": {
+                "score_threshold": 0.5,
+                "hardmax": True,
+            },
+        },
+        {
+            "type": "Examples",
+            "value": {
+                "datum_id": "uid3",
+                "true_positives": [],
+                "false_positives": ["white"],
+                "false_negatives": ["blue"],
+            },
+            "parameters": {
+                "score_threshold": 0.5,
+                "hardmax": True,
+            },
+        },
+    ]
+    for m in actual_metrics:
+        assert m in expected_metrics
+    for m in expected_metrics:
+        assert m in actual_metrics