Compare commits

...
Author SHA1 Message Date
Gatefixer b5158f62ab Merge remote-tracking branch 'origin/main' into gatekeeper/fix-2120-1
# Conflicts:
#	python/python/tests/test_table.py
2026-08-07 10:07:52 +00:00
Gatefixer 191e9eed8d fix(python): cover FP16 GPU index queries 2026-08-05 23:29:00 +00:00
2 changed files with 23 additions and 19 deletions
+2 -14
View File
@@ -91,19 +91,13 @@ def test_quickstart(tmp_path):
} }
) )
# --8<-- [end:alter_columns_vector] # --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] # --8<-- [start:drop_columns]
tbl.drop_columns(["dbl_price"]) tbl.drop_columns(["dbl_price"])
# --8<-- [end:drop_columns] # --8<-- [end:drop_columns]
# --8<-- [start:create_index] # --8<-- [start:create_index]
tbl.create_index(num_sub_vectors=1) tbl.create_index(num_sub_vectors=1)
# --8<-- [end:create_index] # --8<-- [end:create_index]
tbl.search([100, 100]).limit(2).to_pandas()
# --8<-- [start:delete_rows] # --8<-- [start:delete_rows]
tbl.delete('item = "fizz"') tbl.delete('item = "fizz"')
# --8<-- [end:delete_rows] # --8<-- [end:delete_rows]
@@ -185,13 +179,6 @@ async def test_quickstart_async(tmp_path):
} }
) )
# --8<-- [end:alter_columns_async_vector] # --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] # --8<-- [start:drop_columns_async]
await tbl.drop_columns(["dbl_price"]) await tbl.drop_columns(["dbl_price"])
# --8<-- [end:drop_columns_async] # --8<-- [end:drop_columns_async]
@@ -200,6 +187,7 @@ async def test_quickstart_async(tmp_path):
# --8<-- [start:create_index_async] # --8<-- [start:create_index_async]
await tbl.create_index("vector") await tbl.create_index("vector")
# --8<-- [end:create_index_async] # --8<-- [end:create_index_async]
await tbl.vector_search([100, 100]).limit(2).to_pandas()
# --8<-- [start:delete_rows_async] # --8<-- [start:delete_rows_async]
await tbl.delete('item = "fizz"') await tbl.delete('item = "fizz"')
# --8<-- [end:delete_rows_async] # --8<-- [end:delete_rows_async]
+21 -5
View File
@@ -2823,26 +2823,42 @@ def test_create_f16_table_from_arrow_data(mem_db: DBConnection):
assert "s-2" in expected["text"].to_pylist() 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): class MyTable(LanceModel):
text: str text: str
vector: Vector(32, value_type=pa.float16()) vector: Vector(32, value_type=pa.float16())
rng = np.random.default_rng(42) rng = np.random.default_rng(42)
vectors = rng.standard_normal((512, 32)).astype(np.float16)
df = pa.table( df = pa.table(
{ {
"text": [f"s-{i}" for i in range(512)], "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", "f16_tbl",
schema=MyTable, schema=MyTable,
) )
table.add(df) 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() expected = table.search(query).limit(2).to_arrow()
assert "s-2" in expected["text"].to_pylist() assert "s-2" in expected["text"].to_pylist()