Custom data sources

A table provider supplies documents from any storage system. Implement the TableProvider protocol and register it on the session.

The protocol

import pyarrow as pa
from quail.catalog import ScanRequest, TableStatistics


class TableProvider:
    id_col: str

    @property
    def columns(self) -> tuple[str, ...]: ...
    def schema(self) -> pa.Schema: ...
    def content_identity(self) -> str: ...
    def statistics(self) -> TableStatistics: ...
    def scan(self, request: ScanRequest) -> pa.RecordBatchReader: ...

Source: quail/catalog.py.

Example: Arrow rows

class RowsProvider:
    id_col = "id"

    def __init__(self, rows: list[dict]):
        self._table = pa.Table.from_pylist(rows)

    @property
    def columns(self):
        return tuple(self._table.column_names)

    def schema(self):
        return self._table.schema

    def content_identity(self):
        return f"rows:{len(self._table)}"

    def statistics(self):
        return TableStatistics(row_count=len(self._table))

    def scan(self, request: ScanRequest):
        table = self._table.select(list(request.columns))
        if request.limit is not None:
            table = table.slice(0, request.limit)
        return table.to_reader(max_chunksize=request.batch_rows)


session.register("docs", RowsProvider(rows))

Key points

content_identity() must change when the content changes. The session uses it to decide whether to re-tokenize.
scan() returns bounded batches for only the columns in request.columns. The session tokenizes the document column and writes tokens and projected columns to a temporary file. Batches are released as they are written, so the source can be larger than process memory.
On Modal, create providers inside the GPU function. Attach credentials through the function's secrets argument.

Custom data sources

The protocol

Example: Arrow rows

Key points

On this page