mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-13 15:22:30 +00:00
Jack Ye
b21fb2a4dc
The pages describing this code lived in lancedb/docs, a repository with ten stars that neither contributing guide mentioned. A user who hit a wrong page while reading the source had to find a separate, near-invisible repository to fix it. These 102 pages now sit beside the code, and a documentation fix is a pull request in the repository the reader already had open. `docs/web` is a complete Mintlify site: `cd docs/web && mint dev` serves it with no other checkout. Enterprise pages stay private and are merged in when the published site is assembled, so links into them do not resolve here — every open-source page does. `docs/web-tests` holds the tests the code examples are extracted from. They ran against a published wheel before; here they can run against this repository's own SDK, so an example that stops compiling fails the pull request that broke it rather than being noticed later. Both directories sit beside `docs/src` and are invisible to it: mkdocs reads `docs_dir: src`, so the SDK reference builds and deploys exactly as before.
206 lines
6.3 KiB
Python
206 lines
6.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
# --8<-- [start:basic_imports]
|
|
import json
|
|
|
|
import lancedb
|
|
import pandas as pd
|
|
import polars as pl
|
|
import pyarrow as pa
|
|
# --8<-- [end:basic_imports]
|
|
import pytest
|
|
|
|
data_path = "tests/camelot.json"
|
|
|
|
def test_basic_usage(db_path_factory):
|
|
uri = "./basic_usage_db"
|
|
uri = db_path_factory("basic_usage_db")
|
|
db = lancedb.connect(uri)
|
|
|
|
# --8<-- [start:data_load]
|
|
with open(data_path, "r") as f:
|
|
data = json.load(f)
|
|
# --8<-- [end:data_load]
|
|
|
|
# --8<-- [start:basic_create_table]
|
|
table = db.create_table("camelot", data=data, mode="overwrite")
|
|
# --8<-- [end:basic_create_table]
|
|
assert len(table) == 8
|
|
|
|
# --8<-- [start:basic_open_table]
|
|
table = db.open_table("camelot")
|
|
# --8<-- [end:basic_open_table]
|
|
|
|
# --8<-- [start:basic_create_table_pandas]
|
|
pandas_df = pd.DataFrame(data)
|
|
table_pd = db.create_table("camelot_pd", data=pandas_df, mode="overwrite")
|
|
# --8<-- [end:basic_create_table_pandas]
|
|
assert len(table_pd) == 8
|
|
db.drop_table("camelot_pd")
|
|
|
|
# --8<-- [start:basic_create_table_polars]
|
|
polars_df = pl.DataFrame(data)
|
|
table_pl = db.create_table("camelot_pl", data=polars_df, mode="overwrite")
|
|
# --8<-- [end:basic_create_table_polars]
|
|
assert len(table_pl) == 8
|
|
db.drop_table("camelot_pl")
|
|
|
|
# --8<-- [start:basic_create_empty_table]
|
|
schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.uint16()),
|
|
pa.field("name", pa.string()),
|
|
pa.field("role", pa.string()),
|
|
pa.field("description", pa.string()),
|
|
pa.field("vector", pa.list_(pa.float32(), 4)),
|
|
pa.field(
|
|
"stats",
|
|
pa.struct(
|
|
[
|
|
pa.field("strength", pa.int8()),
|
|
pa.field("courage", pa.int8()),
|
|
pa.field("magic", pa.int8()),
|
|
pa.field("wisdom", pa.int8()),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
)
|
|
db.create_table("camelot_pa", schema=schema, mode="overwrite")
|
|
# --8<-- [end:basic_create_empty_table]
|
|
assert "camelot_pa" in db.list_tables().tables
|
|
db.drop_table("camelot_pa")
|
|
|
|
# --8<-- [start:basic_add_data]
|
|
magical_characters = [
|
|
{
|
|
"id": 9,
|
|
"name": "Morgan le Fay",
|
|
"role": "Sorceress",
|
|
"description": "A powerful enchantress, Arthur's half-sister, and a complex figure who oscillates between aiding and opposing Camelot.",
|
|
"vector": [0.10, 0.84, 0.25, 0.70],
|
|
"stats": { "strength": 2, "courage": 3, "magic": 5, "wisdom": 4 }
|
|
},
|
|
{
|
|
"id": 10,
|
|
"name": "The Lady of the Lake",
|
|
"role": "Mystical Guardian",
|
|
"description": "A mysterious supernatural figure associated with Avalon, known for giving Arthur the sword Excalibur.",
|
|
"vector": [0.00, 0.90, 0.58, 0.88],
|
|
"stats": { "strength": 2, "courage": 3, "magic": 5, "wisdom": 5 }
|
|
}
|
|
]
|
|
table.add(magical_characters)
|
|
# --8<-- [end:basic_add_data]
|
|
assert len(table) == 10
|
|
|
|
# --8<-- [start:basic_vector_search]
|
|
query_vector = [0.03, 0.85, 0.61, 0.90]
|
|
result = table.search(query_vector).limit(5).to_polars()
|
|
print(result)
|
|
# --8<-- [end:basic_vector_search]
|
|
|
|
# --8<-- [start:basic_add_columns]
|
|
table.add_columns(
|
|
{
|
|
"power": "cast(((stats.strength + stats.courage + stats.magic + stats.wisdom) / 4.0) as float)"
|
|
}
|
|
)
|
|
# --8<-- [end:basic_add_columns]
|
|
assert "power" in table.schema.names
|
|
|
|
# Run examples to illustrate search
|
|
# --8<-- [start:basic_vector_search_q1]
|
|
# Who are the characters similar to "wizard"?
|
|
query_vector_1 = [0.03, 0.85, 0.61, 0.90]
|
|
r1 = (
|
|
table.search(query_vector_1)
|
|
.limit(5)
|
|
.select(["name", "role", "description"])
|
|
.to_polars()
|
|
)
|
|
print(r1)
|
|
# --8<-- [end:basic_vector_search_q1]
|
|
|
|
# --8<-- [start:basic_vector_search_q2]
|
|
# Who are the characters with high magic stats?
|
|
query_vector_2 = [0.03, 0.85, 0.61, 0.90]
|
|
r2 = (
|
|
table.search(query_vector_2)
|
|
.where("stats.magic > 3")
|
|
.select(["name", "role", "description"])
|
|
.limit(5)
|
|
.to_polars()
|
|
)
|
|
print(r2)
|
|
# --8<-- [end:basic_vector_search_q2]
|
|
|
|
# --8<-- [start:basic_vector_search_q3]
|
|
# Who are the strongest characters?
|
|
r3 = (
|
|
table.search()
|
|
.where("stats.strength > 3")
|
|
.select(["name", "role", "description"])
|
|
.limit(5)
|
|
.to_polars()
|
|
)
|
|
print(r3)
|
|
# --8<-- [end:basic_vector_search_q3]
|
|
|
|
# --8<-- [start:basic_vector_search_q4]
|
|
# Who are the strongest characters?
|
|
r4 = (
|
|
table.search()
|
|
.select(["name", "role", "description", "power"])
|
|
.to_polars()
|
|
)
|
|
print(r4)
|
|
# --8<-- [end:basic_vector_search_q4]
|
|
|
|
# --8<-- [start:basic_sort_polars]
|
|
# Sort Polars DataFrame by power in descending order
|
|
print(r4.sort("power", descending=True).limit(5))
|
|
# --8<-- [end:basic_sort_polars]
|
|
sorted_power = r4.sort("power", descending=True)["power"].to_list()
|
|
assert sorted_power == sorted(sorted_power, reverse=True)
|
|
assert sorted_power[0] == 4.0
|
|
|
|
# --8<-- [start:basic_drop_columns]
|
|
table.drop_columns(["power"])
|
|
# --8<-- [end:basic_drop_columns]
|
|
# --8<-- [start:basic_delete_rows]
|
|
table.delete('role = "Traitor Knight"')
|
|
# --8<-- [end:basic_delete_rows]
|
|
assert len(table) == 9
|
|
# --8<-- [start:basic_drop_table]
|
|
db.drop_table("camelot")
|
|
# --8<-- [end:basic_drop_table]
|
|
assert "camelot" not in db.list_tables().tables
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_basic_usage_async_api(db_path_factory):
|
|
uri = db_path_factory("basic_usage_async_db")
|
|
with open(data_path, "r") as f:
|
|
data = json.load(f)
|
|
|
|
# --8<-- [start:basic_async_api]
|
|
import lancedb
|
|
|
|
async_db = await lancedb.connect_async(uri)
|
|
async_table = await async_db.create_table(
|
|
"camelot_async",
|
|
data=data,
|
|
mode="overwrite",
|
|
)
|
|
|
|
query_vector = [0.03, 0.85, 0.61, 0.90]
|
|
async_results = await (
|
|
await async_table.search(query_vector)
|
|
).limit(5).select(["name", "role", "description"]).to_polars()
|
|
print(async_results)
|
|
# --8<-- [end:basic_async_api]
|
|
|
|
assert async_results.height >= 1
|