pixelquery 0.1.0__py3-none-any.whl

pixelquery/__init__.py

@@ -0,0 +1,118 @@
+"""
+PixelQuery - Turn your COG files into an analysis-ready time-series data cube
+
+Zero-copy virtual references to Cloud-Optimized GeoTIFFs via Icechunk.
+
+Quick Start:
+    >>> import pixelquery as pq
+    >>>
+    >>> # Ingest COGs from a directory
+    >>> result = pq.ingest("./my_cogs/", band_names=["blue", "green", "red", "nir"])
+    >>>
+    >>> # Query as lazy xarray Dataset
+    >>> ds = pq.open_xarray("./warehouse")
+    >>> ndvi = ds.bandmath("(b3 - b2) / (b3 + b2)")  # by band index
+    >>> ndvi = ds.bandmath("(nir - red) / (nir + red)")  # by name
+    >>>
+    >>> # Point time-series
+    >>> ts = pq.timeseries("./warehouse", lon=127.05, lat=37.55)
+"""
+
+# Apply imagecodecs compatibility patch for Icechunk/VirtualTIFF
+from pixelquery._internal.codecs import patch_imagecodecs
+
+patch_imagecodecs()
+
+from pixelquery.core import (
+    DataArray,
+    # Classes
+    Dataset,
+    IngestionError,
+    # Protocols
+    PixelQuery,
+    # Exceptions
+    PixelQueryError,
+    QueryError,
+    QueryResult,
+    TransactionError,
+    ValidationError,
+    compute_evi,
+    compute_ndvi,
+    list_tiles,
+    # Functions
+    open_dataset,
+    open_mfdataset,
+    open_xarray,
+)
+
+# Legacy imports (may fail if deps not installed)
+try:
+    from pixelquery.products import BandInfo
+except ImportError:
+    BandInfo = None  # type: ignore[misc, assignment]
+
+try:
+    from pixelquery.grid import TileGrid
+except ImportError:
+    TileGrid = None  # type: ignore[misc, assignment]
+
+# Register xarray BandMath accessor (ds.bandmath.ndvi(), etc.)
+import pixelquery.core.bandmath
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "DataArray",
+    "Dataset",
+    "IngestionError",
+    "PixelQueryError",
+    "QueryError",
+    "TransactionError",
+    "ValidationError",
+    "__version__",
+    "catalog",
+    "compute_evi",
+    "compute_ndvi",
+    "ingest",
+    "inspect_cog",
+    "inspect_directory",
+    "list_tiles",
+    "open_dataset",
+    "open_mfdataset",
+    "open_xarray",
+    "register_product",
+    "timeseries",
+]
+
+
+# Lazy imports for new Icechunk features (avoids heavy import at startup)
+def __getattr__(name):
+    if name == "ingest":
+        from pixelquery.io.auto_ingest import ingest
+
+        return ingest
+    elif name == "timeseries":
+        from pixelquery.core.timeseries import timeseries
+
+        return timeseries
+    elif name == "inspect_cog":
+        from pixelquery.io.cog_metadata import inspect_cog
+
+        return inspect_cog
+    elif name == "inspect_directory":
+        from pixelquery.io.cog_metadata import inspect_directory
+
+        return inspect_directory
+    elif name == "catalog":
+        from pixelquery.catalog.icechunk_catalog import IcechunkCatalog
+
+        return IcechunkCatalog
+    elif name == "register_product":
+        from pixelquery.catalog.product_profile import register_product
+
+        return register_product
+    elif name == "ProductProfile":
+        from pixelquery.catalog.product_profile import ProductProfile
+
+        return ProductProfile
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

pixelquery/_internal/__init__.py

@@ -0,0 +1,11 @@
+"""
+PixelQuery Internal Module
+
+⚠️ PRIVATE API - Do not use directly!
+
+This module contains internal implementation details.
+The API may change without notice.
+"""
+
+# No exports - internal use only
+__all__: list[str] = []

pixelquery/_internal/codecs.py

@@ -0,0 +1,77 @@
+"""
+Codec compatibility layer for Icechunk/VirtualTIFF.
+
+Fixes imagecodecs numcodecs bug where from_config() passes
+'name'/'configuration' keys that __init__() doesn't accept.
+Applied once at package import time.
+"""
+
+import inspect
+import logging
+
+logger = logging.getLogger(__name__)
+_PATCHED = False
+
+
+def patch_imagecodecs():
+    """
+    Monkeypatch all imagecodecs numcodecs to strip extra keys from from_config().
+
+    Safe to call multiple times; only patches once.
+    """
+    global _PATCHED
+    if _PATCHED:
+        return
+
+    try:
+        import imagecodecs.numcodecs as ic_numcodecs
+    except ImportError:
+        logger.debug("imagecodecs.numcodecs not available, skipping codec patch")
+        return
+
+    patched_count = 0
+    for name in dir(ic_numcodecs):
+        cls = getattr(ic_numcodecs, name)
+        if not isinstance(cls, type) or not hasattr(cls, "from_config"):
+            continue
+
+        try:
+            init_params = set(inspect.signature(cls.__init__).parameters.keys()) - {"self"}  # type: ignore[misc]
+        except (ValueError, TypeError):
+            init_params = set()
+
+        def _make_patched(valid_params):
+            @classmethod  # type: ignore[misc]
+            def patched_from_config(klass, config):
+                config = dict(config)
+                for key in ("id", "name", "configuration"):
+                    config.pop(key, None)
+                if valid_params:
+                    config = {k: v for k, v in config.items() if k in valid_params}
+                return klass(**config)
+
+            return patched_from_config
+
+        cls.from_config = _make_patched(init_params)
+        patched_count += 1
+
+    # Register virtual_tiff codecs under their short names for zarr v3.
+    # Entry points register them as "virtual_tiff.ChunkyCodec" but zarr
+    # metadata stores just "ChunkyCodec".
+    try:
+        import zarr.registry
+        from virtual_tiff.codecs import ChunkyCodec, HorizontalDeltaCodec
+
+        for codec_cls in (ChunkyCodec, HorizontalDeltaCodec):
+            codec_name = codec_cls.__name__
+            try:
+                zarr.registry.get_codec_class(codec_name)
+            except KeyError:
+                zarr.registry.register_codec(codec_name, codec_cls)
+                patched_count += 1
+    except ImportError:
+        pass
+
+    _PATCHED = True
+    if patched_count:
+        logger.debug("Patched %d codecs for imagecodecs/virtual_tiff compatibility", patched_count)

pixelquery/_internal/storage/__init__.py

@@ -0,0 +1,7 @@
+"""
+Internal Storage Module
+
+⚠️ PRIVATE API - Do not use directly!
+"""
+
+__all__: list[str] = []

pixelquery/_internal/storage/arrow_chunk.py

@@ -0,0 +1,412 @@
+"""
+Arrow IPC Chunk Storage
+
+Implements monthly spatiotemporal chunk storage using Arrow IPC format.
+
+NOTE: With Iceberg integration, this module is retained for backwards compatibility
+with existing Arrow-based warehouses. New warehouses should use Iceberg storage
+via IcebergStorageManager.
+
+Performance: Uses Rust extensions when available for 2-35x faster I/O.
+"""
+
+from datetime import UTC, datetime
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.ipc as ipc
+from numpy.typing import NDArray
+
+# Try to use Rust Arrow functions (2-35x faster than Python)
+try:
+    from pixelquery_core import (  # type: ignore[attr-defined]
+        arrow_append_to_chunk,
+        arrow_write_chunk,
+    )
+
+    RUST_ARROW_AVAILABLE = True
+except ImportError:
+    RUST_ARROW_AVAILABLE = False
+
+
+class ArrowChunkWriter:
+    """
+    Arrow IPC chunk writer for monthly spatiotemporal data
+
+    Each chunk represents a single tile-month-band combination:
+    - Tile: Geographic tile (e.g., x0024_y0041)
+    - Month: Temporal partition (e.g., 2024-01)
+    - Band: Spectral band (e.g., red, nir)
+
+    File naming convention:
+        {warehouse}/{table}/data/{tile_id}/{year}-{month:02d}/{band}.arrow
+
+    Arrow Schema:
+        - time: timestamp[ms] - Observation timestamps
+        - pixels: list<uint16> - Variable-length pixel arrays (multi-resolution support)
+        - mask: list<bool> - Cloud/invalid pixel masks
+        - metadata: map<string, string> - Additional metadata
+
+    Examples:
+        >>> writer = ArrowChunkWriter()
+        >>> data = {
+        ...     'time': [datetime(2024, 1, 1), datetime(2024, 1, 15)],
+        ...     'pixels': [np.array([1000, 1100, ...]), np.array([1050, 1150, ...])],
+        ...     'mask': [np.array([False, False, ...]), np.array([True, False, ...])]
+        ... }
+        >>> writer.write_chunk(
+        ...     path="warehouse/table/data/x0024_y0041/2024-01/red.arrow",
+        ...     data=data,
+        ...     product_id="sentinel2_l2a",
+        ...     resolution=10.0
+        ... )
+    """
+
+    # Arrow schema for spatiotemporal chunks
+    SCHEMA = pa.schema(
+        [
+            ("time", pa.timestamp("ms", tz="UTC")),
+            ("pixels", pa.list_(pa.uint16())),  # Variable-length arrays
+            ("mask", pa.list_(pa.bool_())),  # Cloud/invalid mask
+        ]
+    )
+
+    def write_chunk(
+        self,
+        path: str,
+        data: dict[str, list],
+        product_id: str,
+        resolution: float,
+        metadata: dict[str, str] | None = None,
+    ) -> None:
+        """
+        Write spatiotemporal chunk to Arrow IPC file
+
+        Args:
+            path: Output file path
+            data: Dictionary with keys 'time', 'pixels', 'mask'
+            product_id: Product identifier (e.g., "sentinel2_l2a")
+            resolution: Spatial resolution in meters
+            metadata: Additional metadata
+
+        Raises:
+            ValueError: If data is invalid
+            IOError: If write fails
+        """
+        # Validate input data
+        self._validate_data(data)
+
+        # Use Rust implementation if available (2-35x faster)
+        if RUST_ARROW_AVAILABLE:
+            self._write_chunk_rust(path, data, product_id, resolution, metadata)
+            return
+
+        # Fall back to Python implementation
+        # Convert to Arrow arrays
+        time_array = pa.array(data["time"], type=pa.timestamp("ms", tz="UTC"))
+        pixels_array = self._convert_pixels_to_arrow(data["pixels"])
+        mask_array = self._convert_mask_to_arrow(data["mask"])
+
+        # Create record batch
+        batch = pa.RecordBatch.from_arrays(
+            [time_array, pixels_array, mask_array], schema=self.SCHEMA
+        )
+
+        # Prepare metadata
+        chunk_metadata = {
+            "product_id": product_id,
+            "resolution": str(resolution),
+            "num_observations": str(len(data["time"])),
+            "creation_time": datetime.now(UTC).isoformat(),
+        }
+        if metadata:
+            chunk_metadata.update(metadata)
+
+        # Create schema with metadata
+        schema_with_metadata = self.SCHEMA.with_metadata(chunk_metadata)
+
+        # Write to file
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+        with pa.OSFile(path, "wb") as sink, ipc.new_file(sink, schema_with_metadata) as writer:
+            writer.write_batch(batch)
+
+    def _validate_data(self, data: dict[str, list]) -> None:
+        """Validate input data structure"""
+        required_keys = {"time", "pixels", "mask"}
+        if not required_keys.issubset(data.keys()):
+            missing = required_keys - data.keys()
+            raise ValueError(f"Missing required keys: {missing}")
+
+        n_obs = len(data["time"])
+        if len(data["pixels"]) != n_obs:
+            raise ValueError(f"pixels length ({len(data['pixels'])}) != time length ({n_obs})")
+        if len(data["mask"]) != n_obs:
+            raise ValueError(f"mask length ({len(data['mask'])}) != time length ({n_obs})")
+
+    def _convert_pixels_to_arrow(self, pixels: list[NDArray]) -> pa.Array:
+        """Convert list of numpy arrays to Arrow list array"""
+        # Convert each numpy array to list
+        pixels_list = [arr.astype(np.uint16).tolist() for arr in pixels]
+        return pa.array(pixels_list, type=pa.list_(pa.uint16()))
+
+    def _convert_mask_to_arrow(self, masks: list[NDArray]) -> pa.Array:
+        """Convert list of mask arrays to Arrow list array"""
+        masks_list = [arr.astype(bool).tolist() for arr in masks]
+        return pa.array(masks_list, type=pa.list_(pa.bool_()))
+
+    def _write_chunk_rust(
+        self,
+        path: str,
+        data: dict[str, list],
+        product_id: str,
+        resolution: float,
+        metadata: dict[str, str] | None = None,
+    ) -> None:
+        """Write chunk using Rust implementation (2-35x faster)"""
+        # Convert timestamps to milliseconds since epoch
+        times_ms = [int(t.timestamp() * 1000) for t in data["time"]]
+
+        # Convert numpy arrays to lists
+        pixels_list = [arr.astype(np.uint16).flatten().tolist() for arr in data["pixels"]]
+        masks_list = [arr.astype(bool).flatten().tolist() for arr in data["mask"]]
+
+        # Prepare metadata
+        rust_metadata = {
+            "product_id": product_id,
+            "resolution": str(resolution),
+        }
+        if metadata:
+            rust_metadata.update(metadata)
+
+        # Call Rust function
+        arrow_write_chunk(path, times_ms, pixels_list, masks_list, rust_metadata)
+
+    def _append_to_chunk_rust(
+        self,
+        path: str,
+        data: dict[str, list],
+        product_id: str,
+        resolution: float,
+        metadata: dict[str, str] | None = None,
+    ) -> None:
+        """Append to chunk using Rust implementation (2-3x faster)"""
+        # Convert timestamps to milliseconds since epoch
+        times_ms = [int(t.timestamp() * 1000) for t in data["time"]]
+
+        # Convert numpy arrays to lists
+        pixels_list = [arr.astype(np.uint16).flatten().tolist() for arr in data["pixels"]]
+        masks_list = [arr.astype(bool).flatten().tolist() for arr in data["mask"]]
+
+        # Prepare metadata
+        rust_metadata = {
+            "product_id": product_id,
+            "resolution": str(resolution),
+        }
+        if metadata:
+            rust_metadata.update(metadata)
+
+        # Call Rust function
+        arrow_append_to_chunk(path, times_ms, pixels_list, masks_list, rust_metadata)
+
+    def append_to_chunk(
+        self,
+        path: str,
+        data: dict[str, list],
+        product_id: str,
+        resolution: float,
+        metadata: dict[str, str] | None = None,
+    ) -> None:
+        """
+        Append new observations to existing chunk file
+
+        This is more efficient than reading the entire file, concatenating in Python,
+        and rewriting. Uses Rust implementation when available (2-3x faster).
+
+        Args:
+            path: Chunk file path
+            data: Dictionary with keys 'time', 'pixels', 'mask' (single observation each)
+            product_id: Product identifier
+            resolution: Spatial resolution in meters
+            metadata: Additional metadata
+
+        Raises:
+            ValueError: If data is invalid
+            IOError: If read/write fails
+        """
+        path_obj = Path(path)
+
+        # Validate input data
+        self._validate_data(data)
+
+        # Use Rust implementation if available (2-3x faster)
+        if RUST_ARROW_AVAILABLE:
+            self._append_to_chunk_rust(path, data, product_id, resolution, metadata)
+            return
+
+        if not path_obj.exists():
+            # First write - use regular write_chunk
+            self.write_chunk(path, data, product_id, resolution, metadata)
+            return
+
+        # Read existing data as Arrow table (avoid Python conversion)
+        with pa.OSFile(str(path), "rb") as source, ipc.open_file(source) as reader:
+            existing_table = reader.read_all()
+            # Preserve existing metadata
+            existing_metadata = dict(reader.schema.metadata) if reader.schema.metadata else {}
+
+        # Convert new data to Arrow arrays
+        new_time_array = pa.array(data["time"], type=pa.timestamp("ms", tz="UTC"))
+        new_pixels_array = self._convert_pixels_to_arrow(data["pixels"])
+        new_mask_array = self._convert_mask_to_arrow(data["mask"])
+
+        # Create record batch for new data
+        new_batch = pa.RecordBatch.from_arrays(
+            [new_time_array, new_pixels_array, new_mask_array], schema=self.SCHEMA
+        )
+
+        # Convert to table and concatenate with existing
+        new_table = pa.Table.from_batches([new_batch])
+        combined_table = pa.concat_tables([existing_table, new_table])
+
+        # Update metadata
+        chunk_metadata = {
+            k.decode() if isinstance(k, bytes) else k: v.decode() if isinstance(v, bytes) else v
+            for k, v in existing_metadata.items()
+        }
+        chunk_metadata.update(
+            {
+                "product_id": product_id,
+                "resolution": str(resolution),
+                "num_observations": str(len(combined_table)),
+                "last_updated": datetime.now(UTC).isoformat(),
+            }
+        )
+        if metadata:
+            chunk_metadata.update(metadata)
+
+        # Create schema with updated metadata
+        schema_with_metadata = self.SCHEMA.with_metadata(chunk_metadata)
+
+        # Write combined data atomically
+        # Use temporary file + rename for atomic write
+        temp_path = str(path_obj) + ".tmp"
+        try:
+            with (
+                pa.OSFile(temp_path, "wb") as sink,
+                ipc.new_file(sink, schema_with_metadata) as writer,
+            ):
+                writer.write_table(combined_table)
+
+            # Atomic rename
+            import os
+
+            os.replace(temp_path, str(path_obj))
+        except Exception as e:
+            # Clean up temp file on error
+            if Path(temp_path).exists():
+                Path(temp_path).unlink()
+            raise e
+
+
+class ArrowChunkReader:
+    """
+    Arrow IPC chunk reader for monthly spatiotemporal data
+
+    Reads chunks written by ArrowChunkWriter and reconstructs
+    spatiotemporal data.
+
+    Examples:
+        >>> reader = ArrowChunkReader()
+        >>> data, metadata = reader.read_chunk(
+        ...     "warehouse/table/data/x0024_y0041/2024-01/red.arrow"
+        ... )
+        >>> print(len(data['time']))  # Number of observations
+        2
+        >>> print(data['pixels'][0].shape)  # First observation pixel array
+        (256, 256)
+    """
+
+    def read_chunk(
+        self, path: str, reshape: tuple[int, int] | None = None
+    ) -> tuple[dict[str, list], dict[str, str]]:
+        """
+        Read spatiotemporal chunk from Arrow IPC file
+
+        Args:
+            path: Input file path
+            reshape: Optional (height, width) to reshape pixel arrays
+
+        Returns:
+            Tuple of (data, metadata):
+            - data: Dictionary with 'time', 'pixels', 'mask'
+            - metadata: Chunk metadata (product_id, resolution, etc.)
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+            IOError: If read fails
+        """
+        if not Path(path).exists():
+            raise FileNotFoundError(f"Chunk file not found: {path}")
+
+        # Read Arrow file
+        with pa.OSFile(path, "rb") as source, ipc.open_file(source) as reader:
+            # Get metadata from schema
+            metadata = dict(reader.schema.metadata) if reader.schema.metadata else {}
+            # Decode bytes to strings
+            metadata = {k.decode(): v.decode() for k, v in metadata.items()}
+
+            # Read all batches (typically just one per chunk)
+            table = reader.read_all()
+
+        # Convert to Python objects
+        data = {
+            "time": table["time"].to_pylist(),
+            "pixels": self._convert_pixels_from_arrow(table["pixels"], reshape),
+            "mask": self._convert_mask_from_arrow(table["mask"], reshape),
+        }
+
+        return data, metadata
+
+    def _convert_pixels_from_arrow(
+        self, arrow_array: pa.Array, reshape: tuple[int, int] | None
+    ) -> list[NDArray]:
+        """Convert Arrow list array to list of numpy arrays"""
+        pixels_list = arrow_array.to_pylist()
+        arrays = [np.array(pixels, dtype=np.uint16) for pixels in pixels_list]
+
+        if reshape:
+            height, width = reshape
+            arrays = [arr.reshape(height, width) for arr in arrays]
+
+        return arrays
+
+    def _convert_mask_from_arrow(
+        self, arrow_array: pa.Array, reshape: tuple[int, int] | None
+    ) -> list[NDArray]:
+        """Convert Arrow list array to list of mask arrays"""
+        mask_list = arrow_array.to_pylist()
+        arrays = [np.array(mask, dtype=bool) for mask in mask_list]
+
+        if reshape:
+            height, width = reshape
+            arrays = [arr.reshape(height, width) for arr in arrays]
+
+        return arrays
+
+    def read_chunk_metadata(self, path: str) -> dict[str, str]:
+        """
+        Read only chunk metadata without loading data
+
+        Args:
+            path: Input file path
+
+        Returns:
+            Chunk metadata dictionary
+        """
+        if not Path(path).exists():
+            raise FileNotFoundError(f"Chunk file not found: {path}")
+
+        with pa.OSFile(path, "rb") as source, ipc.open_file(source) as reader:
+            metadata = dict(reader.schema.metadata) if reader.schema.metadata else {}
+            return {k.decode(): v.decode() for k, v in metadata.items()}

pixelquery/_internal/storage/base.py

@@ -0,0 +1,44 @@
+"""
+Storage Backend Protocol
+
+Abstract storage interface for multiple backends (local, S3, Azure, GCS).
+"""
+
+from typing import Protocol
+
+
+class StorageBackend(Protocol):
+    """
+    Abstract storage interface
+
+    Enables multiple storage backends (local filesystem, S3, Azure, GCS)
+    while maintaining consistent API.
+    """
+
+    def read_bytes(self, path: str) -> bytes:
+        """Read file contents as bytes"""
+        ...
+
+    def write_bytes(self, path: str, data: bytes) -> None:
+        """Write bytes to file"""
+        ...
+
+    def atomic_rename(self, src: str, dest: str) -> None:
+        """
+        Atomically rename file (critical for transactions)
+
+        This is used for two-phase commit: write to .tmp, then rename.
+        """
+        ...
+
+    def delete(self, path: str) -> None:
+        """Delete file"""
+        ...
+
+    def exists(self, path: str) -> bool:
+        """Check if file exists"""
+        ...
+
+    def list_files(self, prefix: str) -> list[str]:
+        """List files matching prefix (for discovery)"""
+        ...