Merge remote-tracking branch 'origin/main' into gatekeeper/fix-2120-1

# Conflicts:
#	python/python/tests/test_table.py

python/python/lancedb/embeddings/base.py

@@ -183,11 +183,7 @@ class EmbeddingFunction(BaseModel, ABC):
     def VectorField(self, **kwargs):
         """
         Creates a pydantic Field that can automatically annotate
-        the target vector column for this embedding function.
-
-        The field can be annotated as ``list[float]`` for compatibility with
-        static type checkers. LanceDB will infer the fixed vector dimension from
-        this embedding function.
+        the target vector column for this embedding function
         """
         return Field(json_schema_extra={"vector_column_for": self}, **kwargs)
 

python/python/lancedb/pydantic.py

@@ -99,17 +99,6 @@ def Vector(
     ...     pa.field("url", pa.utf8(), False),
     ...     pa.field("embeddings", pa.list_(pa.float32(), 768))
     ... ])
-
-    Notes
-    -----
-    ``Vector`` creates a type dynamically, so calls such as ``Vector(768)`` are
-    not valid static type annotations. For an embedding field, use the standard
-    ``list[float]`` annotation when running mypy; ``VectorField`` supplies the
-    fixed dimension to LanceDB::
-
-        class MyModel(LanceModel):
-            text: str = embeddings.SourceField()
-            vector: list[float] = embeddings.VectorField()
     """
 
     # TODO: make a public parameterized type.
@@ -390,10 +379,6 @@ def _unwrap_optional_annotation(annotation: Any) -> Any | None:
 
 def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
     """Convert a Pydantic FieldInfo to Arrow DataType"""
-    embedding_vector_type = _embedding_vector_to_arrow_type(field)
-    if embedding_vector_type is not None:
-        return embedding_vector_type
-
     unwrapped = _unwrap_optional_annotation(field.annotation)
     if unwrapped is not None:
         return _pydantic_type_to_arrow_type(unwrapped, field)
@@ -407,32 +392,8 @@ def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
     return _pydantic_type_to_arrow_type(field.annotation, field)
 
 
-def _embedding_vector_to_arrow_type(field: FieldInfo) -> pa.DataType | None:
-    """Infer a fixed-size vector type from ``VectorField`` metadata."""
-    if not _is_embedding_vector_annotation(field):
-        return None
-
-    function = get_extras(field, "vector_column_for")
-    return pa.list_(pa.float32(), function.ndims())
-
-
-def _is_embedding_vector_annotation(field: FieldInfo) -> bool:
-    if get_extras(field, "vector_column_for") is None:
-        return False
-
-    annotation = _unwrap_optional_annotation(field.annotation)
-    if annotation is None:
-        annotation = field.annotation
-
-    origin = getattr(annotation, "__origin__", None)
-    args = getattr(annotation, "__args__", ())
-    return origin is list and args == (float,)
-
-
 def is_nullable(field: FieldInfo) -> bool:
     """Check if a Pydantic FieldInfo is nullable."""
-    if _is_embedding_vector_annotation(field):
-        return True
     if _unwrap_optional_annotation(field.annotation) is not None:
         return True
     if isinstance(field.annotation, (_GenericAlias, GenericAlias)):

python/python/tests/docs/test_basic.py

@@ -91,19 +91,13 @@ def test_quickstart(tmp_path):
         }
     )
     # --8<-- [end:alter_columns_vector]
-    # Change it back since we can get a panic with fp16
-    tbl.alter_columns(
-        {
-            "path": "vector",
-            "data_type": pa.list_(pa.float32(), list_size=2),
-        }
-    )
     # --8<-- [start:drop_columns]
     tbl.drop_columns(["dbl_price"])
     # --8<-- [end:drop_columns]
     # --8<-- [start:create_index]
     tbl.create_index(num_sub_vectors=1)
     # --8<-- [end:create_index]
+    tbl.search([100, 100]).limit(2).to_pandas()
     # --8<-- [start:delete_rows]
     tbl.delete('item = "fizz"')
     # --8<-- [end:delete_rows]
@@ -185,13 +179,6 @@ async def test_quickstart_async(tmp_path):
         }
     )
     # --8<-- [end:alter_columns_async_vector]
-    # Change it back since we can get a panic with fp16
-    await tbl.alter_columns(
-        {
-            "path": "vector",
-            "data_type": pa.list_(pa.float32(), list_size=2),
-        }
-    )
     # --8<-- [start:drop_columns_async]
     await tbl.drop_columns(["dbl_price"])
     # --8<-- [end:drop_columns_async]
@@ -200,6 +187,7 @@ async def test_quickstart_async(tmp_path):
     # --8<-- [start:create_index_async]
     await tbl.create_index("vector")
     # --8<-- [end:create_index_async]
+    await tbl.vector_search([100, 100]).limit(2).to_pandas()
     # --8<-- [start:delete_rows_async]
     await tbl.delete('item = "fizz"')
     # --8<-- [end:delete_rows_async]

python/python/tests/test_pydantic.py

@@ -9,7 +9,6 @@ from typing import List, Optional, Tuple
 import pyarrow as pa
 import pydantic
 import pytest
-from lancedb.conftest import MockTextEmbeddingFunction
 from lancedb.pydantic import (
     PYDANTIC_VERSION,
     LanceModel,
@@ -427,25 +426,6 @@ def test_bare_vector_raises_clear_error():
         exec("class TestModel(LanceModel):\n    vector: Vector", namespace)
 
 
-def test_embedding_vector_list_annotation():
-    embedding = MockTextEmbeddingFunction.create()
-
-    class StaticTypingModel(LanceModel):
-        text: str = embedding.SourceField()
-        vector: list[float] = embedding.VectorField()
-
-    schema = pydantic_to_schema(StaticTypingModel)
-    assert schema == pa.schema(
-        [
-            pa.field("text", pa.utf8(), False),
-            pa.field("vector", pa.list_(pa.float32(), embedding.ndims()), True),
-        ]
-    )
-
-    model = StaticTypingModel(text="hello", vector=[0.0] * embedding.ndims())
-    assert model.vector == [0.0] * embedding.ndims()
-
-
 def test_fixed_size_list_field():
     class TestModel(pydantic.BaseModel):
         vec: Vector(16)

python/python/tests/test_table.py

@@ -2823,26 +2823,42 @@ def test_create_f16_table_from_arrow_data(mem_db: DBConnection):
     assert "s-2" in expected["text"].to_pylist()
 
 
-def test_create_f16_table(mem_db: DBConnection):
+@pytest.mark.parametrize("accelerator", [None, "cuda"])
+def test_create_f16_table(tmp_path, accelerator):
+    if accelerator == "cuda":
+        torch = pytest.importorskip("torch")
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+
     class MyTable(LanceModel):
         text: str
         vector: Vector(32, value_type=pa.float16())
 
     rng = np.random.default_rng(42)
+    vectors = rng.standard_normal((512, 32)).astype(np.float16)
     df = pa.table(
         {
             "text": [f"s-{i}" for i in range(512)],
-            "vector": [rng.standard_normal(32).astype(np.float16) for _ in range(512)],
+            "vector": list(vectors),
         }
     )
-    table = mem_db.create_table(
+    db = lancedb.connect(tmp_path)
+    table = db.create_table(
         "f16_tbl",
         schema=MyTable,
     )
     table.add(df)
-    table.create_index(num_partitions=2, num_sub_vectors=2)
+    table.create_index(
+        "vector",
+        config=IvfPq(
+            num_partitions=2,
+            num_sub_vectors=2,
+            accelerator=accelerator,
+        ),
+    )
 
-    query = df["vector"][2].as_py()
+    # Match the issue's float64 query against an explicitly typed float16 column.
+    query = vectors[2].astype(np.float64)
     expected = table.search(query).limit(2).to_arrow()
 
     assert "s-2" in expected["text"].to_pylist()