mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-15 16:22:27 +00:00
fix(python): resolve the single vector column despite a query dimension mismatch (#4526)
\`inf_vector_column_query\`'s \`dim\` filter (added in #3423) is meant to disambiguate between multiple vector columns of different sizes. When a schema has exactly one vector column and the query's dimension doesn't match it, the filter excluded that column too, raising "There is no vector column in the data" instead of resolving the column and letting the dimension mismatch surface at query execution. This was latent until #4409 wired client-side vector-column inference into remote search: a mismatched-dimension query against a single-vector-column remote table now fails client-side with the wrong error, before ever reaching the server's own dimension validation. Fall back to the unfiltered candidate set when the dimension filter leaves zero matches, and only raise when that is still ambiguous or genuinely empty.
This commit is contained in:
@@ -236,6 +236,19 @@ def inf_vector_column_query(schema: pa.Schema, dim: Optional[int] = None) -> str
|
||||
vector_col_names = []
|
||||
for field in schema:
|
||||
vector_col_names.extend(_iter_vector_columns(field, [], dim))
|
||||
|
||||
if len(vector_col_names) == 0 and dim is not None:
|
||||
# The query's dimension matched no column. If the schema is
|
||||
# unambiguous without that filter, the single vector column is still
|
||||
# the right one to search: a dimension mismatch is a query-execution
|
||||
# error, not a missing-column error.
|
||||
unfiltered = []
|
||||
for field in schema:
|
||||
unfiltered.extend(_iter_vector_columns(field, []))
|
||||
if len(unfiltered) == 1:
|
||||
return unfiltered[0]
|
||||
vector_col_names = unfiltered
|
||||
|
||||
if len(vector_col_names) > 1:
|
||||
raise ValueError(
|
||||
"Schema has more than one vector column. "
|
||||
|
||||
@@ -4016,6 +4016,26 @@ def test_search_nested_vector_no_candidates(mem_db: DBConnection):
|
||||
table.search([0.0, 1.0]).limit(1).to_arrow()
|
||||
|
||||
|
||||
def test_search_single_vector_column_wrong_query_dim(mem_db: DBConnection):
|
||||
"""A query vector whose dimension doesn't match the (only) vector column
|
||||
must still resolve that column, surfacing a dimension-mismatch error from
|
||||
query execution rather than a misleading "no vector column" error."""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.uint32(), False),
|
||||
pa.field("vector_col", pa.list_(pa.float32(), 32), False),
|
||||
]
|
||||
)
|
||||
table = mem_db.create_table(
|
||||
"wrong_query_dim",
|
||||
data=[{"id": 1, "vector_col": [0.1] * 32}],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="dim"):
|
||||
table.search([0.01] * 120).limit(1).to_arrow()
|
||||
|
||||
|
||||
def test_compact_cleanup(tmp_db: DBConnection):
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
|
||||
@@ -28,7 +28,13 @@ import pandas as pd
|
||||
import polars as pl
|
||||
import pytest
|
||||
import lancedb
|
||||
from lancedb.util import flatten_columns, get_uri_scheme, join_uri, value_to_sql
|
||||
from lancedb.util import (
|
||||
flatten_columns,
|
||||
get_uri_scheme,
|
||||
inf_vector_column_query,
|
||||
join_uri,
|
||||
value_to_sql,
|
||||
)
|
||||
from utils import exception_output
|
||||
|
||||
|
||||
@@ -1251,3 +1257,41 @@ def test_infer_vector_column_raises_clear_error(tmp_path):
|
||||
# Plain vector search on a table with no vector column should raise
|
||||
# a clear ValueError, not a cryptic TypeError.
|
||||
table.search([1.0, 2.0]).to_list()
|
||||
|
||||
|
||||
def test_inf_vector_column_query_dim_mismatch_still_resolves():
|
||||
"""Regression: a query dimension that matches no column must not hide an
|
||||
otherwise-unambiguous vector column behind "no vector column in the data".
|
||||
|
||||
`dim` exists to disambiguate between *multiple* vector columns of
|
||||
different sizes. When the schema has exactly one vector column, that
|
||||
column is still the only sensible one to search; a dimension mismatch is
|
||||
a query-execution concern, not a missing-column one. This previously
|
||||
broke remote search (#4409 wired up client-side inference there), which
|
||||
expects the server to reject the mismatched dimension itself.
|
||||
"""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.uint32(), False),
|
||||
pa.field("vector_col", pa.list_(pa.float32(), 32), False),
|
||||
pa.field("s", pa.string(), False),
|
||||
]
|
||||
)
|
||||
|
||||
assert inf_vector_column_query(schema, dim=32) == "vector_col"
|
||||
assert inf_vector_column_query(schema, dim=120) == "vector_col"
|
||||
assert inf_vector_column_query(schema) == "vector_col"
|
||||
|
||||
|
||||
def test_inf_vector_column_query_dim_mismatch_with_multiple_columns():
|
||||
"""With more than one vector column, a dimension that matches none of
|
||||
them is still ambiguous and must raise, not silently pick one."""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("vector1", pa.list_(pa.float32(), 10), False),
|
||||
pa.field("vector2", pa.list_(pa.float32(), 20), False),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="more than one vector column"):
|
||||
inf_vector_column_query(schema, dim=99)
|
||||
|
||||
Reference in New Issue
Block a user