feat: support dtype decoding by adding np_number_type to AnalyzedTypeInfo

Author: lemorage

python/cocoindex/convert.py

@@ -127,8 +127,7 @@ def decode(value: Any) -> Any | None:
         return lambda value: uuid.UUID(bytes=value)
 
     if src_type_kind == "Vector":
-        elem_coco_type_info = analyze_type_info(dst_type_info.elem_type)
-        dtype_info = DtypeRegistry.get_by_kind(elem_coco_type_info.kind)
+        dtype_info = DtypeRegistry.get_by_dtype(dst_type_info.np_number_type)
 
         def decode_vector(value: Any) -> Any | None:
             if value is None:

python/cocoindex/tests/test_typing.py

@@ -48,6 +48,7 @@ def test_ndarray_float32_no_dim():
         elem_type=Float32,
         key_type=None,
         struct_type=None,
+        np_number_type=np.float32,
         attrs=None,
         nullable=False,
     )
@@ -62,6 +63,7 @@ def test_vector_float32_no_dim():
         elem_type=Float32,
         key_type=None,
         struct_type=None,
+        np_number_type=np.float32,
         attrs=None,
         nullable=False,
     )
@@ -76,6 +78,7 @@ def test_ndarray_float64_with_dim():
         elem_type=Float64,
         key_type=None,
         struct_type=None,
+        np_number_type=np.float64,
         attrs=None,
         nullable=False,
     )
@@ -90,6 +93,7 @@ def test_vector_float32_with_dim():
         elem_type=Float32,
         key_type=None,
         struct_type=None,
+        np_number_type=np.float32,
         attrs=None,
         nullable=False,
     )
@@ -109,7 +113,7 @@ def test_ndarray_int32_with_dim():
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=10)
-    assert get_args(result.elem_type) == (int, TypeKind("Int32"))
+    assert get_args(result.elem_type) == (int, TypeKind("Int64"))
     assert not result.nullable
 
 
@@ -118,7 +122,7 @@ def test_ndarray_uint8_no_dim():
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=None)
-    assert get_args(result.elem_type) == (int, TypeKind("UInt8"))
+    assert get_args(result.elem_type) == (int, TypeKind("Int64"))
     assert not result.nullable
 
 
@@ -131,6 +135,7 @@ def test_nullable_ndarray():
         elem_type=Float32,
         key_type=None,
         struct_type=None,
+        np_number_type=np.float32,
         attrs=None,
         nullable=True,
     )
@@ -177,6 +182,7 @@ def test_list_of_primitives():
         elem_type=str,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -191,6 +197,7 @@ def test_list_of_structs():
         elem_type=SimpleDataclass,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -205,6 +212,7 @@ def test_sequence_of_int():
         elem_type=int,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -219,6 +227,7 @@ def test_list_with_vector_info():
         elem_type=int,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -233,6 +242,7 @@ def test_dict_str_int():
         elem_type=(str, int),
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -247,6 +257,7 @@ def test_mapping_str_dataclass():
         elem_type=(str, SimpleDataclass),
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -261,6 +272,7 @@ def test_dataclass():
         elem_type=None,
         key_type=None,
         struct_type=SimpleDataclass,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -275,6 +287,7 @@ def test_named_tuple():
         elem_type=None,
         key_type=None,
         struct_type=SimpleNamedTuple,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -289,6 +302,7 @@ def test_tuple_key_value():
         elem_type=None,
         key_type=str,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -303,6 +317,7 @@ def test_str():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -317,6 +332,7 @@ def test_bool():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -331,6 +347,7 @@ def test_bytes():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -345,6 +362,7 @@ def test_uuid():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -359,6 +377,7 @@ def test_date():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -373,6 +392,7 @@ def test_time():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -387,6 +407,7 @@ def test_timedelta():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -401,6 +422,7 @@ def test_float():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -415,6 +437,7 @@ def test_int():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs=None,
         nullable=False,
     )
@@ -429,6 +452,7 @@ def test_type_with_attributes():
         elem_type=None,
         key_type=None,
         struct_type=None,
+        np_number_type=None,
         attrs={"key": "value"},
         nullable=False,
     )

python/cocoindex/typing.py

@@ -119,11 +119,11 @@ class DtypeRegistry:
     _mappings: dict[type, DtypeInfo] = {
         np.float32: DtypeInfo(np.float32, "Float32", float),
         np.float64: DtypeInfo(np.float64, "Float64", float),
-        np.int32: DtypeInfo(np.int32, "Int32", int),
+        np.int32: DtypeInfo(np.int32, "Int64", int),
         np.int64: DtypeInfo(np.int64, "Int64", int),
-        np.uint8: DtypeInfo(np.uint8, "UInt8", int),
-        np.uint16: DtypeInfo(np.uint16, "UInt16", int),
-        np.uint32: DtypeInfo(np.uint32, "UInt32", int),
+        np.uint8: DtypeInfo(np.uint8, "Int64", int),
+        np.uint16: DtypeInfo(np.uint16, "Int64", int),
+        np.uint32: DtypeInfo(np.uint32, "Int64", int),
     }
 
     @classmethod
@@ -135,20 +135,6 @@ def get_by_dtype(cls, dtype: Any) -> DtypeInfo | None:
             )
         return cls._mappings.get(dtype)
 
-    @staticmethod
-    def get_by_kind(kind: str) -> DtypeInfo | None:
-        """Get DtypeInfo by kind."""
-        return next(
-            (info for info in DtypeRegistry._mappings.values() if info.kind == kind),
-            None,
-        )
-
-    @staticmethod
-    def rust_compatible_kind(kind: str) -> str:
-        """Map to a Rust-compatible kind for schema encoding."""
-        # incompatible_integer_kinds = {"Int32", "UInt8", "UInt16", "UInt32", "UInt64"}
-        return "Int64" if "Int" in kind else kind
-
     @staticmethod
     def supported_dtypes() -> KeysView[type]:
         """Get a list of supported NumPy dtypes."""
@@ -167,6 +153,9 @@ class AnalyzedTypeInfo:
 
     key_type: type | None  # For element of KTable
     struct_type: type | None  # For Struct, a dataclass or namedtuple
+    np_number_type: (
+        type | None
+    )  # NumPy dtype for the element type, if represented by numpy.ndarray or a NumPy scalar
 
     attrs: dict[str, Any] | None
     nullable: bool = False
@@ -221,6 +210,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
     struct_type: type | None = None
     elem_type: ElementType | None = None
     key_type: type | None = None
+    np_number_type: type | None = None
     if _is_struct_type(t):
         struct_type = t
 
@@ -254,11 +244,11 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         if not dtype_args:
             raise ValueError("Invalid dtype specification for NDArray")
 
-        numpy_dtype = dtype_args[0]
-        dtype_info = DtypeRegistry.get_by_dtype(numpy_dtype)
+        np_number_type = dtype_args[0]
+        dtype_info = DtypeRegistry.get_by_dtype(np_number_type)
         if dtype_info is None:
             raise ValueError(
-                f"Unsupported numpy dtype for NDArray: {numpy_dtype}. "
+                f"Unsupported numpy dtype for NDArray: {np_number_type}. "
                 f"Supported dtypes: {DtypeRegistry.supported_dtypes()}"
             )
         elem_type = dtype_info.annotated_type
@@ -272,6 +262,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         dtype_info = DtypeRegistry.get_by_dtype(t)
         if dtype_info is not None:
             kind = dtype_info.kind
+            np_number_type = dtype_info.numpy_dtype
         elif t is bytes:
             kind = "Bytes"
         elif t is str:
@@ -301,6 +292,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         elem_type=elem_type,
         key_type=key_type,
         struct_type=struct_type,
+        np_number_type=np_number_type,
         attrs=attrs,
         nullable=nullable,
     )
@@ -355,9 +347,6 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
             raise ValueError("Vector type must have an element type")
         elem_type_info = analyze_type_info(type_info.elem_type)
         encoded_type["element_type"] = _encode_type(elem_type_info)
-        encoded_type["element_type"]["kind"] = DtypeRegistry.rust_compatible_kind(
-            elem_type_info.kind
-        )
         encoded_type["dimension"] = type_info.vector_info.dim
 
     elif type_info.kind in TABLE_TYPES: