# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors # --8<-- [start:import-lancedb] import lancedb # --8<-- [end:import-lancedb] # --8<-- [start:import-numpy] from lancedb.query import BoostQuery, MatchQuery import numpy as np import pyarrow as pa # --8<-- [end:import-numpy] # --8<-- [start:import-datetime] from datetime import datetime # --8<-- [end:import-datetime] # --8<-- [start:import-lancedb-pydantic] from lancedb.pydantic import Vector, LanceModel # --8<-- [end:import-lancedb-pydantic] # --8<-- [start:import-pydantic-base-model] from pydantic import BaseModel # --8<-- [end:import-pydantic-base-model] # --8<-- [start:import-lancedb-fts] from lancedb.index import FTS # --8<-- [end:import-lancedb-fts] # --8<-- [start:import-os] import os # --8<-- [end:import-os] # --8<-- [start:import-embeddings] from lancedb.embeddings import get_registry # --8<-- [end:import-embeddings] import pytest # --8<-- [start:class-definition] class Metadata(BaseModel): source: str timestamp: datetime class Document(BaseModel): content: str meta: Metadata class LanceSchema(LanceModel): id: str vector: Vector(1536) payload: Document # --8<-- [end:class-definition] def test_vector_search(): # --8<-- [start:exhaustive_search] uri = "data/sample-lancedb" db = lancedb.connect(uri) data = [ {"vector": row, "item": f"item {i}"} for i, row in enumerate(np.random.random((10_000, 1536)).astype("float32")) ] tbl = db.create_table("vector_search", data=data, mode="overwrite") tbl.search(np.random.random((1536))).limit(10).to_list() # --8<-- [end:exhaustive_search] # --8<-- [start:exhaustive_search_cosine] tbl.search(np.random.random((1536))).distance_type("cosine").limit(10).to_list() # --8<-- [end:exhaustive_search_cosine] # --8<-- [start:create_table_with_nested_schema] # Let's add 100 sample rows to our dataset data = [ LanceSchema( id=f"id{i}", vector=np.random.randn(1536), payload=Document( content=f"document{i}", meta=Metadata(source=f"source{i % 10}", timestamp=datetime.now()), ), ) for i in range(100) ] # Synchronous client tbl = db.create_table("documents", data=data, mode="overwrite") # --8<-- [end:create_table_with_nested_schema] # --8<-- [start:search_result_as_pyarrow] tbl.search(np.random.randn(1536)).to_arrow() # --8<-- [end:search_result_as_pyarrow] # --8<-- [start:search_result_as_pandas] tbl.search(np.random.randn(1536)).to_pandas() # --8<-- [end:search_result_as_pandas] # --8<-- [start:search_result_as_pandas_flatten_true] tbl.search(np.random.randn(1536)).to_pandas(flatten=True) # --8<-- [end:search_result_as_pandas_flatten_true] # --8<-- [start:search_result_as_pandas_flatten_1] tbl.search(np.random.randn(1536)).to_pandas(flatten=1) # --8<-- [end:search_result_as_pandas_flatten_1] # --8<-- [start:search_result_as_list] tbl.search(np.random.randn(1536)).to_list() # --8<-- [end:search_result_as_list] # --8<-- [start:search_result_as_pydantic] tbl.search(np.random.randn(1536)).to_pydantic(LanceSchema) # --8<-- [end:search_result_as_pydantic] @pytest.mark.asyncio async def test_vector_search_async(): # --8<-- [start:exhaustive_search_async] uri = "data/sample-lancedb" async_db = await lancedb.connect_async(uri) data = [ {"vector": row, "item": f"item {i}"} for i, row in enumerate(np.random.random((10_000, 1536)).astype("float32")) ] async_tbl = await async_db.create_table( "vector_search_async", data=data, mode="overwrite" ) (await (await async_tbl.search(np.random.random((1536)))).limit(10).to_list()) # --8<-- [end:exhaustive_search_async] # --8<-- [start:exhaustive_search_async_cosine] ( await (await async_tbl.search(np.random.random((1536)))) .distance_type("cosine") .limit(10) .to_list() ) # --8<-- [end:exhaustive_search_async_cosine] # --8<-- [start:create_table_async_with_nested_schema] # Let's add 100 sample rows to our dataset data = [ LanceSchema( id=f"id{i}", vector=np.random.randn(1536), payload=Document( content=f"document{i}", meta=Metadata(source=f"source{i % 10}", timestamp=datetime.now()), ), ) for i in range(100) ] async_tbl = await async_db.create_table( "documents_async", data=data, mode="overwrite" ) # --8<-- [end:create_table_async_with_nested_schema] # --8<-- [start:search_result_async_as_pyarrow] await (await async_tbl.search(np.random.randn(1536))).to_arrow() # --8<-- [end:search_result_async_as_pyarrow] # --8<-- [start:search_result_async_as_pandas] await (await async_tbl.search(np.random.randn(1536))).to_pandas() # --8<-- [end:search_result_async_as_pandas] # --8<-- [start:search_result_async_as_list] await (await async_tbl.search(np.random.randn(1536))).to_list() # --8<-- [end:search_result_async_as_list] def test_fts_fuzzy_query(): uri = "data/fuzzy-example" db = lancedb.connect(uri) table = db.create_table( "my_table_fts_fuzzy", data=pa.table( { "text": [ "fa", "fo", # spellchecker:disable-line "fob", "focus", "foo", "food", "foul", ] } ), mode="overwrite", ) table.create_fts_index("text", replace=True) results = table.search(MatchQuery("foo", "text", fuzziness=1)).to_pandas() assert len(results) == 4 assert set(results["text"].to_list()) == { "foo", "fo", # 1 deletion # spellchecker:disable-line "fob", # 1 substitution "food", # 1 insertion } def test_fts_boost_query(): uri = "data/boost-example" db = lancedb.connect(uri) table = db.create_table( "my_table_fts_boost", data=pa.table( { "title": [ "The Hidden Gems of Travel", "Exploring Nature's Wonders", "Cultural Treasures Unveiled", "The Nightlife Chronicles", "Scenic Escapes and Challenges", ], "desc": [ "A vibrant city with occasional traffic jams.", "Beautiful landscapes but overpriced tourist spots.", "Rich cultural heritage but humid summers.", "Bustling nightlife but noisy streets.", "Scenic views but limited public transport options.", ], } ), mode="overwrite", ) table.create_fts_index("desc", replace=True) results = table.search( BoostQuery( MatchQuery("beautiful, cultural, nightlife", "desc"), MatchQuery("bad traffic jams, overpriced", "desc"), ), ).to_pandas() # we will hit 3 results because the positive query has 3 hits assert len(results) == 3 # the one containing "overpriced" will be negatively boosted, # so it will be the last one assert ( results["desc"].to_list()[2] == "Beautiful landscapes but overpriced tourist spots." ) def test_fts_native(): # --8<-- [start:basic_fts] uri = "data/sample-lancedb" db = lancedb.connect(uri) table = db.create_table( "my_table_fts", data=[ {"vector": [3.1, 4.1], "text": "Frodo was a happy puppy"}, {"vector": [5.9, 26.5], "text": "There are several kittens playing"}, ], mode="overwrite", ) table.create_fts_index("text") table.search("puppy").limit(10).select(["text"]).to_list() # [{'text': 'Frodo was a happy puppy', '_score': 0.6931471824645996}] # ... # --8<-- [end:basic_fts] # --8<-- [start:fts_config_stem] table.create_fts_index("text", tokenizer_name="en_stem", replace=True) # --8<-- [end:fts_config_stem] # --8<-- [start:fts_config_folding] table.create_fts_index( "text", language="French", stem=True, ascii_folding=True, replace=True, ) # --8<-- [end:fts_config_folding] # --8<-- [start:fts_prefiltering] table.search("puppy").limit(10).where("text='foo'", prefilter=True).to_list() # --8<-- [end:fts_prefiltering] # --8<-- [start:fts_postfiltering] table.search("puppy").limit(10).where("text='foo'", prefilter=False).to_list() # --8<-- [end:fts_postfiltering] # --8<-- [start:fts_with_position] table.create_fts_index("text", with_position=True, replace=True) # --8<-- [end:fts_with_position] # --8<-- [start:fts_incremental_index] table.add([{"vector": [3.1, 4.1], "text": "Frodo was a happy puppy"}]) table.optimize() # --8<-- [end:fts_incremental_index] @pytest.mark.asyncio async def test_fts_native_async(): # --8<-- [start:basic_fts_async] uri = "data/sample-lancedb" async_db = await lancedb.connect_async(uri) async_tbl = await async_db.create_table( "my_table_fts_async", data=[ {"vector": [3.1, 4.1], "text": "Frodo was a happy puppy"}, {"vector": [5.9, 26.5], "text": "There are several kittens playing"}, ], mode="overwrite", ) # async API uses our native FTS algorithm await async_tbl.create_index("text", config=FTS()) await (await async_tbl.search("puppy")).select(["text"]).limit(10).to_list() # [{'text': 'Frodo was a happy puppy', '_score': 0.6931471824645996}] # ... # --8<-- [end:basic_fts_async] # --8<-- [start:fts_config_stem_async] await async_tbl.create_index( "text", config=FTS(language="English", stem=True, remove_stop_words=True) ) # --8<-- [end:fts_config_stem_async] # --8<-- [start:fts_config_folding_async] await async_tbl.create_index( "text", config=FTS(language="French", stem=True, ascii_folding=True) ) # --8<-- [end:fts_config_folding_async] # --8<-- [start:fts_prefiltering_async] await (await async_tbl.search("puppy")).limit(10).where("text='foo'").to_list() # --8<-- [end:fts_prefiltering_async] # --8<-- [start:fts_postfiltering_async] await ( (await async_tbl.search("puppy")) .limit(10) .where("text='foo'") .postfilter() .to_list() ) # --8<-- [end:fts_postfiltering_async] # --8<-- [start:fts_with_position_async] await async_tbl.create_index("text", config=FTS(with_position=True)) # --8<-- [end:fts_with_position_async] # --8<-- [start:fts_incremental_index_async] await async_tbl.add([{"vector": [3.1, 4.1], "text": "Frodo was a happy puppy"}]) await async_tbl.optimize() # --8<-- [end:fts_incremental_index_async] def _vectors(n, dim, seed=0, column="vector"): rng = np.random.default_rng(seed) return [ {column: rng.random(dim).astype("float32"), "id": i} for i in range(n) ] def _build_indexed_table(db, name, dim=128, n=512): tbl = db.create_table(name, _vectors(n, dim), mode="overwrite") tbl.create_index(metric="cosine", num_partitions=4, num_sub_vectors=8) return tbl def test_vs_distance_metric_and_brute_force(tmp_db): tbl = tmp_db.create_table("vs_metric", _vectors(64, 1536), mode="overwrite") # --8<-- [start:configure_distance_metric] tbl.search(np.random.random((1536))).distance_type("cosine").limit(10).to_list() # --8<-- [end:configure_distance_metric] # --8<-- [start:brute_force_search] tbl.search(np.random.random((1536))).limit(3).to_list() # --8<-- [end:brute_force_search] def test_vs_select_vector_column(tmp_db): db = tmp_db # --8<-- [start:select_vector_column] import pyarrow as pa schema = pa.schema([ pa.field("id", pa.int32()), pa.field( "image", pa.struct([pa.field("embedding", pa.list_(pa.float32(), 2))]), ), ]) table = db.create_table( "nested", data=[{"id": 0, "image": {"embedding": [0.0, 1.0]}}], schema=schema, ) # Inferred: the only vector leaf is `image.embedding`. table.search([0.0, 1.0]).limit(1).to_list() # Explicit: required when more than one vector column matches. table.search([0.0, 1.0], vector_column_name="image.embedding").limit(1).to_list() # --8<-- [end:select_vector_column] def test_vs_index_nested_column(tmp_db): dim = 16 schema = pa.schema( [ pa.field("id", pa.int32()), pa.field( "image", pa.struct([pa.field("embedding", pa.list_(pa.float32(), dim))]), ), ] ) rng = np.random.default_rng(0) data = [ {"id": i, "image": {"embedding": rng.random(dim).astype("float32").tolist()}} for i in range(512) ] table = tmp_db.create_table( "nested_index", data=data, schema=schema, mode="overwrite" ) # --8<-- [start:index_nested_column] table.create_index(vector_column_name="image.embedding") # --8<-- [end:index_nested_column] assert table.list_indices() def test_vs_indexed_queries(tmp_db): table = _build_indexed_table(tmp_db, "vs_indexed", dim=128, n=512) embedding = np.random.random(128) # --8<-- [start:exact_vs_approximate_distances] # Indexed ANN search without refinement (fast, approximate `_distance`) fast_results = ( table.search(embedding) .limit(10) .to_pandas() ) # Recompute distances on full vectors for reranked candidates exact_distance_results = ( table.search(embedding) .limit(10) .refine_factor(1) .to_pandas() ) # Rerank a larger candidate set for better recall (higher latency) higher_recall_results = ( table.search(embedding) .limit(10) .refine_factor(20) .to_pandas() ) # --8<-- [end:exact_vs_approximate_distances] # --8<-- [start:bypass_vector_index] table.search(embedding).bypass_vector_index().limit(5).to_pandas() # --8<-- [end:bypass_vector_index] assert len(fast_results) == 10 assert len(exact_distance_results) == 10 assert len(higher_recall_results) == 10 def test_vs_fast_search(tmp_db, monkeypatch): table = _build_indexed_table(tmp_db, "vs_fast", dim=128, n=512) embedding = np.random.random(128) # `fast_search` is an Enterprise/async query flag; strip it so the snippet # runs against a local OSS table without changing what readers see. real_search = table.search monkeypatch.setattr( table, "search", lambda *a, **k: real_search( *a, **{key: val for key, val in k.items() if key != "fast_search"} ), ) # --8<-- [start:fast_search] table.search(embedding, fast_search=True).limit(5).to_pandas() # --8<-- [end:fast_search] def test_vs_distance_range(tmp_db): tbl = tmp_db.create_table("vs_distance_range", _vectors(256, 256), mode="overwrite") # --8<-- [start:search_distance_range] query = np.random.random(256) # Search for the vectors within the range of [0.1, 0.5) tbl.search(query).distance_range(0.1, 0.5).to_arrow() # Search for the vectors with the distance less than 0.5 tbl.search(query).distance_range(upper_bound=0.5).to_arrow() # Search for the vectors with the distance greater or equal to 0.1 tbl.search(query).distance_range(lower_bound=0.1).to_arrow() # --8<-- [end:search_distance_range] def test_vs_multivector_search(tmp_db): schema = pa.schema( [ pa.field("id", pa.int64()), pa.field("vector", pa.list_(pa.list_(pa.float32(), 256))), ] ) rng = np.random.default_rng(0) data = [{"id": i, "vector": rng.random(size=(2, 256)).tolist()} for i in range(64)] tbl = tmp_db.create_table( "vs_multivector", data=data, schema=schema, mode="overwrite" ) # --8<-- [start:multivector_search] query_multi = np.random.random(size=(2, 256)) results_multi = tbl.search(query_multi).limit(5).to_pandas() # --8<-- [end:multivector_search] assert len(results_multi) <= 5 def test_vs_binary_search(tmp_db): db = tmp_db # --8<-- [start:search_binary_vectors] import numpy as np import pyarrow as pa schema = pa.schema( [ pa.field("id", pa.int64()), # for dim=256, lance stores every 8 bits in a byte # so the vector field should be a list of 256 / 8 = 32 bytes pa.field("vector", pa.list_(pa.uint8(), 32)), ] ) tbl = db.create_table("my_binary_vectors", schema=schema) data = [] for i in range(1024): vector = np.random.randint(0, 2, size=256) # pack the binary vector into bytes to save space packed_vector = np.packbits(vector) data.append( { "id": i, "vector": packed_vector, } ) tbl.add(data) query = np.random.randint(0, 2, size=256) packed_query = np.packbits(query) tbl.search(packed_query).distance_type("hamming").to_arrow() # --8<-- [end:search_binary_vectors] def test_vs_enterprise_filtering(tmp_db, monkeypatch): import sys import types dim = 384 rng = np.random.default_rng(0) rows = [ { "vector": rng.random(dim).astype("float32"), "text": f"story {i}", "keywords": f"kw{i}", "label": i % 4, } for i in range(50) ] tmp_db.create_table("lancedb-enterprise-quickstart", data=rows, mode="overwrite") db = tmp_db # Mock the Hugging Face dataset loader to avoid network downloads. class _FakeDataset: def __init__(self, n): self._emb = [rng.random(dim).astype("float32") for _ in range(n)] self._kw = [f"kw{i}" for i in range(n)] def __getitem__(self, key): if isinstance(key, str): return {"keywords_embeddings": self._emb, "keywords": self._kw}[key] return {"keywords": self._kw[key], "keywords_embeddings": self._emb[key]} fake_datasets = types.ModuleType("datasets") fake_datasets.load_dataset = lambda *a, **k: _FakeDataset(10) monkeypatch.setitem(sys.modules, "datasets", fake_datasets) # --8<-- [start:vector_search_prefilter] from datasets import load_dataset # Load query vector from dataset query_dataset = load_dataset("sunhaozhepy/ag_news_sbert_keywords_embeddings", split="test[5000:5001]") print(f"Query keywords: {query_dataset[0]['keywords']}") query_embed = query_dataset["keywords_embeddings"][0] # Open table and perform search table_name = "lancedb-enterprise-quickstart" table = db.open_table(table_name) # Vector search with filters (pre-filtering is the default) search_results = ( table.search(query_embed) .where("label > 2") .select(["text", "keywords", "label"]) .limit(5) .to_pandas() ) print("Search results (with pre-filtering):") print(search_results) # --8<-- [end:vector_search_prefilter] # --8<-- [start:vector_search_postfilter] results_post_filtered = ( table.search(query_embed) .where("label > 1", prefilter=False) .select(["text", "keywords", "label"]) .limit(5) .to_pandas() ) print("Vector search results with post-filter:") print(results_post_filtered) # --8<-- [end:vector_search_postfilter] # --8<-- [start:batch_search] # Load a batch of query embeddings query_dataset = load_dataset( "sunhaozhepy/ag_news_sbert_keywords_embeddings", split="test[5000:5005]" ) query_embeds = query_dataset["keywords_embeddings"] batch_results = table.search(query_embeds).limit(5).to_pandas() print(batch_results) # --8<-- [end:batch_search] @pytest.mark.skip() def test_hybrid_search(): # --8<-- [start:import-openai] import openai # --8<-- [end:import-openai] # --8<-- [start:openai-embeddings] # Ingest embedding function in LanceDB table # Configuring the environment variable OPENAI_API_KEY if "OPENAI_API_KEY" not in os.environ: # OR set the key here as a variable openai.api_key = "sk-..." embeddings = get_registry().get("openai").create() # --8<-- [end:openai-embeddings] # --8<-- [start:class-Documents] class Documents(LanceModel): vector: Vector(embeddings.ndims()) = embeddings.VectorField() text: str = embeddings.SourceField() # --8<-- [end:class-Documents] # --8<-- [start:basic_hybrid_search] data = [ {"text": "rebel spaceships striking from a hidden base"}, {"text": "have won their first victory against the evil Galactic Empire"}, {"text": "during the battle rebel spies managed to steal secret plans"}, {"text": "to the Empire's ultimate weapon the Death Star"}, ] uri = "data/sample-lancedb" db = lancedb.connect(uri) table = db.create_table("documents", schema=Documents) # ingest docs with auto-vectorization table.add(data) # Create a fts index before the hybrid search table.create_fts_index("text") # hybrid search with default re-ranker table.search("flower moon", query_type="hybrid").to_pandas() # --8<-- [end:basic_hybrid_search] # --8<-- [start:hybrid_search_pass_vector_text] vector_query = [0.1, 0.2, 0.3, 0.4, 0.5] text_query = "flower moon" ( table.search(query_type="hybrid") .vector(vector_query) .text(text_query) .limit(5) .to_pandas() ) # --8<-- [end:hybrid_search_pass_vector_text] @pytest.mark.skip async def test_hybrid_search_async(): import openai # --8<-- [start:openai-embeddings] # Ingest embedding function in LanceDB table # Configuring the environment variable OPENAI_API_KEY if "OPENAI_API_KEY" not in os.environ: # OR set the key here as a variable openai.api_key = "sk-..." embeddings = get_registry().get("openai").create() # --8<-- [end:openai-embeddings] # --8<-- [start:class-Documents] class Documents(LanceModel): vector: Vector(embeddings.ndims()) = embeddings.VectorField() text: str = embeddings.SourceField() # --8<-- [end:class-Documents] # --8<-- [start:basic_hybrid_search_async] uri = "data/sample-lancedb" async_db = await lancedb.connect_async(uri) data = [ {"text": "rebel spaceships striking from a hidden base"}, {"text": "have won their first victory against the evil Galactic Empire"}, {"text": "during the battle rebel spies managed to steal secret plans"}, {"text": "to the Empire's ultimate weapon the Death Star"}, ] async_tbl = await async_db.create_table("documents_async", schema=Documents) # ingest docs with auto-vectorization await async_tbl.add(data) # Create a fts index before the hybrid search await async_tbl.create_index("text", config=FTS()) text_query = "flower moon" # hybrid search with default re-ranker await (await async_tbl.search("flower moon", query_type="hybrid")).to_pandas() # --8<-- [end:basic_hybrid_search_async] # --8<-- [start:hybrid_search_pass_vector_text_async] vector_query = [0.1, 0.2, 0.3, 0.4, 0.5] text_query = "flower moon" await ( async_tbl.query() .nearest_to(vector_query) .nearest_to_text(text_query) .limit(5) .to_pandas() ) # --8<-- [end:hybrid_search_pass_vector_text_async]