mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
test(python): expand minimal dependency coverage
Expand the Linux Python test run without pandas and pylance from two files to the suite while retaining the regular dependency run for optional-package modules. Move optional imports into tests that need them and use Arrow assertions where pandas is not required, preserving the existing behavioral assertions. Validation: - Physical no-pandas/no-pylance broad run: 2884 passed, 316 skipped, 8 existing marker cases deselected. - Regular affected suite: 1687 passed, 39 skipped. - Current-head hosted integration: all 15 checks and 4 workflows passed. Co-authored-by: importcpp <57519936+importcpp@users.noreply.github.com>
This commit is contained in:
@@ -7,12 +7,9 @@ import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
|
||||
import lance
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
import pytest
|
||||
from lance.blob import Blob, BlobType as LanceBlobType
|
||||
|
||||
import lancedb
|
||||
from lancedb import schema as schema_module
|
||||
@@ -114,6 +111,7 @@ def _assert_fetch_apis_reject_missing_row_ids(table, row_ids):
|
||||
|
||||
|
||||
def test_blob_factory_declares_v2_field():
|
||||
LanceBlobType = pytest.importorskip("lance.blob").BlobType
|
||||
field = lancedb.blob("image")
|
||||
assert isinstance(field.type, pa.ExtensionType)
|
||||
assert field.type.extension_name == "lance.blob.v2"
|
||||
@@ -163,6 +161,8 @@ def test_blob_without_thresholds_has_no_threshold_metadata():
|
||||
],
|
||||
)
|
||||
def test_blob_thresholds_match_pylance(monkeypatch, nullable, thresholds):
|
||||
lance = pytest.importorskip("lance")
|
||||
LanceBlobType = pytest.importorskip("lance.blob").BlobType
|
||||
expected = lance.blob_field("payload", nullable=nullable, **thresholds)
|
||||
actual = lancedb.blob("payload", nullable=nullable, **thresholds)
|
||||
assert type(actual.type) is LanceBlobType
|
||||
@@ -304,6 +304,7 @@ def test_local_schema_restores_blob_extensions(tmp_path, use_async, with_pylance
|
||||
|
||||
|
||||
def test_blob_resolves_pylance_type_without_eager_import():
|
||||
pytest.importorskip("lance.blob")
|
||||
script = textwrap.dedent(
|
||||
"""\
|
||||
import sys
|
||||
@@ -375,6 +376,7 @@ def test_blob_fallback_fails_if_name_already_registered():
|
||||
|
||||
|
||||
def test_blob_type_rejects_competing_registration_with_pylance():
|
||||
pytest.importorskip("lance.blob")
|
||||
script = textwrap.dedent(
|
||||
"""\
|
||||
import pyarrow as pa
|
||||
@@ -515,6 +517,7 @@ def _nested_blob_table(db, name):
|
||||
|
||||
|
||||
def test_nested_blob_to_pandas_bytes(tmp_path):
|
||||
pytest.importorskip("pandas")
|
||||
table = _nested_blob_table(lancedb.connect(tmp_path), "nested_pandas_bytes")
|
||||
expected = [
|
||||
{"name": "first", "blob": b"hello"},
|
||||
@@ -538,6 +541,9 @@ def test_nested_blob_to_pandas_bytes(tmp_path):
|
||||
assert table.search().select({"b": "info.blob"}).to_pandas(blob_mode="bytes")[
|
||||
"b"
|
||||
].tolist() == [b"hello", None, None]
|
||||
assert table.search().select([]).to_pandas(blob_mode="bytes")["info"].tolist() == (
|
||||
table.search().select([]).to_arrow(blob_mode="bytes")["info"].to_pylist()
|
||||
)
|
||||
|
||||
|
||||
def test_to_arrow_and_to_list_blob_mode(tmp_path):
|
||||
@@ -573,9 +579,6 @@ def test_to_arrow_and_to_list_blob_mode(tmp_path):
|
||||
assert [
|
||||
row["info"] for row in table.search().select([]).to_list(blob_mode="bytes")
|
||||
] == expected
|
||||
assert table.search().select([]).to_pandas(blob_mode="bytes")["info"].tolist() == (
|
||||
table.search().select([]).to_arrow(blob_mode="bytes")["info"].to_pylist()
|
||||
)
|
||||
assert "_rowid" not in table.search().select([]).to_arrow().column_names
|
||||
# A flat non-blob column named like the nested blob path keeps its values.
|
||||
shadowed = table.search().select({"info": "info", "info.blob": "id"})
|
||||
@@ -694,6 +697,7 @@ def test_vector_query_to_arrow_blob_mode_bytes(tmp_path):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nested_blob_to_pandas_bytes_async(tmp_path):
|
||||
pytest.importorskip("pandas")
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
schema = pa.schema(
|
||||
[
|
||||
@@ -809,6 +813,7 @@ def test_blob_query_explicit_row_id_opt_in():
|
||||
|
||||
|
||||
def test_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
pytest.importorskip("pandas")
|
||||
table = _blob_table("descriptions_no_leak", [{"id": 1, "image": b"x"}])
|
||||
df = table.to_pandas(blob_mode="descriptions")
|
||||
descriptor = df["image"].iloc[0]
|
||||
@@ -818,6 +823,7 @@ def test_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
pytest.importorskip("pandas")
|
||||
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)
|
||||
@@ -983,6 +989,7 @@ def test_to_batches_validates_blob_mode():
|
||||
|
||||
|
||||
def test_to_pandas_keeps_to_arrow_overrides_without_blob_mode():
|
||||
pytest.importorskip("pandas")
|
||||
from lancedb.query import LanceEmptyQueryBuilder
|
||||
|
||||
class ToArrowWithoutBlobMode(LanceEmptyQueryBuilder):
|
||||
@@ -1000,6 +1007,7 @@ def test_to_pandas_keeps_to_arrow_overrides_without_blob_mode():
|
||||
|
||||
|
||||
def test_take_queries_honor_blob_mode():
|
||||
pytest.importorskip("pandas")
|
||||
table = _blob_table(
|
||||
"take_blob_mode", [{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}]
|
||||
)
|
||||
@@ -1449,6 +1457,7 @@ def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path):
|
||||
|
||||
|
||||
def test_merge_insert_blob_array_into_reopened_table(tmp_path):
|
||||
pytest.importorskip("lance")
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("images", schema=schema)
|
||||
@@ -1624,6 +1633,8 @@ def test_row_addressable_paths_exclude_list_children():
|
||||
|
||||
|
||||
def test_merge_insert_writes_pylance_blob_array():
|
||||
lance = pytest.importorskip("lance")
|
||||
LanceBlobType = pytest.importorskip("lance.blob").BlobType
|
||||
table = _blob_table("merge_pylance", [{"id": 1, "image": b"before"}])
|
||||
image = lance.blob_array([b"updated", b"inserted"])
|
||||
assert type(image.type) is LanceBlobType
|
||||
@@ -2085,6 +2096,7 @@ def test_fetch_legacy_v1_blobs():
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_fetch_legacy_v1_blobs(tmp_path):
|
||||
lance = pytest.importorskip("lance")
|
||||
table = _legacy_v1_table("legacy_async_source")
|
||||
# Open a dataset written by Lance independently of the LanceDB connection.
|
||||
lance.write_dataset(
|
||||
@@ -2193,6 +2205,7 @@ def _blob_query_shape(table, shape):
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
@pytest.mark.parametrize("shape", ["scan", "vector", "fts", "hybrid"])
|
||||
def test_blob_mode_for_every_query_shape(shape, blob_mode, with_row_id):
|
||||
pytest.importorskip("pandas")
|
||||
table = _hybrid_blob_table(lancedb.connect("memory:///"))
|
||||
query = _blob_query_shape(table, shape).select(["id", "image"]).limit(2)
|
||||
if with_row_id is not None:
|
||||
@@ -2218,6 +2231,7 @@ def test_blob_mode_for_every_query_shape(shape, blob_mode, with_row_id):
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("shape", ["scan", "vector", "fts", "hybrid"])
|
||||
async def test_async_blob_bytes_for_every_query_shape(shape):
|
||||
pytest.importorskip("pandas")
|
||||
db = await lancedb.connect_async("memory:///")
|
||||
table = await _hybrid_blob_table_async(db)
|
||||
query = table.query()
|
||||
@@ -2750,13 +2764,15 @@ def test_add_external_blob_range_matches_storage_schema(
|
||||
table = db.create_table(
|
||||
"ranges", schema=pa.schema([pa.field("id", pa.int64()), field])
|
||||
)
|
||||
value = (
|
||||
Blob.from_uri(blob_path.as_uri(), position=position, size=6)
|
||||
if as_blob
|
||||
else {"uri": blob_path.as_uri(), "position": position, "size": 6}
|
||||
)
|
||||
if as_blob:
|
||||
Blob = pytest.importorskip("lance.blob").Blob
|
||||
value = Blob.from_uri(blob_path.as_uri(), position=position, size=6)
|
||||
else:
|
||||
value = {"uri": blob_path.as_uri(), "position": position, "size": 6}
|
||||
rows = [{"id": 1, "image": b"inline"}, {"id": 2, "image": value}]
|
||||
data = pd.DataFrame(rows) if as_pandas else rows
|
||||
data = rows
|
||||
if as_pandas:
|
||||
data = pytest.importorskip("pandas").DataFrame(rows)
|
||||
|
||||
if minimal_schema:
|
||||
with pytest.raises(ValueError, match="Blob column 'image'.*storage schema"):
|
||||
@@ -2775,6 +2791,7 @@ def test_add_external_blob_range_matches_storage_schema(
|
||||
def test_add_minimal_blob_schema_preserves_whole_file_and_inline_values(
|
||||
tmp_path, as_pandas
|
||||
):
|
||||
Blob = pytest.importorskip("lance.blob").Blob
|
||||
payload = b"whole-file"
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(payload)
|
||||
@@ -2797,10 +2814,10 @@ def test_add_minimal_blob_schema_preserves_whole_file_and_inline_values(
|
||||
{"uri": uri, "position": None, "size": None},
|
||||
]
|
||||
rows = [{"id": i, "image": value} for i, value in enumerate(values)]
|
||||
table.add(
|
||||
pd.DataFrame(rows) if as_pandas else rows,
|
||||
allow_external_blob_outside_bases=True,
|
||||
)
|
||||
data = rows
|
||||
if as_pandas:
|
||||
data = pytest.importorskip("pandas").DataFrame(rows)
|
||||
table.add(data, allow_external_blob_outside_bases=True)
|
||||
|
||||
by_id = _row_ids_by_id(table)
|
||||
assert table.fetch_blobs(
|
||||
@@ -2828,10 +2845,10 @@ def test_add_bytes_and_uri_in_one_batch_preserves_external_uri(tmp_path, as_pand
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("mixed", schema=schema)
|
||||
|
||||
table.add(
|
||||
pd.DataFrame(rows) if as_pandas else rows,
|
||||
allow_external_blob_outside_bases=True,
|
||||
)
|
||||
data = rows
|
||||
if as_pandas:
|
||||
data = pytest.importorskip("pandas").DataFrame(rows)
|
||||
table.add(data, allow_external_blob_outside_bases=True)
|
||||
|
||||
by_id = _row_ids_by_id(table)
|
||||
assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [
|
||||
@@ -2842,6 +2859,7 @@ def test_add_bytes_and_uri_in_one_batch_preserves_external_uri(tmp_path, as_pand
|
||||
|
||||
@pytest.mark.parametrize("as_pandas", [False, True])
|
||||
def test_add_mixed_python_blob_values_preserves_external_uris(tmp_path, as_pandas):
|
||||
Blob = pytest.importorskip("lance.blob").Blob
|
||||
payload = b"external-payload"
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(payload)
|
||||
@@ -2859,10 +2877,10 @@ def test_add_mixed_python_blob_values_preserves_external_uris(tmp_path, as_panda
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("mixed", schema=schema)
|
||||
|
||||
table.add(
|
||||
pd.DataFrame(rows) if as_pandas else rows,
|
||||
allow_external_blob_outside_bases=True,
|
||||
)
|
||||
data = rows
|
||||
if as_pandas:
|
||||
data = pytest.importorskip("pandas").DataFrame(rows)
|
||||
table.add(data, allow_external_blob_outside_bases=True)
|
||||
|
||||
by_id = _row_ids_by_id(table)
|
||||
assert table.fetch_blobs("image", [by_id[i] for i in range(1, 8)]).to_pylist() == [
|
||||
@@ -2886,18 +2904,22 @@ def test_merge_insert_mixed_python_blobs_validates_external_uri(tmp_path, as_pan
|
||||
{"id": 2, "image": blob_path.as_uri()},
|
||||
]
|
||||
|
||||
data = rows
|
||||
if as_pandas:
|
||||
data = pytest.importorskip("pandas").DataFrame(rows)
|
||||
with pytest.raises(ValueError, match="outside registered external bases"):
|
||||
(
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(pd.DataFrame(rows) if as_pandas else rows)
|
||||
.execute(data)
|
||||
)
|
||||
assert table.count_rows() == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("as_pandas", [False, True])
|
||||
def test_merge_insert_python_blob_dicts_and_objects(as_pandas):
|
||||
Blob = pytest.importorskip("lance.blob").Blob
|
||||
table = _blob_table("merge_python_blobs", [{"id": 1, "image": b"before"}])
|
||||
rows = [
|
||||
{"id": 1, "image": b"updated"},
|
||||
@@ -2905,11 +2927,14 @@ def test_merge_insert_python_blob_dicts_and_objects(as_pandas):
|
||||
{"id": 3, "image": Blob.from_bytes(b"from-blob")},
|
||||
]
|
||||
|
||||
data = rows
|
||||
if as_pandas:
|
||||
data = pytest.importorskip("pandas").DataFrame(rows)
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(pd.DataFrame(rows) if as_pandas else rows)
|
||||
.execute(data)
|
||||
)
|
||||
|
||||
assert result.num_updated_rows == 1
|
||||
@@ -2923,6 +2948,7 @@ def test_merge_insert_python_blob_dicts_and_objects(as_pandas):
|
||||
|
||||
|
||||
def test_create_table_with_python_blob_dict_and_object():
|
||||
Blob = pytest.importorskip("lance.blob").Blob
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table(
|
||||
@@ -2979,6 +3005,7 @@ def _openable_descriptors(table, *, by_address):
|
||||
def test_open_blob_files_preserve_snapshot_order_nulls_and_empty_values(
|
||||
tmp_path, by_address
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
db = lancedb.connect(str(tmp_path))
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("videos", schema=schema)
|
||||
@@ -3049,6 +3076,7 @@ def test_open_blob_files_does_not_open_ahead(monkeypatch):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_blob_files_async_iterates_snapshot_handles(tmp_path):
|
||||
pytest.importorskip("lance")
|
||||
db = lancedb.connect(str(tmp_path))
|
||||
table = db.create_table(
|
||||
"videos",
|
||||
|
||||
@@ -10,6 +10,8 @@ import pytest
|
||||
|
||||
@pytest.mark.parametrize("with_pylance", [False, True])
|
||||
def test_blob_array(with_pylance):
|
||||
if with_pylance:
|
||||
pytest.importorskip("lance")
|
||||
# Use fresh processes so each extension class can own Arrow's registry.
|
||||
script = textwrap.dedent(
|
||||
"""\
|
||||
|
||||
@@ -11,6 +11,8 @@ import pytest
|
||||
@pytest.mark.parametrize("api", ["sync", "async"])
|
||||
@pytest.mark.parametrize("result", ["pandas", "batches", "lazy_error"])
|
||||
def test_blob_query_without_pylance(tmp_path, api, result):
|
||||
if result == "pandas":
|
||||
pytest.importorskip("pandas")
|
||||
script = """\
|
||||
import importlib.abc
|
||||
import sys
|
||||
|
||||
@@ -3668,6 +3668,7 @@ def test_default_checkpoint_resumes_after_compaction(lance_table):
|
||||
on this path -- instead of relying on shuffle_seed/num_splits alone, as it
|
||||
does -- would reject this resume even though the seeded permutation itself
|
||||
is unchanged."""
|
||||
pytest.importorskip("lance")
|
||||
lance_table.add(pa.table({"id": [NUM_ROWS]})) # a second fragment to compact
|
||||
|
||||
ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED)
|
||||
|
||||
@@ -334,11 +334,11 @@ def test_oci_genai_table_round_trip(fake_oci):
|
||||
db = lancedb.connect("memory://")
|
||||
tbl = db.create_table("docs", schema=Docs, mode="overwrite")
|
||||
tbl.add([{"text": "hello"}, {"text": "goodbye world"}])
|
||||
assert len(tbl.to_pandas()["vector"][0]) == DIMS
|
||||
assert len(tbl.to_arrow()["vector"][0].as_py()) == DIMS
|
||||
|
||||
# the fake embeds len(text), so an equally long query lands on that row
|
||||
hits = tbl.search("goodbye world").limit(1).to_pandas()
|
||||
assert hits["text"][0] == "goodbye world"
|
||||
hits = tbl.search("goodbye world").limit(1).to_arrow()
|
||||
assert hits["text"][0].as_py() == "goodbye world"
|
||||
assert _embed_requests(fake_oci)[-1].input_type == "SEARCH_QUERY"
|
||||
|
||||
parsed = tbl.embedding_functions["vector"].function
|
||||
|
||||
@@ -39,7 +39,6 @@ from lancedb.query import (
|
||||
)
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from utils import exception_output
|
||||
@@ -154,7 +153,7 @@ def table(tmp_path) -> ldb.table.LanceTable:
|
||||
count = [local_random.randint(1, 10000) for _ in range(100)]
|
||||
table = db.create_table(
|
||||
"test",
|
||||
data=pd.DataFrame(
|
||||
data=pa.table(
|
||||
{
|
||||
"vector": vectors,
|
||||
"id": [i % 2 for i in range(100)],
|
||||
@@ -239,7 +238,7 @@ async def async_table(tmp_path) -> ldb.table.AsyncTable:
|
||||
count = [local_random.randint(1, 10000) for _ in range(100)]
|
||||
table = await db.create_table(
|
||||
"test",
|
||||
data=pd.DataFrame(
|
||||
data=pa.table(
|
||||
{
|
||||
"vector": vectors,
|
||||
"id": [i % 2 for i in range(100)],
|
||||
@@ -760,9 +759,9 @@ def test_search_order_by_ascending(table):
|
||||
|
||||
def test_create_index_from_table(tmp_path, table):
|
||||
table.create_fts_index("text")
|
||||
df = table.search("puppy").limit(5).select(["text"]).to_pandas()
|
||||
assert len(df) <= 5
|
||||
assert "text" in df.columns
|
||||
result = table.search("puppy").limit(5).select(["text"]).to_arrow()
|
||||
assert len(result) <= 5
|
||||
assert "text" in result.column_names
|
||||
|
||||
# Check whether it can be updated
|
||||
table.add(
|
||||
@@ -782,7 +781,7 @@ def test_create_index_from_table(tmp_path, table):
|
||||
table.create_fts_index("text")
|
||||
|
||||
table.create_fts_index("text", replace=True)
|
||||
assert len(table.search("gorilla").limit(1).to_pandas()) == 1
|
||||
assert len(table.search("gorilla").limit(1).to_arrow()) == 1
|
||||
|
||||
|
||||
def test_create_index_multiple_columns(tmp_path, table):
|
||||
|
||||
@@ -103,6 +103,7 @@ async def test_async_hybrid_query_filters(table: AsyncTable):
|
||||
async def test_hybrid_query_with_stale_fixed_size_binary_prefilter(
|
||||
tmpdir_factory,
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
tmp_path = str(tmpdir_factory.mktemp("stale_scalar_prefilter"))
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
|
||||
|
||||
@@ -15,13 +15,13 @@ import builtins
|
||||
import copy
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
from lance.namespace import DirectoryNamespace, LanceNamespace
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
from utils import start_server_thread
|
||||
from lancedb.listing import AsyncListing, Listing
|
||||
from lancedb.materialized_view import DEFINITION_FORMAT, DEFINITION_META_KEY
|
||||
from lancedb.namespace import (
|
||||
_MAX_QUERY_K,
|
||||
LanceNamespace,
|
||||
LanceNamespaceDBConnection,
|
||||
_query_to_namespace_request,
|
||||
)
|
||||
@@ -38,6 +38,7 @@ class TestNamespaceCreateTableOptions:
|
||||
@pytest.mark.parametrize("version", ["legacy", "stable"])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
def test_deprecated_storage_version(self, tmp_path, version, empty):
|
||||
pytest.importorskip("lance")
|
||||
db = lancedb.connect_namespace("dir", {"root": str(tmp_path / "namespace")})
|
||||
data = None if empty else [{"id": 1}]
|
||||
schema = pa.schema([("id", pa.int64())])
|
||||
@@ -71,6 +72,7 @@ class TestNamespaceCreateTableOptions:
|
||||
@pytest.mark.parametrize("version", ["legacy", "stable"])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
async def test_deprecated_storage_version_async(self, tmp_path, version, empty):
|
||||
pytest.importorskip("lance")
|
||||
root = str(tmp_path / "namespace")
|
||||
db = lancedb.connect_namespace_async("dir", {"root": root})
|
||||
data = None if empty else [{"id": 1}]
|
||||
@@ -451,13 +453,14 @@ class TestNamespaceConnection:
|
||||
assert len(table_names) == 1
|
||||
|
||||
# Verify empty table
|
||||
result = table.to_pandas()
|
||||
result = table.to_arrow()
|
||||
assert len(result) == 0
|
||||
assert list(result.columns) == ["id", "vector", "text"]
|
||||
assert result.column_names == ["id", "vector", "text"]
|
||||
|
||||
def test_table_to_pandas_blob_lazy_through_namespace(self):
|
||||
"""Namespace-backed tables should use Lance blob-aware pandas conversion."""
|
||||
pytest.importorskip("lance")
|
||||
pytest.importorskip("pandas")
|
||||
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
|
||||
db.create_namespace(["test_ns"])
|
||||
data = pa.table(
|
||||
@@ -514,9 +517,9 @@ class TestNamespaceConnection:
|
||||
assert table.id == "test_ns$test_table"
|
||||
|
||||
# Verify empty table with correct schema
|
||||
result = table.to_pandas()
|
||||
result = table.to_arrow()
|
||||
assert len(result) == 0
|
||||
assert list(result.columns) == ["id", "vector"]
|
||||
assert result.column_names == ["id", "vector"]
|
||||
|
||||
def test_drop_table_through_namespace(self):
|
||||
"""Test dropping a table through namespace."""
|
||||
@@ -672,9 +675,9 @@ class TestNamespaceConnection:
|
||||
table = db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Verify empty table was created
|
||||
result = table.to_pandas()
|
||||
result = table.to_arrow()
|
||||
assert len(result) == 0
|
||||
assert list(result.columns) == ["id", "vector", "text"]
|
||||
assert result.column_names == ["id", "vector", "text"]
|
||||
|
||||
# Test add data to the table
|
||||
new_data = [
|
||||
@@ -682,19 +685,19 @@ class TestNamespaceConnection:
|
||||
{"id": 2, "vector": [2.0, 3.0], "text": "item_2"},
|
||||
]
|
||||
table.add(new_data)
|
||||
result = table.to_pandas()
|
||||
result = table.to_arrow()
|
||||
assert len(result) == 2
|
||||
|
||||
# Test delete
|
||||
table.delete("id = 1")
|
||||
result = table.to_pandas()
|
||||
result = table.to_arrow()
|
||||
assert len(result) == 1
|
||||
assert result["id"].values[0] == 2
|
||||
assert result["id"][0].as_py() == 2
|
||||
|
||||
# Test update
|
||||
table.update(where="id = 2", values={"text": "updated"})
|
||||
result = table.to_pandas()
|
||||
assert result["text"].values[0] == "updated"
|
||||
result = table.to_arrow()
|
||||
assert result["text"][0].as_py() == "updated"
|
||||
|
||||
def test_storage_options(self):
|
||||
"""Test passing storage options through namespace connection."""
|
||||
@@ -834,28 +837,28 @@ class TestNamespaceConnection:
|
||||
opened_table_a = db.open_table(
|
||||
"same_name_table", namespace_path=["namespace_a"]
|
||||
)
|
||||
result_a = opened_table_a.to_pandas().sort_values("id").reset_index(drop=True)
|
||||
result_a = opened_table_a.to_arrow().sort_by([("id", "ascending")])
|
||||
assert len(result_a) == 2
|
||||
assert result_a["id"].tolist() == [1, 2]
|
||||
assert result_a["text"].tolist() == [
|
||||
assert result_a["id"].to_pylist() == [1, 2]
|
||||
assert result_a["text"].to_pylist() == [
|
||||
"data_from_namespace_a",
|
||||
"also_from_namespace_a",
|
||||
]
|
||||
assert [v.tolist() for v in result_a["vector"]] == [[1.0, 2.0], [3.0, 4.0]]
|
||||
assert result_a["vector"].to_pylist() == [[1.0, 2.0], [3.0, 4.0]]
|
||||
|
||||
# Verify data in namespace_b table
|
||||
opened_table_b = db.open_table(
|
||||
"same_name_table", namespace_path=["namespace_b"]
|
||||
)
|
||||
result_b = opened_table_b.to_pandas().sort_values("id").reset_index(drop=True)
|
||||
result_b = opened_table_b.to_arrow().sort_by([("id", "ascending")])
|
||||
assert len(result_b) == 3
|
||||
assert result_b["id"].tolist() == [10, 20, 30]
|
||||
assert result_b["text"].tolist() == [
|
||||
assert result_b["id"].to_pylist() == [10, 20, 30]
|
||||
assert result_b["text"].to_pylist() == [
|
||||
"data_from_namespace_b",
|
||||
"also_from_namespace_b",
|
||||
"more_from_namespace_b",
|
||||
]
|
||||
assert [v.tolist() for v in result_b["vector"]] == [
|
||||
assert result_b["vector"].to_pylist() == [
|
||||
[10.0, 20.0],
|
||||
[30.0, 40.0],
|
||||
[50.0, 60.0],
|
||||
@@ -1036,9 +1039,8 @@ class TestAsyncNamespaceConnection:
|
||||
assert result.schema.field("vector").type == pa.list_(pa.float32(), 2)
|
||||
|
||||
# Verify data content
|
||||
result_df = result.to_pandas()
|
||||
assert result_df["id"].tolist() == [1, 2, 3]
|
||||
assert [v.tolist() for v in result_df["vector"]] == [
|
||||
assert result["id"].to_pylist() == [1, 2, 3]
|
||||
assert result["vector"].to_pylist() == [
|
||||
[1.0, 2.0],
|
||||
[3.0, 4.0],
|
||||
[5.0, 6.0],
|
||||
@@ -1047,15 +1049,13 @@ class TestAsyncNamespaceConnection:
|
||||
# Test update operation
|
||||
await table.update({"id": 20}, where="id = 2")
|
||||
result = await table.to_arrow()
|
||||
result_df = result.to_pandas().sort_values("id").reset_index(drop=True)
|
||||
assert result_df["id"].tolist() == [1, 3, 20]
|
||||
assert result.sort_by([("id", "ascending")])["id"].to_pylist() == [1, 3, 20]
|
||||
|
||||
# Test delete operation
|
||||
await table.delete("id = 1")
|
||||
result = await table.to_arrow()
|
||||
assert len(result) == 2
|
||||
result_df = result.to_pandas().sort_values("id").reset_index(drop=True)
|
||||
assert result_df["id"].tolist() == [3, 20]
|
||||
assert result.sort_by([("id", "ascending")])["id"].to_pylist() == [3, 20]
|
||||
|
||||
async def test_drop_table_async(self):
|
||||
"""Test dropping a table asynchronously through namespace."""
|
||||
@@ -1248,23 +1248,24 @@ class TestPushdownOperations:
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
assert db._route_pushdown_to_rust is True
|
||||
|
||||
def test_lance_table_to_arrow_uses_query_pushdown(self):
|
||||
@pytest.mark.parametrize("as_pandas", [False, True], ids=["arrow", "pandas"])
|
||||
def test_lance_table_query_pushdown(self, as_pandas):
|
||||
if as_pandas:
|
||||
pytest.importorskip("pandas")
|
||||
namespace_client = _NamespaceClient()
|
||||
table = _namespace_lance_table(namespace_client)
|
||||
|
||||
assert table.to_arrow().equals(PUSHDOWN_DATA)
|
||||
assert table.to_pandas()["id"].tolist() == list(range(12))
|
||||
assert len(namespace_client.requests) == 2
|
||||
if as_pandas:
|
||||
assert table.to_pandas()["id"].tolist() == list(range(12))
|
||||
else:
|
||||
assert table.to_arrow().equals(PUSHDOWN_DATA)
|
||||
assert len(namespace_client.requests) == 1
|
||||
assert [request.id for request in namespace_client.requests] == [
|
||||
["geneva", "hist"],
|
||||
["geneva", "hist"],
|
||||
["geneva", "hist"]
|
||||
]
|
||||
# Unlimited reads cap k at i32::MAX (the namespace query_table `k`
|
||||
# field is i32); sys.maxsize would overflow the Rust binding.
|
||||
assert [request.k for request in namespace_client.requests] == [
|
||||
_MAX_QUERY_K,
|
||||
_MAX_QUERY_K,
|
||||
]
|
||||
assert [request.k for request in namespace_client.requests] == [_MAX_QUERY_K]
|
||||
assert all(r.k <= 2**31 - 1 for r in namespace_client.requests)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -1345,7 +1346,10 @@ class TestAsyncPushdownOperations:
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
assert len(db._namespace_client_pushdown_operations) == 0
|
||||
|
||||
async def test_async_table_to_arrow_uses_query_pushdown(self):
|
||||
@pytest.mark.parametrize("as_pandas", [False, True], ids=["arrow", "pandas"])
|
||||
async def test_async_table_query_pushdown(self, as_pandas):
|
||||
if as_pandas:
|
||||
pytest.importorskip("pandas")
|
||||
namespace_client = _NamespaceClient()
|
||||
|
||||
table = AsyncTable(
|
||||
@@ -1355,19 +1359,17 @@ class TestAsyncPushdownOperations:
|
||||
pushdown_operations={"QueryTable"},
|
||||
)
|
||||
|
||||
assert (await table.to_arrow()).equals(PUSHDOWN_DATA)
|
||||
assert (await table.to_pandas())["id"].tolist() == list(range(12))
|
||||
assert len(namespace_client.requests) == 2
|
||||
if as_pandas:
|
||||
assert (await table.to_pandas())["id"].tolist() == list(range(12))
|
||||
else:
|
||||
assert (await table.to_arrow()).equals(PUSHDOWN_DATA)
|
||||
assert len(namespace_client.requests) == 1
|
||||
assert [request.id for request in namespace_client.requests] == [
|
||||
["geneva", "hist"],
|
||||
["geneva", "hist"],
|
||||
["geneva", "hist"]
|
||||
]
|
||||
# Unlimited reads cap k at i32::MAX (the namespace query_table `k`
|
||||
# field is i32); sys.maxsize would overflow the Rust binding.
|
||||
assert [request.k for request in namespace_client.requests] == [
|
||||
_MAX_QUERY_K,
|
||||
_MAX_QUERY_K,
|
||||
]
|
||||
assert [request.k for request in namespace_client.requests] == [_MAX_QUERY_K]
|
||||
assert all(r.k <= 2**31 - 1 for r in namespace_client.requests)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -1396,7 +1398,10 @@ class TestAsyncPushdownOperations:
|
||||
assert namespace_client.requests == []
|
||||
|
||||
|
||||
def test_local_table_to_arrow_and_to_pandas_are_unchanged(tmp_path):
|
||||
@pytest.mark.parametrize("as_pandas", [False, True], ids=["arrow", "pandas"])
|
||||
def test_local_table_to_arrow_and_to_pandas_are_unchanged(tmp_path, as_pandas):
|
||||
if as_pandas:
|
||||
pytest.importorskip("pandas")
|
||||
db = lancedb.connect(str(tmp_path / "db"))
|
||||
table = db.create_table(
|
||||
"local",
|
||||
@@ -1406,8 +1411,10 @@ def test_local_table_to_arrow_and_to_pandas_are_unchanged(tmp_path):
|
||||
],
|
||||
)
|
||||
|
||||
assert table.to_arrow().column("id").to_pylist() == [1, 2]
|
||||
assert table.to_pandas()["id"].tolist() == [1, 2]
|
||||
if as_pandas:
|
||||
assert table.to_pandas()["id"].tolist() == [1, 2]
|
||||
else:
|
||||
assert table.to_arrow().column("id").to_pylist() == [1, 2]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("effort", [0.0, 0.5, 1.0])
|
||||
@@ -1432,6 +1439,8 @@ class AttributeAccessNamespace(LanceNamespace):
|
||||
"""
|
||||
|
||||
def __init__(self, root: str):
|
||||
from lance.namespace import DirectoryNamespace
|
||||
|
||||
self._inner = DirectoryNamespace(root=root)
|
||||
self.seen = {}
|
||||
|
||||
@@ -1453,6 +1462,7 @@ class AttributeAccessNamespace(LanceNamespace):
|
||||
|
||||
|
||||
def test_python_namespace_attribute_access(tmp_path):
|
||||
pytest.importorskip("lance.namespace")
|
||||
root = str(tmp_path)
|
||||
data = pa.table({"id": [1, 2, 3], "text": ["alpha", "beta", "gamma"]})
|
||||
lancedb.connect_namespace("dir", {"root": root}).create_table("t", data)
|
||||
|
||||
@@ -946,12 +946,12 @@ def test_iter_with_different_formats(some_permutation: Permutation):
|
||||
python_batches = list(python_perm.iter(batch_size, skip_last_batch=False))
|
||||
assert all(isinstance(batch, list) for batch in python_batches)
|
||||
|
||||
# Test with pandas format
|
||||
|
||||
def test_iter_with_pandas_format(some_permutation: Permutation):
|
||||
pd = pytest.importorskip("pandas")
|
||||
batch_size = 100
|
||||
pandas_perm = some_permutation.with_format("pandas")
|
||||
pandas_batches = list(pandas_perm.iter(batch_size, skip_last_batch=False))
|
||||
# Import pandas to check the type
|
||||
import pandas as pd
|
||||
|
||||
assert all(isinstance(batch, pd.DataFrame) for batch in pandas_batches)
|
||||
|
||||
|
||||
@@ -1053,8 +1053,6 @@ def test_identity_permutation(mem_db):
|
||||
|
||||
|
||||
def test_transform_fn(mem_db):
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import polars as pl
|
||||
|
||||
tbl = mem_db.create_table(
|
||||
@@ -1062,20 +1060,6 @@ def test_transform_fn(mem_db):
|
||||
)
|
||||
permutation = Permutation.identity(tbl)
|
||||
|
||||
np_result = list(permutation.with_format("numpy").iter(10, skip_last_batch=False))[
|
||||
0
|
||||
]
|
||||
assert np_result.shape == (10, 2)
|
||||
assert np_result.dtype == np.int64
|
||||
assert isinstance(np_result, np.ndarray)
|
||||
|
||||
pd_result = list(permutation.with_format("pandas").iter(10, skip_last_batch=False))[
|
||||
0
|
||||
]
|
||||
assert pd_result.shape == (10, 2)
|
||||
assert pd_result.dtypes.tolist() == [np.int64, np.int64]
|
||||
assert isinstance(pd_result, pd.DataFrame)
|
||||
|
||||
pl_result = list(permutation.with_format("polars").iter(10, skip_last_batch=False))[
|
||||
0
|
||||
]
|
||||
@@ -1141,6 +1125,27 @@ def test_transform_fn(mem_db):
|
||||
assert isinstance(arrow_result, pa.RecordBatch)
|
||||
|
||||
|
||||
def test_transform_fn_pandas(mem_db):
|
||||
pd = pytest.importorskip("pandas")
|
||||
import numpy as np
|
||||
|
||||
tbl = mem_db.create_table(
|
||||
"test_table", pa.table({"id": range(10), "value": range(10)})
|
||||
)
|
||||
permutation = Permutation.identity(tbl)
|
||||
numpy_result = list(
|
||||
permutation.with_format("numpy").iter(10, skip_last_batch=False)
|
||||
)[0]
|
||||
assert numpy_result.shape == (10, 2)
|
||||
assert numpy_result.dtype == np.int64
|
||||
assert isinstance(numpy_result, np.ndarray)
|
||||
|
||||
result = list(permutation.with_format("pandas").iter(10, skip_last_batch=False))[0]
|
||||
assert result.shape == (10, 2)
|
||||
assert result.dtypes.tolist() == [np.int64, np.int64]
|
||||
assert isinstance(result, pd.DataFrame)
|
||||
|
||||
|
||||
def test_custom_transform(mem_db):
|
||||
tbl = mem_db.create_table(
|
||||
"test_table", pa.table({"id": range(10), "value": range(10)})
|
||||
|
||||
@@ -13,6 +13,8 @@ from lancedb.namespace import LanceNamespaceDBConnection
|
||||
@pytest.mark.parametrize("seconds", [-1.0, -0.000001])
|
||||
@pytest.mark.parametrize("kind", ["sync", "namespace", "namespace_async", "client"])
|
||||
def test_negative_read_consistency_interval(tmp_path, seconds, kind):
|
||||
if kind == "client":
|
||||
pytest.importorskip("lance")
|
||||
interval = timedelta(seconds=seconds)
|
||||
with pytest.raises(ValueError, match="read_consistency_interval"):
|
||||
if kind == "sync":
|
||||
@@ -44,6 +46,8 @@ async def test_negative_async_read_consistency_interval(tmp_path, seconds):
|
||||
async def test_binding_rejects_invalid_read_consistency_interval(
|
||||
tmp_path, seconds, kind
|
||||
):
|
||||
if kind == "client":
|
||||
pytest.importorskip("lance")
|
||||
with pytest.raises(ValueError, match="read_consistency_interval"):
|
||||
if kind == "connect":
|
||||
await _lancedb.connect(str(tmp_path), read_consistency_interval=seconds)
|
||||
@@ -71,12 +75,23 @@ async def test_valid_read_consistency_intervals(tmp_path, interval):
|
||||
assert await db.get_read_consistency_interval() == interval
|
||||
db.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"interval", [None, timedelta(0), timedelta(microseconds=1), timedelta(seconds=5)]
|
||||
)
|
||||
def test_valid_namespace_read_consistency_intervals(tmp_path, interval):
|
||||
db = lancedb.connect_namespace(
|
||||
"dir", {"root": str(tmp_path / "namespace")}, read_consistency_interval=interval
|
||||
)
|
||||
assert db.read_consistency_interval == interval
|
||||
db.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"interval", [None, timedelta(0), timedelta(microseconds=1), timedelta(seconds=5)]
|
||||
)
|
||||
def test_valid_namespace_client_read_consistency_intervals(tmp_path, interval):
|
||||
pytest.importorskip("lance")
|
||||
client = connect_namespace_client("dir", {"root": str(tmp_path / "client")})
|
||||
db = LanceNamespaceDBConnection(client, read_consistency_interval=interval)
|
||||
assert db.read_consistency_interval == interval
|
||||
|
||||
@@ -1307,7 +1307,7 @@ def test_remote_conversion_blob_mode_errors(method, blob_mode):
|
||||
async def test_remote_to_lance_rejection(
|
||||
tmp_path, version, expose_location, monkeypatch
|
||||
):
|
||||
import lance
|
||||
lance = pytest.importorskip("lance")
|
||||
|
||||
location = tmp_path / "test.lance"
|
||||
lance.write_dataset(pa.table({"id": [1]}), location)
|
||||
@@ -4691,7 +4691,8 @@ def test_remote_blob_columns_and_fetch():
|
||||
def test_sync_remote_blob_pandas_avoids_storage_access(
|
||||
tmp_path, monkeypatch, blob_mode
|
||||
):
|
||||
import lance
|
||||
pytest.importorskip("pandas")
|
||||
lance = pytest.importorskip("lance")
|
||||
|
||||
location = tmp_path / "remote.lance"
|
||||
lance.write_dataset(pa.table({"id": [7], "image": [b"server-owned"]}), location)
|
||||
@@ -4764,6 +4765,7 @@ async def test_async_remote_to_lance_rejects_storage_access(monkeypatch):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_blob_descriptions_to_pandas_uses_server(monkeypatch):
|
||||
pytest.importorskip("pandas")
|
||||
with blob_remote_table() as remote_table:
|
||||
table = remote_table._table
|
||||
|
||||
@@ -4859,6 +4861,7 @@ async def test_async_remote_to_arrow_bytes_refused_before_query(monkeypatch):
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "lazy"])
|
||||
async def test_async_remote_nonblob_projection_to_pandas(blob_mode):
|
||||
pytest.importorskip("pandas")
|
||||
with blob_remote_table() as remote_table:
|
||||
df = await (
|
||||
remote_table._table.query().select(["id"]).to_pandas(blob_mode=blob_mode)
|
||||
@@ -4933,6 +4936,7 @@ def test_remote_blob_query_stashes_row_ids_for_fetch():
|
||||
|
||||
|
||||
def test_remote_blob_query_to_pandas():
|
||||
pytest.importorskip("pandas")
|
||||
with blob_remote_table() as table:
|
||||
query = table.search().select(["id", "image"]).limit(3)
|
||||
df = query.to_pandas(blob_mode="descriptions", split_blocks=True)
|
||||
@@ -4959,6 +4963,7 @@ def test_remote_blob_query_to_pandas_rejects_unpinned_fetch(monkeypatch, blob_mo
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["lazy", "bytes"])
|
||||
def test_remote_nonblob_projection_to_pandas(blob_mode):
|
||||
pytest.importorskip("pandas")
|
||||
with blob_remote_table() as table:
|
||||
df = table.search().select(["id"]).to_pandas(blob_mode=blob_mode)
|
||||
|
||||
@@ -4968,6 +4973,7 @@ def test_remote_nonblob_projection_to_pandas(blob_mode):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_blob_query_to_pandas_flatten():
|
||||
pytest.importorskip("pandas")
|
||||
with blob_remote_table() as table:
|
||||
df = (
|
||||
await table._table.query()
|
||||
|
||||
@@ -127,7 +127,7 @@ def _run_test_reranker(
|
||||
if deterministic:
|
||||
assert result1 == result2
|
||||
|
||||
query_vector = table.to_pandas()["vector"][0]
|
||||
query_vector = table.to_arrow()["vector"][0].as_py()
|
||||
result = (
|
||||
table.search(query_type="hybrid", vector_column_name="vector")
|
||||
.vector(query_vector)
|
||||
@@ -250,7 +250,7 @@ def _run_test_hybrid_reranker(reranker, tmp_path):
|
||||
assert result1 == result3 # 2 & 3 should be the same as they use score as score
|
||||
|
||||
query = "Our father who art in heaven"
|
||||
query_vector = table.to_pandas()["vector"][0]
|
||||
query_vector = table.to_arrow()["vector"][0].as_py()
|
||||
result = (
|
||||
table.search(query_type="hybrid", vector_column_name="vector")
|
||||
.vector(query_vector)
|
||||
|
||||
@@ -302,6 +302,7 @@ def test_table_to_pandas_invalid_blob_mode_non_blob_table(tmp_db: DBConnection):
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"])
|
||||
def test_table_to_pandas_blob_modes(tmp_db: DBConnection, blob_mode):
|
||||
pytest.importorskip("pandas")
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
f"test_to_pandas_blob_{blob_mode}",
|
||||
@@ -324,6 +325,7 @@ def test_table_to_pandas_blob_modes(tmp_db: DBConnection, blob_mode):
|
||||
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
def test_table_to_pandas_blob_bytes(tmp_db: DBConnection, blob_schema):
|
||||
pytest.importorskip("pandas")
|
||||
pytest.importorskip("lance")
|
||||
table = _blob_table(tmp_db, f"test_to_pandas_blob_{blob_schema}_bytes", blob_schema)
|
||||
|
||||
@@ -349,6 +351,7 @@ def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
async def test_async_table_to_pandas_blob_bytes(
|
||||
tmp_db_async: AsyncConnection, blob_schema
|
||||
):
|
||||
pytest.importorskip("pandas")
|
||||
pytest.importorskip("lance")
|
||||
table = await _blob_table_async(
|
||||
tmp_db_async, f"test_async_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
||||
@@ -932,8 +935,8 @@ def test_add_nullability(mem_db: DBConnection):
|
||||
assert table.to_arrow() == expected
|
||||
|
||||
|
||||
def test_add_pydantic_model(mem_db: DBConnection):
|
||||
pytest.importorskip("pandas")
|
||||
@pytest.fixture
|
||||
def pydantic_model_table(mem_db: DBConnection):
|
||||
# https://github.com/lancedb/lancedb/issues/562
|
||||
|
||||
class Metadata(BaseModel):
|
||||
@@ -951,9 +954,6 @@ def test_add_pydantic_model(mem_db: DBConnection):
|
||||
payload: Document
|
||||
|
||||
tbl = mem_db.create_table("mytable", schema=LanceSchema, mode="overwrite")
|
||||
assert tbl.schema == LanceSchema.to_arrow_schema()
|
||||
|
||||
# add works
|
||||
expected = LanceSchema(
|
||||
id="id",
|
||||
vector=[0.0, 0.0],
|
||||
@@ -963,11 +963,21 @@ def test_add_pydantic_model(mem_db: DBConnection):
|
||||
),
|
||||
)
|
||||
add_res = tbl.add([expected])
|
||||
return tbl, LanceSchema, expected, add_res
|
||||
|
||||
|
||||
def test_add_pydantic_model(pydantic_model_table):
|
||||
tbl, LanceSchema, expected, add_res = pydantic_model_table
|
||||
assert tbl.schema == LanceSchema.to_arrow_schema()
|
||||
assert add_res.version == 2
|
||||
|
||||
result = tbl.search([0.0, 0.0]).limit(1).to_pydantic(LanceSchema)[0]
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_add_pydantic_model_pandas(pydantic_model_table):
|
||||
pytest.importorskip("pandas")
|
||||
tbl, _, _, _ = pydantic_model_table
|
||||
flattened = tbl.search([0.0, 0.0]).limit(1).to_pandas(flatten=1)
|
||||
assert len(flattened.columns) == 6 # _distance is automatically added
|
||||
|
||||
|
||||
@@ -4,12 +4,11 @@
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
from datetime import time, timedelta, timezone
|
||||
from decimal import Decimal, localcontext
|
||||
from typing import Optional
|
||||
|
||||
import lance
|
||||
from lance.blob import BlobType as LanceBlobType
|
||||
from lancedb.conftest import MockTextEmbeddingFunction
|
||||
from lancedb.embeddings.base import EmbeddingFunctionConfig
|
||||
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
@@ -26,7 +25,6 @@ from lancedb.table import (
|
||||
sanitize_create_table,
|
||||
)
|
||||
import pyarrow as pa
|
||||
import pandas as pd
|
||||
import polars as pl
|
||||
import pytest
|
||||
import lancedb
|
||||
@@ -180,7 +178,8 @@ def test_value_to_sql_string(tmp_path):
|
||||
)
|
||||
for value in values:
|
||||
table.update(where=f"search = {value_to_sql(value)}", values={"replace": value})
|
||||
assert table.to_pandas().query("search == @value")["replace"].item() == value
|
||||
rows = table.to_arrow().to_pylist()
|
||||
assert next(row["replace"] for row in rows if row["search"] == value) == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -275,10 +274,15 @@ def test_value_to_sql_time_and_timedelta(tmp_path):
|
||||
def test_value_to_sql_keeps_nanoseconds(tmp_path):
|
||||
# Arrow scalars and pandas.Timedelta values can hold nanoseconds, which a
|
||||
# Python timedelta or time would round to microseconds.
|
||||
assert (
|
||||
value_to_sql(pd.Timedelta(1001, "ns"))
|
||||
== "arrow_cast(1001, 'Duration(Nanosecond)')"
|
||||
)
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
pd = None
|
||||
if pd is not None:
|
||||
assert (
|
||||
value_to_sql(pd.Timedelta(1001, "ns"))
|
||||
== "arrow_cast(1001, 'Duration(Nanosecond)')"
|
||||
)
|
||||
assert value_to_sql(pa.scalar(None, pa.duration("ns"))) == "NULL"
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
@@ -781,15 +785,19 @@ class TestModel(lancedb.pydantic.LanceModel):
|
||||
lambda: pa.table({"a": [1], "b": [2]}).to_reader(),
|
||||
lambda: iter(pa.table({"a": [1], "b": [2]}).to_batches()),
|
||||
lambda: tuple(pa.table({"a": [1], "b": [2]}).to_batches()),
|
||||
lambda: lance.write_dataset(
|
||||
lambda: pytest.importorskip("lance").write_dataset(
|
||||
pa.table({"a": [1], "b": [2]}),
|
||||
"memory://test",
|
||||
),
|
||||
lambda: lance.write_dataset(
|
||||
pa.table({"a": [1], "b": [2]}),
|
||||
"memory://test",
|
||||
).scanner(),
|
||||
lambda: pd.DataFrame({"a": [1], "b": [2]}),
|
||||
lambda: (
|
||||
pytest.importorskip("lance")
|
||||
.write_dataset(
|
||||
pa.table({"a": [1], "b": [2]}),
|
||||
"memory://test",
|
||||
)
|
||||
.scanner()
|
||||
),
|
||||
lambda: pytest.importorskip("pandas").DataFrame({"a": [1], "b": [2]}),
|
||||
lambda: pl.DataFrame({"a": [1], "b": [2]}),
|
||||
lambda: pl.LazyFrame({"a": [1], "b": [2]}),
|
||||
lambda: [TestModel(a=1, b=2)],
|
||||
@@ -966,10 +974,10 @@ def test_infer_target_schema_with_vector_embedding_names():
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
[
|
||||
[{"id": 1, "text": "hello"}],
|
||||
pa.RecordBatch.from_pylist([{"id": 1, "text": "hello"}]),
|
||||
pd.DataFrame({"id": [1], "text": ["hello"]}),
|
||||
pl.DataFrame({"id": [1], "text": ["hello"]}),
|
||||
lambda: [{"id": 1, "text": "hello"}],
|
||||
lambda: pa.RecordBatch.from_pylist([{"id": 1, "text": "hello"}]),
|
||||
lambda: pytest.importorskip("pandas").DataFrame({"id": [1], "text": ["hello"]}),
|
||||
lambda: pl.DataFrame({"id": [1], "text": ["hello"]}),
|
||||
],
|
||||
ids=["rows", "pa.RecordBatch", "pd.DataFrame", "pl.DataFrame"],
|
||||
)
|
||||
@@ -1001,6 +1009,7 @@ def test_sanitize_data(
|
||||
schema: Optional[pa.Schema],
|
||||
with_embedding: bool,
|
||||
):
|
||||
data = data()
|
||||
if with_embedding:
|
||||
registry = EmbeddingFunctionRegistry.get_instance()
|
||||
registry.register("test")(MockTextEmbeddingFunction)
|
||||
@@ -1025,7 +1034,9 @@ def test_sanitize_data(
|
||||
# polars uses large_string, pandas 3.0+ uses large_string, others use string
|
||||
if isinstance(data, pl.DataFrame):
|
||||
text_type = pa.large_utf8()
|
||||
elif isinstance(data, pd.DataFrame):
|
||||
elif (pd := sys.modules.get("pandas")) is not None and isinstance(
|
||||
data, pd.DataFrame
|
||||
):
|
||||
text_type = pandas_string_type()
|
||||
else:
|
||||
text_type = pa.string()
|
||||
@@ -1360,6 +1371,8 @@ def test_fixed_size_list_blob_coercion_keeps_null_rows():
|
||||
|
||||
|
||||
def test_cast_to_target_schema_accepts_pylance_blob_v2():
|
||||
lance = pytest.importorskip("lance")
|
||||
LanceBlobType = pytest.importorskip("lance.blob").BlobType
|
||||
target_type = lancedb.BlobType()
|
||||
source = lance.blob_array([b"hello", None])
|
||||
assert type(source.type) is LanceBlobType
|
||||
@@ -1379,6 +1392,8 @@ def test_cast_to_target_schema_accepts_pylance_blob_v2():
|
||||
|
||||
|
||||
def test_cast_to_target_schema_rejects_different_blob_v2_class():
|
||||
lance = pytest.importorskip("lance")
|
||||
|
||||
class OtherBlobType(pa.ExtensionType):
|
||||
def __init__(self):
|
||||
super().__init__(lancedb.BlobType().storage_type, "lance.blob.v2")
|
||||
|
||||
Reference in New Issue
Block a user