metaxy 0.0.0__py3-none-any.whl

metaxy/models/feature.py

@@ -0,0 +1,665 @@
+import hashlib
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from pydantic._internal._model_construction import ModelMetaclass
+from typing_extensions import Self
+
+from metaxy.models.bases import FrozenBaseModel
+from metaxy.models.feature_spec import FeatureSpec
+from metaxy.models.plan import FeaturePlan, FQFieldKey
+from metaxy.models.types import FeatureKey
+
+if TYPE_CHECKING:
+    import narwhals as nw
+
+    from metaxy.data_versioning.diff import (
+        DiffResult,
+        LazyDiffResult,
+        MetadataDiffResolver,
+    )
+    from metaxy.data_versioning.joiners import UpstreamJoiner
+
+
+# Context variable for active graph (module-level)
+_active_graph: ContextVar["FeatureGraph | None"] = ContextVar(
+    "_active_graph", default=None
+)
+
+
+def get_feature_by_key(key: "FeatureKey") -> type["Feature"]:
+    """Get a feature class by its key from the active graph.
+
+    Convenience function that retrieves from the currently active graph.
+
+    Args:
+        key: Feature key to look up
+
+    Returns:
+        Feature class
+
+    Raises:
+        KeyError: If no feature with the given key is registered
+
+    Example:
+        >>> from metaxy import get_feature_by_key, FeatureKey
+        >>> parent_key = FeatureKey(["examples", "parent"])
+        >>> ParentFeature = get_feature_by_key(parent_key)
+    """
+    graph = FeatureGraph.get_active()
+    return graph.get_feature_by_key(key)
+
+
+class FeatureGraph:
+    def __init__(self):
+        self.features_by_key: dict[FeatureKey, type[Feature]] = {}
+        self.feature_specs_by_key: dict[FeatureKey, FeatureSpec] = {}
+
+    def add_feature(self, feature: type["Feature"]) -> None:
+        """Add a feature to the graph.
+
+        Args:
+            feature: Feature class to register
+
+        Raises:
+            ValueError: If a feature with the same key is already registered
+        """
+        if feature.spec.key in self.features_by_key:
+            existing = self.features_by_key[feature.spec.key]
+            raise ValueError(
+                f"Feature with key {feature.spec.key.to_string()} already registered. "
+                f"Existing: {existing.__name__}, New: {feature.__name__}. "
+                f"Each feature key must be unique within a graph."
+            )
+
+        self.features_by_key[feature.spec.key] = feature
+        self.feature_specs_by_key[feature.spec.key] = feature.spec
+
+    def remove_feature(self, key: FeatureKey) -> None:
+        """Remove a feature from the graph.
+
+        Args:
+            key: Feature key to remove
+
+        Raises:
+            KeyError: If no feature with the given key is registered
+        """
+        if key not in self.features_by_key:
+            raise KeyError(
+                f"No feature with key {key.to_string()} found in graph. "
+                f"Available keys: {[k.to_string() for k in self.features_by_key.keys()]}"
+            )
+
+        del self.features_by_key[key]
+        del self.feature_specs_by_key[key]
+
+    def get_feature_by_key(self, key: FeatureKey) -> type["Feature"]:
+        """Get a feature class by its key.
+
+        Args:
+            key: Feature key to look up
+
+        Returns:
+            Feature class
+
+        Raises:
+            KeyError: If no feature with the given key is registered
+
+        Example:
+            >>> graph = FeatureGraph.get_active()
+            >>> parent_key = FeatureKey(["examples", "parent"])
+            >>> ParentFeature = graph.get_feature_by_key(parent_key)
+        """
+        if key not in self.features_by_key:
+            raise KeyError(
+                f"No feature with key {key.to_string()} found in graph. "
+                f"Available keys: {[k.to_string() for k in self.features_by_key.keys()]}"
+            )
+        return self.features_by_key[key]
+
+    def get_feature_plan(self, key: FeatureKey) -> FeaturePlan:
+        feature = self.feature_specs_by_key[key]
+
+        return FeaturePlan(
+            feature=feature,
+            deps=[self.feature_specs_by_key[dep.key] for dep in feature.deps or []]
+            or None,
+        )
+
+    def get_field_version(self, key: "FQFieldKey") -> str:
+        hasher = hashlib.sha256()
+
+        plan = self.get_feature_plan(key.feature)
+        field = plan.feature.fields_by_key[key.field]
+
+        hasher.update(key.to_string().encode())
+        hasher.update(str(field.code_version).encode())
+
+        for k, v in sorted(plan.get_parent_fields_for_field(key.field).items()):
+            hasher.update(self.get_field_version(k).encode())
+
+        return hasher.hexdigest()
+
+    def get_feature_version_by_field(self, key: FeatureKey) -> dict[str, str]:
+        """Computes the feature data version.
+
+        Hash together field data versions versions with the feature code version.
+
+        Returns:
+            dict[str, str]: The data version for each field in the feature plan.
+                Keys are field names as strings.
+        """
+        res = {}
+
+        plan = self.get_feature_plan(key)
+
+        for k, v in plan.feature.fields_by_key.items():
+            res[k.to_string()] = self.get_field_version(
+                FQFieldKey(field=k, feature=key)
+            )
+
+        return res
+
+    def get_feature_version(self, key: FeatureKey) -> str:
+        """Computes the feature version as a single string"""
+        hasher = hashlib.sha256()
+        data_version = self.get_feature_version_by_field(key)
+        for field_key in sorted(data_version):
+            hasher.update(field_key.encode())
+            hasher.update(data_version[field_key].encode())
+
+        return hasher.hexdigest()
+
+    def get_downstream_features(self, sources: list[FeatureKey]) -> list[FeatureKey]:
+        """Get all features downstream of sources, topologically sorted.
+
+        Performs a depth-first traversal of the dependency graph to find all
+        features that transitively depend on any of the source features.
+
+        Args:
+            sources: List of source feature keys
+
+        Returns:
+            List of downstream feature keys in topological order (dependencies first).
+            Does not include the source features themselves.
+
+        Example:
+            >>> # DAG: A -> B -> D
+            >>> #      A -> C -> D
+            >>> graph.get_downstream_features([FeatureKey(["A"])])
+            [FeatureKey(["B"]), FeatureKey(["C"]), FeatureKey(["D"])]
+        """
+        source_set = set(sources)
+        visited = set()
+        post_order = []  # Reverse topological order
+
+        def visit(key: FeatureKey):
+            """DFS traversal."""
+            if key in visited:
+                return
+            visited.add(key)
+
+            # Find all features that depend on this one
+            for feature_key, feature_spec in self.feature_specs_by_key.items():
+                if feature_spec.deps:
+                    for dep in feature_spec.deps:
+                        if dep.key == key:
+                            # This feature depends on 'key', so visit it
+                            visit(feature_key)
+
+            post_order.append(key)
+
+        # Visit all sources
+        for source in sources:
+            visit(source)
+
+        # Remove sources from result, reverse to get topological order
+        result = [k for k in reversed(post_order) if k not in source_set]
+        return result
+
+    @property
+    def snapshot_version(self) -> str:
+        """Generate a snapshot version representing the current topology + versions of the feature graph"""
+        if len(self.feature_specs_by_key) == 0:
+            return "empty"
+
+        hasher = hashlib.sha256()
+        for feature_key in sorted(self.feature_specs_by_key.keys()):
+            hasher.update(feature_key.to_string().encode("utf-8"))
+            hasher.update(self.get_feature_version(feature_key).encode("utf-8"))
+        return hasher.hexdigest()
+
+    def to_snapshot(self) -> dict[str, dict[str, Any]]:
+        """Serialize graph to snapshot format.
+
+        Returns a dict mapping feature_key (string) to feature data dict,
+        including the import path of the Feature class for reconstruction.
+
+        Returns:
+            Dict of feature_key -> {
+                feature_spec: dict,
+                feature_version: str,
+                feature_class_path: str
+            }
+
+        Example:
+            >>> snapshot = graph.to_snapshot()
+            >>> snapshot["video_processing"]["feature_version"]
+            'abc12345'
+            >>> snapshot["video_processing"]["feature_class_path"]
+            'myapp.features.video.VideoProcessing'
+        """
+        snapshot = {}
+
+        for feature_key, feature_cls in self.features_by_key.items():
+            feature_key_str = feature_key.to_string()
+            feature_spec_dict = feature_cls.spec.model_dump(mode="json")  # type: ignore[attr-defined]
+            feature_version = feature_cls.feature_version()  # type: ignore[attr-defined]
+
+            # Get class import path (module.ClassName)
+            class_path = f"{feature_cls.__module__}.{feature_cls.__name__}"
+
+            snapshot[feature_key_str] = {
+                "feature_spec": feature_spec_dict,
+                "feature_version": feature_version,
+                "feature_class_path": class_path,
+            }
+
+        return snapshot
+
+    @classmethod
+    def from_snapshot(
+        cls,
+        snapshot_data: dict[str, dict[str, Any]],
+        *,
+        class_path_overrides: dict[str, str] | None = None,
+        force_reload: bool = False,
+    ) -> "FeatureGraph":
+        """Reconstruct graph from snapshot by importing Feature classes.
+
+        Strictly requires Feature classes to exist at their recorded import paths.
+        This ensures custom methods (like load_input) are available.
+
+        If a feature has been moved/renamed, use class_path_overrides to specify
+        the new location.
+
+        Args:
+            snapshot_data: Dict of feature_key -> {
+                feature_spec: dict,
+                feature_class_path: str,
+                ...
+            } (as returned by to_snapshot() or loaded from DB)
+            class_path_overrides: Optional dict mapping feature_key to new class path
+                                 for features that have been moved/renamed
+            force_reload: If True, reload modules from disk to get current code state.
+
+        Returns:
+            New FeatureGraph with historical features
+
+        Raises:
+            ImportError: If feature class cannot be imported at recorded path
+
+        Example:
+            >>> # Load snapshot from metadata store
+            >>> historical_graph = FeatureGraph.from_snapshot(snapshot_data)
+            >>>
+            >>> # With override for moved feature
+            >>> historical_graph = FeatureGraph.from_snapshot(
+            ...     snapshot_data,
+            ...     class_path_overrides={
+            ...         "video_processing": "myapp.features_v2.VideoProcessing"
+            ...     }
+            ... )
+        """
+        import importlib
+        import sys
+
+        from metaxy.models.feature_spec import FeatureSpec
+
+        graph = cls()
+        class_path_overrides = class_path_overrides or {}
+
+        # If force_reload, collect all module paths first to remove ALL features
+        # from those modules before reloading (modules can have multiple features)
+        modules_to_reload = set()
+        if force_reload:
+            for feature_key_str, feature_data in snapshot_data.items():
+                class_path = class_path_overrides.get(
+                    feature_key_str
+                ) or feature_data.get("feature_class_path")
+                if class_path:
+                    module_path, _ = class_path.rsplit(".", 1)
+                    if module_path in sys.modules:
+                        modules_to_reload.add(module_path)
+
+        # Use context manager to temporarily set the new graph as active
+        # This ensures imported Feature classes register to the new graph, not the current one
+        with graph.use():
+            for feature_key_str, feature_data in snapshot_data.items():
+                # Parse FeatureSpec for validation
+                feature_spec_dict = feature_data["feature_spec"]
+                FeatureSpec.model_validate(feature_spec_dict)
+
+                # Get class path (check overrides first)
+                if feature_key_str in class_path_overrides:
+                    class_path = class_path_overrides[feature_key_str]
+                else:
+                    class_path = feature_data.get("feature_class_path")
+                    if not class_path:
+                        raise ValueError(
+                            f"Feature '{feature_key_str}' has no feature_class_path in snapshot. "
+                            f"Cannot reconstruct historical graph."
+                        )
+
+                # Import the class
+                try:
+                    module_path, class_name = class_path.rsplit(".", 1)
+
+                    # Force reload module from disk if requested
+                    # This is critical for migration detection - when code changes,
+                    # we need fresh imports to detect the changes
+                    if force_reload and module_path in modules_to_reload:
+                        # Before first reload of this module, remove ALL features from this module
+                        # (a module can define multiple features)
+                        if module_path in modules_to_reload:
+                            # Find all features from this module in snapshot and remove them
+                            for fk_str, fd in snapshot_data.items():
+                                fcp = class_path_overrides.get(fk_str) or fd.get(
+                                    "feature_class_path"
+                                )
+                                if fcp and fcp.rsplit(".", 1)[0] == module_path:
+                                    fspec_dict = fd["feature_spec"]
+                                    fspec = FeatureSpec.model_validate(fspec_dict)
+                                    if fspec.key in graph.features_by_key:
+                                        graph.remove_feature(fspec.key)
+
+                            # Mark module as processed so we don't remove features again
+                            modules_to_reload.discard(module_path)
+
+                        module = importlib.reload(sys.modules[module_path])
+                    else:
+                        module = __import__(module_path, fromlist=[class_name])
+
+                    feature_cls = getattr(module, class_name)
+                except (ImportError, AttributeError) as e:
+                    raise ImportError(
+                        f"Cannot import Feature class '{class_path}' for feature graph reconstruction from snapshot. "
+                        f"Feature '{feature_key_str}' is required to reconstruct the graph, but the class "
+                        f"cannot be found at the recorded import path. "
+                    ) from e
+
+                # Validate the imported class matches the stored spec
+                if not hasattr(feature_cls, "spec"):
+                    raise TypeError(
+                        f"Imported class '{class_path}' is not a valid Feature class "
+                        f"(missing 'spec' attribute)"
+                    )
+
+                # Register the imported feature to this graph if not already present
+                # If the module was imported for the first time, the metaclass already registered it
+                # If the module was previously imported, we need to manually register it
+                if feature_cls.spec.key not in graph.features_by_key:
+                    graph.add_feature(feature_cls)
+
+        return graph
+
+    @classmethod
+    def get_active(cls) -> "FeatureGraph":
+        """Get the currently active graph.
+
+        Returns the graph from the context variable if set, otherwise returns
+        the default global graph.
+
+        Returns:
+            Active FeatureGraph instance
+
+        Example:
+            >>> # Normal usage - returns default graph
+            >>> reg = FeatureGraph.get_active()
+            >>>
+            >>> # With custom graph in context
+            >>> with my_graph.use():
+            ...     reg = FeatureGraph.get_active()  # Returns my_graph
+        """
+        return _active_graph.get() or graph
+
+    @classmethod
+    def set_active(cls, reg: "FeatureGraph") -> None:
+        """Set the active graph for the current context.
+
+        This sets the context variable that will be returned by get_active().
+        Typically used in application setup code or test fixtures.
+
+        Args:
+            reg: FeatureGraph to activate
+
+        Example:
+            >>> # In application setup
+            >>> my_graph = FeatureGraph()
+            >>> FeatureGraph.set_active(my_graph)
+            >>>
+            >>> # Now all operations use my_graph
+            >>> FeatureGraph.get_active()  # Returns my_graph
+        """
+        _active_graph.set(reg)
+
+    @contextmanager
+    def use(self):
+        """Context manager to temporarily use this graph as active.
+
+        This is the recommended way to use custom registries, especially in tests.
+        The graph is automatically restored when the context exits.
+
+        Yields:
+            This graph instance
+
+        Example:
+            >>> test_graph = FeatureGraph()
+            >>>
+            >>> with test_graph.use():
+            ...     # All operations use test_graph
+            ...     class TestFeature(Feature, spec=...):
+            ...         pass
+            ...
+            >>> # Outside context, back to previous graph
+        """
+        token = _active_graph.set(self)
+        try:
+            yield self
+        finally:
+            _active_graph.reset(token)
+
+
+# Default global graph
+graph = FeatureGraph()
+
+
+class MetaxyMeta(ModelMetaclass):
+    def __new__(
+        cls,
+        cls_name: str,
+        bases: tuple[type[Any], ...],
+        namespace: dict[str, Any],
+        *,
+        spec: FeatureSpec | None,
+        **kwargs,
+    ) -> type[Self]:  # pyright: ignore[reportGeneralTypeIssues]
+        new_cls = super().__new__(cls, cls_name, bases, namespace, **kwargs)
+
+        if spec:
+            # Get graph from context at class definition time
+            active_graph = FeatureGraph.get_active()
+            new_cls.graph = active_graph  # type: ignore[attr-defined]
+            new_cls.spec = spec  # type: ignore[attr-defined]
+            active_graph.add_feature(new_cls)
+        else:
+            pass  # TODO: set spec to a property that would raise an exception on access
+
+        return new_cls
+
+
+class Feature(FrozenBaseModel, metaclass=MetaxyMeta, spec=None):
+    spec: ClassVar[FeatureSpec]
+    graph: ClassVar[FeatureGraph]
+
+    @classmethod
+    def table_name(cls) -> str:
+        """Get SQL-like table name for this feature.
+
+        Converts feature key to SQL-compatible table name by joining
+        parts with double underscores, consistent with IbisMetadataStore.
+
+        Returns:
+            Table name string (e.g., "my_namespace__my_feature")
+
+        Example:
+            >>> class VideoFeature(Feature, spec=FeatureSpec(
+            ...     key=FeatureKey(["video", "processing"]),
+            ...     ...
+            ... )):
+            ...     pass
+            >>> VideoFeature.table_name()
+            'video__processing'
+        """
+        return cls.spec.table_name()
+
+    @classmethod
+    def feature_version(cls) -> str:
+        """Get hash of feature specification.
+
+        Returns a hash representing the feature's complete configuration:
+        - Feature key
+        - Field definitions and code versions
+        - Dependencies (feature-level and field-level)
+
+        This hash changes when you modify:
+        - Field code versions
+        - Dependencies
+        - Field definitions
+
+        Used to distinguish current vs historical metadata versions.
+        Stored in the 'feature_version' column of metadata DataFrames.
+
+        Returns:
+            SHA256 hex digest (like git short hashes)
+
+        Example:
+            >>> class MyFeature(Feature, spec=FeatureSpec(
+            ...     key=FeatureKey(["my", "feature"]),
+            ...     fields=[FieldSpec(key=FieldKey(["default"]), code_version=1)],
+            ... )):
+            ...     pass
+            >>> MyFeature.feature_version()
+            'a3f8b2c1...'
+        """
+        return cls.graph.get_feature_version(cls.spec.key)
+
+    @classmethod
+    def data_version(cls) -> dict[str, str]:
+        """Get the code-level data version for this feature.
+
+        This returns a static hash based on code versions and dependencies,
+        not sample-level data versions.
+
+        Returns:
+            Dictionary mapping field keys to their data version hashes.
+        """
+        return cls.graph.get_feature_version_by_field(cls.spec.key)
+
+    @classmethod
+    def load_input(
+        cls,
+        joiner: "UpstreamJoiner",
+        upstream_refs: dict[str, "nw.LazyFrame[Any]"],
+    ) -> tuple["nw.LazyFrame[Any]", dict[str, str]]:
+        """Join upstream feature metadata.
+
+        Override for custom join logic (1:many, different keys, filtering, etc.).
+
+        Args:
+            joiner: UpstreamJoiner from MetadataStore
+            upstream_refs: Upstream feature metadata references (lazy where possible)
+
+        Returns:
+            (joined_upstream, upstream_column_mapping)
+            - joined_upstream: All upstream data joined together
+            - upstream_column_mapping: Maps upstream_key -> column name
+        """
+        from metaxy.models.feature_spec import FeatureDep
+
+        # Extract columns and renames from deps
+        upstream_columns: dict[str, tuple[str, ...] | None] = {}
+        upstream_renames: dict[str, dict[str, str] | None] = {}
+
+        if cls.spec.deps:
+            for dep in cls.spec.deps:
+                if isinstance(dep, FeatureDep):
+                    dep_key_str = dep.key.to_string()
+                    upstream_columns[dep_key_str] = dep.columns
+                    upstream_renames[dep_key_str] = dep.rename
+
+        return joiner.join_upstream(
+            upstream_refs=upstream_refs,
+            feature_spec=cls.spec,
+            feature_plan=cls.graph.get_feature_plan(cls.spec.key),
+            upstream_columns=upstream_columns,
+            upstream_renames=upstream_renames,
+        )
+
+    @classmethod
+    def resolve_data_version_diff(
+        cls,
+        diff_resolver: "MetadataDiffResolver",
+        target_versions: "nw.LazyFrame[Any]",
+        current_metadata: "nw.LazyFrame[Any] | None",
+        *,
+        lazy: bool = False,
+    ) -> "DiffResult | LazyDiffResult":
+        """Resolve differences between target and current data versions.
+
+        Override for custom diff logic (ignore certain fields, custom rules, etc.).
+
+        Args:
+            diff_resolver: MetadataDiffResolver from MetadataStore
+            target_versions: Calculated target data_versions (Narwhals LazyFrame)
+            current_metadata: Current metadata for this feature (Narwhals LazyFrame, or None).
+                Should be pre-filtered by feature_version at the store level.
+            lazy: If True, return LazyDiffResult. If False, return DiffResult.
+
+        Returns:
+            DiffResult (eager) or LazyDiffResult (lazy) with added, changed, removed
+
+        Example (default):
+            >>> class MyFeature(Feature, spec=...):
+            ...     pass  # Uses diff resolver's default implementation
+
+        Example (ignore certain field changes):
+            >>> class MyFeature(Feature, spec=...):
+            ...     @classmethod
+            ...     def resolve_data_version_diff(cls, diff_resolver, target_versions, current_metadata, **kwargs):
+            ...         # Get standard diff
+            ...         result = diff_resolver.find_changes(target_versions, current_metadata)
+            ...
+            ...         # Custom: Only consider 'frames' field changes, ignore 'audio'
+            ...         # Users can filter/modify the diff result here
+            ...
+            ...         return result  # Return modified DiffResult
+        """
+        # Diff resolver always returns LazyDiffResult - materialize if needed
+        lazy_result = diff_resolver.find_changes(
+            target_versions=target_versions,
+            current_metadata=current_metadata,
+        )
+
+        # Materialize to DiffResult if lazy=False
+        if not lazy:
+            from metaxy.data_versioning.diff import DiffResult
+
+            return DiffResult(
+                added=lazy_result.added.collect(),
+                changed=lazy_result.changed.collect(),
+                removed=lazy_result.removed.collect(),
+            )
+
+        return lazy_result