From 7952ee2a71d5bb921d097f3cd81f8039bf014818 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:31:19 +0800 Subject: [PATCH] fix(python): cover FP16 GPU index queries (#3819) Keep FP16 vectors through IVF-PQ index creation and search in the synchronous and asynchronous quickstarts. Add native and conditional CUDA regression coverage using a float64 query against an explicitly typed FP16 column. This is test and documentation-example coverage for dtype preservation already present in Lance; it does not change production code. Related #2120. Validation on `bbad20940632b21a9ee85bdcb46b4cc19dbd65e3` with Python 3.12.12, Lance 15.0.0-beta.1, PyArrow 24.0.0, and Torch 2.14.1: - The native FP16 parameter, Arrow FP16 ingestion test, and both quickstarts passed: 4 passed, 1 CUDA-only skip. - The two affected test files passed: 386 passed, 1 CUDA-only skip (`-m "not slow and not s3_test"`). - Ruff 0.9.9 lint and formatting checks passed for the Python package. Real CUDA execution has not been performed. Torch is installed, but `torch.cuda.is_available()` is false; `test_create_f16_table[cuda]` skips with `CUDA not available`. The native CPU result does not establish that the original GPU failure is resolved. Issue #2120 remains open for verification on an actual CUDA device. Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/tests/docs/test_basic.py | 16 ++-------------- python/python/tests/test_table.py | 26 +++++++++++++++++++++----- 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/python/python/tests/docs/test_basic.py b/python/python/tests/docs/test_basic.py index 35d7aac10..ef68114aa 100644 --- a/python/python/tests/docs/test_basic.py +++ b/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] diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 3b3ccd122..97f65b180 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -3965,26 +3965,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()