cocoindex 0.1.72__cp312-cp312-manylinux_2_28_aarch64.whl → 0.1.74__cp312-cp312-manylinux_2_28_aarch64.whl

cocoindex/op.py

@@ -6,11 +6,31 @@ import asyncio
 import dataclasses
 import inspect
 from enum import Enum
-from typing import Any, Awaitable, Callable, Protocol, dataclass_transform, Annotated
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Protocol,
+    dataclass_transform,
+    Annotated,
+    get_args,
+)
 
 from . import _engine  # type: ignore
-from .convert import encode_engine_value, make_engine_value_decoder
-from .typing import TypeAttr, encode_enriched_type, resolve_forward_ref
+from .convert import (
+    encode_engine_value,
+    make_engine_value_decoder,
+    make_engine_struct_decoder,
+)
+from .typing import (
+    TypeAttr,
+    encode_enriched_type,
+    resolve_forward_ref,
+    analyze_type_info,
+    AnalyzedAnyType,
+    AnalyzedBasicType,
+    AnalyzedDictType,
+)
 
 
 class OpCategory(Enum):
@@ -65,6 +85,22 @@ class Executor(Protocol):
     op_category: OpCategory
 
 
+def _load_spec_from_engine(spec_cls: type, spec: dict[str, Any]) -> Any:
+    """
+    Load a spec from the engine.
+    """
+    return spec_cls(**spec)
+
+
+def _get_required_method(cls: type, name: str) -> Callable[..., Any]:
+    method = getattr(cls, name, None)
+    if method is None:
+        raise ValueError(f"Method {name}() is required for {cls.__name__}")
+    if not inspect.isfunction(method):
+        raise ValueError(f"Method {cls.__name__}.{name}() is not a function")
+    return method
+
+
 class _FunctionExecutorFactory:
     _spec_cls: type
     _executor_cls: type
@@ -76,7 +112,7 @@ class _FunctionExecutorFactory:
     def __call__(
         self, spec: dict[str, Any], *args: Any, **kwargs: Any
     ) -> tuple[dict[str, Any], Executor]:
-        spec = self._spec_cls(**spec)
+        spec = _load_spec_from_engine(self._spec_cls, spec)
         executor = self._executor_cls(spec)
         result_type = executor.analyze(*args, **kwargs)
         return (encode_enriched_type(result_type), executor)
@@ -185,7 +221,9 @@ def _register_op_factory(
                     )
                 self._args_decoders.append(
                     make_engine_value_decoder(
-                        [arg_name], arg.value_type["type"], arg_param.annotation
+                        [arg_name],
+                        arg.value_type["type"],
+                        analyze_type_info(arg_param.annotation),
                     )
                 )
                 process_attribute(arg_name, arg)
@@ -217,7 +255,9 @@ def _register_op_factory(
                     )
                 arg_param = expected_arg[1]
                 self._kwargs_decoders[kwarg_name] = make_engine_value_decoder(
-                    [kwarg_name], kwarg.value_type["type"], arg_param.annotation
+                    [kwarg_name],
+                    kwarg.value_type["type"],
+                    analyze_type_info(arg_param.annotation),
                 )
                 process_attribute(kwarg_name, kwarg)
 
@@ -359,3 +399,220 @@ def function(**args: Any) -> Callable[[Callable[..., Any]], FunctionSpec]:
         return _Spec()
 
     return _inner
+
+
+########################################################
+# Custom target connector
+########################################################
+
+
+@dataclasses.dataclass
+class _TargetConnectorContext:
+    target_name: str
+    spec: Any
+    prepared_spec: Any
+    key_decoder: Callable[[Any], Any]
+    value_decoder: Callable[[Any], Any]
+
+
+class _TargetConnector:
+    """
+    The connector class passed to the engine.
+    """
+
+    _spec_cls: type
+    _connector_cls: type
+
+    _get_persistent_key_fn: Callable[[_TargetConnectorContext, str], Any]
+    _apply_setup_change_async_fn: Callable[
+        [Any, dict[str, Any] | None, dict[str, Any] | None], Awaitable[None]
+    ]
+    _mutate_async_fn: Callable[..., Awaitable[None]]
+    _mutatation_type: AnalyzedDictType | None
+
+    def __init__(self, spec_cls: type, connector_cls: type):
+        self._spec_cls = spec_cls
+        self._connector_cls = connector_cls
+
+        self._get_persistent_key_fn = _get_required_method(
+            connector_cls, "get_persistent_key"
+        )
+        self._apply_setup_change_async_fn = _to_async_call(
+            _get_required_method(connector_cls, "apply_setup_change")
+        )
+
+        mutate_fn = _get_required_method(connector_cls, "mutate")
+        self._mutate_async_fn = _to_async_call(mutate_fn)
+
+        # Store the type annotation for later use
+        self._mutatation_type = self._analyze_mutate_mutation_type(
+            connector_cls, mutate_fn
+        )
+
+    @staticmethod
+    def _analyze_mutate_mutation_type(
+        connector_cls: type, mutate_fn: Callable[..., Any]
+    ) -> AnalyzedDictType | None:
+        # Validate mutate_fn signature and extract type annotation
+        mutate_sig = inspect.signature(mutate_fn)
+        params = list(mutate_sig.parameters.values())
+
+        if len(params) != 1:
+            raise ValueError(
+                f"Method {connector_cls.__name__}.mutate(*args) must have exactly one parameter, "
+                f"got {len(params)}"
+            )
+
+        param = params[0]
+        if param.kind != inspect.Parameter.VAR_POSITIONAL:
+            raise ValueError(
+                f"Method {connector_cls.__name__}.mutate(*args) parameter must be *args format, "
+                f"got {param.kind.name}"
+            )
+
+        # Extract type annotation
+        analyzed_args_type = analyze_type_info(param.annotation)
+        if isinstance(analyzed_args_type.variant, AnalyzedAnyType):
+            return None
+
+        if analyzed_args_type.base_type is tuple:
+            args = get_args(analyzed_args_type.core_type)
+            if not args:
+                return None
+            if len(args) == 2:
+                mutation_type = analyze_type_info(args[1])
+                if isinstance(mutation_type.variant, AnalyzedAnyType):
+                    return None
+                if isinstance(mutation_type.variant, AnalyzedDictType):
+                    return mutation_type.variant
+
+        raise ValueError(
+            f"Method {connector_cls.__name__}.mutate(*args) parameter must be a tuple with "
+            f"2 elements (tuple[SpecType, dict[str, ValueStruct]], spec and mutation in dict), "
+            "got {args_type}"
+        )
+
+    def create_export_context(
+        self,
+        name: str,
+        spec: dict[str, Any],
+        key_fields_schema: list[Any],
+        value_fields_schema: list[Any],
+    ) -> _TargetConnectorContext:
+        key_annotation, value_annotation = (
+            (
+                self._mutatation_type.key_type,
+                self._mutatation_type.value_type,
+            )
+            if self._mutatation_type is not None
+            else (None, None)
+        )
+
+        key_type_info = analyze_type_info(key_annotation)
+        if (
+            len(key_fields_schema) == 1
+            and key_fields_schema[0]["type"]["kind"] != "Struct"
+            and isinstance(key_type_info.variant, (AnalyzedAnyType, AnalyzedBasicType))
+        ):
+            # Special case for ease of use: single key column can be mapped to a basic type without the wrapper struct.
+            key_decoder = make_engine_value_decoder(
+                ["(key)"],
+                key_fields_schema[0]["type"],
+                key_type_info,
+            )
+        else:
+            key_decoder = make_engine_struct_decoder(
+                ["(key)"], key_fields_schema, key_type_info
+            )
+
+        value_decoder = make_engine_struct_decoder(
+            ["(value)"], value_fields_schema, analyze_type_info(value_annotation)
+        )
+
+        loaded_spec = _load_spec_from_engine(self._spec_cls, spec)
+        prepare_method = getattr(self._connector_cls, "prepare", None)
+        if prepare_method is None:
+            prepared_spec = loaded_spec
+        else:
+            prepared_spec = prepare_method(loaded_spec)
+
+        return _TargetConnectorContext(
+            target_name=name,
+            spec=loaded_spec,
+            prepared_spec=prepared_spec,
+            key_decoder=key_decoder,
+            value_decoder=value_decoder,
+        )
+
+    def get_persistent_key(self, export_context: _TargetConnectorContext) -> Any:
+        return self._get_persistent_key_fn(
+            export_context.spec, export_context.target_name
+        )
+
+    def describe_resource(self, key: Any) -> str:
+        describe_fn = getattr(self._connector_cls, "describe", None)
+        if describe_fn is None:
+            return str(key)
+        return str(describe_fn(key))
+
+    async def apply_setup_changes_async(
+        self,
+        changes: list[tuple[Any, list[dict[str, Any] | None], dict[str, Any] | None]],
+    ) -> None:
+        for key, previous, current in changes:
+            prev_specs = [
+                _load_spec_from_engine(self._spec_cls, spec)
+                if spec is not None
+                else None
+                for spec in previous
+            ]
+            curr_spec = (
+                _load_spec_from_engine(self._spec_cls, current)
+                if current is not None
+                else None
+            )
+            for prev_spec in prev_specs:
+                await self._apply_setup_change_async_fn(key, prev_spec, curr_spec)
+
+    @staticmethod
+    def _decode_mutation(
+        context: _TargetConnectorContext, mutation: list[tuple[Any, Any | None]]
+    ) -> tuple[Any, dict[Any, Any | None]]:
+        return (
+            context.prepared_spec,
+            {
+                context.key_decoder(key): (
+                    context.value_decoder(value) if value is not None else None
+                )
+                for key, value in mutation
+            },
+        )
+
+    async def mutate_async(
+        self,
+        mutations: list[tuple[_TargetConnectorContext, list[tuple[Any, Any | None]]]],
+    ) -> None:
+        await self._mutate_async_fn(
+            *(
+                self._decode_mutation(context, mutation)
+                for context, mutation in mutations
+            )
+        )
+
+
+def target_connector(spec_cls: type) -> Callable[[type], type]:
+    """
+    Decorate a class to provide a target connector for an op.
+    """
+
+    # Validate the spec_cls is a TargetSpec.
+    if not issubclass(spec_cls, TargetSpec):
+        raise ValueError(f"Expect a TargetSpec, got {spec_cls}")
+
+    # Register the target connector.
+    def _inner(connector_cls: type) -> type:
+        connector = _TargetConnector(spec_cls, connector_cls)
+        _engine.register_target_connector(spec_cls.__name__, connector)
+        return connector_cls
+
+    return _inner

cocoindex/tests/test_convert.py

@@ -1,6 +1,6 @@
 import datetime
 import uuid
-from dataclasses import dataclass, make_dataclass
+from dataclasses import dataclass, make_dataclass, field
 from typing import Annotated, Any, Callable, Literal, NamedTuple
 
 import numpy as np
@@ -19,6 +19,7 @@ from cocoindex.typing import (
     TypeKind,
     Vector,
     encode_enriched_type,
+    analyze_type_info,
 )
 
 
@@ -75,7 +76,9 @@ def build_engine_value_decoder(
     If python_type is not specified, uses engine_type_in_py as the target.
     """
     engine_type = encode_enriched_type(engine_type_in_py)["type"]
-    return make_engine_value_decoder([], engine_type, python_type or engine_type_in_py)
+    return make_engine_value_decoder(
+        [], engine_type, analyze_type_info(python_type or engine_type_in_py)
+    )
 
 
 def validate_full_roundtrip_to(
@@ -103,9 +106,13 @@ def validate_full_roundtrip_to(
     )
 
     for other_value, other_type in decoded_values:
-        decoder = make_engine_value_decoder([], encoded_output_type, other_type)
+        decoder = make_engine_value_decoder(
+            [], encoded_output_type, analyze_type_info(other_type)
+        )
         other_decoded_value = decoder(value_from_engine)
-        assert eq(other_decoded_value, other_value)
+        assert eq(other_decoded_value, other_value), (
+            f"Expected {other_value} but got {other_decoded_value} for {other_type}"
+        )
 
 
 def validate_full_roundtrip(
@@ -362,7 +369,9 @@ def test_decode_scalar_numpy_values() -> None:
         ({"kind": "Float64"}, np.float64, 2.718, np.float64(2.718)),
     ]
     for src_type, dst_type, input_value, expected in test_cases:
-        decoder = make_engine_value_decoder(["field"], src_type, dst_type)
+        decoder = make_engine_value_decoder(
+            ["field"], src_type, analyze_type_info(dst_type)
+        )
         result = decoder(input_value)
         assert isinstance(result, dst_type)
         assert result == expected
@@ -376,7 +385,9 @@ def test_non_ndarray_vector_decoding() -> None:
         "dimension": None,
     }
     dst_type_float = list[np.float64]
-    decoder = make_engine_value_decoder(["field"], src_type, dst_type_float)
+    decoder = make_engine_value_decoder(
+        ["field"], src_type, analyze_type_info(dst_type_float)
+    )
     input_numbers = [1.0, 2.0, 3.0]
     result = decoder(input_numbers)
     assert isinstance(result, list)
@@ -386,7 +397,9 @@ def test_non_ndarray_vector_decoding() -> None:
     # Test list[Uuid]
     src_type = {"kind": "Vector", "element_type": {"kind": "Uuid"}, "dimension": None}
     dst_type_uuid = list[uuid.UUID]
-    decoder = make_engine_value_decoder(["field"], src_type, dst_type_uuid)
+    decoder = make_engine_value_decoder(
+        ["field"], src_type, analyze_type_info(dst_type_uuid)
+    )
     uuid1 = uuid.uuid4()
     uuid2 = uuid.uuid4()
     input_uuids = [uuid1, uuid2]
@@ -396,124 +409,15 @@ def test_non_ndarray_vector_decoding() -> None:
     assert result == [uuid1, uuid2]
 
 
-@pytest.mark.parametrize(
-    "data_type, engine_val, expected",
-    [
-        # All fields match (dataclass)
-        (
-            Order,
-            ["O123", "mixed nuts", 25.0, "default_extra"],
-            Order("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        # All fields match (NamedTuple)
-        (
-            OrderNamedTuple,
-            ["O123", "mixed nuts", 25.0, "default_extra"],
-            OrderNamedTuple("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        # Extra field in engine value (should ignore extra)
-        (
-            Order,
-            ["O123", "mixed nuts", 25.0, "default_extra", "unexpected"],
-            Order("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        (
-            OrderNamedTuple,
-            ["O123", "mixed nuts", 25.0, "default_extra", "unexpected"],
-            OrderNamedTuple("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        # Fewer fields in engine value (should fill with default)
-        (
-            Order,
-            ["O123", "mixed nuts", 0.0, "default_extra"],
-            Order("O123", "mixed nuts", 0.0, "default_extra"),
-        ),
-        (
-            OrderNamedTuple,
-            ["O123", "mixed nuts", 0.0, "default_extra"],
-            OrderNamedTuple("O123", "mixed nuts", 0.0, "default_extra"),
-        ),
-        # More fields in engine value (should ignore extra)
-        (
-            Order,
-            ["O123", "mixed nuts", 25.0, "unexpected"],
-            Order("O123", "mixed nuts", 25.0, "unexpected"),
-        ),
-        (
-            OrderNamedTuple,
-            ["O123", "mixed nuts", 25.0, "unexpected"],
-            OrderNamedTuple("O123", "mixed nuts", 25.0, "unexpected"),
-        ),
-        # Truly extra field (should ignore the fifth field)
-        (
-            Order,
-            ["O123", "mixed nuts", 25.0, "default_extra", "ignored"],
-            Order("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        (
-            OrderNamedTuple,
-            ["O123", "mixed nuts", 25.0, "default_extra", "ignored"],
-            OrderNamedTuple("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        # Missing optional field in engine value (tags=None)
-        (
-            Customer,
-            ["Alice", ["O1", "item1", 10.0, "default_extra"], None],
-            Customer("Alice", Order("O1", "item1", 10.0, "default_extra"), None),
-        ),
-        (
-            CustomerNamedTuple,
-            ["Alice", ["O1", "item1", 10.0, "default_extra"], None],
-            CustomerNamedTuple(
-                "Alice", OrderNamedTuple("O1", "item1", 10.0, "default_extra"), None
-            ),
-        ),
-        # Extra field in engine value for Customer (should ignore)
-        (
-            Customer,
-            ["Alice", ["O1", "item1", 10.0, "default_extra"], [["vip"]], "extra"],
-            Customer(
-                "Alice", Order("O1", "item1", 10.0, "default_extra"), [Tag("vip")]
-            ),
-        ),
-        (
-            CustomerNamedTuple,
-            ["Alice", ["O1", "item1", 10.0, "default_extra"], [["vip"]], "extra"],
-            CustomerNamedTuple(
-                "Alice",
-                OrderNamedTuple("O1", "item1", 10.0, "default_extra"),
-                [Tag("vip")],
-            ),
-        ),
-        # Missing optional field with default
-        (
-            Order,
-            ["O123", "mixed nuts", 25.0],
-            Order("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        (
-            OrderNamedTuple,
-            ["O123", "mixed nuts", 25.0],
-            OrderNamedTuple("O123", "mixed nuts", 25.0, "default_extra"),
-        ),
-        # Partial optional fields
-        (
-            Customer,
-            ["Alice", ["O1", "item1", 10.0]],
-            Customer("Alice", Order("O1", "item1", 10.0, "default_extra"), None),
-        ),
-        (
-            CustomerNamedTuple,
-            ["Alice", ["O1", "item1", 10.0]],
-            CustomerNamedTuple(
-                "Alice", OrderNamedTuple("O1", "item1", 10.0, "default_extra"), None
-            ),
-        ),
-    ],
-)
-def test_struct_decoder_cases(data_type: Any, engine_val: Any, expected: Any) -> None:
-    decoder = build_engine_value_decoder(data_type)
-    assert decoder(engine_val) == expected
+def test_roundtrip_struct() -> None:
+    validate_full_roundtrip(
+        Order("O123", "mixed nuts", 25.0, "default_extra"),
+        Order,
+    )
+    validate_full_roundtrip(
+        OrderNamedTuple("O123", "mixed nuts", 25.0, "default_extra"),
+        OrderNamedTuple,
+    )
 
 
 def test_make_engine_value_decoder_list_of_struct() -> None:
@@ -972,7 +876,9 @@ def test_decode_nullable_ndarray_none_or_value_input() -> None:
         "dimension": None,
     }
     dst_annotation = NDArrayFloat32Type | None
-    decoder = make_engine_value_decoder([], src_type_dict, dst_annotation)
+    decoder = make_engine_value_decoder(
+        [], src_type_dict, analyze_type_info(dst_annotation)
+    )
 
     none_engine_value = None
     decoded_array = decoder(none_engine_value)
@@ -995,7 +901,9 @@ def test_decode_vector_string() -> None:
         "element_type": {"kind": "Str"},
         "dimension": None,
     }
-    decoder = make_engine_value_decoder([], src_type_dict, Vector[str])
+    decoder = make_engine_value_decoder(
+        [], src_type_dict, analyze_type_info(Vector[str])
+    )
     assert decoder(["hello", "world"]) == ["hello", "world"]
 
 
@@ -1006,7 +914,9 @@ def test_decode_error_non_nullable_or_non_list_vector() -> None:
         "element_type": {"kind": "Float32"},
         "dimension": None,
     }
-    decoder = make_engine_value_decoder([], src_type_dict, NDArrayFloat32Type)
+    decoder = make_engine_value_decoder(
+        [], src_type_dict, analyze_type_info(NDArrayFloat32Type)
+    )
     with pytest.raises(ValueError, match="Received null for non-nullable vector"):
         decoder(None)
     with pytest.raises(TypeError, match="Expected NDArray or list for vector"):
@@ -1096,6 +1006,25 @@ def test_full_roundtrip_vector_numeric_types() -> None:
         validate_full_roundtrip(value_u64, Vector[np.uint64, Literal[3]])
 
 
+def test_full_roundtrip_vector_of_vector() -> None:
+    """Test full roundtrip for vector of vector."""
+    value_f32 = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
+    validate_full_roundtrip(
+        value_f32,
+        Vector[Vector[np.float32, Literal[3]], Literal[2]],
+        ([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], list[list[np.float32]]),
+        ([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], list[list[cocoindex.Float32]]),
+        (
+            [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+            list[Vector[cocoindex.Float32, Literal[3]]],
+        ),
+        (
+            value_f32,
+            np.typing.NDArray[np.typing.NDArray[np.float32]],
+        ),
+    )
+
+
 def test_full_roundtrip_vector_other_types() -> None:
     """Test full roundtrip for Vector with non-numeric basic types."""
     uuid_list = [uuid.uuid4(), uuid.uuid4()]
@@ -1216,7 +1145,7 @@ def test_full_roundtrip_scalar_with_python_types() -> None:
         numpy_float: np.float64
         python_float: float
         string: str
-        annotated_int: Annotated[np.int64, TypeKind("int")]
+        annotated_int: Annotated[np.int64, TypeKind("Int64")]
         annotated_float: Float32
 
     instance = MixedStruct(
@@ -1468,3 +1397,41 @@ def test_roundtrip_ktable_with_list_fields() -> None:
 
     # Test Any annotation
     validate_full_roundtrip(teams, dict[str, Team], (expected_dict_dict, Any))
+
+
+def test_auto_default_for_supported_and_unsupported_types() -> None:
+    @dataclass
+    class Base:
+        a: int
+
+    @dataclass
+    class NullableField:
+        a: int
+        b: int | None
+
+    @dataclass
+    class LTableField:
+        a: int
+        b: list[Base]
+
+    @dataclass
+    class KTableField:
+        a: int
+        b: dict[str, Base]
+
+    @dataclass
+    class UnsupportedField:
+        a: int
+        b: int
+
+    validate_full_roundtrip(NullableField(1, None), NullableField)
+
+    validate_full_roundtrip(LTableField(1, []), LTableField)
+
+    # validate_full_roundtrip(KTableField(1, {}), KTableField)
+
+    with pytest.raises(
+        ValueError,
+        match=r"Field 'b' \(type <class 'int'>\) without default value is missing in input: ",
+    ):
+        build_engine_value_decoder(Base, UnsupportedField)