Move Parquet saving into a separate item list collector

codebook.toml

@@ -8,4 +8,8 @@ words = [
   "TPE",
   "hyperparameters",
   "lenskit",
+  "mkdir",
+  "profiler",
+  "recommender",
+  "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)

src/lenskit/data/_collection/__init__.py

@@ -11,6 +11,7 @@
 from ._base import ItemListCollection, ItemListCollector, MutableItemListCollection
 from ._keys import GenericKey, QueryIDKey, UserIDKey, key_dict
 from ._list import ListILC
+from ._parquet import ParquetItemListCollector
 
 __all__ = [
     "GenericKey",
@@ -19,6 +20,7 @@
     "ItemListCollection",
     "ItemListCollector",
     "MutableItemListCollection",
+    "ParquetItemListCollector",
     "ListILC",
     "key_dict",
 ]

src/lenskit/data/_collection/_base.py

@@ -27,8 +27,9 @@
 
 import pandas as pd
 import pyarrow as pa
-from pyarrow.parquet import ParquetDataset, ParquetWriter
+from pyarrow.parquet import ParquetDataset
 from pydantic import JsonValue
+from typing_extensions import Self
 
 from lenskit.diagnostics import DataWarning
 from lenskit.logging import get_logger
@@ -367,20 +368,11 @@ def save_parquet(
         if layout == "flat":
             log.debug("saving flat Parquet file")
             self.to_df().to_parquet(path, compression=compression)
-            return
-
-        writer = None
-        try:
-            for batch in self.record_batches(batch_size):
-                if writer is None:
-                    log.debug("opening Parquet writer", schema=batch.schema)
-                    writer = ParquetWriter(
-                        Path(path), batch.schema, compression=compression or "snappy"
-                    )
-                writer.write_batch(batch)
-        finally:
-            if writer is not None:
-                writer.close()
+        else:
+            from ._parquet import ParquetItemListCollector
+
+            with ParquetItemListCollector(path, self.key_type) as out:
+                out.add_from(self)
 
     @overload
     @classmethod
@@ -595,6 +587,12 @@ def __str__(self):
 class ItemListCollector(Protocol):
     """
     Collect item lists with associated keys, as in :class:`ItemListCollection`.
+
+    An item list collector is also a context manager that yields itself and
+    calls :meth:`close` on exit.
+
+    Stability:
+        Caller
     """
 
     @abstractmethod
@@ -610,21 +608,38 @@ def add(self, list: ItemList, *fields: ID, **kwfields: ID):  # pragma: nocover
         """
         raise NotImplementedError()
 
-    @abstractmethod
     def add_from(self, other: ItemListCollection, **fields: ID):
         """
         Add all collection from another collection to this collection.  If field
         values are supplied, they are used to supplement or overwrite the keys
         in ``other``; a common use case is to add results from multiple
         recommendation runs and save them a single field.
 
+        The default implementation delegates to :meth:`add` in a loop.
+
         Args:
             other:
                 The item list collection to incorporate into this one.
             fields:
                 Additional key fields (must be specified by name).
         """
-        raise NotImplementedError()
+        for key, list in other:
+            kd = fields | key_dict(key)
+            self.add(list, **kd)
+
+    def close(self):
+        """
+        Close this collector. After the collector is closed, no further writes are
+        an error (but may not be proactively checked).
+
+        The default implementation does nothing.
+        """
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
 
 
 class MutableItemListCollection[K: GenericKey](ItemListCollector, ItemListCollection[K]):