mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-19 12:38:38 +00:00
feat(python): blob v2 fetch API (#3578)
Python bindings for blob v2 read on **local** tables. Rust read APIs landed in #3562. This PR wires `fetch_blob_files`, `fetch_blobs`, v2 query/`to_pandas(blob_mode="bytes")`, and hidden `_rowid` metadata so `fetch_*` works from query hits without exposing `_rowid` in the column list. **Cloud:** `RemoteTable.fetch_blobs` / `fetch_blob_files` raise `NotImplementedError` until Phalanx ships the server route (separate track; not blocking local merge). ### Primary path: lazy file handles ```python table = db.create_table("videos", schema=pa.schema([ pa.field("id", pa.int64()), lancedb.blob("video"), ])) table.add([{"id": 1, "video": open("clip.mp4", "rb").read()}]) hits = table.search().select(["id", "video"]).to_arrow() handle = table.fetch_blob_files("video", hits)[0] # seek + partial read — PyAV / decoders can use the handle handle.seek(frame_offset) chunk = handle.read_range(0, 65536) ``` `BlobFile` exposes `seek`, `read`, `read_range`, `read_up_to`, and works with `BufferedReader`. ### When you want full bytes ```python blobs = table.fetch_blobs("video", hits) # eager materialize, null-aligned df = table.to_pandas(blob_mode="bytes") # descriptors → bytes in pandas ``` ### `_rowid` (join key, not user `id`) Fetch needs Lance row ids. For v2 blob queries we auto-inject `_rowid`, stash it in Arrow schema metadata on `to_arrow()`, and drop the visible column unless you pass `.with_row_id(True)`. v1 legacy blobs (`lance-encoding:blob`) unchanged; fetch on v1 raises the migration error. ## Test plan - [x] `./scripts/test-blob.sh python` (105 passed in worktree) - [x] `fetch_blob_files` lazy read, seek, partial read, null alignment, cross-fragment dups - [x] hybrid query → `fetch_blobs` / `fetch_blob_files` - [ ] Will re-review after seek/`BlobFile` commit (`d77ab1a6`) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Drew Gallardo
GitHub
@@ -0,0 +1,562 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import io
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
import pytest
|
||||
|
||||
import lancedb
|
||||
from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids
|
||||
from lancedb.index import FTS
|
||||
from lancedb.schema import blob_column_paths, blob_v2_column_paths
|
||||
|
||||
|
||||
def _blob_table(name, rows):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add(rows)
|
||||
return table
|
||||
|
||||
|
||||
def _blob_array(name, values):
|
||||
blob_type = lancedb.blob(name).type
|
||||
storage_type = blob_type.storage_type
|
||||
storage = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(values, type=pa.large_binary()),
|
||||
pa.array([None] * len(values), type=pa.string()),
|
||||
pa.array([None] * len(values), type=pa.uint64()),
|
||||
pa.array([None] * len(values), type=pa.uint64()),
|
||||
],
|
||||
fields=list(storage_type),
|
||||
)
|
||||
return pa.ExtensionArray.from_storage(blob_type, storage)
|
||||
|
||||
|
||||
def _row_ids_by_id(table):
|
||||
hits = table.search().with_row_id(True).limit(1000).to_arrow()
|
||||
assert "_rowid" in hits.column_names
|
||||
return dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
|
||||
|
||||
def test_blob_factory_declares_v2_field():
|
||||
field = lancedb.blob("image")
|
||||
assert isinstance(field.type, pa.ExtensionType)
|
||||
assert field.type.extension_name == "lance.blob.v2"
|
||||
|
||||
|
||||
def test_blob_v2_column_paths_include_list_children():
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("info", pa.struct([lancedb.blob("blob")])),
|
||||
pa.field("images", pa.list_(lancedb.blob("image"))),
|
||||
pa.field("large_images", pa.large_list(lancedb.blob("large_image"))),
|
||||
pa.field(
|
||||
"fixed_images",
|
||||
pa.list_(lancedb.blob("fixed_image"), list_size=2),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
assert blob_v2_column_paths(schema) == [
|
||||
"info.blob",
|
||||
"images.image",
|
||||
"large_images.large_image",
|
||||
"fixed_images.fixed_image",
|
||||
]
|
||||
|
||||
|
||||
def _legacy_v1_table(name):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field(
|
||||
"legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
|
||||
),
|
||||
]
|
||||
)
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add([{"id": 1, "legacy": b"old"}])
|
||||
return table
|
||||
|
||||
|
||||
def test_blob_v2_column_paths_exclude_legacy_metadata():
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
lancedb.blob("image"),
|
||||
pa.field(
|
||||
"legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
|
||||
),
|
||||
]
|
||||
)
|
||||
assert blob_v2_column_paths(schema) == ["image"]
|
||||
assert blob_column_paths(schema) == ["image", "legacy"]
|
||||
|
||||
|
||||
def test_blob_v2_paths_match_blob_columns():
|
||||
table = _blob_table("paths_match", [{"id": 1, "image": b"x"}])
|
||||
assert blob_v2_column_paths(table.schema) == table.blob_columns()
|
||||
|
||||
db = lancedb.connect("memory:///")
|
||||
info = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first"], type=pa.string()),
|
||||
_blob_array("blob", [b"nested"]),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1], type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
nested = db.create_table("nested_paths", data=data)
|
||||
assert blob_v2_column_paths(nested.schema) == nested.blob_columns()
|
||||
|
||||
|
||||
def test_auto_row_id_stash_round_trip():
|
||||
table = _blob_table(
|
||||
"stash_round_trip",
|
||||
[{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}],
|
||||
)
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
row_ids = hits["_rowid"].to_pylist()
|
||||
|
||||
stashed = stash_auto_row_ids(hits, ["image"])
|
||||
|
||||
assert "_rowid" not in stashed.column_names
|
||||
assert stashed.schema.field("image").metadata == hits.schema.field("image").metadata
|
||||
assert read_row_ids_from_hits(stashed, "image") == row_ids
|
||||
|
||||
|
||||
def test_blob_query_omits_auto_row_id():
|
||||
table = _blob_table("rowid", [{"id": 1, "image": b"x"}])
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
|
||||
def test_blob_query_explicit_row_id_opt_in():
|
||||
table = _blob_table("explicit_rowid", [{"id": 1, "image": b"x"}])
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
assert "_rowid" in hits.column_names
|
||||
|
||||
|
||||
def test_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
table = _blob_table("descriptions_no_leak", [{"id": 1, "image": b"x"}])
|
||||
df = table.to_pandas(blob_mode="descriptions")
|
||||
descriptor = df["image"].iloc[0]
|
||||
assert "_lance_row_id" not in descriptor
|
||||
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
db = await lancedb.connect_async("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = await db.create_table("descriptions_no_leak_async", schema=schema)
|
||||
await table.add([{"id": 1, "image": b"x"}])
|
||||
df = await table.to_pandas(blob_mode="descriptions")
|
||||
descriptor = df["image"].iloc[0]
|
||||
assert "_lance_row_id" not in descriptor
|
||||
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
|
||||
|
||||
|
||||
def test_fetch_blobs_round_trip():
|
||||
table = _blob_table(
|
||||
"round_trip",
|
||||
[{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"]
|
||||
|
||||
|
||||
def test_fetch_blobs_accepts_query_result():
|
||||
table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}])
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"}
|
||||
|
||||
|
||||
def test_fetch_blobs_null_alignment():
|
||||
table = _blob_table(
|
||||
"nulls",
|
||||
[{"id": 1, "image": b"present"}, {"id": 2, "image": None}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
request = [by_id[1], by_id[2], by_id[1]]
|
||||
blobs = table.fetch_blobs("image", request)
|
||||
assert len(blobs) == len(request)
|
||||
assert blobs[0].as_py() == b"present"
|
||||
assert blobs[1].as_py() is None
|
||||
assert blobs[2].as_py() == b"present"
|
||||
|
||||
|
||||
def test_fetch_blobs_nested_path():
|
||||
db = lancedb.connect("memory:///")
|
||||
info = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first", "second"], type=pa.string()),
|
||||
_blob_array("blob", [b"nested-alpha", b"nested-beta"]),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1, 2], type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
table = db.create_table("nested", data=data)
|
||||
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("info.blob", [by_id[1], by_id[2]])
|
||||
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"nested-alpha", b"nested-beta"]
|
||||
|
||||
|
||||
def test_fetch_blob_files_lazy_read():
|
||||
payload = b"lazy-read" * 100
|
||||
table = _blob_table("lazy", [{"id": 1, "image": payload}])
|
||||
by_id = _row_ids_by_id(table)
|
||||
handles = table.fetch_blob_files("image", [by_id[1]])
|
||||
assert len(handles) == 1
|
||||
assert handles[0].read() == payload
|
||||
|
||||
|
||||
def test_fetch_blob_files_null_alignment():
|
||||
table = _blob_table(
|
||||
"lazy_nulls",
|
||||
[{"id": 1, "image": b"here"}, {"id": 2, "image": None}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
handles = table.fetch_blob_files("image", [by_id[2], by_id[1]])
|
||||
assert len(handles) == 2
|
||||
assert handles[0] is None
|
||||
assert handles[1].read() == b"here"
|
||||
|
||||
|
||||
def test_fetch_blobs_rejects_non_blob_column():
|
||||
table = _blob_table("reject", [{"id": 1, "image": b"x"}])
|
||||
with pytest.raises(ValueError, match="not a blob column"):
|
||||
table.fetch_blobs("id", [0])
|
||||
|
||||
|
||||
def test_legacy_v1_query_omits_auto_row_id():
|
||||
table = _legacy_v1_table("legacy_v1")
|
||||
hits = table.search().select(["legacy"]).limit(10).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
|
||||
def test_fetch_blobs_rejects_legacy_v1_column():
|
||||
table = _legacy_v1_table("legacy_fetch")
|
||||
with pytest.raises(ValueError, match="legacy blob column.*blob v2"):
|
||||
table.fetch_blobs("legacy", [0])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_fetch_blob_files_lazy_read():
|
||||
db = await lancedb.connect_async("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = await db.create_table("async_lazy", schema=schema)
|
||||
payload = b"async-lazy" * 100
|
||||
await table.add([{"id": 1, "image": payload}])
|
||||
hits = (
|
||||
await table.query().select({"image_alias": "image"}).limit(10).to_arrow()
|
||||
).combine_chunks()
|
||||
assert "_rowid" not in hits.column_names
|
||||
handles = await table.fetch_blob_files("image", hits)
|
||||
assert len(handles) == 1
|
||||
assert await handles[0].aread() == payload
|
||||
|
||||
|
||||
def test_fetch_blobs_from_query_result_without_row_id_raises():
|
||||
table = _blob_table("no_rowid", [{"id": 1, "image": b"x"}])
|
||||
hits = table.search().select(["id"]).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
with pytest.raises(ValueError, match="_rowid"):
|
||||
table.fetch_blobs("image", hits)
|
||||
|
||||
|
||||
_HYBRID_BLOB_SCHEMA = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("text", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("image"),
|
||||
]
|
||||
)
|
||||
_HYBRID_BLOB_ROWS = [
|
||||
{"id": 1, "text": "hello alpha", "vector": [1.0, 0.0], "image": b"alpha"},
|
||||
{"id": 2, "text": "hello beta", "vector": [0.9, 0.1], "image": b"beta"},
|
||||
{"id": 3, "text": "other", "vector": [0.0, 1.0], "image": b"other"},
|
||||
]
|
||||
|
||||
|
||||
def _hybrid_blob_table(db):
|
||||
table = db.create_table("hybrid_blob_fetch", schema=_HYBRID_BLOB_SCHEMA)
|
||||
table.add(_HYBRID_BLOB_ROWS)
|
||||
table.create_index("text", config=FTS(with_position=False))
|
||||
return table
|
||||
|
||||
|
||||
async def _hybrid_blob_table_async(db):
|
||||
table = await db.create_table("hybrid_blob_fetch_async", schema=_HYBRID_BLOB_SCHEMA)
|
||||
await table.add(_HYBRID_BLOB_ROWS)
|
||||
await table.create_index("text", config=FTS(with_position=False))
|
||||
return table
|
||||
|
||||
|
||||
def test_blob_v2_hybrid_fetch_blobs():
|
||||
table = _hybrid_blob_table(lancedb.connect("memory:///"))
|
||||
hits = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector([1.0, 0.0])
|
||||
.text("hello")
|
||||
.select(["id", "image"])
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert "_rowid" not in hits.column_names
|
||||
assert "_lance_row_id" in hits.schema.field("image").type.names
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blob_v2_hybrid_fetch_blobs_async():
|
||||
db = await lancedb.connect_async("memory:///hybrid_blob_fetch_async")
|
||||
table = await _hybrid_blob_table_async(db)
|
||||
hits = await (
|
||||
table.query()
|
||||
.nearest_to([1.0, 0.0])
|
||||
.nearest_to_text("hello")
|
||||
.select(["id", "image"])
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert "_rowid" not in hits.column_names
|
||||
assert "_lance_row_id" in hits.schema.field("image").type.names
|
||||
blobs = await table.fetch_blobs("image", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
def test_blob_file_seek_read_and_read_range():
|
||||
payload = _identifiable_payload(1024)
|
||||
table = _blob_table("seek_read", [{"id": 1, "image": payload}])
|
||||
by_id = _row_ids_by_id(table)
|
||||
handle = table.fetch_blob_files("image", [by_id[1]])[0]
|
||||
|
||||
assert handle.seek(100) == 100
|
||||
assert handle.read(16) == payload[100:116]
|
||||
handle.seek(100)
|
||||
assert handle.read_range(500, 8) == payload[500:508]
|
||||
assert handle.tell() == 100
|
||||
|
||||
with pytest.raises(ValueError, match="whence"):
|
||||
handle.seek(0, 99)
|
||||
|
||||
|
||||
def test_fetch_blob_files_from_query_partial_read():
|
||||
payload = _identifiable_payload(65536)
|
||||
table = _blob_table("query_partial", [{"id": 1, "image": payload}])
|
||||
hits = table.search().select(["id", "image"]).limit(1).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
assert handle.size() == 65536
|
||||
assert handle.read_range(0, 128) == payload[:128]
|
||||
assert handle.tell() == 0
|
||||
assert handle.seek(40000) == 40000
|
||||
assert handle.read(16) == payload[40000:40016]
|
||||
|
||||
|
||||
def test_blob_file_buffered_reader():
|
||||
payload = _identifiable_payload(4096)
|
||||
table = _blob_table("buffered_reader", [{"id": 1, "image": payload}])
|
||||
hits = table.search().select(["id", "image"]).limit(1).to_arrow()
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
reader = io.BufferedReader(handle)
|
||||
assert reader.read(8) == payload[:8]
|
||||
assert reader.read(8) == payload[8:16]
|
||||
assert reader.read() == payload[16:]
|
||||
|
||||
|
||||
def test_fetch_blob_files_cross_fragment_nulls_and_dups():
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("cross_fragment", schema=schema)
|
||||
table.add([{"id": 1, "image": b"alpha"}])
|
||||
table.add([{"id": 2, "image": None}, {"id": 3, "image": b"beta"}])
|
||||
|
||||
by_id = _row_ids_by_id(table)
|
||||
request = [by_id[3], by_id[2], by_id[1], by_id[3]]
|
||||
handles = table.fetch_blob_files("image", request)
|
||||
assert len(handles) == 4
|
||||
assert handles[1] is None
|
||||
assert handles[0].read() == b"beta"
|
||||
assert handles[2].read() == b"alpha"
|
||||
assert handles[3].seek(1) == 1
|
||||
assert handles[3].read() == b"eta"
|
||||
|
||||
|
||||
def test_blob_file_pyav_decode_seek(tmp_path):
|
||||
av = pytest.importorskip("av")
|
||||
import fractions
|
||||
|
||||
clip = tmp_path / "clip.mp4"
|
||||
with av.open(str(clip), mode="w") as container:
|
||||
stream = container.add_stream("mpeg4", rate=5)
|
||||
stream.width, stream.height, stream.pix_fmt = 32, 32, "yuv420p"
|
||||
stream.time_base = fractions.Fraction(1, 5)
|
||||
for pts in range(5):
|
||||
frame = av.VideoFrame(32, 32, "yuv420p")
|
||||
frame.pts = pts
|
||||
container.mux(stream.encode(frame))
|
||||
container.mux(stream.encode(None))
|
||||
|
||||
table = _blob_table("pyav", [{"id": 1, "image": clip.read_bytes()}])
|
||||
hits = table.search().select(["image"]).limit(1).to_arrow()
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
|
||||
with av.open(handle) as container:
|
||||
stream = container.streams.video[0]
|
||||
container.seek(0)
|
||||
assert next(container.decode(stream)) is not None
|
||||
|
||||
|
||||
def test_blob_v2_hybrid_fetch_blob_files_seek():
|
||||
table = _hybrid_blob_table(lancedb.connect("memory:///"))
|
||||
hits = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector([1.0, 0.0])
|
||||
.text("hello")
|
||||
.select(["id", "image"])
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
handles = table.fetch_blob_files("image", hits)
|
||||
assert len(handles) == 2
|
||||
assert {handle.read_range(0, 2) for handle in handles} == {b"al", b"be"}
|
||||
first = handles[0]
|
||||
assert first.seek(1) == 1
|
||||
assert first.read(2) in {b"lp", b"et"}
|
||||
|
||||
|
||||
def test_blob_file_header_sniff_from_search():
|
||||
payload = b"%PDF-1.7\n" + bytes(4096)
|
||||
table = _blob_table("header_sniff", [{"id": 1, "image": payload}])
|
||||
hits = table.search().select(["id", "image"]).limit(1).to_arrow()
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
assert handle.read_range(0, 4) == b"%PDF"
|
||||
assert handle.tell() == 0
|
||||
|
||||
|
||||
def test_blob_file_multiple_handles_independent_cursors():
|
||||
table = _blob_table(
|
||||
"multi_handle",
|
||||
[{"id": 1, "image": b"first-payload"}, {"id": 2, "image": b"second-payload"}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
first, second = table.fetch_blob_files("image", [by_id[1], by_id[2]])
|
||||
assert first.seek(6) == 6
|
||||
assert second.tell() == 0
|
||||
assert first.read(7) == b"payload"
|
||||
assert second.read(6) == b"second"
|
||||
|
||||
|
||||
def test_fetch_blob_files_nested_path_seek():
|
||||
db = lancedb.connect("memory:///")
|
||||
info = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first", "second"], type=pa.string()),
|
||||
_blob_array("blob", [b"nested-alpha", b"nested-beta"]),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1, 2], type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
table = db.create_table("nested_seek", data=data)
|
||||
by_id = _row_ids_by_id(table)
|
||||
handle = table.fetch_blob_files("info.blob", [by_id[2]])[0]
|
||||
assert handle.seek(7) == 7
|
||||
assert handle.read() == b"beta"
|
||||
|
||||
|
||||
def test_fetch_blobs_survives_sort_after_query():
|
||||
table = _blob_table(
|
||||
"sort_survives",
|
||||
[{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
|
||||
)
|
||||
hits = table.search().select(["id", "image"]).to_arrow()
|
||||
sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
|
||||
sorted_hits = hits.take(sort_idx)
|
||||
|
||||
blobs = table.fetch_blobs("image", sorted_hits)
|
||||
expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == expected
|
||||
|
||||
|
||||
def test_fetch_blobs_survives_filter_and_sort_after_query():
|
||||
table = _blob_table(
|
||||
"filter_sort_survives",
|
||||
[{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
|
||||
)
|
||||
hits = table.search().select(["id", "image"]).to_arrow()
|
||||
filtered = hits.filter(pc.field("id") >= 2)
|
||||
sort_idx = pc.sort_indices(filtered["id"], sort_keys=[("id", "descending")])
|
||||
filtered_sorted = filtered.take(sort_idx)
|
||||
|
||||
blobs = table.fetch_blobs("image", filtered_sorted)
|
||||
expected = [f"payload-{i}".encode() for i in filtered_sorted["id"].to_pylist()]
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == expected
|
||||
|
||||
|
||||
def test_fetch_blob_files_survives_sort_after_query():
|
||||
table = _blob_table(
|
||||
"lazy_sort_survives",
|
||||
[{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
|
||||
)
|
||||
hits = table.search().select(["id", "image"]).to_arrow()
|
||||
sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
|
||||
sorted_hits = hits.take(sort_idx)
|
||||
|
||||
handles = table.fetch_blob_files("image", sorted_hits)
|
||||
expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
|
||||
assert [handle.read() for handle in handles] == expected
|
||||
|
||||
|
||||
def test_fetch_blobs_nested_path_survives_sort_after_query():
|
||||
db = lancedb.connect("memory:///")
|
||||
values = [f"payload-{i}".encode() for i in range(4)]
|
||||
info = pa.StructArray.from_arrays(
|
||||
[pa.array(["row"] * 4, type=pa.string()), _blob_array("blob", values)],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array(range(4), type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
table = db.create_table("nested_sort_survives", data=data)
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
|
||||
sorted_hits = hits.take(sort_idx)
|
||||
|
||||
blobs = table.fetch_blobs("info.blob", sorted_hits)
|
||||
expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == expected
|
||||
|
||||
|
||||
def _identifiable_payload(size: int) -> bytes:
|
||||
block = 256
|
||||
return b"".join(bytes([i % 256]) * block for i in range(size // block))
|
||||
@@ -11,6 +11,7 @@ import lancedb
|
||||
from lancedb.db import AsyncConnection
|
||||
from lancedb.embeddings.base import TextEmbeddingFunction
|
||||
from lancedb.embeddings.registry import get_registry, register
|
||||
from lancedb.expr import col
|
||||
from lancedb.index import FTS, IvfPq
|
||||
import lancedb.pydantic
|
||||
import numpy as np
|
||||
@@ -63,11 +64,71 @@ def _blob_query_data():
|
||||
)
|
||||
|
||||
|
||||
def _create_blob_v2_query_table(db, name):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("tag", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add(
|
||||
[
|
||||
{"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"},
|
||||
{"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"},
|
||||
{"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"},
|
||||
{"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"},
|
||||
]
|
||||
)
|
||||
return table
|
||||
|
||||
|
||||
async def _create_blob_v2_query_table_async(db, name):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("tag", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = await db.create_table(name, schema=schema)
|
||||
await table.add(
|
||||
[
|
||||
{"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"},
|
||||
{"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"},
|
||||
{"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"},
|
||||
{"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"},
|
||||
]
|
||||
)
|
||||
return table
|
||||
|
||||
|
||||
def _assert_lazy_blob(value, expected: bytes):
|
||||
assert hasattr(value, "readall")
|
||||
assert value.readall() == expected
|
||||
|
||||
|
||||
def _assert_blob_bytes_projection(df):
|
||||
assert df["id_alias"].tolist() == [3, 4]
|
||||
assert df["payload"].tolist() == [b"three", b"four"]
|
||||
assert df["double_id"].tolist() == [6, 8]
|
||||
|
||||
|
||||
def _blob_query_table(db, name, blob_schema):
|
||||
if blob_schema == "v1":
|
||||
return db.create_table(name, _blob_query_data())
|
||||
return _create_blob_v2_query_table(db, name)
|
||||
|
||||
|
||||
async def _blob_query_table_async(db, name, blob_schema):
|
||||
if blob_schema == "v1":
|
||||
return await db.create_table(name, _blob_query_data())
|
||||
return await _create_blob_v2_query_table_async(db, name)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def table(tmpdir_factory) -> lancedb.table.Table:
|
||||
tmp_path = str(tmpdir_factory.mktemp("data"))
|
||||
@@ -235,10 +296,11 @@ def test_plain_scan_query_to_pandas_blob_modes(tmp_db, blob_mode):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
def test_plain_scan_query_to_pandas_blob_projection(tmp_db):
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
def test_plain_scan_query_to_pandas_blob_bytes_projection(tmp_db, blob_schema):
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
"test_query_to_pandas_blob_projection", _blob_query_data()
|
||||
table = _blob_query_table(
|
||||
tmp_db, f"test_query_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
||||
)
|
||||
|
||||
df = (
|
||||
@@ -250,9 +312,8 @@ def test_plain_scan_query_to_pandas_blob_projection(tmp_db):
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
assert df["id_alias"].tolist() == [3, 4]
|
||||
assert df["payload"].tolist() == [b"three", b"four"]
|
||||
assert df["double_id"].tolist() == [6, 8]
|
||||
_assert_blob_bytes_projection(df)
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
@@ -348,18 +409,6 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async):
|
||||
assert lazy_df["id"].tolist() == [1]
|
||||
_assert_lazy_blob(lazy_df["blob"].iloc[0], b"one")
|
||||
|
||||
bytes_df = await (
|
||||
table.query()
|
||||
.where("id >= 2")
|
||||
.select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"})
|
||||
.limit(2)
|
||||
.offset(1)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
assert bytes_df["id_alias"].tolist() == [3, 4]
|
||||
assert bytes_df["payload"].tolist() == [b"three", b"four"]
|
||||
assert bytes_df["double_id"].tolist() == [6, 8]
|
||||
|
||||
desc_df = await (
|
||||
table.query()
|
||||
.where("id = 1")
|
||||
@@ -371,6 +420,31 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
async def test_async_plain_scan_query_to_pandas_blob_bytes_projection(
|
||||
tmp_db_async, blob_schema
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
table = await _blob_query_table_async(
|
||||
tmp_db_async,
|
||||
f"test_async_query_to_pandas_blob_{blob_schema}_bytes",
|
||||
blob_schema,
|
||||
)
|
||||
|
||||
df = await (
|
||||
table.query()
|
||||
.where("id >= 2")
|
||||
.select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"})
|
||||
.limit(2)
|
||||
.offset(1)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
_assert_blob_bytes_projection(df)
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
async def test_async_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow(
|
||||
@@ -502,6 +576,18 @@ def test_with_row_id(table: lancedb.table.Table):
|
||||
assert rs["_rowid"].to_pylist() == [0, 1]
|
||||
|
||||
|
||||
def test_blob_v2_query_omits_auto_row_id(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid")
|
||||
|
||||
query_obj = table.search().select(["id", "blob"]).limit(2).to_query_object()
|
||||
assert query_obj.with_row_id is None
|
||||
|
||||
rs = table.search().select(["id", "blob"]).limit(2).to_arrow()
|
||||
|
||||
assert "_rowid" not in rs.column_names
|
||||
assert rs["id"].to_pylist() == [1, 2]
|
||||
|
||||
|
||||
def test_where_repeated_combines_with_and(table: lancedb.table.Table):
|
||||
# Calling where() more than once should AND the filters together instead of
|
||||
# silently replacing the previous one (regression test for #2649).
|
||||
@@ -1946,3 +2032,39 @@ def test_fast_search(tmp_path):
|
||||
# 2. Fast Search -> Should NOT include "LanceScan" (Uses Index)
|
||||
plan = table.search(q).fast_search().explain_plan(True)
|
||||
assert "LanceScan" not in plan
|
||||
|
||||
|
||||
def test_blob_v2_with_row_id_bytes_pandas(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_rowid_bytes_pandas")
|
||||
|
||||
df = (
|
||||
table.search()
|
||||
.with_row_id(True)
|
||||
.select(["id", "blob"])
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
assert "_rowid" in df.columns
|
||||
assert df["id"].tolist() == [1, 2, 3, 4]
|
||||
assert df["blob"].tolist() == [b"one", b"two", b"three", b"four"]
|
||||
|
||||
|
||||
def test_blob_v2_expr_projection_stash(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_expr_projection_stash")
|
||||
|
||||
hits = table.search().select({"blob_alias": col("blob")}).limit(2).to_arrow()
|
||||
|
||||
assert "_rowid" not in hits.column_names
|
||||
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
|
||||
blobs = table.fetch_blobs("blob", hits)
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
|
||||
|
||||
|
||||
def test_blob_v2_to_batches_row_id(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_to_batches_rowid")
|
||||
|
||||
hits = table.search().select(["id", "blob"]).limit(2).to_batches().read_all()
|
||||
|
||||
assert "_rowid" in hits.column_names
|
||||
blobs = table.fetch_blobs("blob", hits)
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
|
||||
|
||||
@@ -45,6 +45,32 @@ def _blob_test_data():
|
||||
)
|
||||
|
||||
|
||||
def _blob_v2_table(db: DBConnection, name: str):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}])
|
||||
return table
|
||||
|
||||
|
||||
async def _blob_v2_table_async(db: AsyncConnection, name: str):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = await db.create_table(name, schema=schema)
|
||||
await table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}])
|
||||
return table
|
||||
|
||||
|
||||
def _blob_table(db: DBConnection, name: str, blob_schema: str):
|
||||
if blob_schema == "v1":
|
||||
return db.create_table(name, data=_blob_test_data())
|
||||
return _blob_v2_table(db, name)
|
||||
|
||||
|
||||
async def _blob_table_async(db: AsyncConnection, name: str, blob_schema: str):
|
||||
if blob_schema == "v1":
|
||||
return await db.create_table(name, data=_blob_test_data())
|
||||
return await _blob_v2_table_async(db, name)
|
||||
|
||||
|
||||
def _assert_lazy_blob(value, expected: bytes):
|
||||
assert hasattr(value, "readall")
|
||||
assert value.readall() == expected
|
||||
@@ -107,6 +133,18 @@ def test_table_to_pandas_blob_modes(tmp_db: DBConnection, blob_mode):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
def test_table_to_pandas_blob_bytes(tmp_db: DBConnection, blob_schema):
|
||||
pytest.importorskip("lance")
|
||||
table = _blob_table(tmp_db, f"test_to_pandas_blob_{blob_schema}_bytes", blob_schema)
|
||||
|
||||
df = table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert list(df.columns) == ["id", "blob"]
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
pd = pytest.importorskip("pandas")
|
||||
data = pa.table({"id": pa.array([1, 2], pa.int64())})
|
||||
@@ -118,15 +156,20 @@ def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_table_to_pandas_blob_bytes(tmp_db_async: AsyncConnection):
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
async def test_async_table_to_pandas_blob_bytes(
|
||||
tmp_db_async: AsyncConnection, blob_schema
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
table = await tmp_db_async.create_table(
|
||||
"test_async_to_pandas_blob_bytes", data=_blob_test_data()
|
||||
table = await _blob_table_async(
|
||||
tmp_db_async, f"test_async_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
||||
)
|
||||
|
||||
df = await table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert list(df.columns) == ["id", "blob"]
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
Reference in New Issue
Block a user