rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-22 03:32:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge main into gatekeeper/fix-4272-1 and preserve query validation

Browse Source
This commit is contained in:
Gatefixer
2026-10-04 04:54:45 +00:00
parent bc0abc7a57 7f5933594f
commit f253b4ba40
30 changed files with 1173 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/js/classes/AutoQuery.md
+2
View File
@@ -417,6 +417,8 @@ Collect the results as an Arrow
ArrowTable.
Empty results retain the query's output schema.
#### Inherited from
`StandardQueryBase.toArrow`
docs/src/js/classes/Query.md
+2
View File
@@ -491,6 +491,8 @@ Collect the results as an Arrow
ArrowTable.
Empty results retain the query's output schema.
#### Inherited from
`StandardQueryBase.toArrow`
docs/src/js/classes/QueryBase.md
+2
View File
@@ -252,6 +252,8 @@ Collect the results as an Arrow
ArrowTable.
Empty results retain the query's output schema.
***
### withRowId()
docs/src/js/classes/TakeQuery.md
+2
View File
@@ -267,6 +267,8 @@ Collect the results as an Arrow
ArrowTable.
Empty results retain the query's output schema.
#### Inherited from
[`QueryBase`](QueryBase.md).[`toArrow`](QueryBase.md#toarrow)
docs/src/js/classes/VectorQuery.md
+4
View File
@@ -447,6 +447,8 @@ nprobes(nprobes): VectorQuery
Set the number of partitions to search (probe)
The number of probes must be greater than 0.
This argument is only used when the vector column has an IVF PQ index.
If there is no index then this value is ignored.
@@ -740,6 +742,8 @@ Collect the results as an Arrow
ArrowTable.
Empty results retain the query's output schema.
#### Inherited from
`StandardQueryBase.toArrow`
nodejs/__test__/query.test.ts
+123
View File
@@ -110,6 +110,129 @@ describe("Query outputSchema", () => {
});
});
describe("Query empty results", () => {
let tmpDir: tmp.DirResult;
let table: Table;
beforeEach(async () => {
tmpDir = tmp.dirSync({ unsafeCleanup: true });
const db = await connect(tmpDir.name);
const schema = new Schema(
[
new Field("id", new Int64(), false),
new Field(
"text",
new Utf8(),
true,
new Map([["description", "searchable text"]]),
),
new Field(
"vector",
new FixedSizeList(2, new Field("item", new Float32())),
true,
),
],
new Map([["source", "empty-query-regression"]]),
);
table = await db.createTable(
"test",
makeArrowTable([{ id: 1n, text: "hello", vector: [1, 2] }], { schema }),
);
});
afterEach(() => {
tmpDir.removeCallback();
});
it.each([
{
name: "filtered scan",
query: (table: Table) => table.query().where("id < 0"),
fields: ["id", "text", "vector"],
},
{
name: "zero limit",
query: (table: Table) => table.query().limit(0),
fields: ["id", "text", "vector"],
},
{
name: "column projection",
query: (table: Table) =>
table.query().where("id < 0").select(["text", "id"]),
fields: ["text", "id"],
},
{
name: "dynamic projection",
query: (table: Table) =>
table.query().where("id < 0").select({ doubled: "id * 2" }),
fields: ["doubled"],
},
{
name: "filtered vector search",
query: (table: Table) =>
table.search([0.5, 0.5]).where("id < 0").limit(3),
fields: ["id", "text", "vector", "_distance"],
},
{
name: "vector projection with row id",
query: (table: Table) =>
table
.vectorSearch([0.5, 0.5])
.where("id < 0")
.select(["id"])
.withRowId(),
fields: ["id", "_distance", "_rowid"],
},
{
name: "fast search without an index",
query: (table: Table) => table.search([0.5, 0.5]).fastSearch(),
fields: ["id", "text", "vector", "_rowid", "_distance"],
},
{
name: "take with no offsets",
query: (table: Table) => table.takeOffsets([]).select(["text"]),
fields: ["text"],
},
])("preserves the output schema for $name", async ({ query, fields }) => {
const result = await query(table).toArrow({ maxBatchLength: 1 });
expect(result.numRows).toBe(0);
expect(result.schema.fields.map((field) => field.name)).toEqual(fields);
for (const name of fields) {
expect(result.getChild(name)).not.toBeNull();
expect(result.getChild(name)?.length).toBe(0);
}
});
it("preserves types, nullability, and metadata", async () => {
const result = await table.query().where("id < 0").toArrow();
const schema = await table.schema();
expect(result.schema).toEqual(schema);
expect(result.getChild("id")?.type.toString()).toBe("Int64");
expect(result.schema.fields[0].nullable).toBe(false);
expect(result.schema.fields[1].metadata.get("description")).toBe(
"searchable text",
);
expect(result.schema.metadata.get("source")).toBe("empty-query-regression");
});
it("preserves the score column for full-text search with no matches", async () => {
await table.createIndex("text", { config: Index.fts() });
const result = await table
.search("missing", "fts")
.select(["id"])
.toArrow();
expect(result.numRows).toBe(0);
expect(result.schema.fields.map((field) => field.name)).toEqual([
"id",
"_score",
]);
expect(result.getChild("_score")?.type.toString()).toBe("Float32");
});
});
describe("Search pagination", () => {
let tmpDir: tmp.DirResult;
let table: Table;
nodejs/__test__/remote.test.ts
+100
View File
@@ -14,6 +14,14 @@ import {
TlsConfig,
connect,
} from "../lancedb";
import {
Table as ArrowTable,
Field,
Int64,
RecordBatch,
Schema,
tableToIPC,
} from "../lancedb/arrow";
import {
HeaderProvider,
OAuthHeaderProvider,
@@ -85,6 +93,98 @@ async function withMockDatabase(
}
describe("remote connection", () => {
it("rejects nprobes(0) before sending a query", async () => {
const requests: string[] = [];
await withMockDatabase(
(req, res) => {
requests.push(req.url ?? "");
if (req.url === "/v1/table/test/describe/") {
res
.writeHead(200, { "Content-Type": "application/json" })
.end(JSON.stringify({ version: 1, schema: { fields: [] } }));
} else {
res.writeHead(404).end();
}
},
async (db) => {
const table = await db.openTable("test");
expect(() => table.vectorSearch([0, 0]).nprobes(0)).toThrow(
"Invalid input, nprobes must be greater than 0",
);
expect(() =>
table.query().nearestTo([0, 0]).fullTextSearch("dog").nprobes(0),
).toThrow("Invalid input, nprobes must be greater than 0");
await expect(
table
.vectorSearch(Promise.resolve([0, 0]))
.nprobes(0)
.toArrow(),
).rejects.toThrow("Invalid input, nprobes must be greater than 0");
},
);
expect(requests).toEqual(["/v1/table/test/describe/"]);
});
it.each([false, true])(
"preserves an empty query's schema with an empty batch: %s",
async (withEmptyBatch) => {
const schema = new Schema(
[new Field("doubled", new Int64(), false)],
new Map([["source", "query-output"]]),
);
const result = new ArrowTable(
schema,
withEmptyBatch ? [new RecordBatch(schema, undefined)] : [],
);
const response = Buffer.from(tableToIPC(result, "stream"));
let queryRequests = 0;
await withMockDatabase(
(req, res) => {
if (req.url?.endsWith("/describe/")) {
res.writeHead(200, { "Content-Type": "application/json" }).end(
JSON.stringify({
name: "items",
version: 1,
schema: {
fields: [
{ name: "id", type: { type: "int64" }, nullable: false },
],
},
}),
);
} else if (req.url?.endsWith("/query/")) {
queryRequests++;
req.resume();
req.on("end", () => {
res
.writeHead(200, {
"Content-Type": "application/vnd.apache.arrow.stream",
})
.end(response);
});
} else {
res.writeHead(404).end();
}
},
async (db) => {
const table = await db.openTable("items");
const result = await table
.query()
.where("id < 0")
.select({ doubled: "id * 2" })
.toArrow();
expect(result.numRows).toBe(0);
expect(result.schema).toEqual(schema);
expect(result.getChild("doubled")?.length).toBe(0);
},
);
expect(queryRequests).toBe(1);
},
);
it("rejects the external blob opt-in without blocking regular adds", async () => {
let insertRequests = 0;
const describeRequests: string[] = [];
nodejs/__test__/table.test.ts
+6
View File
@@ -1211,6 +1211,12 @@ describe("When creating an index", () => {
.toArrow();
expect(rst.numRows).toBe(2);
expect(() => tbl.vectorSearch(queryVec).nprobes(0)).toThrow(
"Invalid input, nprobes must be greater than 0",
);
expect(() =>
tbl.query().nearestTo(queryVec).fullTextSearch("dog").nprobes(0),
).toThrow("Invalid input, nprobes must be greater than 0");
expect(() => tbl.query().nearestTo(queryVec).minimumNprobes(0)).toThrow(
"Invalid input, minimum_nprobes must be greater than 0",
);
nodejs/lancedb/query.ts
+12 -23
View File
@@ -41,26 +41,6 @@ export async function* RecordBatchIterator(
}
}
class RecordBatchIterable<
NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery,
> implements AsyncIterable<RecordBatch>
{
private inner: NativeQueryType;
private options?: QueryExecutionOptions;
constructor(inner: NativeQueryType, options?: QueryExecutionOptions) {
this.inner = inner;
this.options = options;
}
// biome-ignore lint/suspicious/noExplicitAny: skip
[Symbol.asyncIterator](): AsyncIterator<RecordBatch<any>, any, undefined> {
return RecordBatchIterator(
this.inner.execute(this.options?.maxBatchLength, this.options?.timeoutMs),
);
}
}
/**
* Options that control the behavior of a particular query execution
*/
@@ -271,13 +251,20 @@ export class QueryBase<
return RecordBatchIterator(this.nativeExecute());
}
/** Collect the results as an Arrow @see {@link ArrowTable}. */
/**
* Collect the results as an Arrow @see {@link ArrowTable}.
*
* Empty results retain the query's output schema.
*/
async toArrow(options?: Partial<QueryExecutionOptions>): Promise<ArrowTable> {
const batches = [];
const inner = await this.getInner();
for await (const batch of new RecordBatchIterable(inner, options)) {
const iterator = this.nativeExecute(options);
for await (const batch of RecordBatchIterator(iterator)) {
batches.push(batch);
}
if (batches.length === 0) {
return tableFromIPC((await iterator).schema());
}
return new ArrowTable(batches);
}
@@ -532,6 +519,8 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
/**
* Set the number of partitions to search (probe)
*
* The number of probes must be greater than 0.
*
* This argument is only used when the vector column has an IVF PQ index.
* If there is no index then this value is ignored.
*
nodejs/src/iterator.rs
+8
View File
@@ -7,6 +7,8 @@ use lancedb::ipc::batches_to_ipc_file;
use napi::bindgen_prelude::*;
use napi_derive::napi;
use crate::util::schema_to_buffer;
/** Typescript-style Async Iterator over RecordBatches */
#[napi]
pub struct RecordBatchIterator {
@@ -19,6 +21,12 @@ impl RecordBatchIterator {
Self { inner }
}
/// Return the executed stream's output schema as an Arrow IPC file.
#[napi]
pub fn schema(&self) -> napi::Result<Buffer> {
schema_to_buffer(&self.inner.schema())
}
#[napi(catch_unwind)]
pub async unsafe fn next(&mut self) -> napi::Result<Option<Buffer>> {
if let Some(rst) = self.inner.next().await {
nodejs/src/query.rs
+7 -2
View File
@@ -289,8 +289,13 @@ impl VectorQuery {
}
#[napi]
pub fn nprobes(&mut self, nprobe: u32) {
self.inner = self.inner.clone().nprobes(nprobe as usize);
pub fn nprobes(&mut self, nprobe: u32) -> napi::Result<()> {
self.inner = self
.inner
.clone()
.nprobes(nprobe as usize)
.default_error()?;
Ok(())
}
#[napi]
python/python/lancedb/db.py
+55 -6
View File
@@ -122,7 +122,48 @@ def _view_description(
class DBConnection(EnforceOverrides):
"""An active LanceDB connection interface."""
"""An active LanceDB connection interface.
Use [close][lancedb.db.DBConnection.close] to release the connection's
underlying resources, or use the connection as a context manager to close it
automatically when leaving the block, including when an exception is raised.
Examples
--------
>>> import lancedb
>>> with lancedb.connect("memory://") as db:
... assert db.is_open()
>>> db.is_open()
False
"""
def __enter__(self) -> DBConnection:
return self
def __exit__(self, *_) -> None:
self.close()
@abstractmethod
def is_open(self) -> bool:
"""Return True if the connection is open."""
pass
@abstractmethod
def close(self) -> None:
"""Close the connection, releasing any underlying resources.
It is safe to call this method multiple times. Database operations on a
closed connection raise ``RuntimeError: Connection is closed``.
Examples
--------
>>> import lancedb
>>> db = lancedb.connect("memory://")
>>> db.close()
>>> db.is_open()
False
"""
pass
def list_namespaces(
self,
@@ -368,9 +409,10 @@ class DBConnection(EnforceOverrides):
- [LanceModel][lancedb.pydantic.LanceModel]
mode: str; default "create"
The mode to use when creating the table.
Can be either "create" or "overwrite".
Can be "create", "overwrite", or "exist_ok".
By default, if the table already exists, an exception is raised.
If you want to overwrite the table, use mode="overwrite".
To open an existing table without adding data, use mode="exist_ok".
exist_ok: bool, default False
If a table by the same name already exists, then raise an exception
if exist_ok=False. If exist_ok=True, then open the existing table;
@@ -1275,6 +1317,14 @@ class LanceDBConnection(DBConnection):
def __repr__(self) -> str:
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
@override
def is_open(self) -> bool:
return self._conn.is_open()
@override
def close(self) -> None:
self._conn.close()
@override
def serialize(self) -> str:
import json
@@ -1497,8 +1547,6 @@ class LanceDBConnection(DBConnection):
"""
if namespace_path is None:
namespace_path = []
if mode.lower() not in ["create", "overwrite"]:
raise ValueError("mode must be either 'create' or 'overwrite'")
validate_table_name(name)
tbl = LanceTable.create(
@@ -2301,11 +2349,12 @@ class AsyncConnection(object):
- pyarrow.Schema
- [LanceModel][lancedb.pydantic.LanceModel]
mode: Literal["create", "overwrite"]; default "create"
mode: Literal["create", "overwrite", "exist_ok"]; default "create"
The mode to use when creating the table.
Can be either "create" or "overwrite".
Can be "create", "overwrite", or "exist_ok".
By default, if the table already exists, an exception is raised.
If you want to overwrite the table, use mode="overwrite".
To open an existing table without adding data, use mode="exist_ok".
exist_ok: bool, default False
If a table by the same name already exists, then raise an exception
if exist_ok=False. If exist_ok=True, then open the existing table;
python/python/lancedb/namespace.py
+8
View File
@@ -489,6 +489,14 @@ class LanceNamespaceDBConnection(DBConnection):
)
self._uri = self._inner.uri
@override
def is_open(self) -> bool:
return self._inner.is_open()
@override
def close(self) -> None:
self._inner.close()
@override
def serialize(self) -> str:
import json
python/python/lancedb/query.py
+13 -3
View File
@@ -40,7 +40,7 @@ from .rerankers.base import Reranker
from .rerankers.rrf import RRFReranker
from .rerankers.util import check_reranker_result
from .schema import blob_column_paths, is_blob_like_field, schema_has_blob_field
from .util import flatten_columns
from .util import _validate_query_vector, flatten_columns
from . import _wal_hybrid # WAL-PK-FUSION: delete.
from ._blob import (
BLOB_MODE_TO_HANDLING,
@@ -908,6 +908,9 @@ class LanceQueryBuilder(ABC):
fast_search: bool
Skip flat search of unindexed data.
"""
if query_type != "fts":
_validate_query_vector(query)
if ordering_field_name is not None:
import warnings
@@ -1620,6 +1623,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
str_query: Optional[str] = None,
fast_search: bool = None,
):
_validate_query_vector(query)
super().__init__(table)
self._query = query
self._distance_type = None
@@ -1694,13 +1698,15 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceVectorQueryBuilder
The LanceQueryBuilder object.
"""
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
return self
@@ -2496,13 +2502,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceHybridQueryBuilder
The LanceHybridQueryBuilder object.
"""
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
return self
@@ -3686,6 +3694,8 @@ class AsyncVectorQueryBase:
"""
Set the number of partitions to search (probe)
The number of probes must be greater than 0.
This argument is only used when the vector column has an IVF-based index.
If there is no index then this value is ignored.
python/python/lancedb/remote/db.py
+6 -2
View File
@@ -1054,6 +1054,10 @@ class RemoteDBConnection(DBConnection):
"""
return LOOP.run(self._conn.namespace_client())
async def close(self):
"""Close the connection to the database."""
@override
def is_open(self) -> bool:
return self._conn.is_open()
@override
def close(self) -> None:
self._conn.close()
python/python/lancedb/table.py
+34 -14
View File
@@ -2617,6 +2617,10 @@ class Table(ABC):
to nullable. Currently, you cannot change a nullable column to
non-nullable.
Each alteration must specify at least one of "rename", "data_type",
or "nullable". The legacy key "name" is also accepted for renaming.
Unknown keys raise ValueError before any alterations are applied.
Returns
-------
AlterColumnsResult
@@ -4332,9 +4336,9 @@ class LanceTable(Table):
)
if storage_options is None:
storage_options = {}
storage_options["new_table_enable_v2_manifest_paths"] = (
storage_options["new_table_enable_v2_manifest_paths"] = str(
enable_v2_manifest_paths
)
).lower()
self._table = LOOP.run(
self._conn._conn.create_table(
@@ -4766,7 +4770,7 @@ class LanceTable(Table):
[LanceTable.uses_v2_manifest_paths][lancedb.table.LanceTable.uses_v2_manifest_paths]
to check if the table is already using the new path style.
"""
LOOP.run(self._table.migrate_v2_manifest_paths())
LOOP.run(self._table.migrate_manifest_paths_v2())
@deprecation.deprecated(
deprecated_in="0.33.1",
@@ -6122,7 +6126,17 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[str] = None,
query: None = None,
vector_column_name: Optional[str] = None,
query_type: QueryType = "auto",
ordering_field_name: Optional[str] = None,
fts_columns: Optional[Union[str, List[str]]] = None,
) -> AsyncQuery: ...
@overload
async def search(
self,
query: str,
vector_column_name: Optional[str] = None,
query_type: Literal["auto"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6132,7 +6146,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[str] = None,
query: str,
vector_column_name: Optional[str] = None,
query_type: Literal["hybrid"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6142,7 +6156,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[Union[VEC, "PIL.Image.Image", Tuple]] = None,
query: Union[VEC, "PIL.Image.Image", Tuple],
vector_column_name: Optional[str] = None,
query_type: Literal["auto"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6152,7 +6166,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[str] = None,
query: str,
vector_column_name: Optional[str] = None,
query_type: Literal["fts"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6162,9 +6176,7 @@ class AsyncTable:
@overload
async def search(
self,
query: Optional[
Union[VEC, str, "PIL.Image.Image", Tuple, FullTextQuery]
] = None,
query: Union[VEC, str, "PIL.Image.Image", Tuple, FullTextQuery],
vector_column_name: Optional[str] = None,
query_type: Literal["vector"] = ...,
ordering_field_name: Optional[str] = None,
@@ -6180,7 +6192,7 @@ class AsyncTable:
query_type: QueryType = "auto",
ordering_field_name: Optional[str] = None,
fts_columns: Optional[Union[str, List[str]]] = None,
) -> Union[AsyncHybridQuery, AsyncFTSQuery, AsyncVectorQuery]:
) -> Union[AsyncQuery, AsyncHybridQuery, AsyncFTSQuery, AsyncVectorQuery]:
"""Create a search query to find the nearest neighbors
of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/)
and [full-text search](https://lancedb.com/docs/search/full-text-search/).
@@ -6195,8 +6207,9 @@ class AsyncTable:
- *default None*.
Acceptable types are: list, np.ndarray, PIL.Image.Image
- If None then the select/where/limit clauses are applied to filter
the table
- If None then a plain [AsyncQuery][lancedb.query.AsyncQuery] is
returned, equivalent to calling [query][lancedb.table.AsyncTable.query].
The select/where/limit clauses are applied to filter the table.
vector_column_name: str, optional
The name of the vector column to search.
@@ -6224,10 +6237,13 @@ class AsyncTable:
Returns
-------
LanceQueryBuilder
AsyncQuery, AsyncHybridQuery, AsyncFTSQuery, or AsyncVectorQuery
A query builder object representing the query.
"""
if query is None:
return self.query()
def is_embedding(query):
return isinstance(query, (list, np.ndarray, pa.Array, pa.ChunkedArray))
@@ -6877,6 +6893,10 @@ class AsyncTable:
to nullable. Currently, you cannot change a nullable column to
non-nullable.
Each alteration must specify at least one of "rename", "data_type",
or "nullable". The legacy key "name" is also accepted for renaming.
Unknown keys raise ValueError before any alterations are applied.
Returns
-------
AlterColumnsResult
python/python/lancedb/util.py
+11
View File
@@ -280,6 +280,14 @@ def infer_vector_column_dim(data_type: pa.DataType) -> Optional[int]:
return None
def _validate_query_vector(query: Any) -> None:
"""Reject empty vector inputs before vector-column inference or execution."""
if (isinstance(query, list) and not query) or (
isinstance(query, np.ndarray) and query.size == 0
):
raise ValueError("Query vector must not be empty")
def _query_vector_dim(query: Optional[Any]) -> Optional[int]:
if query is None:
return None
@@ -301,6 +309,9 @@ def infer_vector_column_name(
query: Optional[Any], # inferred later in query builder
vector_column_name: Optional[str],
):
if query_type != "fts":
_validate_query_vector(query)
if vector_column_name is not None:
return vector_column_name
python/python/tests/test_db.py
+99
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import contextlib
import inspect
import re
import sys
@@ -350,6 +351,40 @@ async def test_create_table_from_iterator_async(mem_db_async: lancedb.AsyncConne
assert await table.count_rows() == 10
def test_create_exist_ok_mode(tmp_db: lancedb.DBConnection):
db = tmp_db
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
table = db.create_table("t1", data, mode="exist_ok")
assert table.count_rows() == 2
original = table.to_arrow()
reopened = db.create_table("t1", [{"vector": [9.0, 8.0]}], mode="exist_ok")
assert reopened.to_arrow().equals(original)
def test_create_empty_exist_ok_mode(tmp_db: lancedb.DBConnection):
schema = pa.schema([pa.field("id", pa.int64())])
table = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
assert table.schema == schema
assert table.count_rows() == 0
table.add([{"id": 7}])
reopened = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
assert reopened.to_arrow().to_pylist() == [{"id": 7}]
bad_schema = pa.schema([pa.field("other", pa.int64())])
with pytest.raises(ValueError):
tmp_db.create_table("empty", schema=bad_schema, mode="exist_ok")
def test_create_table_invalid_mode(tmp_db: lancedb.DBConnection):
db = tmp_db
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
with pytest.raises(ValueError, match="Invalid mode bogus"):
db.create_table("tb", data, mode="bogus")
assert "tb" not in db
def test_create_exist_ok(tmp_db: lancedb.DBConnection):
from conftest import pandas_string_type
@@ -405,6 +440,49 @@ async def test_connect(tmp_path):
assert str(db) == f"ListingDatabase(uri={tmp_path}, read_consistency_interval=5s)"
@pytest.fixture(params=["memory", "local", "namespace", "from_inner"])
def sync_db(request, tmp_path):
if request.param == "memory":
return lancedb.connect("memory://")
if request.param == "namespace":
return lancedb.connect_namespace("dir", {"root": str(tmp_path)})
db = lancedb.connect(tmp_path)
if request.param == "from_inner":
return lancedb.db.LanceDBConnection.from_inner(db._inner, None)
return db
def test_sync_close(sync_db):
assert sync_db.is_open()
assert sync_db.close() is None
assert not sync_db.is_open()
assert sync_db.close() is None
with pytest.warns(DeprecationWarning, match="table_names"):
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.table_names()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.list_tables()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.open_table("test")
@pytest.mark.parametrize("raise_error", [False, True])
def test_sync_context_manager(sync_db, raise_error):
with contextlib.ExitStack() as stack:
if raise_error:
stack.enter_context(pytest.raises(ValueError, match="test error"))
with sync_db as db:
assert db is sync_db
assert db.is_open()
if raise_error:
raise ValueError("test error")
assert not sync_db.is_open()
with pytest.raises(RuntimeError, match="Connection is closed"):
sync_db.list_tables()
@pytest.mark.asyncio
async def test_close(mem_db_async: lancedb.AsyncConnection):
assert mem_db_async.is_open()
@@ -495,6 +573,27 @@ async def test_create_exist_ok_async(tmp_db_async: lancedb.AsyncConnection):
# await db.create_table("test", schema=bad_schema, exist_ok=True)
@pytest.mark.parametrize("enable_v2_manifest_paths", [False, True])
@pytest.mark.parametrize("empty", [False, True])
def test_create_table_deprecated_v2_manifest_paths(
tmp_path, enable_v2_manifest_paths, empty
):
db = lancedb.connect(tmp_path)
data = None if empty else [{"id": 1}]
with pytest.warns(DeprecationWarning, match="enable_v2_manifest_paths"):
table = db.create_table(
"test",
data=data,
schema=pa.schema([("id", pa.int64())]),
enable_v2_manifest_paths=enable_v2_manifest_paths,
)
assert table.uses_v2_manifest_paths() == enable_v2_manifest_paths
assert table.to_arrow().to_pylist() == ([] if empty else data)
assert db.open_table("test").uses_v2_manifest_paths() == enable_v2_manifest_paths
@pytest.mark.asyncio
async def test_create_table_v2_manifest_paths_async(tmp_path):
db_with_v2_paths = await lancedb.connect_async(
python/python/tests/test_query.py
+53
View File
@@ -24,6 +24,7 @@ from lancedb.pydantic import LanceModel, Vector
from lancedb.query import (
AsyncFTSQuery,
AsyncHybridQuery,
AsyncQuery,
AsyncQueryBase,
AsyncVectorQuery,
ColumnOrdering,
@@ -974,6 +975,19 @@ def test_query_builder_with_filter(table):
assert all(np.array(rs[0]["vector"]) == [3, 4])
@pytest.mark.parametrize("query_type", ["vector", "hybrid"])
@pytest.mark.parametrize("nprobes", [0, -1])
def test_nprobes_nonpositive_sync(table, query_type, nprobes):
if query_type == "hybrid":
query = table.search(query_type="hybrid").vector([0, 0]).text("a")
else:
query = table.search([0, 0])
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(nprobes)
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()
@@ -1005,6 +1019,18 @@ def test_multiple_nprobes_calls_works_sync(table):
).minimum_nprobes(20).to_list()
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
async def test_nprobes_zero_async(table_async: AsyncTable, hybrid):
query = table_async.query().nearest_to([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
@pytest.mark.asyncio
async def test_invalid_nprobes_async(table_async: AsyncTable):
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
@@ -1279,6 +1305,27 @@ async def test_query_to_polars_async(table_async: AsyncTable):
assert df.shape == (0, num_columns)
@pytest.mark.asyncio
@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}])
@pytest.mark.parametrize("with_vector", [False, True])
async def test_async_search_without_query(mem_db_async, search_kwargs, with_vector):
data = [{"id": i} for i in range(10)]
if with_vector:
for row in data:
row["vector"] = [float(row["id"]), 1.0]
table = await mem_db_async.create_table("test", data)
query = await table.search(**search_kwargs)
assert isinstance(query, AsyncQuery)
assert await query.limit(3).to_arrow() == await table.query().limit(3).to_arrow()
query = await table.search(**search_kwargs)
assert await query.where("id >= 8").select(["id"]).limit(3).to_list() == [
{"id": 8},
{"id": 9},
]
@pytest.mark.asyncio
async def test_none_query(table_async: AsyncTable):
with pytest.raises(ValueError):
@@ -2091,6 +2138,12 @@ def test_ensure_vector_query_nested_empty_list():
ensure_vector_query([[]])
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
def test_vector_query_builder_empty_vector(table, query):
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
LanceVectorQueryBuilder(table, query, "vector")
def test_fast_search(tmp_path):
db = lancedb.connect(tmp_path)
python/python/tests/test_remote_db.py
+209 -5
View File
@@ -16,8 +16,9 @@ import uuid
from packaging.version import Version
import lancedb
import numpy as np
from lancedb.conftest import MockTextEmbeddingFunction
from lancedb.query import ColumnOrdering
from lancedb.query import AsyncQuery, ColumnOrdering
from lancedb.remote import ClientConfig
from lancedb.remote.errors import HttpError, RetryError
import pytest
@@ -119,6 +120,77 @@ async def test_async_remote_db():
assert table_names == []
@pytest.mark.parametrize(
"alteration, match",
[
({"path": "id"}, "One of rename, nullable or data_type"),
(
{"path": "id", "nulable": False}, # spellchecker:disable-line
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
{
"path": "id",
"rename": "new_id",
"nulable": False, # spellchecker:disable-line
},
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
],
)
def test_remote_alter_columns_rejects_invalid_before_request(alteration, match):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
with mock_lancedb_connection(handler) as db:
table = db.open_table("test")
requests.clear()
with pytest.raises(ValueError, match=match):
table.alter_columns(alteration)
assert requests == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
"alteration, match",
[
({"path": "id"}, "One of rename, nullable or data_type"),
(
{"path": "id", "nulable": False}, # spellchecker:disable-line
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
],
)
async def test_async_remote_alter_columns_rejects_invalid_before_request(
alteration, match
):
requests = []
def handler(request):
requests.append(request.path)
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
requests.clear()
with pytest.raises(ValueError, match=match):
await table.alter_columns(alteration)
assert requests == []
@pytest.mark.asyncio
async def test_async_checkout():
def handler(request):
@@ -1521,6 +1593,60 @@ def test_query_sync_empty_query():
assert data == expected
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
@pytest.mark.parametrize("vector_column_name", [None, "vector"])
@pytest.mark.parametrize("query_type", ["auto", "vector"])
def test_query_sync_empty_vector(query, vector_column_name, query_type):
def handler(body):
pytest.fail("An empty query vector must be rejected before sending a query")
with query_test_table(handler) as table:
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
table.search(
query, vector_column_name=vector_column_name, query_type=query_type
).limit(3).to_arrow()
@pytest.mark.asyncio
@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}])
async def test_async_search_without_query(search_kwargs):
expected = pa.table({"id": [7, 8, 9]})
def handler(request):
if request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
elif request.path == "/v1/table/test/query/":
body = json.loads(
request.rfile.read(int(request.headers["Content-Length"]))
)
assert body == {
"k": 3,
"filter": "id >= 7",
"vector": [],
"columns": ["id"],
"prefilter": True,
"version": None,
}
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
request.end_headers()
with pa.ipc.new_file(request.wfile, schema=expected.schema) as writer:
writer.write_table(expected)
else:
request.send_response(404)
request.end_headers()
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
query = await table.search(**search_kwargs)
assert isinstance(query, AsyncQuery)
result = await query.where("id >= 7").select(["id"]).limit(3).to_arrow()
assert result == expected
def test_query_sync_maximal():
def handler(body):
assert body == {
@@ -1581,6 +1707,53 @@ def test_query_sync_maximal():
)
@pytest.mark.parametrize("hybrid", [False, True])
def test_query_sync_nprobes_zero(hybrid):
query_requests = []
def handler(body):
query_requests.append(body)
return pa.table({"id": []})
with query_test_table(handler) as table:
if hybrid:
query = table.search(query_type="hybrid").vector([1, 2, 3]).text("dog")
else:
query = table.search([1, 2, 3])
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
assert query_requests == []
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
async def test_query_async_nprobes_zero(hybrid):
requests = []
def handler(request):
requests.append(request.path)
if request.path == "/v1/table/test/describe/":
send_json(request, {"version": 1, "schema": {"fields": []}})
else:
request.send_response(404)
request.end_headers()
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
query = table.query().nearest_to([1, 2, 3])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
assert requests == ["/v1/table/test/describe/"]
def test_query_sync_nprobes():
def handler(body):
assert body == {
@@ -2145,15 +2318,46 @@ async def test_header_provider_overrides_static_headers():
def test_close():
"""Test that close() works without AttributeError."""
import asyncio
def handler(req):
req.send_response(200)
req.end_headers()
with mock_lancedb_connection(handler) as db:
asyncio.run(db.close())
assert db.close() is None
assert not db.is_open()
assert db.close() is None
with pytest.warns(DeprecationWarning, match="table_names"):
with pytest.raises(RuntimeError, match="Connection is closed"):
db.table_names()
with pytest.raises(RuntimeError, match="Connection is closed"):
db.list_tables()
with pytest.raises(RuntimeError, match="Connection is closed"):
db.open_table("test")
@pytest.mark.parametrize("raise_error", [False, True])
def test_sync_context_manager(raise_error):
def handler(req):
req.send_response(200)
req.send_header("Content-Type", "application/json")
req.end_headers()
req.wfile.write(b'{"tables": []}')
with mock_lancedb_connection(handler) as db:
with contextlib.ExitStack() as stack:
if raise_error:
stack.enter_context(pytest.raises(ValueError, match="test error"))
with db as conn:
assert conn is db
assert conn.is_open()
assert conn.list_tables().tables == []
if raise_error:
raise ValueError("test error")
assert not db.is_open()
with pytest.raises(RuntimeError, match="Connection is closed"):
db.list_tables()
@pytest.mark.parametrize("exception", [KeyboardInterrupt, SystemExit, GeneratorExit])
python/python/tests/test_table.py
+109 -4
View File
@@ -113,6 +113,32 @@ def test_basic(mem_db: DBConnection):
assert table.to_arrow() == expected_data
@pytest.mark.parametrize("enable_v2", [False, True])
def test_migrate_v2_manifest_paths(tmp_path, enable_v2):
db = lancedb.connect(
tmp_path,
storage_options={"new_table_enable_v2_manifest_paths": str(enable_v2).lower()},
)
table = db.create_table("calls", [{"id": 1, "vector": [1.0, 1.0]}])
table.add([{"id": 2, "vector": [2.0, 2.0]}])
expected_data = table.to_arrow()
expected_versions = table.list_versions()
assert table.uses_v2_manifest_paths() == enable_v2
# Migration is also safe to repeat on a table already using v2 paths.
for _ in range(2):
table.migrate_v2_manifest_paths()
assert table.uses_v2_manifest_paths()
reopened = db.open_table("calls")
assert reopened.uses_v2_manifest_paths()
assert reopened.to_arrow() == expected_data
assert reopened.list_versions() == expected_versions
manifests = list((tmp_path / "calls.lance" / "_versions").glob("*.manifest"))
assert len(manifests) == len(expected_versions)
assert all(len(path.stem) == 20 and path.stem.isdigit() for path in manifests)
def test_search_preserves_nulls_from_sliced_arrow_table(mem_db: DBConnection):
data = pa.table(
{
@@ -3751,6 +3777,24 @@ def test_empty_query(mem_db: DBConnection):
assert df.num_rows == 42
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
@pytest.mark.parametrize("vector_column_name", [None, "vector"])
@pytest.mark.parametrize("query_type", ["auto", "vector"])
@pytest.mark.parametrize("multiple_vector_columns", [False, True])
def test_search_empty_vector(
mem_db, query, vector_column_name, query_type, multiple_vector_columns
):
fields = [pa.field("vector", pa.list_(pa.float32(), 8))]
if multiple_vector_columns:
fields.append(pa.field("vec2", pa.list_(pa.float32(), 4)))
table = mem_db.create_table("empty_vector_query", schema=pa.schema(fields))
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
table.search(
query, vector_column_name=vector_column_name, query_type=query_type
).limit(3).to_arrow()
def test_search_with_schema_inf_single_vector(mem_db: DBConnection):
class MyTable(LanceModel):
text: str
@@ -4144,14 +4188,74 @@ async def test_add_columns_with_schema(mem_db_async: AsyncConnection):
)
def test_alter_columns(mem_db: DBConnection):
@pytest.mark.parametrize("rename_key", ["rename", "name"])
def test_alter_columns(mem_db: DBConnection, rename_key):
data = pa.table({"id": [0, 1]})
table = mem_db.create_table("my_table", data=data)
alter_columns_res = table.alter_columns({"path": "id", "rename": "new_id"})
alter_columns_res = table.alter_columns({"path": "id", rename_key: "new_id"})
assert alter_columns_res.version == 2
assert table.to_arrow().column_names == ["new_id"]
INVALID_COLUMN_ALTERATIONS = [
(({"path": "id"},), "One of rename, nullable or data_type"),
(
({"path": "id", "nulable": False},), # spellchecker:disable-line
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
(
{
"path": "id",
"rename": "new_id",
"nulable": False, # spellchecker:disable-line
},
),
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
),
(
({"path": "id", "rename": "new_id"}, {"path": "id"}),
"One of rename, nullable or data_type",
),
]
def test_alter_columns_nullable_false(mem_db: DBConnection):
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
result = table.alter_columns({"path": "id", "nullable": False})
assert result.version == 2
assert not table.schema.field("id").nullable
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
def test_alter_columns_rejects_invalid(mem_db: DBConnection, alterations, match):
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
initial_version = table.version
initial_schema = table.schema
with pytest.raises(ValueError, match=match):
table.alter_columns(*alterations)
assert table.version == initial_version
assert table.schema == initial_schema
@pytest.mark.asyncio
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
async def test_alter_columns_rejects_invalid_async(
mem_db_async: AsyncConnection, alterations, match
):
table = await mem_db_async.create_table("my_table", data=pa.table({"id": [0, 1]}))
initial_version = await table.version()
initial_schema = await table.schema()
with pytest.raises(ValueError, match=match):
await table.alter_columns(*alterations)
assert await table.version() == initial_version
assert await table.schema() == initial_schema
def test_update_field_metadata(mem_db: DBConnection):
data = pa.table({"id": [0, 1], "category": ["a", "b"]})
table = mem_db.create_table("my_table", data=data)
@@ -4177,10 +4281,11 @@ def test_update_field_metadata(mem_db: DBConnection):
@pytest.mark.asyncio
async def test_alter_columns_async(mem_db_async: AsyncConnection):
@pytest.mark.parametrize("rename_key", ["rename", "name"])
async def test_alter_columns_async(mem_db_async: AsyncConnection, rename_key):
data = pa.table({"id": [0, 1]})
table = await mem_db_async.create_table("my_table", data=data)
alter_columns_res = await table.alter_columns({"path": "id", "rename": "new_id"})
alter_columns_res = await table.alter_columns({"path": "id", rename_key: "new_id"})
assert alter_columns_res.version == 2
assert (await table.to_arrow()).column_names == ["new_id"]
alter_columns_res = await table.alter_columns(
python/src/query.rs
+5 -4
View File
@@ -995,8 +995,9 @@ impl VectorQuery {
self.inner = self.inner.clone().refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner = self.inner.clone().nprobes(nprobe as usize);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner = self.inner.clone().nprobes(nprobe as usize).infer_error()?;
Ok(())
}
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
@@ -1188,8 +1189,8 @@ impl HybridQuery {
self.inner_vec.refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner_vec.nprobes(nprobe);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner_vec.nprobes(nprobe)
}
pub fn ef(&mut self, ef: u32) {
python/src/table.rs
+11
View File
@@ -1902,6 +1902,17 @@ impl Table {
let alterations = alterations
.iter()
.map(|alteration| {
for key in alteration.keys().iter() {
let key = key.extract::<String>()?;
if !matches!(
key.as_str(),
"path" | "rename" | "name" | "nullable" | "data_type"
) {
return Err(PyValueError::new_err(format!(
"Unknown column alteration key '{key}'"
)));
}
}
let path = alteration
.get_item("path")?
.ok_or_else(|| PyValueError::new_err("Missing path"))?
rust/lancedb/examples/ivf_pq.rs
+1 -1
View File
@@ -142,7 +142,7 @@ async fn search_index(table: &Table) -> Result<()> {
// Override the default of 10 to get more rows
.limit(15)
// Override the default of 20 to search more partitions
.nprobes(30)
.nprobes(30)?
// Override the default of None to apply a refine step
.refine_factor(1)
.execute()
rust/lancedb/src/connection.rs
+58
View File
@@ -568,6 +568,7 @@ impl Connection {
/// Creates a new table by cloning from an existing source table.
/// By default, this performs a shallow clone where the new table shares
/// the underlying data files with the source table.
/// The target table name must be unused. An existing table is left unchanged.
///
/// # Parameters
/// - `target_table_name`: The name of the new table to create
@@ -2193,6 +2194,63 @@ mod tests {
assert_eq!(tables.len(), 0);
}
#[rstest::rstest]
#[tokio::test]
async fn test_clone_table_target_already_exists(
#[values("source", "target")] target: &str,
#[values(false, true)] enable_v2_manifest_paths: bool,
) {
use crate::query::ExecutableQuery;
use futures::TryStreamExt;
let tmp_dir = tempdir().unwrap();
let options = ListingDatabaseOptions::builder()
.enable_v2_manifest_paths(enable_v2_manifest_paths)
.build();
let db = connect(tmp_dir.path().to_str().unwrap())
.database_options(&options)
.execute()
.await
.unwrap();
let source_data = arrow_array::record_batch!(("id", Int32, [0, 1, 2])).unwrap();
let target_data = arrow_array::record_batch!(("id", Int32, [10, 11])).unwrap();
db.create_table("source", source_data.clone())
.execute()
.await
.unwrap();
db.create_table("target", target_data.clone())
.execute()
.await
.unwrap();
let source_uri = tmp_dir.path().join("source.lance");
let err = db
.clone_table(target, source_uri.to_str().unwrap())
.execute()
.await
.unwrap_err();
assert!(
matches!(&err, Error::TableAlreadyExists { name } if name == target),
"unexpected clone error: {err:?}"
);
// Reopen both tables to verify that the failed clone left their data and versions intact.
for (name, expected_data) in [("source", source_data), ("target", target_data)] {
let table = db.open_table(name).execute().await.unwrap();
assert_eq!(table.version().await.unwrap(), 1);
assert_eq!(table.list_versions().await.unwrap().len(), 1);
let batches = table
.query()
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(batches, vec![expected_data]);
}
}
#[tokio::test]
async fn test_clone_table() {
let tmp_dir = tempdir().unwrap();
rust/lancedb/src/database/listing.rs
+15 -1
View File
@@ -1107,6 +1107,21 @@ impl Database for ListingDatabase {
..Default::default()
};
let target_uri = self.table_uri(&request.target_table_name)?;
match DatasetBuilder::from_uri(&target_uri)
.with_read_params(read_params.clone())
.load()
.await
{
Ok(_) => {
return Err(Error::TableAlreadyExists {
name: request.target_table_name,
});
}
Err(lance::Error::DatasetNotFound { .. }) => {}
Err(err) => return Err(err.into()),
}
let mut source_dataset = DatasetBuilder::from_uri(&request.source_uri)
.with_read_params(read_params.clone())
.load()
@@ -1122,7 +1137,6 @@ impl Database for ListingDatabase {
}),
}?;
let target_uri = self.table_uri(&request.target_table_name)?;
source_dataset
.shallow_clone(&target_uri, version_ref, Some(storage_params))
.await
rust/lancedb/src/query.rs
+35 -2
View File
@@ -1252,10 +1252,17 @@ impl VectorQuery {
/// This method sets both the minimum and maximum number of partitions to search.
/// For more fine-grained control see [`VectorQuery::minimum_nprobes`] and
/// [`VectorQuery::maximum_nprobes`].
pub fn nprobes(mut self, nprobes: usize) -> Self {
///
/// Returns an error if `nprobes` is not greater than 0.
pub fn nprobes(mut self, nprobes: usize) -> Result<Self> {
if nprobes == 0 {
return Err(Error::InvalidInput {
message: "nprobes must be greater than 0".to_string(),
});
}
self.request.minimum_nprobes = nprobes;
self.request.maximum_nprobes = Some(nprobes);
self
Ok(self)
}
/// Set the minimum number of partitions to search
@@ -2361,6 +2368,7 @@ mod tests {
.nearest_to(&[9.8, 8.7])
.unwrap()
.nprobes(1000)
.unwrap()
.postfilter()
.distance_type(DistanceType::Cosine)
.approx_mode(ApproxMode::Accurate)
@@ -2386,6 +2394,31 @@ mod tests {
assert_eq!(query.request.refine_factor, Some(999));
}
#[tokio::test]
async fn test_nprobes_validation() {
let conn = connect("memory://").execute().await.unwrap();
let table = conn
.create_table("my_table", make_test_batches())
.execute()
.await
.unwrap();
let query = table.query().nearest_to(&[0.1, 0.2]).unwrap();
assert!(matches!(
query.clone().nprobes(0),
Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0"
));
// Fixed probe counts replace both bounds, even when moving below or
// above the previous range.
let mut query = query;
for nprobes in [30, 1, 50] {
query = query.nprobes(nprobes).unwrap();
assert_eq!(query.request.minimum_nprobes, nprobes);
assert_eq!(query.request.maximum_nprobes, Some(nprobes));
}
}
#[test]
fn test_approx_mode_serde_parse_default_and_display() {
assert_eq!(ApproxMode::default(), ApproxMode::Normal);
rust/lancedb/src/remote/table.rs
+141 -1
View File
@@ -3101,7 +3101,13 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
// loop can re-execute the plan (and re-stream the body) on each retry.
// This mirrors the old `send_streaming(with_retry=true)` path, which
// likewise buffered the reader to support retries.
let batches = new_data.collect::<std::result::Result<Vec<_>, _>>()?;
let schema = RecordBatchReader::schema(new_data.as_ref());
let mut batches = new_data.collect::<std::result::Result<Vec<_>, _>>()?;
// An empty reader still carries a schema. Keep it in an empty batch so
// the buffered source remains scannable and can be replayed on retries.
if batches.is_empty() {
batches.push(RecordBatch::new_empty(schema));
}
let source: Box<dyn Scannable> = Box::new(batches);
let rescannable = source.rescannable();
let input: Arc<dyn ExecutionPlan> =
@@ -4688,6 +4694,25 @@ mod tests {
assert_eq!(result.rows_updated, if old_server { 0 } else { 5 });
}
#[tokio::test]
async fn test_alter_columns_rejects_missing_changes_before_request() {
let table = Table::new_with_handler::<String>("my_table", |request| {
panic!("Unexpected request: {}", request.url().path())
});
for alterations in [
vec![ColumnAlteration::new("id".into())],
vec![
ColumnAlteration::new("id".into()).rename("new_id".into()),
ColumnAlteration::new("id".into()),
],
] {
let err = table.alter_columns(&alterations).await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}");
assert!(err.to_string().contains("path 'id'"));
}
}
#[rstest]
#[case(true)]
#[case(false)]
@@ -4801,6 +4826,104 @@ mod tests {
}
}
#[rstest]
#[case::no_batches_insert(false, false)]
#[case::empty_batch_insert(true, false)]
#[case::no_batches_delete(false, true)]
#[case::empty_batch_delete(true, true)]
#[tokio::test]
async fn test_merge_insert_empty_source(
#[case] has_batch: bool,
#[case] delete_unmatched: bool,
) {
let mut fields = vec![Field::new("id", DataType::Int64, false)];
if !delete_unmatched {
fields.extend([
Field::new("k", DataType::Int64, false),
Field::new(
"vector",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
2,
),
true,
),
Field::new("s", DataType::Utf8, true),
]);
}
let schema = Arc::new(Schema::new_with_metadata(
fields,
HashMap::from([("source".to_string(), "empty".to_string())]),
));
let batches = if has_batch {
vec![Ok(RecordBatch::new_empty(schema.clone()))]
} else {
vec![]
};
let data: Box<dyn RecordBatchReader + Send> =
Box::new(RecordBatchIterator::new(batches, schema.clone()));
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_ref = attempts.clone();
let num_deleted_rows = if delete_unmatched { 3 } else { 0 };
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
assert_eq!(request.headers()[CONTENT_TYPE], ARROW_STREAM_CONTENT_TYPE);
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
assert_eq!(params["on"], "id");
assert_eq!(
params["when_not_matched_insert_all"],
(!delete_unmatched).to_string()
);
assert_eq!(
params["when_not_matched_by_source_delete"],
delete_unmatched.to_string()
);
let body = request.body().unwrap().as_bytes().unwrap();
let reader = StreamReader::try_new(Cursor::new(body), None).unwrap();
assert_eq!(reader.schema(), schema);
for batch in reader {
assert_eq!(batch.unwrap().num_rows(), 0);
}
// The empty source must retain its schema when replayed after a conflict.
if attempts_ref.fetch_add(1, Ordering::SeqCst) == 0 {
http::Response::builder()
.status(409)
.body(String::new())
.unwrap()
} else {
http::Response::builder()
.status(200)
.body(
json!({
"version": 43,
"num_deleted_rows": num_deleted_rows,
"num_inserted_rows": 0,
"num_updated_rows": 0,
})
.to_string(),
)
.unwrap()
}
});
let mut merge = table.merge_insert(&["id"]);
if delete_unmatched {
merge.when_not_matched_by_source_delete(None);
} else {
merge.when_not_matched_insert_all();
}
let result = merge.execute(data).await.unwrap();
assert_eq!(attempts.load(Ordering::SeqCst), 2);
assert_eq!(result.version, 43);
assert_eq!(result.num_deleted_rows, num_deleted_rows);
assert_eq!(result.num_inserted_rows, 0);
assert_eq!(result.num_updated_rows, 0);
}
#[tokio::test]
async fn test_merge_insert_composite_key() {
let batch = RecordBatch::try_new(
@@ -5937,6 +6060,22 @@ mod tests {
assert_eq!(data[0].as_ref().unwrap(), &expected_data);
}
#[tokio::test]
async fn test_query_vector_nprobes_zero() {
let table = Table::new_with_handler::<&str>("my_table", |_| {
panic!("invalid nprobes must be rejected before sending a request")
});
let result = table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.nprobes(0);
assert!(matches!(
result,
Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0"
));
}
#[tokio::test]
async fn test_query_vector_all_params() {
let table = Table::new_with_handler("my_table", |request| {
@@ -6010,6 +6149,7 @@ mod tests {
.postfilter()
.distance_type(crate::DistanceType::Cosine)
.nprobes(12)
.unwrap()
.refine_factor(2)
.bypass_vector_index()
.execute()
rust/lancedb/src/table.rs
+19 -1
View File
@@ -1674,6 +1674,7 @@ impl Table {
/// .unwrap()
/// .refine_factor(5)
/// .nprobes(10)
/// .unwrap()
/// .execute()
/// .await
/// .unwrap();
@@ -1919,11 +1920,28 @@ impl Table {
self.inner.function_errors(&request).await
}
/// Change a column's name or nullability.
/// Change a column's name, data type, or nullability.
///
/// Each alteration must specify at least one of `rename`, `data_type`, or
/// `nullable`. Returns [`Error::InvalidInput`] before applying any changes
/// if an alteration does not specify any of these fields.
pub async fn alter_columns(
&self,
alterations: &[ColumnAlteration],
) -> Result<AlterColumnsResult> {
for alteration in alterations {
if alteration.rename.is_none()
&& alteration.nullable.is_none()
&& alteration.data_type.is_none()
{
return Err(Error::InvalidInput {
message: format!(
"One of rename, nullable or data_type must be specified for path '{}'",
alteration.path
),
});
}
}
self.inner.alter_columns(alterations).await
}
rust/lancedb/src/table/schema_evolution.rs
+23
View File
@@ -807,6 +807,29 @@ mod tests {
// Alter Columns Tests
#[tokio::test]
async fn test_alter_columns_rejects_missing_changes() {
let conn = connect("memory://").execute().await.unwrap();
let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap();
let table = conn.create_table("test", batch).execute().await.unwrap();
let initial_version = table.version().await.unwrap();
let initial_schema = table.schema().await.unwrap();
for alterations in [
vec![ColumnAlteration::new("id".into())],
vec![
ColumnAlteration::new("id".into()).rename("new_id".into()),
ColumnAlteration::new("id".into()),
],
] {
let err = table.alter_columns(&alterations).await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}");
assert!(err.to_string().contains("path 'id'"));
assert_eq!(table.version().await.unwrap(), initial_version);
assert_eq!(table.schema().await.unwrap(), initial_schema);
}
}
#[tokio::test]
async fn test_alter_column_rename() {
let conn = connect("memory://").execute().await.unwrap();