refactor: align target multi-part primary key representation with source

python/cocoindex/convert.py

@@ -9,7 +9,7 @@
 import inspect
 import warnings
 from enum import Enum
-from typing import Any, Callable, Mapping, Type, get_origin
+from typing import Any, Callable, Mapping, Sequence, Type, get_origin
 
 import numpy as np
 
@@ -170,6 +170,37 @@ def encode_basic_value(value: Any) -> Any:
     return encode_basic_value
 
 
+def make_engine_key_decoder(
+    field_path: list[str],
+    key_fields_schema: list[dict[str, Any]],
+    dst_type_info: AnalyzedTypeInfo,
+) -> Callable[[Any], Any]:
+    """
+    Create an encoder closure for a key type.
+    """
+    if len(key_fields_schema) == 1 and isinstance(
+        dst_type_info.variant, (AnalyzedBasicType, AnalyzedAnyType)
+    ):
+        single_key_decoder = make_engine_value_decoder(
+            field_path,
+            key_fields_schema[0]["type"],
+            dst_type_info,
+            for_key=True,
+        )
+
+        def key_decoder(value: list[Any]) -> Any:
+            return single_key_decoder(value[0])
+
+        return key_decoder
+
+    return make_engine_struct_decoder(
+        field_path,
+        key_fields_schema,
+        dst_type_info,
+        for_key=True,
+    )
+
+
 def make_engine_value_decoder(
     field_path: list[str],
     src_type: dict[str, Any],
@@ -244,31 +275,11 @@ def decode(value: Any) -> Any | None:
                     )
 
                 num_key_parts = src_type.get("num_key_parts", 1)
-                key_type_info = analyze_type_info(key_type)
-                key_decoder: Callable[..., Any] | None = None
-                if (
-                    isinstance(
-                        key_type_info.variant, (AnalyzedBasicType, AnalyzedAnyType)
-                    )
-                    and num_key_parts == 1
-                ):
-                    single_key_decoder = make_engine_value_decoder(
-                        field_path,
-                        engine_fields_schema[0]["type"],
-                        key_type_info,
-                        for_key=True,
-                    )
-
-                    def key_decoder(value: list[Any]) -> Any:
-                        return single_key_decoder(value[0])
-
-                else:
-                    key_decoder = make_engine_struct_decoder(
-                        field_path,
-                        engine_fields_schema[0:num_key_parts],
-                        key_type_info,
-                        for_key=True,
-                    )
+                key_decoder = make_engine_key_decoder(
+                    field_path,
+                    engine_fields_schema[0:num_key_parts],
+                    analyze_type_info(key_type),
+                )
                 value_decoder = make_engine_struct_decoder(
                     field_path,
                     engine_fields_schema[num_key_parts:],

python/cocoindex/op.py

@@ -21,6 +21,7 @@
 from .convert import (
     make_engine_value_encoder,
     make_engine_value_decoder,
+    make_engine_key_decoder,
     make_engine_struct_decoder,
 )
 from .typing import (
@@ -29,7 +30,6 @@
     resolve_forward_ref,
     analyze_type_info,
     AnalyzedAnyType,
-    AnalyzedBasicType,
     AnalyzedDictType,
 )
 
@@ -532,24 +532,9 @@ def create_export_context(
             else (Any, Any)
         )
 
-        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,
-                for_key=True,
-            )
-        else:
-            key_decoder = make_engine_struct_decoder(
-                ["(key)"], key_fields_schema, key_type_info, for_key=True
-            )
-
+        key_decoder = make_engine_key_decoder(
+            ["(key)"], key_fields_schema, analyze_type_info(key_annotation)
+        )
         value_decoder = make_engine_struct_decoder(
             ["(value)"], value_fields_schema, analyze_type_info(value_annotation)
         )