ranklens.io

ranklens.io

Streaming readers for runs and relevance judgements.

ClickLogSchema dataclass

Columns of a feedback log: one row per shown document.

Feedback is either a numeric reward column or an event column translated by a FeedbackScale; exactly one of them is set. An empty event means the document was shown and nothing happened.

Source code in src/ranklens/io/clicklog.py

@dataclass(frozen=True, slots=True)
class ClickLogSchema:
    """Columns of a feedback log: one row per shown document.

    Feedback is either a numeric ``reward`` column or an ``event`` column translated
    by a `FeedbackScale`; exactly one of them is set. An empty event means the document
    was shown and nothing happened.
    """

    impression_id: str = "impression_id"
    query_id: str = "query_id"
    doc_id: str = "doc_id"
    position: str = "position"
    reward: str | None = "reward"
    event: str | None = None

    def __post_init__(self) -> None:
        if (self.reward is None) == (self.event is None):
            raise ValueError("set exactly one feedback column: reward or event")

    @property
    def required(self) -> tuple[str, ...]:
        feedback = self.reward or self.event
        assert feedback is not None  # guaranteed by __post_init__
        return (self.impression_id, self.query_id, self.doc_id, self.position, feedback)

ErrorCollector

Accumulates recoverable data errors; read the result with snapshot.

Memory is bounded: at most max_examples messages are kept per error type, no matter how many rows are skipped.

Source code in src/ranklens/io/errors.py

class ErrorCollector:
    """Accumulates recoverable data errors; read the result with `snapshot`.

    Memory is bounded: at most ``max_examples`` messages are kept per error type,
    no matter how many rows are skipped.
    """

    def __init__(self, max_examples: int = 5) -> None:
        if max_examples < 0:
            raise ValueError(f"max_examples must be >= 0, got {max_examples}")
        self._max_examples = max_examples
        self._counts: Counter[str] = Counter()
        self._examples: dict[str, list[str]] = {}

    def record(self, error: DataError) -> None:
        kind = type(error).__name__
        self._counts[kind] += 1
        examples = self._examples.setdefault(kind, [])
        if len(examples) < self._max_examples:
            examples.append(str(error))

    def snapshot(self) -> ErrorSummary:
        """Immutable summary of everything recorded so far."""
        return ErrorSummary(
            total=self._counts.total(),
            counts=tuple(self._counts.most_common()),
            examples=tuple((kind, tuple(msgs)) for kind, msgs in self._examples.items() if msgs),
        )

snapshot

snapshot() -> ErrorSummary

Immutable summary of everything recorded so far.

Source code in src/ranklens/io/errors.py

def snapshot(self) -> ErrorSummary:
    """Immutable summary of everything recorded so far."""
    return ErrorSummary(
        total=self._counts.total(),
        counts=tuple(self._counts.most_common()),
        examples=tuple((kind, tuple(msgs)) for kind, msgs in self._examples.items() if msgs),
    )

FeatureQuery dataclass

Documents of one query: rows[i] are the features of docs[i].

Source code in src/ranklens/io/features.py

@dataclass(frozen=True, slots=True)
class FeatureQuery:
    """Documents of one query: ``rows[i]`` are the features of ``docs[i]``."""

    query_id: QueryId
    docs: tuple[DocId, ...]
    rows: tuple[tuple[float, ...], ...]
    judgements: tuple[tuple[DocId, float], ...]

FeatureSchema dataclass

Id and grade columns; features empty means every other column.

Source code in src/ranklens/io/features.py

@dataclass(frozen=True, slots=True)
class FeatureSchema:
    """Id and grade columns; ``features`` empty means every other column."""

    query_id: str = "query_id"
    doc_id: str = "doc_id"
    relevance: str | None = "relevance"
    features: tuple[str, ...] = ()

ArrowStreamExportable

Bases: Protocol

In-memory table exporting the Arrow C stream: pandas, polars, pyarrow, DuckDB.

Source code in src/ranklens/io/formats.py

@runtime_checkable
class ArrowStreamExportable(Protocol):
    """In-memory table exporting the Arrow C stream: pandas, polars, pyarrow, DuckDB."""

    def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ...

PairedRuns

Merge-join of two runs sorted by query_id: pairs of rankings of the same query.

Written as an iterator rather than built on zip because the two runs have different, only partly overlapping sets of queries: zip would pair unrelated rankings as soon as one run skips a query. The merge keeps one ranking from each side in memory (a lookahead of one), so two huge runs are still streamed.

Both sides must be sorted by query_id as strings — iter_run(check="sorted") guarantees that. Queries present on one side only are counted, not paired. Used as a context manager it closes both streams — handy when the loop breaks early.

Source code in src/ranklens/io/paired.py

class PairedRuns:
    """Merge-join of two runs sorted by query_id: pairs of rankings of the same query.

    Written as an iterator rather than built on ``zip`` because the two runs have
    different, only partly overlapping sets of queries: ``zip`` would pair unrelated
    rankings as soon as one run skips a query. The merge keeps one ranking from each
    side in memory (a lookahead of one), so two huge runs are still streamed.

    Both sides must be sorted by query_id as strings — ``iter_run(check="sorted")``
    guarantees that. Queries present on one side only are counted, not paired.
    Used as a context manager it closes both streams — handy when the loop breaks early.
    """

    def __init__(self, left: Iterable[RankedList], right: Iterable[RankedList]) -> None:
        self._left, self._right = iter(left), iter(right)
        self._ahead_left = next(self._left, None)
        self._ahead_right = next(self._right, None)
        self._previous_left: QueryId | None = None
        self._previous_right: QueryId | None = None
        self.only_left = 0
        """Queries seen so far only in the left run."""
        self.only_right = 0
        """Queries seen so far only in the right run."""

    def __iter__(self) -> "PairedRuns":
        return self

    def __next__(self) -> tuple[RankedList, RankedList]:
        while self._ahead_left is not None and self._ahead_right is not None:
            left, right = self._ahead_left, self._ahead_right
            if left.query_id == right.query_id:
                self._advance_left()
                self._advance_right()
                return left, right
            if left.query_id < right.query_id:
                self.only_left += 1
                self._advance_left()
            else:
                self.only_right += 1
                self._advance_right()
        self._drain()
        raise StopIteration

    def _drain(self) -> None:
        """Count the tail of whichever run is longer; the other one is exhausted."""
        while self._ahead_left is not None:
            self.only_left += 1
            self._advance_left()
        while self._ahead_right is not None:
            self.only_right += 1
            self._advance_right()

    def _advance_left(self) -> None:
        self._previous_left = self._check(self._ahead_left, self._previous_left, "left run")
        self._ahead_left = next(self._left, None)

    def _advance_right(self) -> None:
        self._previous_right = self._check(self._ahead_right, self._previous_right, "right run")
        self._ahead_right = next(self._right, None)

    @staticmethod
    def _check(current: RankedList | None, previous: QueryId | None, side: str) -> QueryId | None:
        # a merge-join on unsorted input silently drops queries, so it is an error
        if current is not None and previous is not None and current.query_id <= previous:
            raise UnsortedInputError(None, previous, current.query_id, side)
        return None if current is None else current.query_id

    def __enter__(self) -> "PairedRuns":
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        for stream in (self._left, self._right):
            close = getattr(stream, "close", None)
            if close is not None:
                close()

only_left instance-attribute

only_left = 0

Queries seen so far only in the left run.

only_right instance-attribute

only_right = 0

Queries seen so far only in the right run.

QrelsSchema dataclass

Names of qrels columns in csv/tsv headers and jsonl keys.

Source code in src/ranklens/io/schema.py

@dataclass(frozen=True, slots=True)
class QrelsSchema:
    """Names of qrels columns in csv/tsv headers and jsonl keys."""

    query_id: str = "query_id"
    doc_id: str = "doc_id"
    relevance: str = "relevance"

    # standard TREC qrels line: qid iteration docno relevance
    TREC_COLUMNS = ("query_id", "iteration", "doc_id", "relevance")

    @classmethod
    def trec(cls) -> "QrelsSchema":
        return cls()

    @property
    def required(self) -> tuple[str, ...]:
        return (self.query_id, self.doc_id, self.relevance)

RunSchema dataclass

Names of run columns in csv/tsv headers and jsonl keys.

Source code in src/ranklens/io/schema.py

@dataclass(frozen=True, slots=True)
class RunSchema:
    """Names of run columns in csv/tsv headers and jsonl keys."""

    query_id: str = "query_id"
    doc_id: str = "doc_id"
    score: str | None = None
    rank: str | None = None
    segments: tuple[str, ...] = ()

    # standard TREC run line: qid Q0 docno rank score tag
    TREC_COLUMNS = ("query_id", "q0", "doc_id", "rank", "score", "run_tag")

    @classmethod
    def trec(cls) -> "RunSchema":
        return cls(score="score", rank="rank")

    @property
    def required(self) -> tuple[str, ...]:
        optional = (self.score, self.rank)
        return (self.query_id, self.doc_id, *(c for c in optional if c), *self.segments)

iter_clicklog

iter_clicklog(
    source: Source,
    *,
    schema: ClickLogSchema | None = None,
    scale: FeedbackScale | None = None,
    fmt: FormatName | None = None,
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> Iterator[Impression]

Stream a feedback log as one Impression per logged showing.

Rows of an impression must be contiguous; within it they may come in any order and are sorted by position. With an event column, scale turns events into rewards. Invalid rows follow the same rules as in iter_run: they raise in strict mode and are skipped and recorded otherwise.

Source code in src/ranklens/io/clicklog.py

def iter_clicklog(
    source: Source,
    *,
    schema: ClickLogSchema | None = None,
    scale: FeedbackScale | None = None,
    fmt: FormatName | None = None,
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> Iterator[Impression]:
    """Stream a feedback log as one `Impression` per logged showing.

    Rows of an impression must be contiguous; within it they may come in any order
    and are sorted by position. With an ``event`` column, ``scale`` turns events into
    rewards. Invalid rows follow the same rules as in `iter_run`: they raise in
    ``strict`` mode and are skipped and recorded otherwise.
    """
    schema = schema or ClickLogSchema()
    if schema.event is not None and scale is None:
        raise ValueError("an event column needs a FeedbackScale to turn events into rewards")
    fmt = fmt or detect_format(source)
    if fmt == "trec":
        raise ValueError("TREC has no feedback log format; use csv, tsv, jsonl or parquet")
    return _impressions(source, fmt, schema, scale, strict, errors)

column_names

column_names(
    source: Source, fmt: FormatName | None = None
) -> list[str]

Columns of a csv/tsv header, of the first jsonl object, or of a parquet/Feather schema.

Source code in src/ranklens/io/features.py

def column_names(source: Source, fmt: FormatName | None = None) -> list[str]:
    """Columns of a csv/tsv header, of the first jsonl object, or of a parquet/Feather schema."""
    fmt = fmt or detect_format(source)
    if fmt in ("parquet", "feather"):
        from ranklens.io.arrow import arrow_column_names

        return arrow_column_names(source, fmt)
    if fmt not in ("csv", "tsv", "jsonl") or not isinstance(source, str | PathLike):
        raise ValueError(f"cannot list the columns of {fmt!r} input: name the features explicitly")
    with open_text(source) as fh:
        if fmt == "jsonl":
            for line in fh:
                if line.strip():
                    return list(json.loads(line))
            return []
        header = next(csv.reader(fh, delimiter="," if fmt == "csv" else "\t"), [])
        return [column.strip() for column in header]

read_feature_table

read_feature_table(
    source: Source,
    *,
    schema: FeatureSchema | None = None,
    fmt: FormatName | None = None,
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> FeatureTable

Load a feature table grouped by query, in the order queries first appear.

Rows of a query need not be contiguous. A non-numeric feature or grade is a malformed row: skipped and recorded, or raised in strict mode, as in the other readers. A feature may be nan — a missing value, which boosting models handle; a grade must be finite, and an empty grade leaves the document without a judgement.

Source code in src/ranklens/io/features.py

def read_feature_table(
    source: Source,
    *,
    schema: FeatureSchema | None = None,
    fmt: FormatName | None = None,
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> FeatureTable:
    """Load a feature table grouped by query, in the order queries first appear.

    Rows of a query need not be contiguous. A non-numeric feature or grade is a malformed
    row: skipped and recorded, or raised in ``strict`` mode, as in the other readers.
    A feature may be ``nan`` — a missing value, which boosting models handle; a grade
    must be finite, and an empty grade leaves the document without a judgement.
    """
    schema = schema or FeatureSchema()
    fmt = fmt or detect_format(source)
    if fmt == "trec":
        raise ValueError("TREC has no feature table format; use csv, tsv, jsonl or parquet")
    ids = [schema.query_id, schema.doc_id, *([schema.relevance] if schema.relevance else [])]
    names = schema.features or tuple(c for c in column_names(source, fmt) if c not in ids)
    if not names:
        raise ValueError("the table has no feature columns")
    name = source_name(source)
    collector = errors if errors is not None else ErrorCollector()
    on_error = _raise if strict else collector.record

    grouped: dict[QueryId, tuple[list[DocId], list[tuple[float, ...]], list[tuple[DocId, float]]]]
    grouped = {}
    seen: set[tuple[QueryId, DocId]] = set()
    for record in iter_records(
        source, fmt, required=[*ids, *names], trec_columns=(), on_error=on_error
    ):
        try:
            query, doc, values, grade = _parse(record, schema, names, name)
        except MalformedRowError as exc:
            on_error(exc)
            continue
        if (query, doc) in seen:
            on_error(DuplicateDocumentError(query, doc))
            continue
        seen.add((query, doc))
        docs, rows, judgements = grouped.setdefault(query, ([], [], []))
        docs.append(doc)
        rows.append(values)
        if grade is not None:
            judgements.append((doc, grade))
    if errors is None:
        _warn_if_skipped(collector)
    return FeatureTable(
        feature_names=tuple(names),
        queries=tuple(
            FeatureQuery(query, tuple(docs), tuple(rows), tuple(judgements))
            for query, (docs, rows, judgements) in grouped.items()
        ),
    )

detect_format

detect_format(source: Source) -> FormatName

Format of a file by extension, or "arrow" for an in-memory table.

A compression suffix (.gz, .bz2, .xz) is skipped for text formats; parquet and Feather compress internally.

Source code in src/ranklens/io/formats.py

def detect_format(source: Source) -> FormatName:
    """Format of a file by extension, or ``"arrow"`` for an in-memory table.

    A compression suffix (``.gz``, ``.bz2``, ``.xz``) is skipped for text formats;
    parquet and Feather compress internally.
    """
    if not isinstance(source, str | PathLike):
        if isinstance(source, ArrowStreamExportable):
            return "arrow"
        raise TypeError(
            f"expected a path or an Arrow-compatible table, got {type(source).__name__}"
        )
    suffixes = [s.lower() for s in Path(source).suffixes]
    compressed = bool(suffixes) and suffixes[-1] in _COMPRESSION
    if compressed:
        suffixes.pop()
    fmt = _EXTENSIONS.get(suffixes[-1]) if suffixes else None
    if fmt is None or (compressed and fmt in _ARROW_FORMATS):
        raise UnsupportedFormatError(str(source), _EXTENSIONS)
    return fmt

iter_run

iter_run(
    source: Source,
    *,
    schema: RunSchema | None = None,
    fmt: FormatName | None = None,
    order: Literal["file", "rank", "score"] | None = None,
    check: Literal["grouped", "sorted"] = "grouped",
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> Iterator[RankedList]

Stream a run as one RankedList per query.

source is a file (csv, tsv, jsonl, TREC, parquet, Feather) or an in-memory table (pandas, polars, pyarrow). Only the rows of the current query are held in memory.

order sets document positions: "file" order; ascending "rank" column (stable for ties); or descending "score" with ties broken by descending doc_id compared as strings — exactly as trec_eval, which ignores the file order. By default TREC runs are ordered by score and other formats by file order. check sets the input contract:

• "grouped": rows of a query are contiguous, queries in any order. Memory grows with the number of queries (ids of finished queries are kept).
• "sorted": query ids ascend as strings. Constant memory; required when two runs are merged side by side.

Invalid rows raise in strict mode. Otherwise they are skipped and recorded in errors; without a collector a SkippedRowsWarning is emitted when the file is exhausted. Grouping and sorting violations always raise.

Source code in src/ranklens/io/readers.py

def iter_run(
    source: Source,
    *,
    schema: RunSchema | None = None,
    fmt: FormatName | None = None,
    order: Literal["file", "rank", "score"] | None = None,
    check: Literal["grouped", "sorted"] = "grouped",
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> Iterator[RankedList]:
    """Stream a run as one `RankedList` per query.

    ``source`` is a file (csv, tsv, jsonl, TREC, parquet, Feather) or an in-memory
    table (pandas, polars, pyarrow). Only the rows of the current query are held
    in memory.

    ``order`` sets document positions: ``"file"`` order; ascending ``"rank"`` column
    (stable for ties); or descending ``"score"`` with ties broken by descending doc_id
    compared as strings — exactly as trec_eval, which ignores the file order.
    By default TREC runs are ordered by score and other formats by file order.
    ``check`` sets the input contract:

    - ``"grouped"``: rows of a query are contiguous, queries in any order.
      Memory grows with the number of *queries* (ids of finished queries are kept).
    - ``"sorted"``: query ids ascend as strings. Constant memory; required
      when two runs are merged side by side.

    Invalid rows raise in ``strict`` mode. Otherwise they are skipped and recorded
    in ``errors``; without a collector a `SkippedRowsWarning` is emitted
    when the file is exhausted. Grouping and sorting violations always raise.
    """
    fmt = fmt or detect_format(source)
    schema = schema or (RunSchema.trec() if fmt == "trec" else RunSchema())
    order = order or ("score" if fmt == "trec" else "file")
    if order == "rank" and schema.rank is None:
        raise ValueError('order="rank" requires a rank column in the schema')
    if order == "score" and schema.score is None:
        raise ValueError('order="score" requires a score column in the schema')
    return _rankings(source, fmt, schema, order, check, strict, errors)

read_qrels

read_qrels(
    source: Source,
    *,
    schema: QrelsSchema | None = None,
    fmt: FormatName | None = None,
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> Qrels

Load relevance judgements into memory: qrels[query_id][doc_id] -> relevance.

source is any source accepted by iter_run.

Unlike runs, qrels are small and are needed for random access, so rows may come in any order. A repeated (query_id, doc_id) pair is an error: the first judgement is kept in non-strict mode. Error handling is the same as in iter_run.

Source code in src/ranklens/io/readers.py

def read_qrels(
    source: Source,
    *,
    schema: QrelsSchema | None = None,
    fmt: FormatName | None = None,
    strict: bool = False,
    errors: ErrorCollector | None = None,
) -> Qrels:
    """Load relevance judgements into memory: ``qrels[query_id][doc_id] -> relevance``.

    ``source`` is any source accepted by `iter_run`.

    Unlike runs, qrels are small and are needed for random access, so rows may come
    in any order. A repeated (query_id, doc_id) pair is an error: the first
    judgement is kept in non-strict mode. Error handling is the same as in
    `iter_run`.
    """
    fmt = fmt or detect_format(source)
    schema = schema or (QrelsSchema.trec() if fmt == "trec" else QrelsSchema())
    name = source_name(source)
    collector = errors if errors is not None else ErrorCollector()
    on_error = _raise if strict else collector.record

    records = iter_records(
        source,
        fmt,
        required=schema.required,
        trec_columns=QrelsSchema.TREC_COLUMNS,
        on_error=on_error,
    )
    qrels: Qrels = {}
    for record in records:
        try:
            row = parse_qrels_row(record, schema, name)
        except MalformedRowError as exc:
            on_error(exc)
            continue
        judgements = qrels.setdefault(row.query_id, {})
        if row.doc_id in judgements:
            previous = judgements[row.doc_id]
            reason = f"repeated judgement for document {row.doc_id!r} (first: {previous:g})"
            on_error(MalformedRowError(row.line_no, record.raw, reason, path=name))
            continue
        judgements[row.doc_id] = row.relevance

    if errors is None:
        _warn_if_skipped(collector)
    return qrels