diff --git a/docs/src/js/classes/AutoQuery.md b/docs/src/js/classes/AutoQuery.md index 79d666124..12f517dbd 100644 --- a/docs/src/js/classes/AutoQuery.md +++ b/docs/src/js/classes/AutoQuery.md @@ -417,6 +417,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/docs/src/js/classes/Query.md b/docs/src/js/classes/Query.md index 0abf49755..3038bf2b3 100644 --- a/docs/src/js/classes/Query.md +++ b/docs/src/js/classes/Query.md @@ -491,6 +491,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/docs/src/js/classes/QueryBase.md b/docs/src/js/classes/QueryBase.md index 664337bad..0ef4fa286 100644 --- a/docs/src/js/classes/QueryBase.md +++ b/docs/src/js/classes/QueryBase.md @@ -252,6 +252,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + *** ### withRowId() diff --git a/docs/src/js/classes/TakeQuery.md b/docs/src/js/classes/TakeQuery.md index 79a53f753..f3d59a2a3 100644 --- a/docs/src/js/classes/TakeQuery.md +++ b/docs/src/js/classes/TakeQuery.md @@ -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) diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index 29749c28f..f49963d60 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -740,6 +740,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/nodejs/__test__/query.test.ts b/nodejs/__test__/query.test.ts index 5f3e68b16..2cfecf823 100644 --- a/nodejs/__test__/query.test.ts +++ b/nodejs/__test__/query.test.ts @@ -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; diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 32d8cccea..f9e50fe51 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -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[] = []; diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 75a787fcd..7c0dc9221 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -41,26 +41,6 @@ export async function* RecordBatchIterator( } } -class RecordBatchIterable< - NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery, -> implements AsyncIterable -{ - 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, 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): Promise { 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); } diff --git a/nodejs/src/iterator.rs b/nodejs/src/iterator.rs index 8d65dfc3f..96d08c765 100644 --- a/nodejs/src/iterator.rs +++ b/nodejs/src/iterator.rs @@ -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 { + schema_to_buffer(&self.inner.schema()) + } + #[napi(catch_unwind)] pub async unsafe fn next(&mut self) -> napi::Result> { if let Some(rst) = self.inner.next().await { diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 60e1780f0..330514484 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -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 diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index cf1910c68..16be63e2d 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -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 diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index a8eb9fce2..a7a689d33 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -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() diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 2bdd57fe9..9d951f559 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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 diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index aeb9feeb8..df5be5a6d 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -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( diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index e3d55917d..6c59a9010 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -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): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 7a3b90125..03595285f 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -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]) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index fa1a494b0..0c12c5c20 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -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( diff --git a/python/src/table.rs b/python/src/table.rs index 5d91d3060..48c6bd0c7 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1903,6 +1903,17 @@ impl Table { let alterations = alterations .iter() .map(|alteration| { + for key in alteration.keys().iter() { + let key = key.extract::()?; + 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"))? diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 985195c07..1d0a4a0ab 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -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::>() + .await + .unwrap(); + assert_eq!(batches, vec![expected_data]); + } + } + #[tokio::test] async fn test_clone_table() { let tmp_dir = tempdir().unwrap(); diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index dc1627116..8255fef4b 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -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 diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index a68facc96..8f28ea3bd 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3111,7 +3111,13 @@ impl BaseTable for RemoteTable { // 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::, _>>()?; + let schema = RecordBatchReader::schema(new_data.as_ref()); + let mut batches = new_data.collect::, _>>()?; + // 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 = Box::new(batches); let rescannable = source.rescannable(); let input: Arc = @@ -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::("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 = + 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::>(); + 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( diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index a99f086ef..a95d8ddf3 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -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 { + 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 } diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index 4d0e2815d..4990901ae 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -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();