metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/metadata_store/lancedb.py

@@ -0,0 +1,391 @@
+"""LanceDB metadata store implementation."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import narwhals as nw
+import polars as pl
+from narwhals.typing import Frame
+from pydantic import Field
+from typing_extensions import Self
+
+from metaxy._utils import collect_to_polars
+from metaxy.metadata_store.base import MetadataStore, MetadataStoreConfig
+from metaxy.metadata_store.types import AccessMode
+from metaxy.metadata_store.utils import is_local_path, sanitize_uri
+from metaxy.models.types import CoercibleToFeatureKey, FeatureKey
+from metaxy.versioning.polars import PolarsVersioningEngine
+from metaxy.versioning.types import HashAlgorithm
+
+logger = logging.getLogger(__name__)
+
+
+class LanceDBMetadataStoreConfig(MetadataStoreConfig):
+    """Configuration for LanceDBMetadataStore.
+
+    Example:
+        ```python
+        config = LanceDBMetadataStoreConfig(
+            uri="/path/to/featuregraph",
+            connect_kwargs={"api_key": "your-api-key"},
+        )
+
+        store = LanceDBMetadataStore.from_config(config)
+        ```
+    """
+
+    uri: str | Path = Field(
+        description="Directory path or URI for LanceDB tables.",
+    )
+    connect_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        description="Extra keyword arguments passed to lancedb.connect().",
+    )
+
+
+class LanceDBMetadataStore(MetadataStore):
+    """
+    [LanceDB](https://lancedb.github.io/lancedb/) metadata store for vector and structured data.
+
+    LanceDB is a columnar database optimized for vector search and multimodal data.
+    Each feature is stored in its own Lance table within the database directory.
+    Uses Polars components for data processing (no native SQL execution).
+
+    Storage layout:
+
+    - Each feature gets its own table: `{namespace}__{feature_name}`
+
+    - Tables are stored as Lance format in the directory specified by the URI
+
+    - LanceDB handles schema evolution, transactions, and compaction automatically
+
+    Example: Local Directory
+        ```py
+        from pathlib import Path
+        from metaxy.metadata_store.lancedb import LanceDBMetadataStore
+
+        # Local filesystem
+        store = LanceDBMetadataStore(Path("/path/to/featuregraph"))
+        ```
+
+    Example: Object Storage (S3, GCS, Azure)
+        ```py
+        # object store (requires credentials)
+        store = LanceDBMetadataStore("s3:///path/to/featuregraph")
+        ```
+
+    Example: LanceDB Cloud
+        ```py
+        import os
+
+        # Option 1: Environment variable
+        os.environ["LANCEDB_API_KEY"] = "your-api-key"
+        store = LanceDBMetadataStore("db://my-database")
+
+        # Option 2: Explicit credentials
+        store = LanceDBMetadataStore(
+            "db://my-database",
+            connect_kwargs={"api_key": "your-api-key", "region": "us-east-1"}
+        )
+        ```
+    """
+
+    _should_warn_auto_create_tables = False
+
+    def __init__(
+        self,
+        uri: str | Path,
+        *,
+        fallback_stores: list[MetadataStore] | None = None,
+        connect_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ):
+        """
+        Initialize [LanceDB](https://lancedb.com/docs/) metadata store.
+
+        The database directory is created automatically if it doesn't exist (local paths only).
+        Tables are created on-demand when features are first written.
+
+        Args:
+            uri: Directory path or URI for LanceDB tables. Supports:
+
+                - **Local path**: `"./metadata"` or `Path("/data/metaxy/lancedb")`
+
+                - **Object stores**: `s3://`, `gs://`, `az://` (requires cloud credentials)
+
+                - **LanceDB Cloud**: `"db://database-name"` (requires API key)
+
+                - **Remote HTTP/HTTPS**: Any URI supported by LanceDB
+
+            fallback_stores: Ordered list of read-only fallback stores.
+                When reading features not found in this store, Metaxy searches
+                fallback stores in order. Useful for local dev → staging → production chains.
+            connect_kwargs: Extra keyword arguments passed directly to
+                [lancedb.connect()](https://lancedb.github.io/lancedb/python/python/#lancedb.connect).
+                Useful for LanceDB Cloud credentials (api_key, region) when you cannot
+                rely on environment variables.
+            **kwargs: Passed to [metaxy.metadata_store.base.MetadataStore][]
+                (e.g., hash_algorithm, hash_truncation_length, prefer_native)
+
+        Note:
+            Unlike SQL stores, LanceDB doesn't require explicit table creation.
+            Tables are created automatically when writing metadata.
+        """
+        self.uri: str = str(uri)
+        self._conn: Any | None = None
+        self._connect_kwargs = connect_kwargs or {}
+        super().__init__(
+            fallback_stores=fallback_stores,
+            auto_create_tables=True,
+            versioning_engine_cls=PolarsVersioningEngine,
+            **kwargs,
+        )
+
+    @contextmanager
+    def _create_versioning_engine(self, plan):
+        """Create Polars versioning engine for LanceDB.
+
+        Args:
+            plan: Feature plan for the feature we're tracking provenance for
+
+        Yields:
+            PolarsVersioningEngine instance
+        """
+        engine = PolarsVersioningEngine(plan=plan)
+        try:
+            yield engine
+        finally:
+            # No cleanup needed for Polars engine
+            pass
+
+    @contextmanager
+    def open(self, mode: AccessMode = "read") -> Iterator[Self]:
+        """Open LanceDB connection.
+
+        For local filesystem paths, creates the directory if it doesn't exist.
+        For remote URIs (S3, LanceDB Cloud, etc.), connects directly.
+        Tables are created on-demand when features are first written.
+
+        Args:
+            mode: Access mode (READ or WRITE). Accepted for consistency but not used
+                by LanceDB (LanceDB handles concurrent access internally).
+
+        Yields:
+            Self: The store instance
+
+        Raises:
+            ConnectionError: If remote connection fails (e.g., invalid credentials)
+        """
+        # Increment context depth to support nested contexts
+        self._context_depth += 1
+
+        try:
+            # Only perform actual open on first entry
+            if self._context_depth == 1:
+                import lancedb
+
+                if is_local_path(self.uri):
+                    Path(self.uri).mkdir(parents=True, exist_ok=True)
+
+                self._conn = lancedb.connect(self.uri, **self._connect_kwargs)
+                self._is_open = True
+                self._validate_after_open()
+
+            yield self
+        finally:
+            # Decrement context depth
+            self._context_depth -= 1
+
+            # Only perform actual close on last exit
+            if self._context_depth == 0:
+                self._conn = None
+                self._is_open = False
+
+    @property
+    def conn(self) -> Any:
+        """Get LanceDB connection.
+
+        Returns:
+            Active LanceDB connection
+
+        Raises:
+            StoreNotOpenError: If store is not open
+        """
+        from metaxy.metadata_store.exceptions import StoreNotOpenError
+
+        if self._conn is None:
+            raise StoreNotOpenError(
+                "LanceDB connection is not open. Store must be used as a context manager."
+            )
+        return self._conn
+
+    # Helpers -----------------------------------------------------------------
+
+    def _table_name(self, feature_key: FeatureKey) -> str:
+        return feature_key.table_name
+
+    def _table_exists(self, table_name: str) -> bool:
+        """Check if a table exists without listing all tables.
+
+        Uses open_table() which is more efficient than listing all tables,
+        especially for remote storage (S3, GCS, etc.) where listing is expensive.
+
+        Args:
+            table_name: Name of the table to check
+
+        Returns:
+            True if table exists, False otherwise
+        """
+        try:
+            self.conn.open_table(table_name)  # type: ignore[attr-defined]
+            return True
+        except (ValueError, FileNotFoundError):
+            # LanceDB raises ValueError when table doesn't exist
+            return False
+
+    def _get_table(self, table_name: str):
+        return self.conn.open_table(table_name)  # type: ignore[attr-defined]
+
+    # ===== MetadataStore abstract methods =====
+
+    def _has_feature_impl(self, feature: CoercibleToFeatureKey) -> bool:
+        """Check if feature exists in LanceDB store.
+
+        Args:
+            feature: Feature to check
+
+        Returns:
+            True if feature exists, False otherwise
+        """
+        feature_key = self._resolve_feature_key(feature)
+        table_name = self._table_name(feature_key)
+        return self._table_exists(table_name)
+
+    def _get_default_hash_algorithm(self) -> HashAlgorithm:
+        """Use XXHASH64 by default to match other non-SQL stores."""
+        return HashAlgorithm.XXHASH64
+
+    # Storage ------------------------------------------------------------------
+
+    def write_metadata_to_store(
+        self,
+        feature_key: FeatureKey,
+        df: Frame,
+        **kwargs: Any,
+    ) -> None:
+        """Append metadata to Lance table.
+
+        Creates the table if it doesn't exist, otherwise appends to existing table.
+        Uses LanceDB's native Polars/Arrow integration for efficient storage.
+
+        Args:
+            feature_key: Feature key to write to
+            df: Narwhals Frame with metadata (already validated by base class)
+        """
+        # Convert Narwhals frame to Polars DataFrame
+        df_polars = collect_to_polars(df)
+
+        table_name = self._table_name(feature_key)
+
+        # LanceDB supports both Polars DataFrames and Arrow tables directly
+        # Try Polars first (native integration), fall back to Arrow if needed
+        try:
+            if self._table_exists(table_name):
+                table = self._get_table(table_name)
+                # Use Polars DataFrame directly - LanceDB handles conversion
+                table.add(df_polars)  # type: ignore[attr-defined]
+            else:
+                # Create table from Polars DataFrame - LanceDB handles schema
+                self.conn.create_table(table_name, data=df_polars)  # type: ignore[attr-defined]
+        except TypeError as exc:
+            if not self._should_fallback_to_arrow(exc):
+                raise
+            # Defensive fallback: Modern LanceDB (>=0.3) accepts Polars DataFrames natively,
+            # but fall back to Arrow if an older version or edge case doesn't support it.
+            # This ensures compatibility across LanceDB versions.
+            logger.debug("Falling back to Arrow format for LanceDB write: %s", exc)
+            arrow_table = df_polars.to_arrow()
+            if self._table_exists(table_name):
+                table = self._get_table(table_name)
+                table.add(arrow_table)  # type: ignore[attr-defined]
+            else:
+                self.conn.create_table(table_name, data=arrow_table)  # type: ignore[attr-defined]
+
+    def _drop_feature_metadata_impl(self, feature_key: FeatureKey) -> None:
+        """Drop Lance table for feature.
+
+        Permanently removes the Lance table from the database directory.
+        Safe to call even if table doesn't exist (no-op).
+
+        Args:
+            feature_key: Feature key to drop metadata for
+        """
+        table_name = self._table_name(feature_key)
+        if self._table_exists(table_name):
+            self.conn.drop_table(table_name)  # type: ignore[attr-defined]
+
+    def read_metadata_in_store(
+        self,
+        feature: CoercibleToFeatureKey,
+        *,
+        filters: Sequence[nw.Expr] | None = None,
+        columns: Sequence[str] | None = None,
+        **kwargs: Any,
+    ) -> nw.LazyFrame[Any] | None:
+        """Read metadata from Lance table.
+
+        Loads data from Lance, converts to Polars, and returns as Narwhals LazyFrame.
+        Applies filters and column selection in memory.
+
+        Args:
+            feature: Feature to read
+            filters: List of Narwhals filter expressions
+            columns: Optional list of columns to select
+            **kwargs: Backend-specific parameters (unused)
+
+        Returns:
+            Narwhals LazyFrame with metadata, or None if table not found
+        """
+        self._check_open()
+        feature_key = self._resolve_feature_key(feature)
+        table_name = self._table_name(feature_key)
+        if not self._table_exists(table_name):
+            return None
+
+        table = self._get_table(table_name)
+        # https://github.com/lancedb/lancedb/issues/1539
+        # Fall back to eager Arrow conversion until LanceDB issue #1539 is resolved.
+        arrow_table = table.to_arrow()
+        pl_lazy = pl.DataFrame(arrow_table).lazy()
+        nw_lazy = nw.from_native(pl_lazy)
+
+        if filters:
+            nw_lazy = nw_lazy.filter(*filters)
+
+        if columns is not None:
+            nw_lazy = nw_lazy.select(columns)
+
+        return nw_lazy
+
+    @staticmethod
+    def _should_fallback_to_arrow(exc: TypeError) -> bool:
+        """Return True when TypeError likely originates from Polars support gaps."""
+        message = str(exc).lower()
+        polars_markers = ("polars", "dataframe", "lazyframe", "data frame")
+        return any(marker in message for marker in polars_markers)
+
+    # Display ------------------------------------------------------------------
+
+    def display(self) -> str:
+        """Human-readable representation with sanitized credentials."""
+        path = sanitize_uri(self.uri)
+        return f"LanceDBMetadataStore(path={path})"
+
+    @classmethod
+    def config_model(cls) -> type[LanceDBMetadataStoreConfig]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return LanceDBMetadataStoreConfig

metaxy/metadata_store/memory.py

@@ -0,0 +1,292 @@
+"""In-memory metadata store implementation."""
+
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+from typing import Any
+
+import narwhals as nw
+import polars as pl
+from narwhals.typing import Frame
+from typing_extensions import Self
+
+from metaxy._utils import collect_to_polars
+from metaxy.metadata_store.base import MetadataStore, MetadataStoreConfig
+from metaxy.metadata_store.types import AccessMode
+from metaxy.models.types import CoercibleToFeatureKey, FeatureKey
+from metaxy.versioning.polars import PolarsVersioningEngine
+from metaxy.versioning.types import HashAlgorithm
+
+
+class InMemoryMetadataStoreConfig(MetadataStoreConfig):
+    """Configuration for InMemoryMetadataStore.
+
+    Example:
+        ```python
+        config = InMemoryMetadataStoreConfig(
+            hash_algorithm=HashAlgorithm.XXHASH64,
+        )
+
+        store = InMemoryMetadataStore.from_config(config)
+        ```
+    """
+
+    pass
+
+
+class InMemoryMetadataStore(MetadataStore):
+    """
+    In-memory metadata store using dict-based storage.
+
+    Features:
+    - Simple dict storage: {FeatureKey: pl.DataFrame}
+    - Fast for testing and prototyping
+    - No persistence (data lost when process exits)
+    - Schema validation on write
+    - Uses Polars components for all operations
+
+    Limitations:
+    - Not suitable for production
+    - Data lost on process exit
+    - No concurrency support across processes
+    - Memory-bound (all data in RAM)
+
+    Notes:
+        Uses Narwhals LazyFrames (nw.LazyFrame) for all operations
+
+    Components:
+        Components are created on-demand in resolve_update().
+        Uses Polars internally but exposes Narwhals interface.
+        Only supports Polars components (no native backend).
+    """
+
+    # Disable auto_create_tables warning for in-memory store
+    # (table creation concept doesn't apply to memory storage)
+    _should_warn_auto_create_tables: bool = False
+
+    def __init__(self, **kwargs: Any):
+        """
+        Initialize in-memory store.
+
+        Args:
+            **kwargs: Passed to MetadataStore.__init__ (e.g., fallback_stores, hash_algorithm)
+        """
+        # Use tuple as key (hashable) instead of string to avoid parsing issues
+        self._storage: dict[tuple[str, ...], pl.DataFrame] = {}
+        super().__init__(**kwargs, versioning_engine_cls=PolarsVersioningEngine)
+
+    def _get_default_hash_algorithm(self) -> HashAlgorithm:
+        """Get default hash algorithm for in-memory store."""
+        return HashAlgorithm.XXHASH64
+
+    def _get_storage_key(self, feature_key: FeatureKey) -> tuple[str, ...]:
+        """Convert feature key to storage key (tuple for hashability)."""
+        return tuple(feature_key)
+
+    @contextmanager
+    def _create_versioning_engine(self, plan) -> Iterator[PolarsVersioningEngine]:
+        """Create Polars provenance engine for in-memory store.
+
+        Args:
+            plan: Feature plan for the feature we're tracking provenance for
+
+        Yields:
+            PolarsVersioningEngine instance
+        """
+        from metaxy.versioning.polars import PolarsVersioningEngine
+
+        # Create engine (only accepts plan parameter)
+        engine = PolarsVersioningEngine(plan=plan)
+
+        try:
+            yield engine
+        finally:
+            # No cleanup needed for Polars engine
+            pass
+
+    def _has_feature_impl(self, feature: CoercibleToFeatureKey) -> bool:
+        feature_key = self._resolve_feature_key(feature)
+        storage_key = self._get_storage_key(feature_key)
+        return storage_key in self._storage
+
+    def write_metadata_to_store(
+        self,
+        feature_key: FeatureKey,
+        df: Frame,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Internal write implementation for in-memory storage.
+
+        Args:
+            feature_key: Feature key to write to
+            df: Narwhals Frame (eager or lazy) with metadata (already validated)
+            **kwargs: Backend-specific parameters (currently unused)
+        """
+        df_polars: pl.DataFrame = collect_to_polars(df)
+
+        storage_key = self._get_storage_key(feature_key)
+
+        # Append or create
+        if storage_key in self._storage:
+            existing_df = self._storage[storage_key]
+
+            # Handle schema evolution: ensure both DataFrames have matching columns
+            # Add missing columns as null to the existing DataFrame
+            for col_name in df_polars.columns:
+                if col_name not in existing_df.columns:
+                    # Get the data type from the new DataFrame
+                    col_dtype = df_polars.schema[col_name]
+                    # Add column with null values of the appropriate type
+                    existing_df = existing_df.with_columns(
+                        pl.lit(None).cast(col_dtype).alias(col_name)
+                    )
+
+            # Add missing columns to the new DataFrame
+            for col_name in existing_df.columns:
+                if col_name not in df_polars.columns:
+                    # Get the data type from the existing DataFrame
+                    col_dtype = existing_df.schema[col_name]
+                    # Add column with null values of the appropriate type
+                    df_polars = df_polars.with_columns(
+                        pl.lit(None).cast(col_dtype).alias(col_name)
+                    )  # type: ignore[arg-type,union-attr]
+
+            # Ensure column order matches by selecting columns in consistent order
+            all_columns = sorted(set(existing_df.columns) | set(df_polars.columns))
+            existing_df = existing_df.select(all_columns)
+            df_polars = df_polars.select(all_columns)
+
+            # Now we can safely concat
+            self._storage[storage_key] = pl.concat(
+                [existing_df, df_polars],
+                how="vertical",
+            )
+        else:
+            # Create new
+            self._storage[storage_key] = df_polars
+
+    def _drop_feature_metadata_impl(self, feature_key: FeatureKey) -> None:
+        """Drop all metadata for a feature from in-memory storage.
+
+        Args:
+            feature_key: Feature key to drop metadata for
+        """
+        storage_key = self._get_storage_key(feature_key)
+
+        # Remove from storage if it exists
+        if storage_key in self._storage:
+            del self._storage[storage_key]
+
+    def read_metadata_in_store(
+        self,
+        feature: CoercibleToFeatureKey,
+        *,
+        feature_version: str | None = None,
+        filters: Sequence[nw.Expr] | None = None,
+        columns: Sequence[str] | None = None,
+        **kwargs: Any,
+    ) -> nw.LazyFrame[Any] | None:
+        """
+        Read metadata from this store only (no fallback).
+
+        Args:
+            feature: Feature to read
+            feature_version: Filter by specific feature_version
+            filters: List of Narwhals filter expressions
+            columns: Optional list of columns to select
+            **kwargs: Backend-specific parameters (currently unused)
+
+        Returns:
+            Narwhals LazyFrame with metadata, or None if not found
+
+        Raises:
+            StoreNotOpenError: If store is not open
+        """
+        self._check_open()
+
+        feature_key = self._resolve_feature_key(feature)
+        storage_key = self._get_storage_key(feature_key)
+
+        if storage_key not in self._storage:
+            return None
+
+        # Start with lazy Polars DataFrame, wrap with Narwhals
+        df_lazy = self._storage[storage_key].lazy()
+        nw_lazy = nw.from_native(df_lazy)
+
+        # Apply feature_version filter
+        if feature_version is not None:
+            nw_lazy = nw_lazy.filter(
+                nw.col("metaxy_feature_version") == feature_version
+            )
+
+        # Apply generic Narwhals filters
+        if filters is not None:
+            for filter_expr in filters:
+                nw_lazy = nw_lazy.filter(filter_expr)
+
+        # Select columns
+        if columns is not None:
+            nw_lazy = nw_lazy.select(columns)
+
+        # Check if result would be empty (we need to check the underlying frame)
+        # For now, return the lazy frame - emptiness check happens when materializing
+        return nw_lazy
+
+    def clear(self) -> None:
+        """
+        Clear all metadata from store.
+
+        Useful for testing.
+        """
+        self._storage.clear()
+
+    # ========== Context Manager Implementation ==========
+
+    @contextmanager
+    def open(self, mode: AccessMode = "read") -> Iterator[Self]:
+        """Open the in-memory store (no-op for in-memory, but accepts mode for consistency).
+
+        Args:
+            mode: Access mode (accepted for consistency but ignored).
+
+        Yields:
+            Self: The store instance
+        """
+        # Increment context depth to support nested contexts
+        self._context_depth += 1
+
+        try:
+            # Only perform actual open on first entry
+            if self._context_depth == 1:
+                # No actual connection needed for in-memory
+                # Mark store as open and validate
+                self._is_open = True
+                self._validate_after_open()
+
+            yield self
+        finally:
+            # Decrement context depth
+            self._context_depth -= 1
+
+            # Only perform actual close on last exit
+            if self._context_depth == 0:
+                # Nothing to clean up
+                self._is_open = False
+
+    def __repr__(self) -> str:
+        """String representation."""
+        num_fallbacks = len(self.fallback_stores)
+        status = "open" if self._is_open else "closed"
+        return (
+            f"InMemoryMetadataStore(status={status}, fallback_stores={num_fallbacks})"
+        )
+
+    def display(self) -> str:
+        """Display string for this store."""
+        status = "open" if self._is_open else "closed"
+        return f"InMemoryMetadataStore(status={status})"
+
+    @classmethod
+    def config_model(cls) -> type[InMemoryMetadataStoreConfig]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return InMemoryMetadataStoreConfig