mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-13 15:22:30 +00:00
feat(search): expose search effort (#4520)
Lance supports `search_effort`, but LanceDB cannot pass it through its query APIs. Expose the parameter in Rust, Python sync/async and TypeScript, including remote requests and hybrid query conversion, using the [Lance v14.0.0-beta.13](https://github.com/lance-format/lance/releases/tag/v14.0.0-beta.13) dependency already on main. The default is `None` (Auto, internally 0.5). Any explicit effort, including 0.5, overrides fixed `nprobes` in either setter order; explicit minimum/maximum bounds remain effective. Invalid values are rejected in the Rust core. Hybrid conversion now retains the vector probe parameters and handles an inferred FTS column without panicking. Namespace queries with explicit effort execute locally because the namespace wire model cannot carry it. MemWAL LSM queries reject it explicitly because Lance's LSM scanner does not yet expose this option. Maturin is included in the dev extra so the documented build command works in the project environment. BREAKING CHANGE: The public Rust `VectorQueryRequest` struct now includes `search_effort: Option<f64>`. Downstream code that initializes every field explicitly must add `search_effort: None` or use `..Default::default()`. Exhaustive destructuring patterns must include the new field or `..`. Query-builder calls and struct literals that already use `..Default::default()` require no source changes. The new effort setters in Python and TypeScript are additive. This PR does not change the persisted file format. Using explicit effort with remote tables requires server support for `search_effort` and its precedence over `nprobes`; the client wire tests do not establish that a deployed server supports it. Server-side support is tracked in lancedb/sophon#8607. Validation on AWS after merging current main (Lance beta.13 / SDK beta.5): - 14 Rust regressions and the API doctest pass. - 60 parameter regressions, the full 47-test namespace suite, and three API doctests pass. The parameter selection includes some namespace cases. - 40 TypeScript query tests and 54 remote tests pass; build, lint and generated docs pass. - Rust Clippy, Python package lint/format, Pyright and lockfile checks pass. Repository-wide Ruff reports 49 lint errors and 11 formatting failures in unchanged `docs/web-tests` files; those files are byte-identical to main and outside this change. All 1,103 Git-tracked files match the AWS validation checkout. No benchmark scripts or result artifacts are included, and no performance improvement is claimed. Uses lance-format/lance#9798. Server-side support is in lancedb/sophon#8607.
This commit is contained in:
Yang Cen
@@ -646,6 +646,38 @@ rerank(reranker): VectorQuery
|
||||
|
||||
***
|
||||
|
||||
### searchEffort()
|
||||
|
||||
```ts
|
||||
searchEffort(searchEffort): VectorQuery
|
||||
```
|
||||
|
||||
Set the initial IVF search effort to a finite value in [0, 1].
|
||||
|
||||
The default, null, uses Auto (internally 0.5). Zero starts at the caller
|
||||
minimum; one searches all partitions within the caller maximum. Later
|
||||
probing may continue when too few candidates are found.
|
||||
Any explicit effort, including 0.5, overrides nprobes in either call order.
|
||||
Explicit minimum and maximum bounds still apply. This is not a recall
|
||||
guarantee and does not change refinement or HNSW ef. The MemWAL LSM
|
||||
scanner does not yet support this option.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **searchEffort**: `null` \| `number`
|
||||
|
||||
#### Returns
|
||||
|
||||
[`VectorQuery`](VectorQuery.md)
|
||||
|
||||
#### Example
|
||||
|
||||
```ts
|
||||
table.vectorSearch([0.1, 0.2]).searchEffort(0.75).limit(10);
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### select()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -527,3 +527,106 @@ describe("Query orderBy", () => {
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("IVF search effort", () => {
|
||||
let directory: tmp.DirResult;
|
||||
let table: Table;
|
||||
const vector = Array.from({ length: 16 }, (_, i) => (i === 0 ? 1 : 0));
|
||||
|
||||
beforeEach(async () => {
|
||||
directory = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(directory.name);
|
||||
const rows = Array.from({ length: 512 }, (_, id) => ({
|
||||
id,
|
||||
vector: Array.from(
|
||||
{ length: 16 },
|
||||
(_, j) =>
|
||||
(Math.floor(id / 32) === j ? 1 : 0) + Math.sin(id * 17 + j) * 0.001,
|
||||
),
|
||||
}));
|
||||
table = await db.createTable("effort", rows.slice(0, 256));
|
||||
await table.add(rows.slice(256));
|
||||
await table.createIndex("vector", {
|
||||
config: Index.ivfFlat({ numPartitions: 16 }),
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
directory.removeCallback();
|
||||
});
|
||||
|
||||
it.each([
|
||||
[null, 6],
|
||||
[0, 1],
|
||||
[0.25, 3],
|
||||
[0.5, 6],
|
||||
[0.75, 10],
|
||||
[1, 16],
|
||||
])(
|
||||
"preserves effort %s through execution and probe precedence",
|
||||
async (effort, expected) => {
|
||||
const partitions = async (query: ReturnType<Table["vectorSearch"]>) => {
|
||||
const plan = await query.limit(8).analyzePlan();
|
||||
const matches = [...plan.matchAll(/partitions_searched=(\d+)/g)];
|
||||
expect(matches.length).toBeGreaterThan(0);
|
||||
return matches.reduce((total, match) => total + Number(match[1]), 0);
|
||||
};
|
||||
expect(
|
||||
await partitions(table.vectorSearch(vector).searchEffort(effort)),
|
||||
).toBe(expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).nprobes(1).searchEffort(effort),
|
||||
),
|
||||
).toBe(effort === null ? 1 : expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).searchEffort(effort).nprobes(1),
|
||||
),
|
||||
).toBe(effort === null ? 1 : expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table
|
||||
.vectorSearch(vector)
|
||||
.searchEffort(effort)
|
||||
.nprobes(1)
|
||||
.searchEffort(null),
|
||||
),
|
||||
).toBe(1);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).searchEffort(effort).maximumNprobes(2),
|
||||
),
|
||||
).toBeLessThanOrEqual(2);
|
||||
const actual = await table
|
||||
.vectorSearch(vector)
|
||||
.searchEffort(effort)
|
||||
.limit(8)
|
||||
.toArray();
|
||||
const exact = await table
|
||||
.vectorSearch(vector)
|
||||
.bypassVectorIndex()
|
||||
.limit(8)
|
||||
.toArray();
|
||||
const ids = new Set(exact.map((row) => row.id));
|
||||
expect(
|
||||
actual.filter((row) => ids.has(row.id)).length / 8,
|
||||
).toBeGreaterThanOrEqual(0.5);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
-0.1,
|
||||
1.1,
|
||||
Number.NaN,
|
||||
Number.POSITIVE_INFINITY,
|
||||
Number.NEGATIVE_INFINITY,
|
||||
])("rejects invalid effort %s before execution", (effort) => {
|
||||
expect(() => table.vectorSearch(vector).searchEffort(effort)).toThrow(
|
||||
"search_effort must be finite and in [0, 1]",
|
||||
);
|
||||
expect(() =>
|
||||
table.vectorSearch(vector).fullTextSearch("dog").searchEffort(effort),
|
||||
).toThrow("search_effort must be finite and in [0, 1]");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -549,6 +549,27 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the initial IVF search effort to a finite value in [0, 1].
|
||||
*
|
||||
* The default, null, uses Auto (internally 0.5). Zero starts at the caller
|
||||
* minimum; one searches all partitions within the caller maximum. Later
|
||||
* probing may continue when too few candidates are found.
|
||||
* Any explicit effort, including 0.5, overrides nprobes in either call order.
|
||||
* Explicit minimum and maximum bounds still apply. This is not a recall
|
||||
* guarantee and does not change refinement or HNSW ef. The MemWAL LSM
|
||||
* scanner does not yet support this option.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* table.vectorSearch([0.1, 0.2]).searchEffort(0.75).limit(10);
|
||||
* ```
|
||||
*/
|
||||
searchEffort(searchEffort: number | null): VectorQuery {
|
||||
this.doVectorCall((inner) => inner.searchEffort(searchEffort));
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the minimum number of probes used.
|
||||
*
|
||||
|
||||
@@ -298,6 +298,16 @@ impl VectorQuery {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub fn search_effort(&mut self, search_effort: Option<f64>) -> napi::Result<()> {
|
||||
self.inner = self
|
||||
.inner
|
||||
.clone()
|
||||
.search_effort(search_effort)
|
||||
.default_error()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub fn minimum_nprobes(&mut self, minimum_nprobe: u32) -> napi::Result<()> {
|
||||
self.inner = self
|
||||
|
||||
@@ -63,7 +63,7 @@ tests = [
|
||||
"polars>=0.19, <=1.32.3",
|
||||
"pyarrow<25",
|
||||
"pyarrow-stubs>=16.0",
|
||||
"pylance==14.0.0b10",
|
||||
"pylance==14.0.0b12",
|
||||
"requests>=2.31.0",
|
||||
"datafusion>=54,<55",
|
||||
"opentelemetry-sdk>=1.30.0",
|
||||
@@ -73,6 +73,7 @@ dev = [
|
||||
"pre-commit>=3.5.0",
|
||||
"pyright>=1.1.350",
|
||||
'typing-extensions>=4.0.0; python_version < "3.11"',
|
||||
"maturin>=1.10",
|
||||
]
|
||||
docs = ["mkdocs", "mkdocs-jupyter", "mkdocs-material", "mkdocstrings-python"]
|
||||
clip = ["torch", "pillow>=12.1.1", "open-clip-torch"]
|
||||
|
||||
@@ -727,6 +727,7 @@ class VectorQuery:
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def refine_factor(self, refine_factor: int): ...
|
||||
def nprobes(self, nprobes: int): ...
|
||||
def search_effort(self, search_effort: Optional[float]): ...
|
||||
def minimum_nprobes(self, minimum_nprobes: int): ...
|
||||
def maximum_nprobes(self, maximum_nprobes: int): ...
|
||||
def bypass_vector_index(self): ...
|
||||
@@ -752,6 +753,7 @@ class HybridQuery:
|
||||
def distance_type(self, distance_type: str): ...
|
||||
def refine_factor(self, refine_factor: int): ...
|
||||
def nprobes(self, nprobes: int): ...
|
||||
def search_effort(self, search_effort: Optional[float]): ...
|
||||
def minimum_nprobes(self, minimum_nprobes: int): ...
|
||||
def maximum_nprobes(self, maximum_nprobes: int): ...
|
||||
def bypass_vector_index(self): ...
|
||||
@@ -781,6 +783,7 @@ class PyQueryRequest:
|
||||
nprobes: Optional[int]
|
||||
minimum_nprobes: Optional[int]
|
||||
maximum_nprobes: Optional[int]
|
||||
search_effort: Optional[float]
|
||||
lower_bound: Optional[float]
|
||||
upper_bound: Optional[float]
|
||||
ef: Optional[int]
|
||||
|
||||
@@ -105,6 +105,11 @@ def _query_to_namespace_request(
|
||||
|
||||
if not isinstance(query, Query):
|
||||
raise TypeError(f"Expected Query, got {type(query)}")
|
||||
if query.search_effort is not None:
|
||||
raise ValueError(
|
||||
"namespace QueryTable does not support search_effort; "
|
||||
"execute the query locally"
|
||||
)
|
||||
|
||||
vector_request = None
|
||||
if query.vector is not None:
|
||||
|
||||
@@ -475,8 +475,8 @@ class MatchQuery(FullTextQuery):
|
||||
----------
|
||||
query : str
|
||||
The query string to match against.
|
||||
column : str
|
||||
The name of the column to match against.
|
||||
column : str, optional
|
||||
The column to match against. None leaves column selection to query planning.
|
||||
boost : float, default 1.0
|
||||
The boost factor for the query.
|
||||
The score of each matching document is multiplied by this value.
|
||||
@@ -505,7 +505,7 @@ class MatchQuery(FullTextQuery):
|
||||
"""
|
||||
|
||||
query: str
|
||||
column: str
|
||||
column: Optional[str] = None
|
||||
boost: float = pydantic.Field(1.0, kw_only=True)
|
||||
fuzziness: Optional[int] = pydantic.Field(0, kw_only=True)
|
||||
max_expansions: int = pydantic.Field(50, kw_only=True)
|
||||
@@ -793,6 +793,9 @@ class Query(pydantic.BaseModel):
|
||||
# and all partitions could be searched if needed to satisfy the limit.
|
||||
maximum_nprobes: Optional[int] = None
|
||||
|
||||
# Optional initial IVF effort; None preserves Auto and explicit nprobes.
|
||||
search_effort: Optional[float] = None
|
||||
|
||||
# lower bound for distance search
|
||||
lower_bound: Optional[float] = None
|
||||
|
||||
@@ -854,6 +857,7 @@ class Query(pydantic.BaseModel):
|
||||
query.nprobes = req.nprobes
|
||||
query.minimum_nprobes = req.minimum_nprobes
|
||||
query.maximum_nprobes = req.maximum_nprobes
|
||||
query.search_effort = req.search_effort
|
||||
query.lower_bound = req.lower_bound
|
||||
query.upper_bound = req.upper_bound
|
||||
query.ef = req.ef
|
||||
@@ -1635,6 +1639,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
self._nprobes = None
|
||||
self._minimum_nprobes = None
|
||||
self._maximum_nprobes = None
|
||||
self._search_effort = None
|
||||
self._lower_bound = None
|
||||
self._upper_bound = None
|
||||
self._refine_factor = None
|
||||
@@ -1716,6 +1721,27 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
self._nprobes = nprobes
|
||||
return self
|
||||
|
||||
def search_effort(self, search_effort: Optional[float]) -> LanceVectorQueryBuilder:
|
||||
"""Set the initial IVF search effort to a finite value in [0, 1].
|
||||
|
||||
None (the default) uses Auto, internally 0.5. Zero starts at the caller
|
||||
minimum; one searches all partitions within the caller maximum. Later
|
||||
probing may continue when too few candidates are found. Any explicit
|
||||
effort, including 0.5, overrides nprobes in either call order. Explicit
|
||||
minimum and maximum bounds still apply. This is not a recall guarantee
|
||||
and does not change refinement or HNSW ef. MemWAL LSM queries do not
|
||||
yet support this option.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> db = lancedb.connect("memory://")
|
||||
>>> table = db.create_table("effort_vectors", [{"vector": [0.1, 0.2]}])
|
||||
>>> query = table.search([0.1, 0.2]).search_effort(0.75)
|
||||
"""
|
||||
self._search_effort = search_effort
|
||||
return self
|
||||
|
||||
def minimum_nprobes(self, minimum_nprobes: int) -> LanceVectorQueryBuilder:
|
||||
"""Set the minimum number of probes to use.
|
||||
|
||||
@@ -1858,6 +1884,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
nprobes=self._nprobes,
|
||||
minimum_nprobes=self._minimum_nprobes,
|
||||
maximum_nprobes=self._maximum_nprobes,
|
||||
search_effort=self._search_effort,
|
||||
lower_bound=self._lower_bound,
|
||||
upper_bound=self._upper_bound,
|
||||
refine_factor=self._refine_factor,
|
||||
@@ -2217,6 +2244,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._nprobes = None
|
||||
self._minimum_nprobes = None
|
||||
self._maximum_nprobes = None
|
||||
self._search_effort = None
|
||||
self._refine_factor = None
|
||||
self._distance_type = None
|
||||
self._phrase_query = None
|
||||
@@ -2524,6 +2552,32 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._nprobes = nprobes
|
||||
return self
|
||||
|
||||
def search_effort(self, search_effort: Optional[float]) -> LanceHybridQueryBuilder:
|
||||
"""Set the initial IVF search effort to a finite value in [0, 1].
|
||||
|
||||
None (the default) uses Auto, internally 0.5. Zero starts at the caller
|
||||
minimum; one searches all partitions within the caller maximum. Later
|
||||
probing may continue when too few candidates are found. Any explicit
|
||||
effort, including 0.5, overrides nprobes in either call order. Explicit
|
||||
minimum and maximum bounds still apply. This is not a recall guarantee
|
||||
and does not change refinement or HNSW ef. MemWAL LSM queries do not
|
||||
yet support this option.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> db = lancedb.connect("memory://")
|
||||
>>> table = db.create_table(
|
||||
... "effort_hybrid", [{"vector": [0.1, 0.2], "text": "dog"}]
|
||||
... )
|
||||
>>> query = (
|
||||
... table.search(query_type="hybrid").vector([0.1, 0.2])
|
||||
... .text("dog").search_effort(0.75)
|
||||
... )
|
||||
"""
|
||||
self._search_effort = search_effort
|
||||
return self
|
||||
|
||||
def minimum_nprobes(self, minimum_nprobes: int) -> LanceHybridQueryBuilder:
|
||||
"""Set the minimum number of probes to use.
|
||||
|
||||
@@ -2802,6 +2856,8 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._vector_query.metric(self._distance_type)
|
||||
if self._nprobes is not None:
|
||||
self._vector_query.nprobes(self._nprobes)
|
||||
if self._search_effort is not None:
|
||||
self._vector_query.search_effort(self._search_effort)
|
||||
if self._minimum_nprobes is not None:
|
||||
self._vector_query.minimum_nprobes(self._minimum_nprobes)
|
||||
if self._maximum_nprobes is not None:
|
||||
@@ -3730,6 +3786,32 @@ class AsyncVectorQueryBase:
|
||||
self._inner.nprobes(nprobes)
|
||||
return self
|
||||
|
||||
def search_effort(self, search_effort: Optional[float]) -> Self:
|
||||
"""Set the initial IVF search effort to a finite value in [0, 1].
|
||||
|
||||
None preserves Auto (internally 0.5) and fixed nprobes. Any explicit
|
||||
effort, including 0.5, overrides nprobes in either call order; explicit
|
||||
minimum and maximum bounds remain effective. Zero starts at the caller
|
||||
minimum and one searches all partitions within the maximum. Later
|
||||
probing may continue to obtain enough candidates. This is not a recall
|
||||
guarantee and does not change refinement or HNSW ef. MemWAL LSM queries
|
||||
do not yet support this option.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import asyncio
|
||||
>>> import lancedb
|
||||
>>> async def build_query():
|
||||
... db = await lancedb.connect_async("memory://")
|
||||
... table = await db.create_table(
|
||||
... "effort_async", [{"vector": [0.1, 0.2]}]
|
||||
... )
|
||||
... return table.vector_search([0.1, 0.2]).search_effort(0.75)
|
||||
>>> query = asyncio.run(build_query())
|
||||
"""
|
||||
self._inner.search_effort(search_effort)
|
||||
return self
|
||||
|
||||
def minimum_nprobes(self, minimum_nprobes: int) -> Self:
|
||||
"""Set the minimum number of probes to use.
|
||||
|
||||
|
||||
@@ -4440,6 +4440,7 @@ class LanceTable(Table):
|
||||
and not self._route_pushdown_to_rust
|
||||
and self.current_branch() is None
|
||||
and query.take_offsets is None
|
||||
and query.search_effort is None
|
||||
):
|
||||
from lancedb.namespace import _execute_server_side_query
|
||||
|
||||
@@ -6441,6 +6442,8 @@ class AsyncTable:
|
||||
async_query = async_query.distance_type(query.distance_type)
|
||||
if query.nprobes is not None:
|
||||
async_query = async_query.nprobes(query.nprobes)
|
||||
if query.search_effort is not None:
|
||||
async_query = async_query.search_effort(query.search_effort)
|
||||
if query.minimum_nprobes is not None and query.maximum_nprobes is not None:
|
||||
async_query = async_query.minimum_nprobes(
|
||||
query.minimum_nprobes
|
||||
@@ -6481,6 +6484,7 @@ class AsyncTable:
|
||||
)
|
||||
and not self._route_pushdown_to_rust
|
||||
and query.take_offsets is None
|
||||
and query.search_effort is None
|
||||
):
|
||||
from lancedb.namespace import _execute_server_side_query
|
||||
|
||||
|
||||
@@ -1209,3 +1209,15 @@ 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]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("effort", [0.0, 0.5, 1.0])
|
||||
def test_namespace_request_rejects_search_effort(effort):
|
||||
from lancedb.query import Query
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="namespace QueryTable does not support search_effort"
|
||||
):
|
||||
_query_to_namespace_request(
|
||||
["test"], Query(vector=[1.0, 0.0], search_effort=effort)
|
||||
)
|
||||
|
||||
@@ -2265,3 +2265,127 @@ def test_blob_v2_to_batches_row_id(tmp_db):
|
||||
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"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"])
|
||||
@pytest.mark.parametrize(
|
||||
"effort,probes", [(None, 4), (0, 1), (0.25, 2), (0.5, 4), (0.75, 8), (1, 16)]
|
||||
)
|
||||
def test_search_effort_probe_budget(tmp_path, monkeypatch, metric, effort, probes):
|
||||
import re
|
||||
from lancedb.index import IvfFlat
|
||||
|
||||
monkeypatch.setenv("LANCE_AUTO_PROBE_MARGIN", "0")
|
||||
monkeypatch.setenv("LANCE_AUTO_MIN_INITIAL_NPROBES", "4")
|
||||
monkeypatch.setenv("LANCE_AUTO_MAX_INITIAL_NPROBES", "4")
|
||||
rng = np.random.default_rng(42)
|
||||
vectors = np.repeat(np.eye(16, dtype=np.float32), 32, axis=0)
|
||||
vectors += rng.normal(0, 0.001, vectors.shape).astype(np.float32)
|
||||
data = pa.table(
|
||||
{
|
||||
"id": np.arange(512),
|
||||
"text": ["dog"] * 512,
|
||||
"vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 16),
|
||||
}
|
||||
)
|
||||
table = lancedb.connect(tmp_path).create_table("effort", data.slice(0, 256))
|
||||
table.add(data.slice(256))
|
||||
table.create_index(
|
||||
"vector", config=IvfFlat(distance_type=metric, num_partitions=16)
|
||||
)
|
||||
vector = np.eye(16, dtype=np.float32)[0]
|
||||
|
||||
def search(value=effort):
|
||||
return table.search(vector).distance_type(metric).search_effort(value).limit(8)
|
||||
|
||||
def partitions(query):
|
||||
vector_plan = query.analyze_plan().split("FTS Search Plan:", 1)[0]
|
||||
values = re.findall(r"partitions_searched=(\d+)", vector_plan)
|
||||
assert values
|
||||
return sum(map(int, values))
|
||||
|
||||
assert partitions(search()) == probes
|
||||
exact = set(
|
||||
table.search(vector)
|
||||
.distance_type(metric)
|
||||
.bypass_vector_index()
|
||||
.limit(8)
|
||||
.to_arrow()["id"]
|
||||
.to_pylist()
|
||||
)
|
||||
actual = set(search().to_arrow()["id"].to_pylist())
|
||||
assert len(actual & exact) / 8 >= 0.5
|
||||
expected = 1 if effort is None else probes
|
||||
assert partitions(search().nprobes(1)) == expected
|
||||
assert (
|
||||
partitions(
|
||||
table.search(vector)
|
||||
.distance_type(metric)
|
||||
.nprobes(1)
|
||||
.search_effort(effort)
|
||||
.limit(8)
|
||||
)
|
||||
== expected
|
||||
)
|
||||
assert partitions(search().nprobes(1).search_effort(None)) == 1
|
||||
assert partitions(search().maximum_nprobes(2)) <= 2
|
||||
assert partitions(search().minimum_nprobes(3).maximum_nprobes(3)) == 3
|
||||
assert partitions(search(0).nprobes(16).maximum_nprobes(0)) == 1
|
||||
if effort is None:
|
||||
assert search().to_arrow() == search(0.5).to_arrow()
|
||||
filtered = search(0).maximum_nprobes(2).where("id % 8 = 0", prefilter=True)
|
||||
assert partitions(filtered) <= 2
|
||||
assert all(row["id"] % 8 == 0 for row in filtered.to_list())
|
||||
|
||||
if metric == "l2" and effort == 0.5:
|
||||
table.create_index("text", config=FTS())
|
||||
hybrid = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector(vector)
|
||||
.text("dog")
|
||||
.nprobes(1)
|
||||
.search_effort(0.5)
|
||||
.limit(8)
|
||||
)
|
||||
assert partitions(hybrid) == 4
|
||||
assert partitions(hybrid.search_effort(None)) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
|
||||
)
|
||||
def test_search_effort_invalid_sync(table, effort):
|
||||
with pytest.raises(ValueError, match="search_effort must be finite and in"):
|
||||
table.search([0, 0]).search_effort(effort).to_list()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hybrid", [False, True])
|
||||
@pytest.mark.parametrize("effort", [None, 0, 0.25, 0.5, 0.75, 1])
|
||||
async def test_search_effort_async_roundtrip(table_async, hybrid, effort):
|
||||
query = table_async.vector_search([0, 0])
|
||||
if hybrid:
|
||||
query = query.nearest_to_text("dog")
|
||||
query = query.nprobes(5).search_effort(effort).minimum_nprobes(2).maximum_nprobes(3)
|
||||
request = query.to_query_object()
|
||||
assert request.search_effort == effort
|
||||
assert request.nprobes == 5
|
||||
assert request.minimum_nprobes == 2
|
||||
assert request.maximum_nprobes == 3
|
||||
assert query.search_effort(None).to_query_object().search_effort is None
|
||||
query.search_effort(effort)
|
||||
if not hybrid:
|
||||
assert len(await query.limit(1).to_list()) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hybrid", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
|
||||
)
|
||||
async def test_search_effort_invalid_async(table_async, hybrid, effort):
|
||||
query = table_async.vector_search([0, 0])
|
||||
if hybrid:
|
||||
query = query.nearest_to_text("dog")
|
||||
with pytest.raises(ValueError, match="search_effort must be finite and in"):
|
||||
query.search_effort(effort)
|
||||
|
||||
+31
-3
@@ -71,7 +71,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
|
||||
match class_name(&ob)?.as_str() {
|
||||
"MatchQuery" => {
|
||||
let query = ob.getattr("query")?.extract()?;
|
||||
let column = ob.getattr("column")?.extract()?;
|
||||
let column = ob.getattr("column")?.extract::<Option<String>>()?;
|
||||
let boost = ob.getattr("boost")?.extract()?;
|
||||
let fuzziness = ob.getattr("fuzziness")?.extract()?;
|
||||
let max_expansions = ob.getattr("max_expansions")?.extract()?;
|
||||
@@ -88,7 +88,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
|
||||
|
||||
let mut query =
|
||||
MatchQuery::new(query)
|
||||
.with_column(Some(column))
|
||||
.with_column(column)
|
||||
.with_boost(boost)
|
||||
.with_fuzziness(fuzziness)
|
||||
.with_max_expansions(max_expansions)
|
||||
@@ -199,7 +199,7 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
|
||||
}
|
||||
namespace
|
||||
.getattr(intern!(py, "MatchQuery"))?
|
||||
.call((query.terms, query.column.unwrap()), Some(&kwargs))
|
||||
.call((query.terms, query.column), Some(&kwargs))
|
||||
}
|
||||
FtsQuery::Phrase(query) => {
|
||||
let kwargs = PyDict::new(py);
|
||||
@@ -339,6 +339,7 @@ pub struct PyQueryRequest {
|
||||
pub nprobes: Option<usize>,
|
||||
pub minimum_nprobes: Option<usize>,
|
||||
pub maximum_nprobes: Option<usize>,
|
||||
pub search_effort: Option<f64>,
|
||||
pub lower_bound: Option<f32>,
|
||||
pub upper_bound: Option<f32>,
|
||||
pub ef: Option<usize>,
|
||||
@@ -371,6 +372,7 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
nprobes: None,
|
||||
minimum_nprobes: None,
|
||||
maximum_nprobes: None,
|
||||
search_effort: None,
|
||||
lower_bound: None,
|
||||
upper_bound: None,
|
||||
ef: None,
|
||||
@@ -399,6 +401,7 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
nprobes: vector_query.nprobes,
|
||||
minimum_nprobes: vector_query.minimum_nprobes,
|
||||
maximum_nprobes: vector_query.maximum_nprobes,
|
||||
search_effort: vector_query.search_effort,
|
||||
lower_bound: vector_query.lower_bound,
|
||||
upper_bound: vector_query.upper_bound,
|
||||
ef: vector_query.ef,
|
||||
@@ -998,6 +1001,15 @@ impl VectorQuery {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn search_effort(&mut self, search_effort: Option<f64>) -> PyResult<()> {
|
||||
self.inner = self
|
||||
.inner
|
||||
.clone()
|
||||
.search_effort(search_effort)
|
||||
.infer_error()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
|
||||
self.inner = self
|
||||
.inner
|
||||
@@ -1191,6 +1203,18 @@ impl HybridQuery {
|
||||
self.inner_vec.nprobes(nprobe)
|
||||
}
|
||||
|
||||
pub fn search_effort(&mut self, search_effort: Option<f64>) -> PyResult<()> {
|
||||
self.inner_vec.search_effort(search_effort)
|
||||
}
|
||||
|
||||
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
|
||||
self.inner_vec.minimum_nprobes(minimum_nprobes)
|
||||
}
|
||||
|
||||
pub fn maximum_nprobes(&mut self, maximum_nprobes: u32) -> PyResult<()> {
|
||||
self.inner_vec.maximum_nprobes(maximum_nprobes)
|
||||
}
|
||||
|
||||
pub fn ef(&mut self, ef: u32) {
|
||||
self.inner_vec.ef(ef);
|
||||
}
|
||||
@@ -1235,6 +1259,10 @@ impl HybridQuery {
|
||||
req.query_vector = vec_req.query_vector;
|
||||
req.column = vec_req.column;
|
||||
req.distance_type = vec_req.distance_type;
|
||||
req.nprobes = vec_req.nprobes;
|
||||
req.minimum_nprobes = vec_req.minimum_nprobes;
|
||||
req.maximum_nprobes = vec_req.maximum_nprobes;
|
||||
req.search_effort = vec_req.search_effort;
|
||||
req.ef = vec_req.ef;
|
||||
req.refine_factor = vec_req.refine_factor;
|
||||
req.lower_bound = vec_req.lower_bound;
|
||||
|
||||
Generated
+34
-8
@@ -1900,6 +1900,7 @@ clip = [
|
||||
{ name = "torch" },
|
||||
]
|
||||
dev = [
|
||||
{ name = "maturin" },
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pyright" },
|
||||
{ name = "ruff" },
|
||||
@@ -1981,6 +1982,7 @@ requires-dist = [
|
||||
{ name = "ibm-watsonx-ai", marker = "python_full_version >= '3.10' and extra == 'embeddings'", specifier = ">=1.1.2" },
|
||||
{ name = "instructorembedding", marker = "extra == 'embeddings'", specifier = ">=1.0.1" },
|
||||
{ name = "lance-namespace", specifier = ">=0.13.0" },
|
||||
{ name = "maturin", marker = "extra == 'dev'", specifier = ">=1.10" },
|
||||
{ name = "mkdocs", marker = "extra == 'docs'" },
|
||||
{ name = "mkdocs-jupyter", marker = "extra == 'docs'" },
|
||||
{ name = "mkdocs-material", marker = "extra == 'docs'" },
|
||||
@@ -2005,7 +2007,7 @@ requires-dist = [
|
||||
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
|
||||
{ name = "pydantic", specifier = ">=2.7.4,<3" },
|
||||
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b10" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" },
|
||||
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },
|
||||
{ name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" },
|
||||
{ name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" },
|
||||
@@ -2156,6 +2158,30 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl", hash = "sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6", size = 9534, upload-time = "2026-05-08T17:33:32.055Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "maturin"
|
||||
version = "1.15.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "tomli", marker = "python_full_version < '3.11'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b9/c8/22e5e21b2679c9bce6415ca578034ca2cc9316be0642ae21e051a2d5198c/maturin-1.15.0.tar.gz", hash = "sha256:94b26cc8e8aba61a5f2099715fe640e18c5f678e9a500408b38761263954228a", size = 385504, upload-time = "2026-08-24T12:11:22.665Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/14/69/5c01b461044eb1f45ddcce006706eb88110c793cdb11c7ae0b5e08492e94/maturin-1.15.0-py3-none-linux_armv6l.whl", hash = "sha256:6bf6dc62e22d4dcfd5a51244ff0d58975fa4979c48209fe84159617648956d82", size = 10206220, upload-time = "2026-08-24T12:10:53.327Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/eb/1f/2b431554e11687cdb1077e0cdadcc118c53f611086b3af00c8545a67c6a5/maturin-1.15.0-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:cd35772633f489841132bc8e71d6fc7f842df30b9c05cd5cdf1ee1ddcb744cc7", size = 19416513, upload-time = "2026-08-24T12:10:56.126Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/51/36/e23a21cb34a648b711036b9b2fe1d4f3f4ee24f8db54215d73f1a9a3a3ec/maturin-1.15.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:c40b4eae7bf5ef1f4b1af8d623fe4105016f93578fb15b764e741d08ec3b92dd", size = 10014962, upload-time = "2026-08-24T12:10:58.486Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8c/80/33b15cb2d8f30f12c807955e8f2fd775027692904e30ec0784744ce8cd83/maturin-1.15.0-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:7eb066372f541f8eb4909c79c5d9bd0b9e8125980bdf1ec9e8aba23c6c8d6c55", size = 10196223, upload-time = "2026-08-24T12:11:00.696Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/fe/91/b495e19e2f5c503b540452b2039115e7b2363867e8c5ad4179eb752fa92c/maturin-1.15.0-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:653020a63525bb224e5ab0adf02e17a2e08bc86dbea7fc1399c9a56d7529b99e", size = 10541186, upload-time = "2026-08-24T12:11:02.857Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5d/2b/2abff58037188d852b124871b1f0d720e1c2bfb3d4f1b03d87c52cd66488/maturin-1.15.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:0ebf9767892725083138e671c34482c660317a2f3d6a29fc0e0f34e9d8c99136", size = 10083468, upload-time = "2026-08-24T12:11:05.012Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/3f/07/b7e9f8be99a6627849e81ac7b6694876bce8f50a92995fe17e3cf2610f0a/maturin-1.15.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:7ab7eebffd7b8debca2265985de4eaeb332141276d24b9560b5ad484d4b3add1", size = 10047786, upload-time = "2026-08-24T12:11:07.042Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ca/ab/167e3cb7accee11b507dbe53e0e87aeccb376d44ae66284c96ee4df3a9fd/maturin-1.15.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:126e12e618b4db42f68c779a56d41f82a390145ba36ac3f621d057eb34f5ad9d", size = 13315332, upload-time = "2026-08-24T12:11:09.433Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/14/4d/801379f646cbc6b00998e5289b0630a886be3a4ee4c75b6bc9b87478a7f1/maturin-1.15.0-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4f9d33e6c3f9615c8caceecbbbd440f8eb25a3ddeb687077682cd5eca2e9ae15", size = 10807183, upload-time = "2026-08-24T12:11:11.73Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/89/27/2e612e1cbd1580e9e94d4722c227b5180dca27b32b955a34b79918aa1292/maturin-1.15.0-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:bf29beddd0c6708f112db51d5275fc28b28b9e9c9c5faae387eaef662918b176", size = 10413274, upload-time = "2026-08-24T12:11:14.04Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/70/d8/202a7b4d75a51f20f84ec9ce3b7345b12b822207072164e2b1c6ef665125/maturin-1.15.0-py3-none-win32.whl", hash = "sha256:da649988be98e87e009e51b1bf0d301b6a301bc0cecbdd60d40d8ba60748d1ca", size = 8928744, upload-time = "2026-08-24T12:11:16.269Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/40/dc/4e90da594986ba78dd3bc8a5921ecdcdb11085b22b02a412caab3b225601/maturin-1.15.0-py3-none-win_amd64.whl", hash = "sha256:552c2be4afd43fe8d5c9f3ec8d4c4756d973b8dcbe94c14084390301f50243e1", size = 10335085, upload-time = "2026-08-24T12:11:18.326Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8b/10/15d4314edf130955edf2dc237aa393a8a7c10f2b9b57b89fa2f61f915659/maturin-1.15.0-py3-none-win_arm64.whl", hash = "sha256:c7dc0c66c78d3debdd9c5aa807e861fbcbf07f3505d34b125df74c03986b0f48", size = 9713795, upload-time = "2026-08-24T12:11:20.83Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mdit-py-plugins"
|
||||
version = "0.6.1"
|
||||
@@ -3912,7 +3938,7 @@ crypto = [
|
||||
|
||||
[[package]]
|
||||
name = "pylance"
|
||||
version = "14.0.0b10"
|
||||
version = "14.0.0b12"
|
||||
source = { registry = "https://pypi.fury.io/lance-format/" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace" },
|
||||
@@ -3922,12 +3948,12 @@ dependencies = [
|
||||
{ name = "pyarrow" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_4hvE/pylance-14.0.0b10-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:415b6b1b9fe735fbae7c755eb9774ae2fe0feddc90bf80bee24087cbfbfdf8c1" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_pM1km/pylance-14.0.0b10-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:87317cad22f4d35fea11eb70aa5c0d71a0e1ea3b96de38f3deca933daa749fff" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_24jdyj/pylance-14.0.0b10-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6da554a2fc2197ba254676b12824d8120b48d65937697e0661a61302458a4740" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_XclaZ/pylance-14.0.0b10-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:36c80a801675bb9dcbd681fa99d0cdd5ca642dfb3ca08d71ac980278a2f7748b" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_L0yh/pylance-14.0.0b10-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:e50ad0b866faefb2fa3ae87328b93fa0d06b2bdd5bb5d0c1a6c74b907cc2fd93" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2gbNr8/pylance-14.0.0b10-cp310-abi3-win_amd64.whl", hash = "sha256:3d86451fe00f71510005bb700737fb1a4f9361554541cbf802f3c5072554d014" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -1127,6 +1127,11 @@ pub struct VectorQueryRequest {
|
||||
/// If unset, Lance's default is used. Zero represents an explicitly unbounded
|
||||
/// maximum on request and wire representations.
|
||||
pub maximum_nprobes: Option<usize>,
|
||||
/// Optional initial IVF search effort in [0, 1].
|
||||
///
|
||||
/// None preserves Auto and fixed `nprobes`; any explicit value overrides
|
||||
/// `nprobes` while retaining explicit minimum and maximum bounds.
|
||||
pub search_effort: Option<f64>,
|
||||
/// The lower bound (inclusive) of the distance to search for.
|
||||
pub lower_bound: Option<f32>,
|
||||
/// The upper bound (exclusive) of the distance to search for.
|
||||
@@ -1153,6 +1158,7 @@ impl Default for VectorQueryRequest {
|
||||
nprobes: None,
|
||||
minimum_nprobes: None,
|
||||
maximum_nprobes: None,
|
||||
search_effort: None,
|
||||
lower_bound: None,
|
||||
upper_bound: None,
|
||||
ef: None,
|
||||
@@ -1165,6 +1171,17 @@ impl Default for VectorQueryRequest {
|
||||
}
|
||||
|
||||
impl VectorQueryRequest {
|
||||
pub(crate) fn validate_search_effort(&self) -> Result<()> {
|
||||
if let Some(effort) = self.search_effort
|
||||
&& (!effort.is_finite() || !(0.0..=1.0).contains(&effort))
|
||||
{
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("search_effort must be finite and in [0, 1], got {effort}"),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn add_query_vector(&mut self, vector: impl IntoQueryVector) -> Result<()> {
|
||||
let vector = vector.to_query_vector(&DataType::Float32, "default")?;
|
||||
let array =
|
||||
@@ -1279,7 +1296,8 @@ impl VectorQuery {
|
||||
/// you the desired recall.
|
||||
///
|
||||
/// LanceDB retains this as `nprobes` through local and remote request
|
||||
/// construction. Lance sets both probe bounds to this value. For more
|
||||
/// construction. Lance sets both probe bounds to this value only when
|
||||
/// [`Self::search_effort`] is unset. For more
|
||||
/// fine-grained control see [`VectorQuery::minimum_nprobes`] and
|
||||
/// [`VectorQuery::maximum_nprobes`].
|
||||
///
|
||||
@@ -1294,6 +1312,31 @@ impl VectorQuery {
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Set the initial IVF search effort to a finite value in [0, 1].
|
||||
///
|
||||
/// None (the default) resolves to 0.5, preserving Auto probing. At 0 the
|
||||
/// search starts at the caller minimum; at 1 it searches all partitions
|
||||
/// within the caller maximum. Intermediate values interpolate around Auto.
|
||||
/// Later probing may continue when too few candidates are found.
|
||||
/// Any explicit value, including 0.5, overrides [`Self::nprobes`] regardless
|
||||
/// of setter order. Explicit minimum and maximum bounds still apply.
|
||||
/// This does not change refinement or HNSW ef and is not a recall guarantee.
|
||||
/// Lance's MemWAL LSM scanner does not yet support this option.
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::{Table, query::{QueryBase, ExecutableQuery}};
|
||||
/// # async fn example(table: &Table) -> lancedb::Result<()> {
|
||||
/// let results = table.query().nearest_to([0.1, 0.2].as_slice())?
|
||||
/// .search_effort(0.75)?.execute().await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn search_effort(mut self, search_effort: impl Into<Option<f64>>) -> Result<Self> {
|
||||
self.request.search_effort = search_effort.into();
|
||||
self.request.validate_search_effort()?;
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Set the minimum number of partitions to search
|
||||
///
|
||||
/// This argument is only used when the vector column has an IVF PQ index.
|
||||
@@ -1333,6 +1376,8 @@ impl VectorQuery {
|
||||
///
|
||||
/// If this value is greater than minimum_nprobes then the excess partitions will
|
||||
/// only be searched if the initial search does not return enough results.
|
||||
/// The bound applies per index segment to initial and later probing without
|
||||
/// disabling Auto probing.
|
||||
///
|
||||
/// This can be useful when there is a narrow filter to allow these queries to
|
||||
/// spend more time searching and avoid potential false negatives.
|
||||
@@ -2448,6 +2493,26 @@ mod tests {
|
||||
assert_eq!(query.request.nprobes, None);
|
||||
assert_eq!(query.request.minimum_nprobes, None);
|
||||
assert_eq!(query.request.maximum_nprobes, None);
|
||||
assert_eq!(query.request.search_effort, None);
|
||||
|
||||
let effort_query = query
|
||||
.clone()
|
||||
.nprobes(20)
|
||||
.unwrap()
|
||||
.search_effort(0.5)
|
||||
.unwrap();
|
||||
assert_eq!(effort_query.request.search_effort, Some(0.5));
|
||||
assert_eq!(effort_query.request.nprobes, Some(20));
|
||||
assert_eq!(effort_query.request.minimum_nprobes, None);
|
||||
assert_eq!(effort_query.request.maximum_nprobes, None);
|
||||
assert_eq!(
|
||||
effort_query
|
||||
.search_effort(None)
|
||||
.unwrap()
|
||||
.request
|
||||
.search_effort,
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
*query
|
||||
.request
|
||||
|
||||
@@ -1137,6 +1137,7 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
mut body: serde_json::Value,
|
||||
query: &VectorQueryRequest,
|
||||
) -> Result<Vec<serde_json::Value>> {
|
||||
query.validate_search_effort()?;
|
||||
self.apply_query_params(&mut body, &query.base)?;
|
||||
|
||||
// Apply general parameters, before we dispatch based on number of query vectors.
|
||||
@@ -1155,6 +1156,9 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
if let Some(maximum_nprobes) = query.maximum_nprobes {
|
||||
body["maximum_nprobes"] = maximum_nprobes.into();
|
||||
}
|
||||
if let Some(search_effort) = query.search_effort {
|
||||
body["search_effort"] = search_effort.into();
|
||||
}
|
||||
body["lower_bound"] = query.lower_bound.into();
|
||||
body["upper_bound"] = query.upper_bound.into();
|
||||
body["ef"] = query.ef.into();
|
||||
@@ -6269,6 +6273,77 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unset(None)]
|
||||
#[case::minimum(Some(0.0))]
|
||||
#[case::low(Some(0.25))]
|
||||
#[case::auto(Some(0.5))]
|
||||
#[case::high(Some(0.75))]
|
||||
#[case::maximum(Some(1.0))]
|
||||
#[tokio::test]
|
||||
async fn test_query_vector_search_effort_wire(#[case] effort: Option<f64>) {
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
let body = request_body_json(&request);
|
||||
assert_eq!(body["nprobes"], 7);
|
||||
assert_eq!(body["minimum_nprobes"], 2);
|
||||
assert_eq!(body["maximum_nprobes"], 8);
|
||||
if let Some(effort) = effort {
|
||||
assert_eq!(body["search_effort"], effort);
|
||||
} else {
|
||||
assert!(body.get("search_effort").is_none());
|
||||
}
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])),
|
||||
vec![Arc::new(Int32Array::from(vec![1]))],
|
||||
)
|
||||
.unwrap();
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
|
||||
.body(write_ipc_file(&batch))
|
||||
.unwrap()
|
||||
});
|
||||
let batches = table
|
||||
.query()
|
||||
.nearest_to(vec![0.1, 0.2])
|
||||
.unwrap()
|
||||
.search_effort(effort)
|
||||
.unwrap()
|
||||
.nprobes(7)
|
||||
.unwrap()
|
||||
.minimum_nprobes(2)
|
||||
.unwrap()
|
||||
.maximum_nprobes(Some(8))
|
||||
.unwrap()
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches[0].num_rows(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::negative(-0.1)]
|
||||
#[case::too_large(1.1)]
|
||||
#[case::nan(f64::NAN)]
|
||||
#[case::infinity(f64::INFINITY)]
|
||||
#[case::negative_infinity(f64::NEG_INFINITY)]
|
||||
#[tokio::test]
|
||||
async fn test_query_vector_search_effort_invalid(#[case] effort: f64) {
|
||||
let table = Table::new_with_handler::<&str>("my_table", |_| {
|
||||
panic!("invalid search_effort must not send a request")
|
||||
});
|
||||
let result = table
|
||||
.query()
|
||||
.nearest_to(vec![0.1, 0.2])
|
||||
.unwrap()
|
||||
.search_effort(effort);
|
||||
assert!(matches!(result, Err(Error::InvalidInput { message })
|
||||
if message.contains("search_effort must be finite and in [0, 1]")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_vector_all_params() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
|
||||
@@ -106,7 +106,7 @@ async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> R
|
||||
}
|
||||
|
||||
fn requires_local_namespace_execution(query: &AnyQuery) -> bool {
|
||||
// The namespace QueryTable request cannot represent approx_mode or use_lsm.
|
||||
// The namespace QueryTable request cannot represent approx_mode, search_effort or use_lsm.
|
||||
// For use_lsm that is worse than a tuning miss: MemWAL read routing lives only
|
||||
// in `create_plan`, so a pushed-down query would return stale base-only data.
|
||||
if query.base().use_lsm.is_some() || query.base().take_offsets.is_some() {
|
||||
@@ -115,7 +115,7 @@ fn requires_local_namespace_execution(query: &AnyQuery) -> bool {
|
||||
matches!(
|
||||
query,
|
||||
AnyQuery::VectorQuery(vector_query)
|
||||
if vector_query.approx_mode.is_some()
|
||||
if vector_query.approx_mode.is_some() || vector_query.search_effort.is_some()
|
||||
)
|
||||
}
|
||||
|
||||
@@ -163,6 +163,7 @@ pub async fn create_plan(
|
||||
AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query),
|
||||
};
|
||||
query.base.check_filter()?;
|
||||
query.validate_search_effort()?;
|
||||
|
||||
let ds_ref = table.dataset.get().await?;
|
||||
|
||||
@@ -293,7 +294,9 @@ pub async fn create_plan(
|
||||
scanner.approx_mode(approx_mode.into());
|
||||
}
|
||||
|
||||
if let Some(nprobes) = query.nprobes {
|
||||
if query.search_effort.is_none()
|
||||
&& let Some(nprobes) = query.nprobes
|
||||
{
|
||||
if query.maximum_nprobes == Some(0) {
|
||||
scanner.minimum_nprobes(nprobes);
|
||||
} else {
|
||||
@@ -306,6 +309,9 @@ pub async fn create_plan(
|
||||
if let Some(maximum_nprobes) = query.maximum_nprobes.filter(|nprobes| *nprobes != 0) {
|
||||
scanner.maximum_nprobes(maximum_nprobes);
|
||||
}
|
||||
if let Some(search_effort) = query.search_effort {
|
||||
scanner.search_effort(search_effort)?;
|
||||
}
|
||||
}
|
||||
|
||||
// For a batch query, `nearest` already applies k to each query vector.
|
||||
@@ -604,6 +610,11 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result<NsQueryTableRequest> {
|
||||
};
|
||||
|
||||
// Check for unsupported features
|
||||
if vq.search_effort.is_some() {
|
||||
return Err(Error::NotSupported {
|
||||
message: "namespace QueryTable does not support search_effort; execute the query locally".to_string(),
|
||||
});
|
||||
}
|
||||
if vq.base.reranker.is_some() {
|
||||
return Err(Error::NotSupported {
|
||||
message: "Reranker is not supported for server-side queries".to_string(),
|
||||
@@ -962,6 +973,19 @@ mod tests {
|
||||
let default_query = AnyQuery::VectorQuery(VectorQueryRequest::default());
|
||||
assert!(!requires_local_namespace_execution(&default_query));
|
||||
|
||||
for effort in [0.0, 0.5, 1.0] {
|
||||
let effort_query = AnyQuery::VectorQuery(VectorQueryRequest {
|
||||
query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))],
|
||||
search_effort: Some(effort),
|
||||
..Default::default()
|
||||
});
|
||||
assert!(requires_local_namespace_execution(&effort_query));
|
||||
assert!(matches!(
|
||||
convert_to_namespace_query(&effort_query),
|
||||
Err(Error::NotSupported { message }) if message.contains("search_effort")
|
||||
));
|
||||
}
|
||||
|
||||
let nprobes_query = AnyQuery::VectorQuery(VectorQueryRequest {
|
||||
query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))],
|
||||
nprobes: Some(20),
|
||||
|
||||
@@ -174,6 +174,9 @@ fn reject_unsupported(query: &VectorQueryRequest) -> Result<()> {
|
||||
// brute-force search the index-only base arm can't do. (ef / approx_mode are
|
||||
// recall/speed knobs and are left to no-op.)
|
||||
if !query.query_vector.is_empty() {
|
||||
if query.search_effort.is_some() {
|
||||
return unsupported("search_effort on vector search");
|
||||
}
|
||||
if query.lower_bound.is_some() || query.upper_bound.is_some() {
|
||||
return unsupported("distance_range on vector search");
|
||||
}
|
||||
@@ -821,6 +824,19 @@ mod tests {
|
||||
use super::*;
|
||||
use lance_index::mem_wal::{CompactedSsTable, IndexCatchupProgress};
|
||||
|
||||
#[test]
|
||||
fn search_effort_is_not_silently_ignored() {
|
||||
let query = VectorQueryRequest {
|
||||
query_vector: vec![Arc::new(arrow_array::Float32Array::from(vec![1.0, 0.0]))],
|
||||
search_effort: Some(0.5),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(
|
||||
matches!(reject_unsupported(&query), Err(Error::NotSupported { message })
|
||||
if message.contains("search_effort") && message.contains("un-compacted MemWAL"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resident_index_checks_distinguish_columns_and_query_types() {
|
||||
use arrow_schema::Field;
|
||||
|
||||
Reference in New Issue
Block a user