metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/models/fields_mapping.py

@@ -0,0 +1,307 @@
+"""Field mapping system for automatic field dependency resolution.
+
+This module provides a flexible system for defining how fields map to upstream
+dependencies, supporting both automatic mapping patterns and explicit configurations.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING, Literal
+
+from pydantic import BaseModel, ConfigDict, TypeAdapter
+from pydantic import Field as PydanticField
+from typing_extensions import Self
+
+from metaxy.models.types import (
+    CoercibleToFieldKey,
+    FeatureKey,
+    FieldKey,
+    ValidatedFieldKeyAdapter,
+)
+
+if TYPE_CHECKING:
+    from metaxy.models.feature_spec import FeatureSpec
+
+
+class FieldsMappingType(str, Enum):
+    """Type of fields mapping between a field key and the upstream field keys."""
+
+    DEFAULT = "default"
+    SPECIFIC = "specific"
+    ALL = "all"
+    NONE = "none"
+
+
+class FieldsMappingResolutionContext(BaseModel):
+    """Context for resolving field mappings.
+
+    This contains all the information needed to resolve field dependencies
+    including the upstream feature being mapped against.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    field_key: FieldKey
+    """The downstream field key being resolved."""
+
+    upstream_feature: FeatureSpec
+    """The upstream feature spec being resolved against."""
+
+    @property
+    def upstream_feature_key(self) -> FeatureKey:
+        """Get the upstream feature key."""
+        return self.upstream_feature.key
+
+    @property
+    def upstream_feature_fields(self) -> set[FieldKey]:
+        """Get the set of field keys from the upstream feature."""
+        return {field.key for field in self.upstream_feature.fields}
+
+
+class BaseFieldsMapping(BaseModel, ABC):  # pyright: ignore[reportUnsafeMultipleInheritance]
+    """Base class for field mapping configurations.
+
+    Field mappings define how a field automatically resolves its dependencies
+    based on upstream feature fields.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    @abstractmethod
+    def resolve_field_deps(
+        self,
+        context: FieldsMappingResolutionContext,
+    ) -> set[FieldKey]:
+        """Resolve automatic field mapping to explicit FieldDep list.
+
+        This method should be overridden by concrete implementations.
+
+        Arguments:
+            context: The resolution context containing field key and upstream feature.
+
+        Returns:
+            Set of [FieldKey][metaxy.models.types.FieldKey] instances for matching fields
+        """
+        raise NotImplementedError
+
+
+class SpecificFieldsMapping(BaseFieldsMapping):
+    """Field mapping that explicitly depends on specific upstream fields."""
+
+    type: Literal[FieldsMappingType.SPECIFIC] = FieldsMappingType.SPECIFIC
+    mapping: dict[FieldKey, set[FieldKey]]
+
+    def resolve_field_deps(
+        self,
+        context: FieldsMappingResolutionContext,
+    ) -> set[FieldKey]:
+        desired_upstream_fields = self.mapping.get(context.field_key, set())
+        return desired_upstream_fields & context.upstream_feature_fields
+
+
+class AllFieldsMapping(BaseFieldsMapping):
+    """Field mapping that explicitly depends on all upstream fields."""
+
+    type: Literal[FieldsMappingType.ALL] = FieldsMappingType.ALL
+
+    def resolve_field_deps(
+        self,
+        context: FieldsMappingResolutionContext,
+    ) -> set[FieldKey]:
+        return context.upstream_feature_fields
+
+
+class NoneFieldsMapping(BaseFieldsMapping):
+    """Field mapping that never matches any upstream fields."""
+
+    type: Literal[FieldsMappingType.NONE] = FieldsMappingType.NONE
+
+    def resolve_field_deps(
+        self,
+        context: FieldsMappingResolutionContext,
+    ) -> set[FieldKey]:
+        return set()
+
+
+class DefaultFieldsMapping(BaseFieldsMapping):
+    """Default automatic field mapping configuration.
+
+    When used, automatically maps fields to matching upstream fields based on field keys.
+
+    Attributes:
+        match_suffix: If True, allows suffix matching (e.g., "french" matches "audio/french")
+        exclude_fields: List of field keys to exclude from auto-mapping
+    """
+
+    type: Literal[FieldsMappingType.DEFAULT] = FieldsMappingType.DEFAULT
+    match_suffix: bool = False
+    exclude_fields: list[FieldKey] = PydanticField(default_factory=list)
+
+    def resolve_field_deps(
+        self,
+        context: FieldsMappingResolutionContext,
+    ) -> set[FieldKey]:
+        res = set()
+
+        for upstream_field_key in context.upstream_feature_fields:
+            # Skip excluded fields
+            if upstream_field_key in self.exclude_fields:
+                continue
+
+            # Check for exact match
+            if upstream_field_key == context.field_key:
+                res.add(upstream_field_key)
+            # Check for suffix match if enabled
+            elif self.match_suffix and self._is_suffix_match(
+                context.field_key, upstream_field_key
+            ):
+                res.add(upstream_field_key)
+
+        # If no fields matched, return ALL fields from this upstream feature
+        # (excluding any explicitly excluded fields)
+        if not res:
+            for upstream_field_key in context.upstream_feature_fields:
+                if upstream_field_key not in self.exclude_fields:
+                    res.add(upstream_field_key)
+
+        return res
+
+    def _is_suffix_match(
+        self, field_key: FieldKey, upstream_field_key: FieldKey
+    ) -> bool:
+        """Check if field_key is a suffix of upstream_field_key.
+
+        For hierarchical keys like "audio/french", this checks if "french"
+        matches the suffix.
+
+        Args:
+            field_key: The field key for which to resolve dependencies.
+            upstream_fields_by_feature_key: Mapping of upstream feature keys to their fields.
+
+        Returns:
+            True if field_key is a suffix of upstream_field_key
+        """
+        # For single-part keys, check if it's the last part of a multi-part key
+        if len(field_key.parts) == 1 and len(upstream_field_key.parts) > 1:
+            return field_key.parts[0] == upstream_field_key.parts[-1]
+
+        # For multi-part keys, check if all parts match as suffix
+        if len(field_key.parts) <= len(upstream_field_key.parts):
+            return upstream_field_key.parts[-len(field_key.parts) :] == field_key.parts
+
+        return False
+
+
+class FieldsMapping(BaseModel):
+    """Base class for field mapping configurations.
+
+    Field mappings define how a field automatically resolves its dependencies
+    based on upstream feature fields. This is separate from explicit field
+    dependencies which are defined directly.
+    """
+
+    model_config = ConfigDict(frozen=True)
+    # mapping: BaseFieldsMapping
+    mapping: (
+        AllFieldsMapping
+        | SpecificFieldsMapping
+        | NoneFieldsMapping
+        | DefaultFieldsMapping
+    ) = PydanticField(..., discriminator="type")
+
+    def resolve_field_deps(
+        self,
+        context: FieldsMappingResolutionContext,
+    ) -> set[FieldKey]:
+        """Resolve field dependencies based on upstream feature fields.
+
+        Invokes the provided mapping to resolve dependencies.
+
+        Args:
+            context: The resolution context containing field key and upstream feature.
+
+        Returns:
+            Set of [FieldKey][metaxy.models.types.FieldKey] instances for matching fields
+        """
+        return self.mapping.resolve_field_deps(context)
+
+    @classmethod
+    def default(
+        cls,
+        *,
+        match_suffix: bool = False,
+        exclude_fields: list[FieldKey] | None = None,
+    ) -> Self:
+        """Create a default field mapping configuration.
+
+        Args:
+            match_suffix: If True, allows suffix matching (e.g., "french" matches "audio/french")
+            exclude_fields: List of field keys to exclude from auto-mapping
+
+        Returns:
+            Configured FieldsMapping instance.
+        """
+        return cls(
+            mapping=DefaultFieldsMapping(
+                match_suffix=match_suffix,
+                exclude_fields=exclude_fields or [],
+            )
+        )
+
+    @classmethod
+    def specific(
+        cls, mapping: dict[CoercibleToFieldKey, set[CoercibleToFieldKey]]
+    ) -> Self:
+        """Create a field mapping that maps downstream field keys into specific upstream field keys.
+
+        Args:
+            mapping: Mapping of downstream field keys to sets of upstream field keys.
+                Keys and values can be strings, sequences, or FieldKey instances.
+
+        Returns:
+            Configured FieldsMapping instance.
+        """
+        # Validate and coerce the mapping keys and values
+        validated_mapping: dict[FieldKey, set[FieldKey]] = {}
+        for key, value_set in mapping.items():
+            validated_key = ValidatedFieldKeyAdapter.validate_python(key)
+            validated_values = {
+                ValidatedFieldKeyAdapter.validate_python(v) for v in value_set
+            }
+            validated_mapping[validated_key] = validated_values
+
+        return cls(mapping=SpecificFieldsMapping(mapping=validated_mapping))
+
+    @classmethod
+    def all(cls) -> Self:
+        """Create a field mapping that explicitly depends on all upstream fields.
+
+        Returns:
+            Configured FieldsMapping instance.
+
+        Examples:
+            >>> # Use in field specifications
+            >>> FieldSpec(
+            ...     key="combined",
+            ...     fields_mapping=FieldsMapping.all()
+            ... )
+        """
+        return cls(mapping=AllFieldsMapping())
+
+    @classmethod
+    def none(cls) -> Self:
+        """Create a field mapping that explicitly depends on no upstream fields.
+
+        This is typically useful when explicitly defining [FieldSpec.deps][metaxy.models.field.FieldSpec] instead.
+
+        Returns:
+            Configured FieldsMapping instance.
+        """
+        return cls(mapping=NoneFieldsMapping())
+
+
+FieldsMappingAdapter = TypeAdapter(
+    AllFieldsMapping | SpecificFieldsMapping | NoneFieldsMapping | DefaultFieldsMapping
+)

metaxy/models/filter_expression.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+
+from typing import Any, NamedTuple
+
+import narwhals as nw
+import sqlglot
+from pydantic import field_serializer, model_validator
+from sqlglot import exp
+from sqlglot.errors import ParseError
+
+from metaxy.models.bases import FrozenBaseModel
+
+LiteralValue = bool | int | float | str | None
+
+
+class FilterParseError(ValueError):
+    """Raised when a filter string cannot be parsed into a supported expression."""
+
+
+class OperandInfo(NamedTuple):
+    expr: nw.Expr
+    is_literal: bool
+    literal_value: LiteralValue
+    is_column: bool
+
+
+class NarwhalsFilter(FrozenBaseModel):
+    """Pydantic model for serializable Narwhals filter expressions."""
+
+    expression: sqlglot.exp.Expression
+    source: str | None = None
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+        "extra": "forbid",
+        "frozen": True,
+    }
+
+    @model_validator(mode="before")
+    @classmethod
+    def _parse_expression_from_string(cls, data: Any) -> Any:
+        if isinstance(data, str):
+            expression = _parse_to_sqlglot_expression(data)
+            return {"expression": expression, "source": data}
+        return data
+
+    @field_serializer("expression")
+    def _serialize_expression(self, expression: sqlglot.exp.Expression) -> str:
+        return expression.sql()
+
+    def to_expr(self) -> nw.Expr:
+        """Convert the stored expression into a Narwhals ``Expr``."""
+        return _expression_to_narwhals(self.expression)
+
+
+def parse_filter_string(filter_string: str) -> nw.Expr:
+    """Parse a SQL WHERE-like string into a Narwhals expression.
+
+    The parser understands SQL `WHERE` clauses composed of comparison operators, logical operators, parentheses,
+    dotted identifiers, and literal values (strings, numbers, booleans, ``NULL``).
+
+    This functionality is implemented with [SQLGlot](https://sqlglot.com/).
+
+    Example:
+        ```python
+        parse_filter_string("NOT (status = 'deleted') AND deleted_at = NULL")
+        # Returns: (~(nw.col("status") == "deleted")) & nw.col("deleted_at").is_null()
+        ```
+    """
+    return NarwhalsFilter.model_validate(filter_string).to_expr()
+
+
+def _parse_to_sqlglot_expression(filter_string: str) -> sqlglot.exp.Expression:
+    if not filter_string or not filter_string.strip():
+        raise FilterParseError("Filter string cannot be empty.")
+
+    try:
+        parsed = sqlglot.parse_one(filter_string)
+    except ParseError as exc:
+        msg = f"Failed to parse filter string: {exc}"
+        raise FilterParseError(msg) from exc
+
+    if parsed is None:
+        raise FilterParseError(
+            f"Failed to parse filter string into an expression for {filter_string}"
+        )
+
+    return parsed
+
+
+def _expression_to_narwhals(node: exp.Expression) -> nw.Expr:
+    """Convert a SQLGlot expression AST node to a Narwhals expression."""
+    node = _strip_parens(node)
+
+    # Logical operators
+    if isinstance(node, exp.Not):
+        operand = node.this
+        if operand is None:
+            raise FilterParseError("NOT operator requires an operand.")
+        return ~_expression_to_narwhals(operand)
+
+    if isinstance(node, exp.And):
+        return _expression_to_narwhals(node.this) & _expression_to_narwhals(
+            node.expression
+        )
+
+    if isinstance(node, exp.Or):
+        return _expression_to_narwhals(node.this) | _expression_to_narwhals(
+            node.expression
+        )
+
+    # Comparison operators - direct mapping to Narwhals operations
+    if isinstance(node, (exp.EQ, exp.NEQ, exp.GT, exp.LT, exp.GTE, exp.LTE)):
+        left = getattr(node, "this", None)
+        right = getattr(node, "expression", None)
+        if left is None or right is None:
+            raise FilterParseError(
+                f"Comparison operator {type(node).__name__} requires two operands."
+            )
+        left_operand = _operand_info(left)
+        right_operand = _operand_info(right)
+
+        # Handle NULL comparisons with IS NULL / IS NOT NULL
+        null_comparison = _maybe_null_comparison(left_operand, right_operand, node)
+        if null_comparison is not None:
+            return null_comparison
+
+        # Apply the appropriate Narwhals operator
+        if isinstance(node, exp.EQ):
+            return left_operand.expr == right_operand.expr
+        elif isinstance(node, exp.NEQ):
+            return left_operand.expr != right_operand.expr
+        elif isinstance(node, exp.GT):
+            return left_operand.expr > right_operand.expr
+        elif isinstance(node, exp.LT):
+            return left_operand.expr < right_operand.expr
+        elif isinstance(node, exp.GTE):
+            return left_operand.expr >= right_operand.expr
+        elif isinstance(node, exp.LTE):
+            return left_operand.expr <= right_operand.expr
+
+    # Terminal nodes (operands)
+    if isinstance(
+        node,
+        (
+            exp.Column,
+            exp.Identifier,
+            exp.Boolean,
+            exp.Literal,
+            exp.Null,
+            exp.Neg,
+        ),
+    ):
+        return _operand_info(node).expr
+
+    raise FilterParseError(f"Unsupported expression: {node.sql()}")
+
+
+def _operand_info(node: exp.Expression) -> OperandInfo:
+    """Extract operand information from a SQLGlot expression node."""
+    node = _strip_parens(node)
+
+    if isinstance(node, exp.Column):
+        return OperandInfo(
+            expr=nw.col(_column_name(node)),
+            is_literal=False,
+            literal_value=None,
+            is_column=True,
+        )
+
+    if isinstance(node, exp.Identifier):
+        return OperandInfo(
+            expr=nw.col(_column_name(node)),
+            is_literal=False,
+            literal_value=None,
+            is_column=True,
+        )
+
+    if isinstance(node, exp.Neg):
+        inner = node.this
+        if inner is None:
+            raise FilterParseError("Unary minus requires an operand.")
+        operand = _operand_info(inner)
+        if not operand.is_literal or not isinstance(
+            operand.literal_value, (int, float)
+        ):
+            raise FilterParseError("Unary minus only supported for numeric literals.")
+        value = -operand.literal_value
+        return OperandInfo(
+            expr=nw.lit(value), is_literal=True, literal_value=value, is_column=False
+        )
+
+    if isinstance(node, exp.Literal):
+        value = _literal_to_python(node)
+        return OperandInfo(
+            expr=nw.lit(value), is_literal=True, literal_value=value, is_column=False
+        )
+
+    if isinstance(node, exp.Boolean):
+        value = _literal_to_python(node)
+        return OperandInfo(
+            expr=nw.lit(value), is_literal=True, literal_value=value, is_column=False
+        )
+
+    if isinstance(node, exp.Null):
+        return OperandInfo(
+            expr=nw.lit(None), is_literal=True, literal_value=None, is_column=False
+        )
+
+    raise FilterParseError(f"Unsupported operand: {node.sql()}")
+
+
+def _maybe_null_comparison(
+    left: OperandInfo,
+    right: OperandInfo,
+    node: exp.Expression,
+) -> nw.Expr | None:
+    """Handle SQL NULL comparisons, converting to IS NULL / IS NOT NULL."""
+    if left.is_literal and left.literal_value is None and right.is_column:
+        column_expr = right.expr
+        if isinstance(node, exp.EQ):
+            return column_expr.is_null()
+        if isinstance(node, exp.NEQ):
+            return ~column_expr.is_null()
+        return None
+
+    if right.is_literal and right.literal_value is None and left.is_column:
+        column_expr = left.expr
+        if isinstance(node, exp.EQ):
+            return column_expr.is_null()
+        if isinstance(node, exp.NEQ):
+            return ~column_expr.is_null()
+        return None
+
+    return None
+
+
+def _literal_to_python(node: exp.Expression) -> LiteralValue:
+    """Convert a SQLGlot literal node to a Python value."""
+    match node:
+        case exp.Null():
+            return None
+        case exp.Boolean():
+            return node.this is True or str(node.this).lower() == "true"
+        case exp.Literal():
+            literal = node
+            if literal.is_string:
+                return literal.name
+            if literal.is_int:
+                return int(literal.this)
+            if literal.is_number:
+                return float(literal.this)
+            return literal.this
+        case _:
+            raise FilterParseError(f"Unsupported literal: {node.sql()}")
+
+
+def _strip_parens(node: exp.Expression) -> exp.Expression:
+    """Remove surrounding parentheses from an expression."""
+    current = node
+    while isinstance(current, exp.Paren) and current.this is not None:
+        current = current.this
+    return current
+
+
+def _identifier_part_to_string(part: exp.Expression | str) -> str:
+    """Convert a column identifier part to a string."""
+    if isinstance(part, exp.Identifier):
+        return part.name
+    if isinstance(part, exp.Star):
+        return "*"
+    if isinstance(part, exp.Expression):
+        return part.sql(dialect="")
+    return str(part)
+
+
+def _column_name(node: exp.Expression) -> str:
+    """Extract the column name from a Column or Identifier node."""
+    if isinstance(node, exp.Column):
+        parts = [_identifier_part_to_string(part) for part in node.parts or ()]
+        name = ".".join(part for part in parts if part)
+    elif isinstance(node, exp.Identifier):
+        name = node.name
+    else:
+        name = node.sql(dialect="")
+
+    name = name.strip()
+    if not name:
+        raise FilterParseError("Column reference is malformed.")
+    return name
+
+
+__all__ = [
+    "FilterParseError",
+    "NarwhalsFilter",
+    "parse_filter_string",
+]