valor-lite 0.36.6__py3-none-any.whl → 0.37.5__py3-none-any.whl

valor_lite/cache/__init__.py

@@ -0,0 +1,11 @@
+from .compute import sort
+from .ephemeral import MemoryCacheReader, MemoryCacheWriter
+from .persistent import FileCacheReader, FileCacheWriter
+
+__all__ = [
+    "FileCacheReader",
+    "FileCacheWriter",
+    "MemoryCacheReader",
+    "MemoryCacheWriter",
+    "sort",
+]

valor_lite/cache/compute.py

@@ -0,0 +1,211 @@
+import heapq
+import tempfile
+from pathlib import Path
+from typing import Callable, Generator
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from valor_lite.cache.ephemeral import MemoryCacheReader, MemoryCacheWriter
+from valor_lite.cache.persistent import FileCacheReader, FileCacheWriter
+
+
+def _merge(
+    source: MemoryCacheReader | FileCacheReader,
+    sink: MemoryCacheWriter | FileCacheWriter,
+    intermediate_sink: MemoryCacheWriter | FileCacheWriter,
+    batch_size: int,
+    sorting: list[tuple[str, str]],
+    columns: list[str] | None = None,
+    table_sort_override: Callable[[pa.Table], pa.Table] | None = None,
+):
+    """Merge locally sorted cache fragments."""
+    for tbl in source.iterate_tables(columns=columns):
+        if table_sort_override is not None:
+            sorted_tbl = table_sort_override(tbl)
+        else:
+            sorted_tbl = tbl.sort_by(sorting)
+        intermediate_sink.write_table(sorted_tbl)
+    intermediate_source = intermediate_sink.to_reader()
+
+    # define merge key
+    def create_sort_key(
+        batches: list[pa.RecordBatch],
+        batch_idx: int,
+        row_idx: int,
+    ):
+        args = [
+            -batches[batch_idx][name][row_idx].as_py()
+            if direction == "descending"
+            else batches[batch_idx][name][row_idx].as_py()
+            for name, direction in sorting
+        ]
+        return (
+            *args,
+            batch_idx,
+            row_idx,
+        )
+
+    # merge sorted rows
+    heap = []
+    batch_iterators = []
+    batches = []
+    for batch_idx, batch_iter in enumerate(
+        intermediate_source.iterate_fragment_batch_iterators(
+            batch_size=batch_size
+        )
+    ):
+        batch_iterators.append(batch_iter)
+        batches.append(next(batch_iterators[batch_idx], None))
+        if batches[batch_idx] is not None and len(batches[batch_idx]) > 0:
+            heap.append(create_sort_key(batches, batch_idx, 0))
+    heapq.heapify(heap)
+
+    while heap:
+        row = heapq.heappop(heap)
+        batch_idx = row[-2]
+        row_idx = row[-1]
+        row_table = batches[batch_idx].slice(row_idx, 1)
+        sink.write_batch(row_table)
+        row_idx += 1
+        if row_idx < len(batches[batch_idx]):
+            heapq.heappush(
+                heap,
+                create_sort_key(batches, batch_idx, row_idx),
+            )
+        else:
+            batches[batch_idx] = next(batch_iterators[batch_idx], None)
+            if batches[batch_idx] is not None and len(batches[batch_idx]) > 0:
+                heapq.heappush(
+                    heap,
+                    create_sort_key(batches, batch_idx, 0),
+                )
+
+    sink.flush()
+
+
+def sort(
+    source: MemoryCacheReader | FileCacheReader,
+    sink: MemoryCacheWriter | FileCacheWriter,
+    batch_size: int,
+    sorting: list[tuple[str, str]],
+    columns: list[str] | None = None,
+    table_sort_override: Callable[[pa.Table], pa.Table] | None = None,
+):
+    """
+    Sort data into new cache.
+
+    Parameters
+    ----------
+    source : MemoryCacheReader | FileCacheReader
+        A read-only cache. If file-based, each file must be locally sorted.
+    sink : MemoryCacheWriter | FileCacheWriter
+        The cache where sorted data will be written.
+    batch_size : int
+        Maximum number of rows allowed to be read into memory per cache file.
+    sorting : list[tuple[str, str]]
+        Sorting arguments in PyArrow format (e.g. [('a', 'ascending'), ('b', 'descending')]).
+        Note that only numeric fields are supported.
+    columns : list[str], optional
+        Option to only read a subset of columns.
+    table_sort_override : Callable[[pa.Table], pa.Table], optional
+        Option to override sort function for singular cache fragments.
+    """
+
+    if source.count_tables() == 1:
+        for tbl in source.iterate_tables(columns=columns):
+            if table_sort_override is not None:
+                sorted_tbl = table_sort_override(tbl)
+            else:
+                sorted_tbl = tbl.sort_by(sorting)
+            sink.write_table(sorted_tbl)
+        sink.flush()
+        return
+
+    if isinstance(sink, FileCacheWriter):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            intermediate_sink = FileCacheWriter.create(
+                path=Path(tmpdir) / "sorting_intermediate",
+                schema=sink.schema,
+                batch_size=sink.batch_size,
+                rows_per_file=sink.rows_per_file,
+                compression=sink.compression,
+                delete_if_exists=False,
+            )
+            _merge(
+                source=source,
+                sink=sink,
+                intermediate_sink=intermediate_sink,
+                batch_size=batch_size,
+                sorting=sorting,
+                columns=columns,
+                table_sort_override=table_sort_override,
+            )
+    else:
+        intermediate_sink = MemoryCacheWriter.create(
+            schema=sink.schema,
+            batch_size=sink.batch_size,
+        )
+        _merge(
+            source=source,
+            sink=sink,
+            intermediate_sink=intermediate_sink,
+            batch_size=batch_size,
+            sorting=sorting,
+            columns=columns,
+            table_sort_override=table_sort_override,
+        )
+
+
+def paginate_index(
+    source: MemoryCacheReader | FileCacheReader,
+    column_key: str,
+    modifier: pc.Expression | None = None,
+    limit: int | None = None,
+    offset: int = 0,
+) -> Generator[pa.Table, None, None]:
+    """
+    Iterate through a paginated cache reader.
+
+    Note this function expects unqiue keys to be fragment-aligned.
+    """
+    total = source.count_rows()
+    limit = limit if limit else total
+
+    # pagination broader than data scope
+    if offset == 0 and limit >= total:
+        for tbl in source.iterate_tables(filter=modifier):
+            yield tbl
+        return
+    elif offset >= total:
+        return
+
+    curr_idx = 0
+    for tbl in source.iterate_tables(filter=modifier):
+        if tbl.num_rows == 0:
+            continue
+
+        # sort the unique keys as they may be out of order
+        unique_values = pc.unique(tbl[column_key]).sort()  # type: ignore[reportAttributeAccessIssue]
+        n_unique = len(unique_values)
+        prev_idx = curr_idx
+        curr_idx += n_unique
+
+        # check for page overlap
+        if curr_idx <= offset:
+            continue
+        elif prev_idx >= (offset + limit):
+            return
+
+        # apply any pagination conditions
+        condition = pc.scalar(True)
+        if prev_idx < offset and curr_idx > offset:
+            condition &= (
+                pc.field(column_key) >= unique_values[offset - prev_idx]
+            )
+        if prev_idx < (offset + limit) and curr_idx > (offset + limit):
+            condition &= (
+                pc.field(column_key) < unique_values[offset + limit - prev_idx]
+            )
+
+        yield tbl.filter(condition)

valor_lite/cache/ephemeral.py

@@ -0,0 +1,302 @@
+from collections.abc import Iterator
+from typing import Any
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+
+
+class MemoryCache:
+    def __init__(
+        self,
+        table: pa.Table,
+        batch_size: int,
+    ):
+        self._table = table
+        self._batch_size = batch_size
+
+    @property
+    def schema(self) -> pa.Schema:
+        return self._table.schema
+
+    @property
+    def batch_size(self) -> int:
+        return self._batch_size
+
+    def count_tables(self) -> int:
+        """Count the number of tables in the cache."""
+        return 1
+
+    def count_rows(self) -> int:
+        """Count the number of rows in the cache."""
+        return self._table.num_rows
+
+
+class MemoryCacheReader(MemoryCache):
+    def iterate_tables(
+        self,
+        columns: list[str] | None = None,
+        filter: pc.Expression | None = None,
+    ) -> Iterator[pa.Table]:
+        """
+        Iterate over tables within the cache.
+
+        Parameters
+        ----------
+        columns : list[str], optional
+            Optionally select columns to be returned.
+        filter : pyarrow.compute.Expression, optional
+            Optionally filter table before returning.
+
+        Returns
+        -------
+        Iterator[pa.Table]
+        """
+        table = self._table
+        if filter is not None:
+            table = table.filter(filter)
+        if columns is not None:
+            table = table.select(columns)
+        yield table
+
+    def iterate_arrays(
+        self,
+        numeric_columns: list[str] | None = None,
+        filter: pc.Expression | None = None,
+    ) -> Iterator[np.ndarray]:
+        """
+        Iterate over chunks within the cache returning arrays.
+
+        Parameters
+        ----------
+        numeric_columns : list[str], optional
+            Optionally select numeric columns to be returned within an array.
+        filter : pyarrow.compute.Expression, optional
+            Optionally filter table before returning.
+
+        Returns
+        -------
+        Iterator[np.ndarray]
+        """
+        for tbl in self.iterate_tables(columns=numeric_columns, filter=filter):
+            yield np.column_stack(
+                [tbl.column(i).to_numpy() for i in range(tbl.num_columns)]
+            )
+
+    def iterate_tables_with_arrays(
+        self,
+        columns: list[str] | None = None,
+        filter: pc.Expression | None = None,
+        numeric_columns: list[str] | None = None,
+    ) -> Iterator[tuple[pa.Table, np.ndarray]]:
+        """
+        Iterate over chunks within the cache returning both tables and arrays.
+
+        Parameters
+        ----------
+        columns : list[str], optional
+            Optionally select columns to be returned.
+        filter : pyarrow.compute.Expression, optional
+            Optionally filter table before returning.
+        numeric_columns : list[str], optional
+            Optionally select numeric columns to be returned within an array.
+
+        Returns
+        -------
+        Iterator[tuple[pa.Table, np.ndarray]]
+
+        """
+        _columns = set(columns) if columns else set()
+        _numeric_columns = set(numeric_columns) if numeric_columns else set()
+        columns = list(_columns.union(_numeric_columns))
+        for tbl in self.iterate_tables(
+            columns=columns,
+            filter=filter,
+        ):
+            table_columns = numeric_columns if numeric_columns else tbl.columns
+            yield tbl, np.column_stack(
+                [tbl[col].to_numpy() for col in table_columns]
+            )
+
+    def iterate_fragment_batch_iterators(
+        self, batch_size: int
+    ) -> Iterator[Iterator[pa.RecordBatch]]:
+        """
+        Yield a table batch iterator.
+
+        This is intended to emulate file-based access patterns.
+
+        Parameters
+        ----------
+        batch_size : int
+            Maximum number of rows allowed to be read per batch.
+
+        Yields
+        ------
+        Iterator[Iterator[pa.RecordBatch]]
+        """
+        yield iter(self._table.to_batches(max_chunksize=batch_size))
+
+
+class MemoryCacheWriter(MemoryCache):
+    def __init__(
+        self,
+        table: pa.Table,
+        batch_size: int,
+    ):
+        super().__init__(
+            table=table,
+            batch_size=batch_size,
+        )
+
+        # internal state
+        self._buffer = []
+
+    @classmethod
+    def create(
+        cls,
+        schema: pa.Schema,
+        batch_size: int,
+    ):
+        """
+        Create an in-memory cache.
+
+        Parameters
+        ----------
+        schema : pa.Schema
+            Cache schema.
+        batch_size : int
+            Target batch size when writing chunks.
+        """
+        return cls(
+            table=schema.empty_table(),
+            batch_size=batch_size,
+        )
+
+    def write_rows(
+        self,
+        rows: list[dict[str, Any]],
+    ):
+        """
+        Write rows to cache.
+
+        Parameters
+        ----------
+        rows : list[dict[str, Any]]
+            A list of rows represented by dictionaries mapping fields to values.
+        """
+        if not rows:
+            return
+        batch = pa.RecordBatch.from_pylist(rows, schema=self.schema)
+        self.write_batch(batch)
+
+    def write_columns(
+        self,
+        columns: dict[str, list | np.ndarray | pa.Array],
+    ):
+        """
+        Write columnar data to cache.
+
+        Parameters
+        ----------
+        columns : dict[str, list | np.ndarray | pa.Array]
+            A mapping of columnar field names to list of values.
+        """
+        if not columns:
+            return
+        batch = pa.RecordBatch.from_pydict(columns)
+        self.write_batch(batch)
+
+    def write_batch(
+        self,
+        batch: pa.RecordBatch,
+    ):
+        """
+        Write a batch to cache.
+
+        Parameters
+        ----------
+        batch : pa.RecordBatch
+            A batch of columnar data.
+        """
+        size = batch.num_rows
+        if self._buffer:
+            size += sum([b.num_rows for b in self._buffer])
+
+        # check size
+        if size < self._batch_size:
+            self._buffer.append(batch)
+            return
+
+        if self._buffer:
+            self._buffer.append(batch)
+            combined_arrays = [
+                pa.concat_arrays([b.column(name) for b in self._buffer])
+                for name in self.schema.names
+            ]
+            batch = pa.RecordBatch.from_arrays(
+                combined_arrays, schema=self.schema
+            )
+            self._buffer = []
+
+        # write batch
+        self.write_table(pa.Table.from_batches([batch]))
+
+    def write_table(
+        self,
+        table: pa.Table,
+    ):
+        """
+        Write a table directly to cache.
+
+        Parameters
+        ----------
+        table : pa.Table
+            A populated table.
+        """
+        self._table = pa.concat_tables([self._table, table])
+
+    def flush(self):
+        """Flush the cache buffer."""
+        if self._buffer:
+            combined_arrays = [
+                pa.concat_arrays([b.column(name) for b in self._buffer])
+                for name in self.schema.names
+            ]
+            batch = pa.RecordBatch.from_arrays(
+                combined_arrays, schema=self.schema
+            )
+            self._table = pa.concat_tables(
+                [self._table, pa.Table.from_batches([batch])]
+            )
+            self._buffer = []
+
+    def sort_by(
+        self,
+        sorting: list[tuple[str, str]],
+    ):
+        """
+        Sort cache in-place.
+
+        Parameters
+        ----------
+        sorting : list[tuple[str, str]]
+            Sorting arguments in PyArrow format (e.g. [('a', 'ascending'), ('b', 'descending')]).
+        """
+        self.flush()
+        self._table = self._table.sort_by(sorting)
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures data is flushed."""
+        self.flush()
+
+    def to_reader(self) -> MemoryCacheReader:
+        """Get cache reader."""
+        self.flush()
+        return MemoryCacheReader(
+            table=self._table, batch_size=self._batch_size
+        )