fix(python): cover FP16 GPU index queries

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_table.py

@@ -2559,25 +2559,42 @@ def test_create_with_embedding_function(mem_db: DBConnection):
     assert actual == expected
 
 
-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": [np.random.randn(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()