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.
233 lines
6.6 KiB
Python
233 lines
6.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import pytest
|
|
try:
|
|
import lancedb
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import io
|
|
from PIL import Image
|
|
except ImportError:
|
|
pass
|
|
|
|
# --8<-- [start:multimodal_imports]
|
|
import lancedb
|
|
import pyarrow as pa
|
|
import pandas as pd
|
|
import numpy as np
|
|
import io
|
|
from PIL import Image
|
|
# --8<-- [end:multimodal_imports]
|
|
|
|
def test_multimodal_ingestion(db_path_factory):
|
|
# Ensure dependencies are available
|
|
pytest.importorskip("PIL")
|
|
pytest.importorskip("lancedb")
|
|
pytest.importorskip("numpy")
|
|
|
|
# --8<-- [start:create_dummy_data]
|
|
# Create some dummy images
|
|
def create_dummy_image(color):
|
|
img = Image.new('RGB', (100, 100), color=color)
|
|
buf = io.BytesIO()
|
|
img.save(buf, format='PNG')
|
|
return buf.getvalue()
|
|
|
|
# Create dataset with metadata, vectors, and image blobs
|
|
data = [
|
|
{
|
|
"id": 1,
|
|
"filename": "red_square.png",
|
|
"vector": np.random.rand(128).astype(np.float32),
|
|
"image_blob": create_dummy_image('red'),
|
|
"label": "red"
|
|
},
|
|
{
|
|
"id": 2,
|
|
"filename": "blue_square.png",
|
|
"vector": np.random.rand(128).astype(np.float32),
|
|
"image_blob": create_dummy_image('blue'),
|
|
"label": "blue"
|
|
}
|
|
]
|
|
# --8<-- [end:create_dummy_data]
|
|
|
|
# --8<-- [start:define_schema]
|
|
# Define schema explictly to ensure image_blob is treated as binary
|
|
schema = pa.schema([
|
|
pa.field("id", pa.int32()),
|
|
pa.field("filename", pa.string()),
|
|
pa.field("vector", pa.list_(pa.float32(), 128)),
|
|
pa.field("image_blob", pa.binary()), # Important: Use pa.binary() for blobs
|
|
pa.field("label", pa.string())
|
|
])
|
|
# --8<-- [end:define_schema]
|
|
|
|
db_uri = db_path_factory("multimodal_db")
|
|
db = lancedb.connect(db_uri)
|
|
|
|
# --8<-- [start:ingest_data]
|
|
tbl = db.create_table("images", data=data, schema=schema, mode="overwrite")
|
|
# --8<-- [end:ingest_data]
|
|
|
|
assert len(tbl) == 2
|
|
|
|
# --8<-- [start:search_data]
|
|
# Search for similar images
|
|
query_vector = np.random.rand(128).astype(np.float32)
|
|
results = tbl.search(query_vector).limit(1).to_pandas()
|
|
# --8<-- [end:search_data]
|
|
|
|
# --8<-- [start:process_results]
|
|
# Convert back to PIL Image
|
|
for _, row in results.iterrows():
|
|
image_bytes = row['image_blob']
|
|
image = Image.open(io.BytesIO(image_bytes))
|
|
print(f"Retrieved image: {row['filename']}, Size: {image.size}")
|
|
# You can now use 'image' with other libraries or display it
|
|
# --8<-- [end:process_results]
|
|
|
|
assert len(results) == 1
|
|
|
|
def test_blob_api_definition(db_path_factory):
|
|
# --8<-- [start:blob_api_schema]
|
|
import pyarrow as pa
|
|
|
|
# Define schema with Blob API metadata for lazy loading
|
|
schema = pa.schema([
|
|
pa.field("id", pa.int64()),
|
|
pa.field(
|
|
"video",
|
|
pa.large_binary(),
|
|
metadata={"lance-encoding:blob": "true"} # Enable Blob API
|
|
),
|
|
])
|
|
# --8<-- [end:blob_api_schema]
|
|
|
|
# --8<-- [start:blob_api_ingest]
|
|
import lancedb
|
|
|
|
db = lancedb.connect(db_path_factory("blob_db"))
|
|
|
|
# Create sample data
|
|
data = [
|
|
{"id": 1, "video": b"fake_video_bytes_1"},
|
|
{"id": 2, "video": b"fake_video_bytes_2"}
|
|
]
|
|
|
|
# Create the table
|
|
tbl = db.create_table("videos", data=data, schema=schema)
|
|
# --8<-- [end:blob_api_ingest]
|
|
assert len(tbl) == 2
|
|
|
|
|
|
def test_blob_api_to_pandas(db_path_factory):
|
|
db = lancedb.connect(db_path_factory("blob_to_pandas_db"))
|
|
schema = pa.schema([
|
|
pa.field("id", pa.int64()),
|
|
pa.field(
|
|
"video",
|
|
pa.large_binary(),
|
|
metadata={"lance-encoding:blob": "true"},
|
|
),
|
|
])
|
|
tbl = db.create_table(
|
|
"videos",
|
|
data=[
|
|
{"id": 1, "video": b"fake_video_bytes_1"},
|
|
{"id": 2, "video": b"fake_video_bytes_2"},
|
|
],
|
|
schema=schema,
|
|
mode="overwrite",
|
|
)
|
|
|
|
# --8<-- [start:blob_api_to_pandas]
|
|
# Default: blob columns come back lazily
|
|
df_lazy = tbl.to_pandas()
|
|
|
|
# Materialize blob bytes eagerly
|
|
df_bytes = tbl.to_pandas(blob_mode="bytes")
|
|
|
|
# Return descriptors instead of payloads
|
|
df_desc = tbl.to_pandas(blob_mode="descriptions")
|
|
|
|
# Forward extra kwargs to PyArrow's to_pandas
|
|
df_typed = tbl.to_pandas(split_blocks=True, self_destruct=True)
|
|
# --8<-- [end:blob_api_to_pandas]
|
|
|
|
assert len(df_lazy) == 2
|
|
assert isinstance(df_bytes["video"].iloc[0], bytes)
|
|
assert df_bytes["video"].tolist() == [
|
|
b"fake_video_bytes_1",
|
|
b"fake_video_bytes_2",
|
|
]
|
|
assert len(df_desc) == 2
|
|
assert len(df_typed) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_to_pandas_kwargs(db_path_factory):
|
|
schema = pa.schema([
|
|
pa.field("id", pa.int64()),
|
|
pa.field("vector", pa.list_(pa.float32(), 128)),
|
|
pa.field(
|
|
"video",
|
|
pa.large_binary(),
|
|
metadata={"lance-encoding:blob": "true"},
|
|
),
|
|
])
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"vector": np.random.rand(128).astype(np.float32),
|
|
"video": f"fake_video_bytes_{i}".encode(),
|
|
}
|
|
for i in range(10)
|
|
]
|
|
|
|
db = lancedb.connect(db_path_factory("query_to_pandas_db"))
|
|
tbl = db.create_table("search_demo", data=data, schema=schema, mode="overwrite")
|
|
|
|
async_db = await lancedb.connect_async(
|
|
str(db_path_factory("query_to_pandas_async_db"))
|
|
)
|
|
tbl_async = await async_db.create_table(
|
|
"search_demo", data=data, schema=schema, mode="overwrite"
|
|
)
|
|
|
|
query_vector = np.random.rand(128).astype(np.float32)
|
|
|
|
# --8<-- [start:query_to_pandas_kwargs]
|
|
# Plain scan query: blob_mode is supported end to end
|
|
df_lazy = (
|
|
tbl.search()
|
|
.where("id = 1")
|
|
.select(["id", "video"])
|
|
.to_pandas(blob_mode="lazy")
|
|
)
|
|
|
|
# Same call shape works on async query builders
|
|
df_bytes = await (
|
|
tbl_async.query()
|
|
.where("id = 1")
|
|
.select(["id", "video"])
|
|
.to_pandas(blob_mode="bytes")
|
|
)
|
|
|
|
# Vector / FTS / hybrid queries can't materialize blob columns,
|
|
# so omit them from the projection
|
|
df_vec = (
|
|
tbl.search(query_vector)
|
|
.limit(10)
|
|
.select(["id", "vector"])
|
|
.to_pandas(split_blocks=True, self_destruct=True)
|
|
)
|
|
# --8<-- [end:query_to_pandas_kwargs]
|
|
|
|
assert len(df_lazy) == 1
|
|
assert len(df_bytes) == 1
|
|
assert df_bytes["video"].iloc[0] == b"fake_video_bytes_1"
|
|
assert len(df_vec) == 10
|
|
assert "video" not in df_vec.columns
|