mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-20 18:52:28 +00:00
Merge main into gatekeeper/fix-4393-1
Keep remote error-parity regressions alongside upstream empty-query and async-search tests. Preserve the combined Python imports and both Node remote test groups.
This commit is contained in:
@@ -417,6 +417,8 @@ Collect the results as an Arrow
|
||||
|
||||
ArrowTable.
|
||||
|
||||
Empty results retain the query's output schema.
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.toArrow`
|
||||
|
||||
@@ -491,6 +491,8 @@ Collect the results as an Arrow
|
||||
|
||||
ArrowTable.
|
||||
|
||||
Empty results retain the query's output schema.
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.toArrow`
|
||||
|
||||
@@ -252,6 +252,8 @@ Collect the results as an Arrow
|
||||
|
||||
ArrowTable.
|
||||
|
||||
Empty results retain the query's output schema.
|
||||
|
||||
***
|
||||
|
||||
### withRowId()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -740,6 +740,8 @@ Collect the results as an Arrow
|
||||
|
||||
ArrowTable.
|
||||
|
||||
Empty results retain the query's output schema.
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.toArrow`
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -14,6 +14,14 @@ import {
|
||||
TlsConfig,
|
||||
connect,
|
||||
} from "../lancedb";
|
||||
import {
|
||||
Table as ArrowTable,
|
||||
Field,
|
||||
Int64,
|
||||
RecordBatch,
|
||||
Schema,
|
||||
tableToIPC,
|
||||
} from "../lancedb/arrow";
|
||||
import {
|
||||
HeaderProvider,
|
||||
OAuthHeaderProvider,
|
||||
@@ -138,6 +146,66 @@ describe("remote connection", () => {
|
||||
);
|
||||
});
|
||||
|
||||
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[] = [];
|
||||
|
||||
+10
-23
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
@@ -1275,6 +1316,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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -1049,6 +1049,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()
|
||||
|
||||
@@ -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",
|
||||
@@ -6107,7 +6111,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,
|
||||
@@ -6117,7 +6131,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,
|
||||
@@ -6127,7 +6141,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,
|
||||
@@ -6137,7 +6151,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,
|
||||
@@ -6147,9 +6161,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,
|
||||
@@ -6165,7 +6177,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/).
|
||||
@@ -6180,8 +6192,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.
|
||||
|
||||
@@ -6209,10 +6222,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))
|
||||
|
||||
@@ -6862,6 +6878,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
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import contextlib
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
@@ -405,6 +406,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 +539,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(
|
||||
|
||||
@@ -24,6 +24,7 @@ from lancedb.pydantic import LanceModel, Vector
|
||||
from lancedb.query import (
|
||||
AsyncFTSQuery,
|
||||
AsyncHybridQuery,
|
||||
AsyncQuery,
|
||||
AsyncQueryBase,
|
||||
AsyncVectorQuery,
|
||||
ColumnOrdering,
|
||||
@@ -1279,6 +1280,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):
|
||||
|
||||
@@ -18,7 +18,7 @@ from packaging.version import Version
|
||||
import lancedb
|
||||
from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction
|
||||
from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry
|
||||
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
|
||||
@@ -120,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):
|
||||
@@ -1679,6 +1750,46 @@ def test_query_sync_empty_query():
|
||||
assert data == expected
|
||||
|
||||
|
||||
@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 == {
|
||||
@@ -2303,15 +2414,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])
|
||||
|
||||
@@ -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(
|
||||
{
|
||||
@@ -4144,14 +4170,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 +4263,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(
|
||||
|
||||
@@ -1903,6 +1903,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"))?
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -3111,7 +3111,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> =
|
||||
@@ -4698,6 +4704,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)]
|
||||
@@ -4811,6 +4836,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(
|
||||
|
||||
@@ -1919,11 +1919,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
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user