metaxy 0.0.0__py3-none-any.whl

metaxy/data_versioning/calculators/duckdb.py

@@ -0,0 +1,186 @@
+"""DuckDB-specific data version calculator with extension management.
+
+This calculator extends IbisDataVersionCalculator to handle DuckDB-specific
+extension loading (e.g., hashfuncs for xxHash support).
+"""
+# pyright: reportImportCycles=false
+
+from typing import TYPE_CHECKING
+
+from metaxy.data_versioning.calculators.ibis import IbisDataVersionCalculator
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+
+if TYPE_CHECKING:
+    import ibis
+
+    from metaxy.data_versioning.calculators.ibis import HashSQLGenerator
+    from metaxy.metadata_store.duckdb import (
+        ExtensionSpec,  # pyright: ignore[reportImportCycles]
+    )
+
+
+class DuckDBDataVersionCalculator(IbisDataVersionCalculator):
+    """DuckDB-specific calculator that manages extensions lazily.
+
+    This calculator:
+    1. Installs and loads DuckDB extensions on first use (lazy loading)
+    2. Supports xxHash64, xxHash32, and MD5 hash functions
+    3. Generates DuckDB-specific SQL for hash computation
+
+    The extension loading happens in __init__, which is only called when
+    native data version calculations are actually needed (not on store open).
+
+    Example:
+        >>> backend = ibis.duckdb.connect("metadata.db")
+        >>> calculator = DuckDBDataVersionCalculator(
+        ...     backend=backend,
+        ...     extensions=["hashfuncs"]
+        ... )
+        >>> # Extensions are now loaded and xxHash64 is available
+    """
+
+    def __init__(
+        self,
+        backend: "ibis.BaseBackend",
+        extensions: "list[ExtensionSpec | str] | None" = None,
+    ):
+        """Initialize DuckDB calculator and load extensions.
+
+        Args:
+            backend: DuckDB Ibis backend connection
+            extensions: List of DuckDB extensions to install/load.
+                Can be strings (from 'community' repo) or dicts with
+                'name' and optional 'repository' keys.
+
+        Example:
+            >>> extensions = ["hashfuncs"]  # Simple form
+            >>> extensions = [{"name": "spatial", "repository": "core_nightly"}]
+        """
+        self._backend = backend
+        self.extensions = extensions or []
+
+        # Load extensions immediately (lazy at calculator creation time)
+        self._load_extensions()
+
+        # Generate hash SQL generators for DuckDB
+        hash_sql_generators = self._generate_hash_sql_generators()
+
+        # Initialize parent with backend and generators
+        super().__init__(
+            backend=backend,
+            hash_sql_generators=hash_sql_generators,
+        )
+
+    def _load_extensions(self) -> None:
+        """Install and load DuckDB extensions.
+
+        This is called once when the calculator is created, which happens
+        lazily when native data version calculations are first needed.
+        """
+        if not self.extensions:
+            return
+
+        # Type narrowing: we know this is a DuckDB backend
+        from typing import Any, cast
+
+        backend = cast(
+            Any, self._backend
+        )  # DuckDB backend has raw_sql but not in ibis.BaseBackend stubs
+
+        for ext_spec in self.extensions:
+            if isinstance(ext_spec, str):
+                # Simple string form - install from community repo
+                ext_name = ext_spec
+                # Install and load extension
+                backend.raw_sql(f"INSTALL {ext_name}")
+                backend.raw_sql(f"LOAD {ext_name}")
+            else:
+                # Dict form with optional repository
+                ext_name = ext_spec.get("name", "")
+                ext_repo = ext_spec.get("repository", "community")
+
+                if ext_repo != "community":
+                    # Set custom repository
+                    backend.raw_sql(f"SET custom_extension_repository='{ext_repo}'")
+
+                # Install and load extension
+                backend.raw_sql(f"INSTALL {ext_name}")
+                backend.raw_sql(f"LOAD {ext_name}")
+
+    def _generate_hash_sql_generators(self) -> dict[HashAlgorithm, "HashSQLGenerator"]:
+        """Generate hash SQL generators for DuckDB.
+
+        DuckDB supports:
+        - MD5: Always available (built-in)
+        - XXHASH32, XXHASH64: Available when 'hashfuncs' extension is loaded
+
+        Returns:
+            Dictionary mapping HashAlgorithm to SQL generator functions
+        """
+        generators: dict[HashAlgorithm, HashSQLGenerator] = {}
+
+        # MD5 is always available
+        def md5_generator(table, concat_columns: dict[str, str]) -> str:
+            hash_selects: list[str] = []
+            for field_key, concat_col in concat_columns.items():
+                hash_col = f"__hash_{field_key}"
+                # md5() in DuckDB returns hex string, cast to VARCHAR for consistency
+                hash_expr = f"CAST(md5({concat_col}) AS VARCHAR)"
+                hash_selects.append(f"{hash_expr} as {hash_col}")
+
+            hash_clause = ", ".join(hash_selects)
+            table_sql = table.compile()
+            return f"SELECT *, {hash_clause} FROM ({table_sql}) AS __metaxy_temp"
+
+        generators[HashAlgorithm.MD5] = md5_generator
+
+        # Check if hashfuncs extension is in the list
+        extension_names = [
+            ext if isinstance(ext, str) else ext.get("name", "")
+            for ext in self.extensions
+        ]
+
+        if "hashfuncs" in extension_names:
+
+            def xxhash32_generator(table, concat_columns: dict[str, str]) -> str:
+                hash_selects: list[str] = []
+                for field_key, concat_col in concat_columns.items():
+                    hash_col = f"__hash_{field_key}"
+                    hash_expr = f"CAST(xxh32({concat_col}) AS VARCHAR)"
+                    hash_selects.append(f"{hash_expr} as {hash_col}")
+
+                hash_clause = ", ".join(hash_selects)
+                table_sql = table.compile()
+                return f"SELECT *, {hash_clause} FROM ({table_sql}) AS __metaxy_temp"
+
+            def xxhash64_generator(table, concat_columns: dict[str, str]) -> str:
+                hash_selects: list[str] = []
+                for field_key, concat_col in concat_columns.items():
+                    hash_col = f"__hash_{field_key}"
+                    hash_expr = f"CAST(xxh64({concat_col}) AS VARCHAR)"
+                    hash_selects.append(f"{hash_expr} as {hash_col}")
+
+                hash_clause = ", ".join(hash_selects)
+                table_sql = table.compile()
+                return f"SELECT *, {hash_clause} FROM ({table_sql}) AS __metaxy_temp"
+
+            generators[HashAlgorithm.XXHASH32] = xxhash32_generator
+            generators[HashAlgorithm.XXHASH64] = xxhash64_generator
+
+        return generators
+
+    @property
+    def supported_algorithms(self) -> list[HashAlgorithm]:
+        """Algorithms supported by this calculator based on loaded extensions."""
+        # Dynamically determine based on what was actually loaded
+        return list(self._hash_sql_generators.keys())
+
+    @property
+    def default_algorithm(self) -> HashAlgorithm:
+        """Default hash algorithm for DuckDB.
+
+        Uses XXHASH64 if hashfuncs extension is loaded, otherwise MD5.
+        """
+        if HashAlgorithm.XXHASH64 in self.supported_algorithms:
+            return HashAlgorithm.XXHASH64
+        return HashAlgorithm.MD5

metaxy/data_versioning/calculators/ibis.py

@@ -0,0 +1,225 @@
+"""Ibis-based data version calculator using native SQL hash functions.
+
+This calculator uses Ibis to generate backend-specific SQL for hash computation,
+executing entirely in the database without pulling data into memory.
+"""
+
+from typing import TYPE_CHECKING, Any, Protocol
+
+import narwhals as nw
+
+from metaxy.data_versioning.calculators.base import DataVersionCalculator
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+
+if TYPE_CHECKING:
+    import ibis
+    import ibis.expr.types
+    import ibis.expr.types.relations
+
+    from metaxy.models.feature_spec import FeatureSpec
+    from metaxy.models.plan import FeaturePlan
+
+
+class HashSQLGenerator(Protocol):
+    """Protocol for backend-specific hash SQL generation.
+
+    Takes an Ibis table with concatenated columns and returns SQL that adds hash columns.
+    """
+
+    def __call__(
+        self, table: "ibis.expr.types.Table", concat_columns: dict[str, str]
+    ) -> str:
+        """Generate SQL query to compute hash columns.
+
+        Args:
+            table: Input Ibis table with concatenated columns
+            concat_columns: Maps field_key -> concat_column_name
+
+        Returns:
+            SQL query string that selects all columns plus hash columns
+        """
+        ...
+
+
+class IbisDataVersionCalculator(DataVersionCalculator):
+    """Calculates data versions using native SQL hash functions via Ibis.
+
+    This calculator:
+    1. Accepts Narwhals LazyFrame as input
+    2. Converts to Ibis table internally
+    3. Builds concatenated columns using Ibis expressions
+    4. Applies backend-specific SQL hash functions
+    5. Returns Narwhals LazyFrame
+
+    Different SQL backends have different hash function names and signatures,
+    so hash functions are provided as SQL template generators per backend.
+
+    Example hash SQL generators:
+        DuckDB: SELECT *, CAST(xxh64(concat_col) AS VARCHAR) as hash FROM table
+        ClickHouse: SELECT *, CAST(xxHash64(concat_col) AS String) as hash FROM table
+        PostgreSQL: SELECT *, MD5(concat_col) as hash FROM table
+    """
+
+    def __init__(
+        self,
+        backend: "ibis.BaseBackend",
+        hash_sql_generators: dict[HashAlgorithm, HashSQLGenerator],
+    ):
+        """Initialize calculator with Ibis backend and hash SQL generators.
+
+        Args:
+            backend: Ibis backend connection for SQL execution
+            hash_sql_generators: Map from HashAlgorithm to SQL generator function
+        """
+        self._backend = backend
+        self._hash_sql_generators = hash_sql_generators
+
+    @property
+    def supported_algorithms(self) -> list[HashAlgorithm]:
+        """Algorithms supported by this calculator."""
+        return list(self._hash_sql_generators.keys())
+
+    @property
+    def default_algorithm(self) -> HashAlgorithm:
+        """Default hash algorithm.
+
+        Base implementation returns XXHASH64 if available, otherwise first available.
+        """
+        if HashAlgorithm.XXHASH64 in self.supported_algorithms:
+            return HashAlgorithm.XXHASH64
+        return self.supported_algorithms[0]
+
+    def calculate_data_versions(
+        self,
+        joined_upstream: nw.LazyFrame[Any],
+        feature_spec: "FeatureSpec",
+        feature_plan: "FeaturePlan",
+        upstream_column_mapping: dict[str, str],
+        hash_algorithm: HashAlgorithm | None = None,
+    ) -> nw.LazyFrame[Any]:
+        """Calculate data_version using SQL hash functions.
+
+        Args:
+            joined_upstream: Narwhals LazyFrame with upstream data joined
+            feature_spec: Feature specification
+            feature_plan: Feature plan
+            upstream_column_mapping: Maps upstream key -> column name
+            hash_algorithm: Hash to use
+
+        Returns:
+            Narwhals LazyFrame with data_version column added
+        """
+        import ibis
+
+        algo = hash_algorithm or self.default_algorithm
+
+        if algo not in self.supported_algorithms:
+            raise ValueError(
+                f"Hash algorithm {algo} not supported by {self.__class__.__name__}. "
+                f"Supported: {self.supported_algorithms}"
+            )
+
+        # Convert Narwhals LazyFrame to Ibis table
+        import ibis.expr.types
+
+        native = joined_upstream.to_native()
+
+        # Validate that we have an Ibis table
+        if not isinstance(native, ibis.expr.types.Table):
+            # Not an Ibis table - this calculator only works with Ibis-backed data
+            raise TypeError(
+                f"IbisDataVersionCalculator requires Ibis-backed data. "
+                f"Got {type(native)} instead. "
+                f"This usually means the metadata store is not using Ibis tables. "
+                f"Use PolarsDataVersionCalculator for non-Ibis stores."
+            )
+
+        ibis_table: ibis.expr.types.Table = native  # type: ignore[assignment]
+
+        # Get the hash SQL generator
+        hash_sql_gen = self._hash_sql_generators[algo]
+
+        # Build concatenated string columns for each field (using Ibis expressions)
+        concat_columns = {}
+
+        for field in feature_spec.fields:
+            field_key_str = (
+                field.key.to_string()
+                if hasattr(field.key, "to_string")
+                else "__".join(field.key)
+            )
+
+            field_deps = feature_plan.field_dependencies.get(field.key, {})
+
+            # Build hash components (same structure as Polars)
+            components = [
+                ibis.literal(field_key_str),
+                ibis.literal(str(field.code_version)),
+            ]
+
+            # Add upstream data versions in deterministic order
+            for upstream_feature_key in sorted(field_deps.keys()):
+                upstream_fields = field_deps[upstream_feature_key]
+                upstream_key_str = (
+                    upstream_feature_key.to_string()
+                    if hasattr(upstream_feature_key, "to_string")
+                    else "__".join(upstream_feature_key)
+                )
+
+                data_version_col_name = upstream_column_mapping.get(
+                    upstream_key_str, "data_version"
+                )
+
+                for upstream_field in sorted(upstream_fields):
+                    upstream_field_str = (
+                        upstream_field.to_string()
+                        if hasattr(upstream_field, "to_string")
+                        else "__".join(upstream_field)
+                    )
+
+                    components.append(
+                        ibis.literal(f"{upstream_key_str}/{upstream_field_str}")
+                    )
+                    # Access struct field for upstream field's hash
+                    components.append(
+                        ibis_table[data_version_col_name][upstream_field_str]
+                    )
+
+            # Concatenate all components with separator
+            concat_expr = components[0]
+            for component in components[1:]:
+                concat_expr = concat_expr.concat(ibis.literal("|")).concat(component)  # pyright: ignore[reportAttributeAccessIssue]
+
+            # Store concat column for this field
+            concat_col_name = f"__concat_{field_key_str}"
+            concat_columns[field_key_str] = concat_col_name
+            ibis_table = ibis_table.mutate(**{concat_col_name: concat_expr})
+
+        # Generate SQL for hashing all concat columns
+        hash_sql = hash_sql_gen(ibis_table, concat_columns)
+
+        # Execute SQL to get table with hash columns
+        result_table = self._backend.sql(hash_sql)  # pyright: ignore[reportAttributeAccessIssue]
+
+        # Build data_version struct from hash columns
+        hash_col_names = [f"__hash_{k}" for k in concat_columns.keys()]
+        field_keys = list(concat_columns.keys())
+
+        # Create struct column from hash columns
+        struct_fields = {
+            field_key: result_table[f"__hash_{field_key}"] for field_key in field_keys
+        }
+
+        # Drop temp columns and add data_version
+        cols_to_keep = [
+            c
+            for c in result_table.columns
+            if c not in concat_columns.values() and c not in hash_col_names
+        ]
+
+        result_table = result_table.select(
+            *cols_to_keep, data_version=ibis.struct(struct_fields)
+        )
+
+        # Convert back to Narwhals LazyFrame
+        return nw.from_native(result_table, eager_only=False)

metaxy/data_versioning/calculators/polars.py

@@ -0,0 +1,135 @@
+"""Polars implementation of data version calculator."""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import narwhals as nw
+import polars as pl
+import polars_hash as plh
+
+from metaxy.data_versioning.calculators.base import DataVersionCalculator
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+
+if TYPE_CHECKING:
+    from metaxy.models.feature_spec import FeatureSpec
+    from metaxy.models.plan import FeaturePlan
+
+
+class PolarsDataVersionCalculator(DataVersionCalculator):
+    """Calculates data versions using polars-hash.
+
+    Accepts Narwhals LazyFrames and converts internally to Polars for hashing.
+    Supports all hash functions available in polars-hash plugin.
+    Default is xxHash64 for cross-database compatibility.
+    """
+
+    # Map HashAlgorithm enum to polars-hash functions
+    _HASH_FUNCTION_MAP: dict[HashAlgorithm, Callable[[pl.Expr], pl.Expr]] = {
+        HashAlgorithm.XXHASH64: lambda expr: expr.nchash.xxhash64(),  # pyright: ignore[reportAttributeAccessIssue]
+        HashAlgorithm.XXHASH32: lambda expr: expr.nchash.xxhash32(),  # pyright: ignore[reportAttributeAccessIssue]
+        HashAlgorithm.WYHASH: lambda expr: expr.nchash.wyhash(),  # pyright: ignore[reportAttributeAccessIssue]
+        HashAlgorithm.SHA256: lambda expr: expr.chash.sha2_256(),  # pyright: ignore[reportAttributeAccessIssue]
+        HashAlgorithm.MD5: lambda expr: expr.nchash.md5(),  # pyright: ignore[reportAttributeAccessIssue]
+    }
+
+    @property
+    def supported_algorithms(self) -> list[HashAlgorithm]:
+        """All algorithms supported by polars-hash."""
+        return list(self._HASH_FUNCTION_MAP.keys())
+
+    @property
+    def default_algorithm(self) -> HashAlgorithm:
+        """xxHash64 - fast and cross-database compatible."""
+        return HashAlgorithm.XXHASH64
+
+    def calculate_data_versions(
+        self,
+        joined_upstream: nw.LazyFrame[Any],
+        feature_spec: "FeatureSpec",
+        feature_plan: "FeaturePlan",
+        upstream_column_mapping: dict[str, str],
+        hash_algorithm: HashAlgorithm | None = None,
+    ) -> nw.LazyFrame[Any]:
+        """Calculate data_version using polars-hash.
+
+        Args:
+            joined_upstream: Narwhals LazyFrame with upstream data joined
+            feature_spec: Feature specification
+            feature_plan: Feature plan
+            upstream_column_mapping: Maps upstream key -> column name
+            hash_algorithm: Hash to use (default: xxHash64)
+
+        Returns:
+            Narwhals LazyFrame with data_version column added
+        """
+        algo = hash_algorithm or self.default_algorithm
+
+        if algo not in self.supported_algorithms:
+            raise ValueError(
+                f"Hash algorithm {algo} not supported by PolarsDataVersionCalculator. "
+                f"Supported: {self.supported_algorithms}"
+            )
+
+        # Convert Narwhals LazyFrame to Polars LazyFrame
+        # Must collect first (LazyFrame doesn't have to_polars, only DataFrame does)
+        pl_lazy = joined_upstream.collect().to_polars().lazy()
+
+        hash_fn = self._HASH_FUNCTION_MAP[algo]
+
+        # Build hash expressions for each field
+        field_exprs = {}
+
+        for field in feature_spec.fields:
+            field_key_str = (
+                field.key.to_string()
+                if hasattr(field.key, "to_string")
+                else "_".join(field.key)
+            )
+
+            field_deps = feature_plan.field_dependencies.get(field.key, {})
+
+            # Build hash components
+            components = [
+                pl.lit(field_key_str),
+                pl.lit(str(field.code_version)),
+            ]
+
+            # Add upstream data versions in deterministic order
+            for upstream_feature_key in sorted(field_deps.keys()):
+                upstream_fields = field_deps[upstream_feature_key]
+                upstream_key_str = (
+                    upstream_feature_key.to_string()
+                    if hasattr(upstream_feature_key, "to_string")
+                    else "_".join(upstream_feature_key)
+                )
+
+                data_version_col_name = upstream_column_mapping.get(
+                    upstream_key_str, "data_version"
+                )
+
+                for upstream_field in sorted(upstream_fields):
+                    upstream_field_str = (
+                        upstream_field.to_string()
+                        if hasattr(upstream_field, "to_string")
+                        else "_".join(upstream_field)
+                    )
+
+                    components.append(
+                        pl.lit(f"{upstream_key_str}/{upstream_field_str}")
+                    )
+                    components.append(
+                        pl.col(data_version_col_name).struct.field(upstream_field_str)
+                    )
+
+            # Concatenate and hash
+            concat_expr = plh.concat_str(*components, separator="|")
+            hashed = hash_fn(concat_expr).cast(pl.Utf8)
+            field_exprs[field_key_str] = hashed
+
+        # Create data_version struct
+        data_version_expr = pl.struct(**field_exprs)  # type: ignore[call-overload]
+
+        result_pl = pl_lazy.with_columns(data_version_expr.alias("data_version"))
+
+        # Convert back to Narwhals LazyFrame
+        return nw.from_native(result_pl, eager_only=False)

metaxy/data_versioning/diff/__init__.py

@@ -0,0 +1,15 @@
+"""Metadata diff resolvers for identifying changed data versions."""
+
+from metaxy.data_versioning.diff.base import (
+    DiffResult,
+    LazyDiffResult,
+    MetadataDiffResolver,
+)
+from metaxy.data_versioning.diff.narwhals import NarwhalsDiffResolver
+
+__all__ = [
+    "DiffResult",
+    "LazyDiffResult",
+    "MetadataDiffResolver",
+    "NarwhalsDiffResolver",
+]