cocoindex 0.1.49__cp311-cp311-macosx_11_0_arm64.whl → 0.1.50__cp311-cp311-macosx_11_0_arm64.whl

cocoindex/__init__.py

@@ -2,7 +2,7 @@
 Cocoindex is a framework for building and running indexing pipelines.
 """
 
-from . import functions, query, sources, storages, cli, utils
+from . import functions, sources, storages, cli, utils
 
 from .auth_registry import AuthEntryReference, add_auth_entry, ref_auth_entry
 from .flow import FlowBuilder, DataScope, DataSlice, Flow, transform_flow

cocoindex/convert.py

@@ -6,6 +6,7 @@ import dataclasses
 import datetime
 import inspect
 import uuid
+import numpy as np
 
 from enum import Enum
 from typing import Any, Callable, get_origin, Mapping
@@ -15,6 +16,7 @@ from .typing import (
     is_namedtuple_type,
     TABLE_TYPES,
     KEY_FIELD_NAME,
+    DtypeRegistry,
 )
 
 
@@ -27,6 +29,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.ndarray):
+        return value
     if isinstance(value, (list, tuple)):
         return [encode_engine_value(v) for v in value]
     if isinstance(value, dict):
@@ -122,6 +126,38 @@ def make_engine_value_decoder(
     if src_type_kind == "Uuid":
         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)
+
+        def decode_vector(value: Any) -> Any | None:
+            if value is None:
+                if dst_type_info.nullable:
+                    return None
+                raise ValueError(
+                    f"Received null for non-nullable vector `{''.join(field_path)}`"
+                )
+
+            if not isinstance(value, (np.ndarray, list)):
+                raise TypeError(
+                    f"Expected NDArray or list for vector `{''.join(field_path)}`, 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"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
+
+        return decode_vector
+
     return lambda value: value
 
 

cocoindex/functions.py

@@ -1,6 +1,8 @@
 """All builtin functions."""
 
-from typing import Annotated, Any, TYPE_CHECKING
+from typing import Annotated, Any, TYPE_CHECKING, Literal
+import numpy as np
+from numpy.typing import NDArray
 import dataclasses
 
 from .typing import Float32, Vector, TypeAttr
@@ -66,11 +68,11 @@ class SentenceTransformerEmbedExecutor:
         self._model = sentence_transformers.SentenceTransformer(self.spec.model, **args)
         dim = self._model.get_sentence_embedding_dimension()
         result: type = Annotated[
-            Vector[Float32, dim],  # type: ignore
+            Vector[np.float32, Literal[dim]],  # type: ignore
             TypeAttr("cocoindex.io/vector_origin_text", text.analyzed_value),
         ]
         return result
 
-    def __call__(self, text: str) -> list[Float32]:
-        result: list[Float32] = self._model.encode(text).tolist()
+    def __call__(self, text: str) -> NDArray[np.float32]:
+        result: NDArray[np.float32] = self._model.encode(text, convert_to_numpy=True)
         return result

cocoindex/lib.py

@@ -6,7 +6,7 @@ import warnings
 from typing import Callable, Any
 
 from . import _engine  # type: ignore
-from . import flow, query, setting
+from . import flow, setting
 from .convert import dump_engine_object
 
 
@@ -24,7 +24,6 @@ def init(settings: setting.Settings | None = None) -> None:
 def start_server(settings: setting.ServerSettings) -> None:
     """Start the cocoindex server."""
     flow.ensure_all_flows_built()
-    query.ensure_all_handlers_built()
     _engine.start_server(settings.__dict__)
 
 

cocoindex/tests/test_convert.py

@@ -1,11 +1,20 @@
 import uuid
 import datetime
 from dataclasses import dataclass, make_dataclass
-from typing import NamedTuple, Literal
+from typing import NamedTuple, Literal, Any, Callable
 import pytest
 import cocoindex
-from cocoindex.typing import encode_enriched_type
-from cocoindex.convert import encode_engine_value, make_engine_value_decoder
+from cocoindex.typing import (
+    encode_enriched_type,
+    Vector,
+)
+from cocoindex.convert import (
+    encode_engine_value,
+    make_engine_value_decoder,
+    dump_engine_object,
+)
+import numpy as np
+from numpy.typing import NDArray
 
 
 @dataclass
@@ -53,7 +62,9 @@ class CustomerNamedTuple(NamedTuple):
     tags: list[Tag] | None = None
 
 
-def build_engine_value_decoder(engine_type_in_py, python_type=None):
+def build_engine_value_decoder(
+    engine_type_in_py: Any, python_type: Any | None = None
+) -> Callable[[Any], Any]:
     """
     Helper to build a converter for the given engine-side type (as represented in Python).
     If python_type is not specified, uses engine_type_in_py as the target.
@@ -62,6 +73,27 @@ def build_engine_value_decoder(engine_type_in_py, python_type=None):
     return make_engine_value_decoder([], engine_type, python_type or engine_type_in_py)
 
 
+def validate_full_roundtrip(
+    value: Any, output_type: Any, input_type: Any | None = None
+) -> None:
+    """
+    Validate the given value doesn't change after encoding, sending to engine (using output_type), receiving back and decoding (using input_type).
+
+    If `input_type` is not specified, uses `output_type` as the target.
+    """
+    from cocoindex import _engine
+
+    encoded_value = encode_engine_value(value)
+    encoded_output_type = encode_enriched_type(output_type)["type"]
+    value_from_engine = _engine.testutil.seder_roundtrip(
+        encoded_value, encoded_output_type
+    )
+    decoded_value = build_engine_value_decoder(input_type or output_type, output_type)(
+        value_from_engine
+    )
+    assert decoded_value == value
+
+
 def test_encode_engine_value_basic_types():
     assert encode_engine_value(123) == 123
     assert encode_engine_value(3.14) == 3.14
@@ -434,57 +466,33 @@ def test_field_position_cases(
     assert decoder(engine_val) == PythonOrder(**expected_dict)
 
 
-def test_roundtrip_ltable():
+def test_roundtrip_ltable() -> None:
     t = list[Order]
     value = [Order("O1", "item1", 10.0), Order("O2", "item2", 20.0)]
-    encoded = encode_engine_value(value)
-    assert encoded == [
-        ["O1", "item1", 10.0, "default_extra"],
-        ["O2", "item2", 20.0, "default_extra"],
-    ]
-    decoded = build_engine_value_decoder(t)(encoded)
-    assert decoded == value
+    validate_full_roundtrip(value, t)
 
     t_nt = list[OrderNamedTuple]
     value_nt = [
         OrderNamedTuple("O1", "item1", 10.0),
         OrderNamedTuple("O2", "item2", 20.0),
     ]
-    encoded = encode_engine_value(value_nt)
-    assert encoded == [
-        ["O1", "item1", 10.0, "default_extra"],
-        ["O2", "item2", 20.0, "default_extra"],
-    ]
-    decoded = build_engine_value_decoder(t_nt)(encoded)
-    assert decoded == value_nt
+    validate_full_roundtrip(value_nt, t_nt)
 
 
-def test_roundtrip_ktable_str_key():
+def test_roundtrip_ktable_str_key() -> None:
     t = dict[str, Order]
     value = {"K1": Order("O1", "item1", 10.0), "K2": Order("O2", "item2", 20.0)}
-    encoded = encode_engine_value(value)
-    assert encoded == [
-        ["K1", "O1", "item1", 10.0, "default_extra"],
-        ["K2", "O2", "item2", 20.0, "default_extra"],
-    ]
-    decoded = build_engine_value_decoder(t)(encoded)
-    assert decoded == value
+    validate_full_roundtrip(value, t)
 
     t_nt = dict[str, OrderNamedTuple]
     value_nt = {
         "K1": OrderNamedTuple("O1", "item1", 10.0),
         "K2": OrderNamedTuple("O2", "item2", 20.0),
     }
-    encoded = encode_engine_value(value_nt)
-    assert encoded == [
-        ["K1", "O1", "item1", 10.0, "default_extra"],
-        ["K2", "O2", "item2", 20.0, "default_extra"],
-    ]
-    decoded = build_engine_value_decoder(t_nt)(encoded)
-    assert decoded == value_nt
+    validate_full_roundtrip(value_nt, t_nt)
 
 
-def test_roundtrip_ktable_struct_key():
+def test_roundtrip_ktable_struct_key() -> None:
     @dataclass(frozen=True)
     class OrderKey:
         shop_id: str
@@ -495,29 +503,17 @@ def test_roundtrip_ktable_struct_key():
         OrderKey("A", 3): Order("O1", "item1", 10.0),
         OrderKey("B", 4): Order("O2", "item2", 20.0),
     }
-    encoded = encode_engine_value(value)
-    assert encoded == [
-        [["A", 3], "O1", "item1", 10.0, "default_extra"],
-        [["B", 4], "O2", "item2", 20.0, "default_extra"],
-    ]
-    decoded = build_engine_value_decoder(t)(encoded)
-    assert decoded == value
+    validate_full_roundtrip(value, t)
 
     t_nt = dict[OrderKey, OrderNamedTuple]
     value_nt = {
         OrderKey("A", 3): OrderNamedTuple("O1", "item1", 10.0),
         OrderKey("B", 4): OrderNamedTuple("O2", "item2", 20.0),
     }
-    encoded = encode_engine_value(value_nt)
-    assert encoded == [
-        [["A", 3], "O1", "item1", 10.0, "default_extra"],
-        [["B", 4], "O2", "item2", 20.0, "default_extra"],
-    ]
-    decoded = build_engine_value_decoder(t_nt)(encoded)
-    assert decoded == value_nt
+    validate_full_roundtrip(value_nt, t_nt)
 
 
-IntVectorType = cocoindex.Vector[int, Literal[5]]
+IntVectorType = cocoindex.Vector[np.int32, Literal[5]]
 
 
 def test_vector_as_vector() -> None:
@@ -525,7 +521,7 @@ def test_vector_as_vector() -> None:
     encoded = encode_engine_value(value)
     assert encoded == [1, 2, 3, 4, 5]
     decoded = build_engine_value_decoder(IntVectorType)(encoded)
-    assert decoded == value
+    assert np.array_equal(decoded, value)
 
 
 ListIntType = list[int]
@@ -536,4 +532,236 @@ def test_vector_as_list() -> None:
     encoded = encode_engine_value(value)
     assert encoded == [1, 2, 3, 4, 5]
     decoded = build_engine_value_decoder(ListIntType)(encoded)
-    assert decoded == value
+    assert np.array_equal(decoded, value)
+
+
+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]]
+NDArrayFloat32Type = NDArray[np.float32]
+NDArrayFloat64Type = NDArray[np.float64]
+NDArrayInt64Type = NDArray[np.int64]
+
+
+def test_encode_engine_value_ndarray():
+    """Test encoding NDArray vectors to lists for the Rust engine."""
+    vec_f32: Float32VectorType = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    assert np.array_equal(encode_engine_value(vec_f32), [1.0, 2.0, 3.0])
+    vec_f64: Float64VectorType = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+    assert np.array_equal(encode_engine_value(vec_f64), [1.0, 2.0, 3.0])
+    vec_i64: Int64VectorType = np.array([1, 2, 3], dtype=np.int64)
+    assert np.array_equal(encode_engine_value(vec_i64), [1, 2, 3])
+    vec_nd_f32: NDArrayFloat32Type = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    assert np.array_equal(encode_engine_value(vec_nd_f32), [1.0, 2.0, 3.0])
+
+
+def test_make_engine_value_decoder_ndarray():
+    """Test decoding engine lists to NDArray vectors."""
+    decoder_f32 = build_engine_value_decoder(Float32VectorType)
+    result_f32 = decoder_f32([1.0, 2.0, 3.0])
+    assert isinstance(result_f32, np.ndarray)
+    assert result_f32.dtype == np.float32
+    assert np.array_equal(result_f32, np.array([1.0, 2.0, 3.0], dtype=np.float32))
+    decoder_f64 = build_engine_value_decoder(Float64VectorType)
+    result_f64 = decoder_f64([1.0, 2.0, 3.0])
+    assert isinstance(result_f64, np.ndarray)
+    assert result_f64.dtype == np.float64
+    assert np.array_equal(result_f64, np.array([1.0, 2.0, 3.0], dtype=np.float64))
+    decoder_i64 = build_engine_value_decoder(Int64VectorType)
+    result_i64 = decoder_i64([1, 2, 3])
+    assert isinstance(result_i64, np.ndarray)
+    assert result_i64.dtype == np.int64
+    assert np.array_equal(result_i64, np.array([1, 2, 3], dtype=np.int64))
+    decoder_nd_f32 = build_engine_value_decoder(NDArrayFloat32Type)
+    result_nd_f32 = decoder_nd_f32([1.0, 2.0, 3.0])
+    assert isinstance(result_nd_f32, np.ndarray)
+    assert result_nd_f32.dtype == np.float32
+    assert np.array_equal(result_nd_f32, np.array([1.0, 2.0, 3.0], dtype=np.float32))
+
+
+def test_roundtrip_ndarray_vector():
+    """Test roundtrip encoding and decoding of NDArray vectors."""
+    value_f32: Float32VectorType = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    encoded_f32 = encode_engine_value(value_f32)
+    np.array_equal(encoded_f32, [1.0, 2.0, 3.0])
+    decoded_f32 = build_engine_value_decoder(Float32VectorType)(encoded_f32)
+    assert isinstance(decoded_f32, np.ndarray)
+    assert decoded_f32.dtype == np.float32
+    assert np.array_equal(decoded_f32, value_f32)
+    value_i64: Int64VectorType = np.array([1, 2, 3], dtype=np.int64)
+    encoded_i64 = encode_engine_value(value_i64)
+    assert np.array_equal(encoded_i64, [1, 2, 3])
+    decoded_i64 = build_engine_value_decoder(Int64VectorType)(encoded_i64)
+    assert isinstance(decoded_i64, np.ndarray)
+    assert decoded_i64.dtype == np.int64
+    assert np.array_equal(decoded_i64, value_i64)
+    value_nd_f64: NDArrayFloat64Type = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+    encoded_nd_f64 = encode_engine_value(value_nd_f64)
+    assert np.array_equal(encoded_nd_f64, [1.0, 2.0, 3.0])
+    decoded_nd_f64 = build_engine_value_decoder(NDArrayFloat64Type)(encoded_nd_f64)
+    assert isinstance(decoded_nd_f64, np.ndarray)
+    assert decoded_nd_f64.dtype == np.float64
+    assert np.array_equal(decoded_nd_f64, value_nd_f64)
+
+
+def test_uint_support():
+    """Test encoding and decoding of unsigned integer vectors."""
+    value_uint8 = np.array([1, 2, 3, 4], dtype=np.uint8)
+    encoded = encode_engine_value(value_uint8)
+    assert np.array_equal(encoded, [1, 2, 3, 4])
+    decoder = make_engine_value_decoder(
+        [], {"kind": "Vector", "element_type": {"kind": "UInt8"}}, NDArray[np.uint8]
+    )
+    decoded = decoder(encoded)
+    assert np.array_equal(decoded, value_uint8)
+    assert decoded.dtype == np.uint8
+    value_uint16 = np.array([1, 2, 3, 4], dtype=np.uint16)
+    encoded = encode_engine_value(value_uint16)
+    assert np.array_equal(encoded, [1, 2, 3, 4])
+    decoder = make_engine_value_decoder(
+        [], {"kind": "Vector", "element_type": {"kind": "UInt16"}}, NDArray[np.uint16]
+    )
+    decoded = decoder(encoded)
+    assert np.array_equal(decoded, value_uint16)
+    assert decoded.dtype == np.uint16
+    value_uint32 = np.array([1, 2, 3], dtype=np.uint32)
+    encoded = encode_engine_value(value_uint32)
+    assert np.array_equal(encoded, [1, 2, 3])
+    decoder = make_engine_value_decoder(
+        [], {"kind": "Vector", "element_type": {"kind": "UInt32"}}, NDArray[np.uint32]
+    )
+    decoded = decoder(encoded)
+    assert np.array_equal(decoded, value_uint32)
+    assert decoded.dtype == np.uint32
+    value_uint64 = np.array([1, 2, 3], dtype=np.uint64)
+    encoded = encode_engine_value(value_uint64)
+    assert np.array_equal(encoded, [1, 2, 3])
+    decoder = make_engine_value_decoder(
+        [], {"kind": "Vector", "element_type": {"kind": "UInt8"}}, NDArray[np.uint64]
+    )
+    decoded = decoder(encoded)
+    assert np.array_equal(decoded, value_uint64)
+    assert decoded.dtype == np.uint64
+
+
+def test_ndarray_dimension_mismatch():
+    """Test dimension enforcement for Vector with specified dimension."""
+    value: Float32VectorType = np.array([1.0, 2.0], dtype=np.float32)
+    encoded = encode_engine_value(value)
+    assert np.array_equal(encoded, [1.0, 2.0])
+    with pytest.raises(ValueError, match="Vector dimension mismatch"):
+        build_engine_value_decoder(Float32VectorType)(encoded)
+
+
+def test_list_vector_backward_compatibility():
+    """Test that list-based vectors still work for backward compatibility."""
+    value: IntVectorType = [1, 2, 3, 4, 5]
+    encoded = encode_engine_value(value)
+    assert encoded == [1, 2, 3, 4, 5]
+    decoded = build_engine_value_decoder(IntVectorType)(encoded)
+    assert isinstance(decoded, np.ndarray)
+    assert decoded.dtype == np.int32
+    assert np.array_equal(decoded, np.array([1, 2, 3, 4, 5], dtype=np.int64))
+    value_list: ListIntType = [1, 2, 3, 4, 5]
+    encoded = encode_engine_value(value_list)
+    assert np.array_equal(encoded, [1, 2, 3, 4, 5])
+    decoded = build_engine_value_decoder(ListIntType)(encoded)
+    assert np.array_equal(decoded, [1, 2, 3, 4, 5])
+
+
+def test_encode_complex_structure_with_ndarray():
+    """Test encoding a complex structure that includes an NDArray."""
+
+    @dataclass
+    class MyStructWithNDArray:
+        name: str
+        data: NDArray[np.float32]
+        value: int
+
+    original = MyStructWithNDArray(
+        name="test_np", data=np.array([1.0, 0.5], dtype=np.float32), value=100
+    )
+    encoded = encode_engine_value(original)
+    expected = [
+        "test_np",
+        [1.0, 0.5],
+        100,
+    ]
+    assert encoded[0] == expected[0]
+    assert np.array_equal(encoded[1], expected[1])
+    assert encoded[2] == expected[2]
+
+
+def test_decode_nullable_ndarray_none_or_value_input():
+    """Test decoding a nullable NDArray with None or value inputs."""
+    src_type_dict = {
+        "kind": "Vector",
+        "element_type": {"kind": "Float32"},
+        "dimension": None,
+    }
+    dst_annotation = NDArrayFloat32Type | None
+    decoder = make_engine_value_decoder([], src_type_dict, dst_annotation)
+
+    none_engine_value = None
+    decoded_array = decoder(none_engine_value)
+    assert decoded_array is None
+
+    engine_value = [1.0, 2.0, 3.0]
+    decoded_array = decoder(engine_value)
+
+    assert isinstance(decoded_array, np.ndarray)
+    assert decoded_array.dtype == np.float32
+    np.testing.assert_array_equal(
+        decoded_array, np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    )
+
+
+def test_decode_vector_string():
+    """Test decoding a vector of strings works for Python native list type."""
+    src_type_dict = {
+        "kind": "Vector",
+        "element_type": {"kind": "Str"},
+        "dimension": None,
+    }
+    decoder = make_engine_value_decoder([], src_type_dict, Vector[str])
+    assert decoder(["hello", "world"]) == ["hello", "world"]
+
+
+def test_decode_error_non_nullable_or_non_list_vector():
+    """Test decoding errors for non-nullable vectors or non-list inputs."""
+    src_type_dict = {
+        "kind": "Vector",
+        "element_type": {"kind": "Float32"},
+        "dimension": None,
+    }
+    decoder = make_engine_value_decoder([], src_type_dict, 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"):
+        decoder("not a list")
+
+
+def test_dump_vector_type_annotation_with_dim():
+    """Test dumping a vector type annotation with a specified dimension."""
+    expected_dump = {
+        "type": {
+            "kind": "Vector",
+            "element_type": {"kind": "Float32"},
+            "dimension": 3,
+        }
+    }
+    assert dump_engine_object(Float32VectorType) == expected_dump
+
+
+def test_dump_vector_type_annotation_no_dim():
+    """Test dumping a vector type annotation with no dimension."""
+    expected_dump_no_dim = {
+        "type": {
+            "kind": "Vector",
+            "element_type": {"kind": "Float64"},
+            "dimension": None,
+        }
+    }
+    assert dump_engine_object(Float64VectorTypeNoDim) == expected_dump_no_dim

cocoindex/tests/test_typing.py

@@ -0,0 +1,499 @@
+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
+import numpy as np
+from numpy.typing import NDArray
+
+from cocoindex.typing import (
+    analyze_type_info,
+    Vector,
+    VectorInfo,
+    TypeKind,
+    TypeAttr,
+    Float32,
+    Float64,
+    encode_enriched_type,
+    AnalyzedTypeInfo,
+)
+
+
+@dataclasses.dataclass
+class SimpleDataclass:
+    name: str
+    value: int
+
+
+class SimpleNamedTuple(NamedTuple):
+    name: str
+    value: Any
+
+
+def test_ndarray_float32_no_dim():
+    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,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_vector_float32_no_dim():
+    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,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_ndarray_float64_with_dim():
+    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,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_vector_float32_with_dim():
+    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,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_ndarray_int64_no_dim():
+    typ = NDArray[np.int64]
+    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
+
+
+def test_ndarray_int32_with_dim():
+    typ = Annotated[NDArray[np.int32], VectorInfo(dim=10)]
+    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 not result.nullable
+
+
+def test_ndarray_uint8_no_dim():
+    typ = NDArray[np.uint8]
+    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 not result.nullable
+
+
+def test_nullable_ndarray():
+    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,
+        attrs=None,
+        nullable=True,
+    )
+
+
+def test_vector_str():
+    typ = Vector[str]
+    result = analyze_type_info(typ)
+    assert result.kind == "Vector"
+    assert result.elem_type == str
+    assert result.vector_info == VectorInfo(dim=None)
+
+
+def test_vector_complex64():
+    typ = Vector[np.complex64]
+    result = analyze_type_info(typ)
+    assert result.kind == "Vector"
+    assert result.elem_type == np.complex64
+    assert result.vector_info == VectorInfo(dim=None)
+
+
+def test_non_numpy_vector():
+    typ = Vector[float, Literal[3]]
+    result = analyze_type_info(typ)
+    assert result.kind == "Vector"
+    assert result.elem_type == float
+    assert result.vector_info == VectorInfo(dim=3)
+
+
+def test_ndarray_any_dtype():
+    typ = NDArray[Any]
+    with pytest.raises(
+        TypeError, match="NDArray for Vector must use a concrete numpy dtype"
+    ):
+        analyze_type_info(typ)
+
+
+def test_list_of_primitives():
+    typ = List[str]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=None),
+        elem_type=str,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_list_of_structs():
+    typ = List[SimpleDataclass]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="LTable",
+        vector_info=None,
+        elem_type=SimpleDataclass,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_sequence_of_int():
+    typ = Sequence[int]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=None),
+        elem_type=int,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_list_with_vector_info():
+    typ = Annotated[List[int], VectorInfo(dim=5)]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=5),
+        elem_type=int,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_dict_str_int():
+    typ = Dict[str, int]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="KTable",
+        vector_info=None,
+        elem_type=(str, int),
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_mapping_str_dataclass():
+    typ = Mapping[str, SimpleDataclass]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="KTable",
+        vector_info=None,
+        elem_type=(str, SimpleDataclass),
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_dataclass():
+    typ = SimpleDataclass
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Struct",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=SimpleDataclass,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_named_tuple():
+    typ = SimpleNamedTuple
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Struct",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=SimpleNamedTuple,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_tuple_key_value():
+    typ = (str, int)
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Int64",
+        vector_info=None,
+        elem_type=None,
+        key_type=str,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_str():
+    typ = str
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Str",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_bool():
+    typ = bool
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Bool",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_bytes():
+    typ = bytes
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Bytes",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_uuid():
+    typ = uuid.UUID
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Uuid",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_date():
+    typ = datetime.date
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Date",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_time():
+    typ = datetime.time
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Time",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_timedelta():
+    typ = datetime.timedelta
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="TimeDelta",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_float():
+    typ = float
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Float64",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_int():
+    typ = int
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Int64",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_type_with_attributes():
+    typ = Annotated[str, TypeAttr("key", "value")]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Str",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        attrs={"key": "value"},
+        nullable=False,
+    )
+
+
+def test_encode_enriched_type_none():
+    typ = None
+    result = encode_enriched_type(typ)
+    assert result is None
+
+
+def test_encode_enriched_type_struct():
+    typ = SimpleDataclass
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Struct"
+    assert len(result["type"]["fields"]) == 2
+    assert result["type"]["fields"][0]["name"] == "name"
+    assert result["type"]["fields"][0]["type"]["kind"] == "Str"
+    assert result["type"]["fields"][1]["name"] == "value"
+    assert result["type"]["fields"][1]["type"]["kind"] == "Int64"
+
+
+def test_encode_enriched_type_vector():
+    typ = NDArray[np.float32]
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Vector"
+    assert result["type"]["element_type"]["kind"] == "Float32"
+    assert result["type"]["dimension"] is None
+
+
+def test_encode_enriched_type_ltable():
+    typ = List[SimpleDataclass]
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "LTable"
+    assert result["type"]["row"]["kind"] == "Struct"
+    assert len(result["type"]["row"]["fields"]) == 2
+
+
+def test_encode_enriched_type_with_attrs():
+    typ = Annotated[str, TypeAttr("key", "value")]
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Str"
+    assert result["attrs"] == {"key": "value"}
+
+
+def test_encode_enriched_type_nullable():
+    typ = str | None
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Str"
+    assert result["nullable"] is True
+
+
+def test_invalid_struct_kind():
+    typ = Annotated[SimpleDataclass, TypeKind("Vector")]
+    with pytest.raises(ValueError, match="Unexpected type kind for struct: Vector"):
+        analyze_type_info(typ)
+
+
+def test_invalid_list_kind():
+    typ = Annotated[List[int], TypeKind("Struct")]
+    with pytest.raises(ValueError, match="Unexpected type kind for list: Struct"):
+        analyze_type_info(typ)
+
+
+def test_unsupported_type():
+    typ = set
+    with pytest.raises(ValueError, match="type unsupported yet: <class 'set'>"):
+        analyze_type_info(typ)

cocoindex/typing.py

@@ -9,14 +9,16 @@ from typing import (
     Annotated,
     NamedTuple,
     Any,
+    KeysView,
     TypeVar,
     TYPE_CHECKING,
     overload,
-    Sequence,
     Generic,
     Literal,
     Protocol,
 )
+import numpy as np
+from numpy.typing import NDArray
 
 
 class VectorInfo(NamedTuple):
@@ -50,7 +52,7 @@ if TYPE_CHECKING:
     Dim_co = TypeVar("Dim_co", bound=int, covariant=True)
 
     class Vector(Protocol, Generic[T_co, Dim_co]):
-        """Vector[T, Dim] is a special typing alias for a list[T] with optional dimension info"""
+        """Vector[T, Dim] is a special typing alias for an NDArray[T] with optional dimension info"""
 
         def __getitem__(self, index: int) -> T_co: ...
         def __len__(self) -> int: ...
@@ -58,25 +60,34 @@ if TYPE_CHECKING:
 else:
 
     class Vector:  # type: ignore[unreachable]
-        """A special typing alias for a list[T] with optional dimension info"""
+        """A special typing alias for an NDArray[T] with optional dimension info"""
 
         def __class_getitem__(self, params):
             if not isinstance(params, tuple):
-                # Only element type provided
-                elem_type = params
-                return Annotated[list[elem_type], VectorInfo(dim=None)]
+                # 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:
+                    return Annotated[NDArray[dtype], VectorInfo(dim=None)]
+                return Annotated[list[dtype], VectorInfo(dim=None)]
             else:
-                # Element type and dimension provided
-                elem_type, dim = params
-                if typing.get_origin(dim) is Literal:
-                    dim = typing.get_args(dim)[0]  # Extract the literal value
-                return Annotated[list[elem_type], VectorInfo(dim=dim)]
+                # Element type and dimension provided, e.g., Vector[np.float32, Literal[3]]
+                dtype, dim_literal = params
+                # Extract the literal value
+                dim_val = (
+                    typing.get_args(dim_literal)[0]
+                    if typing.get_origin(dim_literal) is Literal
+                    else None
+                )
+                if DtypeRegistry.get_by_dtype(dtype) is not None:
+                    return Annotated[NDArray[dtype], VectorInfo(dim=dim_val)]
+                return Annotated[list[dtype], VectorInfo(dim=dim_val)]
 
 
 TABLE_TYPES: tuple[str, str] = ("KTable", "LTable")
 KEY_FIELD_NAME: str = "_key"
 
-ElementType = type | tuple[type, type]
+ElementType = type | tuple[type, type] | Annotated[Any, TypeKind]
 
 
 def is_namedtuple_type(t: type) -> bool:
@@ -89,6 +100,48 @@ 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:
+    _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.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.uint64: DtypeInfo(np.uint64, "UInt64", int),
+    }
+
+    @classmethod
+    def get_by_dtype(cls, dtype: Any) -> DtypeInfo | None:
+        if dtype is Any:
+            raise TypeError(
+                "NDArray for Vector must use a concrete numpy dtype, got `Any`."
+            )
+        return cls._mappings.get(dtype)
+
+    @staticmethod
+    def get_by_kind(kind: str) -> DtypeInfo | None:
+        return next(
+            (info for info in DtypeRegistry._mappings.values() if info.kind == kind),
+            None,
+        )
+
+    @staticmethod
+    def supported_dtypes() -> KeysView[type]:
+        return DtypeRegistry._mappings.keys()
+
+
 @dataclasses.dataclass
 class AnalyzedTypeInfo:
     """
@@ -179,12 +232,34 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
                     vector_info = VectorInfo(dim=None)
         elif not (kind == "Vector" or kind in TABLE_TYPES):
             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")
+
+        numpy_dtype = dtype_args[0]
+        dtype_info = DtypeRegistry.get_by_dtype(numpy_dtype)
+        if dtype_info is None:
+            raise ValueError(
+                f"Unsupported numpy dtype for NDArray: {numpy_dtype}. "
+                f"Supported dtypes: {DtypeRegistry.supported_dtypes()}"
+            )
+        elem_type = dtype_info.annotated_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:
         args = typing.get_args(t)
         elem_type = (args[0], args[1])
         kind = "KTable"
     elif kind is None:
-        if t is bytes:
+        dtype_info = DtypeRegistry.get_by_dtype(t)
+        if dtype_info is not None:
+            kind = dtype_info.kind
+        elif t is bytes:
             kind = "Bytes"
         elif t is str:
             kind = "Str"

{cocoindex-0.1.49.dist-info → cocoindex-0.1.50.dist-info}/METADATA

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

{cocoindex-0.1.49.dist-info → cocoindex-0.1.50.dist-info}/RECORD

@@ -1,27 +1,27 @@
-cocoindex-0.1.49.dist-info/METADATA,sha256=8Fow__E5M04be_uy0-Gz-ojEaokr4scBR4lQdRMEVGo,9875
-cocoindex-0.1.49.dist-info/WHEEL,sha256=p_tvkyHH2UmMBrR2Gemb1ahXJMM2SXUIsCLrWZgJvB8,104
-cocoindex-0.1.49.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.1.49.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-cocoindex/__init__.py,sha256=TkfOR-7upXpzETEvBQUhcnYVuI_A3sqjOdXYeTPAwO0,818
-cocoindex/_engine.cpython-311-darwin.so,sha256=LdbQ9EL9P0L19Xq9Nh07NNZY8MYTcv4O4CiJIpI7p7s,57420944
+cocoindex-0.1.50.dist-info/METADATA,sha256=PdFYx9yT4tBqCXVaP77iQjRZMB26D2I8BWXoA8i7ZqQ,9875
+cocoindex-0.1.50.dist-info/WHEEL,sha256=8hL2oHqulIPF_TJ_OfB-8kD2X9O6HJL0emTmeguFbqc,104
+cocoindex-0.1.50.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.1.50.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+cocoindex/__init__.py,sha256=BC6BvotY0LcKHCUNnZWjKZX7Afl18TDR2Bfv5R6Y4DM,811
+cocoindex/_engine.cpython-311-darwin.so,sha256=hqncg8nnjBaQicxXkJbOf4Hl3-UoE9onm_Q7On3dvQk,57265296
 cocoindex/auth_registry.py,sha256=1XqO7ibjmBBd8i11XSJTvTgdz8p1ptW-ZpuSgo_5zzk,716
 cocoindex/cli.py,sha256=p92s-Ya6VpT4EY5xfbl-R0yyIJv8ySK-2eyxnwC8KqA,17853
-cocoindex/convert.py,sha256=CfsZDRs1kwMeNCeJ3-gifHQm_9GzyqySpUmZyCyWXMo,7316
+cocoindex/convert.py,sha256=dxlMUd6Jweun1xESJv0w-WbIx-NifMXlUF9RqgeggWM,8735
 cocoindex/flow.py,sha256=sC7eg0L9MkZuo7MlvGA4eYe-mEReIyPjKIILFtkWm-Y,30017
-cocoindex/functions.py,sha256=RnrYJFJCfj263wow3rEx89UBtff44ldQ_torgwArPEw,2294
+cocoindex/functions.py,sha256=9A61Jj5a3vQoI2MIAhjXvJrDxSzDhe6VncQWbiVtwcg,2393
 cocoindex/index.py,sha256=j93B9jEvvLXHtpzKWL88SY6wCGEoPgpsQhEGHlyYGFg,540
-cocoindex/lib.py,sha256=W3hPh0QQIjLRRe2tpyLCvL_6ajzXhGGSNJubikkS27s,2985
+cocoindex/lib.py,sha256=BeRUn3RqE_wSsVtsgCzbFFKe1LXgRyRmMOcmwWBuEXo,2940
 cocoindex/llm.py,sha256=KO-R4mrAWtxXD82-Yv5ixpkKMVfkwpbdWwqPVZygLu4,352
 cocoindex/op.py,sha256=hOJoHC8elhLHMNVRTGTBWNSK5uYrQb-FfRV-qt08j8g,11815
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cocoindex/query.py,sha256=UswtEp0wZwwNPeR9x9lRokFpNNZOHkyt5tYfCmGIWWU,3450
 cocoindex/runtime.py,sha256=bAdHYaXFWiiUWyAgzmKTeaAaRR0D_AmaqVCIdPO-v00,1056
 cocoindex/setting.py,sha256=ePqHw1i95rVtqYYRVqzVwtBifRO4SfH1rlMW_AG3Zek,3418
 cocoindex/setup.py,sha256=u5dYZFKfz4yZLiGHD0guNaR0s4zY9JAoZWrWHpAHw_0,773
 cocoindex/sources.py,sha256=JCnOhv1w4o28e03i7yvo4ESicWYAhckkBg5bQlxNH4U,1330
 cocoindex/storages.py,sha256=i9cPzj5LbLJVHjskEWyltUXF9kqCyc6ftMEqll6FDwI,2804
 cocoindex/tests/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-cocoindex/tests/test_convert.py,sha256=tTsDveMCgflhcYmDq-fQq3MRGAmN2xlbtFIEDYlRpGw,17095
-cocoindex/typing.py,sha256=0WQ4F04kFhXwHphv9AHJqc1cThYmLe27eqqGEuGWHAU,9462
+cocoindex/tests/test_convert.py,sha256=EsR7NEGKAypQpr3FxRz8xja_z0sKmYpTjJuCMSUdyYw,26217
+cocoindex/tests/test_typing.py,sha256=P6bnRjxMjKH2B5RQ0ncZQAMyOPSdviFVmML7fCS2iR4,11884
+cocoindex/typing.py,sha256=9366o49R6-YbrG9rc9UOxJjgTYITC-7DjDkJSVIjZL0,12316
 cocoindex/utils.py,sha256=5a3ubVzDowJUJyUl8ecd75Th_OzON3G-MV2vXf0dQSk,503
-cocoindex-0.1.49.dist-info/RECORD,,
+cocoindex-0.1.50.dist-info/RECORD,,

{cocoindex-0.1.49.dist-info → cocoindex-0.1.50.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.8.6)
+Generator: maturin (1.8.7)
 Root-Is-Purelib: false
 Tag: cp311-cp311-macosx_11_0_arm64

cocoindex/query.py

@@ -1,115 +0,0 @@
-from typing import Callable, Any
-from dataclasses import dataclass
-from threading import Lock
-
-from . import flow as fl
-from . import index
-from . import _engine  # type: ignore
-
-_handlers_lock = Lock()
-_handlers: dict[str, _engine.SimpleSemanticsQueryHandler] = {}
-
-
-@dataclass
-class SimpleSemanticsQueryInfo:
-    """
-    Additional information about the query.
-    """
-
-    similarity_metric: index.VectorSimilarityMetric
-    query_vector: list[float]
-    vector_field_name: str
-
-
-@dataclass
-class QueryResult:
-    """
-    A single result from the query.
-    """
-
-    data: dict[str, Any]
-    score: float
-
-
-class SimpleSemanticsQueryHandler:
-    """
-    A query handler that uses simple semantics to query the index.
-    """
-
-    _lazy_query_handler: Callable[[], _engine.SimpleSemanticsQueryHandler]
-
-    def __init__(
-        self,
-        name: str,
-        flow: fl.Flow,
-        target_name: str,
-        query_transform_flow: Callable[..., fl.DataSlice[Any]],
-        default_similarity_metric: index.VectorSimilarityMetric = index.VectorSimilarityMetric.COSINE_SIMILARITY,
-    ) -> None:
-        engine_handler = None
-        lock = Lock()
-
-        def _lazy_handler() -> _engine.SimpleSemanticsQueryHandler:
-            nonlocal engine_handler, lock
-            if engine_handler is None:
-                with lock:
-                    if engine_handler is None:
-                        engine_handler = _engine.SimpleSemanticsQueryHandler(
-                            flow.internal_flow(),
-                            target_name,
-                            fl.TransformFlow(
-                                query_transform_flow, [str]
-                            ).internal_flow(),
-                            default_similarity_metric.value,
-                        )
-                        engine_handler.register_query_handler(name)
-            return engine_handler
-
-        self._lazy_query_handler = _lazy_handler
-
-        with _handlers_lock:
-            _handlers[name] = self
-
-    def internal_handler(self) -> _engine.SimpleSemanticsQueryHandler:
-        """
-        Get the internal query handler.
-        """
-        return self._lazy_query_handler()
-
-    def search(
-        self,
-        query: str,
-        limit: int,
-        vector_field_name: str | None = None,
-        similarity_metric: index.VectorSimilarityMetric | None = None,
-    ) -> tuple[list[QueryResult], SimpleSemanticsQueryInfo]:
-        """
-        Search the index with the given query, limit, vector field name, and similarity metric.
-        """
-        internal_results, internal_info = self.internal_handler().search(
-            query,
-            limit,
-            vector_field_name,
-            similarity_metric.value if similarity_metric is not None else None,
-        )
-        results = [
-            QueryResult(data=result["data"], score=result["score"])
-            for result in internal_results
-        ]
-        info = SimpleSemanticsQueryInfo(
-            similarity_metric=index.VectorSimilarityMetric(
-                internal_info["similarity_metric"]
-            ),
-            query_vector=internal_info["query_vector"],
-            vector_field_name=internal_info["vector_field_name"],
-        )
-        return results, info
-
-
-def ensure_all_handlers_built() -> None:
-    """
-    Ensure all handlers are built.
-    """
-    with _handlers_lock:
-        for handler in _handlers.values():
-            handler.internal_handler()