metaxy 0.0.0__py3-none-any.whl

metaxy/data_versioning/diff/base.py

@@ -0,0 +1,150 @@
+"""Abstract base class for metadata diff resolvers."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+import narwhals as nw
+
+if TYPE_CHECKING:
+    pass
+
+
+class LazyDiffResult(NamedTuple):
+    """Result of diffing with lazy Narwhals LazyFrames (opt-in via lazy=True).
+
+    Contains lazy Narwhals LazyFrames - users decide when/how to materialize.
+
+    Users can:
+    - Keep lazy for further operations: result.added.filter(...)
+    - Materialize to Polars: result.added.collect().to_native()
+    - Materialize to Pandas: result.added.collect().to_pandas()
+    - Materialize to PyArrow: result.added.collect().to_arrow()
+    - Convert to DiffResult: result.collect()
+
+    Backend execution:
+    - SQL stores: All operations stay in SQL until .collect()
+    - Polars stores: Operations stay lazy until .collect()
+
+    Attributes:
+        added: New samples (lazy, never None - empty LazyFrame instead)
+            Columns: [sample_uid, data_version, ...user columns...]
+        changed: Changed samples (lazy, never None)
+            Columns: [sample_uid, data_version, ...user columns...]
+        removed: Removed samples (lazy, never None)
+            Columns: [sample_uid, data_version, ...user columns...]
+
+    Note:
+        May contain additional user columns beyond sample_uid and data_version,
+        depending on what was passed to resolve_update() via align_upstream_metadata.
+    """
+
+    added: nw.LazyFrame[Any]
+    changed: nw.LazyFrame[Any]
+    removed: nw.LazyFrame[Any]
+
+    def collect(self) -> "DiffResult":
+        """Materialize all lazy frames to create a DiffResult.
+
+        Returns:
+            DiffResult with all frames materialized to eager DataFrames.
+        """
+        return DiffResult(
+            added=self.added.collect(),
+            changed=self.changed.collect(),
+            removed=self.removed.collect(),
+        )
+
+
+class DiffResult(NamedTuple):
+    """Result of diffing with eager Narwhals DataFrames (default).
+
+    Contains materialized Narwhals DataFrames - ready to use immediately.
+
+    Users can convert to their preferred format:
+    - Polars: result.added.to_native()
+    - Pandas: result.added.to_pandas()
+    - PyArrow: result.added.to_arrow()
+
+    Attributes:
+        added: New samples (eager, never None - empty DataFrame instead)
+            Columns: [sample_uid, data_version, ...user columns...]
+        changed: Changed samples (eager, never None)
+            Columns: [sample_uid, data_version, ...user columns...]
+        removed: Removed samples (eager, never None)
+            Columns: [sample_uid, data_version, ...user columns...]
+
+    Note:
+        May contain additional user columns beyond sample_uid and data_version,
+        depending on what was passed to resolve_update() via align_upstream_metadata.
+    """
+
+    added: nw.DataFrame[Any]
+    changed: nw.DataFrame[Any]
+    removed: nw.DataFrame[Any]
+
+
+class MetadataDiffResolver(ABC):
+    """Identifies rows with changed data_versions by comparing target with current.
+
+    The diff resolver compares newly calculated data_versions (target) with
+    existing metadata (current) to identify what needs to be written.
+
+    This is Step 3 in the data versioning process:
+    1. Join upstream features → unified upstream view
+    2. Calculate data_version from upstream → target versions
+    3. Diff with current metadata → identify changes ← THIS STEP
+
+    All component boundaries use Narwhals LazyFrames for backend-agnostic processing.
+
+    Examples:
+        - NarwhalsDiffResolver: Backend-agnostic using Narwhals expressions
+        - IbisDiffResolver: Converts to Ibis internally for SQL processing
+
+    Important Design:
+        Takes lazy Narwhals refs as input, returns LazyDiffResult as output.
+        This minimizes query execution:
+        - SQL backends: One query with CTEs computes all three categories
+        - Polars: Uses lazy operations, splits into three LazyFrames
+
+    Users can override Feature.resolve_data_version_diff to customize:
+        - Ignore certain field changes
+        - Apply custom change detection rules
+        - Filter out specific samples
+    """
+
+    @abstractmethod
+    def find_changes(
+        self,
+        target_versions: nw.LazyFrame[Any],
+        current_metadata: nw.LazyFrame[Any] | None,
+    ) -> LazyDiffResult:
+        """Find all changes between target and current metadata.
+
+        Compares target data_versions (newly calculated) with current metadata
+        and categorizes all differences.
+
+        Args:
+            target_versions: Narwhals LazyFrame with newly calculated data_versions
+                Shape: [sample_uid, data_version (calculated), upstream columns...]
+            current_metadata: Narwhals LazyFrame with current metadata, or None
+                Shape: [sample_uid, data_version (existing), feature_version, custom columns...]
+                Should be pre-filtered by feature_version at the caller level if needed.
+
+        Returns:
+            LazyDiffResult with three lazy Narwhals LazyFrames.
+            Caller materializes to DiffResult if needed (for lazy=False).
+
+        Implementation Note:
+            Should build lazy operations without materializing:
+            - SQL backends: Build one lazy query with CTEs for all three categories
+            - Polars: Use lazy operations, no collect() calls
+
+        Note:
+            For immutable append-only storage, typically only 'added' and 'changed'
+            are written. 'removed' is useful for validation/reporting.
+
+            Feature version filtering should happen at the read_metadata() level,
+            not in the diff resolver. The diff resolver just compares whatever
+            metadata is passed to it.
+        """
+        pass

metaxy/data_versioning/diff/narwhals.py

@@ -0,0 +1,108 @@
+"""Narwhals implementation of metadata diff resolver.
+
+Unified diff resolver that works with any backend (Polars, Ibis/SQL) through Narwhals.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import narwhals as nw
+
+from metaxy.data_versioning.diff.base import (
+    LazyDiffResult,
+    MetadataDiffResolver,
+)
+
+if TYPE_CHECKING:
+    pass
+
+
+class NarwhalsDiffResolver(MetadataDiffResolver):
+    """Identifies changed rows using Narwhals operations.
+
+    Uses Narwhals LazyFrames (works with Polars, Ibis, Pandas, PyArrow)
+
+    Strategy:
+    - Categorizes changes into added, changed, and removed
+    - Uses LEFT/RIGHT JOINs to identify each category
+    - Materializes once and splits into three DataFrames (efficient)
+    - Backend-agnostic: same code works for in-memory and SQL backends
+
+    The underlying backend (Polars vs Ibis) determines execution:
+    - Polars backend → operations happen in-memory
+    - Ibis backend → operations happen in SQL database
+    """
+
+    def find_changes(
+        self,
+        target_versions: nw.LazyFrame[Any],
+        current_metadata: nw.LazyFrame[Any] | None,
+    ) -> LazyDiffResult:
+        """Find all changes between target and current.
+
+        Args:
+            target_versions: Narwhals LazyFrame with calculated data_versions
+            current_metadata: Narwhals LazyFrame with current metadata, or None.
+                Should be pre-filtered by feature_version at caller level if needed.
+
+        Returns:
+            LazyDiffResult with three lazy Narwhals frames (caller materializes if needed)
+        """
+        # Select only sample_uid and data_version from target_versions
+        # (it may have intermediate joined columns from upstream)
+        target_versions = target_versions.select(["sample_uid", "data_version"])
+
+        if current_metadata is None:
+            # No existing metadata - all target rows are new
+            # Create empty LazyFrame with proper schema
+            import polars as pl
+
+            empty_lazy = nw.from_native(
+                pl.LazyFrame({"sample_uid": [], "data_version": []})
+            )
+
+            return LazyDiffResult(
+                added=target_versions,
+                changed=empty_lazy,
+                removed=empty_lazy,
+            )
+
+        # Keep only sample_uid and data_version from current for comparison
+        current_comparison = current_metadata.select(
+            "sample_uid", nw.col("data_version").alias("__current_data_version")
+        )
+
+        # LEFT JOIN target with current
+        compared = target_versions.join(
+            current_comparison,
+            on="sample_uid",
+            how="left",
+        )
+
+        # Build lazy queries for each category
+        added_lazy = (
+            compared.filter(nw.col("__current_data_version").is_null())
+            .drop("__current_data_version")
+            .select("sample_uid", "data_version")
+        )
+
+        changed_lazy = (
+            compared.filter(
+                ~nw.col("__current_data_version").is_null()
+                & (nw.col("data_version") != nw.col("__current_data_version"))
+            )
+            .drop("__current_data_version")
+            .select("sample_uid", "data_version")
+        )
+
+        removed_lazy = current_metadata.join(
+            target_versions.select("sample_uid"),
+            on="sample_uid",
+            how="anti",
+        ).select("sample_uid", "data_version")
+
+        # Return lazy frames - caller will materialize if needed
+        return LazyDiffResult(
+            added=added_lazy,
+            changed=changed_lazy,
+            removed=removed_lazy,
+        )

metaxy/data_versioning/hash_algorithms.py

@@ -0,0 +1,19 @@
+"""Hash algorithms supported for data versioning."""
+
+from enum import Enum
+
+
+class HashAlgorithm(Enum):
+    """Supported hash algorithms for data versioning.
+
+    These algorithms are chosen for:
+    - Speed (non-cryptographic hashes preferred)
+    - Cross-database availability
+    - Good collision resistance for data versioning
+    """
+
+    XXHASH64 = "xxhash64"  # Fast, available in DuckDB, ClickHouse, Polars
+    XXHASH32 = "xxhash32"  # Faster for small data, less collision resistant
+    WYHASH = "wyhash"  # Very fast, Polars-specific
+    SHA256 = "sha256"  # Cryptographic, slower, universally available
+    MD5 = "md5"  # Legacy, widely available, not recommended for new code

metaxy/data_versioning/joiners/__init__.py

@@ -0,0 +1,9 @@
+"""Upstream joiners for merging upstream feature metadata."""
+
+from metaxy.data_versioning.joiners.base import UpstreamJoiner
+from metaxy.data_versioning.joiners.narwhals import NarwhalsJoiner
+
+__all__ = [
+    "UpstreamJoiner",
+    "NarwhalsJoiner",
+]

metaxy/data_versioning/joiners/base.py

@@ -0,0 +1,70 @@
+"""Abstract base class for upstream joiners."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import narwhals as nw
+
+if TYPE_CHECKING:
+    from metaxy.models.feature_spec import FeatureSpec
+    from metaxy.models.plan import FeaturePlan
+
+
+class UpstreamJoiner(ABC):
+    """Joins upstream feature metadata together.
+
+    The joiner takes upstream feature metadata (which already has data_version columns)
+    and joins them together to create a unified view of all dependencies.
+
+    This is Step 1 in the data versioning process:
+    1. Join upstream features → unified upstream view
+    2. Calculate data_version from upstream → target versions
+    3. Diff with current metadata → identify changes
+
+    All component boundaries use Narwhals LazyFrames for backend-agnostic processing.
+
+    Examples:
+        - NarwhalsJoiner: Backend-agnostic using Narwhals expressions
+        - IbisJoiner: Converts to Ibis internally for SQL processing
+    """
+
+    @abstractmethod
+    def join_upstream(
+        self,
+        upstream_refs: dict[str, nw.LazyFrame[Any]],
+        feature_spec: "FeatureSpec",
+        feature_plan: "FeaturePlan",
+        upstream_columns: dict[str, tuple[str, ...] | None] | None = None,
+        upstream_renames: dict[str, dict[str, str] | None] | None = None,
+    ) -> tuple[nw.LazyFrame[Any], dict[str, str]]:
+        """Join all upstream features together with optional column selection/renaming.
+
+        Joins upstream feature metadata on sample_uid to create a unified reference
+        containing all upstream data_version columns needed for hash calculation,
+        plus any additional user-specified columns.
+
+        Args:
+            upstream_refs: Upstream feature metadata Narwhals LazyFrames
+                Keys are upstream feature keys (using to_string() format)
+                Values are Narwhals LazyFrames with upstream metadata
+            feature_spec: Specification of the feature being computed
+            feature_plan: Resolved feature plan with dependencies
+            upstream_columns: Optional dict mapping upstream feature keys to tuple of columns to keep.
+                None or missing key = keep all columns. Empty tuple = only system columns.
+            upstream_renames: Optional dict mapping upstream feature keys to rename dicts.
+                Applied after column selection.
+
+        Returns:
+            Tuple of (joined_ref, upstream_column_mapping):
+            - joined_ref: Narwhals LazyFrame with all upstream data joined
+                Contains: sample_uid, data_version columns, and any user columns
+            - upstream_column_mapping: Maps upstream feature key -> data_version column name
+                Example: {"video": "__upstream_video__data_version"}
+
+        Note:
+            - Uses INNER join by default - only sample_uids present in ALL upstream features
+              are included. This ensures we can compute valid data_versions.
+            - System columns (sample_uid, data_version) are always preserved
+            - User columns are preserved based on columns parameter (default: all)
+        """
+        pass

metaxy/data_versioning/joiners/narwhals.py

@@ -0,0 +1,235 @@
+"""Narwhals implementation of upstream joiner.
+
+Unified joiner that works with any backend (Polars, Ibis/SQL) through Narwhals.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import narwhals as nw
+
+from metaxy.data_versioning.joiners.base import UpstreamJoiner
+from metaxy.models.constants import DROPPABLE_SYSTEM_COLUMNS, ESSENTIAL_SYSTEM_COLUMNS
+
+if TYPE_CHECKING:
+    from metaxy.models.feature_spec import FeatureSpec
+    from metaxy.models.plan import FeaturePlan
+
+
+class NarwhalsJoiner(UpstreamJoiner):
+    """Joins upstream features using Narwhals LazyFrames.
+
+    Type Parameters:
+        TRef = nw.LazyFrame (works with Polars, Ibis, Pandas, PyArrow)
+
+    Strategy:
+    - Starts with first upstream feature
+    - Sequentially inner joins remaining upstream features on sample_uid
+    - Renames data_version columns to avoid conflicts
+    - All operations are lazy (no materialization until collect)
+    - Backend-agnostic: same code works for in-memory and SQL backends
+
+    The underlying backend (Polars vs Ibis) is determined by what's wrapped:
+    - nw.from_native(pl.LazyFrame) → stays in Polars
+    - nw.from_native(ibis.Table) → stays in SQL until collect()
+    """
+
+    def join_upstream(
+        self,
+        upstream_refs: dict[str, nw.LazyFrame[Any]],
+        feature_spec: "FeatureSpec",
+        feature_plan: "FeaturePlan",
+        upstream_columns: dict[str, tuple[str, ...] | None] | None = None,
+        upstream_renames: dict[str, dict[str, str] | None] | None = None,
+    ) -> tuple[nw.LazyFrame[Any], dict[str, str]]:
+        """Join upstream Narwhals LazyFrames together with column selection/renaming.
+
+        Args:
+            upstream_refs: Dict of upstream feature key -> Narwhals LazyFrame
+            feature_spec: Feature specification
+            feature_plan: Feature plan
+            upstream_columns: Optional column selection per upstream feature
+            upstream_renames: Optional column renaming per upstream feature
+
+        Returns:
+            (joined Narwhals LazyFrame, column mapping)
+        """
+        if not upstream_refs:
+            # No upstream dependencies - source feature
+            # Return empty LazyFrame with just sample_uid column (with proper type)
+            import polars as pl
+
+            # Create empty frame with explicit Int64 type for sample_uid
+            # This ensures it's not NULL-typed which would fail with Ibis backends
+            empty_df = pl.LazyFrame(
+                {"sample_uid": pl.Series("sample_uid", [], dtype=pl.Int64)}
+            )
+            return nw.from_native(empty_df), {}
+
+        # Initialize parameters if not provided
+        upstream_columns = upstream_columns or {}
+        upstream_renames = upstream_renames or {}
+
+        # Use imported constants for system columns
+        system_cols = ESSENTIAL_SYSTEM_COLUMNS
+        system_cols_to_drop = DROPPABLE_SYSTEM_COLUMNS
+
+        # Track all column names to detect conflicts
+        all_columns: dict[str, str] = {}  # column_name -> source_feature
+
+        # Process and join upstream features
+        upstream_keys = sorted(upstream_refs.keys())
+        first_key = upstream_keys[0]
+        upstream_mapping = {}
+
+        # Process first upstream feature
+        first_ref = upstream_refs[first_key]
+        first_columns_spec = upstream_columns.get(first_key)
+        first_renames_spec = upstream_renames.get(first_key) or {}
+
+        # Get column names from first upstream
+        # We need to collect schema to know available columns
+        # Use lazy evaluation where possible
+        first_schema = first_ref.collect_schema()
+        available_cols = set(first_schema.names())
+
+        # Determine columns to select
+        if first_columns_spec is None:
+            # Keep all columns (new default behavior) except problematic system columns
+            cols_to_select = [c for c in available_cols if c not in system_cols_to_drop]
+        elif first_columns_spec == ():
+            # Keep only essential system columns
+            cols_to_select = [c for c in available_cols if c in system_cols]
+        else:
+            # Keep specified columns plus essential system columns
+            requested = set(first_columns_spec)
+            # Filter out problematic system columns even if requested
+            requested = requested - system_cols_to_drop
+            cols_to_select = list(requested | (available_cols & system_cols))
+
+            # Warn about missing columns
+            missing = requested - available_cols
+            if missing:
+                import warnings
+
+                warnings.warn(
+                    f"Columns {missing} requested but not found in upstream feature {first_key}",
+                    UserWarning,
+                )
+
+        # Build select expressions with renaming for first upstream
+        select_exprs = []
+        for col in cols_to_select:
+            if col == "data_version":
+                # Always rename data_version to avoid conflicts
+                new_name = f"__upstream_{first_key}__data_version"
+                select_exprs.append(nw.col(col).alias(new_name))
+                upstream_mapping[first_key] = new_name
+            elif col in first_renames_spec:
+                # Apply user-specified rename
+                new_name = first_renames_spec[col]
+                if new_name in all_columns:
+                    raise ValueError(
+                        f"Column name conflict: '{new_name}' from {first_key} "
+                        f"conflicts with column from {all_columns[new_name]}. "
+                        f"Use the 'rename' parameter to resolve the conflict."
+                    )
+                select_exprs.append(nw.col(col).alias(new_name))
+                all_columns[new_name] = first_key
+            else:
+                # Keep original name
+                if col != "sample_uid" and col in all_columns:
+                    raise ValueError(
+                        f"Column name conflict: '{col}' appears in both "
+                        f"{first_key} and {all_columns[col]}. "
+                        f"Use the 'rename' parameter to resolve the conflict."
+                    )
+                select_exprs.append(nw.col(col))
+                if col != "sample_uid":
+                    all_columns[col] = first_key
+
+        joined = first_ref.select(select_exprs)
+
+        # Join remaining upstream features
+        for upstream_key in upstream_keys[1:]:
+            upstream_ref = upstream_refs[upstream_key]
+            columns_spec = upstream_columns.get(upstream_key)
+            renames_spec = upstream_renames.get(upstream_key) or {}
+
+            # Get available columns
+            schema = upstream_ref.collect_schema()
+            available_cols = set(schema.names())
+
+            # Determine columns to select
+            if columns_spec is None:
+                # Keep all columns except problematic system columns
+                cols_to_select = [
+                    c for c in available_cols if c not in system_cols_to_drop
+                ]
+            elif columns_spec == ():
+                # Keep only essential system columns
+                cols_to_select = [c for c in available_cols if c in system_cols]
+            else:
+                # Keep specified columns plus essential system columns
+                requested = set(columns_spec)
+                # Filter out problematic system columns even if requested
+                requested = requested - system_cols_to_drop
+                cols_to_select = list(requested | (available_cols & system_cols))
+
+                # Warn about missing columns
+                missing = requested - available_cols
+                if missing:
+                    import warnings
+
+                    warnings.warn(
+                        f"Columns {missing} requested but not found in upstream feature {upstream_key}",
+                        UserWarning,
+                    )
+
+            # Build select expressions with renaming
+            select_exprs = []
+            join_cols = []  # Columns to include in join (exclude sample_uid)
+
+            for col in cols_to_select:
+                if col == "sample_uid":
+                    # Always include sample_uid for joining, but don't duplicate it
+                    select_exprs.append(nw.col(col))
+                elif col == "data_version":
+                    # Always rename data_version to avoid conflicts
+                    new_name = f"__upstream_{upstream_key}__data_version"
+                    select_exprs.append(nw.col(col).alias(new_name))
+                    join_cols.append(new_name)
+                    upstream_mapping[upstream_key] = new_name
+                elif col in renames_spec:
+                    # Apply user-specified rename
+                    new_name = renames_spec[col]
+                    if new_name in all_columns:
+                        raise ValueError(
+                            f"Column name conflict: '{new_name}' from {upstream_key} "
+                            f"conflicts with column from {all_columns[new_name]}. "
+                            f"Use the 'rename' parameter to resolve the conflict."
+                        )
+                    select_exprs.append(nw.col(col).alias(new_name))
+                    join_cols.append(new_name)
+                    all_columns[new_name] = upstream_key
+                else:
+                    # Keep original name
+                    if col in all_columns:
+                        raise ValueError(
+                            f"Column name conflict: '{col}' appears in both "
+                            f"{upstream_key} and {all_columns[col]}. "
+                            f"Use the 'rename' parameter to resolve the conflict."
+                        )
+                    select_exprs.append(nw.col(col))
+                    join_cols.append(col)
+                    all_columns[col] = upstream_key
+
+            upstream_renamed = upstream_ref.select(select_exprs)
+
+            # Join with existing data
+            joined = joined.join(
+                upstream_renamed,
+                on="sample_uid",
+                how="inner",  # Only sample_uids present in ALL upstream
+            )
+
+        return joined, upstream_mapping