Table Reads¶

def scan_table(
    self,
    location: str | Path | tuple[str | Path, ...],
    columns: tuple[str, ...] | None = None,
    filters: pl.Expr | None = None,
    limit: int | None = None,
) -> pl.LazyFrame:
    """Create a lazy parquet scan for one or more table locations.

    Args:
        location: One parquet location, or a tuple of locations to scan as one
            lazy frame.
        columns: Optional columns to select from the scan.
        filters: Optional Polars expression applied to the lazy frame.
        limit: Optional maximum number of rows to read.

    Returns:
        Polars lazy frame for the selected parquet data.

    Examples:
        Load selected manifest rows without materializing unrelated columns:

        >>> manifest = store.scan_table(
        ...     "manifest.parquet",
        ...     columns=("protocol", "segments"),
        ...     filters=pl.col("protocol") == "cycling",
        ... ).collect()
    """
    resolved = (
        [self.resolve(path) for path in location]
        if isinstance(location, tuple)
        else self.resolve(location)
    )
    lf = pl.scan_parquet(resolved)
    if columns is not None:
        lf = lf.select(list(columns))
    if filters is not None:
        lf = lf.filter(filters)
    if limit is not None:
        lf = lf.limit(limit)
    return lf

def iter_table_chunks(
    self,
    location: str | Path | tuple[str | Path, ...],
    chunk_rows: int,
    columns: tuple[str, ...] | None = None,
    filters: pl.Expr | None = None,
) -> Iterator[pl.DataFrame]:
    """Yield sequential parquet batches from one or more complete tables.

    This method streams every row from each table in order, with batches of
    at most `chunk_rows` before any optional filter is applied. Use it when
    a stage should process a complete table without loading it all into
    memory.

    Args:
        location: One parquet location, or a tuple of locations to process in
            order.
        chunk_rows: Maximum number of rows requested per parquet batch before
            filtering.
        columns: Optional columns to read from the parquet file.
        filters: Optional Polars expression applied to each batch after it is
            read.

    Yields:
        Non-empty dataframe chunks.

    Raises:
        ValueError: If `chunk_rows` is less than one.

    Examples:
        Consume a scratch table in bounded-size chunks before deleting it:

        >>> for chunk in store.iter_table_chunks(temp_path, chunk_rows):
        ...     writer.append(chunk, metadata, source_paths)
        >>> store.delete_file(temp_path, missing_ok=True)
    """
    if chunk_rows < 1:
        raise ValueError(f"chunk_rows must be >= 1, got {chunk_rows}")
    resolved = (
        [self.resolve(path) for path in location]
        if isinstance(location, tuple)
        else [self.resolve(location)]
    )
    selected_columns = list(columns) if columns is not None else None
    for path in resolved:
        parquet_file = pq.ParquetFile(path)
        for arrow_batch in parquet_file.iter_batches(
            batch_size=chunk_rows,
            columns=selected_columns,
        ):
            frame = pl.from_arrow(arrow_batch)
            if not isinstance(frame, pl.DataFrame):
                raise TypeError(
                    f"Expected batch conversion to return DataFrame, got {type(frame).__name__}"
                )
            if filters is not None:
                frame = frame.filter(filters)
            if frame.height > 0:
                yield frame

def iter_table_slices(
    self,
    location: str | Path,
    slices: tuple[tuple[int, int], ...],
    chunk_rows: int,
    columns: tuple[str, ...] | None = None,
) -> Iterator[pl.DataFrame]:
    """Yield selected row windows from a single parquet table.

    Unlike `iter_table_chunks`, this method does not scan the whole table. It
    reads only row groups that overlap the requested `slices` and emits the
    selected rows in chunks of at most `chunk_rows`. Use it for
    manifest segments or sharded selections that reference specific row
    ranges.

    Args:
        location: Parquet table location to read from.
        slices: Row windows as `(row_start, row_count)` tuples.
        chunk_rows: Maximum number of rows yielded per chunk.
        columns: Optional columns to read from the parquet file.

    Yields:
        Non-empty dataframe chunks for the requested row windows.

    Raises:
        ValueError: If `chunk_rows` is less than one.

    Examples:
        Read a manifest segment that points to a row range in a shard:

        >>> slices = ((row_start, row_count),)
        >>> for chunk in store.iter_table_slices(path, slices, 100_000):
        ...     process(chunk)
    """
    if chunk_rows < 1:
        raise ValueError(f"chunk_rows must be >= 1, got {chunk_rows}")
    if not slices:
        return
    selected_columns = list(columns) if columns is not None else None
    parquet_file = pq.ParquetFile(self.resolve(location))
    row_group_starts: list[int] = []
    cursor = 0
    for idx in range(parquet_file.metadata.num_row_groups):
        row_group_starts.append(cursor)
        cursor += parquet_file.metadata.row_group(idx).num_rows
    sorted_slices = tuple(sorted(slices))
    for row_group_idx, group_start in enumerate(row_group_starts):
        group_rows = parquet_file.metadata.row_group(row_group_idx).num_rows
        group_end = group_start + group_rows
        overlaps = [
            (max(0, start - group_start), min(group_rows, start + count - group_start))
            for start, count in sorted_slices
            if count > 0 and start < group_end and start + count > group_start
        ]
        if not overlaps:
            continue
        table = parquet_file.read_row_group(row_group_idx, columns=selected_columns)
        frame = pl.from_arrow(table)
        if not isinstance(frame, pl.DataFrame):
            raise TypeError(
                f"Expected row-group conversion to return DataFrame, got {type(frame).__name__}"
            )
        for local_start, local_end in overlaps:
            yield from iter_data_chunks(
                frame.slice(local_start, local_end - local_start),
                chunk_rows,
            )