Expose batch results as a row-by-row iterator

codebook.toml

@@ -11,5 +11,9 @@ words = [
   "TPE",
   "hyperparameters",
   "lenskit",
+  "mkdir",
+  "profiler",
+  "recommender",
   "seealso",
+  "zstd",
 ]

src/lenskit/batch/__init__.py

@@ -16,11 +16,12 @@
 from lenskit.pipeline import Pipeline, PipelineProfiler
 
 from ._queries import BatchInput, BatchRecRequest, TestRequestAdapter
-from ._results import BatchResults
+from ._results import BatchResultRow, BatchResults
 from ._runner import BatchPipelineRunner, InvocationSpec
 
 __all__ = [
     "BatchPipelineRunner",
+    "BatchResultRow",
     "BatchResults",
     "BatchRecRequest",
     "BatchInput",

src/lenskit/batch/_queries.py

@@ -36,6 +36,17 @@
 """
 
 
+@dataclass
+class NormalizedQueryBatch:
+    """
+    Internal representation for a normalized, runnable batch of queries.
+    """
+
+    key_type: type[GenericKey]
+    queries: Iterable[ResolvedBatchRequest]
+    count: int | None
+
+
 class BatchRecRequest(TypedDict, total=False):
     """
     Full recommendation request for batch inference, including candidate items.
@@ -166,7 +177,7 @@ def __iter__(self) -> Iterator[BatchRecRequest]:
 
 def normalize_query_input(
     queries: BatchInput,
-) -> tuple[type[GenericKey], Iterable[ResolvedBatchRequest], int | None]:
+) -> NormalizedQueryBatch:
     kt = None
 
     if isinstance(queries, ItemListCollection):
@@ -197,7 +208,7 @@ def normalize_query_input(
     else:
         raise ValueError("query must have one of query_id, user_id")
 
-    return kt, _iter_queries(q_first, q_iter), n
+    return NormalizedQueryBatch(kt, _iter_queries(q_first, q_iter), n)
 
 
 def _iter_queries(

src/lenskit/batch/_results.py

@@ -6,14 +6,24 @@
 
 from __future__ import annotations
 
-from typing import Sequence
+from typing import NamedTuple, Sequence
 
-from lenskit.data import GenericKey, ItemListCollection
+from lenskit.data import ID, GenericKey, ItemListCollection, key_dict
 
-type BatchResultRow = tuple[GenericKey, dict[str, object]]
-"""
-Results for a single query in the batch recommendations.
-"""
+
+class BatchResultRow(NamedTuple):
+    """
+    Results for a single query in the batch recommendations.
+    """
+
+    key: GenericKey
+    "The key (e.g. user ID) associated with this result row."
+    outputs: dict[str, object]
+    "The outputs associated with this result row."
+
+    def key_dict(self) -> dict[str, ID]:
+        "Get the row's key as a dictionary."
+        return key_dict(self.key)
 
 
 class BatchResults:

src/lenskit/batch/_runner.py

@@ -15,14 +15,15 @@
 from typing import Any, Literal, TypeAlias
 
 from lenskit.data import (
+    GenericKey,
     QueryIDKey,
     UserIDKey,
 )
 from lenskit.logging import Stopwatch, get_logger, item_progress
 from lenskit.parallel import get_parallel_config, is_free_threaded
 from lenskit.pipeline import Pipeline, PipelineProfiler, ProfileSink
 
-from ._queries import BatchInput, ResolvedBatchRequest, normalize_query_input
+from ._queries import BatchInput, NormalizedQueryBatch, ResolvedBatchRequest, normalize_query_input
 from ._results import BatchResultRow, BatchResults
 
 _log = get_logger(__name__)
@@ -159,7 +160,7 @@ def run(
         queries: BatchInput,
     ) -> BatchResults:
         """
-        Run the pipeline and return its results.
+        Run the pipeline and collect its results.
 
         .. note::
 
@@ -177,39 +178,81 @@ def run(
             The batch results, mapping output names to item list collections of
             outputs.
         """
+        if queries is None:  # pragma: nocover
+            raise RuntimeError("no queries specified")
+
+        nqs = normalize_query_input(queries)
+        results = BatchResults(nqs.key_type)
+        with closing(self._run_impl(pipeline, nqs)) as rs:
+            for key, outs in rs:
+                for cn, cr in outs.items():
+                    results.add_result(cn, key, cr)
+
+        return results
+
+    def run_iter(
+        self,
+        pipeline: Pipeline,
+        queries: BatchInput,
+    ) -> Generator[BatchResultRow]:
+        """
+        Run the pipeline and yield its results.
+
+        This generator should be closed when it is done or the run is aborted,
+        to ensure resources are properly cleaned up.  The best way to use this
+        method is as follows::
+
+            with closing(runner.run_iter()) as results:
+                for key, outs in results:
+                    # do something with the data
+                    pass
+
+        .. note::
 
+            The runner does **not** guarantee that results are in the same order
+            as the original inputs — with parallelism, they may be yielded in
+            the order they are ready.
+
+        Args:
+            pipeline:
+                The pipeline to run.
+            queries:
+                The collection of test queries use.  See :ref:`batch-queries`
+                for details on the various input formats.
+        """
         if queries is None:  # pragma: nocover
             raise RuntimeError("no queries specified")
 
+        nqs = normalize_query_input(queries)
+        return self._run_impl(pipeline, nqs)
+
+    def _run_impl(
+        self, pipeline: Pipeline, queries: NormalizedQueryBatch
+    ) -> Generator[BatchResultRow]:
         prof = self.profiler
         if prof is not None:
             prof = prof.multiprocess()
 
-        key_type, q_iter, nq = normalize_query_input(queries)
-
-        log = _log.bind(name=pipeline.name, n_queries=nq, n_jobs=self.n_jobs)
+        log = _log.bind(name=pipeline.name, n_queries=queries.count, n_jobs=self.n_jobs)
 
         log.info("beginning batch run")
 
-        with closing(self._run_results(pipeline, prof, q_iter)) as tasks:
-            with item_progress("Inference", nq) as progress:
+        with closing(self._run_results(pipeline, prof, queries.queries)) as tasks:
+            with item_progress("Inference", queries.count) as progress:
                 # release our reference, will sometimes free the pipeline memory in this process
                 del pipeline
-                results = BatchResults(key_type)
                 timer = Stopwatch()
                 n = 0
                 for key, outs in tasks:
                     n += 1
-                    for cn, cr in outs.items():
-                        results.add_result(cn, key, cr)
+                    yield key, outs
+
                     progress.update()
                 timer.stop()
 
                 rate_ms = timer.elapsed() / n * 1000
                 log.info("finished running in %s", timer, time_per_query="{:.1f}ms".format(rate_ms))
 
-        return results
-
     def _run_results(
         self,
         pipeline: Pipeline,
@@ -268,8 +311,9 @@ def run_pipeline(
     profiler: ProfileSink | None,
     req: ResolvedBatchRequest,
 ) -> BatchResultRow:
+    key: GenericKey
     if isinstance(req.query.query_id, tuple):
-        key = req.query.query_id
+        key = req.query.query_id  # type: ignore
     elif req.query.query_id is not None:
         key = QueryIDKey(req.query.query_id)
     elif req.query.user_id is not None:
@@ -295,4 +339,4 @@ def run_pipeline(
         for cname, oname in inv.components.items():
             result[oname] = outs[cname]
 
-    return key, result  # type: ignore
+    return BatchResultRow(key, result)