mirror of
https://github.com/lancedb/lancedb.git
synced 2025-12-27 07:09:57 +00:00
feat: add maximum and minimum nprobes properties (#2430)
This exposes the maximum_nprobes and minimum_nprobes feature that was added in https://github.com/lancedb/lance/pull/3903 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added support for specifying minimum and maximum probe counts in vector search queries, allowing finer control over search behavior. - Users can now independently set minimum and maximum probes for vector and hybrid queries via new methods and parameters in Python, Node.js, and Rust APIs. - **Bug Fixes** - Improved parameter validation to ensure correct usage of minimum and maximum probe values. - **Tests** - Expanded test coverage to validate correct handling, serialization, and error cases for the new probe parameters. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
@@ -439,6 +439,33 @@ def test_query_builder_with_filter(table):
|
||||
assert all(np.array(rs[0]["vector"]) == [3, 4])
|
||||
|
||||
|
||||
def test_invalid_nprobes_sync(table):
|
||||
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(0).to_list()
|
||||
with pytest.raises(
|
||||
ValueError, match="maximum_nprobes must be greater than minimum_nprobes"
|
||||
):
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector").maximum_nprobes(5).to_list()
|
||||
with pytest.raises(
|
||||
ValueError, match="minimum_nprobes must be less or equal to maximum_nprobes"
|
||||
):
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(100).to_list()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_nprobes_async(table_async: AsyncTable):
|
||||
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
|
||||
await table_async.vector_search([0, 0]).minimum_nprobes(0).to_list()
|
||||
with pytest.raises(
|
||||
ValueError, match="maximum_nprobes must be greater than minimum_nprobes"
|
||||
):
|
||||
await table_async.vector_search([0, 0]).maximum_nprobes(5).to_list()
|
||||
with pytest.raises(
|
||||
ValueError, match="minimum_nprobes must be less or equal to maximum_nprobes"
|
||||
):
|
||||
await table_async.vector_search([0, 0]).minimum_nprobes(100).to_list()
|
||||
|
||||
|
||||
def test_query_builder_with_prefilter(table):
|
||||
df = (
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector")
|
||||
@@ -585,6 +612,21 @@ async def test_query_async(table_async: AsyncTable):
|
||||
table_async.query().nearest_to(pa.array([1, 2])).nprobes(10),
|
||||
expected_num_rows=2,
|
||||
)
|
||||
await check_query(
|
||||
table_async.query().nearest_to(pa.array([1, 2])).minimum_nprobes(10),
|
||||
expected_num_rows=2,
|
||||
)
|
||||
await check_query(
|
||||
table_async.query().nearest_to(pa.array([1, 2])).maximum_nprobes(30),
|
||||
expected_num_rows=2,
|
||||
)
|
||||
await check_query(
|
||||
table_async.query()
|
||||
.nearest_to(pa.array([1, 2]))
|
||||
.minimum_nprobes(10)
|
||||
.maximum_nprobes(20),
|
||||
expected_num_rows=2,
|
||||
)
|
||||
await check_query(
|
||||
table_async.query().nearest_to(pa.array([1, 2])).bypass_vector_index(),
|
||||
expected_num_rows=2,
|
||||
@@ -911,7 +953,39 @@ def test_query_serialization_sync(table: lancedb.table.Table):
|
||||
|
||||
q = table.search([5.0, 6.0]).nprobes(10).refine_factor(5).to_query_object()
|
||||
check_set_props(
|
||||
q, vector_column="vector", vector=[5.0, 6.0], nprobes=10, refine_factor=5
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
minimum_nprobes=10,
|
||||
maximum_nprobes=10,
|
||||
refine_factor=5,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).minimum_nprobes(10).to_query_object()
|
||||
check_set_props(
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
minimum_nprobes=10,
|
||||
maximum_nprobes=None,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).nprobes(50).to_query_object()
|
||||
check_set_props(
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
minimum_nprobes=50,
|
||||
maximum_nprobes=50,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).maximum_nprobes(10).to_query_object()
|
||||
check_set_props(
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
maximum_nprobes=10,
|
||||
minimum_nprobes=None,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).distance_range(0.0, 1.0).to_query_object()
|
||||
@@ -963,7 +1037,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
limit=10,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
nprobes=20,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
)
|
||||
@@ -973,7 +1048,20 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
q,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
nprobes=20,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
)
|
||||
|
||||
q = (await table_async.search([5.0, 6.0])).nprobes(50).to_query_object()
|
||||
check_set_props(
|
||||
q,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
minimum_nprobes=50,
|
||||
maximum_nprobes=50,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -992,7 +1080,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
filter="id = 1",
|
||||
postfilter=True,
|
||||
vector=sample_vector,
|
||||
nprobes=20,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
)
|
||||
@@ -1006,7 +1095,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
check_set_props(
|
||||
q,
|
||||
vector=sample_vector,
|
||||
nprobes=10,
|
||||
minimum_nprobes=10,
|
||||
maximum_nprobes=10,
|
||||
refine_factor=5,
|
||||
postfilter=False,
|
||||
with_row_id=False,
|
||||
@@ -1014,6 +1104,18 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
limit=10,
|
||||
)
|
||||
|
||||
q = (await table_async.search([5.0, 6.0])).minimum_nprobes(5).to_query_object()
|
||||
check_set_props(
|
||||
q,
|
||||
vector=sample_vector,
|
||||
minimum_nprobes=5,
|
||||
maximum_nprobes=20,
|
||||
postfilter=False,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
)
|
||||
|
||||
q = (
|
||||
(await table_async.search([5.0, 6.0]))
|
||||
.distance_range(0.0, 1.0)
|
||||
@@ -1025,7 +1127,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
lower_bound=0.0,
|
||||
upper_bound=1.0,
|
||||
postfilter=False,
|
||||
nprobes=20,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1037,7 +1140,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
distance_type="cosine",
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
nprobes=20,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1049,7 +1153,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
ef=7,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
nprobes=20,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1061,7 +1166,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
bypass_vector_index=True,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
nprobes=20,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
with_row_id=False,
|
||||
limit=10,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user