metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/models/lineage.py

@@ -0,0 +1,285 @@
+"""Lineage relationship types for feature dependencies.
+
+This module defines how features relate to their upstream dependencies in terms of
+cardinality and transformation patterns. These types make explicit the relationship
+between parent and child features, enabling proper provenance aggregation.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+from pydantic import Field as PydanticField
+from typing_extensions import Self
+
+
+class LineageRelationshipType(str, Enum):
+    """Type of lineage relationship between features."""
+
+    IDENTITY = "1:1"
+    AGGREGATION = "N:1"
+    EXPANSION = "1:N"
+
+
+class BaseLineageRelationship(BaseModel, ABC):  # pyright: ignore[reportUnsafeMultipleInheritance]
+    """Base class for lineage relationship configurations.
+
+    Lineage relationships define the cardinality and transformation pattern
+    between a child feature and its upstream dependencies.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    @abstractmethod
+    def get_aggregation_columns(
+        self,
+        target_id_columns: Sequence[str],
+    ) -> Sequence[str] | None:
+        """Get columns to aggregate on for this relationship type.
+
+        Args:
+            target_id_columns: The target feature's ID columns.
+
+        Returns:
+            Columns to group by for aggregation, or None if no aggregation needed.
+        """
+        raise NotImplementedError
+
+
+class IdentityRelationship(BaseLineageRelationship):
+    """One-to-one relationship where each child row maps to exactly one parent row.
+
+    This is the default relationship type. Parent and child features share the same
+    ID columns and have the same cardinality. No aggregation is performed.
+
+    Examples:
+        >>> # Default 1:1 relationship
+        >>> IdentityRelationship()
+
+        >>> # Or use the classmethod
+        >>> LineageRelationship.identity()
+    """
+
+    type: Literal[LineageRelationshipType.IDENTITY] = LineageRelationshipType.IDENTITY
+
+    def get_aggregation_columns(
+        self,
+        target_id_columns: Sequence[str],
+    ) -> Sequence[str] | None:
+        """No aggregation needed for identity relationships."""
+        return None
+
+
+class AggregationRelationship(BaseLineageRelationship):
+    """Many-to-one relationship where multiple parent rows aggregate to one child row.
+
+    Parent features have more granular ID columns than the child. The child aggregates
+    multiple parent rows by grouping on a subset of the parent's ID columns.
+
+    Attributes:
+        on: Columns to group by for aggregation. These should be a subset of the
+            target feature's ID columns. If not specified, uses all target ID columns.
+
+    Examples:
+        >>> # Aggregate sensor readings by hour
+        >>> AggregationRelationship(on=["sensor_id", "hour"])
+        >>> # Parent has: sensor_id, hour, minute
+        >>> # Child has: sensor_id, hour
+
+        >>> # Or use the classmethod
+        >>> LineageRelationship.aggregation(on=["user_id", "session_id"])
+    """
+
+    type: Literal[LineageRelationshipType.AGGREGATION] = (
+        LineageRelationshipType.AGGREGATION
+    )
+    on: Sequence[str] | None = PydanticField(
+        default=None,
+        description="Columns to group by for aggregation. Defaults to all target ID columns.",
+    )
+
+    def get_aggregation_columns(
+        self,
+        target_id_columns: Sequence[str],
+    ) -> Sequence[str]:
+        """Get columns to aggregate on."""
+        return self.on if self.on is not None else target_id_columns
+
+
+class ExpansionRelationship(BaseLineageRelationship):
+    """One-to-many relationship where one parent row expands to multiple child rows.
+
+    Child features have more granular ID columns than the parent. Each parent row
+    generates multiple child rows with additional ID columns.
+
+    Attributes:
+        on: Parent ID columns that identify the parent record. Child records with
+            the same parent IDs will share the same upstream provenance.
+            If not specified, will be inferred from the available columns.
+        id_generation_pattern: Optional pattern for generating child IDs.
+            Can be "sequential", "hash", or a custom pattern. If not specified,
+            the feature's load_input() method is responsible for ID generation.
+
+    Examples:
+        >>> # Video frames from video
+        >>> ExpansionRelationship(
+        ...     on=["video_id"],  # Parent ID
+        ...     id_generation_pattern="sequential"
+        ... )
+        >>> # Parent has: video_id
+        >>> # Child has: video_id, frame_id (generated)
+
+        >>> # Text chunks from document
+        >>> ExpansionRelationship(on=["doc_id"])
+        >>> # Parent has: doc_id
+        >>> # Child has: doc_id, chunk_id (generated in load_input)
+    """
+
+    type: Literal[LineageRelationshipType.EXPANSION] = LineageRelationshipType.EXPANSION
+    on: Sequence[str] = PydanticField(
+        ...,
+        description="Parent ID columns for grouping. Child records with same parent IDs share provenance. Required for expansion relationships.",
+    )
+    id_generation_pattern: str | None = PydanticField(
+        default=None,
+        description="Pattern for generating child IDs. If None, handled by load_input().",
+    )
+
+    def get_aggregation_columns(
+        self,
+        target_id_columns: Sequence[str],
+    ) -> Sequence[str] | None:
+        """Get aggregation columns for the joiner phase.
+
+        For expansion relationships, returns None because aggregation
+        happens during diff resolution, not during joining. The joiner
+        should pass through all parent records without aggregation.
+
+        Args:
+            target_id_columns: The target (child) feature's ID columns.
+
+        Returns:
+            None - no aggregation during join phase for expansion relationships.
+        """
+        # Expansion relationships don't aggregate during join phase
+        # Aggregation happens later during diff resolution
+        return None
+
+
+# Discriminated union type for all lineage relationships
+LineageRelationshipUnion = (
+    IdentityRelationship | AggregationRelationship | ExpansionRelationship
+)
+
+
+class LineageRelationship(BaseModel):
+    """Wrapper class for lineage relationship configurations with convenient constructors.
+
+    This provides a cleaner API for creating lineage relationships while maintaining
+    type safety through discriminated unions.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    relationship: LineageRelationshipUnion = PydanticField(..., discriminator="type")
+
+    @classmethod
+    def identity(cls) -> Self:
+        """Create an identity (1:1) relationship.
+
+        Returns:
+            Configured LineageRelationship for 1:1 relationship.
+
+        Examples:
+            >>> spec = FeatureSpec(
+            ...     key="feature",
+            ...     lineage=LineageRelationship.identity()
+            ... )
+        """
+        return cls(relationship=IdentityRelationship())
+
+    @classmethod
+    def aggregation(cls, on: Sequence[str] | None = None) -> Self:
+        """Create an aggregation (N:1) relationship.
+
+        Args:
+            on: Columns to group by for aggregation. If None, uses all target ID columns.
+
+        Returns:
+            Configured LineageRelationship for N:1 relationship.
+
+        Examples:
+            >>> # Aggregate on specific columns
+            >>> spec = FeatureSpec(
+            ...     key="hourly_stats",
+            ...     id_columns=["sensor_id", "hour"],
+            ...     lineage=LineageRelationship.aggregation(on=["sensor_id", "hour"])
+            ... )
+
+            >>> # Aggregate on all ID columns (default)
+            >>> spec = FeatureSpec(
+            ...     key="user_summary",
+            ...     id_columns=["user_id"],
+            ...     lineage=LineageRelationship.aggregation()
+            ... )
+        """
+        return cls(relationship=AggregationRelationship(on=on))
+
+    @classmethod
+    def expansion(
+        cls,
+        on: Sequence[str],
+        id_generation_pattern: str | None = None,
+    ) -> Self:
+        """Create an expansion (1:N) relationship.
+
+        Args:
+            on: Parent ID columns that identify the parent record. Child records with
+                the same parent IDs will share the same upstream provenance.
+                Required - must explicitly specify which columns link parent to child.
+            id_generation_pattern: Pattern for generating child IDs.
+                Can be "sequential", "hash", or custom. If None, handled by load_input().
+
+        Returns:
+            Configured LineageRelationship for 1:N relationship.
+
+        Examples:
+            >>> # Sequential ID generation with explicit parent ID
+            >>> spec = FeatureSpec(
+            ...     key="video_frames",
+            ...     id_columns=["video_id", "frame_id"],
+            ...     lineage=LineageRelationship.expansion(
+            ...         on=["video_id"],
+            ...         id_generation_pattern="sequential"
+            ...     )
+            ... )
+
+            >>> # Custom ID generation in load_input()
+            >>> spec = FeatureSpec(
+            ...     key="text_chunks",
+            ...     id_columns=["doc_id", "chunk_id"],
+            ...     lineage=LineageRelationship.expansion(on=["doc_id"])
+            ... )
+        """
+        return cls(
+            relationship=ExpansionRelationship(
+                on=on, id_generation_pattern=id_generation_pattern
+            )
+        )
+
+    def get_aggregation_columns(
+        self, target_id_columns: Sequence[str]
+    ) -> Sequence[str] | None:
+        """Get columns to aggregate on for this relationship.
+
+        Args:
+            target_id_columns: The target feature's ID columns.
+
+        Returns:
+            Columns to group by for aggregation, or None if no aggregation needed.
+        """
+        return self.relationship.get_aggregation_columns(target_id_columns)

metaxy/models/plan.py

@@ -0,0 +1,232 @@
+from collections.abc import Mapping
+from functools import cached_property
+
+import pydantic
+
+from metaxy.models.bases import FrozenBaseModel
+from metaxy.models.feature_spec import FeatureDep, FeatureKey, FeatureSpec
+from metaxy.models.field import (
+    FieldDep,
+    FieldKey,
+    FieldSpec,
+    SpecialFieldDep,
+)
+from metaxy.models.fields_mapping import FieldsMappingResolutionContext
+from metaxy.models.types import CoercibleToFieldKey, ValidatedFieldKeyAdapter
+
+# Rebuild the model now that FeatureSpec is available
+FieldsMappingResolutionContext.model_rebuild()
+
+
+class FQFieldKey(FrozenBaseModel):
+    field: FieldKey
+    feature: FeatureKey
+
+    def to_string(self) -> str:
+        return f"{self.feature.to_string()}.{self.field.to_string()}"
+
+    def __repr__(self) -> str:
+        return self.to_string()
+
+    def __lt__(self, other: "FQFieldKey") -> bool:
+        """Enable sorting of FQFieldKey objects."""
+        return self.to_string() < other.to_string()
+
+    def __le__(self, other: "FQFieldKey") -> bool:
+        """Enable sorting of FQFieldKey objects."""
+        return self.to_string() <= other.to_string()
+
+    def __gt__(self, other: "FQFieldKey") -> bool:
+        """Enable sorting of FQFieldKey objects."""
+        return self.to_string() > other.to_string()
+
+    def __ge__(self, other: "FQFieldKey") -> bool:
+        """Enable sorting of FQFieldKey objects."""
+        return self.to_string() >= other.to_string()
+
+
+class FeaturePlan(FrozenBaseModel):
+    """Slice of the feature graph that includes a given feature and its parents"""
+
+    feature: pydantic.SkipValidation[FeatureSpec]
+    deps: pydantic.SkipValidation[list[FeatureSpec] | None]
+    feature_deps: list[FeatureDep] | None = (
+        None  # The actual dependency specifications with field mappings
+    )
+
+    @cached_property
+    def parent_features_by_key(
+        self,
+    ) -> Mapping[FeatureKey, FeatureSpec]:
+        return {feature.key: feature for feature in self.deps or []}
+
+    @cached_property
+    def all_parent_fields_by_key(self) -> Mapping[FQFieldKey, FieldSpec]:
+        res: dict[FQFieldKey, FieldSpec] = {}
+
+        for feature in self.deps or []:
+            for field in feature.fields:
+                res[FQFieldKey(field=field.key, feature=feature.key)] = field
+
+        return res
+
+    @cached_property
+    def parent_fields_by_key(self) -> Mapping[FQFieldKey, FieldSpec]:
+        res: dict[FQFieldKey, FieldSpec] = {}
+
+        for field in self.feature.fields:
+            res.update(self.get_parent_fields_for_field(field.key))
+
+        return res
+
+    @cached_property
+    def parent_fields_by_feature_key(self) -> Mapping[FeatureKey, set[FieldKey]]:
+        res: dict[FeatureKey, set[FieldKey]] = {}
+
+        if self.deps:
+            for feature in self.deps:
+                res[feature.key] = set([f.key for f in feature.fields])
+
+        return res
+
+    def get_parent_fields_for_field(
+        self, key: CoercibleToFieldKey
+    ) -> Mapping[FQFieldKey, FieldSpec]:
+        """Get parent fields for a given field key.
+
+        Args:
+            key: Field key to get parent fields for. Accepts string, sequence, or FieldKey.
+
+        Returns:
+            Mapping of fully qualified field keys to their specs.
+        """
+        # Validate and coerce the key
+        validated_key = ValidatedFieldKeyAdapter.validate_python(key)
+
+        res = {}
+
+        field = self.feature.fields_by_key[validated_key]
+
+        # Get resolved dependencies (combining automatic mapping and explicit deps)
+        resolved_deps = self._resolve_field_deps(field)
+
+        for field_dep in resolved_deps:
+            if field_dep.fields == SpecialFieldDep.ALL:
+                # we depend on all fields of the corresponding upstream feature
+                for parent_field in self.parent_features_by_key[
+                    field_dep.feature
+                ].fields:
+                    res[
+                        FQFieldKey(
+                            field=parent_field.key,
+                            feature=field_dep.feature,
+                        )
+                    ] = parent_field
+
+            elif isinstance(field_dep, FieldDep):
+                #
+                for field_key in field_dep.fields:
+                    fq_key = FQFieldKey(
+                        field=field_key,
+                        feature=field_dep.feature,
+                    )
+                    res[fq_key] = self.all_parent_fields_by_key[fq_key]
+            else:
+                raise ValueError(f"Unsupported dependency type: {type(field_dep)}")
+
+        return res
+
+    def _resolve_field_deps(self, field: FieldSpec) -> list[FieldDep]:
+        """Resolve field dependencies by combining explicit deps and automatic mapping.
+
+        Apply field mappings from the FeatureDep and add explicit deps.
+        """
+
+        if not self.feature_deps:
+            return []
+
+        # Check if field has explicit deps
+        if field.deps and field.deps != []:  # Check for non-empty list
+            if isinstance(field.deps, SpecialFieldDep):
+                # If it's SpecialFieldDep.ALL, return ALL for all upstream features
+                return [
+                    FieldDep(feature=dep.key, fields=SpecialFieldDep.ALL)
+                    for dep in (self.deps or [])
+                ]
+            else:
+                # Use only the explicit deps, no automatic mapping
+                return field.deps
+
+        # No explicit deps - use automatic mapping
+        field_deps = []
+
+        for feature_dep in self.feature_deps:
+            # Resolve field mapping for this specific upstream feature
+            # Get the upstream feature spec
+            upstream_feature = self.parent_features_by_key.get(feature_dep.feature)
+            if not upstream_feature:
+                continue
+
+            # Create resolution context
+            context = FieldsMappingResolutionContext(
+                field_key=field.key, upstream_feature=upstream_feature
+            )
+
+            mapped_deps = feature_dep.fields_mapping.resolve_field_deps(context)
+
+            if mapped_deps:
+                # Add a single FieldDep with all mapped fields
+                field_deps.append(
+                    FieldDep(feature=feature_dep.feature, fields=list(mapped_deps))
+                )
+            # Note: If mapped_deps is empty (e.g., feature excluded),
+            # we don't add any dependency for this feature
+
+        if field_deps:
+            return field_deps
+        else:
+            raise RuntimeError(
+                f"No upstream fields found for field {field} of feature {self.feature}. Please either specify explicit dependencies on it's FieldSpec or ensure that at least one FeatureDep on the FeatureSpec has a valid field mapping."
+            )
+
+    @cached_property
+    def field_dependencies(
+        self,
+    ) -> Mapping[FieldKey, Mapping[FeatureKey, list[FieldKey]]]:
+        """Get dependencies for each field in this feature.
+
+        Returns a mapping from field key to its upstream dependencies.
+        Each dependency maps an upstream feature key to a list of field keys
+        that this field depends on.
+
+        This is the format needed by DataVersionResolver.
+
+        Returns:
+            Mapping of field keys to their dependency specifications.
+            Format: {field_key: {upstream_feature_key: [upstream_field_keys]}}
+        """
+        result: dict[FieldKey, dict[FeatureKey, list[FieldKey]]] = {}
+
+        for field in self.feature.fields:
+            field_deps: dict[FeatureKey, list[FieldKey]] = {}
+
+            # Get resolved dependencies (combining automatic mapping and explicit deps)
+            resolved_deps = self._resolve_field_deps(field)
+
+            # Specific dependencies defined
+            for field_dep in resolved_deps:
+                feature_key = field_dep.feature
+
+                if field_dep.fields == SpecialFieldDep.ALL:
+                    # All fields from this upstream feature
+                    upstream_feature_spec = self.parent_features_by_key[feature_key]
+                    field_deps[feature_key] = [
+                        c.key for c in upstream_feature_spec.fields
+                    ]
+                elif isinstance(field_dep.fields, list):
+                    # Specific fields
+                    field_deps[feature_key] = field_dep.fields
+
+            result[field.key] = field_deps
+
+        return result