opteryx-catalog 0.4.8__py3-none-any.whl

opteryx_catalog/catalog/manifest.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+from typing import Dict
+
+NULL_FLAG = -(1 << 63)
+MIN_K_HASHES = 32
+HISTOGRAM_BINS = 32
+
+
+@dataclass
+class DataFile:
+    file_path: str
+    file_format: str = "PARQUET"
+    record_count: int = 0
+    file_size_in_bytes: int = 0
+    partition: Dict[str, object] = field(default_factory=dict)
+    lower_bounds: Dict[int, bytes] | None = None
+    upper_bounds: Dict[int, bytes] | None = None
+
+
+@dataclass
+class ManifestEntry:
+    snapshot_id: int
+    data_file: DataFile
+    status: str = "added"  # 'added' | 'deleted'
+
+
+@dataclass
+class ParquetManifestEntry:
+    """Represents a single entry in a Parquet manifest with statistics."""
+
+    file_path: str
+    file_format: str
+    record_count: int
+    file_size_in_bytes: int
+    uncompressed_size_in_bytes: int
+    min_k_hashes: list[list[int]]
+    histogram_counts: list[list[int]]
+    histogram_bins: int
+    min_values: list
+    max_values: list
+
+    def to_dict(self) -> dict:
+        return {
+            "file_path": self.file_path,
+            "file_format": self.file_format,
+            "record_count": self.record_count,
+            "file_size_in_bytes": self.file_size_in_bytes,
+            "uncompressed_size_in_bytes": self.uncompressed_size_in_bytes,
+            "min_k_hashes": self.min_k_hashes,
+            "histogram_counts": self.histogram_counts,
+            "histogram_bins": self.histogram_bins,
+            "min_values": self.min_values,
+            "max_values": self.max_values,
+        }
+
+
+def build_parquet_manifest_entry(
+    table: Any, file_path: str, file_size_in_bytes: int
+) -> ParquetManifestEntry:
+    """Build a Parquet manifest entry with statistics for a PyArrow table.
+
+    Args:
+        table: PyArrow table to analyze
+        file_path: Path where the file is stored
+        file_size_in_bytes: Size of the parquet file in bytes
+
+    Returns:
+        ParquetManifestEntry with computed statistics
+    """
+    import pyarrow as pa
+
+    min_k_hashes: list[list[int]] = []
+    histograms: list[list[int]] = []
+    min_values: list[int] = []
+    max_values: list[int] = []
+
+    # Use draken for efficient hashing and compression when available.
+    import heapq
+
+    try:
+        import opteryx.draken as draken  # type: ignore
+
+        for col_idx, col in enumerate(table.columns):
+            # hash column values to 64-bit via draken (new cpdef API)
+            vec = draken.Vector.from_arrow(col)
+            hashes = list(vec.hash())
+
+            # Decide whether to compute min-k/histogram for this column based
+            # on field type and, for strings, average length of values.
+            field_type = table.schema.field(col_idx).type
+            compute_min_k = False
+            if (
+                pa.types.is_integer(field_type)
+                or pa.types.is_floating(field_type)
+                or pa.types.is_decimal(field_type)
+            ):
+                compute_min_k = True
+            elif (
+                pa.types.is_timestamp(field_type)
+                or pa.types.is_date(field_type)
+                or pa.types.is_time(field_type)
+            ):
+                compute_min_k = True
+            elif pa.types.is_string(field_type) or pa.types.is_large_string(field_type):
+                # compute average length from non-null values; only allow
+                # min-k/histogram for short strings (avg <= 16)
+                col_py = None
+                try:
+                    col_py = col.to_pylist()
+                except Exception:
+                    col_py = None
+
+                if col_py is not None:
+                    lens = [len(x) for x in col_py if x is not None]
+                    if lens:
+                        avg_len = sum(lens) / len(lens)
+                        if avg_len <= 16:
+                            compute_min_k = True
+
+            # KMV: take K smallest hashes when allowed; otherwise store an
+            # empty list for this column.
+            if compute_min_k:
+                smallest = heapq.nsmallest(MIN_K_HASHES, hashes)
+                col_min_k = sorted(smallest)
+            else:
+                col_min_k = []
+
+            # For histogram decisions follow the same rule as min-k
+            compute_hist = compute_min_k
+
+            # Use draken.compress() to get canonical int64 per value
+            mapped = list(vec.compress())
+            non_nulls_mapped = [m for m in mapped if m != NULL_FLAG]
+            if non_nulls_mapped:
+                vmin = min(non_nulls_mapped)
+                vmax = max(non_nulls_mapped)
+                col_min = int(vmin)
+                col_max = int(vmax)
+                if compute_hist:
+                    if vmin == vmax:
+                        col_hist = [0] * HISTOGRAM_BINS
+                        col_hist[-1] = len(non_nulls_mapped)
+                    else:
+                        col_hist = [0] * HISTOGRAM_BINS
+                        span = float(vmax - vmin)
+                        for m in non_nulls_mapped:
+                            b = int(((float(m) - float(vmin)) / span) * (HISTOGRAM_BINS - 1))
+                            if b < 0:
+                                b = 0
+                            if b >= HISTOGRAM_BINS:
+                                b = HISTOGRAM_BINS - 1
+                            col_hist[b] += 1
+                else:
+                    col_hist = [0] * HISTOGRAM_BINS
+            else:
+                # no non-null values; histogram via hash buckets
+                col_min = NULL_FLAG
+                col_max = NULL_FLAG
+                if compute_hist:
+                    col_hist = [0] * HISTOGRAM_BINS
+                    for h in hashes:
+                        b = (h >> (64 - 5)) & 0x1F
+                        col_hist[b] += 1
+                else:
+                    col_hist = [0] * HISTOGRAM_BINS
+
+            min_k_hashes.append(col_min_k)
+            histograms.append(col_hist)
+            min_values.append(col_min)
+            max_values.append(col_max)
+    except Exception as exc:
+        print(f"Failed to build full manifest entry: {file_path} - {exc}")
+        return build_parquet_manifest_minmax_entry(table, file_path)
+
+    # Calculate uncompressed size from table
+    uncompressed_size = 0
+    try:
+        for col in table.columns:
+            # Estimate uncompressed size by summing the size of all columns in memory
+            for chunk in col.chunks:
+                for buffer in chunk.buffers():
+                    if buffer is not None:
+                        uncompressed_size += buffer.size
+    except Exception:
+        # Fallback: use compressed size if calculation fails
+        uncompressed_size = file_size_in_bytes
+
+    return ParquetManifestEntry(
+        file_path=file_path,
+        file_format="parquet",
+        record_count=int(table.num_rows),
+        file_size_in_bytes=file_size_in_bytes,
+        uncompressed_size_in_bytes=uncompressed_size,
+        min_k_hashes=min_k_hashes,
+        histogram_counts=histograms,
+        histogram_bins=HISTOGRAM_BINS,
+        min_values=min_values,
+        max_values=max_values,
+    )
+
+
+def build_parquet_manifest_minmax_entry(data: bytes, file_path: str) -> ParquetManifestEntry:
+    """Build a Parquet manifest entry with min/max statistics using fast rugo reader.
+
+    This is much faster than build_parquet_manifest_entry (microseconds per file)
+    and is suitable for bulk file operations where full statistics are not needed.
+
+    Args:
+        data: Raw parquet file bytes
+        file_path: Path where the file is stored
+
+    Returns:
+        ParquetManifestEntry with min/max statistics only (no histograms or k-hashes)
+    """
+    import opteryx.rugo.parquet as parquet_meta
+    from opteryx.compiled.structures.relation_statistics import to_int
+
+    file_size = len(data)
+
+    # Use rugo's fast metadata reader
+    if isinstance(data, memoryview):
+        metadata = parquet_meta.read_metadata_from_memoryview(data, include_statistics=True)
+    else:
+        metadata = parquet_meta.read_metadata_from_memoryview(
+            memoryview(data), include_statistics=True
+        )
+
+    record_count = metadata["num_rows"]
+
+    # Gather min/max per column across all row groups
+    column_stats = {}
+    for row_group in metadata["row_groups"]:
+        for column in row_group["columns"]:
+            column_name = column["name"]
+
+            if column_name not in column_stats:
+                column_stats[column_name] = {"min": None, "max": None}
+
+            min_value = column.get("min")
+            if min_value is not None:
+                # Compress value to int using to_int
+                min_compressed = to_int(min_value)
+                if column_stats[column_name]["min"] is None:
+                    column_stats[column_name]["min"] = min_compressed
+                else:
+                    column_stats[column_name]["min"] = min(
+                        column_stats[column_name]["min"], min_compressed
+                    )
+
+            max_value = column.get("max")
+            if max_value is not None:
+                # Compress value to int using to_int
+                max_compressed = to_int(max_value)
+                if column_stats[column_name]["max"] is None:
+                    column_stats[column_name]["max"] = max_compressed
+                else:
+                    column_stats[column_name]["max"] = max(
+                        column_stats[column_name]["max"], max_compressed
+                    )
+
+    # Extract min/max values (filter out None)
+    min_values = [stats["min"] for stats in column_stats.values() if stats["min"] is not None]
+    max_values = [stats["max"] for stats in column_stats.values() if stats["max"] is not None]
+
+    # Get uncompressed size from metadata
+    uncompressed_size = 0
+    try:
+        for row_group in metadata["row_groups"]:
+            uncompressed_size += row_group.get("total_byte_size", 0)
+    except Exception:
+        # Fallback to compressed size
+        uncompressed_size = file_size
+
+    return ParquetManifestEntry(
+        file_path=file_path,
+        file_format="parquet",
+        record_count=int(record_count),
+        file_size_in_bytes=file_size,
+        uncompressed_size_in_bytes=uncompressed_size,
+        min_k_hashes=[],
+        histogram_counts=[],
+        histogram_bins=0,
+        min_values=min_values,
+        max_values=max_values,
+    )

opteryx_catalog/catalog/metadata.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+from typing import List
+from typing import Optional
+
+
+@dataclass
+class Snapshot:
+    snapshot_id: int
+    timestamp_ms: int
+    author: Optional[str] = None
+    # Indicates whether this snapshot was created by a user (True) or internally (False)
+    user_created: Optional[bool] = None
+    # Monotonic sequence number for writes
+    sequence_number: Optional[int] = None
+    manifest_list: Optional[str] = None
+    # Operation metadata
+    operation_type: Optional[str] = None  # e.g., 'append', 'overwrite', 'compact'
+    parent_snapshot_id: Optional[int] = None
+    schema_id: Optional[str] = None
+    # Commit message for the snapshot
+    commit_message: Optional[str] = None
+    # Summary metrics (store zeros when not applicable)
+    summary: dict = field(
+        default_factory=lambda: {
+            "added-data-files": 0,
+            "added-files-size": 0,
+            "added-records": 0,
+            "deleted-data-files": 0,
+            "deleted-files-size": 0,
+            "deleted-records": 0,
+            "total-data-files": 0,
+            "total-files-size": 0,
+            "total-records": 0,
+        }
+    )
+
+
+@dataclass
+class DatasetMetadata:
+    dataset_identifier: str
+    format_version: int = 2
+    location: str = ""
+    schema: Any = None
+    properties: dict = field(default_factory=dict)
+    # Dataset-level created/updated metadata
+    timestamp_ms: Optional[int] = None
+    author: Optional[str] = None
+    description: Optional[str] = None
+    describer: Optional[str] = None
+    sort_orders: List[int] = field(default_factory=list)
+    # Maintenance policy: retention settings grouped under a single block
+    maintenance_policy: dict = field(
+        default_factory=lambda: {
+            "retained-snapshot-count": None,
+            "retained-snapshot-age-days": None,
+            "compaction-policy": "performance",
+        }
+    )
+    # Compaction policy lives under maintenance_policy as 'compaction-policy'
+    snapshots: List[Snapshot] = field(default_factory=list)
+    current_snapshot_id: Optional[int] = None
+    # Schema management: schemas are stored in a subcollection in Firestore.
+    # `schemas` contains dicts with keys: schema_id, columns (list of {id,name,type}).
+    # Each schema dict may also include `timestamp-ms` and `author`.
+    schemas: List[dict] = field(default_factory=list)
+    current_schema_id: Optional[str] = None
+
+    def current_snapshot(self) -> Optional[Snapshot]:
+        if self.current_snapshot_id is None:
+            return self.snapshots[-1] if self.snapshots else None
+        for s in self.snapshots:
+            if s.snapshot_id == self.current_snapshot_id:
+                return s
+        return None
+
+
+# Dataset terminology: TableMetadata renamed to DatasetMetadata

opteryx_catalog/catalog/metastore.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from typing import Any
+from typing import Iterable
+from typing import Optional
+
+
+class Metastore:
+    """Abstract catalog interface.
+
+    Implementations should provide methods to create, load and manage
+    datasets and views. Terminology in this project follows the mapping:
+    `catalog -> workspace -> collection -> dataset|view`.
+    Signatures are intentionally simple and similar to other catalog
+    implementations to ease future compatibility.
+    """
+
+    def load_dataset(self, identifier: str) -> "Dataset":
+        raise NotImplementedError()
+
+    def create_dataset(
+        self, identifier: str, schema: Any, properties: dict | None = None
+    ) -> "Dataset":
+        raise NotImplementedError()
+
+    def drop_dataset(self, identifier: str) -> None:
+        raise NotImplementedError()
+
+    def list_datasets(self, namespace: str) -> Iterable[str]:
+        raise NotImplementedError()
+
+
+class Dataset:
+    """Abstract dataset interface.
+
+    Minimal methods needed by the Opteryx engine and tests: access metadata,
+    list snapshots, append data, and produce a data scan object.
+    """
+
+    @property
+    def metadata(self) -> Any:
+        raise NotImplementedError()
+
+    def snapshots(self) -> Iterable[Any]:
+        raise NotImplementedError()
+
+    def snapshot(self, snapshot_id: Optional[int] = None) -> Optional[Any]:
+        """Return a specific snapshot by id or the current snapshot when
+        called with `snapshot_id=None`.
+        """
+        raise NotImplementedError()
+
+    def append(self, table):
+        """Append data (implementations can accept pyarrow.Table or similar)."""
+        raise NotImplementedError()
+
+    def scan(
+        self, row_filter=None, snapshot_id: Optional[int] = None, row_limit: Optional[int] = None
+    ) -> Any:
+        raise NotImplementedError()
+
+
+class View:
+    """Abstract view metadata representation."""
+
+    @property
+    def definition(self) -> str:
+        raise NotImplementedError()

opteryx_catalog/catalog/view.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class View:
+    name: str
+    definition: str
+    properties: dict | None = None
+    metadata: Any | None = None

opteryx_catalog/exceptions.py

@@ -0,0 +1,38 @@
+"""Catalog-specific exceptions for opteryx_catalog.
+
+Exceptions mirror previous behavior (they subclass KeyError where callers
+may expect KeyError) but provide explicit types for datasets, views and
+namespaces.
+"""
+
+
+class CatalogError(Exception):
+    """Base class for catalog errors."""
+
+
+class DatasetError(KeyError, CatalogError):
+    pass
+
+
+class DatasetAlreadyExists(DatasetError):
+    pass
+
+
+class DatasetNotFound(DatasetError):
+    pass
+
+
+class ViewError(KeyError, CatalogError):
+    pass
+
+
+class ViewAlreadyExists(ViewError):
+    pass
+
+
+class ViewNotFound(ViewError):
+    pass
+
+
+class CollectionAlreadyExists(KeyError, CatalogError):
+    pass

opteryx_catalog/iops/__init__.py

@@ -0,0 +1,6 @@
+from .base import FileIO
+from .base import InputFile
+from .base import OutputFile
+from .gcs import GcsFileIO
+
+__all__ = ["FileIO", "InputFile", "OutputFile", "GcsFileIO"]

opteryx_catalog/iops/base.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from io import BytesIO
+from typing import BinaryIO
+
+
+class InputFile:
+    def __init__(self, location: str, content: bytes | None = None):
+        self.location = location
+        self._content = content
+
+    def open(self) -> BinaryIO:
+        if self._content is None:
+            raise FileNotFoundError(self.location)
+        return BytesIO(self._content)
+
+
+class OutputFile:
+    def __init__(self, location: str):
+        self.location = location
+
+    def create(self):
+        """Return a file-like object with a `write` method.
+
+        Implementations may return a buffer or a writer that persists on write/close.
+        """
+        raise NotImplementedError()
+
+
+class FileIO:
+    """Minimal FileIO abstraction used by the `opteryx_catalog` layer.
+
+    Concrete implementations should implement `new_input`, `new_output`, and
+    optionally `delete`/`exists`. The abstraction intentionally keeps only the
+    small surface needed by the catalog (read bytes, write bytes).
+    """
+
+    def new_input(self, location: str) -> InputFile:
+        return InputFile(location)
+
+    def new_output(self, location: str) -> OutputFile:
+        return OutputFile(location)

opteryx_catalog/iops/fileio.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from io import BytesIO
+from typing import BinaryIO
+
+
+class InputFile:
+    def __init__(self, location: str, content: bytes | None = None):
+        self.location = location
+        self._content = content
+
+    def open(self) -> BinaryIO:
+        if self._content is None:
+            raise FileNotFoundError(self.location)
+        return BytesIO(self._content)
+
+
+class OutputFile:
+    def __init__(self, location: str):
+        self.location = location
+
+    def create(self):
+        """Return a file-like object with a `write` method.
+
+        Implementations may return a buffer or a writer that persists on write/close.
+        """
+        raise NotImplementedError()
+
+
+class FileIO:
+    """Minimal FileIO abstraction used by the `opteryx_catalog` layer.
+
+    Concrete implementations should implement `new_input`, `new_output`, and
+    optionally `delete`/`exists`. The abstraction intentionally keeps only the
+    small surface needed by the catalog (read bytes, write bytes).
+    """
+
+    def new_input(self, location: str) -> InputFile:
+        return InputFile(location)
+
+    def new_output(self, location: str) -> OutputFile:
+        return OutputFile(location)
+
+
+class _GcsAdapterOutputFile(OutputFile):
+    def __init__(self, location: str, gcs_fileio):
+        super().__init__(location)
+        self._location = location
+        self._gcs_fileio = gcs_fileio
+
+    def create(self):
+        """Return a writer whose `write(data)` uploads the data via the wrapped GCS FileIO.
+
+        We perform the upload on the first write and close the underlying stream
+        afterwards so callers that simply call `out.write(data)` (common pattern
+        in this codebase) will succeed.
+        """
+
+        class _Writer:
+            def __init__(self, location: str, gcs_fileio):
+                self._location = location
+                self._gcs_fileio = gcs_fileio
+                self._stream = None
+
+            def write(self, data: bytes | bytearray):
+                if self._stream is None:
+                    # Create underlying output stream (may be a GcsOutputStream,
+                    # DiscardOutputStream, or CaptureOutputStream depending on
+                    # the wrapped FileIO behaviour).
+                    out = self._gcs_fileio.new_output(self._location)
+                    self._stream = out.create()
+                # Underlying stream implements write/close semantics
+                self._stream.write(data)
+
+            def close(self):
+                if self._stream is not None:
+                    try:
+                        self._stream.close()
+                    except Exception:
+                        pass
+
+        return _Writer(self._location, self._gcs_fileio)
+
+
+class GcsFileIO(FileIO):
+    """GCS-backed FileIO adapter that wraps the existing GCS implementation.
+
+    This adapter delegates to `pyiceberg_firestore_gcs.fileio.gcs_fileio.GcsFileIO`
+    for actual network operations but exposes the small `opteryx_catalog.iops`
+    `FileIO` interface used by the catalog layer.
+    """
+
+    def __init__(self, properties=None):
+        # Lazy import to avoid pulling google libs unless used
+        from pyiceberg_firestore_gcs.fileio.gcs_fileio import GcsFileIO as _GcsImpl
+
+        self._impl = _GcsImpl(properties or {})
+
+    def new_input(self, location: str) -> InputFile:
+        # Read full bytes from the underlying InputFile and return an in-memory InputFile
+        impl_input = self._impl.new_input(location)
+        try:
+            stream = impl_input.open()
+            data = stream.read()
+            return InputFile(location, data)
+        except FileNotFoundError:
+            return InputFile(location, None)
+
+    def new_output(self, location: str) -> OutputFile:
+        return _GcsAdapterOutputFile(location, self._impl)
+
+    def delete(self, location: str) -> None:
+        return self._impl.delete(location)
+
+    def exists(self, location: str) -> bool:
+        try:
+            impl_in = self._impl.new_input(location)
+            # Some implementations provide `exists()`
+            if hasattr(impl_in, "exists"):
+                return impl_in.exists()
+            # Fallback: try to open
+            _ = impl_in.open()
+            return True
+        except Exception:
+            return False