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:
Jack Ye
2026-10-14 17:40:04 -07:00
committed by GitHub
parent 4be1720c51
commit 9f37f39267
3 changed files with 78 additions and 1 deletions
+13
View File
@@ -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. "
+20
View File
@@ -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(
+45 -1
View File
@@ -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)