diff --git a/python/python/tests/docs/test_basic.py b/python/python/tests/docs/test_basic.py index 2a824371f..1eaa4b247 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 069527b21..0771db5d2 100644 --- a/python/python/tests/test_table.py +++ b/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()