lancedb/docs/web-tests/py/test_basic_usage.py

# 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