mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-12 16:22:24 +00:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b5158f62ab | ||
|
|
191e9eed8d |
@@ -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]
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user