cocoindex 0.1.52__cp312-cp312-macosx_11_0_arm64.whl → 0.1.53__cp312-cp312-macosx_11_0_arm64.whl

cocoindex/convert.py

@@ -6,17 +6,19 @@ import dataclasses
 import datetime
 import inspect
 import uuid
+from enum import Enum
+from typing import Any, Callable, Mapping, get_origin
+
 import numpy as np
 
-from enum import Enum
-from typing import Any, Callable, get_origin, Mapping
 from .typing import (
+    KEY_FIELD_NAME,
+    TABLE_TYPES,
+    DtypeRegistry,
     analyze_type_info,
     encode_enriched_type,
+    extract_ndarray_scalar_dtype,
     is_namedtuple_type,
-    TABLE_TYPES,
-    KEY_FIELD_NAME,
-    DtypeRegistry,
 )
 
 
@@ -29,6 +31,8 @@ def encode_engine_value(value: Any) -> Any:
         ]
     if is_namedtuple_type(type(value)):
         return [encode_engine_value(getattr(value, name)) for name in value._fields]
+    if isinstance(value, np.number):
+        return value.item()
     if isinstance(value, np.ndarray):
         return value
     if isinstance(value, (list, tuple)):
@@ -86,6 +90,20 @@ def make_engine_value_decoder(
             field_path, src_type["fields"], dst_type_info.struct_type
         )
 
+    if dst_type_info.np_number_type is not None and src_type_kind != "Vector":
+        numpy_type = dst_type_info.np_number_type
+
+        def decode_numpy_scalar(value: Any) -> Any | None:
+            if value is None:
+                if dst_type_info.nullable:
+                    return None
+                raise ValueError(
+                    f"Received null for non-nullable scalar `{''.join(field_path)}`"
+                )
+            return numpy_type(value)
+
+        return decode_numpy_scalar
+
     if src_type_kind in TABLE_TYPES:
         field_path.append("[*]")
         elem_type_info = analyze_type_info(dst_type_info.elem_type)
@@ -127,33 +145,42 @@ def make_engine_value_decoder(
         return lambda value: uuid.UUID(bytes=value)
 
     if src_type_kind == "Vector":
-        dtype_info = DtypeRegistry.get_by_dtype(dst_type_info.np_number_type)
 
         def decode_vector(value: Any) -> Any | None:
+            field_path_str = "".join(field_path)
+            expected_dim = (
+                dst_type_info.vector_info.dim if dst_type_info.vector_info else None
+            )
+
             if value is None:
                 if dst_type_info.nullable:
                     return None
                 raise ValueError(
-                    f"Received null for non-nullable vector `{''.join(field_path)}`"
+                    f"Received null for non-nullable vector `{field_path_str}`"
                 )
-
             if not isinstance(value, (np.ndarray, list)):
                 raise TypeError(
-                    f"Expected NDArray or list for vector `{''.join(field_path)}`, got {type(value)}"
+                    f"Expected NDArray or list for vector `{field_path_str}`, got {type(value)}"
                 )
-            expected_dim = (
-                dst_type_info.vector_info.dim if dst_type_info.vector_info else None
-            )
             if expected_dim is not None and len(value) != expected_dim:
                 raise ValueError(
-                    f"Vector dimension mismatch for `{''.join(field_path)}`: "
+                    f"Vector dimension mismatch for `{field_path_str}`: "
                     f"expected {expected_dim}, got {len(value)}"
                 )
 
-            # Use NDArray for supported numeric dtypes, else return list
-            if dtype_info is not None:
-                return np.array(value, dtype=dtype_info.numpy_dtype)
-            return value
+            if dst_type_info.np_number_type is None:  # for Non-NDArray vector
+                elem_decoder = make_engine_value_decoder(
+                    field_path + ["[*]"],
+                    src_type["element_type"],
+                    dst_type_info.elem_type,
+                )
+                return [elem_decoder(v) for v in value]
+            else:  # for NDArray vector
+                scalar_dtype = extract_ndarray_scalar_dtype(
+                    dst_type_info.np_number_type
+                )
+                _ = DtypeRegistry.validate_dtype_and_get_kind(scalar_dtype)
+                return np.array(value, dtype=scalar_dtype)
 
         return decode_vector
 

cocoindex/llm.py

@@ -9,6 +9,8 @@ class LlmApiType(Enum):
     OLLAMA = "Ollama"
     GEMINI = "Gemini"
     ANTHROPIC = "Anthropic"
+    LITE_LLM = "LiteLlm"
+    OPEN_ROUTER = "OpenRouter"
 
 
 @dataclass

cocoindex/tests/test_convert.py

@@ -1,22 +1,25 @@
-import uuid
 import datetime
+import uuid
 from dataclasses import dataclass, make_dataclass
-from typing import NamedTuple, Literal, Any, Callable, Union
+from typing import Annotated, Any, Callable, Literal, NamedTuple
+
+import numpy as np
 import pytest
+from numpy.typing import NDArray
+
 import cocoindex
-from cocoindex.typing import (
-    encode_enriched_type,
-    Vector,
-    Float32,
-    Float64,
-)
 from cocoindex.convert import (
+    dump_engine_object,
     encode_engine_value,
     make_engine_value_decoder,
-    dump_engine_object,
 )
-import numpy as np
-from numpy.typing import NDArray
+from cocoindex.typing import (
+    Float32,
+    Float64,
+    TypeKind,
+    Vector,
+    encode_enriched_type,
+)
 
 
 @dataclass
@@ -128,6 +131,19 @@ def test_encode_engine_value_date_time_types() -> None:
     assert encode_engine_value(dt) == dt
 
 
+def test_encode_scalar_numpy_values() -> None:
+    """Test encoding scalar NumPy values to engine-compatible values."""
+    test_cases = [
+        (np.int64(42), 42),
+        (np.float32(3.14), pytest.approx(3.14)),
+        (np.float64(2.718), pytest.approx(2.718)),
+    ]
+    for np_value, expected in test_cases:
+        encoded = encode_engine_value(np_value)
+        assert encoded == expected
+        assert isinstance(encoded, (int, float))
+
+
 def test_encode_engine_value_struct() -> None:
     order = Order(order_id="O123", name="mixed nuts", price=25.0)
     assert encode_engine_value(order) == ["O123", "mixed nuts", 25.0, "default_extra"]
@@ -213,6 +229,47 @@ def test_roundtrip_basic_types() -> None:
     )
 
 
+def test_decode_scalar_numpy_values() -> None:
+    test_cases = [
+        ({"kind": "Int64"}, np.int64, 42, np.int64(42)),
+        ({"kind": "Float32"}, np.float32, 3.14, np.float32(3.14)),
+        ({"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)
+        result = decoder(input_value)
+        assert isinstance(result, dst_type)
+        assert result == expected
+
+
+def test_non_ndarray_vector_decoding() -> None:
+    # Test list[np.float64]
+    src_type = {
+        "kind": "Vector",
+        "element_type": {"kind": "Float64"},
+        "dimension": None,
+    }
+    dst_type_float = list[np.float64]
+    decoder = make_engine_value_decoder(["field"], src_type, dst_type_float)
+    input_numbers = [1.0, 2.0, 3.0]
+    result = decoder(input_numbers)
+    assert isinstance(result, list)
+    assert all(isinstance(x, np.float64) for x in result)
+    assert result == [np.float64(1.0), np.float64(2.0), np.float64(3.0)]
+
+    # 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)
+    uuid1 = uuid.uuid4()
+    uuid2 = uuid.uuid4()
+    input_bytes = [uuid1.bytes, uuid2.bytes]
+    result = decoder(input_bytes)
+    assert isinstance(result, list)
+    assert all(isinstance(x, uuid.UUID) for x in result)
+    assert result == [uuid1, uuid2]
+
+
 @pytest.mark.parametrize(
     "data_type, engine_val, expected",
     [
@@ -565,12 +622,6 @@ Float64VectorTypeNoDim = Vector[np.float64]
 Float32VectorType = Vector[np.float32, Literal[3]]
 Float64VectorType = Vector[np.float64, Literal[3]]
 Int64VectorType = Vector[np.int64, Literal[3]]
-Int32VectorType = Vector[np.int32, Literal[3]]
-UInt8VectorType = Vector[np.uint8, Literal[3]]
-UInt16VectorType = Vector[np.uint16, Literal[3]]
-UInt32VectorType = Vector[np.uint32, Literal[3]]
-UInt64VectorType = Vector[np.uint64, Literal[3]]
-StrVectorType = Vector[str]
 NDArrayFloat32Type = NDArray[np.float32]
 NDArrayFloat64Type = NDArray[np.float64]
 NDArrayInt64Type = NDArray[np.int64]
@@ -767,19 +818,19 @@ def test_full_roundtrip_vector_numeric_types() -> None:
     value_i64: Vector[np.int64, Literal[3]] = np.array([1, 2, 3], dtype=np.int64)
     validate_full_roundtrip(value_i64, Vector[np.int64, Literal[3]])
     value_i32: Vector[np.int32, Literal[3]] = np.array([1, 2, 3], dtype=np.int32)
-    with pytest.raises(ValueError, match="type unsupported yet"):
+    with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_i32, Vector[np.int32, Literal[3]])
     value_u8: Vector[np.uint8, Literal[3]] = np.array([1, 2, 3], dtype=np.uint8)
-    with pytest.raises(ValueError, match="type unsupported yet"):
+    with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u8, Vector[np.uint8, Literal[3]])
     value_u16: Vector[np.uint16, Literal[3]] = np.array([1, 2, 3], dtype=np.uint16)
-    with pytest.raises(ValueError, match="type unsupported yet"):
+    with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u16, Vector[np.uint16, Literal[3]])
     value_u32: Vector[np.uint32, Literal[3]] = np.array([1, 2, 3], dtype=np.uint32)
-    with pytest.raises(ValueError, match="type unsupported yet"):
+    with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u32, Vector[np.uint32, Literal[3]])
     value_u64: Vector[np.uint64, Literal[3]] = np.array([1, 2, 3], dtype=np.uint64)
-    with pytest.raises(ValueError, match="type unsupported yet"):
+    with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
         validate_full_roundtrip(value_u64, Vector[np.uint64, Literal[3]])
 
 
@@ -808,7 +859,88 @@ def test_roundtrip_dimension_mismatch() -> None:
         validate_full_roundtrip(value_f32, Vector[np.float32, Literal[3]])
 
 
-def test_roundtrip_list_backward_compatibility() -> None:
-    """Test full roundtrip for list-based vectors for backward compatibility."""
-    value_list: list[int] = [1, 2, 3]
-    validate_full_roundtrip(value_list, list[int])
+def test_full_roundtrip_scalar_numeric_types() -> None:
+    """Test full roundtrip for scalar NumPy numeric types."""
+    # Test supported scalar types
+    validate_full_roundtrip(np.int64(42), np.int64)
+    validate_full_roundtrip(np.float32(3.14), np.float32)
+    validate_full_roundtrip(np.float64(2.718), np.float64)
+
+    # Test unsupported scalar types
+    for unsupported_type in [np.int32, np.uint8, np.uint16, np.uint32, np.uint64]:
+        with pytest.raises(ValueError, match="Unsupported NumPy dtype"):
+            validate_full_roundtrip(unsupported_type(1), unsupported_type)
+
+
+def test_full_roundtrip_nullable_scalar() -> None:
+    """Test full roundtrip for nullable scalar NumPy types."""
+    # Test with non-null values
+    validate_full_roundtrip(np.int64(42), np.int64 | None)
+    validate_full_roundtrip(np.float32(3.14), np.float32 | None)
+    validate_full_roundtrip(np.float64(2.718), np.float64 | None)
+
+    # Test with None
+    validate_full_roundtrip(None, np.int64 | None)
+    validate_full_roundtrip(None, np.float32 | None)
+    validate_full_roundtrip(None, np.float64 | None)
+
+
+def test_full_roundtrip_scalar_in_struct() -> None:
+    """Test full roundtrip for scalar NumPy types in a dataclass."""
+
+    @dataclass
+    class NumericStruct:
+        int_field: np.int64
+        float32_field: np.float32
+        float64_field: np.float64
+
+    instance = NumericStruct(
+        int_field=np.int64(42),
+        float32_field=np.float32(3.14),
+        float64_field=np.float64(2.718),
+    )
+    validate_full_roundtrip(instance, NumericStruct)
+
+
+def test_full_roundtrip_scalar_in_nested_struct() -> None:
+    """Test full roundtrip for scalar NumPy types in a nested struct."""
+
+    @dataclass
+    class InnerStruct:
+        value: np.float64
+
+    @dataclass
+    class OuterStruct:
+        inner: InnerStruct
+        count: np.int64
+
+    instance = OuterStruct(
+        inner=InnerStruct(value=np.float64(2.718)),
+        count=np.int64(1),
+    )
+    validate_full_roundtrip(instance, OuterStruct)
+
+
+def test_full_roundtrip_scalar_with_python_types() -> None:
+    """Test full roundtrip for structs mixing NumPy and Python scalar types."""
+
+    @dataclass
+    class MixedStruct:
+        numpy_int: np.int64
+        python_int: int
+        numpy_float: np.float64
+        python_float: float
+        string: str
+        annotated_int: Annotated[np.int64, TypeKind("int")]
+        annotated_float: Float32
+
+    instance = MixedStruct(
+        numpy_int=np.int64(42),
+        python_int=43,
+        numpy_float=np.float64(2.718),
+        python_float=3.14,
+        string="hello, world",
+        annotated_int=np.int64(42),
+        annotated_float=2.0,
+    )
+    validate_full_roundtrip(instance, MixedStruct)

cocoindex/tests/test_typing.py

@@ -1,30 +1,21 @@
 import dataclasses
 import datetime
 import uuid
-from typing import (
-    Annotated,
-    List,
-    Dict,
-    Literal,
-    Any,
-    get_args,
-    NamedTuple,
-)
-from collections.abc import Sequence, Mapping
-import pytest
+from collections.abc import Mapping, Sequence
+from typing import Annotated, Any, Dict, List, Literal, NamedTuple, get_args, get_origin
+
 import numpy as np
+import pytest
 from numpy.typing import NDArray
 
 from cocoindex.typing import (
-    analyze_type_info,
+    AnalyzedTypeInfo,
+    TypeAttr,
+    TypeKind,
     Vector,
     VectorInfo,
-    TypeKind,
-    TypeAttr,
-    Float32,
-    Float64,
+    analyze_type_info,
     encode_enriched_type,
-    AnalyzedTypeInfo,
 )
 
 
@@ -42,61 +33,57 @@ class SimpleNamedTuple(NamedTuple):
 def test_ndarray_float32_no_dim() -> None:
     typ = NDArray[np.float32]
     result = analyze_type_info(typ)
-    assert result == AnalyzedTypeInfo(
-        kind="Vector",
-        vector_info=VectorInfo(dim=None),
-        elem_type=Float32,
-        key_type=None,
-        struct_type=None,
-        np_number_type=np.float32,
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
 
 
 def test_vector_float32_no_dim() -> None:
     typ = Vector[np.float32]
     result = analyze_type_info(typ)
-    assert result == AnalyzedTypeInfo(
-        kind="Vector",
-        vector_info=VectorInfo(dim=None),
-        elem_type=Float32,
-        key_type=None,
-        struct_type=None,
-        np_number_type=np.float32,
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
 
 
 def test_ndarray_float64_with_dim() -> None:
     typ = Annotated[NDArray[np.float64], VectorInfo(dim=128)]
     result = analyze_type_info(typ)
-    assert result == AnalyzedTypeInfo(
-        kind="Vector",
-        vector_info=VectorInfo(dim=128),
-        elem_type=Float64,
-        key_type=None,
-        struct_type=None,
-        np_number_type=np.float64,
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=128)
+    assert result.elem_type == np.float64
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float64]
 
 
 def test_vector_float32_with_dim() -> None:
     typ = Vector[np.float32, Literal[384]]
     result = analyze_type_info(typ)
-    assert result == AnalyzedTypeInfo(
-        kind="Vector",
-        vector_info=VectorInfo(dim=384),
-        elem_type=Float32,
-        key_type=None,
-        struct_type=None,
-        np_number_type=np.float32,
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=384)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
 
 
 def test_ndarray_int64_no_dim() -> None:
@@ -104,30 +91,49 @@ def test_ndarray_int64_no_dim() -> None:
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=None)
-    assert get_args(result.elem_type) == (int, TypeKind("Int64"))
-    assert not result.nullable
+    assert result.elem_type == np.int64
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.int64]
 
 
 def test_nullable_ndarray() -> None:
     typ = NDArray[np.float32] | None
     result = analyze_type_info(typ)
-    assert result == AnalyzedTypeInfo(
-        kind="Vector",
-        vector_info=VectorInfo(dim=None),
-        elem_type=Float32,
-        key_type=None,
-        struct_type=None,
-        np_number_type=np.float32,
-        attrs=None,
-        nullable=True,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is True
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
+
+
+def test_scalar_numpy_types() -> None:
+    for np_type, expected_kind in [
+        (np.int64, "Int64"),
+        (np.float32, "Float32"),
+        (np.float64, "Float64"),
+    ]:
+        type_info = analyze_type_info(np_type)
+        assert type_info.kind == expected_kind, (
+            f"Expected {expected_kind} for {np_type}, got {type_info.kind}"
+        )
+        assert type_info.np_number_type == np_type, (
+            f"Expected {np_type}, got {type_info.np_number_type}"
+        )
+        assert type_info.elem_type is None
+        assert type_info.vector_info is None
 
 
 def test_vector_str() -> None:
     typ = Vector[str]
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
-    assert result.elem_type == str
+    assert result.elem_type is str
     assert result.vector_info == VectorInfo(dim=None)
 
 
@@ -143,7 +149,7 @@ def test_non_numpy_vector() -> None:
     typ = Vector[float, Literal[3]]
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
-    assert result.elem_type == float
+    assert result.elem_type is float
     assert result.vector_info == VectorInfo(dim=3)
 
 
@@ -487,6 +493,19 @@ def test_encode_enriched_type_nullable() -> None:
     assert result["nullable"] is True
 
 
+def test_encode_scalar_numpy_types_schema() -> None:
+    for np_type, expected_kind in [
+        (np.int64, "Int64"),
+        (np.float32, "Float32"),
+        (np.float64, "Float64"),
+    ]:
+        schema = encode_enriched_type(np_type)
+        assert schema["type"]["kind"] == expected_kind, (
+            f"Expected {expected_kind} for {np_type}, got {schema['type']['kind']}"
+        )
+        assert not schema.get("nullable", False)
+
+
 def test_invalid_struct_kind() -> None:
     typ = Annotated[SimpleDataclass, TypeKind("Vector")]
     with pytest.raises(ValueError, match="Unexpected type kind for struct: Vector"):

cocoindex/typing.py

@@ -1,22 +1,22 @@
-import typing
 import collections
 import dataclasses
 import datetime
-import types
 import inspect
+import types
+import typing
 import uuid
 from typing import (
+    TYPE_CHECKING,
     Annotated,
-    NamedTuple,
     Any,
-    KeysView,
-    TypeVar,
-    TYPE_CHECKING,
-    overload,
     Generic,
     Literal,
+    NamedTuple,
     Protocol,
+    TypeVar,
+    overload,
 )
+
 import numpy as np
 from numpy.typing import NDArray
 
@@ -67,7 +67,7 @@ else:
                 # No dimension provided, e.g., Vector[np.float32]
                 dtype = params
                 # Use NDArray for supported numeric dtypes, else list
-                if DtypeRegistry.get_by_dtype(dtype) is not None:
+                if dtype in DtypeRegistry._DTYPE_TO_KIND:
                     return Annotated[NDArray[dtype], VectorInfo(dim=None)]
                 return Annotated[list[dtype], VectorInfo(dim=None)]
             else:
@@ -79,7 +79,7 @@ else:
                     if typing.get_origin(dim_literal) is Literal
                     else None
                 )
-                if DtypeRegistry.get_by_dtype(dtype) is not None:
+                if dtype in DtypeRegistry._DTYPE_TO_KIND:
                     return Annotated[NDArray[dtype], VectorInfo(dim=dim_val)]
                 return Annotated[list[dtype], VectorInfo(dim=dim_val)]
 
@@ -90,6 +90,19 @@ KEY_FIELD_NAME: str = "_key"
 ElementType = type | tuple[type, type] | Annotated[Any, TypeKind]
 
 
+def extract_ndarray_scalar_dtype(ndarray_type: Any) -> Any:
+    args = typing.get_args(ndarray_type)
+    _, dtype_spec = args
+    dtype_args = typing.get_args(dtype_spec)
+    if not dtype_args:
+        raise ValueError(f"Invalid dtype specification: {dtype_spec}")
+    return dtype_args[0]
+
+
+def is_numpy_number_type(t: type) -> bool:
+    return isinstance(t, type) and issubclass(t, np.number)
+
+
 def is_namedtuple_type(t: type) -> bool:
     return isinstance(t, type) and issubclass(t, tuple) and hasattr(t, "_fields")
 
@@ -100,41 +113,34 @@ def _is_struct_type(t: ElementType | None) -> bool:
     )
 
 
-class DtypeInfo:
-    """Metadata for a NumPy dtype."""
-
-    def __init__(self, numpy_dtype: type, kind: str, python_type: type) -> None:
-        self.numpy_dtype = numpy_dtype
-        self.kind = kind
-        self.python_type = python_type
-        self.annotated_type = Annotated[python_type, TypeKind(kind)]
-
-
 class DtypeRegistry:
     """
     Registry for NumPy dtypes used in CocoIndex.
-    Provides mappings from NumPy dtypes to CocoIndex's type representation.
+    Maps NumPy dtypes to their CocoIndex type kind.
     """
 
-    _mappings: dict[type, DtypeInfo] = {
-        np.float32: DtypeInfo(np.float32, "Float32", float),
-        np.float64: DtypeInfo(np.float64, "Float64", float),
-        np.int64: DtypeInfo(np.int64, "Int64", int),
+    _DTYPE_TO_KIND: dict[ElementType, str] = {
+        np.float32: "Float32",
+        np.float64: "Float64",
+        np.int64: "Int64",
     }
 
     @classmethod
-    def get_by_dtype(cls, dtype: Any) -> DtypeInfo | None:
-        """Get DtypeInfo by NumPy dtype."""
+    def validate_dtype_and_get_kind(cls, dtype: ElementType) -> str:
+        """
+        Validate that the given dtype is supported, and get its CocoIndex kind by dtype.
+        """
         if dtype is Any:
             raise TypeError(
                 "NDArray for Vector must use a concrete numpy dtype, got `Any`."
             )
-        return cls._mappings.get(dtype)
-
-    @staticmethod
-    def supported_dtypes() -> KeysView[type]:
-        """Get a list of supported NumPy dtypes."""
-        return DtypeRegistry._mappings.keys()
+        kind = cls._DTYPE_TO_KIND.get(dtype)
+        if kind is None:
+            raise ValueError(
+                f"Unsupported NumPy dtype in NDArray: {dtype}. "
+                f"Supported dtypes: {cls._DTYPE_TO_KIND.keys()}"
+            )
+        return kind
 
 
 @dataclasses.dataclass
@@ -214,6 +220,9 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
             kind = "Struct"
         elif kind != "Struct":
             raise ValueError(f"Unexpected type kind for struct: {kind}")
+    elif is_numpy_number_type(t):
+        np_number_type = t
+        kind = DtypeRegistry.validate_dtype_and_get_kind(t)
     elif base_type is collections.abc.Sequence or base_type is list:
         args = typing.get_args(t)
         elem_type = args[0]
@@ -233,21 +242,9 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
             raise ValueError(f"Unexpected type kind for list: {kind}")
     elif base_type is np.ndarray:
         kind = "Vector"
-        args = typing.get_args(t)
-        _, dtype_spec = args
-
-        dtype_args = typing.get_args(dtype_spec)
-        if not dtype_args:
-            raise ValueError("Invalid dtype specification for NDArray")
-
-        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: {np_number_type}. "
-                f"Supported dtypes: {DtypeRegistry.supported_dtypes()}"
-            )
-        elem_type = dtype_info.annotated_type
+        np_number_type = t
+        elem_type = extract_ndarray_scalar_dtype(np_number_type)
+        _ = DtypeRegistry.validate_dtype_and_get_kind(elem_type)
         vector_info = VectorInfo(dim=None) if vector_info is None else vector_info
 
     elif base_type is collections.abc.Mapping or base_type is dict:
@@ -255,11 +252,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         elem_type = (args[0], args[1])
         kind = "KTable"
     elif kind is None:
-        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:
+        if t is bytes:
             kind = "Bytes"
         elif t is str:
             kind = "Str"

{cocoindex-0.1.52.dist-info → cocoindex-0.1.53.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.52
+Version: 0.1.53
 Requires-Dist: sentence-transformers>=3.3.1
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0

{cocoindex-0.1.52.dist-info → cocoindex-0.1.53.dist-info}/RECORD

@@ -1,17 +1,17 @@
-cocoindex-0.1.52.dist-info/METADATA,sha256=8psE6rLZFYrw7cBEU-rVD1JGf4ezIIHMn-Yk2ZVNEKY,9874
-cocoindex-0.1.52.dist-info/WHEEL,sha256=vahFoO0M6DZmYz1cr_3GC_BJXbjcSJJ4jR_JC04LGJc,104
-cocoindex-0.1.52.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.1.52.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+cocoindex-0.1.53.dist-info/METADATA,sha256=HQKpE-cc_Dgi2zpIZOAX8aVCl5Zt8hd__egLbAdlOSU,9874
+cocoindex-0.1.53.dist-info/WHEEL,sha256=vahFoO0M6DZmYz1cr_3GC_BJXbjcSJJ4jR_JC04LGJc,104
+cocoindex-0.1.53.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.1.53.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 cocoindex/__init__.py,sha256=3potlGouaGAznMTEqU5QKuB934qpD1xR6ZH6EL3YMWw,1782
-cocoindex/_engine.cpython-312-darwin.so,sha256=MuIjmOSOeRSPbnynxicEE3YS9XNwJpd_noB_Yc4L4Ws,57344192
+cocoindex/_engine.cpython-312-darwin.so,sha256=7BcD6RlbQ1I1gn8ponJxeRsmwQP4W2t8160TurmWR3k,57344864
 cocoindex/auth_registry.py,sha256=1XqO7ibjmBBd8i11XSJTvTgdz8p1ptW-ZpuSgo_5zzk,716
 cocoindex/cli.py,sha256=joSucqLibjZHs27fwNPzuMPGLeVJJZqaUbeSM-OnEyA,18223
-cocoindex/convert.py,sha256=s5T2IQ1tMBT9JKTzqQtQ-I0sUrfjBGKAnMjDjY6LdrI,8785
+cocoindex/convert.py,sha256=qP_LRTg_B8YPEeHprnX4kl-QV4ks0oqQzH7A1totZRA,9787
 cocoindex/flow.py,sha256=uQKvIyWSysSdhNFI87QD6f2kbF-2R6zFMLZWBEk4rZU,30015
 cocoindex/functions.py,sha256=9A61Jj5a3vQoI2MIAhjXvJrDxSzDhe6VncQWbiVtwcg,2393
 cocoindex/index.py,sha256=j93B9jEvvLXHtpzKWL88SY6wCGEoPgpsQhEGHlyYGFg,540
 cocoindex/lib.py,sha256=BeRUn3RqE_wSsVtsgCzbFFKe1LXgRyRmMOcmwWBuEXo,2940
-cocoindex/llm.py,sha256=KO-R4mrAWtxXD82-Yv5ixpkKMVfkwpbdWwqPVZygLu4,352
+cocoindex/llm.py,sha256=gbjAsD8Y9QZyICCQfSDZYQU2pIA6GqG4jKbRpbnQniA,408
 cocoindex/op.py,sha256=Z7V9Fdz4qeTeozzKmp1Dk1lPUP0GBgVgD_vBEhTJS5Y,11810
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cocoindex/runtime.py,sha256=bAdHYaXFWiiUWyAgzmKTeaAaRR0D_AmaqVCIdPO-v00,1056
@@ -20,9 +20,9 @@ cocoindex/setup.py,sha256=u5dYZFKfz4yZLiGHD0guNaR0s4zY9JAoZWrWHpAHw_0,773
 cocoindex/sources.py,sha256=JCnOhv1w4o28e03i7yvo4ESicWYAhckkBg5bQlxNH4U,1330
 cocoindex/targets.py,sha256=Nfh_tpFd1goTnS_cxBjIs4j9zl3Z4Z1JomAQ1dl3Sic,2796
 cocoindex/tests/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-cocoindex/tests/test_convert.py,sha256=mlCJ1U_T7sPDk1Dw78V1BKaYbfqjDjGAaOZT0zMYOwg,28960
+cocoindex/tests/test_convert.py,sha256=iDA4j-X_r1KeJgaa3osU7SyY8rBc57R0YVW8G4Z83YI,33235
 cocoindex/tests/test_optional_database.py,sha256=snAmkNa6wtOSaxoZE1HgjvL5v_ylitt3Jt_9df4Cgdc,8506
-cocoindex/tests/test_typing.py,sha256=4cy_6kPyGxsM6qX-O8K-jXhsgDr_FXePplERgky22qI,12313
-cocoindex/typing.py,sha256=PEV_ds7AGHAHkZmX_fWg-2mFiEXDx-V-Rlhv11_8rtk,12387
+cocoindex/tests/test_typing.py,sha256=jeZFUyzUTC6hdyo1uERhlKaKAiBKeKCrZ20OLGBkluo,14184
+cocoindex/typing.py,sha256=yLtXdXEF0sYe3cMxTR__g39Gi286CnOKGI0T-hjaCrw,11999
 cocoindex/utils.py,sha256=hUhX-XV6XGCtJSEIpBOuDv6VvqImwPlgBxztBTw7u0U,598
-cocoindex-0.1.52.dist-info/RECORD,,
+cocoindex-0.1.53.dist-info/RECORD,,