metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/versioning/feature_dep_transformer.py

@@ -0,0 +1,151 @@
+from collections.abc import Sequence
+from functools import cached_property
+
+import narwhals as nw
+from narwhals.typing import FrameT
+
+from metaxy.models.constants import (
+    METAXY_DATA_VERSION_BY_FIELD,
+    METAXY_PROVENANCE,
+    METAXY_PROVENANCE_BY_FIELD,
+)
+from metaxy.models.feature_spec import FeatureDep, FeatureSpec
+from metaxy.models.plan import FeaturePlan
+from metaxy.models.types import FeatureKey
+from metaxy.versioning.renamed_df import RenamedDataFrame
+
+
+class FeatureDepTransformer:
+    def __init__(self, dep: FeatureDep, plan: FeaturePlan):
+        """A class responsible for applying transformations that live on the [metaxy.models.feature_spec.FeatureDep][]:
+
+            - Filters (from FeatureDep.filters)
+            - Renames
+            - Selections
+
+        This is supposed to always run before the upstream metadata is joined.
+
+        Will also inject Metaxy system columns.
+        """
+        self.plan = plan
+        self.dep = dep
+
+        # allow adding more in the future
+        self.metaxy_columns_to_load = [
+            METAXY_PROVENANCE_BY_FIELD,
+            METAXY_PROVENANCE,
+            METAXY_DATA_VERSION_BY_FIELD,
+        ]
+
+    @cached_property
+    def upstream_feature_key(self) -> FeatureKey:
+        return self.dep.feature
+
+    @cached_property
+    def upstream_feature_spec(self) -> FeatureSpec:
+        return self.plan.parent_features_by_key[self.dep.feature]
+
+    def transform(
+        self, df: FrameT, filters: Sequence[nw.Expr] | None = None
+    ) -> RenamedDataFrame[FrameT]:
+        """Apply the transformation specified by the feature dependency.
+
+        Args:
+            df: The dataframe to transform, it's expected to represent the raw upstream feature metadata
+                as it resides in the metadata store.
+            filters: Optional sequence of additional filters to apply to the dataframe **after renames**.
+                These are combined with the static filters from FeatureDep.filters.
+
+        Returns:
+            The transformed dataframe coupled with the renamed ID columns
+
+        """
+        # Combine static filters from FeatureDep with any additional filters passed as arguments
+        combined_filters: list[nw.Expr] = []
+        if self.dep.filters is not None:
+            combined_filters.extend(self.dep.filters)
+        if filters:
+            combined_filters.extend(filters)
+
+        return (
+            RenamedDataFrame(
+                df=df, id_columns=list(self.upstream_feature_spec.id_columns)
+            )
+            .rename(self.renames)
+            .filter(combined_filters if combined_filters else None)
+            .select(self.renamed_columns)
+        )
+
+    def rename_upstream_metaxy_column(self, column_name: str) -> str:
+        """Insert the upstream feature key suffix into the column name.
+
+        Is typically applied to Metaxy's system columns since they have to be loaded and do not have user-defined renames."""
+        return f"{column_name}{self.upstream_feature_key.to_column_suffix()}"
+
+    @cached_property
+    def renamed_provenance_col(self) -> str:
+        return self.rename_upstream_metaxy_column(METAXY_PROVENANCE)
+
+    @cached_property
+    def renamed_provenance_by_field_col(self) -> str:
+        return self.rename_upstream_metaxy_column(METAXY_PROVENANCE_BY_FIELD)
+
+    @cached_property
+    def renamed_data_version_by_field_col(self) -> str:
+        return self.rename_upstream_metaxy_column(METAXY_DATA_VERSION_BY_FIELD)
+
+    @cached_property
+    def renamed_metaxy_cols(self) -> list[str]:
+        return list(
+            map(self.rename_upstream_metaxy_column, self.metaxy_columns_to_load)
+        )
+
+    @cached_property
+    def renames(self) -> dict[str, str]:
+        """Get column renames for an upstream feature.
+
+        Returns:
+            Dictionary of column renames
+        """
+        # TODO: potentially include more system columns here?
+        return {
+            **(self.dep.rename or {}),
+            **{
+                col: self.rename_upstream_metaxy_column(col)
+                for col in self.metaxy_columns_to_load
+            },
+        }
+
+    @cached_property
+    def renamed_id_columns(self) -> list[str]:
+        return [
+            self.renames.get(col, col) for col in self.upstream_feature_spec.id_columns
+        ]
+
+    @cached_property
+    def renamed_columns(
+        self,
+    ) -> list[str] | None:
+        """Get columns to select from an upstream feature.
+
+        There include both original and metaxy-injected columns, all already renamed.
+        Users are expected to use renamed column names in their columns specification.
+
+        Returns:
+            List of column names to select, or None to select all columns
+        """
+
+        # If no specific columns requested (None), return None to keep all columns
+        # If empty tuple, return only ID columns and system columns
+        if self.dep.columns is None:
+            return None
+        else:
+            # Apply renames to the selected columns since selection happens after renaming
+            renamed_selected_cols = [
+                self.renames.get(col, col) for col in self.dep.columns
+            ]
+            return [
+                *self.renamed_id_columns,
+                *renamed_selected_cols,
+                *self.renamed_metaxy_cols,
+            ]

metaxy/versioning/ibis.py

@@ -0,0 +1,249 @@
+"""Ibis implementation of VersioningEngine.
+
+CRITICAL: This implementation NEVER materializes lazy expressions.
+All operations stay in the lazy Ibis world for SQL execution.
+"""
+
+from typing import Protocol, cast
+
+import narwhals as nw
+from ibis import Expr as IbisExpr
+from narwhals.typing import FrameT
+
+from metaxy.models.plan import FeaturePlan
+from metaxy.versioning.engine import VersioningEngine
+from metaxy.versioning.types import HashAlgorithm
+
+
+class IbisHashFn(Protocol):
+    def __call__(self, expr: IbisExpr) -> IbisExpr: ...
+
+
+class IbisVersioningEngine(VersioningEngine):
+    """Provenance engine using Ibis for SQL databases.
+
+    Only implements hash_string_column and build_struct_column.
+    All logic lives in the base class.
+
+    CRITICAL: This implementation NEVER leaves the lazy world.
+    All operations stay as Ibis expressions that compile to SQL.
+    """
+
+    def __init__(
+        self,
+        plan: FeaturePlan,
+        hash_functions: dict[HashAlgorithm, IbisHashFn],
+    ) -> None:
+        """Initialize the Ibis engine.
+
+        Args:
+            plan: Feature plan to track provenance for
+            backend: Ibis backend instance (e.g., ibis.duckdb.connect())
+            hash_functions: Mapping from HashAlgorithm to Ibis hash functions.
+                Each function takes an Ibis expression and returns an Ibis expression.
+        """
+        super().__init__(plan)
+        self.hash_functions: dict[HashAlgorithm, IbisHashFn] = hash_functions
+
+    @classmethod
+    def implementation(cls) -> nw.Implementation:
+        return nw.Implementation.IBIS
+
+    def hash_string_column(
+        self,
+        df: FrameT,
+        source_column: str,
+        target_column: str,
+        hash_algo: HashAlgorithm,
+    ) -> FrameT:
+        """Hash a string column using Ibis hash functions.
+
+        Args:
+            df: Narwhals DataFrame backed by Ibis
+            source_column: Name of string column to hash
+            target_column: Name for the new column containing the hash
+            hash_algo: Hash algorithm to use
+
+        Returns:
+            Narwhals DataFrame with new hashed column added, backed by Ibis.
+            The source column remains unchanged.
+        """
+        if hash_algo not in self.hash_functions:
+            raise ValueError(
+                f"Hash algorithm {hash_algo} not supported by this Ibis backend. "
+                f"Supported: {list(self.hash_functions.keys())}"
+            )
+
+        # Import ibis lazily (module-level import restriction)
+        import ibis.expr.types
+
+        # Convert to Ibis table
+        assert df.implementation == nw.Implementation.IBIS, (
+            "Only Ibis DataFrames are accepted"
+        )
+        ibis_table: ibis.expr.types.Table = cast(ibis.expr.types.Table, df.to_native())
+
+        # Get hash function
+        hash_fn = self.hash_functions[hash_algo]
+
+        # Apply hash to source column
+        # Hash functions are responsible for returning strings
+        hashed = hash_fn(ibis_table[source_column])
+
+        # Add new column with the hash
+        result_table = ibis_table.mutate(**{target_column: hashed})  # pyright: ignore[reportArgumentType]
+
+        # Convert back to Narwhals
+        return cast(FrameT, nw.from_native(result_table))
+
+    @staticmethod
+    def build_struct_column(
+        df: FrameT,
+        struct_name: str,
+        field_columns: dict[str, str],
+    ) -> FrameT:
+        """Build a struct column from existing columns.
+
+        Args:
+            df: Narwhals DataFrame backed by Ibis
+            struct_name: Name for the new struct column
+            field_columns: Mapping from struct field names to column names
+
+        Returns:
+            Narwhals DataFrame with new struct column added, backed by Ibis.
+            The source columns remain unchanged.
+        """
+        # Import ibis lazily
+        import ibis.expr.types
+
+        # Convert to Ibis table
+        assert df.implementation == nw.Implementation.IBIS, (
+            "Only Ibis DataFrames are accepted"
+        )
+        ibis_table: ibis.expr.types.Table = cast(ibis.expr.types.Table, df.to_native())
+
+        # Build struct expression - reference columns by name
+        struct_expr = ibis.struct(
+            {
+                field_name: ibis_table[col_name]
+                for field_name, col_name in field_columns.items()
+            }
+        )
+
+        # Add struct column
+        result_table = ibis_table.mutate(**{struct_name: struct_expr})
+
+        # Convert back to Narwhals
+        return cast(FrameT, nw.from_native(result_table))
+
+    @staticmethod
+    def aggregate_with_string_concat(
+        df: FrameT,
+        group_by_columns: list[str],
+        concat_column: str,
+        concat_separator: str,
+        exclude_columns: list[str],
+    ) -> FrameT:
+        """Aggregate DataFrame by grouping and concatenating strings.
+
+        Args:
+            df: Narwhals DataFrame backed by Ibis
+            group_by_columns: Columns to group by
+            concat_column: Column containing strings to concatenate within groups
+            concat_separator: Separator to use when concatenating strings
+            exclude_columns: Columns to exclude from aggregation
+
+        Returns:
+            Narwhals DataFrame with one row per group.
+        """
+        # Import ibis lazily
+        import ibis
+        import ibis.expr.types
+
+        # Convert to Ibis table
+        assert df.implementation == nw.Implementation.IBIS, (
+            "Only Ibis DataFrames are accepted"
+        )
+        ibis_table: ibis.expr.types.Table = cast(ibis.expr.types.Table, df.to_native())
+
+        # Build aggregation expressions
+        agg_exprs = {}
+
+        # Concatenate the concat_column with separator
+        agg_exprs[concat_column] = ibis_table[concat_column].group_concat(
+            concat_separator
+        )
+
+        # Take first value for all other columns (except group_by and exclude)
+        all_columns = set(ibis_table.columns)
+        columns_to_aggregate = (
+            all_columns - set(group_by_columns) - {concat_column} - set(exclude_columns)
+        )
+
+        for col in columns_to_aggregate:
+            agg_exprs[col] = ibis_table[
+                col
+            ].arbitrary()  # Take any value (like first())
+
+        # Perform groupby and aggregate
+        result_table = ibis_table.group_by(group_by_columns).aggregate(**agg_exprs)
+
+        # Convert back to Narwhals
+        return cast(FrameT, nw.from_native(result_table))
+
+    @staticmethod
+    def keep_latest_by_group(
+        df: FrameT,
+        group_columns: list[str],
+        timestamp_column: str,
+    ) -> FrameT:
+        """Keep only the latest row per group based on a timestamp column.
+
+        Uses argmax aggregation to get the value from each column where the
+        timestamp is maximum. This is simpler and more semantically clear than
+        window functions.
+
+        Args:
+            df: Narwhals DataFrame/LazyFrame backed by Ibis
+            group_columns: Columns to group by (typically ID columns)
+            timestamp_column: Column to use for determining "latest" (typically metaxy_created_at)
+
+        Returns:
+            Narwhals DataFrame/LazyFrame with only the latest row per group
+
+        Raises:
+            ValueError: If timestamp_column doesn't exist in df
+        """
+        # Import ibis lazily
+        import ibis.expr.types
+
+        # Convert to Ibis table
+        assert df.implementation == nw.Implementation.IBIS, (
+            "Only Ibis DataFrames are accepted"
+        )
+
+        # Check if timestamp_column exists
+        if timestamp_column not in df.columns:
+            raise ValueError(
+                f"Timestamp column '{timestamp_column}' not found in DataFrame. "
+                f"Available columns: {df.columns}"
+            )
+
+        ibis_table: ibis.expr.types.Table = cast(ibis.expr.types.Table, df.to_native())
+
+        # Use argmax aggregation: for each column, get the value where timestamp is maximum
+        # This directly expresses "get the row with the latest timestamp per group"
+        all_columns = set(ibis_table.columns)
+        non_group_columns = all_columns - set(group_columns)
+
+        # Build aggregation dict: for each non-group column, use argmax(timestamp)
+        agg_exprs = {
+            col: ibis_table[col].argmax(ibis_table[timestamp_column])
+            for col in non_group_columns
+        }
+
+        # Perform groupby and aggregate
+        result_table = ibis_table.group_by(group_columns).aggregate(**agg_exprs)
+
+        # Convert back to Narwhals
+        return cast(FrameT, nw.from_native(result_table))

metaxy/versioning/lineage_handler.py

@@ -0,0 +1,205 @@
+"""Handler for normalizing provenance based on lineage relationships.
+
+This module provides abstractions for handling different lineage relationship types
+(identity, aggregation, expansion) when comparing expected vs current provenance.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import narwhals as nw
+from narwhals.typing import FrameT
+
+from metaxy.models.constants import METAXY_PROVENANCE, METAXY_PROVENANCE_BY_FIELD
+from metaxy.models.lineage import ExpansionRelationship
+from metaxy.utils.hashing import get_hash_truncation_length
+
+if TYPE_CHECKING:
+    from metaxy.models.plan import FeaturePlan
+    from metaxy.versioning.engine import VersioningEngine
+    from metaxy.versioning.types import HashAlgorithm
+
+
+class LineageHandler(ABC):
+    """Base class for handling lineage-based provenance normalization."""
+
+    def __init__(self, feature_plan: FeaturePlan, engine: VersioningEngine):
+        """Initialize handler with feature plan and engine.
+
+        Args:
+            feature_plan: The feature plan containing lineage information
+            engine: The provenance engine instance
+        """
+        self.plan = feature_plan
+        self.feature_spec = feature_plan.feature
+        self.engine = engine
+
+    @abstractmethod
+    def normalize_for_comparison(
+        self,
+        expected: FrameT,
+        current: FrameT,
+        hash_algorithm: HashAlgorithm,
+    ) -> tuple[FrameT, FrameT, list[str]]:
+        """Normalize expected and current DataFrames for provenance comparison.
+
+        Args:
+            expected: Expected metadata computed from upstream
+            current: Current metadata from store
+            hash_algorithm: Hash algorithm to use
+            hash_length: Hash truncation length
+
+        Returns:
+            Tuple of (normalized_expected, normalized_current, join_columns)
+        """
+        pass
+
+
+class IdentityLineageHandler(LineageHandler):
+    """Handler for 1:1 identity lineage relationships.
+
+    No normalization needed - each upstream row maps to exactly one downstream row.
+    """
+
+    def normalize_for_comparison(
+        self,
+        expected: FrameT,
+        current: FrameT,
+        hash_algorithm: HashAlgorithm,
+    ) -> tuple[FrameT, FrameT, list[str]]:
+        """No normalization needed for identity relationships."""
+        id_columns = list(self.feature_spec.id_columns)
+        return expected, current, id_columns
+
+
+class AggregationLineageHandler(LineageHandler):
+    """Handler for N:1 aggregation lineage relationships.
+
+    Multiple upstream rows aggregate to one downstream row. We need to:
+    1. Group expected metadata by aggregation columns (sorted within group)
+    2. Concatenate provenance values deterministically
+    3. Hash the concatenated result using engine's hash method
+    """
+
+    def normalize_for_comparison(
+        self,
+        expected: FrameT,
+        current: FrameT,
+        hash_algorithm: HashAlgorithm,
+    ) -> tuple[FrameT, FrameT, list[str]]:
+        """Aggregate expected provenance by grouping."""
+        id_columns = list(self.feature_spec.id_columns)
+        agg_result = self.feature_spec.lineage.get_aggregation_columns(id_columns)
+        assert agg_result is not None, (
+            "Aggregation relationship must have aggregation columns"
+        )
+        agg_columns = list(agg_result)
+
+        # Aggregate expected provenance
+        expected_agg = self._aggregate_provenance(expected, agg_columns, hash_algorithm)
+
+        return expected_agg, current, agg_columns
+
+    def _aggregate_provenance(
+        self,
+        expected: FrameT,
+        agg_columns: list[str],
+        hash_algorithm: HashAlgorithm,
+    ) -> FrameT:
+        """Aggregate provenance for N:1 relationships.
+
+        Strategy:
+        1. Sort by id_columns within each group for deterministic ordering
+        2. Group by aggregation columns and concatenate provenance with engine's method
+        3. Hash the concatenated result using engine's hash_string_column
+
+        Args:
+            expected: Expected metadata with upstream provenance
+            agg_columns: Columns to group by
+            hash_algorithm: Hash algorithm to use
+            hash_length: Length to truncate hash to
+
+        Returns:
+            Aggregated DataFrame with one row per group
+        """
+        # Sort by all id_columns for deterministic ordering within groups
+        id_columns = list(self.feature_spec.id_columns)
+        expected_sorted = expected.sort(id_columns)
+
+        # Use engine's aggregate_with_string_concat method
+        # This concatenates provenance strings and stores in a temporary column
+        grouped = self.engine.aggregate_with_string_concat(
+            df=expected_sorted,
+            group_by_columns=agg_columns,
+            concat_column=METAXY_PROVENANCE,
+            concat_separator="|",
+            exclude_columns=[METAXY_PROVENANCE_BY_FIELD],
+        )
+
+        # Hash the concatenated provenance using engine's method
+        # Note: the concat column still has name METAXY_PROVENANCE after aggregation
+        hashed = self.engine.hash_string_column(
+            grouped, METAXY_PROVENANCE, "__hashed_prov", hash_algorithm
+        )
+
+        # Replace METAXY_PROVENANCE with truncated hash
+        hashed = hashed.drop(METAXY_PROVENANCE).rename(
+            {"__hashed_prov": METAXY_PROVENANCE}
+        )
+        hashed = hashed.with_columns(
+            nw.col(METAXY_PROVENANCE).str.slice(0, get_hash_truncation_length())
+        )
+
+        # Create placeholder provenance_by_field struct using engine's method
+        field_names = [f.key.to_struct_key() for f in self.plan.feature.fields]
+        field_map = {name: "__aggregated_placeholder" for name in field_names}
+
+        # Add placeholder column
+        hashed = hashed.with_columns(
+            nw.lit("aggregated").alias("__aggregated_placeholder")
+        )
+
+        # Build struct using engine's method
+        result = self.engine.build_struct_column(
+            hashed, METAXY_PROVENANCE_BY_FIELD, field_map
+        )
+
+        # Drop placeholder
+        result = result.drop("__aggregated_placeholder")
+
+        return result
+
+
+class ExpansionLineageHandler(LineageHandler):
+    """Handler for 1:N expansion lineage relationships.
+
+    One upstream row expands to many downstream rows. All downstream rows
+    with the same parent ID should have the same provenance. We group
+    current by parent columns and take any representative row.
+    """
+
+    def normalize_for_comparison(
+        self,
+        expected: FrameT,
+        current: FrameT,
+        hash_algorithm: HashAlgorithm,
+    ) -> tuple[FrameT, FrameT, list[str]]:
+        """Group current by parent ID columns."""
+        # Access the ExpansionRelationship to get the .on attribute
+        assert isinstance(self.feature_spec.lineage.relationship, ExpansionRelationship)
+        parent_columns = list(self.feature_spec.lineage.relationship.on)
+
+        # Group current by parent columns and take any representative row
+        current_grouped = (
+            current.with_columns(nw.lit(True).alias("_dummy"))
+            .filter(
+                nw.col("_dummy")
+                .is_first_distinct()
+                .over(*parent_columns, order_by="_dummy")
+            )
+            .drop("_dummy")
+        )
+
+        return expected, current_grouped, parent_columns