diff --git a/codebook.toml b/codebook.toml index 0e4963f4b..48c77ff4e 100644 --- a/codebook.toml +++ b/codebook.toml @@ -11,5 +11,9 @@ words = [ "TPE", "hyperparameters", "lenskit", + "mkdir", + "profiler", + "recommender", "seealso", + "zstd", ] diff --git a/src/lenskit/batch/__init__.py b/src/lenskit/batch/__init__.py index 5f0861c3c..0c0a55f85 100644 --- a/src/lenskit/batch/__init__.py +++ b/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", diff --git a/src/lenskit/batch/_queries.py b/src/lenskit/batch/_queries.py index 2dff02f12..a3392ca56 100644 --- a/src/lenskit/batch/_queries.py +++ b/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( diff --git a/src/lenskit/batch/_results.py b/src/lenskit/batch/_results.py index 67168293e..6fd966854 100644 --- a/src/lenskit/batch/_results.py +++ b/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: diff --git a/src/lenskit/batch/_runner.py b/src/lenskit/batch/_runner.py index f6ad64be6..1b6aece8d 100644 --- a/src/lenskit/batch/_runner.py +++ b/src/lenskit/batch/_runner.py @@ -15,6 +15,7 @@ from typing import Any, Literal, TypeAlias from lenskit.data import ( + GenericKey, QueryIDKey, UserIDKey, ) @@ -22,7 +23,7 @@ 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) diff --git a/src/lenskit/data/_collection/__init__.py b/src/lenskit/data/_collection/__init__.py index f02a07fb9..0153bf896 100644 --- a/src/lenskit/data/_collection/__init__.py +++ b/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", ] diff --git a/src/lenskit/data/_collection/_base.py b/src/lenskit/data/_collection/_base.py index 7eb5d58ba..9cd0d211f 100644 --- a/src/lenskit/data/_collection/_base.py +++ b/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,7 +608,6 @@ 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 @@ -618,13 +615,31 @@ def add_from(self, other: ItemListCollection, **fields: ID): 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]): diff --git a/src/lenskit/data/_collection/_parquet.py b/src/lenskit/data/_collection/_parquet.py new file mode 100644 index 000000000..1a2e81d54 --- /dev/null +++ b/src/lenskit/data/_collection/_parquet.py @@ -0,0 +1,83 @@ +# This file is part of LensKit. +# Copyright (C) 2018-2023 Boise State University. +# Copyright (C) 2023-2026 Drexel University. +# Licensed under the MIT license, see LICENSE.md for details. +# SPDX-License-Identifier: MIT + +""" +Code for saving Parquet item list collections. +""" + +from __future__ import annotations + +import os +from collections.abc import Sequence +from pathlib import Path +from typing import Generic, Literal + +import pyarrow as pa +from pyarrow.parquet import ParquetWriter + +from lenskit.data import ListILC +from lenskit.logging import get_logger + +from .._items import ItemList +from ._base import ItemListCollector +from ._keys import ID, K + +_log = get_logger(__name__) + + +class ParquetItemListCollector(ItemListCollector, Generic[K]): + """ + Item list collector that saves lists by batches to a Parquet file + (in native format). + """ + + batch_size: int + path: Path + writer: ParquetWriter | None = None + compression: Literal["zstd", "snappy"] | None + _cur_batch: ListILC[K] + + def __init__( + self, + path: Path | os.PathLike[str], + key: type[K] | Sequence[str], + batch_size: int = 5000, + compression: Literal["zstd", "snappy"] | None = "zstd", + ): + self.batch_size = batch_size + self.path = Path(path) + self.compression = compression + self._cur_batch = ListILC(key, index=False) + + def add(self, list: ItemList, *fields: ID, **kwfields: ID): + self._cur_batch.add(list, *fields, **kwfields) + self._maybe_flush() + + def close(self): + self._flush() + assert self.writer is not None + self.writer.close() + + def _maybe_flush(self): + if len(self._cur_batch) >= self.batch_size: + self._flush() + + def _flush(self): + for batch in self._cur_batch.record_batches(self.batch_size): + if self.writer is None: + _log.debug("opening Parquet writer", schema=batch.schema, file=str(self.path)) + self.writer = ParquetWriter( + self.path, batch.schema, compression=self.compression or "none" + ) + self.writer.write_batch(batch) + + if self.writer is None: + _log.warning("creating empty writer", file=str(self.path)) + schema = {k: pa.null() for k in self._cur_batch.key_fields} + schema["item_id"] = pa.null() + self.writer = ParquetWriter( + self.path, pa.schema(schema), compression=self.compression or "none" + ) diff --git a/src/lenskit/data/types.py b/src/lenskit/data/types.py index d88a2d87e..f32b868f9 100644 --- a/src/lenskit/data/types.py +++ b/src/lenskit/data/types.py @@ -35,25 +35,25 @@ "Extent", ] -FeedbackType: TypeAlias = Literal["explicit", "implicit"] +type FeedbackType = Literal["explicit", "implicit"] "Types of feedback supported." -CoreID: TypeAlias = int | str | bytes +type CoreID = int | str | bytes "Core (non-NumPy) identifier types." -NPID: TypeAlias = np.integer[Any] | np.str_ | np.bytes_ | np.object_ +type NPID = np.integer[Any] | np.str_ | np.bytes_ | np.object_ "NumPy entity identifier types." -ID: TypeAlias = CoreID | NPID +type ID = CoreID | NPID "Allowable identifier types." -IDArray: TypeAlias = np.ndarray[tuple[int], np.dtype[NPID]] +type IDArray = np.ndarray[tuple[int], np.dtype[NPID]] "NumPy arrays of identifiers." -IDSequence: TypeAlias = """ +type IDSequence = ( Sequence[ID] | IDArray | pa.StringArray | pa.IntegerArray[Any] | pa.ChunkedArray[Any] | pd.Series[CoreID] - """ +) "Sequences of identifiers." _UIPT = TypeVar("_UIPT") diff --git a/tests/data/test_collection.py b/tests/data/test_collection.py index 9ee1093cd..adc63a9c8 100644 --- a/tests/data/test_collection.py +++ b/tests/data/test_collection.py @@ -380,7 +380,8 @@ def test_save_parquet_with_mkdir(tmpdir: Path): assert (tmpdir / "subdir").exists() f_no_mkdir = tmpdir / "no_mkdir" / "items.parquet" - ilc.save_parquet(f_no_mkdir, mkdir=False) + with raises(FileNotFoundError): + ilc.save_parquet(f_no_mkdir, mkdir=False) assert not (tmpdir / "no_mkdir").exists()