From a0946282ca57a58e18b8cea9a50be39195b3323a Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:42:06 +0100 Subject: [PATCH 01/64] fix(python): return original scores from hybrid search with normalize="rank" (#4343) ## Summary With `rerank(..., normalize="rank")`, a hybrid query returns ranks (1.0, 2.0, 3.0, ...) in the `_distance` and `_score` columns instead of the real vector distances and BM25 scores. `_combine_hybrid_results` converted both result sets to ranks first, and only then captured `original_distances` / `original_scores`. So the step that restores the original values after reranking put the ranks back. With `normalize="score"`, the real values are returned as expected. Repro (RRFReranker(return_score="all"), vector distances 0 / 25 / 100): | normalize | `_distance` returned | `_score` returned | | --- | --- | --- | | `"score"` | 0.0, 25.0, 100.0 | 0.47, 0.67, None | | `"rank"` (before) | 1.0, 2.0, 3.0 | 1.0, 2.0, None | | `"rank"` (after) | 0.0, 25.0, 100.0 | 0.47, 0.67, None | ## Changes - Capture the original `_distance` / `_score` columns before converting them to ranks, then normalize the ranked column. The reranker input is unchanged; only the values restored afterwards are fixed. `AsyncHybridQuery` uses the same function, so both sync and async are covered. ## Test plan - [x] New `test_hybrid_return_all_restores_original_scores` in `test_rerankers.py`, parametrized over `score` / `rank`: the `rank` case fails before the fix and passes after - [x] `pytest python/tests/test_rerankers.py python/tests/test_hybrid_query.py` - [x] `ruff format --check` / `ruff check` I have no local Rust toolchain, so I ran the tests against the released 0.39.0 native module with the changed Python files overlaid. The changed functions are identical between 0.39.0 and main. Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 20 +++++++++----- python/python/tests/test_rerankers.py | 40 +++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index fe8c33861..9d7610f72 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -2575,10 +2575,6 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): with_row_ids: bool, offset: Optional[int] = None, ) -> pa.Table: - if norm == "rank": - vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance") - fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score") - # If both result sets are empty (e.g. after hard filtering), # return early to avoid errors in reranking or score restoration. if vector_results.num_rows == 0 and fts_results.num_rows == 0: @@ -2605,15 +2601,23 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): original_score_row_ids = None # normalize the scores to be between 0 and 1, 0 being most relevant # We check whether the results (vector and FTS) are empty, because when - # they are, they often are missing the _rowid column, which causes an error + # they are, they often are missing the _rowid column, which causes an error. + # The original values are captured before any rank conversion so that + # they can be restored after reranking. if vector_results.num_rows > 0: distance_i = vector_results.column_names.index("_distance") original_distances = vector_results.column(distance_i) original_distance_row_ids = vector_results.column("_rowid") + if norm == "rank": + vector_results = LanceHybridQueryBuilder._rank( + vector_results, "_distance" + ) vector_results = vector_results.set_column( distance_i, vector_results.field(distance_i), - LanceHybridQueryBuilder._normalize_scores(original_distances), + LanceHybridQueryBuilder._normalize_scores( + vector_results.column(distance_i) + ), ) # In fts higher scores represent relevance. Not inverting them here as @@ -2622,10 +2626,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): score_i = fts_results.column_names.index("_score") original_scores = fts_results.column(score_i) original_score_row_ids = fts_results.column("_rowid") + if norm == "rank": + fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score") fts_results = fts_results.set_column( score_i, fts_results.field(score_i), - LanceHybridQueryBuilder._normalize_scores(original_scores), + LanceHybridQueryBuilder._normalize_scores(fts_results.column(score_i)), ) results = reranker.rerank_hybrid(fts_query, vector_results, fts_results) diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index d5ee6b6a4..743669367 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -785,6 +785,46 @@ def test_empty_hybrid_result_reranker(): assert "_rowid" in result.column_names +@pytest.mark.parametrize("norm", ["score", "rank"]) +def test_hybrid_return_all_restores_original_scores(norm): + from lancedb.query import LanceHybridQueryBuilder + + vector_results = pa.table( + { + "text": ["a", "b", "c"], + "_rowid": pa.array([0, 1, 2], type=pa.uint64()), + "_distance": pa.array([0.5, 2.5, 10.0], type=pa.float32()), + } + ) + fts_results = pa.table( + { + "text": ["b", "a"], + "_rowid": pa.array([1, 0], type=pa.uint64()), + "_score": pa.array([3.25, 1.75], type=pa.float32()), + } + ) + + result = LanceHybridQueryBuilder._combine_hybrid_results( + fts_results=fts_results, + vector_results=vector_results, + norm=norm, + fts_query="query", + reranker=RRFReranker(return_score="all"), + limit=10, + with_row_ids=True, + ) + + # The raw distances and scores are returned, not the normalized values + # or ranks used for reranking. + rows = {row["_rowid"]: row for row in result.to_pylist()} + assert rows[0]["_distance"] == 0.5 + assert rows[1]["_distance"] == 2.5 + assert rows[2]["_distance"] == 10.0 + assert rows[0]["_score"] == 1.75 + assert rows[1]["_score"] == 3.25 + assert rows[2]["_score"] is None + + @pytest.mark.parametrize("return_score", ["relevance", "all"]) @pytest.mark.parametrize("first_row_id", [1, 2**63]) @pytest.mark.parametrize("score_type", [pa.float32(), pa.float64()]) From 5f9a8af27c66d8730b9e7cf3c29cac1f2bd237cf Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:42:21 +0100 Subject: [PATCH 02/64] fix(python): keep embedding functions registered without an alias readable (#4356) `EmbeddingFunctionRegistry.register()` documents that the class name is used when no alias is given, and it does register the class under that name. But it stored the alias itself (`None`) as `__embedding_function_registry_alias__`, which is what `function_to_metadata()` writes into the table schema. A table using such a function gets `"name": null` in its `embedding_functions` metadata, and any later read of the functions fails: ```python @get_registry().register() class MyEmbed(TextEmbeddingFunction): ... f = MyEmbed.create() class Schema(LanceModel): text: str = f.SourceField() vector: Vector(4) = f.VectorField() tbl = db.create_table("t", schema=Schema) tbl.add([{"text": "hi"}]) # KeyError: None tbl.search("hi") # KeyError: None ``` This stores the key the class is registered under instead, so the class name round-trips. For aliased classes the key is the alias, so their metadata is unchanged. The TypeScript registry already falls back to `ctor.name` in the same way. Added `test_embedding_function_registered_without_alias`, which creates a table with such a function, reopens it, and adds a row. It fails with `KeyError` before the change. `test_embeddings.py`, `test_pydantic.py`, `test_table.py` and `test_db.py` pass, and `ruff format --check` / `ruff check` are clean. --- python/python/lancedb/embeddings/registry.py | 4 ++- python/python/tests/test_embeddings.py | 31 ++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/embeddings/registry.py b/python/python/lancedb/embeddings/registry.py index 1dc7ffcb8..3b940fe2e 100644 --- a/python/python/lancedb/embeddings/registry.py +++ b/python/python/lancedb/embeddings/registry.py @@ -73,7 +73,9 @@ class EmbeddingFunctionRegistry: raise KeyError(f"{cls.__name__} was already registered") key = alias or cls.__name__ self._functions[key] = cls - cls.__embedding_function_registry_alias__ = alias + # Store the name the class is registered under, so the table + # metadata written for it can be looked up again when reading. + cls.__embedding_function_registry_alias__ = key return cls return decorator diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index c42172f77..51907f141 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -64,6 +64,37 @@ def test_embedding_function(tmp_path): assert np.allclose(actual, expected) +def test_embedding_function_registered_without_alias(tmp_path): + registry = EmbeddingFunctionRegistry.get_instance() + + @registry.register() + class NoAliasEmbeddingFunction(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + return [np.array([1.0, 2.0]) for _ in texts] + + assert registry.get("NoAliasEmbeddingFunction") is NoAliasEmbeddingFunction + + func = NoAliasEmbeddingFunction.create() + + class Schema(LanceModel): + text: str = func.SourceField() + vector: Vector(2) = func.VectorField() + + db = lancedb.connect(tmp_path) + db.create_table("test", schema=Schema) + + # The class name is written to the table metadata, so the function can be + # resolved again when the table is reopened. + table = lancedb.connect(tmp_path).open_table("test") + table.add([{"text": "hello"}]) + conf = table.embedding_functions["vector"] + assert isinstance(conf.function, NoAliasEmbeddingFunction) + assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0]] + + def test_instructor_ndims_uses_instruction(): instructor = get_registry().get("instructor").create() model = MagicMock() From 712a8b86995397de03ae1c0389f68824505f80a5 Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:42:36 +0100 Subject: [PATCH 03/64] fix(python): write plain Enum fields of a LanceModel as their values (#4357) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since #3232, `pydantic_to_schema` maps an `Enum` field to the Arrow type of its values, including enums without a `str`/`int` mixin. Writing such models still failed: `model_to_dict()` returns `model.model_dump()`, which keeps the Enum members, and Arrow can only convert members that are also `str`/`int` instances. ```python class Color(Enum): RED = "red" class Item(LanceModel): color: Color tbl = db.create_table("items", schema=Item) tbl.add([Item(color=Color.RED)]) # pyarrow.lib.ArrowTypeError: Expected bytes, got a 'Color' object ``` `create_table(data=[...])` fails the same way, and int-valued enums raise `ArrowInvalid`. `model_to_dict()` now replaces Enum members with their values, including in nested models and lists. Lists of plain scalars (vectors in particular) are returned without being walked, so the per-row cost stays about the same as `model_dump()` alone (21.3µs vs 20.3µs for a 1536-dim model). Reading back is unchanged: Pydantic already validates the stored value into the Enum member. Added `test_plain_enum_values_are_written`, which covers str and int valued enums, `Optional` and `List` fields, and an Enum inside a nested model, and checks that `to_pydantic` round-trips the rows. It fails with the `ArrowTypeError` above before the change. `test_pydantic.py`, `test_embeddings.py`, `test_table.py` and `test_db.py` pass, and `ruff format --check` / `ruff check` are clean. --- python/python/lancedb/pydantic.py | 18 +++++++++- python/python/tests/test_pydantic.py | 49 ++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index 313c514c3..8fe3c1c93 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -524,4 +524,20 @@ def get_extras(field_info: FieldInfo, key: str) -> Any: def model_to_dict(model: pydantic.BaseModel) -> dict[str, Any]: """Convert a Pydantic model to a dictionary.""" - return model.model_dump() + return _enum_members_to_values(model.model_dump()) + + +def _enum_members_to_values(value: Any) -> Any: + # model_dump() keeps Enum members, and Arrow cannot convert a member of a + # plain Enum (one without a str or int mixin). Write its value instead, + # which is what the Arrow type of the field is derived from. + if isinstance(value, Enum): + return value.value + if isinstance(value, dict): + return {k: _enum_members_to_values(v) for k, v in value.items()} + if isinstance(value, (list, tuple)) and value: + # Skip lists of plain scalars, such as vectors, without walking them. + if isinstance(value[0], (int, float, str, bytes)): + return value + return [_enum_members_to_values(v) for v in value] + return value diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 81674d735..24fb7ee89 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -733,3 +733,52 @@ def test_enum_types(): assert schema.field("priority").type == pa.int64() assert schema.field("opt_status").type == pa.dictionary(pa.int32(), pa.utf8()) assert schema.field("opt_status").nullable + + +def test_plain_enum_values_are_written(mem_db): + """Members of an Enum without a str/int mixin are written as their values.""" + + class Color(Enum): + RED = "red" + BLUE = "blue" + + class Priority(Enum): + LOW = 1 + HIGH = 2 + + class Tag(LanceModel): + color: Color + + class TestModel(LanceModel): + color: Color + priority: Priority + opt_color: Optional[Color] = None + colors: List[Color] + tag: Tag + vector: Vector(2) + + rows = [ + TestModel( + color=Color.RED, + priority=Priority.HIGH, + opt_color=Color.BLUE, + colors=[Color.BLUE, Color.RED], + tag=Tag(color=Color.BLUE), + vector=[1.0, 2.0], + ), + TestModel( + color=Color.BLUE, + priority=Priority.LOW, + colors=[], + tag=Tag(color=Color.RED), + vector=[3.0, 4.0], + ), + ] + tbl = mem_db.create_table("test", schema=TestModel) + tbl.add(rows) + + assert tbl.to_arrow().select(["color", "priority", "opt_color"]).to_pylist() == [ + {"color": "red", "priority": 2, "opt_color": "blue"}, + {"color": "blue", "priority": 1, "opt_color": None}, + ] + assert tbl.search().to_pydantic(TestModel) == rows From 4bdc8c3c0015c54587445409ad10619481843273 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:42:51 +0800 Subject: [PATCH 04/64] fix: wait for the submitted remote index job (#4428) Addresses the client-side waiting portion of #4402. When create_index returns a job ID, wait_timeout now waits for that submission rather than an already-ready index with the same name. It propagates JobFailed / JobCancelled and bounds job polling by the requested timeout. Older servers without job IDs keep name-based waiting, calls without wait_timeout remain fire-and-forget, and the existing two-hour maximum is preserved. The job wrapper is shared with create_index_async; no synchronous server-side config validation is added. Rebased on main 81d52a98, preserving the new materialized-view SQL constructor. Validation on Linux / Rust 1.97, exact head 61122f92: - All five new native HTTP-mock tests pass (DONE/FAILED/CANCELLED, pending deadline, legacy response, no wait, maximum timeout). On unchanged current-main production the job-tracking and deadline regressions fail by returning early against the old index. - All 274 remote-table tests pass, no failures/skips. - cargo fmt --all -- --check and Clippy --lib --tests -- -D warnings pass. Tests and Clippy used --offline --profile ci --locked -p lancedb --no-default-features --features remote. The affected package was rebuilt from each selected source root; dependency caches were reused. This is Rust mocked-HTTP coverage, not live-cloud, Python/Node binding, full-workspace or performance qualification. --- rust/lancedb/src/index/waiter.rs | 2 +- rust/lancedb/src/remote/table.rs | 194 +++++++++++++++++++++++++++++-- 2 files changed, 185 insertions(+), 11 deletions(-) diff --git a/rust/lancedb/src/index/waiter.rs b/rust/lancedb/src/index/waiter.rs index f201ed6c3..eeef22c98 100644 --- a/rust/lancedb/src/index/waiter.rs +++ b/rust/lancedb/src/index/waiter.rs @@ -9,7 +9,7 @@ use std::time::{Duration, Instant}; use tokio::time::sleep; const DEFAULT_SLEEP_MS: u64 = 1000; -const MAX_WAIT: Duration = Duration::from_secs(2 * 60 * 60); +pub(crate) const MAX_WAIT: Duration = Duration::from_secs(2 * 60 * 60); /// Poll the table using list_indices() and index_stats() until all of the indices have 0 un-indexed rows. /// Will return Error::Timeout if the columns are not fully indexed within the timeout. diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 1c569f6e5..c3b6e7b73 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -16,7 +16,7 @@ use crate::expr::expr_to_sql_string; use crate::index::Index; use crate::index::IndexStatistics; use crate::index::scalar::FtsQuery; -use crate::index::waiter::wait_for_index; +use crate::index::waiter::{MAX_WAIT, wait_for_index}; use crate::job::Job; use crate::materialized_view::{ MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult, @@ -650,13 +650,42 @@ impl RemoteTable { .and_then(|body| extract_job_id(&body)); if let Some(wait_timeout) = index.wait_timeout { - let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); - self.wait_for_index(&[&index_name], wait_timeout).await?; + if let Some(job_id) = &job_id { + if wait_timeout > MAX_WAIT { + return Err(Error::InvalidInput { + message: format!("timeout must be less than {:?}", MAX_WAIT), + }); + } + // An older same-named index may still be ready while its + // replacement is building. Wait for this submission instead. + tokio::time::timeout(wait_timeout, self.index_job(job_id.clone()).wait()) + .await + .map_err(|_| Error::Timeout { + message: format!( + "timed out waiting for index job {:?} after {:?}", + job_id, wait_timeout + ), + })??; + } else { + // Older servers do not return a job ID. + let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); + self.wait_for_index(&[&index_name], wait_timeout).await?; + } } Ok(job_id) } + fn index_job(&self, job_id: String) -> Job { + Job::new(Box::new(FreshnessJob { + inner: RemoteJob::new(self.client.clone(), job_id), + freshness: self.freshness.clone(), + version: self.version.clone(), + tracked_result: TrackedJobResult::None, + freshness_request: self.snapshot_freshness_headers(), + })) + } + pub(super) fn new_with_sql_client( client: RestfulLanceDbClient, name: String, @@ -3101,13 +3130,7 @@ impl BaseTable for RemoteTable { async fn create_index_async(&self, index: IndexBuilder) -> Result { Ok(match self.submit_create_index(index).await? { - Some(job_id) => Job::new(Box::new(FreshnessJob { - inner: RemoteJob::new(self.client.clone(), job_id), - freshness: self.freshness.clone(), - version: self.version.clone(), - tracked_result: TrackedJobResult::None, - freshness_request: self.snapshot_freshness_headers(), - })), + Some(job_id) => self.index_job(job_id), None => Job::new_done(), }) } @@ -8000,6 +8023,157 @@ mod tests { assert_eq!(err.to_string(), "Job job-err failed: preempted (in commit)"); } + // The name-based endpoint deliberately exposes an already-ready old index. + // A replacement must wait on the new job, not mistake that index for success. + fn index_wait_fixture( + job_id: Option<&'static str>, + terminal_state: &'static str, + ) -> ( + Table, + Arc, + Arc, + ) { + let job_polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let index_polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let jobs = job_polls.clone(); + let indices = index_polls.clone(); + let table = Table::new_with_handler("my_table", move |request| { + let body = match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + return http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(); + } + "/v1/table/my_table/create_index/" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + // Replacement is the default and omitted from the wire body. + assert!(body.get("replace").is_none()); + match job_id { + Some(id) => json!({"job_id": id}), + None => json!({}), + } + } + "/v1/jobs/describe" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["job_id"], json!(job_id.unwrap())); + let poll = jobs.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + json!({ + "job_id": job_id.unwrap(), + "job_state": if poll == 0 { "IN_PROGRESS" } else { terminal_state }, + "failure": {"phase": "build", "message": "invalid index config", + "retryable": false} + }) + } + "/v1/table/my_table/index/list/" => { + indices.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + json!({"indexes": [{"index_name": "a_idx", "columns": ["a"], + "index_uuid": "3fa85f64-5717-4562-b3fc-2c963f66afa6", + "index_status": "done"}]}) + } + "/v1/table/my_table/index/a_idx/stats/" => json!({ + "num_indexed_rows": 100, "num_unindexed_rows": 0, "index_type": "BTREE" + }), + path => panic!("Unexpected path: {path}"), + }; + http::Response::builder() + .status(200) + .body(body.to_string()) + .unwrap() + }); + (table, job_polls, index_polls) + } + + #[tokio::test] + async fn test_create_index_wait_timeout_tracks_job() { + for state in ["DONE", "FAILED", "CANCELLED"] { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), state); + let result = table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_secs(2)) + .execute() + .await; + match state { + "DONE" => result.unwrap(), + "FAILED" => match result.unwrap_err() { + Error::JobFailed { job_id, failure } => { + assert_eq!(job_id.as_deref(), Some("new-index-job")); + assert_eq!(failure.message.as_deref(), Some("invalid index config")); + assert_eq!(failure.phase.as_deref(), Some("build")); + } + error => panic!("Expected job failure, got {error:?}"), + }, + "CANCELLED" => assert!(matches!(result, Err(Error::JobCancelled { .. }))), + _ => unreachable!(), + } + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 2); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + } + + #[tokio::test] + async fn test_create_index_wait_timeout_bounds_job_polling() { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), "IN_PROGRESS"); + let error = table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_millis(10)) + .execute() + .await + .unwrap_err(); + assert!(matches!(error, Error::Timeout { message } if message.contains("new-index-job"))); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn test_create_index_wait_timeout_legacy_fallback() { + let (table, jobs, indices) = index_wait_fixture(None, "DONE"); + table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_secs(2)) + .execute() + .await + .unwrap(); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 0); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_create_index_wait_timeout_absent_does_not_poll() { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), "IN_PROGRESS"); + table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .execute() + .await + .unwrap(); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 0); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn test_create_index_wait_timeout_preserves_maximum() { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), "DONE"); + let error = table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_secs(2 * 60 * 60 + 1)) + .execute() + .await + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 0); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + /// Servers that return no job id (e.g. an empty create-index response) /// yield an already-done job. #[tokio::test] From e489d0cd7b272bab25066c25a951fd2bac8f8c15 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:43:07 +0800 Subject: [PATCH 05/64] fix: return errors for unsupported remote alteration types (#4430) RemoteTable::alter_columns currently unwraps the JsonDataType conversion. A valid Arrow input that the remote wire format cannot represent (for example Utf8View, BinaryView, or a struct containing Utf8View) therefore panics instead of returning an error. Propagate conversion/serialization failures as Error::InvalidInput, including the column path. Unsupported inputs are rejected before any HTTP request; supported alteration payloads are unchanged. This does not add server support for view types or change server HTTP-error classification. Three new native Rust cases reproduce the panic on unchanged production and return InvalidInput with the fix. The existing supported-type cases remain passing. Validation at head 88184ea, Linux/Rust 1.97.0, locked offline dependencies: - cargo test --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib remote::table::tests:: (269 passed) - cargo clippy --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib --tests -- -D warnings - cargo fmt --all -- --check Current main merges cleanly; the intervening remote-table change is unrelated nprobes validation. Live Cloud, rebuilt Python/Node bindings and full-workspace integration were not tested. AI-assisted; native regression and before/after results independently checked. --- rust/lancedb/src/remote/table.rs | 37 +++++++++++++++++++++++++++----- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index c3b6e7b73..9e2641c22 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3736,16 +3736,25 @@ impl BaseTable for RemoteTable { value["rename"] = serde_json::Value::String(rename.clone()); } if let Some(data_type) = &alteration.data_type { - let json_data_type = JsonDataType::try_from(data_type).unwrap(); - let json_data_type = serde_json::to_value(&json_data_type).unwrap(); - value["data_type"] = json_data_type; + let json_data_type = + JsonDataType::try_from(data_type).map_err(|err| Error::InvalidInput { + message: format!( + "Cannot serialize data type for column '{}': {err}", + alteration.path + ), + })?; + value["data_type"] = serde_json::to_value(&json_data_type).map_err(|err| { + Error::InvalidInput { + message: format!("Cannot serialize column alteration: {err}"), + } + })?; } if let Some(nullable) = &alteration.nullable { value["nullable"] = serde_json::Value::Bool(*nullable); } - value + Ok(value) }) - .collect::>(); + .collect::>>()?; let mut body = serde_json::json!({ "alterations": body }); self.apply_branch_body(&mut body); let request = self @@ -5067,6 +5076,24 @@ mod tests { } } + #[rstest] + #[case::string_view(DataType::Utf8View)] + #[case::binary_view(DataType::BinaryView)] + #[case::nested_view(DataType::Struct(vec![Field::new("value", DataType::Utf8View, true)].into()))] + #[tokio::test] + async fn test_alter_columns_unsupported_type(#[case] data_type: DataType) { + let table = Table::new_with_handler("my_table", |_| -> http::Response { + panic!("an unserializable alteration must not reach the server") + }); + let err = table + .alter_columns(&[ColumnAlteration::new("value".into()).cast_to(data_type)]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("value")); + assert!(err.to_string().contains("Unsupported type")); + } + #[rstest] #[case(true)] #[case(false)] From cc9f01ae2622d2a3726608b4ff875d4478daffd4 Mon Sep 17 00:00:00 2001 From: Preston Zhao Date: Thu, 22 Oct 2026 13:43:21 -0500 Subject: [PATCH 06/64] fix(node): accept exist_ok in createTable (#4432) TypeScript CreateTableOptions.mode omitted exist_ok even though createTable already forwards that mode to the native client. closes: #2900 --------- Co-authored-by: Preston Zhao <195942578+canyang25@users.noreply.github.com> --- docs/src/js/interfaces/CreateTableOptions.md | 5 ++++- nodejs/__test__/connection.test.ts | 14 ++++++++++++++ nodejs/lancedb/connection.ts | 5 ++++- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/docs/src/js/interfaces/CreateTableOptions.md b/docs/src/js/interfaces/CreateTableOptions.md index f46cf3722..e6119e74c 100644 --- a/docs/src/js/interfaces/CreateTableOptions.md +++ b/docs/src/js/interfaces/CreateTableOptions.md @@ -65,7 +65,7 @@ then no error will be raised. ### mode ```ts -mode: "overwrite" | "create"; +mode: "overwrite" | "create" | "exist_ok"; ``` The mode to use when creating the table. @@ -76,6 +76,9 @@ happen. Any provided data will be ignored. If this is set to "overwrite" then any existing table will be replaced. +If this is set to "exist_ok" then an existing table is opened and any +provided data is ignored. + *** ### schema? diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 895c2a7c5..1167ae34a 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -117,6 +117,20 @@ describe("given a connection", () => { await expect(tbl.countRows()).resolves.toBe(1); }); + it("should keep the existing table when mode is exist_ok", async () => { + const created = await db.createTable("test", [{ id: 1 }, { id: 2 }], { + mode: "exist_ok", + }); + await expect(created.countRows()).resolves.toBe(2); + + const reopened = await db.createTable("test", [{ id: 9 }], { + mode: "exist_ok", + }); + await expect(reopened.countRows()).resolves.toBe(2); + const rows = (await reopened.query().toArray()) as { id: number }[]; + expect(rows.map((row) => row.id).sort()).toEqual([1, 2]); + }); + it("should respect limit and page token when listing tables", async () => { const db = await connect(tmpDir.name); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 11341a8c6..4ed511b86 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -53,8 +53,11 @@ export interface CreateTableOptions { * happen. Any provided data will be ignored. * * If this is set to "overwrite" then any existing table will be replaced. + * + * If this is set to "exist_ok" then an existing table is opened and any + * provided data is ignored. */ - mode: "create" | "overwrite"; + mode: "create" | "overwrite" | "exist_ok"; /** * If this is true and the table already exists and the mode is "create" * then no error will be raised. From 13df747f180da94c3a487067e2a8cec8839a5158 Mon Sep 17 00:00:00 2001 From: Preston Zhao Date: Thu, 22 Oct 2026 13:43:36 -0500 Subject: [PATCH 07/64] fix(python): return the session supplied to LanceDBConnection (#4433) LanceDBConnection.session read AsyncConnection.session, but AsyncConnection has no session attribute, so access raised AttributeError. Store the session passed to the constructor and return that object. Assisted by Cursor (Grok). Co-authored-by: Preston Zhao <195942578+canyang25@users.noreply.github.com> --- python/python/lancedb/db.py | 11 ++++++++++- python/python/tests/test_db.py | 3 +++ python/python/tests/test_session.py | 10 ++++++++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index e1e5d270a..a74763afb 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1300,6 +1300,10 @@ class LanceDBConnection(DBConnection): self.storage_options = storage_options self._manifest_enabled = manifest_enabled self._namespace_client_properties = namespace_client_properties + # The native connection does not expose the Python Session wrapper. + # Keep the object supplied by the caller so this property stays + # synchronous for debugger inspection. + self._session = session if _inner is not None: self._conn = _inner # Native-derived wrappers resolve this in their async reconstruction @@ -1366,7 +1370,12 @@ class LanceDBConnection(DBConnection): @property def session(self) -> Optional[Session]: - return self._conn.session + """Return the session supplied when this connection was opened. + + None when the caller did not pass a session, including connections + reconstructed with :meth:`from_inner`. + """ + return self._session @property def uri(self) -> str: diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index 7d37910a3..af63987aa 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -86,6 +86,7 @@ def test_sync_debugger_inspection_does_not_use_background_loop(tmp_path, monkeyp assert members["uri"] == str(tmp_path) assert members["read_consistency_interval"] is None + assert members["session"] is None assert repr(db) == f"LanceDBConnection(uri={str(tmp_path)!r})" assert repr(table) == f"LanceTable(name='test', _conn={db!r})" @@ -105,6 +106,8 @@ def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkey assert db.read_consistency_interval == consistency_interval assert db_from_inner.read_consistency_interval == consistency_interval + assert db.session is None + assert db_from_inner.session is None def test_ingest_pd(tmp_path): diff --git a/python/python/tests/test_session.py b/python/python/tests/test_session.py index 8923477c7..d0bc54210 100644 --- a/python/python/tests/test_session.py +++ b/python/python/tests/test_session.py @@ -36,3 +36,13 @@ def test_session_cache_configuration(tmp_path): assert final_cache_size > initial_cache_size # Cache should have grown assert final_cache_items >= initial_cache_items # Items should not decrease assert initial_cache_size < index_cache_size + metadata_cache_size + + +def test_connection_session_returns_supplied_session(tmp_path): + """Sync connections must expose the session they were opened with.""" + db = lancedb.connect(tmp_path) + assert db.session is None + + session = lancedb.Session() + db = lancedb.connect(tmp_path, session=session) + assert db.session is session From 1e313d3c0d828d04f6c7cb3036d896942c48c86c Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:43:50 +0800 Subject: [PATCH 08/64] fix(python): parse decimal exponent literals exactly (#4449) Decimal literals such as Decimal("1E-7") and Decimal("1E+3") currently raise ValueError while building an expression. Parse the coefficient and exponent exactly, and normalize representable negative scales to zero so both typed filters and predicate.to_sql() queries preserve the value. Checked arithmetic bounds the Decimal128 coefficient to 38 digits, preserves existing scale limits, and rejects nonfinite or overflowing values. This does not add Decimal256 literal support. Added 21 cases covering actual Arrow-backed table queries, SQL queries, zero, signed exponents, precision boundaries, and errors. Both Python 3.11 and 3.13 expression suites pass 123 tests; the baseline fails 10 of the new cases. Python binding Rust tests pass 13; official make doctest passes 77 with 6 skips; Ruff, Rustfmt, hooks, Pyright and workspace/tests/all-features Clippy with -D warnings pass. A fresh environment restored from the original dependency freeze passes the full required suite: 1528 passed, 37 skipped, 68 deselected (143.51s); while optional Torch suites have the same baseline remote mock-server ordering failure. Tested against f4544628; observed official main d864e719 contains one Catalog API commit without expr or Cargo.lock changes. AI assistance: gpt-6.1-sol assisted implementation and local validation; all reported checks were actually executed. --- python/python/tests/test_expr.py | 61 ++++++++++++++++++++++++++++++++ python/src/expr.rs | 58 ++++++++++++++++++++++++------ 2 files changed, 109 insertions(+), 10 deletions(-) diff --git a/python/python/tests/test_expr.py b/python/python/tests/test_expr.py index 0f49231f1..6cfc09922 100644 --- a/python/python/tests/test_expr.py +++ b/python/python/tests/test_expr.py @@ -80,6 +80,67 @@ class TestExprConstruction: e = lit(d) assert isinstance(e, Expr) + @pytest.mark.parametrize( + "value, decimal_type", + [ + ("1E-7", pa.decimal128(12, 7)), + ("-1.23E-7", pa.decimal128(12, 9)), + ("1E+3", pa.decimal128(12, 0)), + ("-1.23E+5", pa.decimal128(12, 0)), + ("0E-7", pa.decimal128(12, 7)), + ("-0E+3", pa.decimal128(12, 0)), + ("1E-38", pa.decimal128(38, 38)), + ("1E+37", pa.decimal128(38, 0)), + ("1.234567890123456789012345678901234567E+37", pa.decimal128(38, 0)), + ("0E+128", pa.decimal128(1, 0)), + ], + ) + def test_lit_decimal_exponent_query(self, tmp_path, value, decimal_type): + expected = Decimal(value) + table = lancedb.connect(tmp_path).create_table( + "decimals", + pa.table( + { + "id": [1, 2], + # Arrow's Python ingestion rejects 0E+128 even though its + # numeric value is representable. Store canonical zero to + # exercise the expression literal independently of ingestion. + "price": pa.array( + [Decimal(0) if expected.is_zero() else expected, None], + type=decimal_type, + ), + } + ), + ) + predicate = col("price") == lit(expected) + result = table.search().where(predicate).to_arrow() + assert result["id"].to_pylist() == [1] + assert result["price"].to_pylist() == [expected] + # Exercise the SQL serializer and parser as well as the typed query path. + sql_result = table.search().where(predicate.to_sql()).to_arrow() + assert sql_result["id"].to_pylist() == [1] + assert sql_result["price"].to_pylist() == [expected] + + @pytest.mark.parametrize( + "value", + [ + "NaN", + "sNaN", + "Infinity", + "-Infinity", + "1E-39", + "1E+38", + "1E+129", + "1E+1000000000", + "0E+129", + "12345678901234567890123456789012345678E+1", + "9" * 39, + ], + ) + def test_lit_decimal_unrepresentable_raises(self, value): + with pytest.raises(ValueError): + lit(Decimal(value)) + class TestExprOperators: def test_eq_operator(self): diff --git a/python/src/expr.rs b/python/src/expr.rs index 79b448fdf..bb779bc1d 100644 --- a/python/src/expr.rs +++ b/python/src/expr.rs @@ -240,22 +240,60 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult { fn parse_decimal(s: &str) -> PyResult<(i128, u8, i8)> { let s = s.trim(); - let dot_pos = s.find('.'); - let scale = if let Some(pos) = dot_pos { - (s.len() - pos - 1) as i8 - } else { - 0 + let (mantissa, exponent) = match s.split_once(['e', 'E']) { + Some((mantissa, exponent)) => ( + mantissa, + exponent.parse::().map_err(|e| { + PyValueError::new_err(format!("failed to parse decimal exponent: {}", e)) + })?, + ), + None => (s, 0), }; + let fractional_digits = mantissa.find('.').map_or(0, |pos| mantissa.len() - pos - 1); + // Keep the coefficient intact instead of expanding an exponent into zeros. + // Arrow Decimal128 has at most 38 digits and an i8 scale (at most 38). + let mut scale = i64::try_from(fractional_digits) + .ok() + .and_then(|fractional_digits| fractional_digits.checked_sub(exponent)) + .and_then(|scale| i8::try_from(scale).ok()) + .filter(|scale| *scale <= 38) + .ok_or_else(|| PyValueError::new_err("decimal scale is outside Decimal128 range"))?; - let digits = s.replace('.', ""); - let val = digits + let digits = mantissa.replace('.', ""); + let mut val = digits .parse::() .map_err(|e| PyValueError::new_err(format!("failed to parse decimal digits: {}", e)))?; - // Precision is total number of digits - let precision = digits.trim_start_matches('-').len() as u8; + // Positive scale cannot exceed precision, including for tiny values such as 1E-38. + let mut precision = digits + .trim_start_matches('-') + .len() + .max(usize::try_from(scale).unwrap_or(0)); + if scale < 0 { + // Arrow cannot cast SQL text to a decimal with negative scale. Normalize + // positive exponents within Decimal128's bounded coefficient range. + if val != 0 { + let shift = -i16::from(scale) as u32; + precision += shift as usize; + if precision > 38 { + return Err(PyValueError::new_err( + "decimal precision exceeds Decimal128 maximum of 38", + )); + } + val = 10_i128 + .checked_pow(shift) + .and_then(|factor| val.checked_mul(factor)) + .ok_or_else(|| PyValueError::new_err("decimal coefficient exceeds i128 range"))?; + } + scale = 0; + } + if precision > 38 { + return Err(PyValueError::new_err( + "decimal precision exceeds Decimal128 maximum of 38", + )); + } - Ok((val, precision, scale)) + Ok((val, precision as u8, scale)) } /// Call an arbitrary registered SQL function by name. From fcd26c425a8657501b9a438fd8ab35f086bc5943 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:44:05 +0800 Subject: [PATCH 09/64] fix(python): preserve integer Vector and MultiVector values (#4450) When `Vector` or `MultiVector` declares an integer Arrow value type, Pydantic currently validates every element as a float. This silently rounds integers above `2**53`, and converts `uint64`'s maximum value into an out-of-range float before table insertion. Use Pydantic's integer schema for integer Arrow value types in both validators. Floating-point types retain the existing validation. Tests exercise model serialization and real local table creation, insertion, and reads for signed and unsigned integer boundaries and floating-point controls, plus dimensions and optional/nullability behavior. Fractional floats are rejected for integer declarations; integral floats remain accepted. This does not extend ANN index support. Validation on macOS arm64, with the current main native extension built using Rust 1.97.0: - Python 3.11.14 and 3.13.11, Pydantic 2.13.5: full `pytest -m "not slow and not s3_test" -q python/tests` suite: **1,537 passed, 37 skipped, 68 deselected** on each runtime. Comparable unchanged-main baselines: **1,521 passed, 37 skipped, 68 deselected** each. - Added cases on unchanged main: **8 failed, 8 passed**; complete `test_pydantic.py` after the fix: **40 passed** on each runtime. - Documentation tests with the existing embeddings dependencies in an isolated environment, Python 3.11.14: **77 passed, 6 skipped**, matching the unchanged-main baseline. The workflow's Linux/Python 3.13 doctest matrix was not reproduced locally. - Repository Ruff lint/format, Pyright, all applicable pre-commit hooks, compile checks, and `git diff --check` passed. Eight public SDK table round trips passed on each final installed wheel runtime. Slow, credential-dependent S3 integration tests, the remaining OS/runtime matrix, and the minimum Pydantic/PyArrow dependency matrix were not run locally. The native extension was built from unchanged current main; the final local wheel repack contains the verified Python change with that identical native extension. A separate candidate maturin packaging attempt triggered Rust recompilation and was stopped, so an official candidate distribution build was not completed. The project requires Pydantic `>=2.7.4,<3`; Pydantic v1 support was removed in #3990. AI assistance: developed and verified with **gpt-6.1-sol**. --- python/python/lancedb/pydantic.py | 12 +++++- python/python/tests/test_pydantic.py | 61 ++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index 8fe3c1c93..bb70c9a67 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -127,7 +127,11 @@ def Vector( core_schema.list_schema( min_length=dim, max_length=dim, - items_schema=core_schema.float_schema(), + items_schema=( + core_schema.int_schema() + if pa.types.is_integer(value_type) + else core_schema.float_schema() + ), ), ) @@ -208,7 +212,11 @@ def MultiVector( items_schema=core_schema.list_schema( min_length=dim, max_length=dim, - items_schema=core_schema.float_schema(), + items_schema=( + core_schema.int_schema() + if pa.types.is_integer(value_type) + else core_schema.float_schema() + ), ), ), ) diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 24fb7ee89..6717bbc32 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -415,6 +415,67 @@ def test_nullable_vector(): assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)]) +@pytest.mark.parametrize("vector_type", [Vector, MultiVector]) +@pytest.mark.parametrize( + "value_type, values", + [ + (pa.int64(), [-(2**53 + 1), 2**53 + 3]), + (pa.uint64(), [2**53 + 1, 2**64 - 1]), + (pa.int8(), [-128, 127]), + (pa.uint8(), [0, 255]), + (pa.float32(), [1.25, 2.5]), + (pa.float64(), [1.25, 2.5]), + ], +) +def test_typed_vector_round_trip(tmp_path, vector_type, value_type, values): + import lancedb + + annotation = vector_type(2, value_type=value_type) + expected = [values, values] if vector_type is MultiVector else values + + class Row(LanceModel): + id: int + payload: annotation + + row = Row(id=1, payload=expected) + table = lancedb.connect(tmp_path).create_table("typed_vectors", [row], schema=Row) + table.add([Row(id=2, payload=expected)]) + assert table.schema == Row.to_arrow_schema() + assert table.search().to_arrow().to_pylist() == [ + {"id": 1, "payload": expected}, + {"id": 2, "payload": expected}, + ] + assert json.loads(row.model_dump_json())["payload"] == expected + + +@pytest.mark.parametrize("vector_type", [Vector, MultiVector]) +@pytest.mark.parametrize("nullable", [False, True]) +def test_integer_vector_validation(vector_type, nullable): + annotation = vector_type(2, value_type=pa.int64(), nullable=nullable) + + class Row(LanceModel): + payload: annotation + + def shape(values): + return [values] if vector_type is MultiVector else values + + assert Row.to_arrow_schema().field("payload").nullable is nullable + assert Row(payload=shape([1.0, -2.0])).model_dump()["payload"] == shape([1, -2]) + with pytest.raises(pydantic.ValidationError, match="valid integer"): + Row(payload=shape([1.5, 2])) + with pytest.raises(pydantic.ValidationError): + Row(payload=shape([1])) + with pytest.raises(pydantic.ValidationError): + Row(payload=shape([1, 2, 3])) + with pytest.raises(pydantic.ValidationError): + Row(payload=None) + + class OptionalRow(LanceModel): + payload: Optional[annotation] = None + + assert OptionalRow(payload=None).payload is None + + def test_bare_vector_raises_clear_error(): namespace = { "__name__": "test_model_without_pyarrow", From 92ada619d2c5aac50c987534a4ebaafd295828c2 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:44:20 +0800 Subject: [PATCH 10/64] fix(python): preserve metadata during vector schema inference (#4451) Creating a table from an Arrow `list` or `large_list` vector field infers a fixed-size vector type, but currently drops that field's metadata. Descriptions and tags such as `lancedb:description` and `lancedb:tag:model` disappear even though the input field carries them. Preserve the original field metadata while replacing its inferred type. Tests use public synchronous and asynchronous table creation, reopening, and reads for floating-point and integer inference, checking the complete schema and row values. Existing dtype inference and nullability remain unchanged. Validation on macOS arm64 with an official native extension built from main `d864e719` using Rust 1.97.0: - Python 3.11.14 and 3.13.11: full `pytest -m "not slow and not s3_test" -q -rs python/tests`: **1,533 passed, 37 skipped, 68 deselected** each, wall times 331.464 s and 416.308 s. - Added public SDK cases on unchanged main: **12 failed**; after the fix: **12 passed** on each runtime. Final installed distributions passed **8 public SDK round trips** each, including fixed-size fields and fields without metadata. - Documentation tests on Python 3.11.14: **77 passed, 6 skipped**. Ruff lint/format, Pyright (**0 errors**), applicable pre-commit hooks (**4 passed, 2 without matching files**), compile checks and `git diff --check` passed. - Comparable unchanged-main full baselines from the same day, native extension and dependency versions: **1,521 passed, 37 skipped, 68 deselected** on each runtime. These were reused rather than rerun. The final local wheel is an explicitly verified repack of the official main wheel, replacing only `table.py` and `RECORD`; all production Python files match the tested source and the native extension is unchanged. No candidate Rust rebuild was performed. Slow/S3 integration tests, the remaining OS/runtime and minimum-dependency matrices, and the workflow's Linux/Python 3.13 doctest job were not run locally. The existing Rust migration draft #2970 and issues #3216/#3222 do not fix this metadata loss in the current Python path; a future migration should preserve the same metadata. During final verification main advanced to `81d52a98` through #4373 and the beta.13 version bump. Neither modified the two files here; the local native and full-suite validation above remain explicitly based on `d864e719`, and the updated native changes have not been tested locally. AI assistance: developed and verified with **gpt-6.1-sol**. --- python/python/lancedb/table.py | 1 + python/python/tests/test_table.py | 52 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 35ec7a466..9b96a1b2a 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5404,6 +5404,7 @@ def _infer_target_schema( field.name, # preserve original field name target_type, nullable=field.nullable, + metadata=field.metadata, ) schema = schema.set(i, new_field) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index b241adb0a..b2d499933 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4507,6 +4507,58 @@ async def test_alter_columns_rejects_invalid_async( assert await table.schema() == initial_schema +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("list_type", [pa.list_, pa.large_list]) +@pytest.mark.parametrize( + "value_type, values, inferred_type", + [ + (pa.float64(), [1.25, 2.5], pa.float32()), + (pa.int32(), [0, 255], pa.uint8()), + (pa.int32(), [-1, 256], pa.float32()), + ], +) +async def test_inferred_vector_field_metadata( + tmp_db: DBConnection, + tmp_db_async: AsyncConnection, + asynchronous, + list_type, + value_type, + values, + inferred_type, +): + metadata = { + "lancedb:description": "Embedding values", + "lancedb:tag:model": "test-model", + } + schema = pa.schema( + [ + pa.field( + "vector", list_type(value_type), nullable=False, metadata=metadata + ), + pa.field("text", pa.string(), metadata={"description": "Source text"}), + ], + metadata={"description": "Source dataset"}, + ) + rows = [{"vector": values, "text": "first"}, {"vector": values, "text": "second"}] + data = pa.Table.from_pylist(rows, schema=schema) + if asynchronous: + await tmp_db_async.create_table("vector_metadata", data=data) + table = await tmp_db_async.open_table("vector_metadata") + actual_schema = await table.schema() + actual_data = await table.to_arrow() + else: + tmp_db.create_table("vector_metadata", data=data) + table = tmp_db.open_table("vector_metadata") + actual_schema = table.schema + actual_data = table.to_arrow() + expected_schema = schema.set( + 0, schema.field("vector").with_type(pa.list_(inferred_type, 2)) + ) + assert actual_schema.equals(expected_schema, check_metadata=True) + assert actual_data.to_pylist() == rows + + 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) From b9ef048abbf62840befb3727efb6f0e2ed5932ca Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:44:34 +0800 Subject: [PATCH 11/64] fix: validate Python read consistency intervals (#4454) Negative `read_consistency_interval` values currently reach Rust's panicking `Duration::from_secs_f64` conversion. For example: ```python from datetime import timedelta import lancedb lancedb.connect("./db", read_consistency_interval=timedelta(seconds=-1)) ``` Validate the interval with `Duration::try_from_secs_f64` at the three native connection entry points and return `ValueError` before constructing the connection or async future. `None`, zero and positive intervals retain their behavior. The raw native binding also rejects NaN, infinity and overflow; the public Python API still accepts `timedelta`, not floats. Tests cover public sync/async connections, native namespace connections, real Python namespace clients, the three raw binding entry points and valid-value controls. Validation on Python 3.12 / Rust 1.97: - Same 26 regression cases against separately built native baseline and candidate: baseline 22 failed / 4 passed; candidate 26 passed. - Existing basic connection and background-loop regression tests pass. - `cargo fmt --all -- --check`, Ruff check/format and `git diff --check` pass. The native builds used the unchanged repository `ci` profile and Cargo.lock with `--no-default-features --features remote,lancedb/metrics-otel,pyo3/extension-module`. Python runtime and focused test dependency versions came from `python/uv.lock`. Source-tree tests used a minimal distribution-version metadata fixture derived from the unchanged Cargo/pyproject manifests; both imported module origins and native versions were checked in fresh processes. This is not a complete Maturin packaging, cloud backend or full Python suite result. No GPU was used. AI-assisted diagnosis, implementation and validation. Signed-off-by: Xucheng Zhou --- .../tests/test_read_consistency_interval.py | 83 +++++++++++++++++++ python/src/connection.rs | 18 +++- 2 files changed, 98 insertions(+), 3 deletions(-) create mode 100644 python/python/tests/test_read_consistency_interval.py diff --git a/python/python/tests/test_read_consistency_interval.py b/python/python/tests/test_read_consistency_interval.py new file mode 100644 index 000000000..bc69775ea --- /dev/null +++ b/python/python/tests/test_read_consistency_interval.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import timedelta + +import lancedb +import pytest +from lance_namespace import connect as connect_namespace_client +from lancedb import _lancedb +from lancedb.namespace import LanceNamespaceDBConnection + + +@pytest.mark.parametrize("seconds", [-1.0, -0.000001]) +@pytest.mark.parametrize("kind", ["sync", "namespace", "namespace_async", "client"]) +def test_negative_read_consistency_interval(tmp_path, seconds, kind): + interval = timedelta(seconds=seconds) + with pytest.raises(ValueError, match="read_consistency_interval"): + if kind == "sync": + lancedb.connect(tmp_path, read_consistency_interval=interval) + elif kind == "client": + client = connect_namespace_client("dir", {"root": str(tmp_path)}) + LanceNamespaceDBConnection(client, read_consistency_interval=interval) + else: + connect = ( + lancedb.connect_namespace + if kind == "namespace" + else lancedb.connect_namespace_async + ) + connect("dir", {"root": str(tmp_path)}, read_consistency_interval=interval) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [-1.0, -0.000001]) +async def test_negative_async_read_consistency_interval(tmp_path, seconds): + with pytest.raises(ValueError, match="read_consistency_interval"): + await lancedb.connect_async( + tmp_path, read_consistency_interval=timedelta(seconds=seconds) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [float("nan"), float("inf"), -float("inf"), 1e30]) +@pytest.mark.parametrize("kind", ["connect", "namespace", "client"]) +async def test_binding_rejects_invalid_read_consistency_interval( + tmp_path, seconds, kind +): + with pytest.raises(ValueError, match="read_consistency_interval"): + if kind == "connect": + await _lancedb.connect(str(tmp_path), read_consistency_interval=seconds) + elif kind == "namespace": + _lancedb.connect_namespace( + "dir", {"root": str(tmp_path)}, read_consistency_interval=seconds + ) + else: + client = connect_namespace_client("dir", {"root": str(tmp_path)}) + _lancedb.connect_namespace_client(client, read_consistency_interval=seconds) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "interval", [None, timedelta(0), timedelta(microseconds=1), timedelta(seconds=5)] +) +async def test_valid_read_consistency_intervals(tmp_path, interval): + db = lancedb.connect(tmp_path / "sync", read_consistency_interval=interval) + assert db.read_consistency_interval == interval + db.close() + + db = await lancedb.connect_async( + tmp_path / "async", read_consistency_interval=interval + ) + assert await db.get_read_consistency_interval() == interval + db.close() + + db = lancedb.connect_namespace( + "dir", {"root": str(tmp_path / "namespace")}, read_consistency_interval=interval + ) + assert db.read_consistency_interval == interval + db.close() + + client = connect_namespace_client("dir", {"root": str(tmp_path / "client")}) + db = LanceNamespaceDBConnection(client, read_consistency_interval=interval) + assert db.read_consistency_interval == interval + db.close() diff --git a/python/src/connection.rs b/python/src/connection.rs index 7e20f5f99..baeb29780 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -49,6 +49,18 @@ impl Connection { } } +fn parse_read_consistency_interval(interval: Option) -> PyResult> { + interval + .map(|seconds| { + Duration::try_from_secs_f64(seconds).map_err(|_| { + PyValueError::new_err( + "read_consistency_interval must be finite, non-negative, and within the supported range", + ) + }) + }) + .transpose() +} + fn parse_namespace_client_pushdown_operations( operations: Option>, ) -> PyResult> { @@ -1077,6 +1089,7 @@ pub fn connect( namespace_client_properties: Option>, oauth_config: Option, ) -> PyResult> { + let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?; future_into_py(py, async move { let mut builder = lancedb::connect(&uri); if let Some(api_key) = api_key { @@ -1095,7 +1108,6 @@ pub fn connect( #[cfg(not(feature = "remote"))] let _ = sql_host_override; if let Some(read_consistency_interval) = read_consistency_interval { - let read_consistency_interval = Duration::from_secs_f64(read_consistency_interval); builder = builder.read_consistency_interval(read_consistency_interval); } if let Some(storage_options) = storage_options { @@ -1144,7 +1156,7 @@ pub fn connect_namespace_client( namespace_client_impl: Option, namespace_client_properties: Option>, ) -> PyResult { - let read_consistency_interval = read_consistency_interval.map(Duration::from_secs_f64); + let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?; let namespace_client_pushdown_operations = parse_namespace_client_pushdown_operations(namespace_client_pushdown_operations)?; let ns_properties = namespace_client_properties.unwrap_or_default(); @@ -1201,7 +1213,7 @@ pub fn connect_namespace( session: Option, namespace_client_pushdown_operations: Option>, ) -> PyResult { - let read_consistency_interval = read_consistency_interval.map(Duration::from_secs_f64); + let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?; let namespace_client_pushdown_operations = parse_namespace_client_pushdown_operations(namespace_client_pushdown_operations)?; From a12aef2a8ce39b7d85c28033bd7eff0d1e60b1da Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:44:48 +0800 Subject: [PATCH 12/64] fix: return a Python error when querying a closed table (#4455) `AsyncTable.query()` unwraps the native closed-table error, so querying a closed table raises `PanicException` instead of the documented `RuntimeError`. `AsyncTable.head()` reaches the same path. Propagate `inner_ref()`'s error through `PyResult` instead. Add regressions for closed `query()` and `head()`, with `count_rows()`, repeated close, and normal open-table queries as controls. Validation against separately built native bindings: - Before: 2 regression failures (`PanicException`), 2 controls passed. - After: all 4 regressions passed. - Existing query tests plus the basic database test: 102 passed, 1 skipped on each arm (Torch is not installed). - `cargo fmt --all -- --check`, Ruff, and `git diff --check` passed. The native build used the locked Cargo dependencies with `remote,lancedb/metrics-otel,pyo3/extension-module`. This is source-tree native testing, not a complete wheel-packaging or cloud-service test. Signed-off-by: Xucheng Zhou --- .../python/tests/test_closed_table_query.py | 32 +++++++++++++++++++ python/src/table.rs | 4 +-- 2 files changed, 34 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_closed_table_query.py diff --git a/python/python/tests/test_closed_table_query.py b/python/python/tests/test_closed_table_query.py new file mode 100644 index 000000000..2ff3a8e98 --- /dev/null +++ b/python/python/tests/test_closed_table_query.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pytest +from lancedb.db import AsyncConnection + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["count_rows", "query", "head"]) +async def test_closed_table_queries_raise_runtime_error( + mem_db_async: AsyncConnection, operation +): + table = await mem_db_async.create_table("closed_query", data=[{"id": 0}]) + assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}] + table.close() + table.close() + assert not table.is_open() + + with pytest.raises(RuntimeError, match="Table closed_query is closed"): + if operation == "query": + table.query() + else: + await getattr(table, operation)() + + +@pytest.mark.asyncio +async def test_open_table_queries(mem_db_async: AsyncConnection): + table = await mem_db_async.create_table("open_query", data=[{"id": 0}]) + assert await table.count_rows() == 1 + assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}] + assert (await table.head()).to_pylist() == [{"id": 0}] + table.close() diff --git a/python/src/table.rs b/python/src/table.rs index c1d23d151..27ed8be8b 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1428,8 +1428,8 @@ impl Table { }) } - pub fn query(&self) -> Query { - Query::new(self.inner_ref().unwrap().query()) + pub fn query(&self) -> PyResult { + Ok(Query::new(self.inner_ref()?.query())) } #[getter] From 7a3709b3768f3d1a496e1baf5832493de25e6c87 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:45:02 +0800 Subject: [PATCH 13/64] fix(remote): preserve request IDs in multipart response errors (#4456) Keep the request ID when a multipart create/complete response has a missing or incorrectly typed `upload_id`/`version`. These errors currently use an empty ID, unlike JSON parse errors, so callers lose the ID needed to correlate the failed request. Add six cases through `Table.add(...).write_parallelism(2).execute()`, capturing the actual outgoing `x-request-id`: four missing/wrong-type cases fail on unchanged main, while two invalid-JSON controls already pass. After the fix all six pass, and all 15 multipart tests pass. Native Linux validation (locked dependencies, `ci` profile): - `cargo test --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib remote::table::tests::test_multipart_write`: 15 passed. - `cargo fmt --all -- --check`: passed. - `cargo clippy --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib --tests -- -D warnings`: passed. The committed source bytes match the tested source. HTTP endpoints are mocked; this does not claim a live Cloud or full-repository test run. --- rust/lancedb/src/remote/table.rs | 69 ++++++++++++++++++++++++++++++-- 1 file changed, 65 insertions(+), 4 deletions(-) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 9e2641c22..84c1b3fef 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1277,7 +1277,7 @@ impl RemoteTable { let body = response.text().await.err_to_http(request_id.clone())?; let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { source: format!("Failed to parse multipart create response: {}", e).into(), - request_id, + request_id: request_id.clone(), status_code: None, })?; parsed["upload_id"] @@ -1285,7 +1285,7 @@ impl RemoteTable { .map(|s| s.to_string()) .ok_or_else(|| Error::Http { source: "Missing upload_id in multipart create response".into(), - request_id: String::new(), + request_id, status_code: None, }) } @@ -1304,12 +1304,12 @@ impl RemoteTable { let body = response.text().await.err_to_http(request_id.clone())?; let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { source: format!("Failed to parse multipart complete response: {}", e).into(), - request_id, + request_id: request_id.clone(), status_code: None, })?; let version = parsed["version"].as_u64().ok_or_else(|| Error::Http { source: "Missing version in multipart complete response".into(), - request_id: String::new(), + request_id, status_code: None, })?; Ok(AddResult { version }) @@ -12596,6 +12596,67 @@ mod tests { ); } + #[rstest] + #[case("create", "{}")] + #[case("create", r#"{"upload_id": 1}"#)] + #[case("create", "not json")] + #[case("complete", "{}")] + #[case("complete", r#"{"version": "invalid"}"#)] + #[case("complete", "not json")] + #[tokio::test] + async fn test_multipart_write_error_preserves_request_id( + #[case] failed_stage: &str, + #[case] body: &str, + ) { + let failed_stage = failed_stage.to_string(); + let body = body.to_string(); + let request_id = Arc::new(std::sync::Mutex::new(String::new())); + let observed_request_id = request_id.clone(); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 4, 0), + move |request| { + let path = request.url().path(); + if path == "/v1/table/my_table/describe/" { + return simple_describe_response(); + } + if path == format!("/v1/table/my_table/multipart_write/{failed_stage}") { + *observed_request_id.lock().unwrap() = request.headers()["x-request-id"] + .to_str() + .unwrap() + .to_string(); + return http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap(); + } + let response = match path { + "/v1/table/my_table/multipart_write/create" => r#"{"upload_id":"upload-1"}"#, + "/v1/table/my_table/insert/" => r#"{"version":1}"#, + "/v1/table/my_table/multipart_write/abort" => "", + _ => panic!("Unexpected request path: {path}"), + }; + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }, + ); + let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap(); + let error = table + .add(vec![batch]) + .write_parallelism(2) + .execute() + .await + .unwrap_err(); + let expected = request_id.lock().unwrap(); + assert!(!expected.is_empty()); + match error { + Error::Http { request_id, .. } => assert_eq!(request_id, *expected), + other => panic!("Expected Http error, got {other:?}"), + } + } + #[tokio::test] async fn test_multipart_write_progress() { let callback_count = Arc::new(AtomicUsize::new(0)); From 9c161869e1b90f3e7733ff851d113f04935c309f Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:45:19 +0800 Subject: [PATCH 14/64] fix(node): respect transformers tokenizer option (#4457) ## Summary `tokenizer` is a documented option, but initialization always passed the model name to `AutoTokenizer.from_pretrained`. Preserve the resolved tokenizer option, fall back to the model name when it is omitted, and identify the selected tokenizer in loading errors. Add five initialization regressions covering default model/tokenizer selection, an explicit tokenizer, registry variables, and tokenizer loading errors. They run in native Node child processes so the real ESM import works without changing Jest/CI flags. Model loaders are mocked and local/remote model access is disabled; these tests do not download models or test inference. ## Validation - Before: three regressions fail and two default-selection controls pass. After: all five pass. - New tests plus the existing registry and embedding-registry suites: 48 passed. - TypeScript compilation, changed-file Biome checks, and `pnpm run docs` pass; generated tracked docs are unchanged. - Checked with Node 24.14.1 / pnpm 11.1.1 on Linux, both with default Jest flags and the Linux CI ESM flag. The native binding and declarations came from the official `0.40.0-beta.13` packages, with registry integrity verified. Rebuilding Rust from source timed out during Cargo metadata resolution, before compilation; no Rust code is changed. --- .../fixtures/transformers_tokenizer.cjs | 66 +++++++++++++++++++ nodejs/__test__/transformers.test.ts | 26 ++++++++ nodejs/lancedb/embedding/transformers.ts | 6 +- 3 files changed, 96 insertions(+), 2 deletions(-) create mode 100644 nodejs/__test__/fixtures/transformers_tokenizer.cjs create mode 100644 nodejs/__test__/transformers.test.ts diff --git a/nodejs/__test__/fixtures/transformers_tokenizer.cjs b/nodejs/__test__/fixtures/transformers_tokenizer.cjs new file mode 100644 index 000000000..bf30cc229 --- /dev/null +++ b/nodejs/__test__/fixtures/transformers_tokenizer.cjs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +const assert = require("node:assert/strict"); +const { + TransformersEmbeddingFunction, +} = require("../../dist/embedding/transformers"); +const { getRegistry } = require("../../dist/embedding/registry"); + +async function main() { + const transformers = await import("@huggingface/transformers"); + // Never fetch a real model, including if a loader mock stops intercepting. + transformers.env.allowRemoteModels = false; + transformers.env.allowLocalModels = false; + const modelCalls = []; + const tokenizerCalls = []; + const testCase = Number(process.argv[2]); + const defaults = "Xenova/all-MiniLM-L6-v2"; + const cases = [ + [{}, defaults, defaults], + [{ model: "test/model" }, "test/model", "test/model"], + [ + { model: "test/model", tokenizer: "test/tokenizer" }, + "test/model", + "test/tokenizer", + ], + [ + { model: "test/model", tokenizer: "$var:tokenizer" }, + "test/model", + "test/tokenizer", + ], + [ + { model: "test/model", tokenizer: "test/tokenizer" }, + "test/model", + "test/tokenizer", + ], + ]; + assert.ok( + Number.isInteger(testCase) && testCase >= 0 && testCase < cases.length, + ); + transformers.AutoModel.from_pretrained = async (...args) => { + modelCalls.push(args); + return {}; + }; + transformers.AutoTokenizer.from_pretrained = async (...args) => { + tokenizerCalls.push(args); + if (testCase === 4) throw new Error("missing tokenizer"); + return {}; + }; + getRegistry().setVar("tokenizer", "test/tokenizer"); + const [options, model, tokenizer] = cases[testCase]; + const fn = new TransformersEmbeddingFunction(options); + if (testCase === 4) { + await assert.rejects(fn.init(), /error loading tokenizer test\/tokenizer/); + } else { + await fn.init(); + } + assert.deepEqual(modelCalls, [[model, { dtype: "fp32" }]]); + assert.deepEqual(tokenizerCalls, [[tokenizer]]); + assert.deepEqual(fn.toJSON(), options); +} + +main().catch((error) => { + console.error(error); + process.exitCode = 1; +}); diff --git a/nodejs/__test__/transformers.test.ts b/nodejs/__test__/transformers.test.ts new file mode 100644 index 000000000..ac66c1499 --- /dev/null +++ b/nodejs/__test__/transformers.test.ts @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import { execFileSync } from "node:child_process"; +import { resolve } from "node:path"; + +describe("Transformers embedding tokenizer selection", () => { + // Run init's native ESM import outside Jest's CommonJS sandbox, like a user. + // The fixture mocks model loading and forbids local/remote model access. + test.each([ + ["default model", 0], + ["default tokenizer", 1], + ["custom tokenizer", 2], + ["tokenizer registry variable", 3], + ["tokenizer loading error", 4], + ])("handles %s", (_name, testCase) => { + execFileSync( + process.execPath, + [ + resolve(__dirname, "fixtures", "transformers_tokenizer.cjs"), + String(testCase), + ], + { stdio: "pipe", timeout: 15_000 }, + ); + }); +}); diff --git a/nodejs/lancedb/embedding/transformers.ts b/nodejs/lancedb/embedding/transformers.ts index 06043ea7c..7e57bc6bc 100644 --- a/nodejs/lancedb/embedding/transformers.ts +++ b/nodejs/lancedb/embedding/transformers.ts @@ -38,6 +38,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction< #model?: import("@huggingface/transformers").PreTrainedModel; #tokenizer?: import("@huggingface/transformers").PreTrainedTokenizer; #modelName: XenovaTransformerOptions["model"]; + #tokenizerName: string; #initialized = false; #tokenizerOptions: XenovaTransformerOptions["tokenizerOptions"]; #ndims?: number; @@ -58,6 +59,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction< this.#ndims = options.ndims; this.#modelName = modelName; + this.#tokenizerName = options.tokenizer ?? modelName; } async init() { @@ -86,11 +88,11 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction< } try { this.#tokenizer = await transformers.AutoTokenizer.from_pretrained( - this.#modelName, + this.#tokenizerName, ); } catch (e) { throw new Error( - `error loading tokenizer for ${this.#modelName}. Make sure you are using a wasm compatible model:\nReason: ${e}`, + `error loading tokenizer ${this.#tokenizerName}. Make sure you are using a wasm compatible model:\nReason: ${e}`, ); } this.#initialized = true; From e1e87e1cae05a39ccb736c516405649c1ce92380 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:45:36 +0800 Subject: [PATCH 15/64] fix(python): preserve logical rows from Hugging Face datasets (#4461) Hugging Face `select`, `filter`, and `shuffle` can retain a physical Arrow table while storing logical row selection/order in an indices table. Ingesting the physical table then silently writes excluded rows or the wrong order and can lose duplicate selections. Read indexed datasets through bounded Arrow-format iteration, and keep the existing direct batch path for datasets without indices. Apply the same conversion to sync/async table creation and additions, including DatasetDict splits. The conversion does not mutate the user's formatting settings. Regression coverage includes reordered/duplicate selections, filtering, shuffling, empty selections, projected numpy/pandas/Arrow formatting, DatasetDicts, rescanning, and bounded batches. Validation on base d864e7198fe55a4e5fbe747fee574737fa102aab with its byte-identical native/config: - Python 3.11 full CI test selection: 1564 passed, 36 skipped, 68 deselected; 162.51s pytest /171.68s wall. - Python 3.13 full CI test selection: 1521 passed, 37 skipped, 68 deselected; 160.32s pytest /163.74s wall. Optional HF tests were skipped in this environment. - Python 3.13 complete doctest command: 71 passed, 6 skipped; 8.42s pytest /12.14s wall. - Python 3.11 HF regression module from a candidate wheel: 43 passed; the same regression module on the unchanged baseline gives 24 failed /19 passed. - Ruff check/format, repository pre-commit hooks, pyright with the source import path, and `git diff --check` passed. This is a Python-only change. Native binaries were reused from the verified baseline; no new native build or external-service integration tests were run. The separate schema-metadata preservation change in #4451 addresses a different conversion defect. The target branch currently includes newer native changes than the verified base. Latest-main native compatibility, Linux CI, and remote-service tests remain for CI/maintainer validation. Prepared with AI assistance (gpt-6.1-sol). --- python/python/lancedb/scannable.py | 20 ++++- python/python/lancedb/table.py | 6 +- python/python/tests/test_huggingface.py | 103 ++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 5 deletions(-) diff --git a/python/python/lancedb/scannable.py b/python/python/lancedb/scannable.py index baf8e0c2b..9a6f71a5d 100644 --- a/python/python/lancedb/scannable.py +++ b/python/python/lancedb/scannable.py @@ -276,6 +276,14 @@ def _from_iterable(data: Iterator) -> Scannable: _registered_modules: set[str] = set() +def _hf_batches(data): + if data._indices is None: + yield from data.data.to_batches() + else: + for table in data.with_format("arrow").iter(batch_size=1000): + yield from table.to_batches() + + def _register_optional_converters(): """Register converters for optional dependencies that are already imported.""" @@ -323,7 +331,15 @@ def _register_optional_converters(): def _from_hf_dataset(data: HFDataset) -> Scannable: table = data.data.table # Access underlying Arrow table return Scannable( - schema=table.schema, num_rows=len(data), reader=table.to_reader + schema=table.schema, + num_rows=len(data), + reader=( + table.to_reader + if data._indices is None + else lambda: pa.RecordBatchReader.from_batches( + table.schema, _hf_batches(data) + ) + ), ) @to_scannable.register(HFDatasetDict) @@ -335,7 +351,7 @@ def _register_optional_converters(): def gen(): for split_name, dataset in data.items(): - for batch in dataset.data.to_batches(): + for batch in _hf_batches(dataset): split_arr = pa.array( [split_name] * len(batch), type=pa.string() ) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 9b96a1b2a..f9cfac6e1 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -28,7 +28,7 @@ from typing import ( ) from urllib.parse import urlparse -from lancedb.scannable import _register_optional_converters, to_scannable +from lancedb.scannable import _hf_batches, _register_optional_converters, to_scannable from . import __version__ from ._blob import ( @@ -400,7 +400,7 @@ def _into_pyarrow_reader( if _check_for_hugging_face(data): if isinstance(data, datasets.Dataset): schema = data.features.arrow_schema - return pa.RecordBatchReader.from_batches(schema, data.data.to_batches()) + return pa.RecordBatchReader.from_batches(schema, _hf_batches(data)) elif isinstance(data, datasets.dataset_dict.DatasetDict): schema = _schema_from_hf(data, None) if "split" not in schema.names: @@ -1195,7 +1195,7 @@ def _to_batches_with_split(data): with an extra `split` column """ for key, dataset in data.items(): - for batch in dataset.data.to_batches(): + for batch in _hf_batches(dataset): table = pa.Table.from_batches([batch]) if "split" not in table.column_names: table = table.append_column( diff --git a/python/python/tests/test_huggingface.py b/python/python/tests/test_huggingface.py index 9b597bbc9..91421b79c 100644 --- a/python/python/tests/test_huggingface.py +++ b/python/python/tests/test_huggingface.py @@ -129,3 +129,106 @@ def test_generator(tmp_path: Path): assert len(tbl) == 2 assert tbl.schema == ds.features.arrow_schema + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["create", "add"]) +@pytest.mark.parametrize( + "selection", + ["normal", "contiguous", "selected", "duplicates", "filtered", "shuffled", "empty"], +) +async def test_write_selected_hf_dataset(tmp_path, asynchronous, operation, selection): + source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")}) + data = { + "normal": source, + "contiguous": source.select(range(1, 4)), + "selected": source.select([4, 1]), + "duplicates": source.select([2, 2, 0]), + "filtered": source.filter(lambda row: row["id"] % 2 == 1), + "shuffled": source.shuffle(seed=42), + "empty": source.select([]), + }[selection] + expected = data.to_list() + if asynchronous: + db = await lancedb.connect_async(tmp_path) + if operation == "create": + table = await db.create_table("selected", data) + else: + table = await db.create_table("selected", schema=data.features.arrow_schema) + await table.add(data) + actual = await table.to_arrow() + else: + db = lancedb.connect(tmp_path) + if operation == "create": + table = db.create_table("selected", data) + else: + table = db.create_table("selected", schema=data.features.arrow_schema) + table.add(data) + actual = table.to_arrow() + assert actual.to_pylist() == expected + assert actual.schema.equals(data.features.arrow_schema, check_metadata=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["create", "add"]) +async def test_write_selected_hf_splits(tmp_path, asynchronous, operation): + source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")}) + data = datasets.DatasetDict( + { + "train": source.select([4, 1]), + "test": source.select([3, 0, 2]), + "empty": source.select([]), + } + ) + expected = [ + dict(row, split=split) + for split, dataset in data.items() + for row in dataset.to_list() + ] + schema = source.features.arrow_schema.append(pa.field("split", pa.string())) + if asynchronous: + db = await lancedb.connect_async(tmp_path) + if operation == "create": + table = await db.create_table("splits", data) + else: + table = await db.create_table("splits", schema=schema) + await table.add(data) + actual = await table.to_arrow() + else: + db = lancedb.connect(tmp_path) + if operation == "create": + table = db.create_table("splits", data) + else: + table = db.create_table("splits", schema=schema) + table.add(data) + actual = table.to_arrow() + assert actual.to_pylist() == expected + assert actual.schema.equals(schema, check_metadata=True) + + +@pytest.mark.parametrize("format_type", [None, "numpy", "pandas", "arrow"]) +@pytest.mark.parametrize("empty", [False, True]) +def test_selected_hf_reader_is_rescannable(format_type, empty): + from lancedb.scannable import _register_optional_converters, to_scannable + + source = datasets.Dataset.from_dict( + {"id": list(range(2105)), "text": [str(i) for i in range(2105)]} + ) + indices = [] if empty else list(reversed(range(2105))) + [1, 1] + data = source.select(indices).with_format(format_type, columns=["id"]) + original_format = data.format + expected = source.select(indices).to_list() + _register_optional_converters() + scannable = to_scannable(data) + assert scannable.rescannable + assert scannable.num_rows == len(expected) + for _ in range(2): + batches = list(scannable.reader()) + if not empty: + assert all(batch.num_rows <= 1000 for batch in batches) + actual = pa.Table.from_batches(batches, schema=scannable.schema) + assert actual.to_pylist() == expected + assert actual.schema.equals(data.features.arrow_schema, check_metadata=True) + assert data.format == original_format From 8552fd2f5b8213ce6be5fdd9d3cdfe7978f10a1d Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:45:51 +0800 Subject: [PATCH 16/64] fix(python): prevent embedding aliases from replacing registered functions (#4462) Registering a different embedding class under an existing alias silently replaced the original class. A deterministic real local SDK table reproduction changed the top result for the same query from alpha to beta after reopening the persisted table. Check the resolved key before inserting a class. Preserve repeated registration of the same class under the same nonempty explicit alias as idempotent, as existing SDK tests rely on it. Repeated implicit class-name registration, including an empty alias, still raises KeyError. Failed registration preserves the mapping and class metadata. Existing embedding tests receive isolated registry mappings; test_util.py remains unchanged. Eight regression cases cover different-class alias collisions, alias/class-name collisions, implicit and empty-alias duplicates, equal class names under distinct aliases, explicit-alias idempotence, and real local creation/reopening/search. This is independent of #4356, which fixes no-alias serialization. Validation on 52a1ccee5a5b8dbfff8568abafb57099c55383b9: - Original registry plus eight new cases: 4 failed, 4 passed; Python 3.11.14; exit 1; 2.902604 s. Candidate plus full utility module: 83 passed, 1 skipped; exit 0; 3.144414 s. - Python 3.11.14 full official non-slow/non-S3 suite: 1532 passed, 36 skipped, 68 deselected; exit 0; 152.631543 s. - Python 3.13.11 full official non-slow/non-S3 suite: 1529 passed, 37 skipped, 68 deselected; exit 0; 163.970571 s. - Official Python 3.13 doctest with existing compatible embeddings/torch/mlx dependencies: 71 passed, 6 skipped; exit 0; 9.635034 s. - Pyright with matching interpreter/source path: 0 errors, 0 warnings; exit 0; 4.299743 s. - Repository Ruff 0.9.9 check and format: exit 0; 157 files already formatted; 0.047418 s and 0.016815 s. Pre-commit changed files: four hooks passed, two inapplicable skipped; exit 0; 0.406136 s. git diff --check passed. Base is official d864e7198fe55a4e5fbe747fee574737fa102aab with its verified matching native extension reused read-only. All 174 native/config paths are unchanged from that base. Latest main bbc4aebe55b69471fa1bd7655599b8d2e3f9352c has byte-identical registry production code; its complete native build was not run. Verification was macOS arm64, with optional dependency/service skips above; Linux/Windows CI remains remote verification. Exact commands, cwd, identities, raw logs and intermediate failures are in the local manifest. The earlier strict same-class check caused 13 utility failures and was corrected in production; it is not described as a baseline failure. Reproduction commands after preparing the repository's declared environment (from repository root unless noted): ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/embeddings/registry.py python/python/tests/test_embedding_registry.py python/python/tests/test_embeddings.py git diff --check ``` These relative commands describe reproducible source/cwd equivalents. The recorded local runs used absolute paths to existing Python 3.11.14/3.13.11 environments, a verified matching reused native extension, and read-only compatible optional-dependency cache paths for the 3.13 doctest, rather than a newly installed `.venv`. Full tests used separate task-private working directories and `--basetemp` directories to avoid cross-run database collisions. Implemented with assistance from gpt-6.1-sol; the final diff was reviewed locally by the collaborating agents. --- python/python/lancedb/embeddings/registry.py | 6 +- .../python/tests/test_embedding_registry.py | 88 +++++++++++++++++++ python/python/tests/test_embeddings.py | 6 ++ 3 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_embedding_registry.py diff --git a/python/python/lancedb/embeddings/registry.py b/python/python/lancedb/embeddings/registry.py index 3b940fe2e..728ff36da 100644 --- a/python/python/lancedb/embeddings/registry.py +++ b/python/python/lancedb/embeddings/registry.py @@ -69,9 +69,11 @@ class EmbeddingFunctionRegistry: def decorator(cls): if not issubclass(cls, EmbeddingFunction): raise TypeError("Must be a subclass of EmbeddingFunction") - if cls.__name__ in self._functions: - raise KeyError(f"{cls.__name__} was already registered") key = alias or cls.__name__ + if key in self._functions and ( + not alias or self._functions[key] is not cls + ): + raise KeyError(f"{key} was already registered") self._functions[key] = cls # Store the name the class is registered under, so the table # metadata written for it can be looked up again when reading. diff --git a/python/python/tests/test_embedding_registry.py b/python/python/tests/test_embedding_registry.py new file mode 100644 index 000000000..07c10bd7e --- /dev/null +++ b/python/python/tests/test_embedding_registry.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import lancedb +import pytest +from lancedb.embeddings.base import TextEmbeddingFunction +from lancedb.embeddings.registry import EmbeddingFunctionRegistry, get_registry +from lancedb.pydantic import LanceModel, Vector + + +class FirstEmbedding(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts] + + +class SecondEmbedding(FirstEmbedding): + def generate_embeddings(self, texts): + return [[0.0, 1.0] if text == "alpha" else [1.0, 0.0] for text in texts] + + +@pytest.mark.parametrize( + "first_alias,second_alias,second_class", + [ + ("shared", "shared", SecondEmbedding), + (None, "FirstEmbedding", SecondEmbedding), + ("SecondEmbedding", None, SecondEmbedding), + (None, None, FirstEmbedding), + ("", "", FirstEmbedding), + ], +) +def test_duplicate_registration_preserves_original( + first_alias, second_alias, second_class +): + registry = EmbeddingFunctionRegistry() + registry.register(first_alias)(FirstEmbedding) + key = first_alias or "FirstEmbedding" + original_alias = FirstEmbedding.__embedding_function_registry_alias__ + second_attributes = dict(second_class.__dict__) + with pytest.raises(KeyError, match=f"{key} was already registered"): + registry.register(second_alias)(second_class) + assert registry.get(key) is FirstEmbedding + assert FirstEmbedding.__embedding_function_registry_alias__ == original_alias + assert dict(second_class.__dict__) == second_attributes + + +def test_class_name_does_not_collide_with_distinct_alias(): + registry = EmbeddingFunctionRegistry() + registry.register()(FirstEmbedding) + same_name = type("FirstEmbedding", (SecondEmbedding,), {}) + registry.register("separate")(same_name) + assert registry.get("FirstEmbedding") is FirstEmbedding + assert registry.get("separate") is same_name + + +def test_same_class_explicit_alias_registration_is_idempotent(): + registry = EmbeddingFunctionRegistry() + assert registry.register("shared")(FirstEmbedding) is FirstEmbedding + attributes = dict(FirstEmbedding.__dict__) + assert registry.register("shared")(FirstEmbedding) is FirstEmbedding + assert registry.get("shared") is FirstEmbedding + assert dict(FirstEmbedding.__dict__) == attributes + + +def test_alias_collision_does_not_change_persisted_table(tmp_path, monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + alias = "test-alias-collision-persisted-table" + registry.register(alias)(FirstEmbedding) + function = FirstEmbedding.create(max_retries=0) + + class Row(LanceModel): + text: str = function.SourceField() + vector: Vector(2) = function.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table( + "rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}] + ) + assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + assert registry.register(alias)(FirstEmbedding) is FirstEmbedding + with pytest.raises(KeyError, match=f"{alias} was already registered"): + registry.register(alias)(SecondEmbedding) + reopened = lancedb.connect(tmp_path).open_table("rows") + assert isinstance(reopened.embedding_functions["vector"].function, FirstEmbedding) + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 51907f141..d6d128845 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -23,6 +23,12 @@ from lancedb.embeddings.utils import retry from lancedb.pydantic import LanceModel, Vector +@pytest.fixture(autouse=True) +def isolated_embedding_registry(monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + + def mock_embed_func(input_data): return [np.random.randn(128).tolist() for _ in range(len(input_data))] From 554dc01f4d653f95bf044059e1f88ad7e3e71457 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:46:06 +0800 Subject: [PATCH 17/64] fix(python): preserve custom embedding field metadata (#4463) `SourceField(json_schema_extra={...})` and `VectorField(json_schema_extra={...})` currently raise `TypeError` because the helper supplies the same `Field` keyword twice. Explicit `None` also fails. This prevents users from adding ordinary Pydantic JSON schema metadata to embedding fields. Copy and merge the supplied dictionary before attaching the helper's internal embedding identity. Caller dictionaries remain unchanged, and reserved embedding keys retain the real function. Other `Field` arguments keep their existing behavior. This change covers dictionaries and `None`; callable `json_schema_extra` remains unsupported by the existing embedding metadata reader and is outside this fix. The 11 regression cases exercise both helpers, empty/custom/None metadata, reserved-key collisions, other Field arguments, and a deterministic local embedding model through JSON schema generation, table creation, reopening, and search. On the original source, 9 fail and 2 controls pass; the candidate passes all 11 on Python 3.11 and 3.13. Validation uses read-only reused environments and the native extension built for d864e719. The production helper is byte-identical to current main bbc4aebe before this fix; no current-main native rebuild was performed. The branch contains only this helper fix and its regression test, independently based on d864e719. Reproducible checkout commands below assume the documented environment is installed as `.venv`; the actual local runs used existing interpreters through absolute paths and isolated generated-data working directories: ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) git diff --check .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/embeddings/base.py python/python/tests/test_embedding_field_metadata.py ``` Doctests require the repository's tests/dev/embeddings dependencies, including mlx on this platform. All checks below ran against frozen head 3fa6820. Implemented with assistance from gpt-6.1-sol; the final diff was reviewed locally. Frozen-head validation (all exit 0): - Python 3.11 full: 1535 passed, 36 skipped, 68 deselected; 190.272415 seconds wall time. - Python 3.13 full: 1532 passed, 37 skipped, 68 deselected; 249.289778 seconds wall time. - Python 3.13 doctests: 71 passed, 6 skipped; 10.391002 seconds wall time. - Pyright: 0 errors, 0 warnings; 4.884641 seconds wall time. - pre-commit: 4 passed, 2 skipped; 0.747764 seconds wall time. - Repository Ruff check: passed; 0.078177 seconds wall time. - Repository Ruff format: 157 files already formatted; 0.034561 seconds wall time. --- python/python/lancedb/embeddings/base.py | 8 +- .../tests/test_embedding_field_metadata.py | 84 +++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_embedding_field_metadata.py diff --git a/python/python/lancedb/embeddings/base.py b/python/python/lancedb/embeddings/base.py index f26a940df..96ca6556c 100644 --- a/python/python/lancedb/embeddings/base.py +++ b/python/python/lancedb/embeddings/base.py @@ -239,7 +239,9 @@ class EmbeddingFunction(BaseModel, ABC): Creates a pydantic Field that can automatically annotate the source column for this embedding function """ - return Field(json_schema_extra={"source_column_for": self}, **kwargs) + json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {}) + json_schema_extra["source_column_for"] = self + return Field(json_schema_extra=json_schema_extra, **kwargs) def VectorField(self, **kwargs): """ @@ -250,7 +252,9 @@ class EmbeddingFunction(BaseModel, ABC): static type checkers. LanceDB will infer the fixed vector dimension from this embedding function. """ - return Field(json_schema_extra={"vector_column_for": self}, **kwargs) + json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {}) + json_schema_extra["vector_column_for"] = self + return Field(json_schema_extra=json_schema_extra, **kwargs) def __eq__(self, __value: object) -> bool: if not hasattr(__value, "__dict__"): diff --git a/python/python/tests/test_embedding_field_metadata.py b/python/python/tests/test_embedding_field_metadata.py new file mode 100644 index 000000000..68754ee35 --- /dev/null +++ b/python/python/tests/test_embedding_field_metadata.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import copy + +import lancedb +import pytest +from lancedb.embeddings.base import TextEmbeddingFunction +from lancedb.embeddings.registry import get_registry +from lancedb.pydantic import LanceModel, Vector + + +class LocalEmbedding(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts] + + +@pytest.mark.parametrize( + "method,tag", + [("SourceField", "source_column_for"), ("VectorField", "vector_column_for")], +) +@pytest.mark.parametrize("extra", [None, {}, {"x-purpose": "retrieval"}]) +def test_embedding_field_preserves_extra_metadata(method, tag, extra): + function = LocalEmbedding.create(max_retries=0) + original = copy.deepcopy(extra) + field = getattr(function, method)(json_schema_extra=extra) + assert extra == original + assert field.json_schema_extra[tag] is function + for key, value in (extra or {}).items(): + assert field.json_schema_extra[key] == value + + +@pytest.mark.parametrize( + "method,tag", + [("SourceField", "source_column_for"), ("VectorField", "vector_column_for")], +) +def test_extra_metadata_cannot_replace_embedding_identity(method, tag): + function = LocalEmbedding.create(max_retries=0) + supplied = {tag: "caller value", "x-purpose": "retrieval"} + field = getattr(function, method)(json_schema_extra=supplied) + assert supplied[tag] == "caller value" + assert field.json_schema_extra[tag] is function + assert field.json_schema_extra["x-purpose"] == "retrieval" + + +@pytest.mark.parametrize("method", ["SourceField", "VectorField"]) +def test_embedding_field_preserves_other_field_arguments(method): + function = LocalEmbedding.create(max_retries=0) + field = getattr(function, method)(title="Custom title", description="Custom field") + assert field.title == "Custom title" + assert field.description == "Custom field" + + +def test_extra_metadata_preserves_schema_and_local_embedding(tmp_path, monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + registry.register("test-embedding-field-extra")(LocalEmbedding) + function = LocalEmbedding.create(max_retries=0) + + class Row(LanceModel): + text: str = function.SourceField(json_schema_extra={"x-purpose": "source"}) + vector: Vector(2) = function.VectorField( + json_schema_extra={"x-purpose": "embedding"} + ) + + json_schema = Row.model_json_schema() + assert json_schema["properties"]["text"]["x-purpose"] == "source" + assert json_schema["properties"]["vector"]["x-purpose"] == "embedding" + configs = Row.parse_embedding_functions() + assert len(configs) == 1 + assert configs[0].function is function + assert configs[0].source_column == "text" + assert configs[0].vector_column == "vector" + db = lancedb.connect(tmp_path) + table = db.create_table( + "rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}] + ) + assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + reopened = lancedb.connect(tmp_path).open_table("rows") + assert isinstance(reopened.embedding_functions["vector"].function, LocalEmbedding) + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" From 694942a1e0dcc85927ecb6c6362fb4003af8a50a Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:46:21 +0800 Subject: [PATCH 18/64] fix(python): convert tuple and NumPy datetime update values (#4464) `Table.update` and `AsyncTable.update` convert Python values into SQL literals. Tuple arrays and `numpy.datetime64` values currently raise `NotImplementedError`, even when the same values can be ingested into the corresponding local Arrow columns. Register tuples with the existing recursive list converter and format NumPy timestamps as complete nanosecond ISO strings with an explicit UTC offset. `NaT` becomes SQL NULL; subnanosecond units raise an explicit error rather than silently losing precision. Regression tests exercise local synchronous and asynchronous create/update/Arrow reads, date and timestamp units, nanosecond precision, NULL, and existing list/ndarray/Python datetime behavior. Nested and empty tuple arrays are compared with equivalent lists on the actual SDK: successful results must match, and the matching d864 native engine's existing list limitations must produce the same stable error without changing either table. This does not change SQL map handling. Implemented with assistance from gpt-6.1-sol; final diff reviewed locally. Local verification uses Python 3.11.14 and 3.13.11 in existing read-only development environments with the native extension matching base d864e7198fe55a4e5fbe747fee574737fa102aab. The current main `util.py` is byte-identical to that base; current-main native code has not been rebuilt. Slow tests, S3 tests, remote production services and other operating systems are not validated here. Reproducible commands after preparing the repository's documented development environment (these relative `.venv` examples describe equivalent commands; the local runs use absolute read-only environment paths and fresh task-private working directories): ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests/test_util_update_values.py -q -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/util.py python/python/tests/test_util_update_values.py git diff --check d864e7198fe55a4e5fbe747fee574737fa102aab HEAD ``` Frozen commit `33e74d0ed3d90a20419816a3f1fc15ff038db452` passed all local gates: Python 3.11.14 full suite 1,590 passed / 36 skipped / 68 deselected; Python 3.13.11 full suite 1,587 passed / 37 skipped / 68 deselected. All 66 added SDK cases passed on both runtimes. Python 3.13 doctests: 71 passed / 6 skipped. Pyright: 0 errors / 0 warnings. Ruff, formatting, applicable pre-commit hooks, two-file syntax/import checks and the PR merge-base diff check exited 0. Exact local argv/runtime/wall metadata and raw logs are retained in the acceptance manifest. The follow-up addresses the non-UTC timestamp issue reported in review: an update previously changed the epoch ingested by Arrow (New York +5 hours, Tokyo -9 hours). On the published original head, 4 non-UTC sync/async SDK regressions failed and 12 naive/UTC/NaT controls passed (exit 1; wall 3.713 seconds). Formatting with `timezone="UTC"` emits a Z suffix and preserves the ingested nanosecond instant independently of the destination timezone. This is consistent with the Arrow-aligned epoch convention in merged #3775. The original 50 cases also remain green. Full-suite wall times were 158.619 seconds (3.11) and 163.631 seconds (3.13), both exit 0. --- python/python/lancedb/util.py | 12 +- .../python/tests/test_util_update_values.py | 215 ++++++++++++++++++ 2 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 python/python/tests/test_util_update_values.py diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index f1a652924..a043f0c8d 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -426,7 +426,8 @@ def _(value: date): @value_to_sql.register(list) -def _(value: list): +@value_to_sql.register(tuple) +def _(value: Union[list, tuple]): return "[" + ", ".join(map(value_to_sql, value)) + "]" @@ -471,6 +472,15 @@ def _(value: np.floating): return value_to_sql(float(value)) +@value_to_sql.register(np.datetime64) +def _(value: np.datetime64): + if np.isnat(value): + return value_to_sql(None) + if np.datetime_data(value.dtype)[0] in ("ps", "fs", "as"): + raise ValueError("SQL timestamps cannot preserve sub-nanosecond precision") + return value_to_sql(str(np.datetime_as_string(value, unit="ns", timezone="UTC"))) + + def deprecated(func): """This is a decorator which can be used to mark functions as deprecated. It will result in a warning being emitted diff --git a/python/python/tests/test_util_update_values.py b/python/python/tests/test_util_update_values.py new file mode 100644 index 000000000..fdef867da --- /dev/null +++ b/python/python/tests/test_util_update_values.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import date, datetime + +import lancedb +import numpy as np +import pyarrow as pa +import pytest + + +ARRAY_VALUES = [ + ((1.0, 2.0), [1.0, 2.0]), + ([1.0, 2.0], [1.0, 2.0]), + (np.array([1.0, 2.0]), [1.0, 2.0]), + (None, None), +] + + +@pytest.mark.parametrize("value,expected", ARRAY_VALUES) +def test_update_array_values(tmp_path, value, expected): + dtype = pa.list_(pa.float64()) + db = lancedb.connect(tmp_path) + table = db.create_table( + "values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)}) + ) + table.update(values={"array": value}) + assert table.to_arrow()["array"].to_pylist() == [expected] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value,expected", ARRAY_VALUES) +async def test_update_array_values_async(tmp_path, value, expected): + dtype = pa.list_(pa.float64()) + db = await lancedb.connect_async(tmp_path) + table = await db.create_table( + "values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)}) + ) + await table.update({"array": value}) + assert (await table.to_arrow())["array"].to_pylist() == [expected] + + +TIME_VALUES = [ + np.datetime64("2025", "Y"), + np.datetime64("2025-01", "M"), + np.datetime64("2025-01-02", "D"), + np.datetime64("2025-01-02T03", "h"), + np.datetime64("2025-01-02T03:04", "m"), + np.datetime64("2025-01-02T03:04:05", "s"), + np.datetime64("2025-01-02T03:04:05.123", "ms"), + np.datetime64("2025-01-02T03:04:05.123456", "us"), + np.datetime64("2025-01-02T03:04:05.123456789", "ns"), + np.datetime64("NaT", "ns"), + datetime(2025, 1, 2, 3, 4, 5, 123456), + None, +] + + +def expected_time(value): + if value is None or isinstance(value, np.datetime64) and np.isnat(value): + return None + return int(np.datetime64(value, "ns").astype(np.int64)) + + +@pytest.mark.parametrize("value", TIME_VALUES) +def test_update_timestamp_values(tmp_path, value): + db = lancedb.connect(tmp_path) + data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))}) + table = db.create_table("times", data) + table.update(values={"time": value}) + assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [ + expected_time(value) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", TIME_VALUES) +async def test_update_timestamp_values_async(tmp_path, value): + db = await lancedb.connect_async(tmp_path) + data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))}) + table = await db.create_table("times", data) + await table.update({"time": value}) + assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [ + expected_time(value) + ] + + +@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []]) +def test_nested_array_tuple_matches_list(tmp_path, values): + db = lancedb.connect(tmp_path) + data = pa.table( + {"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))} + ) + control = db.create_table("list_values", data) + candidate = db.create_table("tuple_values", data) + tuples = tuple(tuple(value) for value in values) + try: + control.update(values={"array": values}) + except ValueError as error: + reason = ( + "Expected a literal value in array" if values else "concat requires input" + ) + assert reason in str(error) + with pytest.raises(ValueError, match=reason): + candidate.update(values={"array": tuples}) + assert control.to_arrow().equals(data) + assert candidate.to_arrow().equals(data) + else: + candidate.update(values={"array": tuples}) + assert candidate.to_arrow().equals(control.to_arrow()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []]) +async def test_nested_array_tuple_matches_list_async(tmp_path, values): + db = await lancedb.connect_async(tmp_path) + data = pa.table( + {"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))} + ) + control = await db.create_table("list_values", data) + candidate = await db.create_table("tuple_values", data) + tuples = tuple(tuple(value) for value in values) + try: + await control.update({"array": values}) + except ValueError as error: + reason = ( + "Expected a literal value in array" if values else "concat requires input" + ) + assert reason in str(error) + with pytest.raises(ValueError, match=reason): + await candidate.update({"array": tuples}) + assert (await control.to_arrow()).equals(data) + assert (await candidate.to_arrow()).equals(data) + else: + await candidate.update({"array": tuples}) + assert (await candidate.to_arrow()).equals(await control.to_arrow()) + + +@pytest.mark.parametrize("unit", ["ps", "fs", "as"]) +def test_subnanosecond_update_does_not_round(tmp_path, unit): + db = lancedb.connect(tmp_path) + data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))}) + table = db.create_table("precise", data) + with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"): + table.update(values={"time": np.datetime64(1, unit)}) + assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unit", ["ps", "fs", "as"]) +async def test_subnanosecond_update_does_not_round_async(tmp_path, unit): + db = await lancedb.connect_async(tmp_path) + data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))}) + table = await db.create_table("precise", data) + with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"): + await table.update({"time": np.datetime64(1, unit)}) + assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [0] + + +@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"]) +def test_update_numpy_date_column(tmp_path, unit): + value = np.datetime64("2025-01-02", unit) + db = lancedb.connect(tmp_path) + table = db.create_table( + "dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())}) + ) + table.update(values={"day": value}) + assert table.to_arrow()["day"].to_pylist() == [value.astype("datetime64[D]").item()] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"]) +async def test_update_numpy_date_column_async(tmp_path, unit): + value = np.datetime64("2025-01-02", unit) + db = await lancedb.connect_async(tmp_path) + table = await db.create_table( + "dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())}) + ) + await table.update({"day": value}) + assert (await table.to_arrow())["day"].to_pylist() == [ + value.astype("datetime64[D]").item() + ] + + +@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"]) +@pytest.mark.parametrize( + "value", + [np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")], +) +def test_update_numpy_timestamp_matches_arrow_ingestion(tmp_path, timezone, value): + dtype = pa.timestamp("ns", tz=timezone) + ingested = pa.array(np.array([value]), type=dtype) + db = lancedb.connect(tmp_path) + table = db.create_table("timestamps", pa.table({"time": ingested})) + expected = table.to_arrow()["time"].cast(pa.int64()).to_pylist() + table.update(values={"time": value}) + assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"]) +@pytest.mark.parametrize( + "value", + [np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")], +) +async def test_update_numpy_timestamp_matches_arrow_ingestion_async( + tmp_path, timezone, value +): + dtype = pa.timestamp("ns", tz=timezone) + ingested = pa.array(np.array([value]), type=dtype) + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("timestamps", pa.table({"time": ingested})) + expected = (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() + await table.update({"time": value}) + assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == expected From 57a6cce6819913adfca2852bb795a1ec0dd4641c Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:46:34 +0800 Subject: [PATCH 19/64] fix(python): omit excluded model fields from persistence (#4465) LanceDB serializes Pydantic rows with `model_dump()`, which omits `Field(exclude=True)`, but currently constructs an Arrow schema containing those same fields. Creating a local table therefore fails when an excluded field is non-nullable, including defaulted transient fields and nested models. Filter unconditionally excluded fields when converting model schemas, including nested structs. Skip excluded source/vector fields when deriving embedding configuration so persistent metadata does not refer to absent columns. Keep `field_names()` as model-field reflection. Tests exercise implicit and explicit schemas, synchronous and asynchronous create/add/reopen/readback, nested fields, defaults and default factories, retained fields, and local deterministic embeddings. Excluded values are not persisted; `to_pydantic` restores declared defaults. Conditional exclusion and readback of required excluded fields without defaults are outside this change. Implemented with assistance from gpt-6.1-sol; final diff reviewed locally. Local validation reuses the native extension matching official base d864e7198fe55a4e5fbe747fee574737fa102aab. The current-main Pydantic module is byte-identical to that base, but current-main native code has not been rebuilt. Slow/S3 tests, remote services, and other operating systems were not validated. Reproducible commands after preparing the documented development environment (relative `.venv` examples are equivalent commands; actual local runs use absolute read-only environments and fresh task-private working directories): ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests/test_pydantic_excluded_fields.py -q -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/pydantic.py python/python/tests/test_pydantic_excluded_fields.py git diff --check d864e7198fe55a4e5fbe747fee574737fa102aab HEAD ``` Final frozen-source results: - EXCLUDE_TARGET_COMPLETE_ORIGINAL_RED: Python 3.11.14, exit 1, 1.900076000019908 seconds, {"failed"=>7}. - EXCLUDE_TARGET_EMBEDDING: Python 3.11.14, exit 0, 1.3686820000002626 seconds, {"passed"=>7}. - EXCLUDE_TARGET313: Python 3.13.11, exit 0, 1.7214069999754429 seconds, {"passed"=>7}. - EXCLUDE_FULL311: Python 3.11.14, exit 0, 217.56785399999353 seconds, {"passed"=>1531, "skipped"=>36, "deselected"=>68, "warnings"=>251}. - EXCLUDE_FULL313: Python 3.13.11, exit 0, 197.93847399999504 seconds, {"passed"=>1528, "skipped"=>37, "deselected"=>68, "warnings"=>251}. - EXCLUDE_DOCTEST313: Python 3.13.11, exit 0, 11.8951390000002 seconds, {"passed"=>71, "skipped"=>6, "warnings"=>2}. - EXCLUDE_PYRIGHT: Python 3.13.11, exit 0, 5.808195000019623 seconds, {"warnings"=>0, "errors"=>0}. - EXCLUDE_PRECOMMIT: Python 3.11.14, exit 0, 0.5624079999979585 seconds, {"passed"=>4, "skipped"=>2}. - EXCLUDE_RUFF_FINAL: Ruff 0.9.9, exit 0, 0.06722099997568876 seconds, {}. - EXCLUDE_FORMAT_FINAL: Ruff 0.9.9, exit 0, 0.033322000002954155 seconds, {"files_formatted"=>157}. - EXCLUDE_SYNTAX: Python 3.11.14, exit 0, 0.031629999983124435 seconds, {"files_compiled"=>2}. - EXCLUDE_DIFF_CHECK: Git, exit 0, 0.024460999993607402 seconds, {}. --- python/python/lancedb/pydantic.py | 9 +- .../tests/test_pydantic_excluded_fields.py | 128 ++++++++++++++++++ 2 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 python/python/tests/test_pydantic_excluded_fields.py diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index bb70c9a67..db8723a99 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -268,7 +268,9 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType: def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]: return [ - _pydantic_to_field(name, field) for name, field in model.model_fields.items() + _pydantic_to_field(name, field) + for name, field in model.model_fields.items() + if field.exclude is not True ] @@ -501,6 +503,9 @@ class LanceModel(pydantic.BaseModel): vec_and_function = [] for name, field_info in cls.safe_get_fields().items(): + if field_info.exclude is True: + continue + func = get_extras(field_info, "vector_column_for") if func is not None: vec_and_function.append([name, func]) @@ -508,6 +513,8 @@ class LanceModel(pydantic.BaseModel): configs = [] for vec, func in vec_and_function: for source, field_info in cls.safe_get_fields().items(): + if field_info.exclude is True: + continue src_func = get_extras(field_info, "source_column_for") if src_func is func: # note we can't use == here since the function is a pydantic diff --git a/python/python/tests/test_pydantic_excluded_fields.py b/python/python/tests/test_pydantic_excluded_fields.py new file mode 100644 index 000000000..5f1fd1adb --- /dev/null +++ b/python/python/tests/test_pydantic_excluded_fields.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import lancedb +import pytest +from pydantic import BaseModel, Field +from lancedb.pydantic import LanceModel + + +class Details(BaseModel): + label: str + transient: str = Field(default="nested default", exclude=True) + + +class Row(LanceModel): + id: int + details: Details + transient: str = Field(default="local default", exclude=True) + generated: list[str] = Field(default_factory=lambda: ["local"], exclude=True) + retained: str = Field(default="persisted", exclude=False) + + +def rows(): + return [ + Row( + id=i, + details=Details(label=str(i), transient="private"), + transient="private", + generated=["private"], + ) + for i in (1, 2) + ] + + +def assert_schema(schema): + assert schema.names == ["id", "details", "retained"] + assert [field.name for field in schema.field("details").type] == ["label"] + assert Row.field_names() == ["id", "details", "transient", "generated", "retained"] + + +def assert_restored(restored): + assert sorted(row.id for row in restored) == [1, 2] + for row in restored: + assert row.transient == "local default" + assert row.generated == ["local"] + assert row.details.transient == "nested default" + assert row.details.label == str(row.id) + assert row.retained == "persisted" + + +@pytest.mark.parametrize("explicit_schema", [False, True]) +def test_excluded_fields_round_trip(tmp_path, explicit_schema): + data = rows() + kwargs = {"schema": Row} if explicit_schema else {} + db = lancedb.connect(tmp_path) + table = db.create_table("rows", data=data[:1], **kwargs) + table.add(data[1:]) + assert_schema(table.schema) + reopened = lancedb.connect(tmp_path).open_table("rows") + assert_restored(reopened.search().to_pydantic(Row)) + assert data[0].transient == "private" + assert data[0].generated == ["private"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("explicit_schema", [False, True]) +async def test_excluded_fields_round_trip_async(tmp_path, explicit_schema): + data = rows() + kwargs = {"schema": Row} if explicit_schema else {} + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("rows", data=data[:1], **kwargs) + await table.add(data[1:]) + assert_schema(await table.schema()) + reopened_db = await lancedb.connect_async(tmp_path) + reopened = await reopened_db.open_table("rows") + assert_restored(await reopened.query().to_pydantic(Row)) + assert data[0].details.transient == "private" + + +@pytest.mark.parametrize("excluded", ["transient", "source", "vector"]) +def test_excluded_fields_keep_embedding_metadata_consistent( + tmp_path, monkeypatch, excluded +): + from lancedb.embeddings.base import TextEmbeddingFunction + from lancedb.embeddings.registry import get_registry + from lancedb.pydantic import Vector + + calls = [] + + class LocalEmbedding(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + calls.extend(texts) + return [[1.0, 0.0] for _ in texts] + + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + registry.register("test-excluded-fields")(LocalEmbedding) + function = LocalEmbedding.create(max_retries=0) + + class EmbeddedRow(LanceModel): + text: str = function.SourceField(default="alpha", exclude=excluded == "source") + vector: Vector(2) = function.VectorField( + default_factory=lambda: [1.0, 0.0], exclude=excluded == "vector" + ) + transient: str = Field(default="local", exclude=True) + + configs = EmbeddedRow.parse_embedding_functions() + assert len(configs) == (1 if excluded == "transient" else 0) + db = lancedb.connect(tmp_path) + data = [{"text": "alpha"}] if excluded == "transient" else [EmbeddedRow()] + table = db.create_table("embeddings", schema=EmbeddedRow, data=data) + assert "transient" not in table.schema.names + assert ("text" in table.schema.names) == (excluded != "source") + assert ("vector" in table.schema.names) == (excluded != "vector") + reopened = lancedb.connect(tmp_path).open_table("embeddings") + assert len(reopened.embedding_functions) == len(configs) + restored = reopened.search().to_pydantic(EmbeddedRow) + assert restored[0].text == "alpha" + assert restored[0].vector == [1.0, 0.0] + assert restored[0].transient == "local" + if excluded == "transient": + assert calls == ["alpha"] + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + else: + assert calls == [] From 036f4ec53fac660b7703fe707dbcd7c5b405e020 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:13:31 +0800 Subject: [PATCH 20/64] fix(node): preserve options when namespace is omitted (#4479) The namespace overloads allow `undefined` to select the root namespace, but overload dispatch treats it as the legacy options argument and discards the trailing options. For example, `createTable(name, data, undefined, { mode: "overwrite" })` raises an already-exists error instead of overwriting. Preserve the trailing options for `createTable`, `createEmptyTable`, `tableNames`, and `listTables` when the namespace is omitted. Legacy options-object calls and explicit namespace arrays keep their existing behavior. Adds native-database regressions for overwrite, existOk, and both listing APIs, parameterized over `undefined` and `[]`. Validation on Windows / Node 24.19.0: - Before the fix: all five omitted-namespace cases failed; all five explicit-root controls passed. - Current TypeScript source with official, integrity-verified 0.41.0-beta.0 native binaries: the entire connection test file passed, 47/47 (transpile-only). - The same fix replayed onto the matching 0.41.0-beta.0 source: TypeScript compilation, strict ts-jest connection tests (47/47), and TypeDoc passed. - Biome check/format on both changed files and diff whitespace checks passed. No Rust changes. A complete current-main native rebuild remains for CI: released declarations predate the unrelated lazy-iterator APIs, so current-main typechecking against those declarations reports six pre-existing iterator signature errors. Signed-off-by: Xucheng Zhou Co-authored-by: Xuanwo --- nodejs/__test__/connection.test.ts | 60 ++++++++++++++++++++++++++++++ nodejs/lancedb/connection.ts | 8 ++-- 2 files changed, 64 insertions(+), 4 deletions(-) diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 1167ae34a..355b86ccb 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -117,6 +117,66 @@ describe("given a connection", () => { await expect(tbl.countRows()).resolves.toBe(1); }); + describe.each([undefined, []])("root namespace %p", (namespacePath) => { + it("should preserve createTable options after the namespace", async () => { + await db.createTable("test", [{ id: 1 }, { id: 2 }]); + + const table = await db.createTable("test", [{ id: 3 }], namespacePath, { + mode: "overwrite", + }); + + const rows = await table.query().toArray(); + expect(rows.map((row) => row.id)).toEqual([3]); + }); + + it("should preserve createEmptyTable overwrite after the namespace", async () => { + await db.createTable("test", [{ id: 1 }]); + const schema = new Schema([new Field("id", new Float64(), true)]); + + const table = await db.createEmptyTable("test", schema, namespacePath, { + mode: "overwrite", + }); + + await expect(table.countRows()).resolves.toBe(0); + }); + + it("should preserve createEmptyTable existOk after the namespace", async () => { + await db.createTable("test", [{ id: 1 }]); + const schema = new Schema([new Field("id", new Float64(), true)]); + + const table = await db.createEmptyTable("test", schema, namespacePath, { + existOk: true, + }); + + await expect(table.countRows()).resolves.toBe(1); + }); + + it("should preserve tableNames options after the namespace", async () => { + for (const name of ["a", "b", "c"]) { + await db.createTable(name, [{ id: 1 }]); + } + + await expect( + db.tableNames(namespacePath, { limit: 1, startAfter: "a" }), + ).resolves.toEqual(["b"]); + }); + + it("should preserve listTables options after the namespace", async () => { + for (const name of ["a", "b", "c"]) { + await db.createTable(name, [{ id: 1 }]); + } + + const first = await db.listTables(namespacePath, { limit: 1 }); + expect(first.tables).toEqual(["a"]); + expect(first.pageToken).toBeDefined(); + + const second = await db.listTables(namespacePath, { + limit: 1, + pageToken: first.pageToken, + }); + expect(second.tables).toEqual(["b"]); + }); + }); it("should keep the existing table when mode is exist_ok", async () => { const created = await db.createTable("test", [{ id: 1 }, { id: 2 }], { mode: "exist_ok", diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 4ed511b86..f866d1ecd 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -720,7 +720,7 @@ export class LocalConnection extends Connection { } else { // First argument is options object (backwards compatibility) namespacePath = undefined; - tableNamesOptions = namespacePathOrOptions; + tableNamesOptions = namespacePathOrOptions ?? options; } return this.inner.tableNames( @@ -831,7 +831,7 @@ export class LocalConnection extends Connection { : undefined; const listTablesOptions = Array.isArray(namespacePathOrOptions) ? options - : namespacePathOrOptions; + : (namespacePathOrOptions ?? options); return this.inner.listTables( namespacePath ?? [], @@ -942,7 +942,7 @@ export class LocalConnection extends Connection { } else { // Third argument is options object (backwards compatibility) namespacePath = undefined; - createOptions = namespacePathOrOptions; + createOptions = namespacePathOrOptions ?? options; } return this._createTableImpl(name, data, namespacePath, createOptions); @@ -989,7 +989,7 @@ export class LocalConnection extends Connection { } else { // Third argument is options object (backwards compatibility) namespacePath = undefined; - createOptions = namespacePathOrOptions; + createOptions = namespacePathOrOptions ?? options; } let mode: string = createOptions?.mode ?? "create"; From 13d54db59411029f6c80755b7453090e17beb417 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:13:46 +0800 Subject: [PATCH 21/64] fix(node): preserve caller-owned storage options (#4480) `getStorageOptions` writes typed creation options back into the caller's `storageOptions` object. This makes frozen options throw and causes overrides to leak into later calls that reuse the storage map. For example, using `{ storageOptions: shared, enableV2ManifestPaths: true }` changes `shared.newTableEnableV2ManifestPaths` to `"true"`. A subsequent creation with just `{ storageOptions: shared }` then uses V2 manifest paths even if `shared` originally requested `"false"`. Merge the typed overrides into a local map instead. Their existing precedence, key normalization, and undefined/default behavior remain unchanged. This affects both `createTable` and `createEmptyTable`; no native/Rust changes. Adds six native-database regressions covering frozen creation options, frozen nested storage options with conflicting defaults, and cross-call reuse for both APIs. Validation (Linux x64 / Node 24.19.0 / official 0.41.0-beta.0 native library): - Before: all six new regressions fail; an existing namespace storage-options control passes. - After: the full connection test file passes, 43/43 (current TypeScript source, transpile-only). - Same fix on the matching 0.41.0-beta.0 source: TypeScript compilation, strict ts-jest connection tests (43/43), and TypeDoc with warnings-as-errors pass. - Biome on changed files and whitespace checks pass. A current-main native rebuild was not run. Released native declarations predate six unrelated lazy-iterator signatures; current-main typechecking reports the same six errors before and after this fix (only downstream line coordinates shift). ### Connection entry-point follow-up The same ownership problem occurs in `connect(uri, options)` and `connectNamespace`: normalization assigns a new storage map into the caller's options. Copy the options before normalization instead. The `connect(options)` overload already copies its input and is retained as a control. Nine added native-database cases cover frozen outer options (with/without storage options) and repeated reuse, including table creation and row-count checks. Before this entry-point fix, the URI/namespace cases fail (6), while the object-overload controls pass (3). Afterward the complete connection file passes 52/52 on current TypeScript with the official 0.41.0-beta.0 native library. Applying the entry-point fix to matching-release source also passes strict TypeScript, strict ts-jest (52/52), and TypeDoc. Biome passes on both newly changed files. The same six current-main/released-native lazy-iterator declaration errors remain unchanged. No Rust/native rebuild or GPU claim is added. --------- Signed-off-by: Xucheng Zhou --- nodejs/__test__/connection.test.ts | 144 +++++++++++++++++++++++++++++ nodejs/lancedb/connection.ts | 23 ++--- nodejs/lancedb/index.ts | 17 ++-- 3 files changed, 164 insertions(+), 20 deletions(-) diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 355b86ccb..e32e7e626 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -36,6 +36,87 @@ describe("when connecting", () => { }); }); +describe.each([ + [ + "URI overload", + (uri: string, options: { storageOptions?: Record }) => + connect(uri, options), + ], + [ + "object overload", + async ( + uri: string, + options: { storageOptions?: Record }, + ) => { + const input = { uri, ...options }; + if (Object.isFrozen(options)) { + Object.freeze(input); + } + const db = await connect(input); + expect(input).toEqual({ uri, ...options }); + expect(input.storageOptions).toBe(options.storageOptions); + return db; + }, + ], + [ + "namespace", + (uri: string, options: { storageOptions?: Record }) => + connectNamespace("dir", { root: uri }, options), + ], +])("%s connection options", (_name, open) => { + let tmpDir: tmp.DirResult; + beforeEach(() => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + }); + afterEach(() => tmpDir.removeCallback()); + + it("should accept frozen options without storage options", async () => { + const options = Object.freeze({}); + const db = await open(tmpDir.name, options); + try { + await db.createTable("frozen", [{ id: 1 }]); + await expect(db.tableNames()).resolves.toContain("frozen"); + expect(options).toEqual({}); + } finally { + await db.close(); + } + }); + + it("should accept frozen options with storage options", async () => { + const storageOptions = Object.freeze({ + newTableDataStorageVersion: "stable", + }); + const options = Object.freeze({ storageOptions }); + const db = await open(tmpDir.name, options); + try { + const table = await db.createTable("frozen", [{ id: 1 }]); + await expect(table.countRows()).resolves.toBe(1); + expect(options.storageOptions).toBe(storageOptions); + expect(storageOptions).toEqual({ newTableDataStorageVersion: "stable" }); + } finally { + await db.close(); + } + }); + + it("should preserve reused options and their storage options", async () => { + const storageOptions = { newTableDataStorageVersion: "stable" }; + const options = { storageOptions }; + for (let i = 0; i < 2; i++) { + const db = await open(tmpDir.name, options); + try { + const table = await db.createTable(`reused_${i}`, [{ id: i }]); + await expect(table.countRows()).resolves.toBe(1); + expect(options.storageOptions).toBe(storageOptions); + expect(storageOptions).toEqual({ + newTableDataStorageVersion: "stable", + }); + } finally { + await db.close(); + } + } + }); +}); + describe("given a connection", () => { let tmpDir: tmp.DirResult; let db: Connection; @@ -367,6 +448,69 @@ describe("given a connection", () => { }); }); + describe.each(["createTable", "createEmptyTable"] as const)( + "%s storage options", + (method) => { + const schema = new Schema([new Field("id", new Float64(), true)]); + + it("should accept frozen creation options without storage options", async () => { + const options = Object.freeze({ + dataStorageVersion: "stable", + enableV2ManifestPaths: true, + }); + const table = + method === "createTable" + ? await db.createTable("frozen_options", [{ id: 1 }], options) + : await db.createEmptyTable("frozen_options", schema, options); + expect(await (table as LocalTable).usesV2ManifestPaths()).toBe(true); + expect(options).not.toHaveProperty("storageOptions"); + }); + + it("should accept frozen storage options and preserve override precedence", async () => { + const storageOptions = Object.freeze({ + newTableDataStorageVersion: "legacy", + newTableEnableV2ManifestPaths: "false", + }); + const options = Object.freeze({ + storageOptions, + dataStorageVersion: "stable", + enableV2ManifestPaths: true, + }); + const table = + method === "createTable" + ? await db.createTable("frozen_storage", [{ id: 1 }], options) + : await db.createEmptyTable("frozen_storage", schema, options); + expect(await (table as LocalTable).usesV2ManifestPaths()).toBe(true); + expect(options.storageOptions).toBe(storageOptions); + expect(storageOptions).toEqual({ + newTableDataStorageVersion: "legacy", + newTableEnableV2ManifestPaths: "false", + }); + }); + + it("should not retain a previous override when reusing storage options", async () => { + const storageOptions = { newTableEnableV2ManifestPaths: "false" }; + const create = async (name: string, enableV2ManifestPaths?: boolean) => + method === "createTable" + ? db.createTable(name, [{ id: 1 }], { + storageOptions, + enableV2ManifestPaths, + }) + : db.createEmptyTable(name, schema, { + storageOptions, + enableV2ManifestPaths, + }); + const first = await create("override", true); + expect(await (first as LocalTable).usesV2ManifestPaths()).toBe(true); + const second = await create("without_override"); + expect(await (second as LocalTable).usesV2ManifestPaths()).toBe(false); + expect(storageOptions).toEqual({ + newTableEnableV2ManifestPaths: "false", + }); + }); + }, + ); + it("should be able to migrate tables to the V2 manifest paths", async () => { const db = await connect(tmpDir.name); const table = (await db.createEmptyTable( diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index f866d1ecd..44fde96f1 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -893,23 +893,24 @@ export class LocalConnection extends Connection { private getStorageOptions( options?: Partial, ): Record | undefined { + let storageOptions = options?.storageOptions; if (options?.dataStorageVersion !== undefined) { - if (options.storageOptions === undefined) { - options.storageOptions = {}; - } - options.storageOptions["newTableDataStorageVersion"] = - options.dataStorageVersion; + storageOptions = { + ...storageOptions, + newTableDataStorageVersion: options.dataStorageVersion, + }; } if (options?.enableV2ManifestPaths !== undefined) { - if (options.storageOptions === undefined) { - options.storageOptions = {}; - } - options.storageOptions["newTableEnableV2ManifestPaths"] = - options.enableV2ManifestPaths ? "true" : "false"; + storageOptions = { + ...storageOptions, + newTableEnableV2ManifestPaths: options.enableV2ManifestPaths + ? "true" + : "false", + }; } - return cleanseStorageOptions(options?.storageOptions); + return cleanseStorageOptions(storageOptions); } async createTable( diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 1057e2d53..4ad8b2a93 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -420,10 +420,10 @@ export async function connect( throw new Error("uri is required"); } - finalOptions = (finalOptions as ConnectionOptions) ?? {}; - (finalOptions).storageOptions = cleanseStorageOptions( - (finalOptions).storageOptions, - ); + finalOptions = { + ...finalOptions, + storageOptions: cleanseStorageOptions(finalOptions.storageOptions), + }; // Create native header provider if one was provided let nativeProvider: NativeJsHeaderProvider | undefined; @@ -630,11 +630,10 @@ export async function connectNamespace( properties = configOrProperties as Record; } - const finalOptions: ConnectNamespaceOptions = (options ?? - {}) as ConnectNamespaceOptions; - finalOptions.storageOptions = cleanseStorageOptions( - finalOptions.storageOptions, - ); + const finalOptions: ConnectNamespaceOptions = { + ...options, + storageOptions: cleanseStorageOptions(options?.storageOptions), + }; const nativeConn = await LanceDbConnection.newWithNamespace( implName, From e256b3c2cfc49f45c74d4ac5e31a2e60b6a40efa Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:03 +0800 Subject: [PATCH 22/64] fix(python): honor namespace table storage options (#4483) Follow-up on current main `7b3ba4c4` (merge commit `0b0475c1`): resolved the test import conflict, preserving both upstream imports and the original `warnings` import. The final PR diff is still limited to the same storage-option helper/call sites and regression class. AST comparisons preserve every upstream test definition and the original regression class/helper; both changed files compile and pass Ruff 0.15.1 check/format. Fresh locked-environment runtime validation was attempted with `uv run --locked --extra tests --extra dev pytest python/tests/test_namespace.py -q`, but dependency setup failed downloading the locked `pylance==14.0.0b12` wheel due a network connection error. No current-main native test result is claimed. Repository-wide Ruff reports 49 findings and 11 formatting differences in unchanged upstream docs/web-tests files; these were left outside this PR. CI approval and a human review are still pending. --- Original implementation and historical validation (before the main merge): Namespace-backed `create_table` accepts `data_storage_version` and `enable_v2_manifest_paths`, but neither sync nor async creation forwards them. The requested format/manifest settings are silently ignored. Normalize these deprecated arguments into the existing `new_table_*` storage options, matching ordinary connections. Explicit arguments take precedence; the caller's dictionary is not mutated. No Rust or native interface changes. Regression tests inspect persisted tables, including sync/async creation, populated/empty tables, stable/legacy formats, both manifest settings and explicit-option precedence. Validation on Linux/Python 3.12, based on current main `d3390a9`: - Rebuilt the exact-main native extension and reused that identical extension in separate baseline/candidate environments; only the candidate Python namespace wrapper differs. Dependencies came from the unchanged `uv.lock`. - Same complete `test_namespace.py`: baseline 16 regression failures and 39 passing controls; candidate **55 passed**, no skips/errors. - Ruff check/format and diff whitespace checks passed. This validates the local directory-namespace path, not a live Cloud service or the complete repository CI suite. Signed-off-by: Xucheng Zhou --- python/python/lancedb/namespace.py | 32 +++++++ python/python/tests/test_namespace.py | 120 ++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index f8675ebb5..a5d88eb1f 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -11,6 +11,7 @@ through a namespace abstraction. from __future__ import annotations import sys +import warnings from typing import Any, Dict, Iterable, List, Optional, Union from uuid import UUID @@ -82,6 +83,31 @@ from ._lancedb import Session _MAX_QUERY_K = 2**31 - 1 +def _create_table_storage_options( + storage_options: Optional[Dict[str, str]], + data_storage_version: Optional[str], + enable_v2_manifest_paths: Optional[bool], +) -> Optional[Dict[str, str]]: + if data_storage_version is None and enable_v2_manifest_paths is None: + return storage_options + options = dict(storage_options or {}) + for name, value in ( + ("data_storage_version", data_storage_version), + ("enable_v2_manifest_paths", enable_v2_manifest_paths), + ): + if value is not None: + warnings.warn( + f"setting {name} directly on create_table is deprecated. " + "Use storage_options instead.", + DeprecationWarning, + stacklevel=3, + ) + options[f"new_table_{name}"] = ( + str(value).lower() if isinstance(value, bool) else value + ) + return options + + def _query_to_namespace_request( table_id: List[str], query: "Query", @@ -587,6 +613,9 @@ class LanceNamespaceDBConnection(DBConnection): if mode.lower() not in ["create", "overwrite"]: raise ValueError("mode must be either 'create' or 'overwrite'") validate_table_name(name) + storage_options = _create_table_storage_options( + storage_options, data_storage_version, enable_v2_manifest_paths + ) async_table = LOOP.run( self._inner.create_table( name, @@ -1230,6 +1259,9 @@ class AsyncLanceNamespaceDBConnection: if mode.lower() not in ["create", "overwrite"]: raise ValueError("mode must be either 'create' or 'overwrite'") validate_table_name(name) + storage_options = _create_table_storage_options( + storage_options, data_storage_version, enable_v2_manifest_paths + ) table = await self._inner.create_table( name, data, diff --git a/python/python/tests/test_namespace.py b/python/python/tests/test_namespace.py index 520a6992d..b55da4d57 100644 --- a/python/python/tests/test_namespace.py +++ b/python/python/tests/test_namespace.py @@ -6,6 +6,7 @@ import tempfile import shutil import importlib +import warnings import http.server import json import threading @@ -25,6 +26,125 @@ PUSHDOWN_DATA = pa.table( ) +class TestNamespaceCreateTableOptions: + @pytest.mark.parametrize("version", ["legacy", "stable"]) + @pytest.mark.parametrize("empty", [False, True]) + def test_deprecated_storage_version(self, tmp_path, version, empty): + db = lancedb.connect_namespace("dir", {"root": str(tmp_path / "namespace")}) + data = None if empty else [{"id": 1}] + schema = pa.schema([("id", pa.int64())]) + other_version = "stable" if version == "legacy" else "legacy" + options = {"new_table_data_storage_version": other_version} + control = lancedb.connect(tmp_path / "control").create_table( + "test", + data=data, + schema=schema, + storage_options={"new_table_data_storage_version": version}, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = db.create_table( + "test", + data=data, + schema=schema, + storage_options=options, + data_storage_version=version, + ) + + assert ( + table.to_lance().data_storage_version + == control.to_lance().data_storage_version + ) + assert table.to_arrow().to_pylist() == ([] if empty else data) + assert options == {"new_table_data_storage_version": other_version} + assert any("data_storage_version" in str(w.message) for w in caught) + + @pytest.mark.asyncio + @pytest.mark.parametrize("version", ["legacy", "stable"]) + @pytest.mark.parametrize("empty", [False, True]) + async def test_deprecated_storage_version_async(self, tmp_path, version, empty): + root = str(tmp_path / "namespace") + db = lancedb.connect_namespace_async("dir", {"root": root}) + data = None if empty else [{"id": 1}] + schema = pa.schema([("id", pa.int64())]) + other_version = "stable" if version == "legacy" else "legacy" + options = {"new_table_data_storage_version": other_version} + control = lancedb.connect(tmp_path / "control").create_table( + "test", + data=data, + schema=schema, + storage_options={"new_table_data_storage_version": version}, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = await db.create_table( + "test", + data=data, + schema=schema, + storage_options=options, + data_storage_version=version, + ) + + reopened = lancedb.connect_namespace("dir", {"root": root}).open_table("test") + assert ( + reopened.to_lance().data_storage_version + == control.to_lance().data_storage_version + ) + assert (await table.to_arrow()).to_pylist() == ([] if empty else data) + assert options == {"new_table_data_storage_version": other_version} + assert any("data_storage_version" in str(w.message) for w in caught) + + @pytest.mark.parametrize("enabled", [False, True]) + @pytest.mark.parametrize("empty", [False, True]) + def test_deprecated_manifest_paths(self, tmp_path, enabled, empty): + db = lancedb.connect_namespace("dir", {"root": str(tmp_path)}) + data = None if empty else [{"id": 1}] + options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = db.create_table( + "test", + data=data, + schema=pa.schema([("id", pa.int64())]), + storage_options=options, + enable_v2_manifest_paths=enabled, + ) + + assert table.uses_v2_manifest_paths() == enabled + assert table.to_arrow().to_pylist() == ([] if empty else data) + assert db.open_table("test").uses_v2_manifest_paths() == enabled + assert options == { + "new_table_enable_v2_manifest_paths": str(not enabled).lower() + } + assert any("enable_v2_manifest_paths" in str(w.message) for w in caught) + + @pytest.mark.asyncio + @pytest.mark.parametrize("enabled", [False, True]) + @pytest.mark.parametrize("empty", [False, True]) + async def test_deprecated_manifest_paths_async(self, tmp_path, enabled, empty): + db = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)}) + data = None if empty else [{"id": 1}] + options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = await db.create_table( + "test", + data=data, + schema=pa.schema([("id", pa.int64())]), + storage_options=options, + enable_v2_manifest_paths=enabled, + ) + + assert await table.uses_v2_manifest_paths() == enabled + assert (await table.to_arrow()).to_pylist() == ([] if empty else data) + reopened = await db.open_table("test") + assert await reopened.uses_v2_manifest_paths() == enabled + assert options == { + "new_table_enable_v2_manifest_paths": str(not enabled).lower() + } + assert any("enable_v2_manifest_paths" in str(w.message) for w in caught) + + def _ipc_file(table: pa.Table = PUSHDOWN_DATA) -> bytes: sink = pa.BufferOutputStream() with pa.ipc.new_file(sink, table.schema) as writer: From b205e7c628c327a10f778a8d6db217c63da96e97 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:17 +0800 Subject: [PATCH 23/64] fix(node): prevent table listing limits from wrapping (#4492) ## Summary Prevent large JavaScript table-listing limits from wrapping during the native `u32` conversion. With two tables, `limit: 2 ** 32` currently returns no tables and `2 ** 32 + 1` returns only one. Validate nonnegative safe integers using the existing helper and clamp valid limits above `u32::MAX` before calling the native binding. Apply the same normalization to `tableNames` and `listTables`, including namespace overloads. Undefined, zero, small limits and `u32::MAX` keep their existing behavior. ## Validation Built the real Node native extension from main `0aa842b17f13d1f8ffd2851bc57c222bc5a376ba` with the repository's `pnpm build`, using locked dependencies and reusable build caches. No mocked/released native extension or disabled TypeScript diagnostics. With identical regression tests and that native library in fresh Jest processes: - Baseline: 40 passed, 8 failed; candidate: **48 passed**, no failures. - Both arms: `pnpm run tsc` and `pnpm run docs` passed. - Changed-file Biome check passed. Tests cover wrapping limits, invalid numbers, zero/small/u32 limits, defaults and both overload forms. Validation is the focused connection suite, not the complete repository suite. No performance claim. AI-assisted implementation and verification; the patch and exact tested source bytes were reviewed before publication. Signed-off-by: aden Co-authored-by: aden --- nodejs/__test__/connection.test.ts | 50 ++++++++++++++++++++++++++++++ nodejs/lancedb/connection.ts | 11 +++++-- 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index e32e7e626..bbd0c3c52 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -336,6 +336,56 @@ describe("given a connection", () => { expect(seen).toEqual(created); }); + it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])( + "should not wrap a large table listing limit %s", + async (limit) => { + await db.createTable("a", [{ id: 1 }]); + await db.createTable("b", [{ id: 2 }]); + for (const names of [ + db.tableNames({ limit }), + db.tableNames([], { limit }), + ]) { + await expect(names).resolves.toEqual(["a", "b"]); + } + for (const page of [ + db.listTables({ limit }), + db.listTables([], { limit }), + ]) { + await expect(page).resolves.toEqual({ tables: ["a", "b"] }); + } + }, + ); + + it.each([ + -1, + 1.5, + Number.NaN, + Number.POSITIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.MAX_SAFE_INTEGER + 1, + ])("should reject an invalid table listing limit %s", async (limit) => { + await expect(db.tableNames({ limit })).rejects.toThrow(/limit/); + await expect(db.tableNames([], { limit })).rejects.toThrow(/limit/); + await expect(db.listTables({ limit })).rejects.toThrow(/limit/); + await expect(db.listTables([], { limit })).rejects.toThrow(/limit/); + }); + + it("should preserve zero, small and unsigned-32-bit table listing limits", async () => { + await db.createTable("a", [{ id: 1 }]); + await db.createTable("b", [{ id: 2 }]); + for (const [limit, names] of [ + [0, []], + [1, ["a"]], + [2 ** 32 - 1, ["a", "b"]], + ] as const) { + await expect(db.tableNames({ limit })).resolves.toEqual(names); + const page = await db.listTables({ limit }); + expect(page.tables).toEqual(names); + } + await expect(db.tableNames()).resolves.toEqual(["a", "b"]); + await expect(db.listTables()).resolves.toEqual({ tables: ["a", "b"] }); + }); + it("should list all names and use a default page size of 100", async () => { const names = Array.from( { length: 130 }, diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 44fde96f1..c93414fab 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -726,7 +726,7 @@ export class LocalConnection extends Connection { return this.inner.tableNames( namespacePath ?? [], tableNamesOptions?.startAfter, - tableNamesOptions?.limit, + normalizeTableListingLimit(tableNamesOptions?.limit), ); } @@ -836,7 +836,7 @@ export class LocalConnection extends Connection { return this.inner.listTables( namespacePath ?? [], listTablesOptions?.pageToken, - listTablesOptions?.limit, + normalizeTableListingLimit(listTablesOptions?.limit), ); } @@ -1107,6 +1107,13 @@ export class LocalConnection extends Connection { } } +function normalizeTableListingLimit(limit?: number): number | undefined { + validateNonNegativeInteger(limit, "limit"); + // The native binding takes u32: larger JS numbers otherwise wrap to zero or + // a smaller page size. A listing limit is an upper bound, so saturate it. + return limit === undefined ? undefined : Math.min(limit, 0xffffffff); +} + /** * Takes storage options and makes all the keys snake case. */ From 3d883af674924a7e7266da19d9ce1d774298f74c Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:31 +0800 Subject: [PATCH 24/64] fix(python): propagate invalid index wait_timeout errors (#4495) Invalid `wait_timeout` values currently reach `Duration::extract(...).unwrap()` in both native index-creation bindings. A normal Python argument error becomes `pyo3_runtime.PanicException`, which is not caught by `except Exception`. Propagate the extraction error with `map(...).transpose()?` before constructing the index operation. Valid `None` and `timedelta` behavior is unchanged. Tests cover both `create_index` and `create_index_async`, three invalid types, valid `None`/`timedelta`, and no index created after rejection. Validation against main `0aa842b`, Python 3.12, uv-locked dependencies, freshly compiled PyO3 extension for each arm: - Unmodified source: 6 new regressions fail with `PanicException`; 6 controls pass. - Fixed source: 12/12 pass; broader index-creation regression selection: 18/18 pass. - `cargo fmt --all --check`, Ruff check/format, and diff check pass. The broader selection includes the 12 cases above. No cloud, full-suite, or performance claim. Signed-off-by: Xucheng Zhou --- python/python/tests/test_table.py | 19 +++++++++++++++++++ python/src/table.rs | 8 ++++++-- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index b2d499933..238ee5330 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1633,6 +1633,25 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection): job.cancel() +@pytest.mark.parametrize("method", ["create_index", "create_index_async"]) +@pytest.mark.parametrize("wait_timeout", ["invalid", 1, []]) +def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout): + table = mem_db.create_table("invalid_timeout", [{"id": 1}]) + with pytest.raises(TypeError): + getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout) + assert table.list_indices() == [] + + +@pytest.mark.parametrize("method", ["create_index", "create_index_async"]) +@pytest.mark.parametrize("wait_timeout", [None, timedelta(seconds=1)]) +def test_create_index_valid_wait_timeout(mem_db: DBConnection, method, wait_timeout): + table = mem_db.create_table("valid_timeout", [{"id": 1}]) + result = getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout) + if result is not None: + result.wait() + assert len(table.list_indices()) == 1 + + @patch("lancedb.table.AsyncTable.create_index") def test_create_index_method(mock_create_index, mem_db: DBConnection): table = mem_db.create_table( diff --git a/python/src/table.rs b/python/src/table.rs index 27ed8be8b..e77c58578 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1134,7 +1134,9 @@ impl Table { train: Option, ) -> PyResult> { let index = extract_index_params(&index)?; - let timeout = wait_timeout.map(|t| t.extract::().unwrap()); + let timeout = wait_timeout + .map(|t| t.extract::()) + .transpose()?; let mut op = self_ .inner_ref()? .create_index_with_timeout(&[column], index, timeout); @@ -1165,7 +1167,9 @@ impl Table { train: Option, ) -> PyResult> { let index = extract_index_params(&index)?; - let timeout = wait_timeout.map(|t| t.extract::().unwrap()); + let timeout = wait_timeout + .map(|t| t.extract::()) + .transpose()?; let mut op = self_ .inner_ref()? .create_index_with_timeout(&[column], index, timeout); From 93f054bcc5b02927720cdcd52fd7b1cc13c9a731 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:45 +0800 Subject: [PATCH 25/64] fix(rust): finish query streams after a zero-duration timeout (#4506) A positive query timeout emits one error and then ends, but `Duration::ZERO` leaves `TimeoutStream` in `NotStarted`, returning the same error on every subsequent poll. Transition the zero-duration path to `Completed` before returning its error. The initial error is unchanged. Extend both timeout tests to check that subsequent polls return `None`; the normal-completion test is unchanged. This affects callers that continue polling after an error, not `try_collect`, which already stops at the first error. Validation: - `cargo fmt --all` and `cargo fmt --all --check`: pass. - Mechanically extracted the unchanged implementation and three upstream timeout tests into a native Rust test binary using real cached Arrow 58.4.0, DataFusion 54.1.0, futures 0.3.32 and Tokio 1.53.1 libraries (matching the current lock). With the new assertions, baseline: 2 pass / 1 fail; fix: 3 pass. - Full-crate `cargo test --locked --offline --features remote -p lancedb --lib test_timeout_stream` could not start because this environment lacks the newly locked Lance `v14.0.0-beta.10` checkout. The component comparison is not a claim that the full Cargo suite passed; CI still needs to validate integration. Signed-off-by: Xucheng Zhou --- rust/lancedb/src/utils/mod.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index 20f209f2e..ebf3f697b 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -686,7 +686,9 @@ impl Stream for TimeoutStream { match &mut self.state { TimeoutState::NotStarted { timeout } => { if timeout.is_zero() { - return std::task::Poll::Ready(Some(Err(Self::timeout_error(timeout)))); + let err = Self::timeout_error(timeout); + self.state = TimeoutState::Completed; + return std::task::Poll::Ready(Some(Err(err))); } let deadline = Box::pin(tokio::time::sleep(*timeout)); self.state = TimeoutState::Started { @@ -1159,6 +1161,8 @@ mod tests { .to_string() .contains("Query timeout") ); + assert!(timeout_stream.next().await.is_none()); + assert!(timeout_stream.next().await.is_none()); } #[tokio::test] @@ -1178,6 +1182,8 @@ mod tests { let result = timeout_stream.next().await.unwrap(); assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("Query timeout")); + assert!(timeout_stream.next().await.is_none()); + assert!(timeout_stream.next().await.is_none()); } #[tokio::test] From d826b6115a0ca30998d9aa1bfb83d01e8f65bf4b Mon Sep 17 00:00:00 2001 From: Aviral Saxena Date: Fri, 23 Oct 2026 00:45:01 +0530 Subject: [PATCH 26/64] fix(rust): explain client-side model files when remote tokenize fails (#4519) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses option 2 of #4440. tokenize() builds the FTS analyzer on the client. For jieba/* and lindera/* that needs model files under LANCE_LANGUAGE_MODEL_HOME, and when they’re missing you only got lance’s raw Invalid directory path: ... error, even though server-side search works. Now, if a jieba/* or lindera/* tokenizer can’t be built, the error keeps the original cause and adds that tokenize() runs on the client, which directory the files need to be in, and that for remote tables they must match the server’s. I also updated the tokenize docs (Rust, Python, TS, generated JS docs) to say a mismatch can’t be detected, so tokens may differ from the server’s index. Behavior when the files exist is unchanged. Detecting a mismatch needs a server-side tokenize endpoint (option 1), so I left it out. The hint isn’t limited to remote tables, since Table::as_native is slated for removal and the wording fits local tables too. Tests: Rust test with mocked jieba/default index_details and an empty model dir, one checking other tokenizer errors don’t get the hint, and Python sync/async tests. --- docs/src/js/classes/Table.md | 6 +- nodejs/lancedb/table.ts | 6 +- python/python/lancedb/remote/table.py | 8 +- python/python/lancedb/table.py | 18 +++- python/python/tests/test_remote_db.py | 73 +++++++++++++++ rust/lancedb/src/remote/table.rs | 125 ++++++++++++++++++++++++++ rust/lancedb/src/table.rs | 34 +++++-- 7 files changed, 257 insertions(+), 13 deletions(-) diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index b41ade302..7cf87c405 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -1298,7 +1298,11 @@ Specify exactly one of `column` or `indexName`. Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in the client process from index metadata. For remote tables, this means the -same tokenizer model files must also exist locally. +same tokenizer model files must also exist locally (under +`LANCE_LANGUAGE_MODEL_HOME`) and must match the files the server uses. A +mismatch cannot be detected: the returned tokens may differ from what the +server's index produces. If the model files are missing, an error naming +the tokenizer is raised. #### Parameters diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 0c79e6a7e..4ae1aefb8 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -989,7 +989,11 @@ export abstract class Table { * * Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in * the client process from index metadata. For remote tables, this means the - * same tokenizer model files must also exist locally. + * same tokenizer model files must also exist locally (under + * `LANCE_LANGUAGE_MODEL_HOME`) and must match the files the server uses. A + * mismatch cannot be detected: the returned tokens may differ from what the + * server's index produces. If the model files are missing, an error naming + * the tokenizer is raised. */ abstract tokenize( query: string, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 96b988911..86e8fb8a0 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -279,8 +279,12 @@ class RemoteTable(Table): """Tokenize a query using the tokenizer configured on an FTS index. Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are - rebuilt in the client process from index metadata, so the same tokenizer - model files must exist locally. + rebuilt in the client process from index metadata. For remote tables, + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ return LOOP.run( self._table.tokenize(query, column=column, index_name=index_name) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index f9cfac6e1..124fa34b5 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2558,7 +2558,11 @@ class Table(ABC): Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are rebuilt in the client process from index metadata. For remote tables, - this means the same tokenizer model files must also exist locally. + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ @abstractmethod @@ -4813,7 +4817,11 @@ class LanceTable(Table): Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are rebuilt in the client process from index metadata. For remote tables, - this means the same tokenizer model files must also exist locally. + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ return LOOP.run( self._table.tokenize(query, column=column, index_name=index_name) @@ -7505,7 +7513,11 @@ class AsyncTable: Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are rebuilt in the client process from index metadata. For remote tables, - this means the same tokenizer model files must also exist locally. + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ return await self._inner.tokenize(query, column=column, index_name=index_name) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 80bd10289..a45e1aaef 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1793,6 +1793,79 @@ def test_table_wait_for_index_timeout(): table.wait_for_index(["id_idx"], timedelta(seconds=1)) +def _jieba_fts_handler(request): + """Serve a table with a ``jieba/default`` FTS index and no other state.""" + if request.path == "/v1/table/test/describe/": + body = { + "version": 1, + "schema": { + "fields": [ + {"name": "text", "type": {"type": "string"}, "nullable": True} + ] + }, + } + elif request.path == "/v1/table/test/index/list/": + index_details = { + "base_tokenizer": "jieba/default", + "language": "English", + "with_position": False, + "max_token_length": 40, + "lower_case": True, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + } + body = { + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": json.dumps(index_details), + } + ] + } + else: + request.send_response(404) + request.end_headers() + return + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(body).encode()) + + +_MODEL_FILES_ERROR = ( + r"(?s)Invalid directory path.*" + r"full text search index 'text_idx'.*" + r"tokenize\(\) runs on the client.*'jieba/default'.*" + r"LANCE_LANGUAGE_MODEL_HOME.*match the files the server uses" +) + + +def test_remote_tokenize_model_backed_without_model_files(tmp_path, monkeypatch): + monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path)) + with mock_lancedb_connection(_jieba_fts_handler) as db: + table = db.open_table("test") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + table.tokenize("我们是中国人", index_name="text_idx") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + table.tokenize("我们是中国人", column="text") + + +@pytest.mark.asyncio +async def test_remote_tokenize_model_backed_without_model_files_async( + tmp_path, monkeypatch +): + monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path)) + async with mock_lancedb_connection_async(_jieba_fts_handler) as db: + table = await db.open_table("test") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + await table.tokenize("我们是中国人", index_name="text_idx") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + await table.tokenize("我们是中国人", column="text") + + def test_stats(): stats = { "total_bytes": 38, diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 84c1b3fef..b722dc4a1 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -8871,6 +8871,131 @@ mod tests { ); } + #[tokio::test] + #[serial_test::serial(language_model_env)] + async fn test_tokenize_model_backed_tokenizer_without_model_files() { + let model_home = tempfile::tempdir().unwrap(); + // SAFETY: This is only called in tests, serialized on the env key. + unsafe { + std::env::set_var("LANCE_LANGUAGE_MODEL_HOME", model_home.path()); + } + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let index_details = serde_json::json!({ + "base_tokenizer": "jieba/default", + "language": "English", + "with_position": false, + "max_token_length": 40, + "lower_case": true, + "stem": false, + "remove_stop_words": false, + "ascii_folding": false, + }) + .to_string(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": index_details, + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let result = table.tokenize("我们是中国人", "text_idx").await; + // SAFETY: This is only called in tests, serialized on the env key. + unsafe { + std::env::remove_var("LANCE_LANGUAGE_MODEL_HOME"); + } + + let err = result.unwrap_err(); + let Error::InvalidInput { message } = err else { + panic!("expected InvalidInput, got {err:?}"); + }; + // The original cause and the index context are preserved... + assert!(message.contains("Invalid directory path"), "{message}"); + assert!( + message.contains("full text search index 'text_idx'"), + "{message}" + ); + // ...and the message explains that this is a client-side limitation. + assert!( + message.contains("tokenize() runs on the client"), + "{message}" + ); + assert!(message.contains("'jieba/default'"), "{message}"); + assert!(message.contains("LANCE_LANGUAGE_MODEL_HOME"), "{message}"); + assert!( + message.contains("match the files the server uses"), + "{message}" + ); + } + + #[tokio::test] + async fn test_tokenize_unknown_tokenizer_error_has_no_model_files_hint() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let index_details = serde_json::json!({ + "base_tokenizer": "no_such_tokenizer", + "language": "English", + "with_position": false, + "max_token_length": 40, + "lower_case": true, + "stem": false, + "remove_stop_words": false, + "ascii_folding": false, + }) + .to_string(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": index_details, + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let err = table.tokenize("hello", "text_idx").await.unwrap_err(); + let message = err.to_string(); + assert!( + !message.contains("tokenize() runs on the client"), + "{message}" + ); + } + #[tokio::test] async fn test_tokenize_requires_existing_index_name() { let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 39e6f1911..3123eb776 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -2264,7 +2264,11 @@ impl Table { /// /// Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in /// the client process from index metadata. For remote tables, this means the - /// same tokenizer model files must also exist locally. + /// same tokenizer model files must also exist locally, and must match the + /// files the server uses: the client cannot detect a mismatch, so it may + /// return tokens the server's index never produces. If the model files are + /// missing, this returns an error naming the tokenizer and + /// `LANCE_LANGUAGE_MODEL_HOME`. pub async fn tokenize(&self, query: &str, index_name: &str) -> Result> { let indices = self.inner.list_indices().await?; let matches = indices @@ -2298,7 +2302,8 @@ impl Table { /// The column must have exactly one FTS index. Model-backed tokenizers such /// as `jieba/*` and `lindera/*` are rebuilt in the client process from /// index metadata. For remote tables, this means the same tokenizer model - /// files must also exist locally. + /// files must also exist locally, and must match the files the server uses + /// (see [`Self::tokenize`]). pub async fn tokenize_with_column(&self, query: &str, column: &str) -> Result> { let schema = self.inner.schema().await?; let (column, _) = resolve_arrow_field_path(schema.as_ref(), column)?; @@ -2354,13 +2359,30 @@ impl Table { ), } })?; + let model_backed_tokenizer = serde_json::from_str::(details) + .ok() + .and_then(|details| { + let tokenizer = details.get("base_tokenizer")?.as_str()?.to_string(); + (tokenizer.starts_with("jieba/") || tokenizer.starts_with("lindera/")) + .then_some(tokenizer) + }); tokenize(query, ¶ms).map_err(|err| match err { - Error::InvalidInput { message } => Error::InvalidInput { - message: format!( + Error::InvalidInput { message } => { + let mut message = format!( "{} for full text search index '{}' for {}", message, index.name, selector_description - ), - }, + ); + if let Some(tokenizer) = model_backed_tokenizer { + message.push_str(&format!( + ". tokenize() runs on the client: the '{}' tokenizer is rebuilt \ + in this process from the index metadata, so its model files must \ + be available under LANCE_LANGUAGE_MODEL_HOME and, for remote \ + tables, match the files the server uses", + tokenizer + )); + } + Error::InvalidInput { message } + } err => err, }) } From b6d15558ec8df6c15f25b1012e330802b4d1f830 Mon Sep 17 00:00:00 2001 From: Christopher Krah Date: Thu, 22 Oct 2026 21:18:52 +0200 Subject: [PATCH 27/64] fix(python): stop retrying a missing API key or provider package (#4521) `retry_with_exponential_backoff` retries every exception except an auth error (`python/python/lancedb/embeddings/utils.py:298-318`). For Cohere, Gemini, Jina and VoyageAI a missing API key raises from `api_key_not_found_help`, and a missing provider package raises `ImportError` everywhere. Neither can succeed on a retry, yet both sleep through all seven rounds and then surface wrapped in `Exception("Maximum number of retries...")`. Before, on `142d032`, a VoyageAI function with no key set (script below, `time.sleep` recorded, `random.seed(0)`): ``` $ python lance_retry_deterministic.py python/python raised: Exception('Maximum number of retries (7) exceeded.', ValueError('Please set the VOYAGEAI_API_KEY environment v attempts: 8 sleeps: [4, 13, 37, 93, 280, 788, 2811] slept 4026 s before a missing-API-key error surfaced ``` After: ``` $ python lance_retry_deterministic.py python/python raised: MissingApiKeyError('Please set the VOYAGEAI_API_KEY environment variable.') attempts: 1 sleeps: [] ```
lance_retry_deterministic.py ```python # loads the pure-Python embeddings package from the tree, native extension stubbed import os, random, sys, time, types TREE = sys.argv[1] pkg = types.ModuleType("lancedb"); pkg.__path__ = [TREE + "/lancedb"]; sys.modules["lancedb"] = pkg native = types.ModuleType("lancedb._lancedb"); native.__getattr__ = lambda name: None sys.modules["lancedb._lancedb"] = native import pyarrow as pa from lancedb.embeddings import get_registry random.seed(0) os.environ.pop("VOYAGE_API_KEY", None) slept = [] time.sleep = slept.append func = get_registry().get("voyageai").create(name="voyage-3") try: func.compute_source_embeddings_with_retry(pa.array(["hello"])) except Exception as e: print("raised:", repr(e)[:110]) print("attempts:", len(slept) + 1, "sleeps:", [round(s) for s in slept]) print("slept %.0f s before a missing-API-key error surfaced" % sum(slept)) ```
The message names the wrong variable: `voyageai.py:501` reads `VOYAGE_API_KEY`. I left that for a separate fix. `api_key_not_found_help` now raises `MissingApiKeyError`, a `ValueError` subclass, so existing `except ValueError` callers still catch it. That error and `ImportError` join the 401 exemption from #2995. A plain `ValueError` is still retried, because `bedrock.py:171` wraps boto3 throttling in one. No issue exists for this. OpenAI is not covered: with no key, `openai.OpenAI()` raises `OpenAIError`, which has no status code, so it still retries. I can follow up on that. ### Verification ``` # lance_stub: an empty stand-in for lancedb._lancedb, loaded from the source tree $ pytest -p lance_stub --noconftest python/tests/test_embeddings.py -k "retry or 401 or deterministic" -q 5 passed, 23 deselected in 0.59s # utils.py from 142d032, tests kept: 2 failed, 3 passed, 23 deselected in 0.74s $ pytest -p lance_stub --noconftest python/tests/test_voyageai_embeddings.py -q 18 passed in 0.18s # not run: the full suite or pyright; the native extension did not build here ``` Baseline on `142d032` has 2 of these tests. The new ones cover a missing key, a missing package, and a wrapped throttling error that must still retry. ruff 0.16.10 check and format are clean. @wjones127 you merged the 401 change. Is a `ValueError` subclass the way you want to mark this? Co-authored-by: Xuanwo --- python/python/lancedb/embeddings/utils.py | 13 ++++++- python/python/tests/test_embeddings.py | 45 +++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/embeddings/utils.py b/python/python/lancedb/embeddings/utils.py index f76c9b1a5..fb2a94ecb 100644 --- a/python/python/lancedb/embeddings/utils.py +++ b/python/python/lancedb/embeddings/utils.py @@ -320,6 +320,11 @@ def _is_non_retryable_error(error: Exception) -> bool: if error.status_code in (401, 403): return True + # A missing API key or a missing optional dependency fails the same way on + # every attempt + if isinstance(error, (MissingApiKeyError, ImportError)): + return True + return False @@ -337,9 +342,15 @@ def url_retrieve(url: str): raise ConnectionError("could not download {} due to {}".format(url, err)) +class MissingApiKeyError(ValueError): + """Raised when an embedding provider's API key is not set.""" + + def api_key_not_found_help(provider): logging.error("Could not find API key for %s", provider) - raise ValueError(f"Please set the {provider.upper()}_API_KEY environment variable.") + raise MissingApiKeyError( + f"Please set the {provider.upper()}_API_KEY environment variable." + ) def is_flash_attn_2_available(): diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index d6d128845..2438bfe80 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -633,6 +633,51 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +@patch("time.sleep") +@pytest.mark.parametrize( + "error", + [ + pytest.param("missing_key", id="missing-api-key"), + pytest.param(ImportError("No module named 'voyageai'"), id="missing-package"), + ], +) +def test_no_retry_on_deterministic_errors(mock_sleep, error): + from lancedb.embeddings.utils import ( + api_key_not_found_help, + retry_with_exponential_backoff, + ) + + def func(): + if error == "missing_key": + api_key_not_found_help("voyageai") + raise error + + mock_func = MagicMock(side_effect=func) + wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3) + + with pytest.raises((ValueError, ImportError)): + wrapped_func() + + assert mock_func.call_count == 1 + assert mock_sleep.call_count == 0 + + +@patch("time.sleep") +def test_retry_on_wrapped_provider_error(mock_sleep): + """Providers such as Bedrock wrap throttling in a plain ValueError.""" + from lancedb.embeddings.utils import retry_with_exponential_backoff + + mock_func = MagicMock( + side_effect=[ + ValueError("Error raised by boto3 client: ThrottlingException"), + [0.1], + ] + ) + wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3) + + assert wrapped_func() == [0.1] + assert mock_func.call_count == 2 + assert mock_sleep.call_count == 1 def test_openai_bad_request_partial_batch_recovery(): # Regression test for issue #1677 From 9f4cd1618bddf1041c8ebdaa0b0ccb33656dc57c Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:19:06 +0800 Subject: [PATCH 28/64] fix(python): handle empty batch iterables with a schema (#4538) ## Summary `create_table(data=iter(()), schema=...)` currently raises `RuntimeError: coroutine raised StopIteration`, although an empty list with the same schema works. A tuple of batches also fails because the helper calls `next()` before obtaining an iterator. Normalize the iterable once and use the supplied schema for an empty reader. Without a schema, raise a clear `ValueError`. Nonempty schema inference and batch casting are unchanged. Tests cover synchronous/asynchronous creation, empty/nonempty iterators and tuples, and schema-aware `add()` sanitization. The default schema-less `add()` rejection is deliberately unchanged. ## Validation - Built the current Rust/PyO3 extension from base `9f37f392` with the repository's `ci` profile; both Python-only arms use the same verified extension and locked Python dependency versions. - Complete `python/tests/test_table.py` + `python/tests/test_scannable.py`: **268 passed, 0 skipped**. The unchanged production code with identical tests has **24 failures / 244 passes**. - Both arms used four-CPU affinity. An initial unrestricted-affinity run exhausted thread creation under the private memory limit; that failed attempt is not counted as a pass. - Package-wide Ruff check/format pass. Repository-wide checks still report pre-existing findings in unchanged `docs/web-tests`; those files are not modified here. AI-assisted change, reproduced against unchanged production code and validated with a native source build. This does not claim the full repository/platform matrix or a performance improvement. Signed-off-by: Xucheng Zhou --- python/python/lancedb/table.py | 15 +++++-- python/python/tests/test_table.py | 75 +++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 124fa34b5..c7de4161d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -469,7 +469,7 @@ def _into_pyarrow_reader( ): return data.collect().to_arrow().to_reader() elif isinstance(data, Iterable): - return _iterator_to_reader(data) + return _iterator_to_reader(data, schema) else: raise TypeError( f"Unknown data type {type(data)}. " @@ -479,10 +479,19 @@ def _into_pyarrow_reader( ) -def _iterator_to_reader(data: Iterable) -> pa.RecordBatchReader: +def _iterator_to_reader( + data: Iterable, schema: Optional[pa.Schema] = None +) -> pa.RecordBatchReader: # Each batch is treated as it's own reader, mainly so we can # re-use the _into_pyarrow_reader logic. - first = _into_pyarrow_reader(next(data)) + data = iter(data) + try: + first_batch = next(data) + except StopIteration: + if schema is None: + raise ValueError("Cannot create table from empty iterable without a schema") + return pa.RecordBatchReader.from_batches(schema, []) + first = _into_pyarrow_reader(first_batch) schema = first.schema def gen(): diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 238ee5330..e23b478f5 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4847,6 +4847,81 @@ def test_stats(mem_db: DBConnection): assert stats_with_index["fragment_stats"]["num_fragments"] == 1 +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.parametrize("with_schema", [False, True]) +def test_create_table_recordbatch_iterable( + mem_db: DBConnection, iterable_type, empty, with_schema +): + schema = pa.schema([pa.field("id", pa.int64())]) + batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)] + data = iterable_type(batches) + if empty and not with_schema: + with pytest.raises(ValueError, match="empty iterable without a schema"): + mem_db.create_table("test", data=data) + return + + table = mem_db.create_table( + "test", data=data, schema=schema if with_schema else None + ) + assert table.schema == schema + assert table.to_arrow().to_pylist() == ([] if empty else [{"id": 1}]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.parametrize("with_schema", [False, True]) +async def test_create_table_recordbatch_iterable_async( + mem_db_async: AsyncConnection, iterable_type, empty, with_schema +): + schema = pa.schema([pa.field("id", pa.int64())]) + batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)] + data = iterable_type(batches) + if empty and not with_schema: + with pytest.raises(ValueError, match="empty iterable without a schema"): + await mem_db_async.create_table("test", data=data) + return + + table = await mem_db_async.create_table( + "test", data=data, schema=schema if with_schema else None + ) + assert await table.schema() == schema + assert (await table.to_arrow()).to_pylist() == ([] if empty else [{"id": 1}]) + + +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +def test_add_empty_recordbatch_iterable( + mem_db: DBConnection, iterable_type, on_bad_vectors +): + table = mem_db.create_table("test", [{"id": 1}]) + if on_bad_vectors == "error": + # The default scannable path still requires an input schema. This + # fix does not change its validation policy for schema-less input. + with pytest.raises(ValueError, match="empty iterator"): + table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + else: + # Sanitization has the existing table schema available. + table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + assert table.to_arrow().to_pylist() == [{"id": 1}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +async def test_add_empty_recordbatch_iterable_async( + mem_db_async: AsyncConnection, iterable_type, on_bad_vectors +): + table = await mem_db_async.create_table("test", [{"id": 1}]) + if on_bad_vectors == "error": + with pytest.raises(ValueError, match="empty iterator"): + await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + else: + await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + assert (await table.to_arrow()).to_pylist() == [{"id": 1}] + + def test_create_table_empty_list_with_schema(mem_db: DBConnection): """Test creating table with empty list data and schema From 2d54630a3c8cfc1faaec67289c864bba66578d2a Mon Sep 17 00:00:00 2001 From: Avanish Garg Date: Fri, 23 Oct 2026 00:52:24 +0530 Subject: [PATCH 29/64] fix(python): stop retrying permanent embedding errors and cap the backoff delay (#4542) Permanent embedding errors (HTTP 400/404/422 and other 4xx) were retried the full `max_retries` times with an uncapped, multiplicatively jittered backoff, so a typo'd model name looked like a hang of about an hour (median 50 min of sleeping with a 404; see the issue for the repro and numbers). Changes (`python/python/lancedb/embeddings/utils.py`, `embeddings/base.py`): - An HTTP status in 4xx other than 408/409/429 is raised immediately. The status is read from `status_code`, `http_status`, `response.status_code`, an integer `code` (google-api-core, urllib) or a botocore `response["ResponseMetadata"]["HTTPStatusCode"]`. AWS throttling reported as 400 stays retryable (matched by error code). `AuthenticationError` by name is unchanged. - Still retried: 408, 409, 429, >=500, and anything without a recognisable status (connection errors, timeouts, custom exceptions), so existing custom embedding functions behave as before. - New keyword `max_delay=60` (last in the signature, default keeps call sites working). `delay = min(delay * base, max_delay)` with equal jitter `delay/2 + random() * delay/2`. With defaults the waits are [1,2], [2,4], [4,8], [8,16], [16,32], [30,60], [30,60] s, worst case about 182 s instead of up to 2.5 h. - The "Maximum number of retries" exception now has `from e`. The `max_retries` comment in `base.py` describes the rules. Not changed: `embeddings/openai.py` (it swallows `BadRequestError`; open #4154/#4005/#3152 cover it), Retry-After handling (the code never used it), and the "C Data interface error" wrapping from the native layer (the original message is now included in it). Behaviour changes to be aware of: every 4xx except 408/409/429 fails fast, including 401/403 coming through `response.status_code` (requests/httpx), which used to be retried; a provider that signals transient failures with a 4xx and no AWS-style code is no longer retried; the first retry waits 1-2 s instead of 2-4 s. Tests (`python/python/tests/test_embeddings.py`): 48 new cases (23 test functions): permanent errors are not retried, transient errors are, botocore-style errors, the delay is capped, exact exponential growth without jitter, jitter bounds, no float overflow for a large `max_retries`, the original error is attached, and end to end through `table.add` and `table.search` with a stub function (one call and no sleep). On `main` 23 of them fail with the change stashed; the rest are guards that must not change. `test_embeddings.py`: 61 passed, 4 skipped. `ruff format`/`ruff check` are clean on the touched files. Verification caveat: `main`'s Python layer needs a newer native module than the 0.40.0 wheel, so the tests were run with the wheel's native module preloaded through a small shim (the same shim ran the baseline on `main`); a wider run of embeddings/pydantic/query/table tests has the same 36 failures before and after (that mismatch). Not run: real provider APIs (a fake OpenAI-style server was used), Node or Rust code. Fixes #4541 Co-authored-by: Xuanwo --- python/python/lancedb/embeddings/base.py | 8 +- python/python/lancedb/embeddings/utils.py | 103 +++++++-- python/python/tests/test_embeddings.py | 263 ++++++++++++++++++++++ 3 files changed, 354 insertions(+), 20 deletions(-) diff --git a/python/python/lancedb/embeddings/base.py b/python/python/lancedb/embeddings/base.py index 96ca6556c..c11914fab 100644 --- a/python/python/lancedb/embeddings/base.py +++ b/python/python/lancedb/embeddings/base.py @@ -38,9 +38,11 @@ class EmbeddingFunction(BaseModel, ABC): 3. ndims() which returns the number of dimensions of the vector column """ - max_retries: int = ( - 7 # Setting 0 disables retires. Maybe this should not be enabled by default, - ) + # Maximum number of retries for transient errors (rate limits, 5xx responses, + # connection errors, ...), with exponential backoff capped at 60 seconds between + # attempts. Permanent errors (HTTP 4xx other than 408/409/429) are never retried. + # Setting 0 disables retries. + max_retries: int = 7 _ndims: int = PrivateAttr() _original_args: dict = PrivateAttr() diff --git a/python/python/lancedb/embeddings/utils.py b/python/python/lancedb/embeddings/utils.py index fb2a94ecb..ff8fc6a18 100644 --- a/python/python/lancedb/embeddings/utils.py +++ b/python/python/lancedb/embeddings/utils.py @@ -241,15 +241,30 @@ def retry_with_exponential_backoff( exponential_base: float = 2, jitter: bool = True, max_retries: int = 7, + max_delay: float = 60, ): """Retry a function with exponential backoff. + Only errors that may succeed when repeated are retried. Errors that carry an + HTTP status code are retried when the status is 408, 409, 429 or >= 500; any + other 4xx status (for example 400, 401, 403, 404 or 422) is a permanent failure + that is raised immediately. Errors without a recognisable status code (for + example connection errors, timeouts, or exceptions raised by a custom embedding + function) are retried. + + The delay before the n-th retry is ``initial_delay * exponential_base ** n``, + never more than ``max_delay``. With ``jitter`` enabled, a random amount of up to + half of that delay is removed ("equal jitter"), so the actual wait is always in + the range ``[delay / 2, delay]``. + Args: func (function): The function to be retried. initial_delay (float): Initial delay in seconds (default is 1). exponential_base (float): The base for exponential backoff (default is 2). jitter (bool): Whether to add jitter to the delay (default is True). max_retries (int): Maximum number of retries (default is 7). + max_delay (float): Maximum delay between two attempts, in seconds + (default is 60). Returns: function: The decorated function. @@ -265,13 +280,12 @@ def retry_with_exponential_backoff( try: return func(*args, **kwargs) - # Currently retrying on all exceptions as there is no way to know the - # format of the error msgs used by different APIs. We'll log the error - # and say that it is assumed that if this portion errors out, it's due - # to rate limit but the user should check the error message to be sure. + # Errors are retried unless they are known to be permanent (see + # `_is_non_retryable_error`). Custom embedding functions can raise + # arbitrary exceptions, so unknown errors are assumed to be transient. except Exception as e: # noqa: PERF203 - # Don't retry on authentication errors (e.g., OpenAI 401) - # These are permanent failures that won't be fixed by retrying + # Don't retry permanent failures (e.g., OpenAI 401 or 404). They + # won't be fixed by retrying and would only delay the real error. if _is_non_retryable_error(e): raise @@ -280,24 +294,72 @@ def retry_with_exponential_backoff( if num_retries > max_retries: raise Exception( f"Maximum number of retries ({max_retries}) exceeded.", e - ) + ) from e - delay *= exponential_base * (1 + jitter * random.random()) + delay = min(delay * exponential_base, max_delay) + sleep_for = delay + if jitter: + sleep_for = delay / 2 + random.random() * delay / 2 logging.warning( "Error occurred: %s \n Retrying in %s seconds (retry %s of %s) \n", e, - delay, + sleep_for, num_retries, max_retries, ) - time.sleep(delay) + time.sleep(sleep_for) return wrapper +# HTTP statuses in the 4xx range that are worth retrying: request timeout, +# conflict (the OpenAI SDK treats it as transient) and rate limiting. +_RETRYABLE_CLIENT_STATUSES = (408, 409, 429) + +# AWS error codes that are reported with a 4xx status but are transient. +_RETRYABLE_AWS_ERROR_CODES = ( + "Throttl", + "TooManyRequests", + "RequestLimitExceeded", + "ProvisionedThroughputExceeded", + "ModelNotReady", + "ModelTimeout", +) + + +def _get_http_status(error: Exception): + """Best-effort extraction of an HTTP status code from an exception. + + Provider SDKs expose it in different ways: ``status_code`` (openai, cohere), + ``http_status`` (voyageai), ``response.status_code`` (requests, httpx), + ``code`` (google-api-core, urllib) or ``response["ResponseMetadata"]`` + (botocore). Returns None if no status can be found. + """ + candidates = [ + getattr(error, "status_code", None), + getattr(error, "http_status", None), + ] + response = getattr(error, "response", None) + if isinstance(response, dict): + candidates.append(response.get("ResponseMetadata", {}).get("HTTPStatusCode")) + else: + candidates.append(getattr(response, "status_code", None)) + candidates.append(getattr(error, "code", None)) + + for status in candidates: + if isinstance(status, int) and not isinstance(status, bool): + if 100 <= status <= 599: + return status + return None + + def _is_non_retryable_error(error: Exception) -> bool: """Check if an error should not be retried. + An error is non-retryable if it is an authentication error or if it carries an + HTTP client-error status (4xx) other than 408, 409 and 429. Errors without a + recognisable status are treated as retryable. + Args: error: The exception to check @@ -314,18 +376,25 @@ def _is_non_retryable_error(error: Exception) -> bool: if error_type == "AuthenticationError": return True - # Check for other common non-retryable HTTP status codes - # 401 Unauthorized, 403 Forbidden - if hasattr(error, "status_code"): - if error.status_code in (401, 403): - return True - # A missing API key or a missing optional dependency fails the same way on # every attempt if isinstance(error, (MissingApiKeyError, ImportError)): return True - return False + status = _get_http_status(error) + if status is None or not 400 <= status < 500: + return False + if status in _RETRYABLE_CLIENT_STATUSES: + return False + + # AWS reports throttling with a 400 status; the error code tells them apart. + response = getattr(error, "response", None) + if isinstance(response, dict): + aws_code = str(response.get("Error", {}).get("Code", "")) + if aws_code.startswith(_RETRYABLE_AWS_ERROR_CODES): + return False + + return True def url_retrieve(url: str): diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 2438bfe80..129fa6742 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -633,6 +633,269 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +class _HttpStatusError(Exception): + """Stand-in for provider SDK errors (openai, cohere, ...) with a status code.""" + + def __init__(self, status_code, message="error"): + super().__init__(message) + self.status_code = status_code + + +class _ResponseStatusError(Exception): + """Stand-in for requests/httpx errors, which expose `response.status_code`.""" + + def __init__(self, status_code): + super().__init__(f"HTTP {status_code}") + self.response = MagicMock(status_code=status_code) + + +def _run_with_retry(side_effect, **kwargs): + """Run a function through the retry helper with a mocked `time.sleep`. + + Returns (mock_func, sleeps, error). + """ + from lancedb.embeddings.utils import retry_with_exponential_backoff + + mock_func = MagicMock(side_effect=side_effect) + sleeps = [] + error = None + with patch("time.sleep", side_effect=sleeps.append): + try: + retry_with_exponential_backoff(mock_func, **kwargs)() + except Exception as e: + error = e + return mock_func, sleeps, error + + +@pytest.mark.parametrize("status", [400, 401, 403, 404, 405, 410, 413, 422]) +@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError]) +def test_retry_does_not_retry_permanent_http_errors(status, make_error): + err = make_error(status) + mock_func, sleeps, error = _run_with_retry(err, max_retries=7) + + assert error is err + assert mock_func.call_count == 1 + assert sleeps == [] + + +@pytest.mark.parametrize("status", [408, 409, 429, 500, 502, 503, 504]) +@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError]) +def test_retry_retries_transient_http_errors(status, make_error): + mock_func, sleeps, error = _run_with_retry( + [make_error(status), make_error(status), "result"], max_retries=7 + ) + + assert error is None + assert mock_func.call_count == 3 + assert len(sleeps) == 2 + + +@pytest.mark.parametrize( + "exc", + [ + ConnectionError("connection reset"), + TimeoutError("timed out"), + OSError("network is unreachable"), + ValueError("custom embedding function failure"), + RuntimeError("boom"), + Exception("plain exception"), + ], +) +def test_retry_retries_errors_without_http_status(exc): + mock_func, sleeps, error = _run_with_retry([exc, exc, "result"], max_retries=7) + + assert error is None + assert mock_func.call_count == 3 + assert len(sleeps) == 2 + + +def test_retry_error_status_attributes_of_other_types_are_ignored(): + # `code` is a string for openai errors ("model_not_found"), and arbitrary + # exceptions may carry unrelated numeric attributes. Neither is an HTTP status. + class WeirdError(Exception): + code = "model_not_found" + status_code = None + + class ExitCodeError(Exception): + code = 2 + + for exc in (WeirdError("a"), ExitCodeError("b")): + mock_func, sleeps, error = _run_with_retry(exc, max_retries=2) + assert mock_func.call_count == 3 + assert len(sleeps) == 2 + + +def test_retry_botocore_style_errors(): + # botocore ClientError exposes `response["ResponseMetadata"]["HTTPStatusCode"]`. + # Bedrock reports throttling with a 400, which must stay retryable. + class ClientError(Exception): + def __init__(self, status, code): + super().__init__(code) + self.response = { + "ResponseMetadata": {"HTTPStatusCode": status}, + "Error": {"Code": code}, + } + + throttled = ClientError(400, "ThrottlingException") + mock_func, sleeps, error = _run_with_retry([throttled, "result"], max_retries=3) + assert error is None + assert mock_func.call_count == 2 + + invalid = ClientError(400, "ValidationException") + mock_func, sleeps, error = _run_with_retry(invalid, max_retries=3) + assert error is invalid + assert mock_func.call_count == 1 + assert sleeps == [] + + +def test_retry_authentication_error_by_name_not_retried(): + class AuthenticationError(Exception): + pass + + mock_func, sleeps, error = _run_with_retry(AuthenticationError("bad key")) + assert isinstance(error, AuthenticationError) + assert mock_func.call_count == 1 + assert sleeps == [] + + +def test_retry_gives_up_after_max_retries_with_original_error_attached(): + original = _HttpStatusError(503, "overloaded") + mock_func, sleeps, error = _run_with_retry(original, max_retries=3) + + assert mock_func.call_count == 4 + assert len(sleeps) == 3 + assert "Maximum number of retries (3) exceeded" in str(error) + assert error.args[1] is original + assert error.__cause__ is original + + +def test_retry_delay_is_capped_by_max_delay(): + # Default settings (7 retries, 60s cap): the old implementation could sleep for + # up to ~47 minutes in a single retry and hours in total. + exc = _HttpStatusError(503) + for seed in range(50): + import random + + random.seed(seed) + _, sleeps, _ = _run_with_retry(exc) + assert len(sleeps) == 7 + assert max(sleeps) <= 60 + assert sum(sleeps) <= 7 * 60 + + # Custom cap + _, sleeps, _ = _run_with_retry(exc, max_retries=10, max_delay=5) + assert len(sleeps) == 10 + assert max(sleeps) <= 5 + + +def test_retry_delay_without_jitter_is_exact_exponential_then_capped(): + _, sleeps, _ = _run_with_retry( + _HttpStatusError(500), + initial_delay=1, + exponential_base=2, + jitter=False, + max_retries=8, + max_delay=30, + ) + assert sleeps == [2, 4, 8, 16, 30, 30, 30, 30] + + +def test_retry_jitter_is_bounded(): + # Each wait is within [delay / 2, delay], where delay grows by exponential_base + # per retry (no random factor is compounded across retries). + exc = _HttpStatusError(500) + for seed in range(50): + import random + + random.seed(seed) + _, sleeps, _ = _run_with_retry(exc, max_retries=8, max_delay=100) + for i, slept in enumerate(sleeps): + delay = min(2 ** (i + 1), 100) + assert delay / 2 <= slept <= delay + + # Extremes of the random source map to the ends of the range + with patch("random.random", return_value=0.0): + _, sleeps, _ = _run_with_retry(exc, max_retries=3) + assert sleeps == [1, 2, 4] + with patch("random.random", return_value=1.0): + _, sleeps, _ = _run_with_retry(exc, max_retries=3) + assert sleeps == [2, 4, 8] + + +def test_retry_large_max_retries_does_not_overflow(): + _, sleeps, _ = _run_with_retry(_HttpStatusError(503), max_retries=1500) + assert len(sleeps) == 1500 + assert max(sleeps) <= 60 + + +def _register_failing_embedding(name, error): + """Register an embedding function that always raises `error`. + + Returns the class and a list that records every call. The function is + re-created from table metadata, so the call count has to live outside of it. + """ + calls = [] + + @register(name) + class FailingEmbeddingFunction(TextEmbeddingFunction): + def ndims(self): + return 4 + + def generate_embeddings(self, texts): + calls.append(len(texts)) + raise error + + return FailingEmbeddingFunction, calls + + +@pytest.mark.parametrize("status", [400, 404, 422]) +def test_permanent_embedding_error_is_not_retried(tmp_path, status): + """A permanent provider error (e.g. a misspelled model name) must surface + immediately from both `table.add` and `table.search` instead of being retried + with minutes of backoff.""" + error = _HttpStatusError(status, "The model `typo-embed` does not exist") + cls, calls = _register_failing_embedding(f"failing-embedding-{status}", error) + func = cls.create() # default max_retries=7 + + class Schema(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table("t", schema=Schema) + + with patch("time.sleep") as mock_sleep: + with pytest.raises(Exception, match="The model `typo-embed` does not exist"): + table.add([{"text": "hello"}]) + assert len(calls) == 1 + + with pytest.raises(_HttpStatusError, match="does not exist"): + table.search("hello").limit(1).to_list() + assert len(calls) == 2 + + mock_sleep.assert_not_called() + + +def test_transient_embedding_error_is_retried_in_table_add(tmp_path): + cls, calls = _register_failing_embedding( + "failing-embedding-503", _HttpStatusError(503) + ) + func = cls.create(max_retries=2) + + class Schema(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table("t", schema=Schema) + + with patch("time.sleep") as mock_sleep: + with pytest.raises(Exception, match="Maximum number of retries"): + table.add([{"text": "hello"}]) + + assert len(calls) == 3 + assert mock_sleep.call_count == 2 + assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list) @patch("time.sleep") @pytest.mark.parametrize( "error", From bc4c32b779de70c4c62a8411ea7eca1c165c9f7e Mon Sep 17 00:00:00 2001 From: Avanish Garg Date: Fri, 23 Oct 2026 00:52:39 +0530 Subject: [PATCH 30/64] fix: keep tagged versions when pruning in optimize (#4544) `optimize()` fails once a tagged version is older than the cleanup cutoff, although the docs say tagged versions are exempt from cleanup, and the failed prune also skips the index-optimize step (repro and numbers in the issue). Change: `OptimizeAction::Prune` already exposes `error_if_tagged_old_versions: Option`, so no new API. `cleanup_old_versions` and `cleanup_old_versions_before` in `rust/lancedb/src/table/optimize.rs` now use `unwrap_or(false)`, which covers `OptimizeAction::Prune`, `OptimizeAction::All` and `optimize_prune_before`. The Python and Node bindings pass `None`, which now means "retain tagged versions"; I left `python/src/table.rs` and `nodejs/src/table.rs` untouched (this also avoids a conflict with #3870, which edits that region). `Some(true)` still opts in to the error. Doc comments updated; the sentence in `docs/web/management/table.mdx` is accurate once this lands. Tests: Rust `test_optimize_prune_retains_tagged_versions` (`Some(true)` still errors; the default retains the tag; `optimize_prune_before`; `checkout_tag` plus a data read), a Python sync test with a scalar index (index step runs, a second optimize works, untagged versions are pruned, `checkout("release-1")` works), a Python async test and a Node test. **What I could and could not verify (please read):** - I could NOT compile Rust here (not enough disk), so `cargo check`, `cargo test`, the Python/Node binding builds and therefore the fixed-path assertions of the new Rust, Python and Node tests have not been run. This is why the PR is a draft; I will mark it ready once CI is green and fix whatever it finds. - Ran: both new Python tests fail on the released 0.40.0 wheel with the exact error above (so they do catch the bug); pylance 13.0.0 confirms the semantics (`cleanup_old_versions(error_if_tagged_old_versions=False)` removes the untagged versions and keeps the tagged one, still readable with `lance.dataset(version="keep")`); in Lance v14.0.0-beta.7 `cleanup.rs` the policy field is `bool`, the builder default is `true`, and Lance's own `auto_cleanup_hook` uses `Some(false)`; `rustfmt --edition 2024 --check` passes on `optimize.rs`. - The Rust and Node tests rely on timestamp ordering (`nextMillisecond()`, `days(0)`/`now()`), so they could be flaky: tell me if you prefer another approach. Behaviour change: anyone who relied on the error to learn that a tag blocks cleanup now gets a silent retain. Not changed: the deprecated sync `LanceTable.cleanup_old_versions` calls pylance directly and still errors on tagged versions (the keyword may not exist in the oldest supported pylance); `freshness.rs:1377` and `refresh.rs:1077` still pass `None` (I did not check what those paths do with tagged versions). Fixes #4543 --- nodejs/__test__/table.test.ts | 15 +++++ python/python/tests/test_table.py | 51 ++++++++++++++ rust/lancedb/src/table/optimize.rs | 105 +++++++++++++++++++++++++++-- 3 files changed, 167 insertions(+), 4 deletions(-) diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 7d10f8b59..84340c72c 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3742,6 +3742,21 @@ describe("when optimizing a dataset", () => { expect(stats.prune.oldVersionsRemoved).toBe(2); }); + it("retains tagged versions older than the cleanup cutoff", async () => { + await (await table.tags()).create("keep", 1); + await table.add([{ id: 3 }]); + + const cutoff = await nextMillisecond(); + const stats = await table.optimize({ cleanupOlderThan: cutoff }); + // Versions 2 and 3 are pruned, tagged version 1 is retained + expect(stats.prune.oldVersionsRemoved).toBe(2); + + await table.checkout("keep"); + expect(await table.countRows()).toBe(1); + await table.checkoutLatest(); + expect(await table.countRows()).toBe(3); + }); + it("delete unverified", async () => { const version = await table.version(); const versionFile = `${tmpDir.name}/${table.name}.lance/_versions/${String( diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index e23b478f5..b55abfa7e 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4792,6 +4792,57 @@ def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection): assert all("concurrent" in str(w.message) for w in seen) +def test_optimize_retains_tagged_versions(tmp_db: DBConnection): + # A tagged version older than the cleanup cutoff must be retained, not make + # optimize() fail part-way (after compaction, before the index step). + table = tmp_db.create_table( + "test", data=pa.table({"id": list(range(100)), "text": ["a"] * 100}) + ) + table.create_index("id", config=BTree()) + table.tags.create("release-1", table.version) + tagged_version = table.version + for i in range(3): + table.add(pa.table({"id": [1000 + i], "text": ["new"]})) + assert table.stats()["fragment_stats"]["num_fragments"] == 4 + assert table.index_stats("id_idx").num_unindexed_rows == 3 + sleep(1.1) + + table.optimize(cleanup_older_than=timedelta(seconds=1)) + + # The index step ran after the prune step + assert table.index_stats("id_idx").num_unindexed_rows == 0 + # It stays optimizable + table.optimize(cleanup_older_than=timedelta(seconds=1)) + assert table.stats()["fragment_stats"]["num_fragments"] == 1 + + # Untagged old versions were pruned, the tagged one was retained + versions = [v["version"] for v in table.list_versions()] + assert tagged_version in versions + for added in range(tagged_version + 1, tagged_version + 4): + assert added not in versions + assert len(table) == 103 + + table.checkout("release-1") + assert len(table) == 100 + assert table.to_arrow().num_rows == 100 + + +@pytest.mark.asyncio +async def test_optimize_retains_tagged_versions_async(tmp_db_async: AsyncConnection): + table = await tmp_db_async.create_table("test", data=[{"x": [1]}]) + await table.tags.create("keep", await table.version()) + await table.add(data=[{"x": [2]}]) + await table.add(data=[{"x": [3]}]) + sleep(1.1) + + stats = await table.optimize(cleanup_older_than=timedelta(seconds=1)) + assert stats.prune.old_versions_removed >= 1 + + await table.checkout("keep") + assert await table.count_rows() == 1 + assert await table.query().to_arrow() == pa.table({"x": [[1]]}) + + def test_replace_field_metadata(tmp_path): db = lancedb.connect(tmp_path) table = db.create_table("my_table", data=[{"x": 0}]) diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index 2a502b059..ebe03359c 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -69,7 +69,11 @@ pub enum OptimizeAction { /// **WARNING**: This should only be set to true if you can guarantee that no other process is /// currently working on this dataset. Otherwise the dataset could be put into a corrupted state. delete_unverified: Option, + /// Tagged versions are never removed, regardless of their age. + /// /// If true, an error will be returned if there are any old versions that are still tagged. + /// If false or `None` (the default), tagged versions are silently retained and all other + /// old versions are removed. error_if_tagged_old_versions: Option, }, /// Optimize the indices @@ -125,6 +129,12 @@ pub(crate) async fn optimize_indices(table: &NativeTable, options: &OptimizeOpti /// no other process is currently working on this dataset. Otherwise the /// dataset could be put into a corrupted state. /// +/// * `error_if_tagged_old_versions` - Tagged versions are always retained. If +/// this is true, an error is returned when a tagged version is older than the +/// cutoff. Defaults to false (the tagged versions are retained silently), +/// unlike Lance's own default, so that `optimize` does not fail (and skip the +/// steps after pruning) once a tagged version ages past the cutoff. +/// /// This calls into [lance::dataset::Dataset::cleanup_old_versions] and /// returns the result. pub(crate) async fn cleanup_old_versions( @@ -136,7 +146,11 @@ pub(crate) async fn cleanup_old_versions( table.dataset.ensure_mutable()?; let dataset = table.dataset.get().await?; let stats = dataset - .cleanup_old_versions(older_than, delete_unverified, error_if_tagged_old_versions) + .cleanup_old_versions( + older_than, + delete_unverified, + Some(error_if_tagged_old_versions.unwrap_or(false)), + ) .await?; // Computed-column signature sidecars live outside lance's directories; // drop the ones the surviving versions no longer reference. @@ -161,9 +175,8 @@ pub(crate) async fn cleanup_old_versions_before( if let Some(delete_unverified) = delete_unverified { policy = policy.delete_unverified(delete_unverified); } - if let Some(error_if_tagged_old_versions) = error_if_tagged_old_versions { - policy = policy.error_if_tagged_old_versions(error_if_tagged_old_versions); - } + // Tagged versions are retained rather than raising an error unless asked. + policy = policy.error_if_tagged_old_versions(error_if_tagged_old_versions.unwrap_or(false)); let stats = dataset.cleanup_with_policy(policy.build()).await?; // Computed-column signature sidecars live outside lance's directories; // drop the ones the surviving versions no longer reference. @@ -511,6 +524,90 @@ mod tests { assert_eq!(all_values, expected); } + #[tokio::test] + async fn test_optimize_prune_retains_tagged_versions() { + let tmpdir = tempfile::tempdir().unwrap(); + let conn = connect(tmpdir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let make_batch = |start: i32| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(start..start + 10))], + ) + .unwrap() + }; + + let table = conn + .create_table("test_prune_tagged", make_batch(0)) + .execute() + .await + .unwrap(); + // Tag the first version, then create more versions after it. + table.tags().await.unwrap().create("keep", 1).await.unwrap(); + for i in 1..4 { + table.add(make_batch(i * 10)).execute().await.unwrap(); + } + assert_eq!(table.list_versions().await.unwrap().len(), 4); + + // Explicitly asking for an error still errors, and leaves everything in place. + let err = table + .optimize(OptimizeAction::Prune { + older_than: Some(chrono::Duration::try_days(0).unwrap()), + delete_unverified: Some(true), + error_if_tagged_old_versions: Some(true), + }) + .await + .unwrap_err(); + assert!(err.to_string().contains("tagged version"), "{err}"); + assert_eq!(table.list_versions().await.unwrap().len(), 4); + + // By default, tagged versions are retained and the other old versions are removed. + let stats = table + .optimize(OptimizeAction::Prune { + older_than: Some(chrono::Duration::try_days(0).unwrap()), + delete_unverified: Some(true), + error_if_tagged_old_versions: None, + }) + .await + .unwrap(); + assert_eq!(stats.prune.unwrap().old_versions, 2); + let mut versions = table + .list_versions() + .await + .unwrap() + .into_iter() + .map(|v| v.version) + .collect::>(); + versions.sort(); + assert_eq!(versions, vec![1, 4]); + + // Same through the absolute-cutoff entry point used by the Node binding. + table.add(make_batch(40)).execute().await.unwrap(); + let stats = table + .optimize_prune_before(chrono::Utc::now(), Some(true), None) + .await + .unwrap(); + assert_eq!(stats.prune.unwrap().old_versions, 1); + + // The tagged version is still readable, including its data files. + assert_eq!(table.count_rows(None).await.unwrap(), 50); + table.checkout_tag("keep").await.unwrap(); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 10); + } + #[tokio::test] async fn test_optimize_index() { let conn = connect("memory://").execute().await.unwrap(); From 5fa30ae404fc7df335c8031ae474400b834c62b2 Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:22:53 -0400 Subject: [PATCH 31/64] fix(python): honor explicit Cohere and VoyageAI reranker API keys (#4546) `CohereReranker` and `VoyageAIReranker` read the environment variable before the `api_key` argument. When `COHERE_API_KEY` or `VOYAGE_API_KEY` is set, the reranker ignores an explicit `api_key` and uses the account from the environment. The docs say the opposite. `docs/web/integrations/reranking/cohere.mdx` and `voyageai.mdx` use the environment variable only when `api_key` is not provided. `JinaReranker` and `WatsonxReranker` already put the explicit key first, and #4515 makes the same change for `OpenaiReranker`. ## Fix `CohereReranker._client` and `VoyageAIReranker._client` now use `self.api_key` first and the environment variable as the fallback. ## Tests `python/python/tests/test_reranker_credentials.py` replaces the SDK module with a stub whose `Client` records the API key. It covers both keys set, each key alone, and no key. With the stub, the test runs without `cohere` or `voyageai` installed. - On `main`, the 2 cases with both keys set fail and the other 6 pass. With the fix, all 8 pass. - `test_rerankers.py` with the new file: 24 passed, 13 skipped (the skipped tests need real API keys). - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed files copied over it. In that wheel, both reranker files are byte-identical to `main` before this change. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/rerankers/cohere.py | 2 +- python/python/lancedb/rerankers/voyageai.py | 2 +- .../python/tests/test_reranker_credentials.py | 74 +++++++++++++++++++ 3 files changed, 76 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_reranker_credentials.py diff --git a/python/python/lancedb/rerankers/cohere.py b/python/python/lancedb/rerankers/cohere.py index efd752c22..a1e797f15 100644 --- a/python/python/lancedb/rerankers/cohere.py +++ b/python/python/lancedb/rerankers/cohere.py @@ -62,7 +62,7 @@ class CohereReranker(Reranker): "COHERE_API_KEY not set. Either set it in your environment or \ pass it as `api_key` argument to the CohereReranker." ) - return cohere.Client(os.environ.get("COHERE_API_KEY") or self.api_key) + return cohere.Client(self.api_key or os.environ.get("COHERE_API_KEY")) def _rerank(self, result_set: pa.Table, query: str): result_set = self._handle_empty_results(result_set) diff --git a/python/python/lancedb/rerankers/voyageai.py b/python/python/lancedb/rerankers/voyageai.py index d3dbd59ac..28847564b 100644 --- a/python/python/lancedb/rerankers/voyageai.py +++ b/python/python/lancedb/rerankers/voyageai.py @@ -64,7 +64,7 @@ class VoyageAIReranker(Reranker): pass it as `api_key` argument to the VoyageAIReranker." ) return voyageai.Client( - api_key=os.environ.get("VOYAGE_API_KEY") or self.api_key, + api_key=self.api_key or os.environ.get("VOYAGE_API_KEY"), ) def _rerank(self, result_set: pa.Table, query: str): diff --git a/python/python/tests/test_reranker_credentials.py b/python/python/tests/test_reranker_credentials.py new file mode 100644 index 000000000..c8f63a9ef --- /dev/null +++ b/python/python/tests/test_reranker_credentials.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import sys +import types + +import pytest + +from lancedb.rerankers import CohereReranker, VoyageAIReranker + +RERANKERS = [ + pytest.param("cohere", "COHERE_API_KEY", CohereReranker, {}, id="cohere"), + pytest.param( + "voyageai", + "VOYAGE_API_KEY", + VoyageAIReranker, + {"model_name": "rerank-2.5"}, + id="voyageai", + ), +] + + +def fake_sdk(monkeypatch, module_name): + """Replace the SDK module with one whose Client records the API key.""" + keys = [] + + def client(api_key=None): + keys.append(api_key) + return object() + + module = types.ModuleType(module_name) + module.Client = client + monkeypatch.setitem(sys.modules, module_name, module) + return keys + + +@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS) +@pytest.mark.parametrize( + "explicit_key, environment_key, expected_key", + [ + ("explicit-test-key", "environment-test-key", "explicit-test-key"), + ("explicit-test-key", None, "explicit-test-key"), + (None, "environment-test-key", "environment-test-key"), + ], +) +def test_reranker_api_key_precedence( + monkeypatch, + module_name, + env_var, + reranker_cls, + kwargs, + explicit_key, + environment_key, + expected_key, +): + keys = fake_sdk(monkeypatch, module_name) + if environment_key is None: + monkeypatch.delenv(env_var, raising=False) + else: + monkeypatch.setenv(env_var, environment_key) + + reranker_cls(api_key=explicit_key, **kwargs)._client + assert keys == [expected_key] + + +@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS) +def test_reranker_requires_api_key( + monkeypatch, module_name, env_var, reranker_cls, kwargs +): + keys = fake_sdk(monkeypatch, module_name) + monkeypatch.delenv(env_var, raising=False) + with pytest.raises(ValueError, match=f"{env_var} not set"): + reranker_cls(**kwargs)._client + assert keys == [] From 353ce4d80c9e66f8f1ac9e49dbcec0ec899f1f9d Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:23:07 -0400 Subject: [PATCH 32/64] fix(python): give each JinaEmbeddings instance its own session (#4547) `JinaEmbeddings` stores its `requests` session in a `ClassVar`. The first instance builds the session with its own API key, and every later instance reuses that session. So a second instance with a different `api_key` sends the first instance's key. ## Fix The session is now a per-instance `cached_property`, the same pattern `OpenAIEmbeddings` uses for `_openai_client`. `_init_client` and the `ClassVar` are gone. An instance with no `api_key` and no `JINA_API_KEY` now raises the missing-key error. Before, it reused the session of an earlier instance and sent that instance's key. ## Tests `test_jina_embeddings_use_their_own_api_key` creates two instances with different keys through the registry. It stubs `requests.Session.post` and records the `Authorization` header of each request. - On `main`, the test fails because both requests send `Bearer key-a`. With the fix, it passes. - `test_embeddings.py`: 18 passed, 8 skipped. - An instance still pickles after the session exists, and `safe_model_dump()` still returns only `api_key`. - On pydantic 2.7.4, the minimum version, a model with the same `cached_property` gets one session per instance and leaves it out of `model_dump()`. - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed file copied over it. In that wheel, `embeddings/jinaai.py` is byte-identical to `main` before this change. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/embeddings/jinaai.py | 26 +++++++++++----------- python/python/tests/test_embeddings.py | 25 +++++++++++++++++++++ 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/python/python/lancedb/embeddings/jinaai.py b/python/python/lancedb/embeddings/jinaai.py index f6ab601b3..25b5a940b 100644 --- a/python/python/lancedb/embeddings/jinaai.py +++ b/python/python/lancedb/embeddings/jinaai.py @@ -5,9 +5,10 @@ import os import io import base64 +from functools import cached_property from urllib.parse import urlparse from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, List, Union, Optional, Any, Dict +from typing import TYPE_CHECKING, List, Union, Optional, Any, Dict import numpy as np import pyarrow as pa @@ -51,7 +52,6 @@ class JinaEmbeddings(EmbeddingFunction): name: str = "jina-clip-v1" api_key: Optional[str] = None - _session: ClassVar = None def ndims(self): # TODO: fix hardcoding @@ -206,8 +206,7 @@ class JinaEmbeddings(EmbeddingFunction): texts: list[str] or np.ndarray (of str) The texts to embed """ - self._init_client() - resp = JinaEmbeddings._session.post( # type: ignore + resp = self._session.post( API_URL, json={"input": input, "model": self.name} ).json() if "data" not in resp: @@ -220,14 +219,15 @@ class JinaEmbeddings(EmbeddingFunction): return [result["embedding"] for result in sorted_embeddings] - def _init_client(self): + @cached_property + def _session(self): import requests - if JinaEmbeddings._session is None: - if self.api_key is None and os.environ.get("JINA_API_KEY") is None: - api_key_not_found_help("jina") - api_key = self.api_key or os.environ.get("JINA_API_KEY") - JinaEmbeddings._session = requests.Session() - JinaEmbeddings._session.headers.update( - {"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"} - ) + if self.api_key is None and os.environ.get("JINA_API_KEY") is None: + api_key_not_found_help("jina") + api_key = self.api_key or os.environ.get("JINA_API_KEY") + session = requests.Session() + session.headers.update( + {"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"} + ) + return session diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 129fa6742..e319d3048 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -1053,3 +1053,28 @@ def test_jina_generate_image_input_dict_local_path(tmp_path): image_dict = JinaEmbeddings._generate_image_input_dict(image) assert "image" in image_dict assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0 + + +def test_jina_embeddings_use_their_own_api_key(monkeypatch): + """Each JinaEmbeddings instance must send its own API key.""" + pytest.importorskip("requests") + monkeypatch.delenv("JINA_API_KEY", raising=False) + registry = get_registry() + registry.set_var("jina_key_a", "key-a") + registry.set_var("jina_key_b", "key-b") + func_a = registry.get("jina").create(api_key="$var:jina_key_a") + func_b = registry.get("jina").create(api_key="$var:jina_key_b") + + sent = [] + + def fake_post(session, url, json=None, **kwargs): + sent.append(session.headers["Authorization"]) + response = MagicMock() + response.json.return_value = {"data": [{"index": 0, "embedding": [0.0]}]} + return response + + with patch("requests.Session.post", fake_post): + func_a.compute_source_embeddings(["hello"]) + func_b.compute_source_embeddings(["hello"]) + + assert sent == ["Bearer key-a", "Bearer key-b"] From 2af7d88ec7e974484e1be0d9f7bfba53732e0ef6 Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:23:21 -0400 Subject: [PATCH 33/64] fix(python): keep OpenaiReranker scores on the rows they rate (#4549) `OpenaiReranker._rerank` asks the chat model for a relevance score per document. The model can return the documents in any order, for example sorted by score. The reranker then replaces the `text` column with the documents in that order and appends the scores in that order. The other columns keep the input order. So after a rerank, the `text` and `_relevance_score` of a row belong to a different row. The `id`, the vector and every other column still describe the original row. If the model leaves out a document, `append_column` raises `ArrowInvalid`, because the new column has fewer rows than the table. ## Fix `_rerank` keeps the `text` column unchanged, so it also keeps its position and type. It maps each score back to its row by the document content. If the model leaves out a document, `_rerank` raises a `ValueError` that gives the number of documents with no score. ## Tests `python/python/tests/test_openai_reranker.py` replaces the chat client with a stub that returns a fixed reply. The tests do not need `openai` or an API key. - One test returns the documents sorted by score. It checks that each `id` keeps its own `text` and score. On `main`, the row with id 1 gets the text `cats purr`, which belongs to id 2. - One test returns 2 of 3 documents. It checks for the `ValueError` with the count of missing documents. On `main`, the call raises `ArrowInvalid: Added column's length must match table's length`. - Both tests fail on `main` and pass with the fix. - `test_rerankers.py`: 16 passed, 13 skipped, on the lancedb 0.40.0 wheel with and without the fix. The skipped tests need real API keys. - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed file copied over it. In that wheel, `rerankers/openai.py` and `rerankers/base.py` are byte-identical to `main`. The Node SDK has no OpenAI reranker, so there is no parity change. This change does not conflict with #4515, which edits `_client` in the same file. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/rerankers/openai.py | 21 ++++--- python/python/tests/test_openai_reranker.py | 66 +++++++++++++++++++++ 2 files changed, 78 insertions(+), 9 deletions(-) create mode 100644 python/python/tests/test_openai_reranker.py diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index 4cbbb0822..da0394d1c 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -70,15 +70,18 @@ class OpenaiReranker(Reranker): ], ) results = json.loads(response.choices[0].message.content)["documents"] - docs, scores = list( - zip(*[(result["content"], result["relevance_score"]) for result in results]) - ) # tuples - # replace the self.column column with the docs - result_set = result_set.drop(self.column) - result_set = result_set.append_column( - self.column, pa.array(docs, type=pa.string()) - ) - # add the scores + # The model returns the documents in its own order, so match each score + # back to its row by content. + scores_by_content = { + result["content"]: result["relevance_score"] for result in results + } + missing = sum(doc not in scores_by_content for doc in docs) + if missing: + raise ValueError( + f"The model returned no relevance_score for {missing} of " + f"{len(docs)} documents" + ) + scores = [scores_by_content[doc] for doc in docs] result_set = result_set.append_column( "_relevance_score", pa.array(scores, type=pa.float32()) ) diff --git a/python/python/tests/test_openai_reranker.py b/python/python/tests/test_openai_reranker.py new file mode 100644 index 000000000..43475103f --- /dev/null +++ b/python/python/tests/test_openai_reranker.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json +from types import SimpleNamespace + +import pyarrow as pa +import pytest + +from lancedb.rerankers import OpenaiReranker + + +def reranker_replying(documents): + """Return an OpenaiReranker whose chat client replies with these documents.""" + reply = json.dumps({"documents": documents}) + + def create(**kwargs): + message = SimpleNamespace(content=reply) + return SimpleNamespace(choices=[SimpleNamespace(message=message)]) + + reranker = OpenaiReranker(api_key="test-key") + reranker.__dict__["_client"] = SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ) + return reranker + + +def vector_results(): + return pa.table( + { + "id": [1, 2, 3], + "text": ["dogs bark", "cats purr", "birds sing"], + "_distance": [0.1, 0.2, 0.3], + } + ) + + +def test_scores_stay_with_their_rows_when_the_model_reorders(): + reranker = reranker_replying( + [ + {"content": "cats purr", "relevance_score": 0.9}, + {"content": "birds sing", "relevance_score": 0.5}, + {"content": "dogs bark", "relevance_score": 0.1}, + ] + ) + + result = reranker.rerank_vector("which animal purrs?", vector_results()) + + assert result.select(["id", "text"]).to_pylist() == [ + {"id": 2, "text": "cats purr"}, + {"id": 3, "text": "birds sing"}, + {"id": 1, "text": "dogs bark"}, + ] + assert result["_relevance_score"].to_pylist() == pytest.approx([0.9, 0.5, 0.1]) + + +def test_a_row_the_model_leaves_out_raises(): + reranker = reranker_replying( + [ + {"content": "cats purr", "relevance_score": 0.9}, + {"content": "dogs bark", "relevance_score": 0.1}, + ] + ) + + with pytest.raises(ValueError, match="1 of 3"): + reranker.rerank_vector("which animal purrs?", vector_results()) From 0cb5bd0a5d7b96ebafa83dc0f98f98f6af70a189 Mon Sep 17 00:00:00 2001 From: Captain <42566883+po-et@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:23:36 +0800 Subject: [PATCH 34/64] fix(python): accept VoyageAI text batch containers (#4558) Direct VoyageAI `compute_source_embeddings()` calls reject `list[str]` and NumPy text batches before reaching the provider, although both belong to the declared `TEXT` input type. For example, `["first", "second"]` raises `ValueError: Input type not allowed with text model.` Normalize these containers to a list before the existing string-entry validation. The regression cases exercise the public embedding API with a stub provider: input order, list/NumPy batches, Arrow controls, empty batches, and rejection of non-text entries and ordinary two-dimensional arrays. The VoyageAI integration page now documents the accepted text containers. Verified on current main `7b3ba4c` with a locally compiled native extension (0.41.0b7, Python 3.12.12): - VoyageAI suite: 7 failed / 20 passed before the fix; 27 passed afterward. - Adjacent embedding suite: 13 passed, 4 skipped, 8 slow cases deselected. - Python CI's pinned Ruff 0.9.9 lint/format checks pass; configured Pyright reports zero errors/warnings; all applicable pre-commit hooks pass. The provider is stubbed; no live Voyage API request or model-quality evaluation was performed. AI disclosure: I used Codex for investigation, implementation, tests, review, and drafting this description. --- docs/web/integrations/embedding/voyageai.mdx | 4 +- python/python/lancedb/embeddings/voyageai.py | 2 + .../python/tests/test_voyageai_embeddings.py | 57 +++++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/docs/web/integrations/embedding/voyageai.mdx b/docs/web/integrations/embedding/voyageai.mdx index cfa795841..c0741b39b 100644 --- a/docs/web/integrations/embedding/voyageai.mdx +++ b/docs/web/integrations/embedding/voyageai.mdx @@ -45,6 +45,8 @@ Supported parameters (to be passed in `create` method) are: | `output_dimension` | `int` | `None` | Output embedding dimension. Only supported by `voyage-multimodal-3.5`. Valid options: 256, 512, 1024 (default), 2048. | +For text models, `compute_source_embeddings()` accepts a single string, a list of strings, a one-dimensional NumPy array, or a PyArrow `Array` or `ChunkedArray`. Batch inputs retain their order. + Usage Example: @@ -61,4 +63,4 @@ The `voyage-multimodal-3.5` model can embed text alongside images. You can use i {PyEmbeddingVoyageaiMultimodal} - \ No newline at end of file + diff --git a/python/python/lancedb/embeddings/voyageai.py b/python/python/lancedb/embeddings/voyageai.py index 4a721c998..95525e07b 100644 --- a/python/python/lancedb/embeddings/voyageai.py +++ b/python/python/lancedb/embeddings/voyageai.py @@ -145,6 +145,8 @@ def sanitize_text_input(inputs: TEXT) -> List[str]: """ if isinstance(inputs, str): inputs = [inputs] + elif isinstance(inputs, (list, np.ndarray)): + inputs = list(inputs) elif isinstance(inputs, pa.Array): inputs = inputs.to_pylist() elif isinstance(inputs, pa.ChunkedArray): diff --git a/python/python/tests/test_voyageai_embeddings.py b/python/python/tests/test_voyageai_embeddings.py index 040cade1f..d49978fb3 100644 --- a/python/python/tests/test_voyageai_embeddings.py +++ b/python/python/tests/test_voyageai_embeddings.py @@ -6,6 +6,8 @@ These tests verify model registration and configuration without requiring API calls. """ +import numpy as np +import pyarrow as pa import pytest from unittest.mock import MagicMock, patch @@ -122,3 +124,58 @@ class TestVoyageAIModelRegistration: assert "voyage-4" not in func.multimodal_embedding_models assert "voyage-4-lite" not in func.multimodal_embedding_models assert "voyage-4-large" not in func.multimodal_embedding_models + + +@pytest.mark.parametrize( + "inputs", + [ + ["first", "second"], + np.array(["first", "second"]), + pa.array(["first", "second"]), + pa.chunked_array([["first"], ["second"]]), + ], + ids=["list", "numpy", "arrow", "chunked-arrow"], +) +def test_text_source_embeddings_accept_batch_containers(inputs): + from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + client = MagicMock() + client.tokenize.return_value = [[1], [2]] + client.embed.return_value.embeddings = [[0.1, 0.2], [0.3, 0.4]] + VoyageAIEmbeddingFunction.client = client + func = get_registry().get("voyageai").create(name="voyage-3") + + assert func.compute_source_embeddings(inputs) == [[0.1, 0.2], [0.3, 0.4]] + client.embed.assert_called_once_with( + texts=["first", "second"], model="voyage-3", input_type="document" + ) + + +@pytest.mark.parametrize("inputs", [[], np.array([], dtype=str)]) +def test_text_source_embeddings_accept_empty_batch(inputs): + from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + client = MagicMock() + VoyageAIEmbeddingFunction.client = client + func = get_registry().get("voyageai").create(name="voyage-3") + + assert func.compute_source_embeddings(inputs) == [] + client.tokenize.assert_not_called() + client.embed.assert_not_called() + + +@pytest.mark.parametrize( + "inputs", + [["text", 42], np.array(["text", None]), np.array([["first", "second"]])], +) +def test_text_source_embeddings_reject_non_text_batch_entries(inputs): + from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + client = MagicMock() + VoyageAIEmbeddingFunction.client = client + func = get_registry().get("voyageai").create(name="voyage-3") + + with pytest.raises(ValueError, match="Each input should be str"): + func.compute_source_embeddings(inputs) + client.tokenize.assert_not_called() + client.embed.assert_not_called() From 118e606e8e06b6aa3ad57839184f655d996ae831 Mon Sep 17 00:00:00 2001 From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:53:50 +0530 Subject: [PATCH 35/64] fix(python): align vector parent indices after flattening null lists (#4572) Arrow's `flatten()` skips child values hidden under null list parents, but `list_parent_indices()` can retain their indices. This makes `has_nan_values()` raise `ArrowInvalid` because the parent indices and NaN mask have different lengths. Filter out indices belonging to null parents before applying the mask. The regression test covers `ListArray` and `LargeListArray`, float16/32/64, and chunked inputs: 12 cases. NaN and null-component policies stay the same. The crash also affects #4503. This patch fixes it separately from that PR's changes to bad-vector policies. Validation: - All 12 regression cases failed before the fix. The utility suite now passes with 87 passed and 1 skipped. - Ruff lint and formatting pass for both changed files, as does `git diff --check`. - The patch applies cleanly to main at `ca275b6d`. The patched helper from that revision also passes all 12 cases. - A separate AI reviewer checked commit `343082a6` and found no blocking issues. They reran the utility suite and compared 4,059 generated cases plus 12 explicit null-parent cases with a row-by-row reference implementation. Those checks include slices, chunked inputs, and fixed-size lists. SDK tests used base `07e626ec` with the existing compiled extension. The broader Python run stopped after 474 passed and 23 skipped at `test_local_module_outside_code_warns_at_registration` with `ModuleNotFoundError`. That test also fails with the original vector helper restored. Repository-wide Ruff reports 49 lint errors and 11 files needing formatting, all in untouched documentation examples. Full SDK integration tests against current main require a newer native build and have not been run. --- python/python/lancedb/table.py | 5 +++++ python/python/tests/test_util.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index c7de4161d..5eea55760 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5333,6 +5333,11 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray else: values_has_nan = pc.is_nan(values) values_indices = pc.list_parent_indices(arr) + if arr.null_count: + # Match flatten(), which excludes child spans under null parent lists. + values_indices = pc.filter( + values_indices, pc.take(pc.is_valid(arr), values_indices) + ) has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan)) indices = pa.array(range(len(arr)), type=pa.uint32()) return pc.is_in(indices, has_nan_indices) diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 4671109b8..4fb16bb31 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -21,6 +21,7 @@ from lancedb.table import ( _infer_target_schema, _merge_metadata, _sanitize_data, + has_nan_values, sanitize_create_table, ) import pyarrow as pa @@ -489,6 +490,20 @@ def test_handle_bad_vectors_nan(on_bad_vectors): assert output["vector"].combine_chunks() == expected +@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()]) +@pytest.mark.parametrize("large", [False, True]) +@pytest.mark.parametrize("chunked", [False, True]) +def test_has_nan_values_skips_hidden_children_of_null_lists(dtype, large, chunked): + array_type = pa.LargeListArray if large else pa.ListArray + values = pa.array([float("nan"), None, float("nan"), 3.0, 4.0, 5.0], type=dtype) + arr = array_type.from_arrays( + [0, 2, 4, 6], values, mask=pa.array([True, False, False]) + ) + if chunked: + arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2, 1)]) + assert has_nan_values(arr).to_pylist() == [False, True, False] + + def test_handle_bad_vectors_noop(): # ChunkedArray should be preserved as-is vector = pa.chunked_array( From f2c6aa35c236e04b012c150ce94dcc6fc09fd2f7 Mon Sep 17 00:00:00 2001 From: Captain <42566883+po-et@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:24:04 +0800 Subject: [PATCH 36/64] fix(python): normalize VoyageAI image modes before JPEG encoding (#4576) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PIL images, PNG bytes and local image paths in VoyageAI source embeddings are saved directly as JPEG. RGBA, grayscale-with-alpha (LA) and palette (P) images therefore raise `OSError: cannot write mode ... as JPEG` before the provider is called. Convert to RGB at all three JPEG encoding sites. The caller’s original PIL image remains unchanged. Document that alpha is discarded and callers should composite onto a chosen background when it matters. Validation: - New image regressions + existing VoyageAI tests: **18 failed / 26 passed before → 44 passed after**. Tests cover both multimodal models, PIL/bytes/Path, RGBA/LA/P/RGB, decoded JPEG pixels and invalid images. - Independent local LanceDB table ingestion smoke: three PNG modes stored with 1024-dimensional vectors, original bytes preserved. Uses the freshly built current native extension. - Adjacent local embedding tests: 15 passed, 1 skipped (no OpenAI key), 9 deselected. The separate Flickr download test failed with `RemoteDisconnected`; it is excluded from the local result. - Pyright: 0 errors; CI-pinned Ruff 0.9.9 lint and Python formatting passed. External VoyageAI calls are stubbed; real Pillow conversion and LanceDB ingestion execute locally. No live provider or full-repository test result is claimed. AI assistance was used for implementation and validation. --- docs/web/integrations/embedding/voyageai.mdx | 7 +- python/python/lancedb/embeddings/voyageai.py | 6 +- .../tests/test_voyageai_image_inputs.py | 88 +++++++++++++++++++ 3 files changed, 97 insertions(+), 4 deletions(-) create mode 100644 python/python/tests/test_voyageai_image_inputs.py diff --git a/docs/web/integrations/embedding/voyageai.mdx b/docs/web/integrations/embedding/voyageai.mdx index c0741b39b..957ab2b4d 100644 --- a/docs/web/integrations/embedding/voyageai.mdx +++ b/docs/web/integrations/embedding/voyageai.mdx @@ -57,7 +57,12 @@ Usage Example: ### Multimodal Example {#multimodal-example} -The `voyage-multimodal-3.5` model can embed text alongside images. You can use image URLs, file paths, or PIL Image objects: +The `voyage-multimodal-3.5` model can embed text alongside images. You can use image URLs, local `pathlib.Path` objects, raw image bytes, or PIL Image objects. + +For source embeddings, PIL images and images loaded from bytes or local paths are +converted to RGB and encoded as JPEG. This supports RGBA, grayscale with alpha, +and palette images; transparency is discarded. Composite transparent images onto +your chosen background before embedding if the background color matters. diff --git a/python/python/lancedb/embeddings/voyageai.py b/python/python/lancedb/embeddings/voyageai.py index 95525e07b..bcb2fd5a3 100644 --- a/python/python/lancedb/embeddings/voyageai.py +++ b/python/python/lancedb/embeddings/voyageai.py @@ -75,7 +75,7 @@ def transform_input(input_data: Union[str, bytes, Path]): content = {"type": "text", "text": input_data} elif isinstance(input_data, PIL_Image.Image): buffered = BytesIO() - input_data.save(buffered, format="JPEG") + input_data.convert("RGB").save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") content = { "type": "image_base64", @@ -84,7 +84,7 @@ def transform_input(input_data: Union[str, bytes, Path]): elif isinstance(input_data, bytes): img = PIL_Image.open(BytesIO(input_data)) buffered = BytesIO() - img.save(buffered, format="JPEG") + img.convert("RGB").save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") content = { "type": "image_base64", @@ -103,7 +103,7 @@ def transform_input(input_data: Union[str, bytes, Path]): else: img = PIL_Image.open(input_data) buffered = BytesIO() - img.save(buffered, format="JPEG") + img.convert("RGB").save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") content = { "type": "image_base64", diff --git a/python/python/tests/test_voyageai_image_inputs.py b/python/python/tests/test_voyageai_image_inputs.py new file mode 100644 index 000000000..885f94821 --- /dev/null +++ b/python/python/tests/test_voyageai_image_inputs.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import base64 +from io import BytesIO +from unittest.mock import MagicMock + +import pytest + +from lancedb.embeddings import get_registry +from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + +@pytest.mark.parametrize("model_name", ["voyage-multimodal-3", "voyage-multimodal-3.5"]) +@pytest.mark.parametrize("input_kind", ["pil", "bytes", "path"]) +@pytest.mark.parametrize( + "mode,color,expected_rgb", + [ + ("RGBA", (64, 128, 192, 128), (64, 128, 192)), + ("LA", (96, 128), (96, 96, 96)), + ("P", 0, (64, 128, 192)), + ("RGB", (64, 128, 192), (64, 128, 192)), + ], +) +def test_multimodal_source_embeddings_encode_image_modes( + model_name, input_kind, mode, color, expected_rgb, tmp_path, monkeypatch +): + Image = pytest.importorskip("PIL.Image") + image = Image.new(mode, (8, 6), color) + if mode == "P": + image.putpalette([64, 128, 192] + [0, 0, 0] * 255) + image.info["transparency"] = 0 + original_pixels = image.tobytes() + image_input = image + if input_kind != "pil": + png = BytesIO() + image.save(png, format="PNG") + if input_kind == "bytes": + image_input = png.getvalue() + else: + image_input = tmp_path / "image.png" + image_input.write_bytes(png.getvalue()) + + client = MagicMock() + client.multimodal_embed.return_value.embeddings = [[0.1, 0.2]] + monkeypatch.setattr(VoyageAIEmbeddingFunction, "client", client) + func = get_registry().get("voyageai").create(name=model_name) + + assert func.compute_source_embeddings(image_input) == [[0.1, 0.2]] + + client.multimodal_embed.assert_called_once() + request = client.multimodal_embed.call_args.kwargs + assert request["model"] == model_name + assert request["input_type"] == "document" + assert len(request["inputs"]) == 1 + (content,) = request["inputs"][0]["content"] + assert content["type"] == "image_base64" + prefix, encoded = content["image_base64"].split(",", 1) + assert prefix == "data:image/jpeg;base64" + with Image.open(BytesIO(base64.b64decode(encoded, validate=True))) as decoded: + assert decoded.format == "JPEG" + assert decoded.mode == "RGB" + assert decoded.size == (8, 6) + assert decoded.getpixel((0, 0)) == pytest.approx(expected_rgb, abs=3) + assert image.mode == mode + assert image.tobytes() == original_pixels + client.embed.assert_not_called() + + +@pytest.mark.parametrize("input_kind", ["bytes", "path"]) +def test_multimodal_source_embeddings_reject_invalid_image( + input_kind, tmp_path, monkeypatch +): + PIL = pytest.importorskip("PIL") + image_input = b"not an image" + if input_kind == "path": + image_input = tmp_path / "invalid.png" + image_input.write_bytes(b"not an image") + + client = MagicMock() + monkeypatch.setattr(VoyageAIEmbeddingFunction, "client", client) + func = get_registry().get("voyageai").create(name="voyage-multimodal-3") + + with pytest.raises(PIL.UnidentifiedImageError): + func.compute_source_embeddings(image_input) + + client.multimodal_embed.assert_not_called() + client.embed.assert_not_called() From 566517be05ff39d3fc6445c5031abccceecf11eb Mon Sep 17 00:00:00 2001 From: Raphael Fakhri Date: Thu, 22 Oct 2026 22:24:23 +0300 Subject: [PATCH 37/64] fix(python): apply pyarrow expression filters in PyarrowDatasetAdapter (#4592) `PyarrowDatasetAdapter` lets DuckDB query a LanceDB table, but DuckDB pushes `WHERE` clauses into `Dataset.scanner(filter=...)` as `pyarrow.compute.Expression` objects. The adapter forwarded them to `Table.count_rows` and `LanceQueryBuilder.where`, which accept only SQL strings, so every filtered DuckDB query failed with `TypeError: argument 'filter': 'Expression' object is not an instance of 'str'`. This change: - Applies a `pyarrow.compute.Expression` filter locally with `pyarrow.dataset.Scanner.from_batches`. The scan reads all columns and pyarrow applies the filter and projection, because DuckDB can filter on a column it does not project. - Applies `head` limits after the filter, so the limit counts matching rows. - Routes `PyarrowDatasetAdapter.count_rows(filter)` through the scanner, so it accepts both filter types. SQL string filters keep the existing path and still run inside LanceDB. Fixes #4591 ### Testing - New `test_dataset_adapter_expression_filter` in `test_pyarrow.py` covers `count_rows`, `to_table`, `head` and a projection that excludes the filtered column. - New `test_query_with_filter` in `test_duckdb.py` runs DuckDB queries with a `WHERE` clause. - Both tests fail without the fix and pass with it. `ruff check` and `ruff format --check` pass. --- python/python/lancedb/integrations/pyarrow.py | 24 ++++++++++++++++++- python/python/tests/test_duckdb.py | 16 +++++++++++++ python/python/tests/test_pyarrow.py | 17 +++++++++++++ 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/integrations/pyarrow.py b/python/python/lancedb/integrations/pyarrow.py index 085964690..1626a6b57 100644 --- a/python/python/lancedb/integrations/pyarrow.py +++ b/python/python/lancedb/integrations/pyarrow.py @@ -52,6 +52,9 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): raise NotImplementedError("memory_pool not supported") def count_rows(self): + if isinstance(self.filter, pa.compute.Expression): + # stream the batches so memory does not grow with the matching rows + return sum(batch.num_rows for batch in self.to_reader()) return self.table.count_rows(self.filter) def from_batches(self, **kwargs): @@ -83,6 +86,8 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): return self.to_reader().read_all() def to_reader(self, *, limit: Optional[int] = None): + if isinstance(self.filter, pa.compute.Expression): + return self._to_reader_with_expression(limit=limit) query = self.table.search() # Disable the builtin limit if limit is None: @@ -98,13 +103,30 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): query = query.where(self.filter, prefilter=True) return query.to_batches(batch_size=self.batch_size) + def _to_reader_with_expression(self, *, limit: Optional[int] = None): + # LanceDB filters are SQL strings, but engines such as DuckDB push down + # pyarrow compute expressions. Scan the table and apply the filter and + # projection with pyarrow, since the filter may reference columns that + # are not part of the projection. + if limit is not None and limit <= 0: + raise ValueError("limit must be positive") + query = self.table.search().limit(self.table.count_rows()) + scanner = pa.dataset.Scanner.from_batches( + query.to_batches(batch_size=self.batch_size), + columns=self.columns, + filter=self.filter, + ) + if limit is None: + return scanner.to_reader() + return scanner.head(limit).to_reader() + class PyarrowDatasetAdapter(pa.dataset.Dataset): def __init__(self, table: Table): self.table = table def count_rows(self, filter: Optional[Filter] = None): - return self.table.count_rows(filter) + return self.scanner(filter=filter).count_rows() def get_fragments(self, filter: Optional[Filter] = None): raise NotImplementedError diff --git a/python/python/tests/test_duckdb.py b/python/python/tests/test_duckdb.py index a307149ea..a24933210 100644 --- a/python/python/tests/test_duckdb.py +++ b/python/python/tests/test_duckdb.py @@ -22,3 +22,19 @@ def test_basic_query(tmp_path): results = duck_conn.sql("SELECT SUM(y) FROM adapter").fetchall() assert results[0][0] == 26 + + +def test_query_with_filter(tmp_path): + data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}) + conn = lancedb.connect(tmp_path) + tbl = conn.create_table("test", data) + + adapter = PyarrowDatasetAdapter(tbl) # noqa: F841 + + duck_conn = duckdb.connect() + + results = duck_conn.sql("SELECT SUM(y) FROM adapter WHERE x > 2").fetchall() + assert results[0][0] == 15 + + results = duck_conn.sql("SELECT y FROM adapter WHERE x > 2 ORDER BY y").fetchall() + assert results == [(7,), (8,)] diff --git a/python/python/tests/test_pyarrow.py b/python/python/tests/test_pyarrow.py index 7d99d19c6..272eee2ea 100644 --- a/python/python/tests/test_pyarrow.py +++ b/python/python/tests/test_pyarrow.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import pyarrow as pa +import pyarrow.compute as pc import lancedb from lancedb.integrations.pyarrow import PyarrowDatasetAdapter @@ -48,3 +49,19 @@ def test_dataset_adapter(tmp_path): assert adapter.head(10).num_rows == 0 assert adapter.scanner().projected_schema == pa.schema({"x": pa.int64()}) + + +def test_dataset_adapter_expression_filter(tmp_path): + data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}) + conn = lancedb.connect(tmp_path) + tbl = conn.create_table("test", data) + + adapter = PyarrowDatasetAdapter(tbl) + filt = pc.field("x") > 2 + + assert adapter.count_rows(filt) == 2 + assert adapter.to_table(filter=filt) == data.slice(2, 2) + assert adapter.head(1, filter=filt) == data.slice(2, 1) + # The filter can reference a column that is not projected + assert adapter.to_table(columns=["y"], filter=filt) == pa.table({"y": [7, 8]}) + assert adapter.scanner(filter=filt).count_rows() == 2 From a6011dcd9501a2b50a3e9d753bd0b343140d5662 Mon Sep 17 00:00:00 2001 From: Captain <42566883+po-et@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:24:40 +0800 Subject: [PATCH 38/64] fix(python): accept NumPy text arrays in Ollama embeddings (#4605) `OllamaEmbeddings.compute_source_embeddings(np.array(["hello", "world"]))` raises a validation error before sending a request. The embedding interface accepts NumPy arrays, but the adapter passes the array directly to Ollama's string-or-sequence request field. Convert NumPy text arrays to lists before the batched request. This preserves input order, the single request, model options and `keep_alive`. The tests cover Unicode/object/sliced arrays, empty input, existing list/Arrow/query inputs and API error propagation. The added input-contract tests also cover two-dimensional arrays and byte strings, with equivalent list inputs as controls. The real Ollama SDK rejects 2-D inputs and invalid UTF-8 locally, before transport. Valid UTF-8 bytes are decoded normally. No extra adapter validation is needed for these cases. Validation on Python 3.12.14 / macOS arm64: - Built the Rust/PyO3 bindings using the repository's uv/maturin environment. - The initial regression run had 5 failures and 4 passing controls on the base; all 9 passed with the fix. With the 6 additional input-contract cases, the related embedding and trust suites now total 55 passed / 11 skipped. - All Python Ruff lint and formatting checks pass. Root checks still report the same 49 lint errors and 11 documentation formatting files reproduced on the untouched base. - Separate review exercised 13 public API cases. Another 7 groups passed through a real loopback HTTP transport, including local validation, UTF-8 serialization and empty-batch model-error propagation. The regression tests use the real Ollama 0.6.2 SDK with HTTPX MockTransport. The loopback check uses an HTTP surrogate, not an Ollama daemon. No live model inference was run. A live Ollama server was unavailable locally. [The current server handler](https://github.com/ollama/ollama/blob/2cd63f7a15039ccd790e97b31edc091fa1707c56/server/routes.go#L1212) explicitly returns HTTP 200 with empty embeddings for empty input after model lookup and runner setup. The adapter preserves the existing empty-list request and its model/setup errors. This is a source check, not a live-server test. Developed with OpenAI Codex, with separate agent review and executed local checks. --- docs/web/integrations/embedding/ollama.mdx | 3 + python/python/lancedb/embeddings/ollama.py | 2 + python/python/tests/test_ollama_embeddings.py | 126 ++++++++++++++++++ 3 files changed, 131 insertions(+) create mode 100644 python/python/tests/test_ollama_embeddings.py diff --git a/docs/web/integrations/embedding/ollama.mdx b/docs/web/integrations/embedding/ollama.mdx index 120cde8f2..bd161a692 100644 --- a/docs/web/integrations/embedding/ollama.mdx +++ b/docs/web/integrations/embedding/ollama.mdx @@ -10,6 +10,9 @@ Generate embeddings via the [ollama](https://github.com/ollama/ollama-python) py - [Ollama docs on embeddings](https://github.com/ollama/ollama/blob/main/docs/api.md#generate-embeddings) - [Ollama blog on embeddings](https://ollama.com/blog/embedding-models) +For direct calls, `compute_source_embeddings` accepts a list of strings or a +one-dimensional NumPy string array. Both inputs are sent to Ollama in one batch. + | Parameter | Type | Default Value | Description | |------------------------|----------------------------|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| | `name` | `str` | `nomic-embed-text` | The name of the model. | diff --git a/python/python/lancedb/embeddings/ollama.py b/python/python/lancedb/embeddings/ollama.py index aad8742de..3151c7c72 100644 --- a/python/python/lancedb/embeddings/ollama.py +++ b/python/python/lancedb/embeddings/ollama.py @@ -70,6 +70,8 @@ class OllamaEmbeddings(TextEmbeddingFunction): The texts to embed """ # TODO retry, rate limit, token limit + if isinstance(texts, np.ndarray): + texts = texts.tolist() embeddings = self._compute_embedding(texts) return list(embeddings) diff --git a/python/python/tests/test_ollama_embeddings.py b/python/python/tests/test_ollama_embeddings.py new file mode 100644 index 000000000..9d8568ca7 --- /dev/null +++ b/python/python/tests/test_ollama_embeddings.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json + +import httpx +import numpy as np +import pyarrow as pa +import pytest +from pydantic import ValidationError + +from lancedb.embeddings import get_registry + + +@pytest.fixture +def ollama_embeddings(): + pytest.importorskip("ollama") + requests = [] + + def respond(request): + body = json.loads(request.content) + requests.append(body) + if body["model"] == "missing-model": + return httpx.Response(404, json={"error": "model not found"}) + return httpx.Response( + 200, + json={ + "embeddings": [ + [float(index), float(len(text))] + for index, text in enumerate(body["input"]) + ] + }, + ) + + function = ( + get_registry() + .get("ollama") + .create( + name="nomic-embed-text", + options={"num_ctx": 4096}, + keep_alive="10m", + ollama_client_kwargs={ + "transport": httpx.MockTransport(respond), + "trust_env": False, + }, + ) + ) + yield function, requests + function._ollama_client._client.close() + + +@pytest.mark.parametrize( + "texts", + [ + pytest.param(["hello", "世界"], id="list"), + pytest.param(np.array(["hello", "世界"]), id="numpy-unicode"), + pytest.param(np.array(["hello", "世界"], dtype=object), id="numpy-object"), + pytest.param( + np.array([b"hello", "世界".encode("utf-8")]), id="numpy-utf8-bytes" + ), + pytest.param([b"hello", "世界".encode("utf-8")], id="list-utf8-bytes"), + pytest.param(np.array(["hello", "unused", "世界"])[::2], id="numpy-slice"), + pytest.param(pa.array(["hello", "世界"]), id="arrow"), + pytest.param(pa.chunked_array([["hello"], ["世界"]]), id="chunked-arrow"), + ], +) +def test_ollama_source_embeddings_inputs(ollama_embeddings, texts): + function, requests = ollama_embeddings + + assert function.compute_source_embeddings(texts) == [[0.0, 5.0], [1.0, 2.0]] + assert requests == [ + { + "model": "nomic-embed-text", + "input": ["hello", "世界"], + "options": {"num_ctx": 4096}, + "keep_alive": "10m", + } + ] + + +@pytest.mark.parametrize( + "texts", + [ + pytest.param(np.array([["hello"], ["world"]]), id="numpy-two-dimensional"), + pytest.param([["hello"], ["world"]], id="list-two-dimensional"), + pytest.param(np.array([b"\xff"]), id="numpy-invalid-utf8"), + pytest.param([b"\xff"], id="list-invalid-utf8"), + ], +) +def test_ollama_source_embeddings_rejects_invalid_input_locally( + ollama_embeddings, texts +): + function, requests = ollama_embeddings + + with pytest.raises(ValidationError, match="valid string"): + function.compute_source_embeddings(texts) + + assert requests == [] + + +def test_ollama_generate_embeddings_empty_array(ollama_embeddings): + function, requests = ollama_embeddings + + assert function.generate_embeddings(np.array([], dtype=str)) == [] + assert len(requests) == 1 + assert requests[0]["input"] == [] + + +def test_ollama_query_embeddings(ollama_embeddings): + function, requests = ollama_embeddings + + assert function.compute_query_embeddings("hello") == [[0.0, 5.0]] + assert len(requests) == 1 + assert requests[0]["input"] == ["hello"] + + +def test_ollama_source_embeddings_preserves_api_error(ollama_embeddings): + ollama = pytest.importorskip("ollama") + function, requests = ollama_embeddings + function.name = "missing-model" + + with pytest.raises(ollama.ResponseError, match="model not found") as error: + function.compute_source_embeddings(np.array(["hello"])) + + assert error.value.status_code == 404 + assert len(requests) == 1 From e5873d976350d74d7d276794abb8bde1c45fcb47 Mon Sep 17 00:00:00 2001 From: Ishita Singh Date: Fri, 23 Oct 2026 00:54:56 +0530 Subject: [PATCH 39/64] fix(node): support nested blobs with blobMode 'bytes' (#4610) Closes #4609. `blobMode: "bytes"` only resolved top-level blob columns, so a blob nested in a struct came back as its descriptor. Now struct projections that contain blob columns are resolved too: the bytes are fetched per blob column and the struct is rebuilt with `LargeBinary` leaves, preserving struct validity so a null struct stays null. Sibling fields pass through untouched, aliases work for both `info.image AS image` and `info AS details`, and blobs nested in lists still keep their descriptors. Tested with `npx tsc`, `biome check`, and the full `table.test.ts` suite (580 tests, including new regressions for select-all, struct selection, aliases, null structs, null and empty blobs, sibling fields, empty results, and single-row batches). Not run against remote tables. --- .../js/interfaces/QueryExecutionOptions.md | 11 +- nodejs/__test__/table.test.ts | 125 ++++++++++++++ nodejs/lancedb/arrow.ts | 1 + nodejs/lancedb/query.ts | 154 ++++++++++++++---- 4 files changed, 258 insertions(+), 33 deletions(-) diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index 8d3749f38..c0b33dcf0 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -20,12 +20,13 @@ How to return blob v2 columns. Applies to `toArrow()` and `toArray()`. - `"descriptions"` (default): each blob is a descriptor. Read the bytes with [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles). -- `"bytes"`: each top-level blob column holds the blob bytes as - `LargeBinary`. The query also reads `_rowid` to fetch the bytes, and - leaves it out of the result unless [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was - called. All the bytes are held in memory. +- `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, + including blobs nested in structs. The query also reads `_rowid` to + fetch the bytes, and leaves it out of the result unless + [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was called. All the bytes are held in + memory. -Blobs nested in a struct or a list keep their descriptors. +Blobs nested in a list keep their descriptors. The rows and the bytes come from the same table version. If the table changes while the query runs, the query runs again, up to three times. diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 84340c72c..92e3df456 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3385,6 +3385,106 @@ describe("when dealing with blob columns", () => { expect(bytes[0]!.equals(payload)).toBe(true); }); + it("returns bytes for blobs nested in structs", async () => { + const { table, alpha, empty } = await openNestedBlobTable(); + const rows = await table.query().toArray({ blobMode: "bytes" }); + const byId = new Map(rows.map((row) => [Number(row.id), row])); + expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha); + expect(byId.get(1)!.info.caption).toBe("one"); + expect(byId.get(2)!.info.image).toBeNull(); + expect(byId.get(2)!.info.caption).toBe("two"); + expect(byId.get(3)!.info).toBeNull(); + expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty); + }); + + it("lines up nested rows and bytes across single-row batches", async () => { + const { table, alpha, empty } = await openNestedBlobTable(); + const result = await table + .query() + .toArrow({ blobMode: "bytes", maxBatchLength: 1 }); + expect(result.batches.length).toBe(4); + const byId = new Map(result.toArray().map((row) => [Number(row.id), row])); + expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha); + expect(byId.get(1)!.info.caption).toBe("one"); + expect(byId.get(2)!.info.image).toBeNull(); + expect(byId.get(2)!.info.caption).toBe("two"); + expect(byId.get(3)!.info).toBeNull(); + expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty); + }); + + it("handles struct fields named like Object prototypes", async () => { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + new Field( + "info", + new Struct([ + new Field("constructor", new Struct([blob("image")]), true), + ]), + true, + ), + ]); + const payload = Buffer.from("proto"); + const table = await db.createTable( + "prototype_named_blobs", + [{ id: 1n, info: { constructor: { image: payload } } }], + { schema }, + ); + expect(await table.blobColumns()).toEqual(["info.constructor.image"]); + const [row] = await table.query().toArray({ blobMode: "bytes" }); + // Arrow's StructRow shadows the `constructor` property, so read the + // field through the JSON form instead of property access. + const info = JSON.parse(JSON.stringify(row.info)); + expect(Buffer.from(Object.values(info.constructor.image))).toEqual( + payload, + ); + }); + + it("returns nested blob bytes for a struct selection", async () => { + const { table, alpha } = await openNestedBlobTable(); + const [row] = await table + .query() + .where("id = 1") + .select(["info"]) + .toArray({ blobMode: "bytes" }); + expect(Object.keys(row)).toEqual(["info"]); + expect(Buffer.from(row.info.image)).toEqual(alpha); + expect(row.info.caption).toBe("one"); + }); + + it("returns nested blob bytes through aliases", async () => { + const { table, alpha } = await openNestedBlobTable(); + const [row] = await table + .query() + .where("id = 1") + .select( + new Map([ + ["picture", "info.image"], + ["details", "info"], + ]), + ) + .toArray({ blobMode: "bytes" }); + expect(Buffer.from(row.picture)).toEqual(alpha); + expect(Buffer.from(row.details.image)).toEqual(alpha); + expect(row.details.caption).toBe("one"); + }); + + it("returns a nested bytes schema for an empty result", async () => { + const { table } = await openNestedBlobTable(); + const result = await table + .query() + .where("id > 100") + .toArrow({ blobMode: "bytes" }); + expect(result.numRows).toBe(0); + const field = result.schema.fields.find((f) => f.name === "info")!; + expect(arrow.DataType.isStruct(field.type)).toBe(true); + const image = (field.type as Struct).children.find( + (f) => f.name === "image", + )!; + expect(arrow.DataType.isLargeBinary(image.type)).toBe(true); + expect(isBlobField(image)).toBe(false); + }); + it("creates list blobs without advertising unfetchable elements", async () => { const db = await connect(tmpDir.name); const schema = new Schema([ @@ -3593,6 +3693,31 @@ describe("when dealing with blob columns", () => { return { table, rowIds, alpha, beta }; } + async function openNestedBlobTable() { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + new Field( + "info", + new Struct([blob("image"), new Field("caption", new Utf8(), true)]), + true, + ), + ]); + const alpha = Buffer.from("alpha"); + const empty = Buffer.alloc(0); + const table = await db.createTable( + "nested_blobs", + [{ id: 1n, info: { image: alpha, caption: "one" } }], + { schema }, + ); + await table.add([ + { id: 2n, info: { image: null, caption: "two" } }, + { id: 3n, info: null }, + { id: 4n, info: { image: empty, caption: "four" } }, + ]); + return { table, alpha, empty }; + } + async function openExternalBlobTable() { const db = await connect(path.join(tmpDir.name, "db")); const payload = Buffer.from("external-payload"); diff --git a/nodejs/lancedb/arrow.ts b/nodejs/lancedb/arrow.ts index 68b35e256..56c8e1897 100644 --- a/nodejs/lancedb/arrow.ts +++ b/nodejs/lancedb/arrow.ts @@ -61,6 +61,7 @@ import { import { inferSchema } from "./schema"; export * from "apache-arrow"; +export type { Data as ArrowData } from "apache-arrow"; export type SchemaLike = | Schema | { diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 2abe6db97..823239ef9 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors import { + type ArrowData, Table as ArrowTable, + DataType, Field, type IntoVector, LargeBinary, @@ -70,16 +72,80 @@ function blobSources( blobColumns: string[], projection: [string, string][] | undefined, ): Map { - const columns = new Set(blobColumns); + const holds = (source: string) => + blobColumns.some( + (column) => column === source || column.startsWith(`${source}.`), + ); const pairs = projection ?? - blobColumns.map((column): [string, string] => [column, column]); - return new Map(pairs.filter(([, source]) => columns.has(source))); + Array.from(new Set(blobColumns.map((column) => column.split(".")[0]))).map( + (name): [string, string] => [name, name], + ); + return new Map(pairs.filter(([, source]) => holds(source))); +} + +// Fetched bytes for the blobs under one output column. A blob column is a +// leaf; a struct holds the blobs nested in it, keyed by field name. A Map +// keeps struct field names from reaching shared JavaScript prototypes. +type BlobBytes = (Buffer | null)[] | Map; + +function isLeaf(bytes: BlobBytes): bytes is (Buffer | null)[] { + return Array.isArray(bytes); +} + +function nestedFieldType(type: Struct, bytes: Map): Struct { + return new Struct( + type.children.map((field) => { + const inner = bytes.get(field.name); + if (inner === undefined) { + return field; + } + if (isLeaf(inner)) { + return new Field(field.name, new LargeBinary(), true); + } + return new Field( + field.name, + nestedFieldType(field.type as Struct, inner), + field.nullable, + ); + }), + ); +} + +// Rebuilds a struct with the blobs under it as bytes. The struct keeps its own +// validity, so a null struct stays null. +function nestedData( + data: ArrowData, + bytes: Map, + offset: number, +): ArrowData { + const type = data.type as Struct; + const children = data.children.map((child, index) => { + const inner = bytes.get(type.children[index].name); + if (inner === undefined) { + return child; + } + if (isLeaf(inner)) { + return vectorFromArray( + inner.slice(offset, offset + data.length), + new LargeBinary(), + ).data[0]; + } + return nestedData(child, inner, offset); + }); + return makeData({ + type: nestedFieldType(type, bytes), + length: data.length, + nullCount: data.nullCount, + nullBitmap: data.nullBitmap, + children, + }); } async function readBlobBytes( table: ArrowTable, sources: Map, + blobColumns: string[], snapshot: NativeTable, dropRowId: boolean, ): Promise { @@ -91,7 +157,11 @@ async function readBlobBytes( rowIdColumn === null ? [] : Array.from(rowIdColumn.toArray() as BigUint64Array); - const replacements = new Map(); + const fetch = async (column: string) => + (rowIds.length === 0 ? [] : await snapshot.fetchBlobs(column, rowIds)).map( + (value) => value ?? null, + ); + const replacements = new Map(); for (const [output, source] of sources) { const index = table.schema.fields.findIndex( (field) => field.name === output, @@ -99,12 +169,30 @@ async function readBlobBytes( if (index < 0) { continue; } - const values = - rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds); - replacements.set( - index, - values.map((value) => value ?? null), - ); + if (blobColumns.includes(source)) { + replacements.set(index, await fetch(source)); + continue; + } + if (!DataType.isStruct(table.schema.fields[index].type)) { + continue; + } + const nested = new Map(); + for (const column of blobColumns.filter((c) => + c.startsWith(`${source}.`), + )) { + const path = column.slice(source.length + 1).split("."); + let level = nested; + for (const name of path.slice(0, -1)) { + let next = level.get(name); + if (next === undefined) { + next = new Map(); + level.set(name, next); + } + level = next as Map; + } + level.set(path[path.length - 1], await fetch(column)); + } + replacements.set(index, nested); } const result = withBinaryColumns(table, replacements); if (!dropRowId) { @@ -121,26 +209,37 @@ async function readBlobBytes( // the bytes. function withBinaryColumns( table: ArrowTable, - replacements: Map, + replacements: Map, ): ArrowTable { if (replacements.size === 0) { return table; } - const fields = table.schema.fields.map((field, index) => - replacements.has(index) + const fields = table.schema.fields.map((field, index) => { + const bytes = replacements.get(index); + if (bytes === undefined) { + return field; + } + return isLeaf(bytes) ? new Field(field.name, new LargeBinary(), true) - : field, - ); + : new Field( + field.name, + nestedFieldType(field.type as Struct, bytes), + field.nullable, + ); + }); const schema = new Schema(fields, table.schema.metadata); let offset = 0; const batches = table.batches.map((batch) => { const children = batch.data.children.map((child, index) => { - const values = replacements.get(index); - if (values === undefined) { + const bytes = replacements.get(index); + if (bytes === undefined) { return child; } + if (!isLeaf(bytes)) { + return nestedData(child, bytes, offset); + } return vectorFromArray( - values.slice(offset, offset + batch.numRows), + bytes.slice(offset, offset + batch.numRows), new LargeBinary(), ).data[0]; }); @@ -175,12 +274,13 @@ export interface QueryExecutionOptions { * * - `"descriptions"` (default): each blob is a descriptor. Read the bytes * with {@link Table.fetchBlobs} or {@link Table.fetchBlobFiles}. - * - `"bytes"`: each top-level blob column holds the blob bytes as - * `LargeBinary`. The query also reads `_rowid` to fetch the bytes, and - * leaves it out of the result unless {@link QueryBase.withRowId} was - * called. All the bytes are held in memory. + * - `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, + * including blobs nested in structs. The query also reads `_rowid` to + * fetch the bytes, and leaves it out of the result unless + * {@link QueryBase.withRowId} was called. All the bytes are held in + * memory. * - * Blobs nested in a struct or a list keep their descriptors. + * Blobs nested in a list keep their descriptors. * * The rows and the bytes come from the same table version. If the table * changes while the query runs, the query runs again, up to three times. @@ -429,17 +529,15 @@ export class QueryBase< // unchanged version afterwards proves it read the snapshot's version. const snapshot = await context.table.querySnapshot(); const version = await snapshot.version(); - const sources = blobSources( - await snapshot.blobColumns(), - context.projection, - ); + const blobColumns = await snapshot.blobColumns(); + const sources = blobSources(blobColumns, context.projection); if (sources.size === 0) { return this.collect(options, false); } const addRowId = !context.withRowId; const table = await this.collect(options, addRowId); if ((await context.table.version()) === version) { - return readBlobBytes(table, sources, snapshot, addRowId); + return readBlobBytes(table, sources, blobColumns, snapshot, addRowId); } if (attempt === BLOB_BYTES_ATTEMPTS) { throw new Error( From f5f4463b15bf1e6b08725da26f05cf8dd5188888 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:42:22 +0800 Subject: [PATCH 40/64] fix(node): materialize nested struct blobs in bytes mode (#4611) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `blobMode: "bytes"` left blobs inside structs as descriptors because projection matching ignored parent structs and byte replacement only rebuilt top-level fields. Expand direct struct projections and aliases to their blob leaves, then recursively replace the corresponding Arrow fields and data with `LargeBinary`. Track literal output names separately from descendant paths so scalar aliases stay intact and every overlapping blob projection is converted. Preserve null structs, null and empty blobs, sibling fields, schema metadata, and the existing row ID and snapshot behavior. Blobs inside lists retain their descriptors. Update the API docs to describe struct support. Regression coverage includes full and parent-struct projections, leaf and struct aliases, dotted output names, overlapping projections in both orders, nested structs, nulls, empty blobs and results, multiple batches, list siblings, query reuse, row IDs, and concurrent writes. Regression tests reproduce the nested-blob and projection-collision failures before the corresponding changes. Validation: - `pnpm build` - `pnpm lint` - `pnpm run docs` - `pnpm test __test__/nested_blob.test.ts __test__/blob.test.ts __test__/query.test.ts --runInBand --verbose=false` — 142 passed - `pnpm test __test__/table.test.ts --runInBand --verbose=false -t 'when dealing with blob columns'` — 60 passed Fixes #4609 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- .../js/interfaces/QueryExecutionOptions.md | 13 +- nodejs/__test__/nested_blob.test.ts | 327 ++++++++++++++++++ nodejs/__test__/table.test.ts | 125 ++++--- nodejs/lancedb/query.ts | 259 +++++++------- 4 files changed, 528 insertions(+), 196 deletions(-) create mode 100644 nodejs/__test__/nested_blob.test.ts diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index c0b33dcf0..f18ada7aa 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -20,13 +20,14 @@ How to return blob v2 columns. Applies to `toArrow()` and `toArray()`. - `"descriptions"` (default): each blob is a descriptor. Read the bytes with [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles). -- `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - including blobs nested in structs. The query also reads `_rowid` to - fetch the bytes, and leaves it out of the result unless - [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was called. All the bytes are held in - memory. +- `"bytes"`: each blob column, including blobs nested in structs, holds + the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + leaves it out of the result unless [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was + called. All the bytes are held in memory. -Blobs nested in a list keep their descriptors. +Selecting a parent struct or directly aliasing a blob or struct also returns +bytes. Blobs inside lists keep their descriptors. Arbitrary projection +expressions are not materialized as bytes. The rows and the bytes come from the same table version. If the table changes while the query runs, the query runs again, up to three times. diff --git a/nodejs/__test__/nested_blob.test.ts b/nodejs/__test__/nested_blob.test.ts new file mode 100644 index 000000000..508b57c77 --- /dev/null +++ b/nodejs/__test__/nested_blob.test.ts @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import * as tmp from "tmp"; +import { type Table, blob, connect, isBlobField } from "../lancedb"; +import { + DataType, + Field, + Int64, + List, + Schema, + Struct, + Utf8, +} from "../lancedb/arrow"; + +describe('nested blobs with blobMode "bytes"', () => { + let tmpDir: tmp.DirResult; + let table: Table; + + const caption = new Field( + "caption", + new Utf8(), + false, + new Map([["label", "caption"]]), + ); + const info = new Field( + "info", + new Struct([ + blob("image"), + caption, + new Field("preview", new Struct([blob("image")]), true), + ]), + true, + new Map([["label", "info"]]), + ); + const schema = new Schema( + [new Field("id", new Int64(), false), info, blob("cover")], + new Map([["source", "nested-blobs"]]), + ); + + beforeEach(async () => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + table = await db.createTable( + "nested_blobs", + [ + { + id: 1n, + info: { + image: Buffer.from("hello"), + caption: "first", + preview: { image: Buffer.from("preview") }, + }, + cover: Buffer.from("cover"), + }, + { id: 2n, info: null, cover: null }, + { + id: 3n, + info: { image: null, caption: "third", preview: null }, + cover: null, + }, + { + id: 4n, + info: { + image: Buffer.alloc(0), + caption: "fourth", + preview: { image: null }, + }, + cover: Buffer.alloc(0), + }, + ], + { schema }, + ); + }); + + afterEach(() => tmpDir.removeCallback()); + + it.each([false, true])( + "reads struct leaves across batches (select parent: %s)", + async (selectParent) => { + const query = table + .query() + .orderBy({ columnName: "id", ascending: true }); + if (selectParent) { + query.select(["info"]); + } + const result = await query.toArrow({ + blobMode: "bytes", + maxBatchLength: 1, + }); + expect(result.batches.length).toBeGreaterThan(1); + expect(result.schema.fields.map((field) => field.name)).toEqual( + selectParent ? ["info"] : ["id", "info", "cover"], + ); + expect(result.schema.metadata).toEqual(schema.metadata); + const infoField = result.schema.fields.find( + (field) => field.name === "info", + )!; + expect(infoField.nullable).toBe(true); + expect(infoField.metadata).toEqual(info.metadata); + const fields = (infoField.type as Struct).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[0].metadata.size).toBe(0); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + + const rows = result.toArray(); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(Buffer.from(rows[0].info.preview.image).toString()).toBe( + "preview", + ); + expect(rows[1].info).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(rows[2].info.caption).toBe("third"); + expect(rows[2].info.preview).toBeNull(); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + expect(rows[3].info.preview.image).toBeNull(); + if (!selectParent) { + expect(Buffer.from(rows[0].cover).toString()).toBe("cover"); + expect(rows[1].cover).toBeNull(); + } + }, + ); + + it("reads direct leaf and struct aliases", async () => { + const rows = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select({ + image: "info.image", + details: "info", + preview: "info.preview", + caption: "info.caption", + }) + .toArray({ blobMode: "bytes", maxBatchLength: 1 }); + expect(Object.keys(rows[0])).toEqual([ + "image", + "details", + "preview", + "caption", + ]); + expect(Buffer.from(rows[0].image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.preview.image).toString()).toBe( + "preview", + ); + expect(Buffer.from(rows[0].preview.image).toString()).toBe("preview"); + expect(rows[0].caption).toBe("first"); + expect(rows[1].details).toBeNull(); + expect(rows[1].image).toBeNull(); + expect(rows[2].image).toBeNull(); + expect(Buffer.from(rows[3].image)).toEqual(Buffer.alloc(0)); + }); + + it.each([false, true])( + "preserves a literal scalar alias matching a nested blob path (reverse: %s)", + async (reverse) => { + const projection: [string, string][] = [ + ["info", "info"], + ["info.image", "id"], + ]; + if (reverse) { + projection.reverse(); + } + const [row] = await table + .query() + .where("id = 1") + .select(new Map(projection)) + .toArray({ blobMode: "bytes" }); + expect(row["info.image"]).toBe(1n); + expect(Buffer.from(row.info.image).toString()).toBe("hello"); + }, + ); + + it.each([false, true])( + "materializes every overlapping blob projection (reverse: %s)", + async (reverse) => { + const columns = ["info.image", "info"]; + if (reverse) { + columns.reverse(); + } + const result = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select(columns) + .toArrow({ blobMode: "bytes", maxBatchLength: 1 }); + expect(result.schema.fields.map((field) => field.name)).toEqual(columns); + expect(DataType.isLargeBinary(result.getChild("info.image")!.type)).toBe( + true, + ); + expect( + DataType.isLargeBinary( + result.getChild("info")!.getChild("image")!.type, + ), + ).toBe(true); + const rows = result.toArray(); + expect(Buffer.from(rows[0]["info.image"]).toString()).toBe("hello"); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(rows[1]["info.image"]).toBeNull(); + expect(rows[1].info).toBeNull(); + expect(rows[2]["info.image"]).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(Buffer.from(rows[3]["info.image"])).toEqual(Buffer.alloc(0)); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + }, + ); + + it("materializes dotted aliases independently of nested fields", async () => { + const [row] = await table + .query() + .where("id = 1") + .select({ + details: "info", + "details.image": "info.image", + "details.preview": "info.preview", + "details.preview.image": "info.preview.image", + }) + .toArray({ blobMode: "bytes" }); + expect(Buffer.from(row.details.image).toString()).toBe("hello"); + expect(Buffer.from(row["details.image"]).toString()).toBe("hello"); + expect(Buffer.from(row.details.preview.image).toString()).toBe("preview"); + expect(Buffer.from(row["details.preview"].image).toString()).toBe( + "preview", + ); + expect(Buffer.from(row["details.preview.image"]).toString()).toBe( + "preview", + ); + }); + + it.each([false, true])( + "converts the empty result schema (alias: %s)", + async (alias) => { + const query = table.query().where("id < 0"); + if (alias) { + query.select({ image: "info.image", details: "info" }); + } + const result = await query.toArrow({ blobMode: "bytes" }); + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual( + alias ? ["image", "details"] : ["id", "info", "cover"], + ); + const fields = ( + result.schema.fields.find( + (field) => field.name === (alias ? "details" : "info"), + )!.type as Struct + ).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + if (alias) { + expect(DataType.isLargeBinary(result.schema.fields[0].type)).toBe(true); + } + }, + ); + + it("preserves requested row IDs and leaves subsequent descriptors intact", async () => { + const query = table.query().where("id = 1").select(["info"]); + const [bytes] = await query.toArray({ blobMode: "bytes" }); + expect(Object.keys(bytes)).toEqual(["info"]); + const [description] = await query.toArray(); + expect(Object.keys(description)).toEqual(["info"]); + expect(Number(description.info.image.size)).toBe(5); + const [withRowId] = await query.withRowId().toArray({ blobMode: "bytes" }); + expect(typeof withRowId._rowid).toBe("bigint"); + expect(await table.fetchBlobs("info.image", [withRowId._rowid])).toEqual([ + Buffer.from(withRowId.info.image), + ]); + }); + + it("keeps descriptors in description mode", async () => { + const [description] = await table + .query() + .where("id = 1") + .toArray({ blobMode: "descriptions" }); + expect(Number(description.info.image.size)).toBe(5); + }); + + it("keeps list elements as descriptors alongside struct blob bytes", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable( + "list_blobs", + [ + { + info: { image: Buffer.from("hello"), images: [Buffer.from("list")] }, + }, + ], + { + schema: new Schema([ + new Field( + "info", + new Struct([ + blob("image"), + new Field("images", new List(blob("item")), true), + ]), + true, + ), + ]), + }, + ); + const [bytes] = await table.query().toArray({ blobMode: "bytes" }); + expect(Buffer.from(bytes.info.image).toString()).toBe("hello"); + expect(Number(bytes.info.images.get(0).size)).toBe(4); + const [row] = await table + .query() + .select(["info.images"]) + .toArray({ blobMode: "bytes" }); + expect(Number(row["info.images"].get(0).size)).toBe(4); + expect(Object.keys(row)).toEqual(["info.images"]); + const result = await table.query().toArrow({ blobMode: "bytes" }); + const fields = (result.schema.fields[0].type as Struct).children; + const descriptions = await table.query().toArrow(); + expect(fields[1]).toEqual( + (descriptions.schema.fields[0].type as Struct).children[1], + ); + expect(DataType.isStruct((fields[1].type as List).children[0].type)).toBe( + true, + ); + }); +}); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 92e3df456..d848676cd 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3314,51 +3314,64 @@ describe("when dealing with blob columns", () => { ).rejects.toThrow(/blobMode must be/); }); - it("reruns the query when the table changes before the version check", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - await update(2n, "new"); - return liveVersion(); - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(2n); - expect(Buffer.from(row.image).toString()).toBe("new"); - }); + it.each([false, true])( + "reruns the query when the table changes before the version check (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + await update(2n, "new"); + return liveVersion(); + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(2n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "new", + ); + }, + ); - it("reads bytes from the query's version when a write lands before the fetch", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - const version = await liveVersion(); - await update(2n, "new"); - return version; - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(1n); - expect(Buffer.from(row.image).toString()).toBe("old"); - }); + it.each([false, true])( + "reads bytes from the query's version when a write lands before the fetch (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + const version = await liveVersion(); + await update(2n, "new"); + return version; + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(1n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "old", + ); + }, + ); - it("fails when the table changes during every attempt", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - let generation = 1n; - native.version = async () => { - generation += 1n; - await update(generation, `generation ${generation}`); - return liveVersion(); - }; - try { - await expect( - table.query().toArray({ blobMode: "bytes" }), - ).rejects.toThrow(/changed during each of 3 attempts/); - expect(generation).toBe(4n); - } finally { - native.version = liveVersion; - } - }); + it.each([false, true])( + "fails when the table changes during every attempt (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + let generation = 1n; + native.version = async () => { + generation += 1n; + await update(generation, `generation ${generation}`); + return liveVersion(); + }; + try { + await expect( + table.query().toArray({ blobMode: "bytes" }), + ).rejects.toThrow(/changed during each of 3 attempts/); + expect(generation).toBe(4n); + } finally { + native.version = liveVersion; + } + }, + ); it("rejects fetchBlobs on a non-blob column", async () => { const { table, rowIds } = await openBlobTable(); @@ -3642,18 +3655,28 @@ describe("when dealing with blob columns", () => { ).map((value) => (value == null ? null : Number(value.size))); } - async function openVersionedBlobTable() { + async function openVersionedBlobTable(nested = false) { const db = await connect(tmpDir.name, { storageOptions: { newTableEnableStableRowIds: "true" }, }); const schema = new Schema([ new Field("id", new Int64(), false), new Field("generation", new Int64(), true), - blob("image"), + nested + ? new Field("info", new Struct([blob("image")]), true) + : blob("image"), ]); const table = await db.createTable( "versions", - [{ id: 1n, generation: 1n, image: Buffer.from("old") }], + [ + { + id: 1n, + generation: 1n, + ...(nested + ? { info: { image: Buffer.from("old") } } + : { image: Buffer.from("old") }), + }, + ], { schema }, ); const native = ( @@ -3663,7 +3686,15 @@ describe("when dealing with blob columns", () => { await table .mergeInsert("id") .whenMatchedUpdateAll() - .execute([{ id: 1n, generation, image: Buffer.from(image) }]); + .execute([ + { + id: 1n, + generation, + ...(nested + ? { info: { image: Buffer.from(image) } } + : { image: Buffer.from(image) }), + }, + ]); }; return { table, native, update }; } diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 823239ef9..b96e4fda6 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +import type { Data as ArrowData } from "apache-arrow"; import { - type ArrowData, Table as ArrowTable, DataType, Field, @@ -14,7 +14,6 @@ import { extractVectorBuffer, fromBufferToRecordBatch, fromRecordBatchToBuffer, - makeData, tableFromIPC, vectorFromArray, } from "./arrow"; @@ -68,84 +67,42 @@ export interface QueryBlobContext { const BLOB_BYTES_ATTEMPTS = 3; +interface BlobSource { + // Output names are literal: an alias containing dots is still one field. + output: string; + path: string[]; + source: string; +} + function blobSources( blobColumns: string[], projection: [string, string][] | undefined, -): Map { - const holds = (source: string) => - blobColumns.some( - (column) => column === source || column.startsWith(`${source}.`), - ); - const pairs = - projection ?? - Array.from(new Set(blobColumns.map((column) => column.split(".")[0]))).map( - (name): [string, string] => [name, name], - ); - return new Map(pairs.filter(([, source]) => holds(source))); -} - -// Fetched bytes for the blobs under one output column. A blob column is a -// leaf; a struct holds the blobs nested in it, keyed by field name. A Map -// keeps struct field names from reaching shared JavaScript prototypes. -type BlobBytes = (Buffer | null)[] | Map; - -function isLeaf(bytes: BlobBytes): bytes is (Buffer | null)[] { - return Array.isArray(bytes); -} - -function nestedFieldType(type: Struct, bytes: Map): Struct { - return new Struct( - type.children.map((field) => { - const inner = bytes.get(field.name); - if (inner === undefined) { - return field; +): BlobSource[] { + if (projection === undefined) { + return blobColumns.map((source) => { + const [output, ...path] = source.split("."); + return { output, path, source }; + }); + } + const sources: BlobSource[] = []; + for (const [output, source] of projection) { + for (const column of blobColumns) { + if (column === source || column.startsWith(`${source}.`)) { + sources.push({ + output, + path: + column === source ? [] : column.slice(source.length + 1).split("."), + source: column, + }); } - if (isLeaf(inner)) { - return new Field(field.name, new LargeBinary(), true); - } - return new Field( - field.name, - nestedFieldType(field.type as Struct, inner), - field.nullable, - ); - }), - ); -} - -// Rebuilds a struct with the blobs under it as bytes. The struct keeps its own -// validity, so a null struct stays null. -function nestedData( - data: ArrowData, - bytes: Map, - offset: number, -): ArrowData { - const type = data.type as Struct; - const children = data.children.map((child, index) => { - const inner = bytes.get(type.children[index].name); - if (inner === undefined) { - return child; } - if (isLeaf(inner)) { - return vectorFromArray( - inner.slice(offset, offset + data.length), - new LargeBinary(), - ).data[0]; - } - return nestedData(child, inner, offset); - }); - return makeData({ - type: nestedFieldType(type, bytes), - length: data.length, - nullCount: data.nullCount, - nullBitmap: data.nullBitmap, - children, - }); + } + return sources; } async function readBlobBytes( table: ArrowTable, - sources: Map, - blobColumns: string[], + sources: BlobSource[], snapshot: NativeTable, dropRowId: boolean, ): Promise { @@ -157,42 +114,27 @@ async function readBlobBytes( rowIdColumn === null ? [] : Array.from(rowIdColumn.toArray() as BigUint64Array); - const fetch = async (column: string) => - (rowIds.length === 0 ? [] : await snapshot.fetchBlobs(column, rowIds)).map( - (value) => value ?? null, + const outputFields = new Map( + table.schema.fields.map((field) => [field.name, field]), + ); + const replacements = new Map(); + for (const { output, path, source } of sources) { + let field = outputFields.get(output); + for (const name of path) { + field = + field !== undefined && DataType.isStruct(field.type) + ? field.type.children.find((child) => child.name === name) + : undefined; + } + if (field === undefined) { + continue; + } + const values = + rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds); + replacements.set( + field, + values.map((value) => value ?? null), ); - const replacements = new Map(); - for (const [output, source] of sources) { - const index = table.schema.fields.findIndex( - (field) => field.name === output, - ); - if (index < 0) { - continue; - } - if (blobColumns.includes(source)) { - replacements.set(index, await fetch(source)); - continue; - } - if (!DataType.isStruct(table.schema.fields[index].type)) { - continue; - } - const nested = new Map(); - for (const column of blobColumns.filter((c) => - c.startsWith(`${source}.`), - )) { - const path = column.slice(source.length + 1).split("."); - let level = nested; - for (const name of path.slice(0, -1)) { - let next = level.get(name); - if (next === undefined) { - next = new Map(); - level.set(name, next); - } - level = next as Map; - } - level.set(path[path.length - 1], await fetch(column)); - } - replacements.set(index, nested); } const result = withBinaryColumns(table, replacements); if (!dropRowId) { @@ -205,48 +147,76 @@ async function readBlobBytes( ); } -// A fresh field, not a clone: the descriptor field's metadata must not follow -// the bytes. function withBinaryColumns( table: ArrowTable, - replacements: Map, + replacements: Map, ): ArrowTable { if (replacements.size === 0) { return table; } - const fields = table.schema.fields.map((field, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return field; + const replaceField = (field: Field): Field => { + if (replacements.has(field)) { + // The descriptor's extension and encoding metadata must not follow bytes. + return new Field(field.name, new LargeBinary(), true); } - return isLeaf(bytes) - ? new Field(field.name, new LargeBinary(), true) - : new Field( - field.name, - nestedFieldType(field.type as Struct, bytes), - field.nullable, - ); - }); + if (DataType.isStruct(field.type)) { + const children = field.type.children.map(replaceField); + if (children.some((child, i) => child !== field.type.children[i])) { + return field.clone({ type: new Struct(children) }); + } + } + return field; + }; + const fields = table.schema.fields.map(replaceField); const schema = new Schema(fields, table.schema.metadata); let offset = 0; const batches = table.batches.map((batch) => { - const children = batch.data.children.map((child, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return child; + const replaceData = ( + data: ArrowData, + field: Field, + updated: Field, + ): ArrowData => { + if (field === updated) { + return data; } - if (!isLeaf(bytes)) { - return nestedData(child, bytes, offset); + const values = replacements.get(field); + if (values !== undefined) { + return vectorFromArray( + values.slice(offset, offset + batch.numRows), + new LargeBinary(), + ).data[0]; } - return vectorFromArray( - bytes.slice(offset, offset + batch.numRows), - new LargeBinary(), - ).data[0]; - }); + const children = data.children.map((child, index) => + replaceData( + child, + field.type.children[index], + updated.type.children[index], + ), + ); + // Keep the struct's validity bitmap and offset, including null parents. + return data.clone( + updated.type, + undefined, + undefined, + undefined, + undefined, + children, + ); + }; + const children = batch.data.children.map((child, index) => + replaceData(child, table.schema.fields[index], fields[index]), + ); offset += batch.numRows; return new RecordBatch( schema, - makeData({ type: new Struct(fields), length: batch.numRows, children }), + batch.data.clone( + new Struct(fields), + undefined, + undefined, + undefined, + undefined, + children, + ), ); }); return new ArrowTable(schema, batches); @@ -274,13 +244,14 @@ export interface QueryExecutionOptions { * * - `"descriptions"` (default): each blob is a descriptor. Read the bytes * with {@link Table.fetchBlobs} or {@link Table.fetchBlobFiles}. - * - `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - * including blobs nested in structs. The query also reads `_rowid` to - * fetch the bytes, and leaves it out of the result unless - * {@link QueryBase.withRowId} was called. All the bytes are held in - * memory. + * - `"bytes"`: each blob column, including blobs nested in structs, holds + * the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + * leaves it out of the result unless {@link QueryBase.withRowId} was + * called. All the bytes are held in memory. * - * Blobs nested in a list keep their descriptors. + * Selecting a parent struct or directly aliasing a blob or struct also returns + * bytes. Blobs inside lists keep their descriptors. Arbitrary projection + * expressions are not materialized as bytes. * * The rows and the bytes come from the same table version. If the table * changes while the query runs, the query runs again, up to three times. @@ -529,15 +500,17 @@ export class QueryBase< // unchanged version afterwards proves it read the snapshot's version. const snapshot = await context.table.querySnapshot(); const version = await snapshot.version(); - const blobColumns = await snapshot.blobColumns(); - const sources = blobSources(blobColumns, context.projection); - if (sources.size === 0) { + const sources = blobSources( + await snapshot.blobColumns(), + context.projection, + ); + if (sources.length === 0) { return this.collect(options, false); } const addRowId = !context.withRowId; const table = await this.collect(options, addRowId); if ((await context.table.version()) === version) { - return readBlobBytes(table, sources, blobColumns, snapshot, addRowId); + return readBlobBytes(table, sources, snapshot, addRowId); } if (attempt === BLOB_BYTES_ATTEMPTS) { throw new Error( From 88203d25156a0dd7fc7d2c51aee66410d546ae18 Mon Sep 17 00:00:00 2001 From: Christopher Krah Date: Thu, 22 Oct 2026 21:42:36 +0200 Subject: [PATCH 41/64] fix(python): keep a renamed Permutation column in its position (#4626) The `Permutation.rename_column` docstring says it "only changes the name of the column that is visible from this permutation". But `permutation.py:868-870` adds the new name at the end of the selection and then deletes the old key, so the renamed column also moves to the last position. Hugging Face `datasets`, which `Permutation` mirrors (#2725), keeps it in place. Positional formats feel it most: `with_format("numpy")` swaps the feature order of every row. Before, on `d97df0c`: ``` $ python repro.py /python/python # lancedb 0.41.0b12 extension, Python sources from the tree ['id', 'value'] ['value', 'row_id'] [{'value': 0, 'row_id': 0}] [[10 1]] ``` After: ``` $ python repro.py /python/python ['id', 'value'] ['row_id', 'value'] [{'row_id': 0, 'value': 0}] [[ 1 10]] ```
repro.py (its first output line, lancedb.__file__, is cut above) ```python import sys; sys.path.insert(0, sys.argv[1]) import tempfile import lancedb from lancedb.permutation import Permutation print(lancedb.__file__) db = lancedb.connect(tempfile.mkdtemp()) tbl = db.create_table("t", [{"id": i, "value": i * 10} for i in range(3)]) perm = Permutation.identity(tbl) print(perm.column_names) renamed = perm.rename_column("id", "row_id") print(renamed.column_names) print(renamed.__getitems__([0])) print(renamed.with_format("numpy").__getitems__([1])) ```
For comparison, `datasets` 5.1.0 gives `['new_id', 'value']` for `Dataset.from_dict({'id': [1, 2], 'value': [3, 4]}).rename_column('id', 'new_id').column_names`. No issue exists for this. The fix rebuilds the selection dict in order. The cost: code that renamed a column and then read it by position in numpy batches (say `[:, -1]`) gets a different column after this, with no error. `rename_columns` calls `rename_column` once per name, so it gets the fix too. The tests that check schema order only rename `value`, the last column, and `test_rename_rowid` checks membership but not order. The new test renames `id`. ### Verification ``` # Python 3.12.6; lancedb 0.41.0b12 wheel from pypi.fury.io, its _lancedb.abi3.so copied into the tree; pylance 14.0.0-rc.1 $ PYTHONPATH=python python -m pytest python/tests/test_permutation.py -q -p no:cacheprovider d97df0c: 72 passed in 2.53s this branch: 73 passed in 2.36s permutation.py reverted, test kept: FAILED python/tests/test_permutation.py::test_rename_column_keeps_position assert value: int64\nnew_id: int64 == new_id: int64\nvalue: int64 1 failed, 72 passed in 2.41s $ ruff format --check . && ruff check . # ruff 0.9.9, as python.yml pins it 152 files already formatted All checks passed! # not run: pyright (its include list covers neither permutation.py nor tests/), the Rust build ``` @westonpace, this is your class from #2725. Do you want the docstring to say outright that the column keeps its position? If so, I'll add that line to both rename methods in this PR. Signed-off-by: Christopher Krah <48338417+chrikrah@users.noreply.github.com> --- python/python/lancedb/permutation.py | 8 +++++--- python/python/tests/test_permutation.py | 7 +++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 355ebf798..e95db27f6 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -909,9 +909,11 @@ class Permutation: f"Cannot rename column `{old_name}` to `{new_name}` because a column " "with that name already exists" ) - new_selection = self.selection.copy() - new_selection[new_name] = new_selection[old_name] - del new_selection[old_name] + # Rebuild the mapping so the renamed column keeps its position + new_selection = { + (new_name if name == old_name else name): value + for name, value in self.selection.items() + } return self._with_selection(new_selection) def rename_columns(self, column_map: dict[str, str]) -> "Permutation": diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index ed1894340..50ca53061 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -811,6 +811,13 @@ def test_rename_column(some_permutation: Permutation): some_permutation.rename_column("non_existent", "new_value") +def test_rename_column_keeps_position(some_permutation: Permutation): + renamed = some_permutation.rename_column("id", "new_id") + assert renamed.schema == pa.schema([("new_id", pa.int64()), ("value", pa.int64())]) + batch = renamed.with_format("arrow").__getitems__([0]) + assert batch.schema.names == ["new_id", "value"] + + def test_rename_columns(some_permutation: Permutation): assert some_permutation.rename_columns({"value": "new_value"}).schema == pa.schema( [("id", pa.int64()), ("new_value", pa.int64())] From 631ea41a635890218bf7ee1ab80df6caccc7606c Mon Sep 17 00:00:00 2001 From: Colin Patrick McCabe Date: Thu, 22 Oct 2026 15:41:30 -0700 Subject: [PATCH 42/64] feat!: return lazy iterators for table and namespace listings (#4599) Make table and namespace listings consistent with other resource listings across Rust, Python, and TypeScript. Callers can iterate over individual names across all pages, with configurable page sizes and continuation tokens. In order to avoid breaking existing Python users, this change adds some compatibility methods to the returned Python iterators. Code such as: ``` for t in db.list_tables().tables: print(t) ``` will continue to work as expected because the iterator has a `tables` field and a `page_token` field, which function as expected to reveal the first page of results. In other words, it is never necessary to use the iterator __as an iterator__. It may be used as if it were an old-style response object. (As the iterator progresses, the `tables` list may switch to the second, third, etc. page of results.) Similar compatibility code allows the async python API to keep working as well. --- docs/src/js/classes/Connection.md | 80 ++------ docs/src/js/globals.md | 4 - .../js/interfaces/ListNamespacesOptions.md | 27 --- .../js/interfaces/ListNamespacesResponse.md | 23 --- docs/src/js/interfaces/ListTablesOptions.md | 35 ---- docs/src/js/interfaces/ListTablesResponse.md | 23 --- docs/src/js/interfaces/TableNamesOptions.md | 2 +- docs/src/python/python.md | 23 ++- docs/web-tests/ts/connection.test.ts | 16 +- docs/web-tests/ts/quickstart.test.ts | 7 +- nodejs/__test__/connection.test.ts | 156 ++++++++-------- nodejs/lancedb/connection.ts | 133 ++++---------- nodejs/lancedb/index.ts | 4 - nodejs/src/connection.rs | 101 ++--------- python/python/lancedb/_lancedb.pyi | 11 +- python/python/lancedb/db.py | 70 ++++--- python/python/lancedb/listing.py | 171 ++++++++++++++++++ python/python/lancedb/namespace.py | 39 ++-- python/python/lancedb/remote/db.py | 16 +- python/python/tests/test_embeddings.py | 4 + python/python/tests/test_listing.py | 137 ++++++++++++++ python/python/tests/test_ollama_embeddings.py | 2 +- python/src/connection.rs | 56 ++---- python/src/listing.rs | 30 ++- rust/lancedb/src/connection.rs | 161 +++++++++-------- rust/lancedb/src/listing.rs | 130 +++++++++---- rust/lancedb/src/remote/db.rs | 66 ++----- 27 files changed, 814 insertions(+), 713 deletions(-) delete mode 100644 docs/src/js/interfaces/ListNamespacesOptions.md delete mode 100644 docs/src/js/interfaces/ListNamespacesResponse.md delete mode 100644 docs/src/js/interfaces/ListTablesOptions.md delete mode 100644 docs/src/js/interfaces/ListTablesResponse.md diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index d5ab8c3d1..c110eb38c 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -653,94 +653,48 @@ See [Listing](Listing.md) for pagination state and error behavior. ### listNamespaces() ```ts -abstract listNamespaces(namespacePath?, options?): Promise +abstract listNamespaces(namespacePath?, options?): Listing ``` -List the immediate child namespaces under the given parent. - -Results may be paginated. To retrieve subsequent pages, pass the -`pageToken` returned by a previous call. +Lazily list immediate child namespaces, 100 per page by default. #### Parameters * **namespacePath?**: `string`[] - The parent namespace path. Defaults - to the root namespace if omitted. -* **options?**: `Partial`<[`ListNamespacesOptions`](../interfaces/ListNamespacesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) #### Returns -`Promise`<[`ListNamespacesResponse`](../interfaces/ListNamespacesResponse.md)> - -Child namespace names and - an optional token for fetching the next page. +[`Listing`](Listing.md)<`string`> *** ### listTables() -#### listTables(options) - ```ts -abstract listTables(options?): Promise +abstract listTables(namespacePath?, options?): Listing ``` -List a page of the tables in this database. +Lazily list tables in a namespace, requesting 100 per page by default. -To retrieve the tables after the page, pass the `pageToken` the response -carries back in. A page can be shorter than `limit` without being the last -one, so walk until a response carries no page token: - -```ts -const names = []; -let pageToken = undefined; -do { - const page = await conn.listTables({ pageToken, limit: 100 }); - names.push(...page.tables); - pageToken = page.pageToken; -} while (pageToken); -``` - -##### Parameters - -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). - -##### Returns - -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> - -A page of table names and an - optional token for the tables after it. - -#### listTables(namespacePath, options) - -```ts -abstract listTables(namespacePath?, options?): Promise -``` - -List a page of the tables in this database. - -##### Parameters +#### Parameters * **namespacePath?**: `string`[] - The namespace path to list tables from - (defaults to root namespace) -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) -##### Returns +#### Returns -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> +[`Listing`](Listing.md)<`string`> -A page of table names and an - optional token for the tables after it. +#### Example + +```ts +for await (const name of conn.listTables([], { pageLimit: 20 })) { + console.log(name); +} +``` *** diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 646341190..e3830e9ad 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -108,10 +108,6 @@ - [JobEventsOptions](interfaces/JobEventsOptions.md) - [JobFailureInfo](interfaces/JobFailureInfo.md) - [JobInfo](interfaces/JobInfo.md) -- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) -- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) -- [ListTablesOptions](interfaces/ListTablesOptions.md) -- [ListTablesResponse](interfaces/ListTablesResponse.md) - [ListingOptions](interfaces/ListingOptions.md) - [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md) - [MemtableStats](interfaces/MemtableStats.md) diff --git a/docs/src/js/interfaces/ListNamespacesOptions.md b/docs/src/js/interfaces/ListNamespacesOptions.md deleted file mode 100644 index 79680417f..000000000 --- a/docs/src/js/interfaces/ListNamespacesOptions.md +++ /dev/null @@ -1,27 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesOptions - -# Interface: ListNamespacesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An optional limit to the number of results to return. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response for pagination. diff --git a/docs/src/js/interfaces/ListNamespacesResponse.md b/docs/src/js/interfaces/ListNamespacesResponse.md deleted file mode 100644 index a99ab480f..000000000 --- a/docs/src/js/interfaces/ListNamespacesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesResponse - -# Interface: ListNamespacesResponse - -## Properties - -### namespaces - -```ts -namespaces: string[]; -``` - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` diff --git a/docs/src/js/interfaces/ListTablesOptions.md b/docs/src/js/interfaces/ListTablesOptions.md deleted file mode 100644 index b8dd8b5ef..000000000 --- a/docs/src/js/interfaces/ListTablesOptions.md +++ /dev/null @@ -1,35 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesOptions - -# Interface: ListTablesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An upper bound on how many tables to return. -Defaults to 100. Zero returns an empty page without a continuation token. - -A page may hold fewer than this and still not be the last one, so keep -going while the response carries a page token rather than while pages are -full. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response, to resume listing where it left off. - -The token is opaque: it carries whatever the database needs to resume, and -callers should not construct or interpret one. diff --git a/docs/src/js/interfaces/ListTablesResponse.md b/docs/src/js/interfaces/ListTablesResponse.md deleted file mode 100644 index 76cac2b23..000000000 --- a/docs/src/js/interfaces/ListTablesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesResponse - -# Interface: ListTablesResponse - -## Properties - -### pageToken? - -```ts -optional pageToken: string; -``` - -*** - -### tables - -```ts -tables: string[]; -``` diff --git a/docs/src/js/interfaces/TableNamesOptions.md b/docs/src/js/interfaces/TableNamesOptions.md index 123aa6c9f..e1fd8c382 100644 --- a/docs/src/js/interfaces/TableNamesOptions.md +++ b/docs/src/js/interfaces/TableNamesOptions.md @@ -8,7 +8,7 @@ ## Deprecated -Use [ListTablesOptions](ListTablesOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) +Use [ListingOptions](ListingOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) instead. ## Properties diff --git a/docs/src/python/python.md b/docs/src/python/python.md index f4c697497..c465d0ac4 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -26,14 +26,17 @@ is also an [asynchronous API client](#connections-asynchronous). ::: lancedb.db.DBConnection -Resource listings (`list_secrets`, `list_views`, `list_jobs`, `list_functions`, and -`list_materialized_views`) return lazy iterators. Use `list(db.list_views())` to -collect synchronous results, or `[name async for name in db.list_views()]` with -an asynchronous connection. `page_limit` controls each request, and `page_token` -resumes a listing. Drain the iterator's cached results before saving its token. +Resource listings return iterators. Use `list(db.list_tables())` to collect all +table names. `db.list_tables().tables` is a list of the first page; `limit` +controls its page size and `page_token` resumes at a saved token. The same +patterns apply to `list_namespaces()`. Other resource listings use `page_limit`. ::: lancedb.listing.Listing +::: lancedb.listing.TableListing + +::: lancedb.listing.NamespaceListing + ::: lancedb.Session ## Catalogs (Synchronous) @@ -631,6 +634,16 @@ can be used to create, list, or open tables. ::: lancedb.listing.AsyncListing +`async for name in db.list_tables()` loads pages as needed. To inspect the +first page, use `page = await db.list_tables()` and then `page.tables`. Reading +the page list does not consume the iterator. `page.page_token` is the token for +the next request; drain cached results before using it to resume in a new +listing. The same behavior applies to namespaces. + +::: lancedb.listing.AsyncTableListing + +::: lancedb.listing.AsyncNamespaceListing + ## Namespaces (Asynchronous) ::: lancedb.connect_namespace_async diff --git a/docs/web-tests/ts/connection.test.ts b/docs/web-tests/ts/connection.test.ts index 891501d4c..651be7d36 100644 --- a/docs/web-tests/ts/connection.test.ts +++ b/docs/web-tests/ts/connection.test.ts @@ -4,6 +4,12 @@ import { afterEach, describe, expect, jest, test } from "@jest/globals"; import * as lancedb from "@lancedb/lancedb"; import { withTempDirectory } from "./util.ts"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + test("connect to a local database", async () => { await withTempDirectory(async (tempDir) => { const cwd = process.cwd(); @@ -13,12 +19,12 @@ test("connect to a local database", async () => { const db = await lancedb.connect("ex_lancedb"); // --8<-- [end:connect] await db.createTable("probe", [{ id: 1 }]); - expect((await db.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(db.listTables())).toEqual(["probe"]); } finally { process.chdir(cwd); } const reopened = await lancedb.connect(`${tempDir}/ex_lancedb`); - expect((await reopened.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(reopened.listTables())).toEqual(["probe"]); }); }); @@ -109,8 +115,8 @@ async function namespaceTableOpsExample() { ); // Verify - console.log((await db.listNamespaces()).namespaces); // ["prod"] - console.log((await db.listNamespaces(["prod"])).namespaces); // ["recommendations", "search"] + console.log(await collectNames(db.listNamespaces())); // ["prod"] + console.log(await collectNames(db.listNamespaces(["prod"]))); // ["recommendations", "search"] console.log(await db.tableNames(["prod", "search"])); // ["user"] console.log(await db.tableNames(["prod", "recommendations"])); // ["user"] // --8<-- [end:namespace_table_ops] @@ -124,7 +130,7 @@ async function namespaceAdminOpsExample() { await db.createNamespace(["prod"]); await db.createNamespace(["prod", "search"]); - const childNamespaces = (await db.listNamespaces(["prod"])).namespaces; + const childNamespaces = await collectNames(db.listNamespaces(["prod"])); console.log(`Child namespaces under ${JSON.stringify(namespace)}:`, childNamespaces); // Child namespaces under ["prod","search"]: [ 'search' ] diff --git a/docs/web-tests/ts/quickstart.test.ts b/docs/web-tests/ts/quickstart.test.ts index e32a14f86..2c3ee8252 100644 --- a/docs/web-tests/ts/quickstart.test.ts +++ b/docs/web-tests/ts/quickstart.test.ts @@ -20,12 +20,7 @@ const connections = async function tableNames(db: lancedb.Connection) { const names: string[] = []; - let pageToken: string | undefined; - do { - const page = await db.listTables({ pageToken }); - names.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); + for await (const name of db.listTables()) names.push(name); return names; } diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index bbd0c3c52..73f99803a 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -5,15 +5,15 @@ import { mkdirSync, readdirSync } from "fs"; import { join } from "path"; import { Field, Float64, Schema } from "apache-arrow"; import * as tmp from "tmp"; -import { - Connection, - ListTablesResponse, - Table, - connect, - connectNamespace, -} from "../lancedb"; +import { Connection, Table, connect, connectNamespace } from "../lancedb"; import { LocalTable } from "../lancedb/table"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + describe("when connecting", () => { let tmpDir: tmp.DirResult; beforeEach(() => { @@ -135,7 +135,7 @@ describe("given a connection", () => { await db.close(); expect(db.isOpen()).toBe(false); await expect(db.tableNames()).rejects.toThrow("Connection is closed"); - await expect(db.listTables()).rejects.toThrow("Connection is closed"); + expect(() => db.listTables()).toThrow("Connection is closed"); await expect(db.renameTable("a", "b")).rejects.toThrow( "Connection is closed", ); @@ -247,15 +247,15 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const first = await db.listTables(namespacePath, { limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables(namespacePath, { pageLimit: 1 }); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables(namespacePath, { - limit: 1, - pageToken: first.pageToken, + const second = db.listTables(namespacePath, { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); }); it("should keep the existing table when mode is exist_ok", async () => { @@ -299,19 +299,18 @@ describe("given a connection", () => { await db.createTable("a", [{ id: 1 }]); await db.createTable("c", [{ id: 1 }]); - const all = await db.listTables(); - expect(all.tables).toEqual(["a", "b", "c"]); - expect(all.pageToken).toBeUndefined(); + expect(await collectNames(db.listTables())).toEqual(["a", "b", "c"]); - const first = await db.listTables({ limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables({ - limit: 1, - pageToken: first.pageToken, + const second = db.listTables([], { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); it("should visit every table exactly once when walking pages", async () => { @@ -322,18 +321,9 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const seen: string[] = []; - let pageToken: string | undefined = undefined; - do { - const page: ListTablesResponse = await db.listTables({ - limit: 2, - pageToken, - }); - seen.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); - - expect(seen).toEqual(created); + expect(await collectNames(db.listTables([], { pageLimit: 2 }))).toEqual( + created, + ); }); it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])( @@ -347,11 +337,10 @@ describe("given a connection", () => { ]) { await expect(names).resolves.toEqual(["a", "b"]); } - for (const page of [ - db.listTables({ limit }), - db.listTables([], { limit }), - ]) { - await expect(page).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + expect(() => + db.listTables(namespacePath, { pageLimit: limit }), + ).toThrow(/limit/); } }, ); @@ -366,8 +355,11 @@ describe("given a connection", () => { ])("should reject an invalid table listing limit %s", async (limit) => { await expect(db.tableNames({ limit })).rejects.toThrow(/limit/); await expect(db.tableNames([], { limit })).rejects.toThrow(/limit/); - await expect(db.listTables({ limit })).rejects.toThrow(/limit/); - await expect(db.listTables([], { limit })).rejects.toThrow(/limit/); + for (const namespacePath of [undefined, []]) { + expect(() => db.listTables(namespacePath, { pageLimit: limit })).toThrow( + /limit/, + ); + } }); it("should preserve zero, small and unsigned-32-bit table listing limits", async () => { @@ -379,11 +371,21 @@ describe("given a connection", () => { [2 ** 32 - 1, ["a", "b"]], ] as const) { await expect(db.tableNames({ limit })).resolves.toEqual(names); - const page = await db.listTables({ limit }); - expect(page.tables).toEqual(names); } await expect(db.tableNames()).resolves.toEqual(["a", "b"]); - await expect(db.listTables()).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + for (const pageLimit of [1, 2 ** 31 - 1]) { + expect( + await collectNames(db.listTables(namespacePath, { pageLimit })), + ).toEqual(["a", "b"]); + } + for (const pageLimit of [0, 2 ** 31, 2 ** 32 - 1]) { + expect(() => db.listTables(namespacePath, { pageLimit })).toThrow( + /limit/, + ); + } + } + expect(await collectNames(db.listTables())).toEqual(["a", "b"]); }); it("should list all names and use a default page size of 100", async () => { @@ -396,20 +398,17 @@ describe("given a connection", () => { } await expect(db.tableNames()).resolves.toEqual(names); await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]); - const first = await db.listTables(); - expect(first.tables).toEqual(names.slice(0, 100)); - expect(first.pageToken).toBeDefined(); - const second = await db.listTables({ pageToken: first.pageToken }); - expect(second.tables).toEqual(names.slice(100)); - expect(second.pageToken).toBeUndefined(); - await expect(db.listTables({ limit: 200 })).resolves.toEqual({ - tables: names, - pageToken: undefined, - }); - await expect(db.listTables({ limit: 0 })).resolves.toEqual({ - tables: [], - pageToken: undefined, - }); + const first = db.listTables(); + const firstPage: string[] = []; + for (let i = 0; i < 100; i++) firstPage.push((await first.next()).value!); + expect(firstPage).toEqual(names.slice(0, 100)); + expect(first.pageToken()).toBeDefined(); + expect(first.numPageResults()).toBe(0); + expect(await collectNames(first)).toEqual(names.slice(100)); + expect(await collectNames(db.listTables([], { pageLimit: 200 }))).toEqual( + names, + ); + expect(() => db.listTables([], { pageLimit: 0 })).toThrow(/limit/); }); it("should list tables in a namespace", async () => { @@ -420,12 +419,8 @@ describe("given a connection", () => { await db.createNamespace(["child"]); await db.createTable("nested", [{ id: 1 }], ["child"]); - await expect(db.listTables(["child"])).resolves.toEqual( - expect.objectContaining({ tables: ["nested"] }), - ); - await expect(db.listTables()).resolves.toEqual( - expect.objectContaining({ tables: [] }), - ); + expect(await collectNames(db.listTables(["child"]))).toEqual(["nested"]); + expect(await collectNames(db.listTables())).toEqual([]); }); it("should create tables in v2 mode", async () => { @@ -719,22 +714,35 @@ describe("namespaces", () => { it("should list namespaces created at the root", async () => { await db.createNamespace(["alpha"]); await db.createNamespace(["beta"]); - const list = await db.listNamespaces(); - expect(list.namespaces).toEqual(expect.arrayContaining(["alpha", "beta"])); + const list = await collectNames(db.listNamespaces()); + expect(list).toEqual(expect.arrayContaining(["alpha", "beta"])); }); it("should list child namespaces under a parent", async () => { await db.createNamespace(["parent"]); await db.createNamespace(["parent", "child"]); - const list = await db.listNamespaces(["parent"]); - expect(list.namespaces).toContain("child"); + const list = await collectNames(db.listNamespaces(["parent"])); + expect(list).toContain("child"); + }); + + it("should resume a namespace listing from its page token", async () => { + await db.createNamespace(["alpha"]); + await db.createNamespace(["beta"]); + const first = db.listNamespaces([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect((await first.next()).value).toBe("alpha"); + const second = db.listNamespaces([], { + pageLimit: 1, + pageToken: first.pageToken(), + }); + expect(await collectNames(second)).toEqual(["beta"]); }); it("should drop a namespace", async () => { await db.createNamespace(["ephemeral"]); await db.dropNamespace(["ephemeral"]); - const list = await db.listNamespaces(); - expect(list.namespaces).not.toContain("ephemeral"); + const list = await collectNames(db.listNamespaces()); + expect(list).not.toContain("ephemeral"); }); it("should raise an error on any namespace op after close", async () => { @@ -742,7 +750,7 @@ describe("namespaces", () => { await expect(db.describeNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); - await expect(db.listNamespaces()).rejects.toThrow("Connection is closed"); + expect(() => db.listNamespaces()).toThrow("Connection is closed"); await expect(db.createNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index c93414fab..5f067eeef 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -30,16 +30,12 @@ import type { DescribeNamespaceResponse, DropNamespaceResponse, JobInfo, - ListNamespacesResponse, - ListTablesResponse, } from "./native"; import { ViewDescription, viewDescriptionFromNative } from "./view"; export type { CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, }; import { sanitizeTable } from "./sanitize"; import { LocalTable, Table } from "./table"; @@ -141,7 +137,7 @@ export interface OpenTableOptions { } /** - * @deprecated Use {@link ListTablesOptions} with {@link Connection.listTables} + * @deprecated Use {@link ListingOptions} with {@link Connection.listTables} * instead. */ export interface TableNamesOptions { @@ -157,32 +153,6 @@ export interface TableNamesOptions { limit?: number; } -export interface ListTablesOptions { - /** - * Token from a previous response, to resume listing where it left off. - * - * The token is opaque: it carries whatever the database needs to resume, and - * callers should not construct or interpret one. - */ - pageToken?: string; - /** - * An upper bound on how many tables to return. - * Defaults to 100. Zero returns an empty page without a continuation token. - * - * A page may hold fewer than this and still not be the last one, so keep - * going while the response carries a page token rather than while pages are - * full. - */ - limit?: number; -} - -export interface ListNamespacesOptions { - /** Token from a previous response for pagination. */ - pageToken?: string; - /** An optional limit to the number of results to return. */ - limit?: number; -} - export interface CreateNamespaceOptions { /** Creation mode. */ mode?: "create" | "exist_ok" | "overwrite"; @@ -278,45 +248,18 @@ export abstract class Connection { options?: Partial, ): Promise; - /** - * List a page of the tables in this database. - * - * To retrieve the tables after the page, pass the `pageToken` the response - * carries back in. A page can be shorter than `limit` without being the last - * one, so walk until a response carries no page token: - * + /** Lazily list tables in a namespace, requesting 100 per page by default. + * @example * ```ts - * const names = []; - * let pageToken = undefined; - * do { - * const page = await conn.listTables({ pageToken, limit: 100 }); - * names.push(...page.tables); - * pageToken = page.pageToken; - * } while (pageToken); + * for await (const name of conn.listTables([], { pageLimit: 20 })) { + * console.log(name); + * } * ``` - * - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. - */ - abstract listTables( - options?: Partial, - ): Promise; - /** - * List a page of the tables in this database. - * - * @param {string[]} namespacePath - The namespace path to list tables from - * (defaults to root namespace) - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. */ abstract listTables( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Open a table in the database. @@ -542,23 +485,11 @@ export abstract class Connection { namespacePath: string[], ): Promise; - /** - * List the immediate child namespaces under the given parent. - * - * Results may be paginated. To retrieve subsequent pages, pass the - * `pageToken` returned by a previous call. - * - * @param {string[]} namespacePath - The parent namespace path. Defaults - * to the root namespace if omitted. - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} Child namespace names and - * an optional token for fetching the next page. - */ + /** Lazily list immediate child namespaces, 100 per page by default. */ abstract listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Create a new namespace at the given path. @@ -821,22 +752,17 @@ export class LocalConnection extends Connection { ); } - async listTables( - namespacePathOrOptions?: string[] | Partial, - options?: Partial, - ): Promise { - // Detect if first argument is namespacePath array or options object - const namespacePath = Array.isArray(namespacePathOrOptions) - ? namespacePathOrOptions - : undefined; - const listTablesOptions = Array.isArray(namespacePathOrOptions) - ? options - : (namespacePathOrOptions ?? options); - - return this.inner.listTables( - namespacePath ?? [], - listTablesOptions?.pageToken, - normalizeTableListingLimit(listTablesOptions?.limit), + listTables( + namespacePath?: string[], + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listTables( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } @@ -1039,12 +965,15 @@ export class LocalConnection extends Connection { listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise { - return this.inner.listNamespaces( - namespacePath ?? [], - options?.pageToken, - options?.limit, + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listNamespaces( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 4ad8b2a93..0f1b5f333 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -97,13 +97,9 @@ export { Connection, CreateTableOptions, TableNamesOptions, - ListTablesOptions, OpenTableOptions, - ListNamespacesOptions, CreateNamespaceOptions, DropNamespaceOptions, - ListNamespacesResponse, - ListTablesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 67c4706e4..914f9e4c9 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -16,8 +16,7 @@ use crate::table::Table; use lancedb::connection::{ConnectBuilder, Connection as LanceDBConnection, connect_namespace}; use lance_namespace::models::{ - CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest, - ListTablesRequest, + CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, }; use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema}; @@ -31,18 +30,6 @@ pub struct DescribeNamespaceResponse { pub properties: Option>, } -#[napi(object)] -pub struct ListNamespacesResponse { - pub namespaces: Vec, - pub page_token: Option, -} - -#[napi(object)] -pub struct ListTablesResponse { - pub tables: Vec, - pub page_token: Option, -} - #[napi(object)] pub struct CreateNamespaceResponse { pub properties: Option>, @@ -239,31 +226,20 @@ impl Connection { op.execute().await.default_error() } - /// List a page of tables in the database. + /// Lazily list tables in the database. #[napi(catch_unwind)] - pub async fn list_tables( + pub fn list_tables( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let request = ListTablesRequest { - // The root namespace is an empty path, not an absent one: a namespace-backed - // database rejects a request that names no namespace. - id: Some(namespace_path.unwrap_or_default()), - page_token, - limit: limit.map(|limit| i32::try_from(limit).unwrap_or(i32::MAX)), - ..Default::default() - }; - let response = self - .get_inner()? - .list_tables(request) - .await - .default_error()?; - Ok(ListTablesResponse { - tables: response.tables, - page_token: response.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Create table from a Apache Arrow IPC (file) buffer. @@ -663,28 +639,19 @@ impl Connection { } #[napi(catch_unwind)] - /// List child namespaces under the given namespace path - pub async fn list_namespaces( + /// Lazily list child namespaces under the given namespace path + pub fn list_namespaces( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let req = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: checked_limit(limit)?, - ..Default::default() - }; - let resp = self - .get_inner()? - .list_namespaces(req) - .await - .default_error()?; - Ok(ListNamespacesResponse { - namespaces: resp.namespaces, - page_token: resp.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[napi(catch_unwind)] @@ -787,33 +754,3 @@ impl Connection { .default_error() } } - -fn checked_limit(limit: Option) -> napi::Result> { - limit - .map(i32::try_from) - .transpose() - .map_err(|_| napi::Error::from_reason("limit must not exceed 2147483647")) -} - -#[cfg(test)] -mod tests { - use super::checked_limit; - - #[test] - fn checked_limit_rejects_values_above_i32_max() { - // u32 values above i32::MAX wrapped negative through `as i32`, - // sending a negative page size to the server. - assert!(checked_limit(Some(3_000_000_000)).is_err()); - assert!(checked_limit(Some(u32::MAX)).is_err()); - } - - #[test] - fn checked_limit_passes_boundary_and_missing_values() { - assert_eq!( - checked_limit(Some(i32::MAX as u32)).unwrap(), - Some(i32::MAX) - ); - assert_eq!(checked_limit(Some(0)).unwrap(), Some(0)); - assert_eq!(checked_limit(None).unwrap(), None); - } -} diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 96eef064c..6ef23abc7 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -19,11 +19,9 @@ from .index import ( FTS, ) from lance_namespace import ( - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, ) from .remote import ClientConfig from .types import BaseTokenizerType @@ -114,12 +112,12 @@ class Connection(object): uri: str async def is_open(self): ... async def close(self): ... - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: ... + ) -> NameListing: ... async def create_namespace( self, namespace_path: List[str], @@ -136,12 +134,12 @@ class Connection(object): self, namespace_path: List[str], ) -> DescribeNamespaceResponse: ... - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: ... + ) -> NameListing: ... async def table_names( self, namespace_path: Optional[List[str]], @@ -1010,6 +1008,7 @@ async def connect_catalog( class NameListing: def num_page_results(self) -> int: ... def page_token(self) -> Optional[str]: ... + async def next_page(self) -> Optional[dict[str, Any]]: ... def __aiter__(self) -> NameListing: ... async def __anext__(self) -> str: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index a74763afb..90dc1cccf 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -32,18 +32,24 @@ if sys.version_info >= (3, 12): else: from overrides import EnforceOverrides, override # type: ignore -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, + validate_page_limit, +) from lancedb.embeddings.registry import EmbeddingFunctionRegistry from lancedb.common import data_to_reader, sanitize_uri, validate_schema from lancedb.background_loop import LOOP from lance_namespace import ( LanceNamespace, - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, connect as namespace_connect, ) from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError @@ -171,7 +177,7 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -187,12 +193,12 @@ class DBConnection(EnforceOverrides): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return ListNamespacesResponse(namespaces=[], page_token=None) + return NamespaceListing(AsyncNamespaceListing(None, empty=True)) def create_namespace( self, @@ -317,8 +323,8 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -334,7 +340,7 @@ class DBConnection(EnforceOverrides): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ raise NotImplementedError( @@ -1436,7 +1442,7 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -1452,16 +1458,14 @@ class LanceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( - namespace_path=namespace_path, - page_token=page_token, - limit=limit, + namespace_path=namespace_path, page_token=page_token, limit=limit ) ) @@ -1512,8 +1516,8 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -1529,12 +1533,12 @@ class LanceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -2286,12 +2290,12 @@ class AsyncConnection(object): else: return None - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -2306,15 +2310,18 @@ class AsyncConnection(object): Returns ------- - ListNamespacesResponse + AsyncNamespaceListing Response containing namespace names and optional pagination token """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_namespaces( + validate_page_limit(limit) + inner = self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListNamespacesResponse(**result) + if limit == 0: + return AsyncNamespaceListing(None, empty=True) + return AsyncNamespaceListing(inner, page_token) async def create_namespace( self, @@ -2393,13 +2400,13 @@ class AsyncConnection(object): result = await self._inner.describe_namespace(namespace_path) return DescribeNamespaceResponse(**result) - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> AsyncTableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -2415,15 +2422,18 @@ class AsyncConnection(object): Returns ------- - ListTablesResponse + AsyncTableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_tables( + validate_page_limit(limit) + inner = self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListTablesResponse(**result) + if limit == 0: + return AsyncTableListing(None, empty=True) + return AsyncTableListing(inner, page_token) async def table_names( self, diff --git a/python/python/lancedb/listing.py b/python/python/lancedb/listing.py index 0620aa7d8..8abc9d85b 100644 --- a/python/python/lancedb/listing.py +++ b/python/python/lancedb/listing.py @@ -5,6 +5,7 @@ from __future__ import annotations +import asyncio from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar from .background_loop import LOOP @@ -83,3 +84,173 @@ class Listing(Iterator[T], Generic[T]): for initial, cached, final-page, and error behavior. """ return self._inner.page_token() + + +def validate_page_limit(limit: Optional[int]) -> None: + """Validate the legacy limit argument before passing it to the bindings.""" + if limit is not None and ( + not isinstance(limit, int) or limit < 0 or limit > 2**31 - 1 + ): + raise ValueError("limit must be between 0 and 2147483647") + + +class _AsyncPagedNameListing(AsyncIterator[str]): + """An async listing with an inspectable copy of its current backend page.""" + + def __init__( + self, inner: Any, page_token: Optional[str] = None, *, empty: bool = False + ): + self._inner = inner + self._token = page_token + self._page: Optional[list[str]] = [] if empty else None + self._position = 0 + self._done = empty + self._failure: Optional[BaseException] = None + self._lock = asyncio.Lock() + + @property + def page_token(self) -> Optional[str]: + """The continuation token for the next backend request.""" + return self._token + + def num_page_results(self) -> int: + """Count the remaining names in the current page.""" + return 0 if self._page is None else len(self._page) - self._position + + def _page_copy(self) -> list[str]: + if self._page is None: + raise RuntimeError("Await or iterate the listing before reading its page") + return self._page.copy() + + async def _fetch_page(self) -> None: + try: + page = await self._inner.next_page() + except asyncio.CancelledError: + raise + except BaseException as exc: + self._failure = exc + self._done = True + raise + if page is None: + self._page = [] + self._token = None + self._done = True + else: + self._page = page["items"] + self._token = page["page_token"] + self._done = self._token is None + self._position = 0 + + async def _load_initial(self): + async with self._lock: + if self._failure is not None: + raise self._failure + if self._page is None: + await self._fetch_page() + return self + + def __await__(self): + return self._load_initial().__await__() + + def __aiter__(self): + return self + + async def __anext__(self) -> str: + async with self._lock: + if self._failure is not None: + raise StopAsyncIteration + if self._page is None: + await self._fetch_page() + assert self._page is not None + while self._position >= len(self._page): + if self._done: + raise StopAsyncIteration + await self._fetch_page() + assert self._page is not None + item = self._page[self._position] + self._position += 1 + return item + + +class AsyncTableListing(_AsyncPagedNameListing): + """Awaitable table iterator; awaiting loads one page. + + For example:: + + page = await db.list_tables(limit=20) + print(page.tables, page.page_token) + async for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + """A copy of every table name in the most recently fetched page.""" + return self._page_copy() + + +class AsyncNamespaceListing(_AsyncPagedNameListing): + """Awaitable namespace iterator; awaiting loads one page. + + For example:: + + page = await db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + """A copy of every name in the most recently fetched page.""" + return self._page_copy() + + +class _PagedNameListing(Iterator[str]): + def __init__(self, inner: _AsyncPagedNameListing): + self._inner = inner + LOOP.run(inner._load_initial()) + + def __iter__(self): + return self + + def __next__(self) -> str: + try: + return LOOP.run(self._inner.__anext__()) + except StopAsyncIteration: + raise StopIteration from None + + @property + def page_token(self) -> Optional[str]: + return self._inner.page_token + + def num_page_results(self) -> int: + return self._inner.num_page_results() + + +class TableListing(_PagedNameListing): + """Table iterator retaining the current page as a list. + + For example:: + + page = db.list_tables(limit=20) + print(page.tables, page.page_token) + for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + return self._inner.tables + + +class NamespaceListing(_PagedNameListing): + """Namespace iterator retaining the current page as a list. + + For example:: + + page = db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + return self._inner.namespaces diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index a5d88eb1f..55baa39a3 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -52,7 +52,14 @@ from lancedb.background_loop import LOOP from lancedb.arrow import AsyncRecordBatchReader from lancedb.db import AsyncConnection, DBConnection from lancedb.job import AsyncJob, Job -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, +) from lancedb.sql import AsyncQuery as AsyncSqlQuery from lancedb.sql import QueryDescription, QueryParameters from lance_namespace import ( @@ -61,8 +68,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, NamespaceExistsRequest, TableExistsRequest, ) @@ -854,7 +859,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """ List child namespaces under the given namespace. @@ -871,12 +876,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1023,7 +1028,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """ List all tables in this database with pagination support. @@ -1040,12 +1045,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1459,12 +1464,12 @@ class AsyncLanceNamespaceDBConnection: namespace_path = [] await self._inner.drop_all_tables(namespace_path=namespace_path) - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """ List child namespaces under the given namespace. @@ -1481,12 +1486,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_namespaces( + return self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1616,12 +1621,12 @@ class AsyncLanceNamespaceDBConnection: except TableNotFoundError: return False - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> AsyncTableListing: """ List all tables in this database with pagination support. @@ -1638,12 +1643,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_tables( + return self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 89c7f8a78..08125cc42 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -36,7 +36,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig import pyarrow as pa from ..common import DATA -from lancedb.listing import Listing +from lancedb.listing import Listing, NamespaceListing, TableListing from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job @@ -54,8 +54,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, ) from ..pydantic import LanceModel from ..table import Table @@ -268,7 +266,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -284,12 +282,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -379,7 +377,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """List a page of tables in this database. Parameters @@ -396,12 +394,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index e319d3048..207a3d253 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -896,6 +896,8 @@ def test_transient_embedding_error_is_retried_in_table_add(tmp_path): assert len(calls) == 3 assert mock_sleep.call_count == 2 assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list) + + @patch("time.sleep") @pytest.mark.parametrize( "error", @@ -941,6 +943,8 @@ def test_retry_on_wrapped_provider_error(mock_sleep): assert wrapped_func() == [0.1] assert mock_func.call_count == 2 assert mock_sleep.call_count == 1 + + def test_openai_bad_request_partial_batch_recovery(): # Regression test for issue #1677 diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py index 4a2745eb8..b0b973bbf 100644 --- a/python/python/tests/test_listing.py +++ b/python/python/tests/test_listing.py @@ -88,6 +88,143 @@ async def advance(items, asynchronous): KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"] +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_compatible_name_listing_pages(listing_server, kind, asynchronous): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a", "b"], "next")), + (200, response_page(kind, ["c"], None)), + ] + ) + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)( + namespace_path=["team"], page_token="start", limit=2 + ) + field = kind + assert len(requests) == (0 if asynchronous else 1) + if asynchronous: + with pytest.raises(RuntimeError, match="Await or iterate"): + getattr(listing, field) + assert listing.page_token == "start" + assert await listing is listing + assert len(requests) == 1 + page_names = getattr(listing, field) + assert type(page_names) is list + assert page_names == ["a", "b"] + page_names.append("changed") + assert getattr(listing, field) == ["a", "b"] + assert listing.page_token == "next" + assert listing.num_page_results() == 2 + assert await advance(listing, asynchronous) == "a" + assert getattr(listing, field) == ["a", "b"] + assert await advance(listing, asynchronous) == "b" + assert len(requests) == 1 + assert await advance(listing, asynchronous) == "c" + assert getattr(listing, field) == ["c"] + assert listing.page_token is None + assert len(requests) == 2 + assert ([item async for item in listing] if asynchronous else list(listing)) == [] + if asynchronous: + assert await listing is listing + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_skips_empty_page_only_during_iteration( + listing_server, kind +): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, [], "next")), + (200, response_page(kind, ["a"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + await listing + assert getattr(listing, kind) == [] + assert listing.page_token == "next" + assert len(requests) == 1 + assert [name async for name in listing] == ["a"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_name_listing_zero_limit_uses_no_request( + listing_server, kind, asynchronous +): + endpoint, requests, _ = listing_server + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)(limit=0) + if asynchronous: + assert await listing is listing + assert [name async for name in listing] == [] + else: + assert list(listing) == [] + assert getattr(listing, kind) == [] + assert listing.page_token is None + assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_failure_retains_token(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((401, {})) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(page_token="retry") + with pytest.raises(HttpError): + await listing + assert listing.page_token == "retry" + with pytest.raises(HttpError): + await listing + assert [name async for name in listing] == [] + assert len(requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_direct_async_iteration(listing_server, kind): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a"], "next")), + (200, response_page(kind, ["b"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(limit=1) + assert requests == [] + assert [name async for name in listing] == ["a", "b"] + assert getattr(listing, kind) == ["b"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_serializes_concurrent_advances(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((200, response_page(kind, ["a", "b"], None))) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + results = await asyncio.gather( + listing.__anext__(), + listing.__anext__(), + listing.__anext__(), + return_exceptions=True, + ) + assert results[:2] == ["a", "b"] + assert isinstance(results[2], StopAsyncIteration) + assert len(requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", KINDS) @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/python/python/tests/test_ollama_embeddings.py b/python/python/tests/test_ollama_embeddings.py index 9d8568ca7..de36dc24b 100644 --- a/python/python/tests/test_ollama_embeddings.py +++ b/python/python/tests/test_ollama_embeddings.py @@ -3,7 +3,6 @@ import json -import httpx import numpy as np import pyarrow as pa import pytest @@ -14,6 +13,7 @@ from lancedb.embeddings import get_registry @pytest.fixture def ollama_embeddings(): + httpx = pytest.importorskip("httpx") pytest.importorskip("ollama") requests = [] diff --git a/python/src/connection.rs b/python/src/connection.rs index baeb29780..2cdf71f4a 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -590,29 +590,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_namespaces( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListNamespacesRequest; - let request = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_namespaces(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("namespaces", response.namespaces)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[pyo3(signature = (namespace_path, mode=None, properties=None))] @@ -725,29 +713,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_tables( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListTablesRequest; - let request = ListTablesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_tables(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("tables", response.tables)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Get the configuration for constructing an equivalent namespace client. diff --git a/python/src/listing.rs b/python/src/listing.rs index 2d1a3916e..79a7effdd 100644 --- a/python/src/listing.rs +++ b/python/src/listing.rs @@ -3,10 +3,11 @@ use crate::error::PythonErrorExt; use crate::runtime::future_into_py; -use futures::{StreamExt, future::poll_fn}; +use futures::{Future, StreamExt, future::poll_fn}; use lancedb::listing::{Listing, ListingOptions}; use pyo3::exceptions::PyStopAsyncIteration; -use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyDict, PyDictMethods}; +use pyo3::{Bound, Py, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; use std::sync::{Arc, Mutex}; pub fn options(page_token: Option, page_limit: Option) -> ListingOptions { @@ -52,6 +53,31 @@ impl NameListing { Ok(value) }) } + + fn next_page(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + let next_lock = self_.next_lock.clone(); + future_into_py(self_.py(), async move { + let _guard = next_lock.lock().await; + let page = futures::future::poll_fn(|cx| { + let mut inner = inner.lock().unwrap(); + let future = inner.next_page(); + futures::pin_mut!(future); + future.poll(cx) + }) + .await + .infer_error()?; + Python::attach(|py| -> PyResult>> { + page.map(|page| { + let dict = PyDict::new(py); + dict.set_item("items", page.items)?; + dict.set_item("page_token", page.page_token)?; + Ok(dict.unbind()) + }) + .transpose() + }) + }) + } } #[pyclass] diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 8a0e489c2..124f6a54a 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -13,7 +13,7 @@ use lance::dataset::refs::MAIN_BRANCH; use lance_namespace::models::{ CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest, DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest, - ListNamespacesResponse, ListTablesRequest, ListTablesResponse, + ListTablesRequest, }; #[cfg(feature = "aws")] use object_store::aws::AwsCredential; @@ -28,7 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; -use crate::listing::{Listing, ListingOptions}; +use crate::listing::{Listing, ListingOptions, ListingPage}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -1126,12 +1126,43 @@ impl Connection { self.internal.drop_all_tables(namespace_path).await } - /// List immediate child namespace names in the given namespace - pub async fn list_namespaces( + /// Lazily list immediate child namespace names. Each request asks for at + /// most 100 names unless `options.page_limit` is set. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_namespaces(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_namespaces( &self, - request: ListNamespacesRequest, - ) -> Result { - self.internal.list_namespaces(request).await + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_namespaces(ListNamespacesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.namespaces, + page_token: response.page_token, + }) + } + }) } /// Create a new namespace @@ -1182,23 +1213,43 @@ impl Connection { self.internal.namespace_client_config().await } - /// List tables with pagination support + /// Lazily list tables. Each request asks for at most 100 names unless + /// `options.page_limit` is set. /// - /// The default limit is 100 tables per page for both local and remote connections, - /// including namespaces. Zero returns an empty page without a continuation token. - /// Follow the response's opaque `page_token` until it is absent to retrieve every table; - /// a page can contain fewer than the limit even when more tables remain. - pub async fn list_tables(&self, mut request: ListTablesRequest) -> Result { - // Apply the SDK's page size here so every backend receives the same request. - let limit = request.limit.get_or_insert(100); - if *limit <= 0 { - return Ok(ListTablesResponse { - context: None, - tables: Vec::new(), - page_token: None, - }); - } - self.internal.list_tables(request).await + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_tables(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_tables( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_tables(ListTablesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.tables, + page_token: response.page_token, + }) + } + }) } /// Get the in-memory embedding registry. @@ -2260,26 +2311,11 @@ mod tests { } names.sort(); - // Walking in pages has to reach every table exactly once, with nothing lost at a - // page boundary. - let mut seen = Vec::with_capacity(names.len()); - let mut page_token = None; - loop { - let page = db - .list_tables(ListTablesRequest { - id: Some(Vec::new()), - limit: Some(2), - page_token, - ..Default::default() - }) - .await - .unwrap(); - seen.extend(page.tables); - page_token = page.page_token.filter(|token| !token.is_empty()); - if page_token.is_none() { - break; - } - } + let seen = futures::TryStreamExt::try_collect::>( + db.list_tables(&[], ListingOptions::default().page_limit(2)), + ) + .await + .unwrap(); assert_eq!(seen, names); } @@ -2305,37 +2341,20 @@ mod tests { .is_empty() ); - let first = db.list_tables(ListTablesRequest::default()).await.unwrap(); - assert_eq!(first.tables, names[..100]); + let mut listing = db.list_tables(&[], ListingOptions::default()); + assert_eq!(listing.num_page_results(), 0); + let first = listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items, names[..100]); assert!(first.page_token.is_some()); - let second = db - .list_tables(ListTablesRequest { - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(second.tables, names[100..]); + let second = listing.next_page().await.unwrap().unwrap(); + assert_eq!(second.items, names[100..]); assert!(second.page_token.is_none()); + assert!(listing.next_page().await.unwrap().is_none()); - let all = db - .list_tables(ListTablesRequest { - limit: Some(200), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(all.tables, names); - assert!(all.page_token.is_none()); - let empty = db - .list_tables(ListTablesRequest { - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(empty.tables.is_empty()); - assert!(empty.page_token.is_none()); + let mut all = db.list_tables(&[], ListingOptions::default().page_limit(200)); + assert_eq!(all.next_page().await.unwrap().unwrap().items, names); + let mut invalid = db.list_tables(&[], ListingOptions::default().page_limit(0)); + assert!(invalid.next_page().await.is_err()); } #[tokio::test] diff --git a/rust/lancedb/src/listing.rs b/rust/lancedb/src/listing.rs index 710202d9d..a0f6e5d5e 100644 --- a/rust/lancedb/src/listing.rs +++ b/rust/lancedb/src/listing.rs @@ -51,8 +51,9 @@ type FetchPage = /// A lazy stream of resources with inspectable pagination state. /// -/// No requests are made until the stream is polled. Empty pages with continuation -/// tokens are skipped. Errors terminate the stream and retain the failed request's +/// No requests are made until the stream is polled or [`Self::next_page`] is called. +/// Item iteration skips empty pages with continuation tokens; `next_page` returns +/// them. Errors terminate the listing and retain the failed request's /// token, allowing a new stream to retry it. Repeated tokens are rejected. /// /// A token advances when its page is fetched. Drain [`Self::num_page_results`] @@ -100,6 +101,64 @@ impl Listing { pub fn page_token(&self) -> Option<&str> { self.options.page_token.as_deref() } + + /// Return exactly one page without consuming a following page. + /// + /// If item iteration has started, this returns the remaining items in its + /// cached page. An empty page with a continuation token is returned as is. + pub async fn next_page(&mut self) -> Result>> { + if self.items.len() == 0 { + if self.done { + return Ok(None); + } + futures::future::poll_fn(|cx| self.poll_fetch_page(cx)).await?; + } + let items = std::mem::replace(&mut self.items, Vec::new().into_iter()).collect(); + Ok(Some(ListingPage { + items, + page_token: self.options.page_token.clone(), + })) + } + + fn poll_fetch_page(&mut self, cx: &mut Context<'_>) -> Poll> { + if self + .options + .page_limit + .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) + { + self.done = true; + return Poll::Ready(Err(Error::InvalidInput { + message: "Listing page limit must be between 1 and 2147483647".into(), + })); + } + let pending = self + .pending + .get_or_insert_with(|| (self.fetch)(self.options.clone())); + let response = futures::ready!(pending.as_mut().poll(cx)); + self.pending = None; + match response { + Ok(response) => { + let token = response.page_token.filter(|token| !token.is_empty()); + if token + .as_ref() + .is_some_and(|token| !self.seen_tokens.insert(token.clone())) + { + self.done = true; + return Poll::Ready(Err(Error::Runtime { + message: "Listing response repeated a page_token".into(), + })); + } + self.done = token.is_none(); + self.options.page_token = token; + self.items = response.items.into_iter(); + Poll::Ready(Ok(())) + } + Err(error) => { + self.done = true; + Poll::Ready(Err(error)) + } + } + } } impl Stream for Listing { @@ -114,41 +173,8 @@ impl Stream for Listing { if this.done { return Poll::Ready(None); } - if this - .options - .page_limit - .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) - { - this.done = true; - return Poll::Ready(Some(Err(Error::InvalidInput { - message: "Listing page limit must be between 1 and 2147483647".into(), - }))); - } - let pending = this - .pending - .get_or_insert_with(|| (this.fetch)(this.options.clone())); - let response = futures::ready!(pending.as_mut().poll(cx)); - this.pending = None; - match response { - Ok(response) => { - let token = response.page_token.filter(|token| !token.is_empty()); - if token - .as_ref() - .is_some_and(|token| !this.seen_tokens.insert(token.clone())) - { - this.done = true; - return Poll::Ready(Some(Err(Error::Runtime { - message: "Listing response repeated a page_token".into(), - }))); - } - this.done = token.is_none(); - this.options.page_token = token; - this.items = response.items.into_iter(); - } - Err(error) => { - this.done = true; - return Poll::Ready(Some(Err(error))); - } + if let Err(error) = futures::ready!(this.poll_fetch_page(cx)) { + return Poll::Ready(Some(Err(error))); } } } @@ -224,6 +250,36 @@ mod tests { assert!(listing.next().await.is_none()); } + #[tokio::test] + async fn next_page_preserves_empty_pages_and_remaining_items() { + let requests = Arc::new(Mutex::new(0)); + let count = requests.clone(); + let mut listing = Listing::new(ListingOptions::default(), move |options| { + *count.lock().unwrap() += 1; + async move { + let (items, page_token) = match options.page_token.as_deref() { + None => (vec![], Some("second".to_string())), + Some("second") => (vec!["a", "b"], Some("last".to_string())), + Some("last") => (vec!["c"], None), + _ => unreachable!(), + }; + Ok(ListingPage { items, page_token }) + } + }); + assert_eq!(*requests.lock().unwrap(), 0); + assert!(listing.next_page().await.unwrap().unwrap().items.is_empty()); + assert_eq!(*requests.lock().unwrap(), 1); + assert_eq!(listing.try_next().await.unwrap(), Some("a")); + assert_eq!(*requests.lock().unwrap(), 2); + let remaining = listing.next_page().await.unwrap().unwrap(); + assert_eq!(remaining.items, vec!["b"]); + assert_eq!(remaining.page_token.as_deref(), Some("last")); + assert_eq!(*requests.lock().unwrap(), 2); + assert_eq!(listing.next_page().await.unwrap().unwrap().items, vec!["c"]); + assert!(listing.next_page().await.unwrap().is_none()); + assert_eq!(*requests.lock().unwrap(), 3); + } + #[tokio::test] async fn invalid_limits_do_not_fetch() { for limit in [0, u32::MAX] { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 4a8cbf32e..8a456ab3b 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -2349,16 +2349,8 @@ mod tests { .unwrap() .is_empty() ); - let page = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(page.tables.is_empty()); - assert!(page.page_token.is_none()); + let mut listing = conn.list_tables(&namespace, ListingOptions::default().page_limit(0)); + assert!(listing.next_page().await.is_err()); } #[rstest::rstest] @@ -2393,25 +2385,17 @@ mod tests { ) .unwrap() }); - let first = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(first.tables.len(), 100); + let mut first_listing = conn.list_tables(&namespace, ListingOptions::default()); + let first = first_listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items.len(), 100); assert_eq!(first.page_token.as_deref(), Some("opaque-token")); - let second = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); + let mut second_listing = conn.list_tables( + &namespace, + ListingOptions::default().page_token(first.page_token.unwrap()), + ); + let second = second_listing.next_page().await.unwrap().unwrap(); let expected: Vec<_> = (100..130).map(|i| format!("t{i:03}")).collect(); - assert_eq!(second.tables, expected); + assert_eq!(second.items, expected); assert!(second.page_token.is_none()); } @@ -3680,14 +3664,9 @@ mod tests { assert!(table.is_ok(), "Failed to create table: {:?}", table.err()); // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .expect("Failed to list tables"); - assert_eq!(list_response.tables, vec!["test_table"]); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names, vec!["test_table"]); // Get namespace client and verify it can also list tables let namespace_client = conn.namespace_client().await.unwrap(); @@ -3754,17 +3733,12 @@ mod tests { } // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(list_response.tables.len(), 3); - assert!(list_response.tables.contains(&"table1".to_string())); - assert!(list_response.tables.contains(&"table2".to_string())); - assert!(list_response.tables.contains(&"table3".to_string())); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names.len(), 3); + assert!(names.contains(&"table1".to_string())); + assert!(names.contains(&"table2".to_string())); + assert!(names.contains(&"table3".to_string())); } } From 444f7f3bde8d18aa6ef39cba714dbe1bb448fb1d Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 23 Oct 2026 12:45:25 +0800 Subject: [PATCH 43/64] ci: refresh apt lists before installing protobuf in docs build (#4636) The docs deploy workflow installs protobuf-compiler without refreshing apt lists first. Ubuntu noble-updates replaced `protobuf 3.21.12-8.2ubuntu0.3` with `.0.4`, so stale package lists in older runner images make `apt install` fail with 404 and exit code 100. Run `apt update` before installing, matching the other workflows. --- .github/workflows/docs.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index d0ec583bf..fb3fb1d6a 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -38,6 +38,7 @@ jobs: uses: actions/checkout@v6 - name: Install dependencies needed for ubuntu run: | + sudo apt update sudo apt install -y protobuf-compiler libssl-dev rustup update && rustup default - name: Set up Python From 45c4b159cab990a749d592b42e21e030f7479192 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Thu, 22 Oct 2026 22:51:40 -0700 Subject: [PATCH 44/64] fix: cap the remote job poll interval at one second (#4629) job.wait() polled on a backoff that doubled up to 5 s, so a job's completion was seen 2.5 s late on average, as long as a small refresh itself takes. Cap the interval at 1 s; the first polls still back off from 200 ms, so a quick job is seen quickly and a long one costs one request a second. Linear: OSS-2704 https://linear.app/lancedb/issue/OSS-2704 --- rust/lancedb/src/remote/job.rs | 36 ++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 41d8971ac..7d2e70ac9 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -18,7 +18,9 @@ use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. const INITIAL_POLL_INTERVAL: Duration = Duration::from_millis(200); -const MAX_POLL_INTERVAL: Duration = Duration::from_secs(5); +/// A job's completion is seen at most this long after it happens. A 5 s cap +/// added 2.5 s on average to every short refresh, as much as the job itself. +const MAX_POLL_INTERVAL: Duration = Duration::from_secs(1); #[derive(Debug, Clone, PartialEq, Eq)] enum JobState { @@ -268,13 +270,18 @@ impl JobHandle for RemoteJob { #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + use async_trait::async_trait; use crate::Result; use crate::function::{FunctionVersion, RefreshColumnResult}; use crate::job::{Job, JobHandle, TerminalResult}; + use crate::remote::client::test_utils::client_with_handler; - use super::DescribeJobResponse; + use super::{DescribeJobResponse, RemoteJob}; const FUNCTION_JOB: &str = include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); @@ -333,6 +340,31 @@ mod tests { ); } + /// Five in-progress polls then done: 0.2 + 0.4 + 0.8 + 1 + 1 s of waiting. + /// Under the old 5 s cap the same job would have waited 10 s. + #[tokio::test(start_paused = true)] + async fn wait_polls_at_most_every_second() { + let describes = Arc::new(AtomicUsize::new(0)); + let seen = describes.clone(); + let client = client_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/jobs/describe"); + let n = seen.fetch_add(1, Ordering::SeqCst); + let body = if n < 5 { + r#"{"job_state":"IN_PROGRESS"}"#.to_string() + } else { + UNIT_JOB.to_string() + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let started = tokio::time::Instant::now(); + RemoteJob::new(client, "job-1".to_string()) + .wait() + .await + .expect("unit job completes"); + assert_eq!(describes.load(Ordering::SeqCst), 6); + assert_eq!(started.elapsed(), Duration::from_millis(3400)); + } + #[test] fn remote_wire_unknown_fields_are_forward_decodable() { let response: DescribeJobResponse = From a9535684cc1d9901c54c852f103dd0d492d61c3f Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 23 Oct 2026 20:17:38 +0800 Subject: [PATCH 45/64] chore(deps): bump ctutils to 0.4.4 for RUSTSEC-2026-0385 (#4638) Fixes RUSTSEC-2026-0385: `ctutils` < 0.4.4 generates masks with branches on secret values (`u128` mask generation / comparisons are not constant-time on some targets). - Advisory: https://rustsec.org/advisories/RUSTSEC-2026-0385 - Affected: `ctutils 0.4.2` pinned in `Cargo.lock`, reached transitively via `digest 0.11` (RustCrypto) - Change: `cargo update -p ctutils --precise 0.4.4`; single lockfile-only change (patch bump, no API impact) The `deny` CI job started failing for every PR head after the advisory was published at 06:45Z today; `main`'s last `deny` run (05:51Z) predates it and does not yet cover this advisory. --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 34ee05cb1..8bad6ca92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2230,9 +2230,9 @@ dependencies = [ [[package]] name = "ctutils" -version = "0.4.2" +version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +checksum = "12218662ccac37d9714f2191252553d67585b97df8c839cbe49e0bde278dc865" dependencies = [ "cmov", ] From b762ec221ba5d99ab368f2959f50381c589c762b Mon Sep 17 00:00:00 2001 From: LuQQiu Date: Fri, 23 Oct 2026 10:06:20 -0700 Subject: [PATCH 46/64] docs: add Enterprise distributed search reference page (#4625) Adds the Enterprise reference page at `docs/web/enterprise/distributed-search.mdx`, following the same arrangement as Enterprise architecture. The page describes distributed search and the embedded execution model, with links to shared search guides. Enterprise architecture links to it, and the official sidebar lists it immediately after Architecture in the existing Deployment group. Adds stable anchors for the RaBitQ and `approx_mode` sections linked by the Enterprise guide. Companion: https://github.com/lancedb/sophon/pull/8986 supplies the detailed Enterprise overlay at the same path. Merge this reference-page PR first so the overlay has an existing route to replace. Validation: - Combined-site assembly with both PR branches passed (309 files). - Mintlify build validation and link, anchor, and redirect checks passed. --- docs/web/docs.json | 1 + docs/web/enterprise/architecture.mdx | 3 ++ docs/web/enterprise/distributed-search.mdx | 33 ++++++++++++++++++++++ docs/web/indexing/vector-index.mdx | 2 +- docs/web/search/vector-search.mdx | 2 +- 5 files changed, 39 insertions(+), 2 deletions(-) create mode 100644 docs/web/enterprise/distributed-search.mdx diff --git a/docs/web/docs.json b/docs/web/docs.json index 9c9468a3b..54203f1a0 100644 --- a/docs/web/docs.json +++ b/docs/web/docs.json @@ -350,6 +350,7 @@ "pages": [ "enterprise/index", "enterprise/architecture", + "enterprise/distributed-search", "enterprise/deployment/index", "deployment/storage", "deployment/monitoring", diff --git a/docs/web/enterprise/architecture.mdx b/docs/web/enterprise/architecture.mdx index c47dcbd2e..8f887c142 100644 --- a/docs/web/enterprise/architecture.mdx +++ b/docs/web/enterprise/architecture.mdx @@ -18,3 +18,6 @@ run on machines that can be replaced without any of them owning the data. [Architecture](/architecture) covers this alongside the embedded form, which is the comparison most readers want first. + +See [Distributed search](/enterprise/distributed-search) for how Enterprise +distributes index execution and table-data reads. diff --git a/docs/web/enterprise/distributed-search.mdx b/docs/web/enterprise/distributed-search.mdx new file mode 100644 index 000000000..ebac960ba --- /dev/null +++ b/docs/web/enterprise/distributed-search.mdx @@ -0,0 +1,33 @@ +--- +title: "Distributed search" +sidebarTitle: "Distributed search" +description: "How LanceDB Enterprise distributes index execution and table-data reads." +icon: "network-wired" +availability: + oss: unavailable + enterprise: available + summary: >- + LanceDB Enterprise distributes index execution and table-data reads across + plan executors. Embedded LanceDB executes searches in the application process. +--- + +Distributed search spreads index execution, table-data reads, and caching across +multiple serving machines in LanceDB Enterprise. Query nodes plan queries and +combine results. Plan executors cache index structures and table data and execute +the remote work. See [Enterprise architecture](/enterprise/architecture). + +## Embedded search {#embedded-search} + +LanceDB OSS runs inside your application process. It searches indexes and reads +table data using that process's resources, including its index cache. Reading a +table from object storage does not by itself distribute query execution across +machines. + +## Search interfaces {#search-interfaces} + +The shared search guides describe the query APIs and index settings: + +- [Vector search](/search/vector-search) and [vector indexes](/indexing/vector-index). +- [Full-text search](/search/full-text-search) and [FTS indexes](/indexing/fts-index). +- [Filtering](/search/filtering) and [scalar indexes](/indexing/scalar-index). +- [Query plans](/search/explain) for inspecting execution and runtime measurements. diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index 2742a992f..b26e0c295 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -384,7 +384,7 @@ The inner workings of RaBitQ quantization are rather mathematically dense. It ge applying a uniformly random, approximately distance-preserving orthogonal transformation of the vertices of the $d$-dimensional hypercube, where $d$ is the dimensionality of the dataset. We defer the details, and an elegant theoretical error bound, to [the original paper.](https://arxiv.org/pdf/2405.12497) -##### Using RaBitQ +##### Using RaBitQ {#using-rabitq} Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, call `create_index` with `index_type="IVF_RQ"`. Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`. diff --git a/docs/web/search/vector-search.mdx b/docs/web/search/vector-search.mdx index 3519df553..cc6ae5421 100644 --- a/docs/web/search/vector-search.mdx +++ b/docs/web/search/vector-search.mdx @@ -235,7 +235,7 @@ The table below summarizes the behavior of `_distance` in search results based o See [vector indexes](/indexing/vector-index) for index configuration. The following sections cover query-time settings. -### Tuning `approx_mode` +### Tuning `approx_mode` {#tuning-approx-mode} Use `approx_mode` when you want to adjust the speed/recall trade-off for approximate vector search at query time. This setting currently applies only to RQ-quantized indexes, such as `IVF_RQ`; other index types ignore it. From 5a79c7d22da9f9729d27adc7e9d0b166767bbd4e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 24 Oct 2026 01:52:45 +0800 Subject: [PATCH 47/64] ci: run aarch64 Windows tests on main only (#4637) Run the aarch64 Windows build/test cycle only on pushes to main so PRs wait for a single Windows platform (x86_64). Main keeps the same native aarch64 coverage as before. GitHub evaluates job-level `if` before the matrix strategy is applied, so the aarch64 target is split into its own `windows-arm` job instead of being filtered inside the existing matrix. The job keeps the historical check name. --- .github/workflows/rust.yml | 43 ++++++++++++++++++++++++++++++++------ 1 file changed, 37 insertions(+), 6 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index cbfefc83d..930ae67e2 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -245,12 +245,6 @@ jobs: include: - target: x86_64-pc-windows-msvc runner: windows-2022 - # windows-11-arm is a standard runner, so it is free on public repos. - # Running natively lets the aarch64 tests actually execute -- this - # job used to cross-compile them and then skip the test step, paying - # full codegen and link cost for a compile check. - - target: aarch64-pc-windows-msvc - runner: windows-11-arm runs-on: ${{ matrix.runner }} defaults: run: @@ -279,6 +273,43 @@ jobs: # dependency graph a second time. cargo test --profile ci --features aws,remote --locked --target ${{ matrix.target }} + # Run the aarch64 Windows tests only on main so PRs wait for a single + # Windows platform. windows-11-arm is a standard runner, so it is free on + # public repos. Running natively lets the aarch64 tests actually execute -- + # this job used to cross-compile them and then skip the test step, paying + # full codegen and link cost for a compile check. + windows-arm: + # Keep the historical name so the check shows up as before on main. + name: windows (aarch64-pc-windows-msvc, windows-11-arm) + if: github.ref == 'refs/heads/main' + runs-on: windows-11-arm + defaults: + run: + working-directory: rust/lancedb + steps: + - uses: actions/checkout@v6 + - name: Set target + run: rustup target add aarch64-pc-windows-msvc + - uses: Swatinem/rust-cache@v2 + with: + # Restore everywhere, but only save from main. Per-PR saves are + # unreadable outside their own branch anyway, since GitHub scopes + # caches to the creating ref. + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Install Protoc v21.12 + run: choco install --no-progress protoc + - name: Build + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + cargo build --profile ci --features aws,remote --tests --locked --target aarch64-pc-windows-msvc + - name: Run tests + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + # `--target` has to match the build step above. Without it cargo uses + # target/ci/ rather than target//ci/ and rebuilds the entire + # dependency graph a second time. + cargo test --profile ci --features aws,remote --locked --target aarch64-pc-windows-msvc + msrv: # Check the minimum supported Rust version name: MSRV Check - Rust v${{ matrix.msrv }} From 702389654e11d1af5cd120d64866f028f41d6b90 Mon Sep 17 00:00:00 2001 From: KarthikS Date: Fri, 23 Oct 2026 15:27:32 -0400 Subject: [PATCH 48/64] fix(python): align pylance extra and preview install indexes (#4556) Fixes #4248. The preview index still contains pylance 1.0.0b6. uv's default first-index resolution can select it, or backtrack to an old LanceDB release when installing the extra. Meanwhile, the extra still permits pylance 5 even though current main bundles Lance 14. Raise the optional dependency to `pylance>=14.0.0b12`, the version already pinned in the test extra, and keep the lockfile metadata consistent. Document both Fury indexes and uv's `unsafe-best-match` strategy in the Python README and the source-owned quickstart (`docs/web/quickstart.mdx`). Explain the index trust requirement and recommend installing the extra to enforce the version floor. Add CI checks that reject older Lance majors and ensure the pylance and tests extras remain compatible. Run these checks when either dependency manifest changes, so a future Lance major bump cannot silently leave the Python minimum behind. ### Validation - Recreated the reported install on Ubuntu 26.04 / CPython 3.12.15 with uv 0.12.23: `lancedb==0.40.0b5` plus the old index selects `pylance==1.0.0b6`. Reproduced both the missing `blob_array` import and the unexpected `blob_handling` scanner argument using a real blob table. - Resolved the modified `python/pyproject.toml` with the pylance extra against both indexes: selects `pylance==14.0.0b18`, matching the current Rust pin. - Ran both documented pip and uv installations in isolated environments: `lancedb==0.41.0b7`, `pylance==14.0.0b18`; blob-array construction and a blob-bytes round trip through `to_pandas(blob_mode="bytes")` pass. Both package dependency checks pass. - New regression fails on the original dependency bound (27 rejected-version cases incorrectly accepted), then passes with the fix. Full CI-script suite: 3 tests pass. - `uv lock --check` with both indexes passes; focused Ruff checks pass; `git diff --check` passes. - Repository-wide Ruff reports 49 existing errors in untouched `docs/web-tests` files. Unrelated formatting changes were discarded. Runtime smoke checks used released native wheels; source dependency resolution was tested directly from the modified pyproject. No Rust or Python runtime code changes. --------- Co-authored-by: Xuanwo --- .../lancedb-update-lance-dependency/SKILL.md | 27 +++++++++++ .github/workflows/ci-scripts.yml | 6 +++ ci/tests/test_python_dependencies.py | 47 +++++++++++++++++++ docs/web/quickstart.mdx | 27 +++++++++++ python/README.md | 27 +++++++++++ python/pyproject.toml | 5 +- python/uv.lock | 18 +++---- 7 files changed, 146 insertions(+), 11 deletions(-) create mode 100644 ci/tests/test_python_dependencies.py diff --git a/.agents/skills/lancedb-update-lance-dependency/SKILL.md b/.agents/skills/lancedb-update-lance-dependency/SKILL.md index e97fae7ac..3cba2ae75 100644 --- a/.agents/skills/lancedb-update-lance-dependency/SKILL.md +++ b/.agents/skills/lancedb-update-lance-dependency/SKILL.md @@ -48,6 +48,23 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r This updates the Rust workspace Lance dependencies through `ci/set_lance_version.py`, updates `java/pom.xml`, refreshes Cargo metadata, and prints JSON metadata containing `branch_name`, `commit_message`, and `pr_title`. + Also synchronize the Python requirements in `python/pyproject.toml`: set the + `tests` extra's exact `pylance` pin to the matching published Python release, + and update the `pylance` extra's minimum when needed to exclude older bundled + Lance majors. The minimum must still accept the test pin. Python prerelease + spelling differs from Rust (for example, `15.0.0-beta.1` becomes `15.0.0b1`). + Confirm the matching Python release is published before changing the pin. + + Refresh and check `python/uv.lock` from the repository root: + + ```bash + uv lock --project python --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + uv lock --project python --check --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + ``` + + This strategy considers versions across indexes; use it only with trusted + indexes. Commit the Python requirements and lockfile together with the bump. + 5. Run validation: ```bash @@ -57,6 +74,16 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r Fix real diagnostics and rerun clippy until it succeeds. Do not skip warnings. + Run the Python dependency guards with Python 3.11+ and `packaging` installed: + + ```bash + python3 -m unittest discover -s ci/tests -p test_python_dependencies.py -v + ``` + + These checks reject older Lance majors and ensure the `pylance` and `tests` + extras remain compatible. The full CI-script suite uses + `python3 -m unittest discover -s ci/tests -v`. + 6. Inspect `git status --short` and `git diff` to ensure only the Lance dependency update and required compatibility fixes are present. 7. If the task only asks to prepare local changes, stop here and report the changed files and validation result. diff --git a/.github/workflows/ci-scripts.yml b/.github/workflows/ci-scripts.yml index f7f0dbb07..d45289c2d 100644 --- a/.github/workflows/ci-scripts.yml +++ b/.github/workflows/ci-scripts.yml @@ -5,6 +5,8 @@ on: branches: - main paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -12,6 +14,8 @@ on: - .github/workflows/ci-scripts.yml pull_request: paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -30,5 +34,7 @@ jobs: - uses: actions/setup-python@v6 with: python-version: "3.13" + - name: Install test dependencies + run: python -m pip install packaging - name: Run tests run: python -m unittest discover -s ci/tests -v diff --git a/ci/tests/test_python_dependencies.py b/ci/tests/test_python_dependencies.py new file mode 100644 index 000000000..97fb672c6 --- /dev/null +++ b/ci/tests/test_python_dependencies.py @@ -0,0 +1,47 @@ +import tomllib +import unittest +from pathlib import Path + +from packaging.requirements import Requirement +from packaging.version import Version + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +class PythonDependenciesTest(unittest.TestCase): + def setUp(self): + with (REPO_ROOT / "python/pyproject.toml").open("rb") as f: + self.project = tomllib.load(f)["project"] + self.pylance = Requirement(self.project["optional-dependencies"]["pylance"][0]) + + def test_pylance_extra_excludes_older_lance_majors(self): + with (REPO_ROOT / "Cargo.toml").open("rb") as f: + lance = tomllib.load(f)["workspace"]["dependencies"]["lance"] + if "version" not in lance: + self.skipTest("Local Lance checkout has no version requirement") + major = Version(lance["version"].lstrip("=")).major + for older_major in range(major): + for suffix in (".0.0", ".0.0b6", ".999.999"): + version = f"{older_major}{suffix}" + with self.subTest(version=version): + self.assertFalse( + self.pylance.specifier.contains(version, prereleases=True), + f"The pylance extra must exclude Lance {version}; " + f"the bundled Lance is major {major}", + ) + + def test_pylance_extra_accepts_tested_version(self): + requirements = [ + Requirement(dep) for dep in self.project["optional-dependencies"]["tests"] + ] + tested = next(dep for dep in requirements if dep.name == "pylance") + pin = next(spec for spec in tested.specifier if spec.operator == "==") + self.assertTrue( + self.pylance.specifier.contains(pin.version, prereleases=True), + f"The tests and pylance extras must be installable together: {tested}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/web/quickstart.mdx b/docs/web/quickstart.mdx index 788dd6585..3f466041d 100644 --- a/docs/web/quickstart.mdx +++ b/docs/web/quickstart.mdx @@ -58,6 +58,33 @@ uv add --prerelease allow --index https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra and +include the Lance preview index: + + +```bash Python (pip) icon="terminal" +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +```bash Python (uv) icon="terminal" +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + Pre-release builds receive the same level of testing as stable releases, but their availability is not guaranteed for more than 6 months after release. For real-world workloads, we recommend you use the latest stable release diff --git a/python/README.md b/python/README.md index f941c9e20..907f7c099 100644 --- a/python/README.md +++ b/python/README.md @@ -40,6 +40,33 @@ Stable releases are created about every 2 weeks. For the latest features and bug pip install --pre --extra-index-url https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra +and include the Lance preview index: + +```bash +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +With uv, also select the best matching version across the indexes: + +```bash +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + ### Threading in CPU-limited containers LanceDB uses separate pools for compute work and storage I/O. On a container with diff --git a/python/pyproject.toml b/python/pyproject.toml index ed54ae536..a0de6509c 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -45,7 +45,8 @@ repository = "https://github.com/lancedb/lancedb" [project.optional-dependencies] pylance = [ - "pylance>=5.0.0b5", + # Keep the minimum in the bundled Lance major and compatible with tests below. + "pylance>=15.0.0b1", ] # A library only needs the OpenTelemetry API; the application supplies and # configures the SDK (the actual exporter/reader). See @@ -63,7 +64,7 @@ tests = [ "polars>=0.19, <=1.32.3", "pyarrow<25", "pyarrow-stubs>=16.0", - "pylance==14.0.0b12", + "pylance==15.0.0b1", "requests>=2.31.0", "datafusion>=54,<55", "opentelemetry-sdk>=1.30.0", diff --git a/python/uv.lock b/python/uv.lock index e808f5f9e..99bcb439c 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -2006,8 +2006,8 @@ requires-dist = [ { name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" }, { name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" }, { name = "pydantic", specifier = ">=2.7.4,<3" }, - { name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" }, - { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" }, + { name = "pylance", marker = "extra == 'pylance'", specifier = ">=15.0.0b1" }, + { name = "pylance", marker = "extra == 'tests'", specifier = "==15.0.0b1" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" }, { name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" }, { name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" }, @@ -3938,7 +3938,7 @@ crypto = [ [[package]] name = "pylance" -version = "14.0.0b12" +version = "15.0.0b1" source = { registry = "https://pypi.fury.io/lance-format/" } dependencies = [ { name = "lance-namespace" }, @@ -3948,12 +3948,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" }, - { url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" }, - { url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" }, - { url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" }, - { url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" }, - { url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" }, + { url = "https://pypi.fury.io/lance-format/-/ver_23L1Nd/pylance-15.0.0b1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:19605d3daca1224cc0db079054be5d4eec403dcf286bf4617430bfc259ab0a52" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1Ho3RK/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be22b0b6f1e1d250ffaa82b449e982c817db33103c483d1005020285310623cf" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1QmyF0/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5b2761d60b1b9db112058213f1f7b2de82b08cbe65888fab08b0a56e702482e" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1in2Ph/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:294a0fedae312ed32d1b66cddbf7e69770ec9b4c82a873af60c619ed88450718" }, + { url = "https://pypi.fury.io/lance-format/-/ver_21xpLH/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ff20a16703d4d3f3bc675b25ce1871ba1d54137929dc46305bffdd07f7b35a6b" }, + { url = "https://pypi.fury.io/lance-format/-/ver_AAGzF/pylance-15.0.0b1-cp310-abi3-win_amd64.whl", hash = "sha256:3b134e24a5becfc940ca134cf78aa3bae8b80ec40f8927a1b856273ab76213f2" }, ] [[package]] From e74a7c15b5f322f8ac15b96bdfc391616518d711 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 24 Oct 2026 03:47:16 +0800 Subject: [PATCH 49/64] fix: bump tokio to address RUSTSEC-2026-0386 (#4642) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upgrade the locked Tokio version from 1.53.1 to 1.53.3 to fix [RUSTSEC-2026-0386](https://rustsec.org/advisories/RUSTSEC-2026-0386) and restore the Rust workflow's `deny` job. The diff changes only Tokio's version and checksum. Validation: - `cargo deny check advisories` — passed. - `cargo deny check advisories bans licenses sources` — passed with the repository's `deny.toml`. - `cargo check --quiet -p lancedb` — passed (dead-code warnings). - PR `deny` job — [passed](https://github.com/lancedb/lancedb/actions/runs/40000920646/job/120576455117). --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8bad6ca92..fac89c4da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10446,9 +10446,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.53.1" +version = "1.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "ce3335fa71841cda333a58d7615b03901380ecf09d59b3296d21f8bbac0dde4e" dependencies = [ "bytes", "libc", From 3321860e9fc7248d9c97c7513bc0ac0d258738a0 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Sat, 24 Oct 2026 04:22:30 +0800 Subject: [PATCH 50/64] fix(python): preserve Arrow map scalars in updates (#4448) Arrow map scalars returned by a table query cannot currently be passed back to `update`: the generic Arrow scalar conversion turns them into a list of tuples, which the SQL literal encoder cannot handle. Empty map scalars become an untyped SQL list and also fail to update a map column. Add a dedicated `MapScalar` conversion that renders its keys and values as a SQL map while preserving the existing null conversion. The regression tests exercise real synchronous and asynchronous table updates, quoted keys, nullable values, empty typed maps, and nested maps. Validation: the complete documented unit-test scope passed locally (1,515 passed, 37 skipped; unchanged baseline: 1,507 passed, 37 skipped). Full doctests passed (71 passed, 6 skipped), and repository-wide Python formatting/lint, Rust formatting/clippy, and Python 3.13 type checking passed. The locally built native wheel also passed all eight regression cases in a separate Python 3.13 environment. Full Linux/S3 integration and all platform matrices remain for CI. The SQL map constructor's existing backend constraints still apply: duplicate keys are rejected, and a nested map containing both nonempty and empty inner maps can fail SQL type unification. This change does not claim arbitrary Arrow map roundtrip support. AI assistance: implementation and local verification used Codex with gpt-6.1-sol agents; the diff and evidence were reviewed before submission. --------- Co-authored-by: Xuanwo --- python/python/lancedb/util.py | 9 ++++++ python/python/tests/test_table.py | 51 +++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index a043f0c8d..36536f499 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -452,6 +452,15 @@ def _(value: pa.Scalar): return value_to_sql(value.as_py()) +@value_to_sql.register(pa.MapScalar) +def _(value: pa.MapScalar): + if not value.is_valid: + return value_to_sql(None) + keys = value_to_sql(list(value.values.field(0))) + values = value_to_sql(list(value.values.field(1))) + return f"map({keys}, {values})" + + @value_to_sql.register(np.ndarray) def _(value: np.ndarray): return value_to_sql(value.tolist()) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index b55abfa7e..18d3c16db 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -2724,6 +2724,57 @@ def test_update_with_arrow_scalar(mem_db: DBConnection): assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]] +@pytest.mark.asyncio +@pytest.mark.parametrize("async_api", [False, True]) +@pytest.mark.parametrize("entries", [[("o'brien", 4), ("nil", None)], [], None]) +async def test_update_map_arrow_scalar(tmp_path, async_api, entries): + map_type = pa.map_(pa.string(), pa.int64()) + data = pa.table( + { + "id": [1, 2], + "attrs": pa.array([[("original", 1)], entries], type=map_type), + } + ) + if async_api: + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("maps", data) + source = await table.query().where("id = 2").to_arrow() + result = await table.update({"attrs": source["attrs"][0]}, where="id = 1") + actual = await table.query().to_arrow() + else: + db = lancedb.connect(tmp_path) + table = db.create_table("maps", data) + source = table.search().where("id = 2").to_arrow() + result = table.update(values={"attrs": source["attrs"][0]}, where="id = 1") + actual = table.to_arrow() + assert result.rows_updated == 1 + assert actual.schema.field("attrs").type == map_type + assert actual.sort_by("id")["attrs"].to_pylist() == [entries, entries] + + +@pytest.mark.parametrize( + "map_type, entries", + [ + (pa.map_(pa.int64(), pa.string()), []), + ( + pa.map_(pa.string(), pa.map_(pa.string(), pa.int64())), + [("outer", [("inner", 3)]), ("null", None)], + ), + ], +) +def test_update_typed_map_arrow_scalar(mem_db: DBConnection, map_type, entries): + table = mem_db.create_table( + "maps", + pa.table({"id": [1, 2], "attrs": pa.array([None, entries], type=map_type)}), + ) + scalar = table.search().where("id = 2").to_arrow()["attrs"][0] + result = table.update(values={"attrs": scalar}, where="id = 1") + actual = table.to_arrow().sort_by("id")["attrs"] + assert result.rows_updated == 1 + assert actual.type == map_type + assert actual.to_pylist() == [entries, entries] + + @pytest.mark.parametrize("use_arrow_scalar", [False, True]) @pytest.mark.parametrize( "decimal_type, value", From 9da4658d57cbdcd037827f12effd9a4dc869a095 Mon Sep 17 00:00:00 2001 From: Joaquin Hui Gomez <132194176+joaquinhuigomez@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:10:07 +0100 Subject: [PATCH 51/64] fix(python): accept with_row_id in take_offsets/take_row_ids on every Table implementation (#4374) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Table.take_offsets` and `Table.take_row_ids` declare a keyword-only `with_row_id: bool = False`, but none of the three implementations accept it — `LanceTable`, `AsyncTable` and `RemoteTable` all raise `TypeError: ... got an unexpected keyword argument 'with_row_id'` on the documented call. The builders they return already support it (`LanceTakeQueryBuilder.with_row_id()` / `AsyncTakeQuery.with_row_id()`), so this is a signature that never reached the implementations. The base docstrings also don't mention the parameter and still say `take_offsets` returns a `pa.RecordBatch` (it returns a `LanceTakeQueryBuilder`), so those are corrected too. Each implementation now accepts the kwarg and applies `.with_row_id()` to the builder it returns. Tests: sync and async assert `_rowid` appears in the result columns; the remote test uses the suite's request-capturing fixture and asserts `with_row_id: true` reaches the wire (and is absent by default). All three fail on `main` with the `TypeError`. ruff format/check clean. A related mismatch I left alone: the deprecated `LanceTable.create_scalar_index`/`create_fts_index` also drop a `wait_timeout` the base declares — deprecated since 0.25, so not worth touching here. Co-authored-by: Xuanwo --- python/python/lancedb/remote/table.py | 18 +++++++-- python/python/lancedb/table.py | 58 ++++++++++++++++++++------- python/python/tests/test_query.py | 27 +++++++++++++ python/python/tests/test_remote_db.py | 17 ++++++++ 4 files changed, 102 insertions(+), 18 deletions(-) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 86e8fb8a0..0744410b3 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1154,11 +1154,21 @@ class RemoteTable(Table): """ return LOOP.run(self._table.uri()) - def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + if with_row_id: + builder.with_row_id() + return builder - def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + if with_row_id: + builder.with_row_id() + return builder def uses_v2_manifest_paths(self) -> bool: return LOOP.run(self._table.uses_v2_manifest_paths()) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 5eea55760..b9557f219 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2103,11 +2103,14 @@ class Table(ABC): ---------- offsets: list[int] The offsets to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - pa.RecordBatch - A record batch containing the rows at the given offsets. + LanceTakeQueryBuilder + A query builder that can be executed to get the rows at the given + offsets. """ def __getitems__(self, offsets: list[int]) -> pa.RecordBatch: @@ -2188,11 +2191,13 @@ class Table(ABC): ---------- row_ids: list[int] The row ids to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - AsyncTakeQuery - A query object that can be executed to get the rows. + LanceTakeQueryBuilder + A query builder that can be executed to get the rows. """ @abstractmethod @@ -3182,11 +3187,21 @@ class LanceTable(Table): """Get the current version of the table""" return LOOP.run(self._table.version()) - def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + if with_row_id: + builder.with_row_id() + return builder - def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + if with_row_id: + builder.with_row_id() + return builder def blob_columns(self) -> list[str]: return LOOP.run(self._table.blob_columns()) @@ -7325,7 +7340,9 @@ class AsyncTable: """ await self._inner.restore(version) - def take_offsets(self, offsets: list[int]) -> AsyncTakeQuery: + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> AsyncTakeQuery: """ Take a list of offsets from the table. @@ -7343,15 +7360,23 @@ class AsyncTable: ---------- offsets: list[int] The offsets to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - pa.RecordBatch - A record batch containing the rows at the given offsets. + AsyncTakeQuery + A query object that can be executed to get the rows at the given + offsets. """ - return AsyncTakeQuery(self._inner.take_offsets(offsets), self) + query = AsyncTakeQuery(self._inner.take_offsets(offsets), self) + if with_row_id: + query.with_row_id() + return query - def take_row_ids(self, row_ids: list[int]) -> AsyncTakeQuery: + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> AsyncTakeQuery: """ Take a list of row ids from the table. @@ -7372,13 +7397,18 @@ class AsyncTable: ---------- row_ids: list[int] The row ids to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- AsyncTakeQuery A query object that can be executed to get the rows. """ - return AsyncTakeQuery(self._inner.take_row_ids(row_ids), self) + query = AsyncTakeQuery(self._inner.take_row_ids(row_ids), self) + if with_row_id: + query.with_row_id() + return query async def blob_columns(self) -> list[str]: return await self._inner.blob_columns() diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 58ea7deb5..cef985aa6 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -2324,6 +2324,33 @@ def test_take_queries(tmp_path): ] +def test_take_queries_with_row_id(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("test", pa.table({"idx": range(100)})) + + by_row_id = table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_row_id.column_names + by_offset = table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_offset.column_names + + # The column is opt-in and must stay out of the default projection. + assert "_rowid" not in table.take_row_ids([5, 2, 17]).to_arrow().column_names + + +@pytest.mark.asyncio +async def test_take_queries_with_row_id_async(tmp_path): + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("test", pa.table({"idx": range(100)})) + + by_row_id = await table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_row_id.column_names + by_offset = await table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_offset.column_names + + default = await table.take_row_ids([5, 2, 17]).to_arrow() + assert "_rowid" not in default.column_names + + def test_take_queries_to_batches(tmp_path): # Regression test for the sync take-query path: `to_batches` previously # raised ``AttributeError: 'AsyncTakeQuery' object has no attribute diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index a45e1aaef..9655f5807 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2225,6 +2225,23 @@ def test_head(): assert data == pa.table({"id": [1, 2, 3]}) +def test_remote_take_row_ids_with_row_id(): + bodies = [] + + def handler(body): + bodies.append(body) + return pa.table({"id": [1, 2], "_rowid": pa.array([0, 1], pa.uint64())}) + + with query_test_table(handler) as table: + result = table.take_row_ids([0, 1], with_row_id=True).to_arrow() + assert "_rowid" in result.column_names + assert bodies[-1]["with_row_id"] is True + + with query_test_table(handler) as table: + table.take_row_ids([0, 1]).to_arrow() + assert "with_row_id" not in bodies[-1] + + def test_query_sync_minimal(): def handler(body): assert body == { From 27b65064d0166feabe1c307eba07f66ca2601857 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:37:07 -0700 Subject: [PATCH 52/64] fix(python): support custom blob thresholds without pylance (#4634) --- python/python/lancedb/schema.py | 62 ++++++++++------ python/python/tests/test_blob.py | 37 ++++++++++ python/python/tests/test_table.py | 118 ++++++++++++++++++++++++++++++ 3 files changed, 193 insertions(+), 24 deletions(-) diff --git a/python/python/lancedb/schema.py b/python/python/lancedb/schema.py index 5edda5610..8bec17058 100644 --- a/python/python/lancedb/schema.py +++ b/python/python/lancedb/schema.py @@ -4,6 +4,7 @@ """Schema helpers for Lance blob columns.""" +import ctypes import importlib from typing import TYPE_CHECKING, Optional @@ -16,6 +17,7 @@ if TYPE_CHECKING: _BLOB_EXTENSION_NAME = "lance.blob.v2" _BLOB_V1_KEY = "lance-encoding:blob" _ARROW_EXT_NAME_KEY = "ARROW:extension:name" +_MAX_RUST_USIZE = ctypes.c_size_t(-1).value _BLOB_V2_STORAGE_TYPE = pa.struct( [ pa.field("data", pa.large_binary(), nullable=True), @@ -209,32 +211,44 @@ def blob( Maximum size in bytes of a single packed blob sidecar file before a new one is started. Must be positive. - The thresholds are forwarded to ``lance.blob_field``, which validates - them and stores them as field metadata, so pylance is required whenever - any threshold is set. When all thresholds are omitted the behavior is - unchanged. + Thresholds are validated and stored as Arrow field metadata, with the + same format as ``lance.blob_field``. They do not require pylance. """ - if ( - inline_size_threshold is None - and dedicated_size_threshold is None - and pack_file_size_threshold is None + metadata = {} + for threshold, value, allow_zero, key in ( + ( + "inline_size_threshold", + inline_size_threshold, + True, + "lance-encoding:blob-inline-size-threshold", + ), + ( + "dedicated_size_threshold", + dedicated_size_threshold, + False, + "lance-encoding:blob-dedicated-size-threshold", + ), + ( + "pack_file_size_threshold", + pack_file_size_threshold, + False, + "lance-encoding:blob-pack-file-size-threshold", + ), ): - blob_type = _resolve_blob_type() - return pa.field(name, blob_type(), nullable=nullable) - try: - lance = importlib.import_module("lance") - except ModuleNotFoundError as err: - raise ValueError( - "Setting blob tier thresholds requires pylance to be installed" - ) from err - _resolve_blob_type() - return lance.blob_field( - name, - nullable=nullable, - inline_size_threshold=inline_size_threshold, - dedicated_size_threshold=dedicated_size_threshold, - pack_file_size_threshold=pack_file_size_threshold, - ) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{threshold} must be an int, got {type(value).__name__}") + if value < 0 or (value == 0 and not allow_zero): + bound = "non-negative" if allow_zero else "positive" + raise ValueError(f"{threshold} must be {bound}") + if value > _MAX_RUST_USIZE: + raise OverflowError(f"{threshold} must fit in a Rust usize") + metadata[key] = str(value) + + blob_type = _resolve_blob_type() + field = pa.field(name, blob_type(), nullable=nullable) + return field.with_metadata(metadata) if metadata else field def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType: diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index ec1b8b27e..c2dc2a265 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import ctypes import io import subprocess import sys @@ -14,6 +15,7 @@ import pytest from lance.blob import Blob, BlobType as LanceBlobType import lancedb +from lancedb import schema as schema_module from lancedb._blob import ( _set_blob_column, blob_v2_projection_sources, @@ -139,6 +141,41 @@ def test_blob_without_thresholds_has_no_threshold_metadata(): assert not any(b"threshold" in key for key in metadata) +@pytest.mark.parametrize("nullable", [False, True]) +@pytest.mark.parametrize( + "thresholds", + [ + {}, + {"inline_size_threshold": 0}, + {"inline_size_threshold": 8192}, + {"dedicated_size_threshold": 1 << 20}, + {"pack_file_size_threshold": 1 << 26}, + { + "inline_size_threshold": 8192, + "dedicated_size_threshold": 1 << 20, + "pack_file_size_threshold": 1 << 26, + }, + { + "inline_size_threshold": ctypes.c_size_t(-1).value, + "dedicated_size_threshold": ctypes.c_size_t(-1).value, + "pack_file_size_threshold": ctypes.c_size_t(-1).value, + }, + ], +) +def test_blob_thresholds_match_pylance(monkeypatch, nullable, thresholds): + expected = lance.blob_field("payload", nullable=nullable, **thresholds) + actual = lancedb.blob("payload", nullable=nullable, **thresholds) + assert type(actual.type) is LanceBlobType + assert actual.equals(expected, check_metadata=True) + + # Both extension classes must encode the same Arrow field on the wire. + monkeypatch.setattr( + schema_module, "_resolve_blob_type", lambda: schema_module._FallbackBlobType + ) + fallback = lancedb.blob("payload", nullable=nullable, **thresholds) + assert pa.schema([fallback]).serialize() == pa.schema([expected]).serialize() + + def test_blob_thresholds_are_validated(): with pytest.raises(ValueError, match="inline_size_threshold must be non-negative"): lancedb.blob("payload", inline_size_threshold=-1) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 18d3c16db..6ccba606f 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -59,6 +59,124 @@ def _blob_v2_table(db: DBConnection, name: str): return table +@pytest.mark.parametrize( + "threshold", + [ + "inline_size_threshold", + "dedicated_size_threshold", + "pack_file_size_threshold", + ], +) +@pytest.mark.parametrize( + "value,error", + [ + (-1, ValueError), + (True, TypeError), + (False, TypeError), + (1.5, TypeError), + ("8192", TypeError), + (ctypes.c_size_t(-1).value + 1, OverflowError), + ], +) +def test_blob_thresholds_reject_invalid_values(threshold, value, error): + # Also runs in CI with pylance actually uninstalled. + with pytest.raises(error, match=f"{threshold} must"): + lancedb.blob("blob", **{threshold: value}) + + +@pytest.mark.parametrize( + "threshold", ["dedicated_size_threshold", "pack_file_size_threshold"] +) +def test_blob_thresholds_reject_zero_for_positive_thresholds(threshold): + with pytest.raises(ValueError, match=f"{threshold} must be positive"): + lancedb.blob("blob", **{threshold: 0}) + + +def test_blob_thresholds_accept_integer_boundaries(): + maximum = ctypes.c_size_t(-1).value + field = lancedb.blob( + "blob", + nullable=False, + inline_size_threshold=0, + dedicated_size_threshold=maximum, + pack_file_size_threshold=maximum, + ) + assert not field.nullable + assert field.metadata == { + b"lance-encoding:blob-inline-size-threshold": b"0", + b"lance-encoding:blob-dedicated-size-threshold": str(maximum).encode(), + b"lance-encoding:blob-pack-file-size-threshold": str(maximum).encode(), + } + + +def test_blob_thresholds_control_writes_after_reopen(tmp_path): + # Keep this in test_table.py: CI runs it without pylance or pandas installed. + field = lancedb.blob( + "blob", + inline_size_threshold=4, + dedicated_size_threshold=12, + pack_file_size_threshold=8, + ) + metadata = { + b"lance-encoding:blob-inline-size-threshold": b"4", + b"lance-encoding:blob-dedicated-size-threshold": b"12", + b"lance-encoding:blob-pack-file-size-threshold": b"8", + } + assert field.metadata == metadata + schema = pa.schema([pa.field("id", pa.int64()), field]) + db = lancedb.connect(tmp_path) + table = db.create_table("blobs", schema=schema) + # Every payload fits in the default inline tier; custom thresholds must + # instead send the larger values to packed and dedicated storage. + rows = [ + {"id": 0, "blob": b"a" * 2}, + {"id": 1, "blob": b"b" * 8}, + {"id": 2, "blob": b"c" * 16}, + {"id": 3, "blob": b"d" * 8}, + ] + + def check_writes(table, expected): + assert table.schema.field("blob").metadata.items() >= metadata.items() + hits = table.search().with_row_id(True).limit(10).to_arrow() + actual = {row["id"]: row for row in hits.to_pylist()} + assert actual.keys() == {row["id"] for row in expected} + for row in expected: + size = len(row["blob"]) + descriptor = actual[row["id"]]["blob"] + assert descriptor["size"] == size + assert descriptor["kind"] == (2 if size > 12 else 1 if size > 4 else 0) + row_ids = [actual[row["id"]]["_rowid"] for row in expected] + assert table.fetch_blobs("blob", row_ids).to_pylist() == [ + row["blob"] for row in expected + ] + # Two packed values reach the configured sidecar limit independently. + packed = [ + actual[row["id"]]["blob"] for row in expected if len(row["blob"]) == 8 + ] + assert len({desc["blob_id"] for desc in packed}) == len(packed) + + table.add(rows) + check_writes(table, rows) + table = lancedb.connect(tmp_path).open_table("blobs") + check_writes(table, rows) + updates = [ + {"id": 0, "blob": b"e" * 16}, + {"id": 1, "blob": b"f" * 2}, + {"id": 2, "blob": b"g" * 8}, + {"id": 3, "blob": b"h" * 8}, + {"id": 4, "blob": b"i" * 16}, + ] + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(updates) + ) + assert result.num_updated_rows == 4 + assert result.num_inserted_rows == 1 + check_writes(lancedb.connect(tmp_path).open_table("blobs"), updates) + + async def _blob_v2_table_async(db: AsyncConnection, name: str): schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")]) table = await db.create_table(name, schema=schema) From 8cf115b081d6a4b48db13a11c8fcbde8deadb750 Mon Sep 17 00:00:00 2001 From: alexcastgonz <82875882+pcbeingused333@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:25:44 -0400 Subject: [PATCH 53/64] fix(python): replace the retired gpt-4-turbo-preview default in OpenaiReranker (#4613) OpenAI shuts down `gpt-4-turbo` and its snapshots on 2026-10-23 (https://developers.openai.com/api/docs/deprecations), so `OpenaiReranker` fails for every user who relies on the default `gpt-4-turbo-preview` after that date. This changes the default to `gpt-5.6-sol`, the replacement OpenAI names for `gpt-4-turbo` on that page, in the class, its docstring and the integration docs page. `_rerank` also stops sending `temperature=0`: the GPT-5 models only accept the default temperature, so keeping it would fail the new default on its first call. The reranker still asks for `json_object` output and the rest of the request is unchanged. The new test calls `_rerank` with a mocked client and checks that no `temperature` is sent. It fails on main and passes with the change. No request was made to OpenAI. This change was made with Claude Code. Co-authored-by: Xuanwo --- docs/web/integrations/reranking/openai.mdx | 2 +- python/python/lancedb/rerankers/openai.py | 5 ++--- python/python/tests/test_rerankers.py | 19 +++++++++++++++++++ 3 files changed, 22 insertions(+), 4 deletions(-) diff --git a/docs/web/integrations/reranking/openai.mdx b/docs/web/integrations/reranking/openai.mdx index 6ceb8a030..71fe63117 100644 --- a/docs/web/integrations/reranking/openai.mdx +++ b/docs/web/integrations/reranking/openai.mdx @@ -26,7 +26,7 @@ credentials from environment variables or a secrets manager. ## Accepted Arguments {#accepted-arguments} | Argument | Type | Default | Description | | --- | --- | --- | --- | -| `model_name` | `str` | `"gpt-4-turbo-preview"` | The name of the reranker model to use.| +| `model_name` | `str` | `"gpt-5.6-sol"` | The name of the reranker model to use.| | `column` | `str` | `"text"` | The name of the column to use as input to the cross encoder model. | | `return_score` | `str` | `"relevance"` | Options are "relevance" or "all". The type of score to return. If "relevance", will return only the `_relevance_score. If "all" is supported, will return relevance score along with the vector and/or fts scores depending on query type. | | `api_key` | `str` | `None` | The API key to use. If None, will use the OPENAI_API_KEY environment variable. diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index da0394d1c..49b067f79 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -21,7 +21,7 @@ class OpenaiReranker(Reranker): Parameters ---------- - model_name : str, default "gpt-4-turbo-preview" + model_name : str, default "gpt-5.6-sol" The name of the cross encoder model to use. column : str, default "text" The name of the column to use as input to the cross encoder model. @@ -33,7 +33,7 @@ class OpenaiReranker(Reranker): def __init__( self, - model_name: str = "gpt-4-turbo-preview", + model_name: str = "gpt-5.6-sol", column: str = "text", return_score="relevance", api_key: Optional[str] = None, @@ -54,7 +54,6 @@ class OpenaiReranker(Reranker): response = self._client.chat.completions.create( model=self.model_name, response_format={"type": "json_object"}, - temperature=0, messages=[ { "role": "system", diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 743669367..dd6cb4ab5 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import os +from unittest.mock import MagicMock import random import lancedb @@ -547,6 +548,24 @@ def test_openai_reranker(tmp_path): _run_test_reranker(reranker, table, "single player experience", None, schema) +def test_openai_reranker_does_not_send_temperature(): + reranker = OpenaiReranker(api_key="test-key") + create = MagicMock() + create.return_value.choices = [ + MagicMock( + message=MagicMock( + content='{"documents": [{"content": "a", "relevance_score": 0.5}]}' + ) + ) + ] + reranker.__dict__["_client"] = MagicMock() + reranker.__dict__["_client"].chat.completions.create = create + + reranker._rerank(pa.table({"text": ["a"]}), "query") + + assert "temperature" not in create.call_args.kwargs + + @pytest.mark.skipif( os.environ.get("JINA_API_KEY") is None, reason="JINA_API_KEY not set" ) From 0dd78656b966233569d2f44a2d8833dffa9003e1 Mon Sep 17 00:00:00 2001 From: Kerman Lau <107780822+WitBlue6@users.noreply.github.com> Date: Sat, 24 Oct 2026 07:13:09 +0800 Subject: [PATCH 54/64] fix(python): honor explicit OpenAI reranker API keys (#4515) Fixes #4514 ## Summary Honor an explicitly supplied `OpenaiReranker(api_key=...)` before falling back to `OPENAI_API_KEY`, matching the constructor documentation. Previously the environment key always won when both were set, so an instance configured for a different account still used the process-wide account. ## Tests Added four cases using real OpenAI client construction without sending requests: explicit/environment conflict, explicit key alone, environment fallback, and missing credentials. The conflict case fails before the change; all four pass afterward. Clients are closed after each assertion. - Targeted source-module tests: 4 passed. - Ruff check and format check: passed for both changed files. - Full native source build was attempted but failed because the environment has no Rust toolchain; bootstrap then attempted a cache path outside the writable workspace. The targeted tests load the current Python reranker module against installed dependencies and the released LanceDB package; no database/native behavior is exercised. Full native/integration suite not run. This change and its tests were prepared with an AI coding agent. The diff was reviewed and the failure/pass behavior verified locally. Co-authored-by: Xuanwo --- python/python/lancedb/rerankers/openai.py | 2 +- .../tests/test_openai_reranker_credentials.py | 38 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) create mode 100644 python/python/tests/test_openai_reranker_credentials.py diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index 49b067f79..75efbbfc7 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -131,4 +131,4 @@ class OpenaiReranker(Reranker): "OPENAI_API_KEY not set. Either set it in your environment or \ pass it as `api_key` argument to the CohereReranker." ) - return openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY") or self.api_key) + return openai.OpenAI(api_key=self.api_key or os.environ.get("OPENAI_API_KEY")) diff --git a/python/python/tests/test_openai_reranker_credentials.py b/python/python/tests/test_openai_reranker_credentials.py new file mode 100644 index 000000000..aa11239b2 --- /dev/null +++ b/python/python/tests/test_openai_reranker_credentials.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pytest + +from lancedb.rerankers import OpenaiReranker + + +@pytest.mark.parametrize( + "explicit_key, environment_key, expected_key", + [ + ("explicit-test-key", "environment-test-key", "explicit-test-key"), + ("explicit-test-key", None, "explicit-test-key"), + (None, "environment-test-key", "environment-test-key"), + ], +) +def test_openai_reranker_api_key_precedence( + monkeypatch, explicit_key, environment_key, expected_key +): + pytest.importorskip("openai") + if environment_key is None: + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + else: + monkeypatch.setenv("OPENAI_API_KEY", environment_key) + + reranker = OpenaiReranker(api_key=explicit_key) + client = reranker._client + try: + assert client.api_key == expected_key + finally: + client.close() + + +def test_openai_reranker_requires_api_key(monkeypatch): + pytest.importorskip("openai") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(ValueError, match="OPENAI_API_KEY not set"): + OpenaiReranker()._client From dd4795028b390302f8edad3b1cdef9b2ed3d9b9f Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 08:25:16 +0800 Subject: [PATCH 55/64] fix(python): align sync remote methods with async behavior (#4411) Sync remote table and connection methods inherited bare base-class stubs, so unsupported conversion and bulk-drop operations raised `NotImplementedError` with an empty message and differed from the async API. Delegate `to_lance()`, `drop_all_tables()`, and `drop_database()` to the async implementation. Sync and async remote Lance conversion share #4306's native-table guard and reject before requesting a URI or opening client storage. Conversion documentation directs remote reads through server queries. Add an explicit Polars error, suggest query conversions in Arrow/Pandas/Polars errors, and give the remaining base-class errors an operation name. Preserve upstream `blob_mode` validation for remote Arrow and Pandas conversion. The shared remote bulk-drop error suggests dropping tables individually. Regression tests cover sync/async rejection parity, accessible and missing advertised locations, pinned versions, and bulk-drop rejection without requests. Sync blob query tests verify that descriptions come from the server and unsupported modes cannot open client storage. Keep the remote closed-blob-handle assertion consistent with the upstream `ValueError` contract. Validation: - All 30 focused regression cases passed. - 222 Python tests passed with rebuilt bindings in the locked uv environment: the full remote and error suites, plus selected local conversion, blob, and deletion tests. - Python lint and formatting, Rust formatting, `cargo check --quiet --features remote --tests --examples`, and `cargo clippy --quiet --features remote --tests --examples` passed. - Root Ruff checking reports 49 existing errors in upstream `docs/web-tests`; those files are unchanged relative to `main`. Fixes #4392 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/db.py | 31 +++-- python/python/lancedb/remote/db.py | 14 +++ python/python/lancedb/remote/table.py | 29 ++++- python/python/lancedb/table.py | 56 +++++---- python/python/tests/test_errors.py | 21 ++++ python/python/tests/test_remote_db.py | 157 ++++++++++++++++++++++++-- rust/lancedb/src/remote/db.rs | 4 +- 7 files changed, 271 insertions(+), 41 deletions(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 90dc1cccf..74acc28ba 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -548,7 +548,9 @@ class DBConnection(EnforceOverrides): LanceTable(name='table4', ...) """ - raise NotImplementedError + raise NotImplementedError( + "create_table is not supported for this connection type" + ) def __getitem__(self, name: str) -> LanceTable: return self.open_table(name) @@ -603,7 +605,9 @@ class DBConnection(EnforceOverrides): ------- A LanceTable object representing the table. """ - raise NotImplementedError + raise NotImplementedError( + "open_table is not supported for this connection type" + ) def create_materialized_view( self, @@ -743,13 +747,17 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_table is not supported for this connection type" + ) def drop_table_async( self, name: str, namespace_path: Optional[List[str]] = None ) -> Job: """Start dropping a table and return its cleanup job.""" - raise NotImplementedError + raise NotImplementedError( + "drop_table_async is not supported for this connection type" + ) def rename_table( self, @@ -777,19 +785,26 @@ class DBConnection(EnforceOverrides): cur_namespace_path = [] if new_namespace_path is None: new_namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "rename_table is not supported for this connection type" + ) def drop_database(self): """ Drop database This is the same thing as dropping all the tables """ - raise NotImplementedError + raise NotImplementedError( + "drop_database is not supported for this connection type" + ) def drop_all_tables(self, namespace_path: Optional[List[str]] = None): """ Drop all tables from the database + Remote connections do not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + Parameters ---------- namespace_path: List[str], optional @@ -798,7 +813,9 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_all_tables is not supported for this connection type" + ) @property def uri(self) -> str: diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 08125cc42..5e9288471 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -841,6 +841,20 @@ class RemoteDBConnection(DBConnection): job = LOOP.run(self._conn.drop_table_async(name, namespace_path=namespace_path)) return Job(job if isinstance(job, AsyncJob) else AsyncJob(job)) + @override + def drop_all_tables(self, namespace_path: Optional[List[str]] = None): + """Drop all tables from the database. + + The remote API does not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + """ + LOOP.run(self._conn.drop_all_tables(namespace_path=namespace_path)) + + @override + def drop_database(self): + """Deprecated alias for ``drop_all_tables()``.""" + LOOP.run(self._conn.drop_database()) + @override def rename_table( self, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 0744410b3..53591182e 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -7,6 +7,7 @@ import logging from functools import cached_property import os from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -82,6 +83,10 @@ from ..table import ( from ..types import BaseTokenizerType from ..util import infer_vector_column_name +if TYPE_CHECKING: + import lance + import polars as pl + class RemoteTable(Table): def __init__( @@ -243,12 +248,32 @@ class RemoteTable(Table): def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """to_arrow() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_arrow() is not supported for remote tables.") + raise NotImplementedError( + "to_arrow() is not supported for remote tables. " + "Use table.search().to_arrow() instead." + ) def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs): """to_pandas() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_pandas() is not supported for remote tables.") + raise NotImplementedError( + "to_pandas() is not supported for remote tables. " + "Use table.search().to_pandas() instead." + ) + + def to_lance(self, **kwargs) -> "lance.LanceDataset": + """Direct Lance dataset access is not supported for remote tables. + + Use ``table.search().to_arrow()`` to read query results through the server. + """ + return LOOP.run(self._table.to_lance(**kwargs)) + + def to_polars(self, **kwargs) -> "pl.LazyFrame": + """Use ``table.search().to_polars()`` to read remote query results.""" + raise NotImplementedError( + "to_polars() is not supported for remote tables. " + "Use table.search().to_polars() instead." + ) def checkout(self, version: Union[int, str]): result = LOOP.run(self._table.checkout(version)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index b9557f219..c08c6390e 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1356,13 +1356,13 @@ class Table(ABC): @abstractmethod def name(self) -> str: """The name of this Table""" - raise NotImplementedError + raise NotImplementedError("name is not supported for this table type") @property @abstractmethod def version(self) -> int: """The version of this Table""" - raise NotImplementedError + raise NotImplementedError("version is not supported for this table type") @property @abstractmethod @@ -1371,7 +1371,7 @@ class Table(ABC): of this Table """ - raise NotImplementedError + raise NotImplementedError("schema is not supported for this table type") @property @abstractmethod @@ -1400,7 +1400,7 @@ class Table(ABC): tags = table.tags.list() """ - raise NotImplementedError + raise NotImplementedError("tags is not supported for this table type") @property def branches(self) -> "Branches": @@ -1409,11 +1409,11 @@ class Table(ABC): Branches are isolated, writable lines of history forked from another branch (or version). Writes on a branch do not affect ``main``. """ - raise NotImplementedError + raise NotImplementedError("branches is not supported for this table type") def current_branch(self) -> Optional[str]: """The branch this table handle is scoped to, or ``None`` for ``main``.""" - raise NotImplementedError + raise NotImplementedError("current_branch is not supported for this table type") def __len__(self) -> int: """The number of rows in this Table""" @@ -1436,7 +1436,7 @@ class Table(ABC): filter: str, optional A SQL where clause to filter the rows to count. """ - raise NotImplementedError + raise NotImplementedError("count_rows is not supported for this table type") def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs) -> "pandas.DataFrame": """Return the table as a pandas DataFrame. @@ -1471,20 +1471,26 @@ class Table(ABC): ------- pa.Table """ - raise NotImplementedError + raise NotImplementedError("to_arrow is not supported for this table type") def to_lance(self, **kwargs) -> lance.LanceDataset: - """Return the table as a lance.LanceDataset. + """Return a local table as a lance.LanceDataset. + + Remote tables cannot be opened as Lance datasets by the client. Use + ``table.search().to_arrow()`` to read them through the server. Returns ------- lance.LanceDataset """ - raise NotImplementedError + raise NotImplementedError("to_lance is not supported for this table type") def to_polars(self, **kwargs) -> "pl.LazyFrame": """Return the table as a Polars LazyFrame. + For remote tables, use ``table.search().to_polars()`` to convert query + results to a Polars DataFrame. + Note ---- The Polars streaming engine is not supported because it does not currently @@ -1495,7 +1501,7 @@ class Table(ABC): ------- polars.LazyFrame """ - raise NotImplementedError + raise NotImplementedError("to_polars is not supported for this table type") # New unified API overload @overload @@ -1600,7 +1606,7 @@ class Table(ABC): ... "l2", vector_column_name="vector" ... ) """ - raise NotImplementedError + raise NotImplementedError("create_index is not supported for this table type") def create_index_async( self, @@ -1618,7 +1624,9 @@ class Table(ABC): be complete when returned; callers must not assume the index exists until :meth:`Job.wait` returns. """ - raise NotImplementedError + raise NotImplementedError( + "create_index_async is not supported for this table type" + ) def drop_index(self, name: str) -> None: """ @@ -1638,7 +1646,7 @@ class Table(ABC): Use [list_indices][lancedb.table.Table.list_indices] to find the names of the indices. """ - raise NotImplementedError + raise NotImplementedError("drop_index is not supported for this table type") def wait_for_index( self, index_names: Iterable[str], timeout: timedelta = timedelta(seconds=300) @@ -1655,14 +1663,14 @@ class Table(ABC): timeout: timedelta Timeout to wait for asynchronous indexing. The default is 5 minutes. """ - raise NotImplementedError + raise NotImplementedError("wait_for_index is not supported for this table type") @abstractmethod def stats(self) -> TableStatistics: """ Retrieve table and fragment statistics. """ - raise NotImplementedError + raise NotImplementedError("stats is not supported for this table type") @abstractmethod def create_scalar_index( @@ -1727,7 +1735,9 @@ class Table(ABC): ``my_col = 0 OR not_indexed = 1`` will not be able to use any scalar index on ``my_col``. """ - raise NotImplementedError + raise NotImplementedError( + "create_scalar_index is not supported for this table type" + ) def create_fts_index( self, @@ -1853,7 +1863,9 @@ class Table(ABC): ``LANCE_LANGUAGE_MODEL_HOME`` to override the default platform data directory under ``lance/language_models``. """ - raise NotImplementedError + raise NotImplementedError( + "create_fts_index is not supported for this table type" + ) @abstractmethod def add( @@ -1927,7 +1939,7 @@ class Table(ABC): AddResult An object containing the new version number of the table after adding data. """ - raise NotImplementedError + raise NotImplementedError("add is not supported for this table type") def merge_insert(self, on: Union[str, Iterable[str]]) -> LanceMergeInsertBuilder: """ @@ -2079,7 +2091,7 @@ class Table(ABC): - and also the "_distance" column which is the distance between the query vector and the returned vector. """ - raise NotImplementedError + raise NotImplementedError("search is not supported for this table type") @abstractmethod def take_offsets( @@ -2365,7 +2377,7 @@ class Table(ABC): x vector 0 3 [5.0, 6.0] """ - raise NotImplementedError + raise NotImplementedError("delete is not supported for this table type") @abstractmethod def update( @@ -2432,7 +2444,7 @@ class Table(ABC): 1 4 [5.0, 6.0] 2 3 [10.0, 10.0] """ - raise NotImplementedError + raise NotImplementedError("update is not supported for this table type") @abstractmethod def cleanup_old_versions( diff --git a/python/python/tests/test_errors.py b/python/python/tests/test_errors.py index 338b6a0d7..0e362e5d8 100644 --- a/python/python/tests/test_errors.py +++ b/python/python/tests/test_errors.py @@ -3,7 +3,28 @@ import pickle +import pytest + +from lancedb.db import DBConnection from lancedb.remote.errors import HttpError, LanceDBClientError, RetryError +from lancedb.table import Table + + +@pytest.mark.parametrize( + "base, method", + [ + (Table, "to_lance"), + (Table, "to_polars"), + (DBConnection, "drop_all_tables"), + (DBConnection, "drop_database"), + ], +) +def test_base_unsupported_methods_have_messages(base, method): + with pytest.raises(NotImplementedError) as exc_info: + getattr(base, method)(None) + message = str(exc_info.value) + assert method in message + assert "not supported" in message def test_pickle_lancedb_client_error(): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 9655f5807..c1db92948 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -857,10 +857,116 @@ def test_table_unimplemented_functions(): with mock_lancedb_connection(handler) as db: table = db.create_table("test", [{"id": 1}]) - with pytest.raises(NotImplementedError): - table.to_arrow() - with pytest.raises(NotImplementedError): - table.to_pandas() + for method in ["to_arrow", "to_pandas", "to_polars"]: + with pytest.raises(NotImplementedError) as exc_info: + getattr(table, method)() + message = str(exc_info.value) + assert method in message + assert "not" in message and "supported" in message + assert "remote tables" in message + assert "cloud" not in message.lower() + assert f"search().{method}()" in message + + +@pytest.mark.parametrize("method", ["to_arrow", "to_pandas"]) +@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes", "invalid"]) +def test_remote_conversion_blob_mode_errors(method, blob_mode): + with blob_remote_table() as table: + if blob_mode == "invalid": + with pytest.raises(ValueError, match="blob_mode must be one of"): + getattr(table, method)(blob_mode=blob_mode) + else: + with pytest.raises(NotImplementedError) as error: + getattr(table, method)(blob_mode=blob_mode) + assert f"search().{method}()" in str(error.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [None, 1]) +@pytest.mark.parametrize("expose_location", [False, True]) +async def test_remote_to_lance_rejection( + tmp_path, version, expose_location, monkeypatch +): + import lance + + location = tmp_path / "test.lance" + lance.write_dataset(pa.table({"id": [1]}), location) + lance.write_dataset(pa.table({"id": [2]}), location, mode="append") + requests = [] + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) + + def handler(request): + requests.append(request.path) + assert request.path == "/v1/table/test/describe/" + content_len = int(request.headers.get("Content-Length", 0)) + body = json.loads(request.rfile.read(content_len)) if content_len else {} + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + response = { + "version": body.get("version") or 2, + "schema": {"fields": []}, + } + if expose_location: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) + + options = {"default_scan_options": {"with_row_id": True}} + with mock_lancedb_connection(handler) as db: + table = db.open_table("test", version=version) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as sync_error: + table.to_lance(**options) + assert requests == opened_requests + + async with mock_lancedb_connection_async(handler) as db: + table = await db.open_table("test", version=version) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as async_error: + await table.to_lance(**options) + assert requests == opened_requests + + assert "to_lance() is not supported for remote tables" in str(sync_error.value) + assert "query the server" in str(sync_error.value) + assert str(sync_error.value) == str(async_error.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method, kwargs", + [ + ("drop_all_tables", {}), + ("drop_all_tables", {"namespace_path": []}), + ("drop_all_tables", {"namespace_path": ["team"]}), + ("drop_database", {}), + ], +) +async def test_remote_drop_all_tables_not_supported(method, kwargs): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(500) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError) as sync_error: + getattr(db, method)(**kwargs) + + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError) as async_error: + await getattr(db, method)(**kwargs) + + assert "Dropping all tables is not currently supported in the remote API" in str( + sync_error.value + ) + assert "drop_table" in str(sync_error.value) + assert str(sync_error.value) == str(async_error.value) + assert requests == [] def test_table_to_pandas_not_supported(): @@ -1681,10 +1787,13 @@ def test_remote_python_only_messages(caplog): table = db.open_table("test", storage_options={}, index_cache_size=1) assert "storage_options is ignored for remote tables" in caplog.text assert "index_cache_size is ignored for remote tables" in caplog.text - for method in ["to_arrow", "to_pandas"]: + for method in ["to_arrow", "to_pandas", "to_polars"]: with pytest.raises(NotImplementedError) as error: getattr(table, method)() - assert str(error.value) == f"{method}() is not supported for remote tables." + assert str(error.value) == ( + f"{method}() is not supported for remote tables. " + f"Use table.search().{method}() instead." + ) for method in ["compact_files", "cleanup_old_versions"]: with pytest.warns(UserWarning) as warning: getattr(table, method)() @@ -3490,14 +3599,17 @@ def blob_descriptor_take_table(): @contextlib.contextmanager -def blob_remote_table(*, server_version=Version("0.5.0")): +def blob_remote_table(*, server_version=Version("0.5.0"), location=None): def handler(request): if request.path == "/v1/table/test/describe/": request.send_response(200) request.send_header("Content-Type", "application/json") request.send_header("phalanx-version", str(server_version)) request.end_headers() - request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode()) + response = dict(BLOB_DESCRIBE_RESPONSE) + if location is not None: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) elif request.path.startswith("/v1/table/test/blob/image/"): path = request.path.partition("?")[0] row_id = int(path.split("/")[-2]) @@ -3532,7 +3644,6 @@ def blob_remote_table(*, server_version=Version("0.5.0")): if columns == ["id"]: response_table = blob_query_response_table().select(["id"]) else: - assert body["with_row_id"] is True if columns == ["image"]: # Blob handles are sized from a descriptor take. assert body["filter"].startswith("_rowid IN") @@ -3545,6 +3656,8 @@ def blob_remote_table(*, server_version=Version("0.5.0")): ) else: response_table = blob_query_response_table() + if not body.get("with_row_id", False): + response_table = response_table.drop(["_rowid"]) request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.file") request.send_header("x-lancedb-version", "1") @@ -3579,6 +3692,32 @@ def test_remote_blob_columns_and_fetch(): assert blobs.to_pylist() == [b"alpha", None, b"gamma"] +@pytest.mark.parametrize("blob_mode", ["descriptions", "bytes", "lazy"]) +def test_sync_remote_blob_pandas_avoids_storage_access( + tmp_path, monkeypatch, blob_mode +): + import lance + + location = tmp_path / "remote.lance" + lance.write_dataset(pa.table({"id": [7], "image": [b"server-owned"]}), location) + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) + with blob_remote_table(location=location) as table: + monkeypatch.setattr(table, "to_lance", fail_dataset) + query = table.search().select(["id", "image"]) + if blob_mode == "descriptions": + df = query.to_pandas(blob_mode=blob_mode) + assert df["id"].tolist() == [1, 2, 3] + assert "_rowid" not in df.columns + assert df["image"][0]["size"] == 5 + else: + with pytest.raises(NotImplementedError, match="stable table snapshot"): + query.to_pandas(blob_mode=blob_mode) + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True]) async def test_remote_fetch_blob_ranges(use_async): diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 8a456ab3b..cad60ab7d 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1707,7 +1707,9 @@ impl Database for RemoteDatabase { // TODO: Implement namespace-aware drop_all_tables let _namespace_path = namespace_path; // Suppress unused warning for now Err(crate::Error::NotSupported { - message: "Dropping all tables is not currently supported in the remote API".to_string(), + message: "Dropping all tables is not currently supported in the remote API. \ + List tables and call drop_table for each table instead." + .to_string(), }) } From 6777a40dc280be8fd8f55e67750cbc6ecf7a2a60 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 09:51:10 +0800 Subject: [PATCH 56/64] fix: support remote tables in storage root (#3881) ## Summary - allow `db://` to select the remote storage root when a host override supplies the endpoint - preserve the empty `x-lancedb-database` header used by the legacy clients for root storage - reject root-storage connections without a host override before constructing an invalid Cloud URL - document the supported URI and cover both accepted and rejected configurations ## Root cause The shared Rust remote URL parser required every `db://` URI to contain a host. That validation was introduced when the Node and Python clients moved to the Rust implementation, but the legacy Node client intentionally treated the empty host in `db://` as the storage-root selector. ## Validation - `cargo fmt --all` - `cargo test --quiet --features remote -p lancedb test_connect_remote_storage_root` - `cargo check --quiet --features remote --tests --examples` - `cargo clippy --quiet --features remote --tests --examples` - `cargo test --quiet --features remote --tests` Fixes #1944 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- rust/lancedb/src/connection.rs | 34 +++++++++++++++++++++++++++++++ rust/lancedb/src/lib.rs | 1 + rust/lancedb/src/remote/client.rs | 31 +++++++++++++++++----------- rust/lancedb/src/remote/db.rs | 7 +++++++ 4 files changed, 61 insertions(+), 12 deletions(-) diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 124f6a54a..482acf89a 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -1270,6 +1270,7 @@ pub struct ConnectRequest { /// - `/path/to/database` - local database on file system. /// - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud object store /// - `db://dbname` - LanceDB Cloud + /// - `db://` with a host override - remote tables in the storage root pub uri: String, #[cfg(feature = "remote")] @@ -1382,6 +1383,8 @@ impl ConnectBuilder { /// /// This option is only used when connecting to LanceDB Cloud (db:// URIs) /// and will be ignored for other URIs. + /// Use the URI `db://` together with a host override to connect to remote + /// tables stored directly in the storage root. /// /// # Arguments /// @@ -2007,6 +2010,37 @@ mod tests { } } + #[cfg(feature = "remote")] + #[tokio::test] + async fn test_connect_remote_storage_root() { + let conn = ConnectBuilder::new("db://") + .region("us-east-1") + .api_key("my-api-key") + .host_override("https://example.com") + .execute() + .await + .unwrap(); + + let (impl_name, properties) = conn.namespace_client_config().await.unwrap(); + assert_eq!(impl_name, "rest"); + assert_eq!(properties["uri"], "https://example.com"); + assert!( + !properties.contains_key("header.x-lancedb-database"), + "storage-root connections must not send the x-lancedb-database header" + ); + + let result = ConnectBuilder::new("db://") + .region("us-east-1") + .api_key("my-api-key") + .execute() + .await; + assert!(matches!( + result, + Err(Error::InvalidInput { message }) + if message.contains("A host override is required") + )); + } + #[cfg(feature = "remote")] #[tokio::test] async fn test_connect_rejects_header_provider_with_oauth_config() { diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index e15017e7f..68521ff2c 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -54,6 +54,7 @@ //! - `/path/to/database` - local database on file system. //! - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud object store //! - `db://dbname` - Lance Cloud +//! - `db://` with a host override - remote tables in the storage root //! //! You can also use [`ConnectBuilder`] to configure the connection to the database. //! diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 783d2ef84..fe4a2072e 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -378,23 +378,30 @@ pub struct ParsedDbUrl { /// Parse a database URL and extract the database name and optional prefix. /// -/// Expected format: `db://db_name` or `db://db_name/prefix` +/// Expected format: `db://db_name`, `db://db_name/prefix`, or `db://` when +/// connecting to the storage root through a host override. pub fn parse_db_url(db_url: &str) -> Result { let parsed_url = url::Url::parse(db_url).map_err(|err| Error::InvalidInput { message: format!("db_url is not a valid URL. '{db_url}'. Error: {err}"), })?; debug_assert_eq!(parsed_url.scheme(), "db"); - if !parsed_url.has_host() { - return Err(Error::InvalidInput { - message: format!("Invalid database URL (missing host) '{}'", db_url), - }); - } - let db_name = urlencoding::decode(parsed_url.host_str().unwrap()) - .map_err(|err| Error::InvalidInput { - message: format!("Invalid encoded database name: {err}"), - })? - .into_owned(); - validate_database_name(&db_name)?; + let db_name = match parsed_url.host_str() { + Some(db_name) => { + let db_name = urlencoding::decode(db_name) + .map_err(|err| Error::InvalidInput { + message: format!("Invalid encoded database name: {err}"), + })? + .into_owned(); + validate_database_name(&db_name)?; + db_name + } + None if matches!(parsed_url.path(), "" | "/") => String::new(), + None => { + return Err(Error::InvalidInput { + message: format!("Invalid database URL (missing host) '{}'", db_url), + }); + } + }; let db_prefix = { let prefix = parsed_url.path().trim_start_matches('/'); if prefix.is_empty() { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index cad60ab7d..f07d20fc3 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -318,6 +318,13 @@ impl RemoteDatabase { read_consistency_interval: Option, ) -> Result { let parsed = super::client::parse_db_url(uri)?; + if parsed.db_name.is_empty() && host_overrides.rest.is_none() { + return Err(Error::InvalidInput { + message: + "A host override is required when connecting to the storage root with 'db://'" + .to_string(), + }); + } Self::try_new_with_identity( uri, api_key, From a8305e072246d21a97e8857c480e3bbdb9a4178f Mon Sep 17 00:00:00 2001 From: Romin Parekh <3831468+rominparekh@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:08:40 -0700 Subject: [PATCH 57/64] fix(rust): accept usize::MAX as a local query limit (#4567) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #2211. Local queries converted `limit`/`offset` with `as i64`, so `limit(usize::MAX)` wrapped to `-1` and Lance failed with `Limit must be non-negative`. The remote client already worked around this (#2354); this fixes the local path. - Saturate scanner offsets to `i64::MAX` and cap limits at `i64::MAX - offset`, keeping both conversion and full-text pagination addition safe. - Use `saturating_add` for vector `top_k = limit + offset`, then cap candidates at the dataset row count to avoid indexed-search capacity and refinement overflows. Keep a positive candidate bound for empty tables because Lance rejects `k = 0`. - Leave public APIs and remote query handling unchanged. ### Testing - Regression coverage includes plain scans, indexed IVF search, refinement, index bypass, empty vector tables, and indexed full-text pagination with finite and extreme limits/offsets. - Reproduced the indexed-search and FTS overflows before applying the follow-up fix. - `cargo test --quiet --features remote -p lancedb --lib query::` — 81 passed. - `cargo clippy --quiet --features remote -p lancedb --tests --examples` — clean. - `cargo fmt --all` and `git diff --check` — clean. Vector planning now reads the dataset row count (metadata fast path) to bound candidate allocation; this was not performance-benchmarked. --------- Co-authored-by: Romin Parekh Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 29 +++--- rust/lancedb/src/query.rs | 153 ++++++++++++++++++++++++++++++++ rust/lancedb/src/table/query.rs | 21 ++++- 3 files changed, 188 insertions(+), 15 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 9d7610f72..52f38d077 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -1561,9 +1561,10 @@ class LanceQueryBuilder(ABC): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -1619,8 +1620,8 @@ class LanceQueryBuilder(ABC): GlobalLimitExec: elapsed=..., skip=0, fetch=10, metrics=[output_rows=..., elapsed_compute=..., output_bytes=...] FilterExec: elapsed=..., _distance@2 IS NOT NULL, metrics=[...] - SortExec: elapsed=..., TopK(fetch=10), expr=[...], - preserve_partitioning=[...], + SortExec: elapsed=..., TopK(fetch=1), expr=[...], + preserve_partitioning=[...], filter=[...], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., row_replacements=...] KNNVectorDistance: elapsed=..., metric=l2, @@ -1633,6 +1634,7 @@ class LanceQueryBuilder(ABC): metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., fragments_scanned=..., ranges_scanned=1, rows_scanned=1, bytes_read=..., iops=..., requests=..., task_wait_time=...] + Parameters ---------- @@ -2959,9 +2961,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -3601,9 +3604,10 @@ class AsyncQueryBase(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters @@ -4521,9 +4525,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): LanceRead: uri=..., projection=[text], source=stream(_rowid) GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score, _rowid@0 as _rowid] LanceRead: uri=..., projection=[vector, text], source=stream(_rowid) GlobalLimitExec: skip=0, fetch=10 @@ -4877,9 +4881,10 @@ class BaseQueryBuilder(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 00a2b88f3..9bac51209 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -2714,6 +2714,159 @@ mod tests { } } + #[tokio::test] + async fn test_limit_usize_max() { + let conn = connect("memory://foo").execute().await.unwrap(); + let table = conn + .create_table("my_table", make_non_empty_batches()) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .offset(1) + .nearest_to(&[0.1; 4]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected - 1); + } + + #[tokio::test] + async fn test_indexed_vector_limit_usize_max() { + use crate::index::vector::IvfFlatIndexBuilder; + + let conn = connect("memory://indexed_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_non_empty_batches()) + .execute() + .await + .unwrap(); + table + .create_index( + &["vector"], + Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(1)), + ) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + for limit in [expected, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, usize::MAX] { + for (bypass, refine_factor) in [(false, None), (false, Some(2)), (true, None)] { + let query = table + .query() + .limit(limit) + .offset(offset) + .nearest_to(&[0.1; 4]) + .unwrap(); + let query = if bypass { + query.bypass_vector_index() + } else { + query + }; + let query = if let Some(factor) = refine_factor { + query.refine_factor(factor) + } else { + query + }; + let batches: Vec = + query.execute().await.unwrap().try_collect().await.unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected.saturating_sub(offset), + "limit={limit}, offset={offset}, bypass={bypass}, refine={refine_factor:?}" + ); + } + } + } + } + + #[tokio::test] + async fn test_empty_vector_limit_usize_max() { + let conn = connect("memory://empty_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_test_batches()) + .execute() + .await + .unwrap(); + let batches: Vec = table + .query() + .limit(usize::MAX) + .nearest_to(&[0.1; 128]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + + #[tokio::test] + async fn test_fts_limit_usize_max_with_offset() { + let conn = connect("memory://fts_limit").execute().await.unwrap(); + let data = RecordBatch::try_from_iter([( + "text", + Arc::new(StringArray::from_iter_values((0..32).map(|_| "hello"))) + as arrow_array::ArrayRef, + )]) + .unwrap(); + let table = conn + .create_table("documents", data) + .execute() + .await + .unwrap(); + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + + for limit in [32, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, i64::MAX as usize, usize::MAX] { + let batches: Vec = table + .query() + .full_text_search(FullTextSearchQuery::new("hello".into())) + .limit(limit) + .offset(offset) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 32usize.saturating_sub(offset), + "limit={limit}, offset={offset}" + ); + } + } + } + #[tokio::test] async fn test_repeated_only_if_combines_with_and() { use crate::expr::{col, lit}; diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index f625a4aff..4ad840350 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -280,7 +280,13 @@ pub async fn create_plan( let (_, element_type) = lance::index::vector::utils::get_vector_type(schema, &column)?; let is_binary = matches!(element_type, DataType::UInt8); - let top_k = query.base.limit.unwrap_or(DEFAULT_TOP_K) + query.base.offset.unwrap_or(0); + let top_k = query + .base + .limit + .unwrap_or(DEFAULT_TOP_K) + .saturating_add(query.base.offset.unwrap_or(0)) + // Lance allocates candidates from k; no search can return more rows than the table. + .min(ds_ref.count_rows(None).await?.max(1)); if is_binary { let query_vector = arrow::compute::cast(&query_vector, &DataType::UInt8)?; @@ -317,9 +323,18 @@ pub async fn create_plan( // For a batch query, `nearest` already applies k to each query vector. // Adding Scanner's global limit would truncate the combined result to k rows. if !is_batch_query { + let offset = query + .base + .offset + .map(|offset| i64::try_from(offset).unwrap_or(i64::MAX)); scanner.limit( - query.base.limit.map(|limit| limit as i64), - query.base.offset.map(|offset| offset as i64), + // Lance's FTS planner adds limit + offset as i64. + query.base.limit.map(|limit| { + i64::try_from(limit) + .unwrap_or(i64::MAX) + .min(i64::MAX - offset.unwrap_or(0)) + }), + offset, )?; } From cf5c9b5453e5abb30341d33ca1b67a7812b58ecb Mon Sep 17 00:00:00 2001 From: J Ganesh Kumar Reddy Date: Sat, 24 Oct 2026 10:01:49 +0530 Subject: [PATCH 58/64] fix: validate cached remote tables on open (#4557) ## Summary Fixes #4401. - Validate cached remote table entries with the existing `describe` endpoint before returning a handle. - Preserve cached server-version behavior after a successful validation. - Add a regression test for a table dropped by another connection after its entry has been cached. A cached entry can otherwise outlive an external drop, allowing `open_table()` to return a handle that fails only on its first later operation. ## Verification [Linux validation](https://github.com/Ganesh1729-ui/lancedb/actions/runs/39420560974) passed against the exact runtime and test sources at PR head `d39092b01f6b49921b5d66ccedf29709a7ddc23d`, using the repository's Rust toolchain and `ci` profile: - `cargo fmt --all -- --check` - `cargo test --profile ci --locked -p lancedb --features remote --lib remote::db::tests::test_open_table -- --nocapture`: **8 passed**, including cached-table deletion, independent handle checkout, branch/version handling, schema seeding and describe-response handling. - Removing only the cached-hit existence guard makes `test_open_table_checks_that_a_cached_table_still_exists` fail at its `TableNotFound` assertion; restoring the fix makes the same test pass again. - `git diff --check` The validation workflow is on a separate fork branch and is excluded from this PR. These are focused, mock-backed remote unit tests; the full upstream matrix and live Cloud integration tests still need upstream CI approval. Local Windows compilation was unavailable because the MSVC environment has no `link.exe`; the linked Linux run completed successfully. --------- Co-authored-by: Ganesh172919 Co-authored-by: Xuanwo --- deny.toml | 10 +++++++++ rust/lancedb/src/remote/db.rs | 42 +++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/deny.toml b/deny.toml index 7d60ae63f..77412709f 100644 --- a/deny.toml +++ b/deny.toml @@ -110,8 +110,18 @@ ignore = [ # compact_str/smol_str, so clearing this requires polars to migrate. # https://rustsec.org/advisories/RUSTSEC-2026-0249 { id = "RUSTSEC-2026-0249", reason = "smartstring unmaintained via polars; no fixed upstream release" }, + # crypto-bigint < 0.7.6: timing side-channels in constant-time ops, + # incl. subtraction, zero counting, division and Montgomery multiplication. + # Reached only through aws-sdk-s3 sigv4a -> aws-sigv4 (p256/ecdsa); signed + # requests carry the SigV4a credential-derived key. The patched 0.7.6 was + # published 2026-10-23, but no published aws-sigv4 release uses it yet + # (latest 1.6.3 pins crypto-bigint ^0.5), so clearing requires smithy-rs to + # ship an updated aws-sigv4. + # https://rustsec.org/advisories/RUSTSEC-2026-0387 + { id = "RUSTSEC-2026-0387", reason = "crypto-bigint via aws-sigv4 (sigv4a); no fixed upstream aws-sigv4 release yet" }, ] + # --------------------------------------------------------------------------- # Licenses: only allow licenses we've reviewed as compatible with Apache-2.0. # --------------------------------------------------------------------------- diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index f07d20fc3..768de79d2 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1614,6 +1614,17 @@ impl Database for RemoteDatabase { // Every open gets its own checkout, schema cache, and freshness state. if let Some(version) = self.table_cache.get(&cache_key).await { + // A cached version can outlive a drop from another connection. Verify that + // the table still exists before returning a handle for it. + let req = self + .client + .post(&format!("/v1/table/{}/describe/", identifier)) + .json(&serde_json::json!({})); + let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?; + let rsp = + RemoteTable::::handle_table_not_found(&request.name, rsp, &request_id).await?; + self.client.check_response(&request_id, rsp).await?; + Ok(Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), request.name, @@ -2557,6 +2568,37 @@ mod tests { assert_eq!(table.name(), "table1"); } + #[tokio::test] + async fn test_open_table_checks_that_a_cached_table_still_exists() { + let describe_calls = Arc::new(AtomicUsize::new(0)); + let calls = describe_calls.clone(); + let mut db = super::RemoteDatabase::new_mock(move |request| { + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + match calls.fetch_add(1, Ordering::SeqCst) { + 0 => http::Response::builder() + .status(200) + .body(r#"{"version": 1}"#.to_string()) + .unwrap(), + 1 => http::Response::builder() + .status(404) + .body("table not found".to_string()) + .unwrap(), + call => panic!("unexpected describe call: {call}"), + } + }); + db.table_cache = moka::future::Cache::new(10); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + + conn.open_table("table1").execute().await.unwrap(); + let result = conn.open_table("table1").execute().await; + + assert!(matches!(result, Err(Error::TableNotFound { .. }))); + assert_eq!(describe_calls.load(Ordering::SeqCst), 2); + } + #[tokio::test] async fn test_open_table_checkout_is_independent_per_handle() { let latest = Arc::new(AtomicUsize::new(2)); From 2c9f2e75a068a960863f3c45c36334d965bf2c1c Mon Sep 17 00:00:00 2001 From: Yeluru Nitish <141549775+Nitish-1303@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:21:48 -0500 Subject: [PATCH 59/64] fix: keep FTS indexes in list_indices when their files are missing (#4251) Fixes #4250 `list_indices()` loads each FTS index's segment params from the index files to fill in `index_details`. That call used `?`, so a single FTS index with missing files made the whole listing fail. A BTree or Bitmap index with missing files still lists fine, because it never takes that path. Now, if the params can't be loaded, it logs a warning and keeps the entry with the details from the manifest. Callers can still see the damaged index and rebuild or drop it. Added a Python test based on the repro in the issue. It fails on 0.39.0 with the same RuntimeError. --------- Co-authored-by: Xuanwo --- python/python/tests/test_fts.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index d7fa2e0ec..1dc2fd9a1 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -391,6 +391,20 @@ def test_custom_stop_words_list(table): ) +def test_list_indices_with_missing_fts_files(table): + # Regression test for https://github.com/lancedb/lancedb/issues/4250 + table.create_fts_index("text") + table.create_scalar_index("id") + + fts = next(i for i in table.list_indices() if i.index_type == "FTS") + shutil.rmtree(Path(table.uri) / "_indices" / fts.index_uuid) + + indices = {i.name: i for i in table.list_indices()} + assert set(indices) == {"text_idx", "id_idx"} + assert indices["text_idx"].index_type == "FTS" + assert indices["text_idx"].columns == ["text"] + + def test_search_fts(table): table.create_fts_index("text") results = table.search("puppy").select(["id", "text"]).limit(5).to_list() From 0ea932dafdac1ba3f86135cded890012d9733fdc Mon Sep 17 00:00:00 2001 From: Shiduo Li Date: Fri, 23 Oct 2026 23:47:41 -0700 Subject: [PATCH 60/64] fix(rust): compile jobs without remote feature (#4103) ## What Move typed job decode error construction behind a feature-aware helper so builds without the `remote` feature do not reference `Error::Http`. Local decode failures continue to surface as `Error::Runtime`, while remote builds retain HTTP errors with request IDs. Add a regression test for local typed-result decode failures. Closes #4096. ## Testing - `cargo fmt --all -- --check` - `cargo check -p lancedb --no-default-features` - `cargo test -p lancedb --no-default-features job::tests::local_typed_result_decode_errors_are_runtime_errors` Co-authored-by: Xuanwo --- rust/lancedb/src/job.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 6865e59ee..d671365f9 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -636,6 +636,14 @@ mod tests { use super::*; + #[test] + fn local_typed_result_decode_errors_are_runtime_errors() { + assert!(matches!( + TerminalResult::local(Value::Null).decode::(), + Err(Error::Runtime { .. }) + )); + } + #[tokio::test] async fn mapped_spawned_job_reuses_outcome() { let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1); From 6b6778e68495466e78cc786ebc912601c57c8b5d Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sat, 24 Oct 2026 18:50:27 -0700 Subject: [PATCH 61/64] feat: infer lazy blob reader Function inputs (#4663) Infer Blob V2 input access from the Python Function signature. `video: BlobReader` requests a lazy, seekable reader; `video: bytes` requests the complete payload. Users do not need a separate `blob_inputs` registration option. ## Python experience For a video table joined to a table of frame timestamps, the Function can pass its reader directly to a media decoder: ```python import io import av from lancedb import BlobReader, udf @udf(pip=["av", "Pillow"]) def extract_frame(video: BlobReader, timestamp_s: float) -> bytes: with av.open(video) as container: stream = container.streams.video[0] container.seek(int(timestamp_s / stream.time_base), stream=stream) for frame in container.decode(stream): if frame.time is not None and frame.time >= timestamp_s: output = io.BytesIO() frame.to_image().save(output, format="JPEG") return output.getvalue() raise ValueError("No frame at or after the requested timestamp") ``` After registering `extract_frame`, the query keeps the usual Function call: ```sql SELECT f.video_id, f.timestamp_s, extract_frame(v.video, f.timestamp_s) AS image FROM videos AS v JOIN frames AS f ON v.id = f.video_id ``` The returned JPEG can feed a subsequent image-model step. The reader provides `read`, `readinto`, `seek`, `tell`, `size()`, and `read_range(offset, length)`; range reads leave the sequential cursor unchanged. Each row invocation owns its reader, which closes after the call. Decoder seeks can request the parts of the MP4 needed to locate and decode the requested frame; the amount read depends on the file layout and codec. Functions that need the whole value can keep a bytes signature: ```python @udf def video_length(video: bytes) -> int: return len(video) ``` ## Execution contract - Signature inference maps `BlobReader` to a Blob V2 field and records `blob_access="reader"`. `bytes` keeps its binary Arrow type and records `blob_access="bytes"`; it can also bind to a Blob V2 source. Explicit schemas support the same choice, including a Blob V2 field annotated as `bytes`. - Reader execution carries a compact blob reference into the Function worker. Payload reads happen when Python calls the reader, with requests split into at most 1 MiB ranges. The full video does not need to enter the Arrow input batch just to invoke the Function. - Bytes execution materializes the full source before the Function runs and remains subject to eager-input memory admission. Calling `reader.read()` without a size also materializes the remaining contents in Python; lazy access depends on how the Function consumes it. - Input access metadata survives Python/Rust Function serialization and registration. Packaged Functions import the worker-provided reader implementation. This PR defines the SDK contract; execution requires corresponding server support. - This changes blob inputs. The example still returns ordinary image bytes; it does not add streamed outputs or run an image model. Adds coverage for signature inference, explicit schemas, nullable readers, signature round trips, seeking, bounded range reads, buffering, and invalid responses. --- docs/src/python/python.md | 8 ++ docs/web/management/function.mdx | 34 +++++ python/python/lancedb/__init__.py | 1 + python/python/lancedb/_blob_reader.py | 120 ++++++++++++++++++ python/python/lancedb/functions.py | 80 +++++++++--- python/python/tests/test_blob_reader.py | 59 +++++++++ .../tests/test_first_class_function_slice2.py | 55 ++++++++ rust/lancedb/src/function.rs | 23 ++++ 8 files changed, 359 insertions(+), 21 deletions(-) create mode 100644 python/python/lancedb/_blob_reader.py create mode 100644 python/python/tests/test_blob_reader.py diff --git a/docs/src/python/python.md b/docs/src/python/python.md index c465d0ac4..5d1393924 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -304,6 +304,14 @@ listing a storage directory. ## Functions and Jobs +Annotate a blob input with `lancedb.BlobReader` for lazy, seekable access, or +with `bytes` to materialize its contents before the Function runs. Reader +inputs infer the Blob V2 type without an explicit input schema. Readers are +scoped to each invocation and support `read`, `readinto`, `seek`, `tell`, +`size`, and cursor-independent `read_range`. + +::: lancedb.BlobReader + ::: lancedb.functions.FunctionArtifact ::: lancedb.functions.FunctionParameter diff --git a/docs/web/management/function.mdx b/docs/web/management/function.mdx index 8e5b53e4a..6df99b5fd 100644 --- a/docs/web/management/function.mdx +++ b/docs/web/management/function.mdx @@ -205,6 +205,40 @@ Function output blobs follow the same compaction and cleanup lifecycle as ingested Managed blobs. Retained versions continue to own their payloads; cleanup reclaims an object only after its retained references are gone. +### Draft: incremental blob inputs (upcoming) + +This feature requires an SDK that exports `BlobReader` and a compatible +Enterprise runtime with reader input support enabled. Coordinate the runtime +upgrade before registering or executing reader Functions. + +Annotate a row Function parameter with `BlobReader` to read a Blob v2 input +incrementally. The signature selects the access mode without another +registration option. For example, read the first MP4 box type without loading +the whole video: + +```python +from lancedb import BlobReader, udf + + +@udf +def first_box_type(video: BlobReader) -> bytes: + video.seek(4) + return video.read(4) +``` + +The worker passes a seekable reader backed by the stored blob. `read(size)` +reads from the current cursor; `read_range(offset, length)` reads a byte range +without moving it. `size()` returns the payload size without reading it. +Each row call receives its own reader, which closes when the call returns or +raises. Do not retain it for another invocation. + +Use `video: bytes` when the Function needs the complete payload. Bytes inputs +materialize before invocation and remain subject to eager input limits. +Reader inputs keep the unread payload out of the input batch, but calling +`video.read()` without a size loads all remaining bytes in Python. Reader +inputs are supported in row mode. Return values remain ordinary values such +as `bytes`; this feature does not add streamed outputs. + ## Declare the environment {#environment} The function runs remotely, so its dependencies are part of its definition: diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 6b902c466..a3aa55209 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -24,6 +24,7 @@ from .remote import ClientConfig from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector +from ._blob_reader import BlobReader as BlobReader from .job import AsyncJob, Job from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery diff --git a/python/python/lancedb/_blob_reader.py b/python/python/lancedb/_blob_reader.py new file mode 100644 index 000000000..2fd0f00ca --- /dev/null +++ b/python/python/lancedb/_blob_reader.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Seekable binary input for Functions that consume Blob V2 values lazily.""" + +import io +import operator +import urllib.request +from collections.abc import Callable +from http import HTTPStatus + +_READ_CHUNK_BYTES = 1024 * 1024 +_READ_TIMEOUT_SECONDS = 300 + + +class BlobReader(io.RawIOBase): + """Read a blob incrementally inside a Function. + + Annotating a Function input with ``BlobReader`` requests a lazy, seekable + reader. Annotating it with ``bytes`` requests materialized bytes instead. + ``read_range`` does not change the sequential cursor. Readers are scoped + to the invocation and must not be retained for a later call. + + For local calls, construct a reader from a range-reading callable and the + blob's size. The callable accepts an offset and length and returns bytes. + """ + + def __init__(self, read_range: Callable[[int, int], bytes], size: int): + super().__init__() + self._read_range = read_range + self._size = operator.index(size) + if self._size < 0: + raise ValueError("blob size must be nonnegative") + self._position = 0 + + @classmethod + def _from_uri(cls, uri: str, size: int): + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + + def read_range(offset, length): + request = urllib.request.Request( + uri, headers={"Range": f"bytes={offset}-{offset + length - 1}"} + ) + with opener.open(request, timeout=_READ_TIMEOUT_SECONDS) as response: + if response.status != HTTPStatus.PARTIAL_CONTENT: + raise OSError( + "blob range service did not return a partial response" + ) + data = response.read(length + 1) + if len(data) != length: + raise OSError( + "blob range service returned an invalid payload length" + ) + return data + + return cls(read_range, size) + + def readable(self) -> bool: + return not self.closed + + def seekable(self) -> bool: + return not self.closed + + def size(self) -> int: + """Return the blob size without reading its payload.""" + self._checkClosed() + return self._size + + def tell(self) -> int: + self._checkClosed() + return self._position + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + self._checkClosed() + offset = operator.index(offset) + if whence == io.SEEK_SET: + position = offset + elif whence == io.SEEK_CUR: + position = self._position + offset + elif whence == io.SEEK_END: + position = self._size + offset + else: + raise ValueError(f"invalid whence: {whence}") + if position < 0: + raise ValueError("negative seek position") + self._position = position + return position + + def read_range(self, offset: int, length: int) -> bytes: + """Read a byte range without changing the cursor; truncate at EOF.""" + self._checkClosed() + offset, length = operator.index(offset), operator.index(length) + if offset < 0 or length < 0: + raise ValueError("offset and length must be nonnegative") + remaining = min(length, max(0, self._size - offset)) + chunks = [] + while remaining: + count = min(remaining, _READ_CHUNK_BYTES) + chunk = self._read_range(offset, count) + if len(chunk) != count: + raise OSError("blob range reader returned an invalid payload length") + chunks.append(chunk) + offset += count + remaining -= count + return b"".join(chunks) + + def read(self, size: int = -1) -> bytes: + self._checkClosed() + size = -1 if size is None else operator.index(size) + if size < 0: + size = max(0, self._size - self._position) + data = self.read_range(self._position, size) + self._position += len(data) + return data + + def readinto(self, buffer) -> int: + view = memoryview(buffer).cast("B") + data = self.read(len(view)) + view[: len(data)] = data + return len(data) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 7f5776b32..42c4b27e1 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -33,6 +33,7 @@ from typing import ( Annotated, Any, Callable, + Literal, Optional, Union, get_args, @@ -52,7 +53,9 @@ from pydantic import ( model_validator, ) +from .schema import blob as _blob_field from .schema import is_blob_v2_field as _is_blob_v2_field +from ._blob_reader import BlobReader from .secrets import EnvVarSecret _Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1) @@ -213,6 +216,7 @@ class FunctionParameter(_RemoteValue): name: str arrow_type: str nullable: bool + blob_access: Optional[Literal["bytes", "reader"]] = None class FunctionResultField(_OpenRemoteValue): @@ -1096,6 +1100,43 @@ def _infer_signature( if (input_schema is None) != (output_schema is None): raise ValueError("input_schema and output_schema must be provided together") + try: + annotations = get_type_hints(function, include_extras=True) + except Exception as error: + raise TypeError(f"failed to resolve Function annotations: {error}") from error + + def input_parameter(field): + annotation = annotations.get(field.name) + if get_origin(annotation) in (Union, types.UnionType): + choices = [ + value for value in get_args(annotation) if value is not type(None) + ] + if len(choices) == 1: + annotation = choices[0] + arrow_type = _canonical_arrow_field(field) + if annotation is BlobReader: + if arrow_type != _FUNCTION_BLOB_V2_TYPE: + raise TypeError( + f"BlobReader input '{field.name}' requires a Blob V2 field" + ) + access = "reader" + elif annotation is bytes: + if arrow_type != _FUNCTION_BLOB_V2_TYPE and not ( + pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type) + ): + raise TypeError( + f"bytes input '{field.name}' requires binary or Blob V2" + ) + access = "bytes" + else: + access = "bytes" if arrow_type == _FUNCTION_BLOB_V2_TYPE else None + return FunctionParameter( + name=field.name, + arrow_type=arrow_type, + nullable=field.nullable, + blob_access=access, + ) + if input_schema is not None: if not isinstance(input_schema, pa.Schema): raise TypeError("input_schema must be a PyArrow Schema") @@ -1110,20 +1151,9 @@ def _infer_signature( "input_schema fields must exactly match callable parameters in order: " f"expected {expected!r}, got {actual!r}" ) - inputs = tuple( - FunctionParameter( - name=field.name, - arrow_type=_canonical_arrow_field(field), - nullable=field.nullable, - ) - for field in input_schema - ) + inputs = tuple(input_parameter(field) for field in input_schema) return FunctionSignature(inputs=inputs, output=_function_output(output_schema)) - try: - annotations = get_type_hints(function, include_extras=True) - except Exception as error: - raise TypeError(f"failed to resolve Function annotations: {error}") from error missing = [ parameter.name for parameter in parameters if parameter.name not in annotations ] @@ -1132,15 +1162,21 @@ def _infer_signature( raise TypeError(f"missing Function annotations: {names!r}") inputs = [] for parameter in parameters: - data_type, nullable = _annotation_type(annotations[parameter.name]) - inputs.append( - FunctionParameter( - name=parameter.name, - arrow_type=_canonical_arrow_field( - pa.field(parameter.name, data_type, nullable=nullable) - ), - nullable=nullable, + annotation = annotations[parameter.name] + base = annotation + nullable = False + if get_origin(base) in (Union, types.UnionType): + choices = [value for value in get_args(base) if value is not type(None)] + if len(choices) == 1: + base, nullable = choices[0], True + if base is BlobReader: + inputs.append( + input_parameter(_blob_field(parameter.name, nullable=nullable)) ) + continue + data_type, nullable = _annotation_type(annotation) + inputs.append( + input_parameter(pa.field(parameter.name, data_type, nullable=nullable)) ) output_type, output_nullable = _annotation_type(annotations["return"]) if output_nullable: @@ -1239,7 +1275,7 @@ _PYTHON_CALLABLE_ARTIFACT = "python_callable" _PYTHON_BUNDLE_ARTIFACT = "python_bundle" # Modules the Function image places beside user code. _RESERVED_SOURCE_MODULES = frozenset( - {"user_function", "lance_source", "lance_callable"} + {"user_function", "lance_source", "lance_callable", "lance_blob"} ) # The service accepts bundle paths made of these components only. _SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") @@ -1499,6 +1535,8 @@ class _SourcePackager: """One module-level line that rebinds `name` to `value` in the artifact: an import for modules and importable classes/functions, a literal otherwise.""" + if value is BlobReader: + return f"from lance_blob import BlobReader as {name}" if isinstance(value, types.ModuleType): if value.__name__.split(".")[0] in _NAMESPACE_MODULES: raise ValueError( diff --git a/python/python/tests/test_blob_reader.py b/python/python/tests/test_blob_reader.py new file mode 100644 index 000000000..dd12409af --- /dev/null +++ b/python/python/tests/test_blob_reader.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import io + +import pytest + +from lancedb import BlobReader + + +def test_reader_is_lazy_and_seekable(): + requests = [] + size = 64 * 1024 * 1024 + + def read_range(offset, length): + requests.append((offset, length)) + return bytes((index % 251 for index in range(offset, offset + length))) + + with BlobReader(read_range, size) as reader: + assert requests == [] + assert reader.size() == size + reader.seek(-4, io.SEEK_END) + assert reader.read(10) == bytes(index % 251 for index in range(size - 4, size)) + assert reader.read(1) == b"" + assert requests == [(size - 4, 4)] + assert reader.read_range(1, 3) == b"\x01\x02\x03" + assert reader.tell() == size + reader.seek(2) + target = bytearray(3) + assert reader.readinto(target) == 3 + assert target == b"\x02\x03\x04" + reader.seek(-1, io.SEEK_CUR) + assert reader.tell() == 4 + with pytest.raises(ValueError, match="closed"): + reader.read(1) + + +def test_reader_supports_buffered_io_and_bounds_range_requests(): + requests = [] + payload = b"x" * (3 * 1024 * 1024 + 7) + + def read_range(offset, length): + requests.append((offset, length)) + return payload[offset : offset + length] + + with io.BufferedReader(BlobReader(read_range, len(payload))) as reader: + assert reader.read() == payload + assert requests + assert all(length <= 1024 * 1024 for _, length in requests) + + +def test_reader_rejects_invalid_ranges_and_short_reads(): + reader = BlobReader(lambda offset, length: b"", 10) + with pytest.raises(ValueError): + reader.seek(-1) + with pytest.raises(ValueError): + reader.read_range(-1, 1) + with pytest.raises(OSError, match="payload length"): + reader.read(1) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 9ee3e388a..874034564 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -23,6 +23,7 @@ import pyarrow as pa import pytest import lancedb +from lancedb import BlobReader from lancedb.functions import ( PythonRuntimeSpec, SecretBinding, @@ -966,6 +967,60 @@ def test_blob_fields_use_the_scalar_function_semantic_type(): assert signature.output.arrow_type == "blob_v2" +def test_blob_input_access_is_inferred_from_annotations(): + @udf + def read_blob(video: BlobReader, timestamp_s: float) -> bytes: + video.seek(int(timestamp_s)) + return video.read(4) + + @udf + def read_bytes(video: bytes) -> int: + return len(video) + + reader = read_blob.registration_request.signature + assert reader.inputs[0].arrow_type == "blob_v2" + assert reader.inputs[0].blob_access == "reader" + assert reader.inputs[1].blob_access is None + assert reader.output.arrow_type == "binary" + eager = read_bytes.registration_request.signature.inputs[0] + assert eager.arrow_type == "binary" + assert eager.blob_access == "bytes" + encoded = json.loads(read_blob.registration_request.to_canonical_json()) + assert encoded["signature"]["inputs"][0]["blob_access"] == "reader" + + +def test_nullable_blob_reader_and_explicit_blob_bytes(): + @udf + def nullable(video: Optional[BlobReader]) -> int: + return 0 if video is None else video.size() + + @udf( + input_schema=pa.schema([lancedb.blob("video", nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def eager(video: bytes) -> int: + return len(video) + + parameter = nullable.registration_request.signature.inputs[0] + assert parameter.nullable is True + assert parameter.blob_access == "reader" + parameter = eager.registration_request.signature.inputs[0] + assert parameter.arrow_type == "blob_v2" + assert parameter.blob_access == "bytes" + + +def test_blob_reader_rejects_an_incompatible_explicit_schema(): + def read(video: BlobReader) -> int: + return video.size() + + with pytest.raises(TypeError, match="requires a Blob V2 field"): + udf( + read, + input_schema=pa.schema([pa.field("video", pa.int64(), nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + + def test_whole_named_struct_function_can_include_a_blob_result_field(): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=False)]), diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 98d915118..36e6c30c8 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -156,11 +156,20 @@ pub struct FunctionArtifact { } /// One ordered Arrow input parameter. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BlobAccess { + Bytes, + Reader, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionParameter { pub name: String, pub arrow_type: String, pub nullable: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub blob_access: Option, } /// One field of an ordered named-struct result. @@ -956,6 +965,20 @@ mod conda_environment_tests { mod tests { use super::*; + #[test] + fn blob_reader_access_survives_signature_round_trip() { + let parameter: FunctionParameter = serde_json::from_value(serde_json::json!({ + "name": "video", "arrow_type": "blob_v2", "nullable": false, + "blob_access": "reader" + })) + .unwrap(); + assert_eq!(parameter.blob_access, Some(BlobAccess::Reader)); + assert_eq!( + serde_json::to_value(parameter).unwrap()["blob_access"], + "reader" + ); + } + /// Canonical form is what the FunctionVersion hash is taken over, so key /// order must come from the keys and not from however serde happened to /// emit them. Nesting is included because the sort is recursive. From 24c2669a50d31ad14fb763bbdd1e236a5087b0b7 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 25 Oct 2026 00:47:29 -0700 Subject: [PATCH 62/64] feat(sql): follow adaptive streaming and polling responses (#4664) SQL execution now starts with one Flight call. The server can stream rows immediately or return a polling continuation, which the SDK follows through the existing query handle and reader APIs. The same behavior applies with and without parameters. Both initial transports share direct-stream lifecycle handling. Handoffs are explicitly negotiated; existing servers can continue returning ordinary streams. Add coverage for streaming, deferred results, cancellation, malformed decisions, and legacy responses. ```mermaid flowchart TD A[execute_query or execute_query_async] --> B{Parameters supplied?} B -->|No| C[DoGet: SQL command ticket] B -->|Yes| D[DoExchange: SQL and Arrow parameter row] C --> E[Server chooses execution path] D --> E E -->|Stream| F[Initial response: streaming decision and Arrow batches] E -->|Defer| G[Initial response: PollInfo continuation] G --> H[SDK follows PollFlightInfo] H -->|Still running| H H -->|Result endpoints available| I[SDK fetches tickets with DoGet] F --> J[Query reader yields Arrow batches] I --> J ``` `execute_query` opens the reader; `execute_query_async` returns its query handle. The caller does not select the transport or execution path. --- docs/src/python/python.md | 13 +- docs/web/basics/sql.mdx | 2 + python/python/lancedb/db.py | 13 +- rust/lancedb/src/connection.rs | 10 +- rust/lancedb/src/remote/sql.rs | 86 ++++--- .../src/remote/sql/{exchange.rs => stream.rs} | 191 +++++++++++----- rust/lancedb/src/remote/sql_test.rs | 212 +++++++++++++++++- rust/lancedb/src/sql.rs | 20 +- 8 files changed, 414 insertions(+), 133 deletions(-) rename rust/lancedb/src/remote/sql/{exchange.rs => stream.rs} (69%) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 5d1393924..41d3e16d7 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -221,10 +221,19 @@ The connected database and `default_namespace_path=["public"]` are used for unqualified tables. Fully qualified references can still query other databases and namespaces available to the same deployment. `execute_query` returns a reader as soon as its initial result stream is available. `execute_query_async` -returns a query handle immediately; use it to inspect progress, open a reader, +returns a query handle once submission is accepted; use it to inspect progress, open a reader, or cancel the query. The SQL client is initialized by the first query and retained for the lifetime of the remote connection. Query ids are random, -connection-scoped references rather than encoded SQL or durable resume tokens: +connection-scoped references rather than encoded SQL or durable resume tokens. + +Both APIs start with one Flight call. The server can stream results on that +call or return a polling continuation that the SDK follows automatically. +This behavior is the same with and without parameters. Open the reader directly; +a directly streamed query needs its reader to make progress. Dropping an unread +handle or its reader stops direct streaming. Deferred execution continues until +explicitly cancelled. + +For example: ```python import lancedb diff --git a/docs/web/basics/sql.mdx b/docs/web/basics/sql.mdx index 1c4ad46cd..c89bd7a20 100644 --- a/docs/web/basics/sql.mdx +++ b/docs/web/basics/sql.mdx @@ -82,6 +82,8 @@ In application code, load credentials from environment variables or a secrets manager. The [connection guide](/basics/connect#enterprise) shows that setup. `execute_query()` returns a blocking Arrow reader when results start arriving. +The server can stream rows on the first call or return result tickets that the +SDK follows automatically, with the same behavior for parameterized queries. The example uses `read_all()` to collect the small result in memory. For larger results, consume the reader in batches. To submit a query and check its status separately, use `execute_query_async()`. See [Queries](/basics/query) for query diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 74acc28ba..8e53eac04 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1253,12 +1253,13 @@ class DBConnection(EnforceOverrides): parameters=[np.array([0.1, 0.2, 0.3], dtype=np.float32), 10], ) - Parameters are supported in queries, not in DDL or DML. A - parameterized query runs on the call that returns its rows: - cancelling it, or dropping its handle before reading, stops it on the - server. Its status follows that stream as well: ``describe()`` - reports ``running`` until the reader has received every row, so open - ``reader()`` directly rather than waiting for ``finished``. + Parameters are supported in queries, not in DDL or DML. With or + without parameters, the server chooses whether to stream rows on the + initial call or return a polling continuation. The reader follows + either path transparently. Dropping an unread handle or its reader + stops a directly streamed query; deferred execution requires explicit + cancellation. Open ``reader()`` directly: a streamed query needs its + reader to make progress and reports ``running`` until the stream ends. Local connections do not support SQL. """ diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 482acf89a..26afef5e8 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -372,12 +372,10 @@ impl ExecuteQueryAsyncBuilder { /// Do not follow a named placeholder directly with `$`: `$a$` starts a /// dollar-quoted string. /// - /// A parameterized statement runs on the call that returns its rows rather - /// than detached from it: cancelling the query, or dropping its handle - /// before reading it, stops the statement on the server. Its status follows - /// that stream too -- it reports `Running` until the reader has received - /// every row -- so open the reader directly rather than waiting for - /// `Finished`. + /// With or without parameters, the server may stream rows on the initial + /// call or return a polling continuation. The reader follows either path + /// transparently. Open it directly instead of waiting for `Finished`: + /// directly streamed queries need their reader to make progress. /// /// # Example /// diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index a58a9afad..d8f04119f 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -34,9 +34,9 @@ use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; -mod exchange; +mod stream; -use exchange::{ExchangeQuery, ExchangeQueryHandle}; +use stream::{StreamingQuery, StreamingQueryHandle, Submission}; const DEFAULT_SQL_PORT: u16 = 10025; const DEFAULT_SQL_TLS_PORT: u16 = 10026; @@ -191,20 +191,39 @@ impl SqlClient { } pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result { - match request.parameters { - Some(parameters) => { - self.submit_with_parameters( - &request.query, - parameters, - &request.default_namespace_path, - ) - .await + let timeout = self.inner.overall_timeout()?; + with_overall_timeout(timeout, "SQL query submission", async { + validate_namespace_path(&request.default_namespace_path)?; + let query_id = Uuid::now_v7(); + let submission = StreamingQuery::open( + self.inner.clone(), + &request.query, + request.parameters, + &request.default_namespace_path, + ) + .await?; + match submission { + Submission::Stream(response) => { + let query = + Arc::new(StreamingQuery::new(query_id, self.inner.clone(), response)); + self.queries + .insert(query_id, RegisteredQuery::Streamed(query.clone())); + Ok(Query::new(Arc::new(StreamingQueryHandle::new(query)))) + } + Submission::Deferred(poll_info) => { + let query = Arc::new(RemoteQuery::new( + query_id, + self.inner.clone(), + request.default_namespace_path, + poll_info, + )?); + self.queries + .insert(query_id, RegisteredQuery::Polled(query.clone())); + Ok(Query::new(Arc::new(RemoteQueryHandle::new(query)))) + } } - None => { - self.submit(&request.query, &request.default_namespace_path) - .await - } - } + }) + .await } pub(super) async fn submit( @@ -271,35 +290,6 @@ impl SqlClient { }) } - /// Values cannot ride on a poll, so a statement that has them is sent on - /// its own exchange instead. - pub(super) async fn submit_with_parameters( - &self, - query: &str, - parameters: RecordBatch, - default_namespace_path: &[String], - ) -> Result { - let timeout = self.inner.overall_timeout()?; - with_overall_timeout(timeout, "SQL query submission", async { - validate_namespace_path(default_namespace_path)?; - let query_id = Uuid::now_v7(); - let query = Arc::new( - ExchangeQuery::open( - query_id, - self.inner.clone(), - query, - parameters, - default_namespace_path, - ) - .await?, - ); - self.queries - .insert(query_id, RegisteredQuery::Exchanged(query.clone())); - Ok(Query::new(Arc::new(ExchangeQueryHandle::new(query)))) - }) - .await - } - pub(super) async fn describe(&self, query_id: Uuid) -> Result { let query = self .queries @@ -586,14 +576,14 @@ impl SqlClientInner { #[derive(Clone)] enum RegisteredQuery { Polled(Arc), - Exchanged(Arc), + Streamed(Arc), } impl RegisteredQuery { fn touch(&self) { match self { Self::Polled(query) => query.touch(), - Self::Exchanged(query) => query.touch(), + Self::Streamed(query) => query.touch(), } } @@ -602,14 +592,14 @@ impl RegisteredQuery { fn registry_expired(&self) -> bool { match self { Self::Polled(query) => query.registry_expired(Arc::strong_count(query) == 1), - Self::Exchanged(query) => query.registry_expired(Arc::strong_count(query) == 1), + Self::Streamed(query) => query.registry_expired(Arc::strong_count(query) == 1), } } async fn describe(&self) -> Result { match self { Self::Polled(query) => query.describe().await, - Self::Exchanged(query) => query.describe(), + Self::Streamed(query) => query.describe(), } } } diff --git a/rust/lancedb/src/remote/sql/exchange.rs b/rust/lancedb/src/remote/sql/stream.rs similarity index 69% rename from rust/lancedb/src/remote/sql/exchange.rs rename to rust/lancedb/src/remote/sql/stream.rs index 523aaf5d7..ed85c6c92 100644 --- a/rust/lancedb/src/remote/sql/exchange.rs +++ b/rust/lancedb/src/remote/sql/stream.rs @@ -1,26 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! SQL statements sent with parameters. -//! -//! A statement without parameters is polled: `PollFlightInfo` starts it -//! detached from the connection and `DoGet` reads its endpoints. Neither call -//! carries Arrow data, so a statement's values would have to be spelled into -//! its text. A statement with parameters is sent on one `DoExchange` instead: -//! the command and the parameter schema, then one row, and the result rows -//! come back on the same call. The call is the statement's lifetime, so -//! cancelling it -- or dropping the handle before reading -- stops the -//! statement on the server. +//! SQL submitted on one call, with either streamed rows or a polling handoff. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::Instant; use arrow_array::RecordBatch; -use arrow_flight::FlightDescriptor; +use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::encode::FlightDataEncoderBuilder; -use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; +use arrow_flight::sql::{Any, CommandStatementQuery, ProstMessageExt}; +use arrow_flight::{FlightDescriptor, PollInfo, Ticket}; use arrow_schema::SchemaRef; +use futures::TryStreamExt; use prost::Message; use tokio::sync::{Notify, mpsc}; use uuid::Uuid; @@ -51,8 +44,8 @@ fn failure(error: &Error) -> Lifecycle { }) } -/// A statement running on its own `DoExchange` call. -pub(super) struct ExchangeQuery { +/// A statement streaming on its initial `DoGet` or `DoExchange` call. +pub(super) struct StreamingQuery { id: Uuid, client: Arc, /// The call's rows, until a reader takes them or a cancellation drops them. @@ -63,7 +56,7 @@ pub(super) struct ExchangeQuery { last_accessed: StdMutex, } -struct PreparedExchange { +struct PreparedStream { schema: SchemaRef, first: Option, rest: Option, @@ -76,7 +69,7 @@ struct PreparedExchange { /// this one did. A query cancelled meanwhile has its rows dropped instead, /// which ends the call. struct Borrowed<'a> { - query: &'a ExchangeQuery, + query: &'a StreamingQuery, rows: Option, } @@ -104,18 +97,37 @@ impl Drop for Borrowed<'_> { } } -impl ExchangeQuery { - /// Send the statement and its parameters, returning once the server has - /// planned it: a statement it refuses is refused here, as a polled one is - /// on its first poll. +pub(super) const PROTOCOL_HEADER: &str = "lancedb-sql-protocol"; +pub(super) const PROTOCOL_VERSION: &str = "adaptive-v1"; +const RESPONSE_TYPE_URL: &str = "type.googleapis.com/lancedb.sql.QueryResponse"; + +#[derive(Clone, PartialEq, Message)] +pub(super) struct QueryResponse { + #[prost(message, optional, tag = "1")] + pub poll_info: Option, +} + +#[derive(Clone, PartialEq, Message)] +struct CommandTicket { + #[prost(bytes = "vec", tag = "2")] + command: Vec, +} + +pub(super) enum Submission { + Stream(ResultEndpointStream), + Deferred(PollInfo), +} + +impl StreamingQuery { pub(super) async fn open( - id: Uuid, client: Arc, query: &str, - parameters: RecordBatch, + parameters: Option, default_namespace_path: &[String], - ) -> Result { - if parameters.num_rows() != 1 { + ) -> Result { + if let Some(parameters) = ¶meters + && parameters.num_rows() != 1 + { return Err(Error::InvalidInput { message: format!( "SQL query parameters must be exactly one row, got {}", @@ -130,39 +142,102 @@ impl ExchangeQuery { Some(DEFAULT_READ_TIMEOUT), )? .unwrap(); - let descriptor = FlightDescriptor::new_cmd( - CommandStatementQuery { - query: query.to_string(), - transaction_id: None, + let command = CommandStatementQuery { + query: query.to_string(), + transaction_id: None, + } + .as_any() + .encode_to_vec(); + let connection = client.connection(&request_id).await?; + let headers = client.headers(default_namespace_path, &request_id).await?; + let mut metadata = super::client_with_headers(connection.client.clone(), &headers)? + .metadata() + .clone(); + metadata.insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + let mut flight = connection.client.clone(); + let response = if let Some(parameters) = parameters { + let messages = FlightDataEncoderBuilder::new() + .with_flight_descriptor(Some(FlightDescriptor::new_cmd(command))) + .with_schema(parameters.schema()) + .build(futures::stream::iter([Ok(parameters)])) + .try_collect::>() + .await + .map_err(|error| sql_error(&request_id, error))?; + let mut request = tonic::Request::new(futures::stream::iter(messages)); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_exchange(request)).await + } else { + let mut request = + tonic::Request::new(Ticket::new(CommandTicket { command }.encode_to_vec())); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_get(request)).await + } + .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? + .map_err(|error| sql_error(&request_id, error))?; + let negotiated = response + .metadata() + .get(PROTOCOL_HEADER) + .is_some_and(|value| value == PROTOCOL_VERSION); + let mut response_stream = response.into_inner(); + if negotiated { + let first = tokio::time::timeout(read_timeout, response_stream.message()) + .await + .map_err(|_| sql_error(&request_id, "SQL query handoff timed out"))? + .map_err(|error| sql_error(&request_id, error))? + .ok_or_else(|| { + sql_error(&request_id, "SQL response omitted its execution decision") + })?; + if !first.data_header.is_empty() || !first.data_body.is_empty() { + return Err(sql_error( + &request_id, + "SQL response sent rows before its execution decision", + )); } - .as_any() - .encode_to_vec(), + let envelope = + Any::decode(first.app_metadata).map_err(|error| sql_error(&request_id, error))?; + if envelope.type_url != RESPONSE_TYPE_URL { + return Err(sql_error( + &request_id, + "SQL response returned an unknown execution decision", + )); + } + let decision = QueryResponse::decode(envelope.value) + .map_err(|error| sql_error(&request_id, error))?; + if let Some(info) = decision.poll_info { + if info.info.is_none() && info.flight_descriptor.is_none() { + return Err(sql_error( + &request_id, + "SQL handoff omitted result and continuation information", + )); + } + return Ok(Submission::Deferred(info)); + } + } + let stream = FlightRecordBatchStream::new_from_flight_data( + response_stream + .map_err(|status| arrow_flight::error::FlightError::Tonic(Box::new(status))), ); - // The descriptor rides on the schema message, and the row follows it. - let request = FlightDataEncoderBuilder::new() - .with_flight_descriptor(Some(descriptor)) - .with_schema(parameters.schema()) - .build(futures::stream::iter([Ok(parameters)])); - let mut flight = client - .client_with_headers(default_namespace_path, &request_id) - .await?; - let stream = tokio::time::timeout(read_timeout, flight.do_exchange(request)) - .await - .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? - .map_err(|err| sql_error(&request_id, err))?; - Ok(Self { + Ok(Submission::Stream(ResultEndpointStream { + stream, + request_id, + read_timeout, + })) + } + + pub(super) fn new( + id: Uuid, + client: Arc, + response: ResultEndpointStream, + ) -> Self { + Self { id, client, - response: StdMutex::new(Some(ResultEndpointStream { - stream, - request_id, - read_timeout, - })), + response: StdMutex::new(Some(response)), lifecycle: StdMutex::new(Lifecycle::Running), cancelled: Notify::new(), terminal_at: OnceLock::new(), last_accessed: StdMutex::new(Instant::now()), - }) + } } pub(super) fn touch(&self) { @@ -242,7 +317,7 @@ impl ExchangeQuery { } /// Take the rows and wait for the first batch, which carries the schema. - async fn prepare(&self) -> Result { + async fn prepare(&self) -> Result { let Some(rows) = self.response.lock().unwrap().take() else { // Taken by an earlier attempt that failed, or dropped by a cancel. return Err(match &*self.lifecycle.lock().unwrap() { @@ -275,7 +350,7 @@ impl ExchangeQuery { if first.is_none() { self.settle(Lifecycle::Finished); } - Ok(PreparedExchange { + Ok(PreparedStream { schema, rest: first.is_some().then_some(rows), first, @@ -284,7 +359,7 @@ impl ExchangeQuery { async fn run( self: Arc, - mut prepared: PreparedExchange, + mut prepared: PreparedStream, sender: mpsc::Sender>, ) -> Result<()> { if let Some(batch) = prepared.first.take() @@ -343,13 +418,13 @@ impl ExchangeQuery { } } -pub(super) struct ExchangeQueryHandle { - query: Arc, +pub(super) struct StreamingQueryHandle { + query: Arc, result_started: AtomicBool, } -impl ExchangeQueryHandle { - pub(super) fn new(query: Arc) -> Self { +impl StreamingQueryHandle { + pub(super) fn new(query: Arc) -> Self { Self { query, result_started: AtomicBool::new(false), @@ -357,7 +432,7 @@ impl ExchangeQueryHandle { } } -impl Drop for ExchangeQueryHandle { +impl Drop for StreamingQueryHandle { // Rows nobody took are rows nobody will read. Keeping the call open would // hold the statement running on the server, stalled on flow control, // until the registry forgot it. @@ -369,7 +444,7 @@ impl Drop for ExchangeQueryHandle { } #[async_trait::async_trait] -impl QueryHandle for ExchangeQueryHandle { +impl QueryHandle for StreamingQueryHandle { fn id(&self) -> Uuid { self.query.touch(); self.query.id diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 79ca339db..4837baab7 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -19,6 +19,7 @@ use futures::{StreamExt, TryStreamExt}; use tonic::metadata::MetadataMap; use tonic::{Request, Response, Status, Streaming}; +use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse}; use super::*; use crate::database::Database; use crate::remote::RemoteCatalogOptions; @@ -70,6 +71,9 @@ type ReceivedExchanges = Arc)> #[derive(Clone)] struct TestSqlService { + adaptive: bool, + defer: bool, + invalid_decision: bool, query_count: Arc, do_get_count: Arc, cancel_count: Arc, @@ -120,6 +124,51 @@ impl TestSqlService { .unwrap_or_default(), }); } + fn adaptive_response( + &self, + query: &str, + rows: BoxStream<'static, std::result::Result>, + ) -> Response>> { + if !self.adaptive { + return Response::new(rows); + } + let poll_info = + self.defer.then(|| PollInfo { + info: Some(FlightInfo::new().with_endpoint( + FlightEndpoint::new().with_ticket(Ticket::new(query.to_string())), + )), + flight_descriptor: Some(FlightDescriptor::new_cmd(format!("poll:1:{query}"))), + progress: Some(0.25), + expiration_time: None, + }); + let first = FlightData { + app_metadata: if self.invalid_decision { + b"invalid protobuf".to_vec().into() + } else { + Any { + type_url: "type.googleapis.com/lancedb.sql.QueryResponse".to_string(), + value: QueryResponse { poll_info }.encode_to_vec().into(), + } + .encode_to_vec() + .into() + }, + ..Default::default() + }; + let rest = if self.defer { + futures::stream::empty().boxed() + } else { + rows + }; + let mut response = Response::new( + futures::stream::once(async { Ok(first) }) + .chain(rest) + .boxed(), + ); + response + .metadata_mut() + .insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + response + } } impl Default for TestSqlService { @@ -154,6 +203,9 @@ impl Default for TestSqlService { let dictionary_result = RecordBatch::try_new(dictionary_schema, vec![Arc::new(dictionary)]).unwrap(); Self { + adaptive: false, + defer: false, + invalid_decision: false, query_count: Arc::new(AtomicUsize::new(0)), do_get_count: Arc::new(AtomicUsize::new(0)), cancel_count: Arc::new(AtomicUsize::new(0)), @@ -308,7 +360,22 @@ impl FlightService for TestSqlService { request: Request, ) -> std::result::Result::DoGetStream>, Status> { self.do_get_count.fetch_add(1, Ordering::SeqCst); - let ticket = request.get_ref().ticket.as_ref(); + let initial = Any::decode(request.get_ref().ticket.clone()) + .ok() + .and_then(|envelope| Any::decode(envelope.value).ok()) + .and_then(|command| command.unpack::().ok().flatten()); + if initial.is_some() { + self.capture_headers(request.metadata()); + assert_eq!( + request.metadata().get(PROTOCOL_HEADER).unwrap(), + PROTOCOL_VERSION + ); + } + let query = initial + .as_ref() + .map(|command| command.query.as_bytes()) + .unwrap_or(request.get_ref().ticket.as_ref()); + let ticket = query; let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; @@ -337,7 +404,11 @@ impl FlightService for TestSqlService { encoder = encoder.with_max_flight_data_size(8 * 1024 * 1024); } let stream = encoder.build(input).map_err(Status::from); - Ok(Response::new(Box::pin(stream))) + if let Some(initial) = initial { + Ok(self.adaptive_response(&initial.query, stream.boxed())) + } else { + Ok(Response::new(stream.boxed())) + } } async fn do_put( @@ -486,7 +557,7 @@ impl FlightService for TestSqlService { let _dropped = &dropped; message }); - Ok(Response::new(Box::pin(stream))) + Ok(self.adaptive_response(&query, stream.boxed())) } } @@ -1609,3 +1680,138 @@ fn validates_metadata_with_header_map() { assert!(insert_header(&mut headers, "bad header", "value").is_err()); assert!(insert_header(&mut headers, "valid-header", "bad\nvalue").is_err()); } + +#[tokio::test] +async fn adaptive_streams_start_on_one_rpc_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = parameters.clone().unwrap_or_else(|| service.result.clone()); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT echo") + }) + .await + .unwrap(); + assert_eq!(query.describe().await.unwrap().status, QueryStatus::Running); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(gets.load(Ordering::SeqCst), usize::from(!with_parameters)); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_share_polling_and_survive_handle_drop() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = service.result.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT adaptive") + }) + .await + .unwrap(); + let id = query.id(); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + drop(query); + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 0, + "the original SQL must not be resubmitted through polling" + ); + assert_eq!( + gets.load(Ordering::SeqCst), + 1 + usize::from(!with_parameters) + ); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_use_server_cancellation_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let cancels = service.cancel_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT cancel") + }) + .await + .unwrap(); + query.cancel().await.unwrap(); + assert_eq!(cancels.load(Ordering::SeqCst), 1); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Cancelled + ); + assert!(query.reader().await.is_err()); + } +} + +#[tokio::test] +async fn malformed_adaptive_decision_never_resubmits_the_statement() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + invalid_decision: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + assert!( + client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT invalid") + }) + .await + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + } +} diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql.rs index c2ef13ba5..fee2d5d79 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql.rs @@ -18,9 +18,9 @@ pub enum QueryStatus { Running, /// The server has made the complete result available. /// - /// A query submitted with parameters streams its rows on the call that - /// runs it, so it reports `Finished` only once its reader has received - /// the last batch. Read it rather than waiting for this status. + /// Directly streamed queries report `Finished` when their stream reaches + /// its end. Deferred queries report the server's execution state. Open the + /// reader directly instead of waiting for this status before reading. Finished, /// The server accepted cancellation but has not confirmed it yet. Cancelling, @@ -64,9 +64,10 @@ pub(crate) trait QueryHandle: Send + Sync { /// A handle to a submitted SQL query. /// /// The handle can be inspected, opened as an Arrow reader, or cancelled. -/// Dropping it does not cancel the server-side query, unless the query was -/// submitted with parameters: those run on the call that returns their rows, -/// so dropping the handle before reading, or dropping the reader, stops them. +/// The server may stream results on the initial call or defer execution and +/// return result tickets. Dropping an unread handle or its reader stops a +/// directly streamed query; deferred execution continues until explicitly +/// cancelled. This applies equally with and without parameters. /// Identifier lookup is scoped to the connection that submitted the query and /// is not a durable resume mechanism. pub struct Query { @@ -95,10 +96,9 @@ impl Query { /// Get a point-in-time description of the query. /// - /// For a query submitted with parameters the status follows its result - /// stream: it stays [`QueryStatus::Running`] until the reader has received - /// every row, so open [`Self::reader`] directly instead of polling for - /// [`QueryStatus::Finished`] first. + /// Directly streamed queries stay [`QueryStatus::Running`] until the result + /// stream ends. Deferred queries expose the server's execution state. Open + /// [`Self::reader`] directly instead of waiting for [`QueryStatus::Finished`]. pub async fn describe(&self) -> Result { self.handle.describe().await } From 9af77570532c57029a47c7d47481b5cc12a9e2fe Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 25 Oct 2026 01:20:13 -0700 Subject: [PATCH 63/64] fix: execute remote materialized view creation through SQL (#4665) Remote materialized-view creation currently expects a REST response containing a registry job ID, while refresh already submits SQL. Route both through the SQL query backend and preserve Python's synchronous and asynchronous return types. ### Query to Job mapping `SqlClient::submit_as_job()` submits an `ExecuteQueryRequest` through the same backend as `execute_query_async()`. Once submission succeeds, it retains the returned `Query` and starts one client-side task that consumes its result reader to completion. `Job::from_query_task()` wraps that task, its cached outcome, a single ID and a cancellation callback. Python exposes this as `Job[T]` or `AsyncJob[T]`. The existing `job_type` property maps SQL-backed handles to the new `"query"` kind. The shared `id` contains that query's UUID string, so Python callers can inspect it with `db.describe_query(job.id)` on the submitting connection. No separate `query_id`, backend enum or `type()` API is added. | Handle | `job_type` | `id` | Available lookup | | --- | --- | --- | --- | | SQL-backed | `"query"`, known at submission | Connection-scoped query UUID string | `describe_query(job.id)` | | Registry-backed | Existing server-reported kind, after `status()` or `refresh()` | Registry job ID string | Registry job APIs, descriptions and events | | Local | `None` | `None` | The current in-process handle | | Job concept | SQL-backed behavior | | --- | --- | | Submission | Await initial SQL submission before returning the handle. Submission errors propagate immediately. | | `status()` | Tracks the adapter task: running during result consumption/finalization, finished after both succeed, failed if either fails, or cancelled after cancellation. | | `wait()` | Returns the operation's typed terminal result. One task consumes the query reader once; repeated waits reuse its cached outcome without resubmitting SQL or reopening the reader. | | `cancel()` | Delegates to `Query.cancel()`. Rejected or unsupported cancellation leaves the task running; an accepted request stops the task. Cancelling a finished operation is a no-op. | | Registry details | Query-backed handles expose no registry description, event history, pause/resume or reattachment after a client restart. `refresh()` updates cached lifecycle state. | The Job tracks command completion rather than exposing SQL batches. A wait timeout does not cancel the query. Query IDs remain scoped to the submitting connection and cannot be used with registry APIs or as durable resume tokens. Python `describe_query()` accepts both UUID objects and the string `id` exposed by query-backed jobs. Rust, Python and TypeScript expose the same classification through their existing job-type accessors. ### Materialized-view operations - **CREATE:** Render `CREATE MATERIALIZED VIEW` with the view namespace and optional auto-refresh settings. After command completion, resolve the job to `None`. Synchronous creation waits on the job, then opens and returns the materialized view. - **REFRESH:** Submit `REFRESH MATERIALIZED VIEW` through the same adapter. After command completion, read the refreshed table's version and row count to construct the existing typed refresh result. Synchronous refresh waits on the same path. Remote creation requires the SQL service. Execution placement remains a server decision. REST endpoints are unchanged. The job guide documents the ID/job-type contract and query-backed limitations. Regression coverage exercises SQL submission, namespaces, refresh options, job-type/ID preservation through typed result mapping, repeated waits, failures and cancellation rejection. No local builds or tests were run. --- docs/src/js/classes/Job.md | 13 +- docs/web/basics/jobs.mdx | 37 ++++- nodejs/__test__/materialized_view.test.ts | 1 + nodejs/lancedb/job.ts | 13 +- nodejs/src/job.rs | 12 +- python/python/lancedb/_lancedb.pyi | 2 + python/python/lancedb/db.py | 25 ++- python/python/lancedb/job.py | 30 ++-- python/python/lancedb/namespace.py | 2 +- python/python/lancedb/remote/db.py | 2 +- .../tests/test_first_class_function_slice2.py | 4 + .../python/tests/test_materialized_views.py | 87 +++++++--- python/python/tests/test_sql.py | 20 ++- python/src/job.rs | 5 + rust/lancedb/src/job.rs | 76 +++++++-- rust/lancedb/src/materialized_view.rs | 3 +- rust/lancedb/src/remote/db.rs | 144 ++++++----------- rust/lancedb/src/remote/sql.rs | 45 +++--- rust/lancedb/src/remote/sql_test.rs | 152 +++++++++++++++++- rust/lancedb/src/remote/table.rs | 8 +- 20 files changed, 465 insertions(+), 216 deletions(-) diff --git a/docs/src/js/classes/Job.md b/docs/src/js/classes/Job.md index cfbe454ec..74e1c00e2 100644 --- a/docs/src/js/classes/Job.md +++ b/docs/src/js/classes/Job.md @@ -17,7 +17,8 @@ extra round trip on every call: - [Job.refresh](Job.md#refresh) and [Job.status](Job.md#status) fetch the whole record. - [Job.wait](Job.md#wait) records the terminal state it establishes, but not the rest of the record. -- Everything is null until one of those runs. +- Registry details are null until one of those runs. Query handles report + jobType "query" immediately. ## Accessors @@ -57,8 +58,10 @@ get id(): null | string Identifies the operation on the server that is running it. -Operations that run in this process have no server id. The value is -opaque: parsing it or storing it to resume the job later is not supported. +Registry handles expose a registry job ID; query handles expose a +connection-scoped SQL query UUID string. Local handles have no ID. +Query handles report jobType "query". Their IDs are not registry IDs or +durable resume tokens. #### Returns @@ -72,8 +75,8 @@ opaque: parsing it or storing it to resume the job later is not supported. get jobType(): null | string ``` -The job's type, as the server names it. Null for an in-process job, which -has no server-side record. +"query" immediately for SQL-backed handles; otherwise the cached +server-reported job type. Null for in-process handles. #### Returns diff --git a/docs/web/basics/jobs.mdx b/docs/web/basics/jobs.mdx index 0577cd120..a87663a51 100644 --- a/docs/web/basics/jobs.mdx +++ b/docs/web/basics/jobs.mdx @@ -20,19 +20,20 @@ failures. Methods such as Python's `create_index_async` and TypeScript's `createIndexAsync` return job handles. In OSS, the work runs in your application -process; in Enterprise, it runs on server workers. A returned handle can -represent an operation that has already finished, so check its status before +process; in Enterprise, it runs on server workers or through the SQL query +backend. A returned handle can represent an operation that has already finished, +so check its status before waiting or requesting cancellation. See the operation's API reference for when the method returns. ## Submit and wait for a job {#waiting} 1. **Submit the operation.** Keep the returned job handle. For an Enterprise - job, also keep its ID if another process needs to look it up later. + registry job, also keep its ID if another process needs to look it up later. 2. **Check progress.** Read the current status while the operation runs. 3. **Wait for completion.** Check that the job succeeded before using its - result. If it failed or was cancelled, inspect the failure details and any - results it committed before stopping. + result. If it failed or was cancelled, inspect the failure details and any + results it committed before stopping. In Python, `job.status()` reports the current state and `job.wait()` waits for completion, raising if the job fails or is cancelled. If `job.wait()` times out, the client stops waiting. @@ -59,13 +60,37 @@ history. ## OSS and Enterprise execution {#in-process} -| | OSS | Enterprise | +The following table describes local operations and Enterprise registry jobs. +SQL-backed job handles follow the query lifecycle described below. + +| | OSS | Enterprise registry job | | --- | --- | --- | | Execution | In your application process | On server workers | | Job identity | No server ID | A server ID for looking up the job | | Event history | No server event history | Available through the jobs API | | Reattaching after a client restart | No server job to reopen | Reopen by ID while the record is retained | +### SQL-backed job handles + +Remote materialized-view CREATE and REFRESH use the SQL query backend. Python's +`create_materialized_view_async()` and `mv.refresh_async()` still return job +handles with `status()`, `wait()` and `cancel()`. Synchronous creation and refresh +wait on the same SQL execution before returning. + +The existing `job.job_type` identifies query-backed handles as `"query"` from +submission, and `job.id` contains the connection-scoped SQL query UUID string. +Use `describe_query(job.id)` on the submitting connection to inspect the query. +Registry handles retain their server-reported job types (available after +`status()` or `refresh()`) and registry IDs. Local handles have no ID or job type. + +SQL-backed handles do not expose registry records or event history and cannot +be reopened through `open_job()` or reattached after a client restart. Their IDs +are not durable resume tokens. Cancellation follows the SQL server's +capabilities and can be unsupported. + +The server decides whether SQL executes on its serving fleet or starts SQL job +workers. That placement does not change the handle's query-backed identity. + ## Learn more {#next} - [Build and repair indexes](/indexing/reindexing): submit index work. diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index 3e47a00e9..3852af717 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -129,6 +129,7 @@ describe("materialized views", () => { const job = await db.dropMaterializedViewAsync("adults"); expect(job.id).toBeNull(); + expect(job.jobType).toBeNull(); await job.wait(); expect(await collect(db.listMaterializedViews())).toEqual([]); }); diff --git a/nodejs/lancedb/job.ts b/nodejs/lancedb/job.ts index 0baa0f571..1c8d74ad2 100644 --- a/nodejs/lancedb/job.ts +++ b/nodejs/lancedb/job.ts @@ -24,7 +24,8 @@ export interface JobEventsOptions { * - {@link Job.refresh} and {@link Job.status} fetch the whole record. * - {@link Job.wait} records the terminal state it establishes, but not the * rest of the record. - * - Everything is null until one of those runs. + * - Registry details are null until one of those runs. Query handles report + * jobType "query" immediately. * * @hideconstructor */ @@ -38,8 +39,10 @@ export class Job { /** * Identifies the operation on the server that is running it. * - * Operations that run in this process have no server id. The value is - * opaque: parsing it or storing it to resume the job later is not supported. + * Registry handles expose a registry job ID; query handles expose a + * connection-scoped SQL query UUID string. Local handles have no ID. + * Query handles report jobType "query". Their IDs are not registry IDs or + * durable resume tokens. */ get id(): string | null { return this.inner.id ?? null; @@ -51,8 +54,8 @@ export class Job { } /** - * The job's type, as the server names it. Null for an in-process job, which - * has no server-side record. + * "query" immediately for SQL-backed handles; otherwise the cached + * server-reported job type. Null for in-process handles. */ get jobType(): string | null { return this.inner.jobType ?? null; diff --git a/nodejs/src/job.rs b/nodejs/src/job.rs index 214e687b4..320812128 100644 --- a/nodejs/src/job.rs +++ b/nodejs/src/job.rs @@ -29,8 +29,8 @@ impl Job { #[napi] impl Job { - /// Identifies the operation on the server that is running it. Operations - /// that run in this process have no server id. The value is opaque. + /// The operation's opaque ID. Interpret registry and query IDs using job_type(). + /// Local handles have no ID. #[napi(getter)] pub fn id(&self) -> Option { self.inner.id().map(str::to_string) @@ -62,8 +62,8 @@ impl Job { /// Ask the backend for this job's current state, and for a server-side job /// its full record, then cache it for the getters below. /// - /// They are all null until this runs, because submitting an operation - /// returns only a job id. {@link Job.status} fetches the whole record too; + /// Registry details are null until this runs; query handles already expose + /// job_type "query". {@link Job.status} fetches the whole registry record too; /// {@link Job.wait} records only the terminal state it establishes. #[napi(catch_unwind)] pub async fn refresh(&self) -> napi::Result<()> { @@ -76,8 +76,8 @@ impl Job { self.inner.state() } - /// The job's type, as the server names it. Null for an in-process job, - /// which has no server-side record. + /// "query" immediately for SQL-backed handles; otherwise the cached + /// server-reported job type. Null for in-process handles. #[napi(getter)] pub fn job_type(&self) -> Option { self.inner.job_type() diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 6ef23abc7..9e602cce8 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -320,6 +320,8 @@ class Job: @property def id(self) -> Optional[str]: ... @property + def job_type(self) -> Optional[str]: ... + @property def _state(self) -> Optional[str]: ... @property def _description(self) -> Optional[JobDescription]: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 8e53eac04..1558dd52a 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -674,10 +674,11 @@ class DBConnection(EnforceOverrides): ) -> Job[None]: """Submit materialized-view creation and return its job. - The job may already be complete for a local database. On LanceDB - Cloud and Enterprise, its ``id`` is the server job identifier from - the ``202 Accepted`` create response. Wait for the job before opening - or querying the view. + The job may already be complete for a local database. Remote creation + submits CREATE MATERIALIZED VIEW through the SQL server. Its job has + ``job_type == "query"`` and ``id`` identifies the connection-scoped + query. Registry event history is unavailable. + Wait for the job before opening or querying the view. """ raise NotImplementedError( "materialized views are not supported on this connection type" @@ -1265,7 +1266,7 @@ class DBConnection(EnforceOverrides): """ raise NotImplementedError("SQL is not supported for this connection type") - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" raise NotImplementedError("SQL is not supported for this connection type") @@ -2869,7 +2870,9 @@ class AsyncConnection(object): ) -> AsyncJob[None]: """Submit materialized-view creation and return its job. - Wait for the returned job before opening or querying the view. + Wait for the returned job before opening or querying the view. Remote + creation uses SQL; ``job_type`` is ``"query"`` and ``id`` + identifies that query on this connection. """ source, native_options = await self._materialized_view_source( source, select, where, limit @@ -3409,8 +3412,14 @@ class AsyncConnection(object): ) ) - async def describe_query(self, query_id: UUID) -> QueryDescription: - """Describe a submitted SQL query by its connection-scoped id.""" + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: + """Describe a submitted SQL query by its UUID or string ID. + + For a handle with ``job_type == "query"``, pass its ``id`` on the + connection that submitted it. Registry job IDs are not query IDs. + """ + if isinstance(query_id, str): + query_id = UUID(query_id) return await self._inner.describe_query(query_id) async def namespace_client(self) -> LanceNamespace: diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index e57fde0ea..e61a188f2 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -"""Handles to operations a server may run asynchronously.""" +"""Handles to asynchronous registry jobs, SQL queries, and local operations.""" import asyncio import json @@ -44,12 +44,12 @@ class AsyncJob(Generic[T]): @property def id(self) -> Optional[str]: - """Identifies the operation on the server that is running it. + """The operation's opaque ID, interpreted according to :attr:`job_type`. - Returned for correlating with server logs or the jobs API. Operations - that run in this process have no server id and return `None`. The value - is opaque: parsing it or storing it to resume the job later is not - supported. + Registry jobs expose a registry job ID; SQL-backed jobs expose a + connection-scoped query ID as a string. Local operations return ``None``. + Query IDs cannot be used with registry job APIs or to resume after a + client restart. """ return self._inner.id if self._inner is not None else None @@ -94,9 +94,11 @@ class AsyncJob(Generic[T]): """Ask the backend for this job's current state, and for a server-side job its full record, then cache it for the properties below. - The properties are all `None` until this runs, because submitting an - operation returns only a job id. `status` fetches the whole record too; - `wait` records only the terminal state it establishes. + Registry details are `None` until this runs, because submitting an + operation does not fetch its full registry record. Query-backed jobs + expose only lifecycle state through this method. For registry jobs, + `status` fetches the whole record too; `wait` records only the terminal + state it establishes. """ if self._inner is None: return @@ -114,11 +116,13 @@ class AsyncJob(Generic[T]): @property def job_type(self) -> Optional[str]: - """The job's type, as the server names it. + """The job's kind of work, without contacting the backend. - `None` for an in-process job, which has no server-side record. + SQL-backed handles report ``"query"`` immediately. Registry handles + report the server's job type after :meth:`refresh` or :meth:`status`. + In-process handles return ``None``. """ - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: @@ -294,7 +298,7 @@ class Job(Generic[T]): @property def job_type(self) -> Optional[str]: """The job's type. See :attr:`AsyncJob.job_type`.""" - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 55baa39a3..992926769 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -1683,7 +1683,7 @@ class AsyncLanceNamespaceDBConnection: parameters=parameters, ) - async def describe_query(self, query_id: UUID) -> QueryDescription: + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query when supported.""" return await self._inner.describe_query(query_id) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 5e9288471..62d00f284 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -1134,7 +1134,7 @@ class RemoteDBConnection(DBConnection): ) @override - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" return LOOP.run( self._conn.describe_query( diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 874034564..61e511345 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1729,7 +1729,11 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip(): ) registration = db.create_function_async(normalize_score) assert registration.id == "job-register" + assert registration.job_type is None + assert registration.status() == "finished" + assert registration.job_type == "create_function" created = registration.wait() + assert registration.job_type == "create_function" reopened = db.get_function("normalize_score", version=created.version) assert created == reopened diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index 46972a15c..e2be9c0d4 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -8,6 +8,7 @@ import threading import lancedb import pytest +from lancedb.db import AsyncConnection from lancedb.materialized_view import MaterializedViewDefinition from lancedb.remote.db import RemoteDBConnection @@ -56,7 +57,7 @@ def mock_remote_materialized_views(): @contextlib.contextmanager -def mock_remote_materialized_view_create(): +def mock_remote_materialized_view_drop(): requests = [] class Handler(http.server.BaseHTTPRequestHandler): @@ -67,7 +68,8 @@ def mock_remote_materialized_view_create(): length = int(self.headers.get("Content-Length", 0)) body = json.loads(self.rfile.read(length) or b"{}") requests.append((self.path, body)) - job_id = "mv-drop-123" if self.path.endswith("/drop") else "mv-create-123" + assert self.path.endswith("/drop") + job_id = "mv-drop-123" encoded = json.dumps({"job_id": job_id}).encode() self.send_response(202) self.send_header("Content-Type", "application/json") @@ -98,7 +100,7 @@ def test_remote_list_uses_namespace_route(): def test_remote_drop_async_returns_server_job(): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect( "db://dev", api_key="fake", @@ -192,11 +194,13 @@ def test_create_and_refresh_jobs(tmp_path): "adults", "people", where="age >= 18", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert create_job.wait() is None view = db.open_materialized_view("adults") refresh_job = view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 2 @@ -337,11 +341,13 @@ async def test_async_create_and_refresh_jobs(tmp_path): "adults", "people", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert await create_job.wait() is None view = await db.open_materialized_view("adults") refresh_job = await view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = await refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 1 @@ -567,6 +573,27 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): assert list(resumed) == [] +class QueryBackedCreateJob: + id = "0198f1b2-c3d4-7e5f-8123-456789abcdef" + + job_type = "query" + + async def wait(self): + return None + + async def status(self): + return "finished" + + +class CapturingCreateConnection: + def __init__(self): + self.calls = [] + + async def create_materialized_view_async(self, name, source, **options): + self.calls.append((name, source, options)) + return QueryBackedCreateJob() + + @pytest.mark.parametrize( "options", [ @@ -576,16 +603,28 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): ], ) def test_remote_create_forwards_refresh_options_without_interval_default(options): - with mock_remote_materialized_view_create() as (host, requests): - db = lancedb.connect("db://dev", api_key="fake", host_override=host) - job = db.create_materialized_view_async("adults", "people", **options) - assert job.id == "mv-create-123" - assert requests == [ + native = CapturingCreateConnection() + db = RemoteDBConnection.__new__(RemoteDBConnection) + db._conn = AsyncConnection(native) + job = db.create_materialized_view_async("adults", "people", **options) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert job.wait() is None + assert job.status() == "finished" + assert native.calls == [ ( - "/v1/materialized_view/adults/create", + "adults", + "people", { - "query": 'SELECT * FROM "people"', - **options, + "projections": None, + "filter": None, + "limit": None, + "with_no_data": False, + "auto_refresh_enabled": options.get("auto_refresh_enabled"), + "auto_refresh_interval_seconds": options.get( + "auto_refresh_interval_seconds" + ), }, ) ] @@ -593,18 +632,18 @@ def test_remote_create_forwards_refresh_options_without_interval_default(options @pytest.mark.asyncio async def test_async_remote_create_keeps_interval_unspecified(): - with mock_remote_materialized_view_create() as (host, requests): - db = await lancedb.connect_async("db://dev", api_key="fake", host_override=host) - job = await db.create_materialized_view_async( - "adults", - "people", - auto_refresh_enabled=True, - ) - assert job.id == "mv-create-123" - assert requests[0][1] == { - "query": 'SELECT * FROM "people"', - "auto_refresh_enabled": True, - } + native = CapturingCreateConnection() + db = AsyncConnection(native) + job = await db.create_materialized_view_async( + "adults", "people", auto_refresh_enabled=True + ) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert await job.wait() is None + assert await job.status() == "finished" + assert native.calls[0][2]["auto_refresh_enabled"] is True + assert native.calls[0][2]["auto_refresh_interval_seconds"] is None @pytest.mark.parametrize( @@ -625,7 +664,7 @@ def test_embedded_create_refuses_automatic_refresh_before_creating_view( @pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf")]) def test_remote_create_rejects_invalid_interval_without_request(seconds): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect("db://dev", api_key="fake", host_override=host) with pytest.raises(ValueError, match="positive, finite"): db.create_materialized_view_async( diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index e6f5fd3c7..539b88647 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -137,9 +137,23 @@ async def test_async_namespace_connection_rejects_sql(tmp_path): await connection.describe_query(NIL_QUERY_ID) -def test_describe_query_requires_uuid(): - with pytest.raises(TypeError, match="UUID"): - remote_connection().describe_query(str(NIL_QUERY_ID)) +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", [NIL_QUERY_ID, str(NIL_QUERY_ID)]) +async def test_describe_query_accepts_uuid_and_job_id_string(identity): + calls = [] + + class NativeConnection: + async def describe_query(self, query_id): + calls.append(query_id) + + connection = AsyncConnection(NativeConnection()) + await connection.describe_query(identity) + assert calls == [NIL_QUERY_ID] + + +def test_describe_query_rejects_registry_job_id(): + with pytest.raises(ValueError, match="UUID"): + remote_connection().describe_query("registry-job-id") @pytest.mark.parametrize( diff --git a/python/src/job.rs b/python/src/job.rs index e22b2f897..1740b3e00 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -81,6 +81,11 @@ impl Job { self.inner.id().map(str::to_string) } + #[getter] + pub fn job_type(&self) -> Option { + self.inner.job_type() + } + pub fn status(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); future_into_py( diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index d671365f9..5c27587dc 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -13,6 +13,7 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; use tokio::sync::watch; use tokio::task::{AbortHandle, JoinHandle}; +use uuid::Uuid; use crate::database::JobDescription; use crate::error::{Error, JobFailure, Result}; @@ -55,10 +56,13 @@ fn job_detail_not_supported(what: &str) -> Result { /// Backend-specific tracking for an asynchronous operation. #[async_trait] pub(crate) trait JobHandle: Send + Sync { - /// Server-assigned id, when the backend has one. + /// Backend-specific id, when the operation has one. fn id(&self) -> Option<&str> { None } + fn job_type(&self) -> Option<&str> { + None + } async fn status(&self) -> Result; async fn wait(&self) -> Result; async fn cancel(&self) -> Result<()>; @@ -167,14 +171,15 @@ struct JobCache { /// the endpoint's successful terminal result; unit-result operations use the /// default `Job<()>`. /// -/// The detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the +/// Registry detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the /// handle last observed. Submitting an operation returns only a job id, so /// populating them eagerly would cost an extra round trip on every call: /// /// - [`Job::refresh`] and [`Job::status`] fetch the whole record. /// - [`Job::wait`] records the terminal state it establishes, but not the rest /// of the record; call [`Job::refresh`] for that. -/// - Everything is `None` until one of those runs. +/// - Registry details are `None` until one of those runs. Query-backed handles +/// report [`Job::job_type`] as `"query"` immediately. pub struct Job where T: Clone + Send + Sync + 'static, @@ -273,13 +278,18 @@ where Self::new_typed(Box::new(SpawnedJob::new(task))) } - pub(crate) fn spawned_with_cancellation(task: JoinHandle>, cancel: F) -> Self + pub(crate) fn from_query_task( + query_id: Uuid, + task: JoinHandle>, + cancel: F, + ) -> Self where F: Fn() -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { let cancellation: Cancellation = Arc::new(move || Box::pin(cancel())); Self::new_typed(Box::new(SpawnedJob::new_with_cancellation( + query_id, task, cancellation, ))) @@ -290,12 +300,10 @@ impl Job where T: Clone + Send + Sync + 'static, { - /// Identifies the operation on the server that is running it. - /// - /// Returned for correlating with server logs or the jobs API. Operations - /// that run in this process have no server id and return `None`. The - /// value is opaque: parsing it or storing it to resume the job later is - /// not supported. + /// The operation's opaque identity, interpreted according to [`Job::job_type`]. + /// Registry handles expose a registry job ID; query handles expose the + /// connection-scoped SQL query UUID as a string. Local handles return None. + /// A query ID cannot be used with registry APIs or to resume after restart. pub fn id(&self) -> Option<&str> { match &self.inner { JobInner::Handle { handle, .. } => handle.id(), @@ -369,8 +377,14 @@ where self.cache_read().description.clone() } - /// The job's type, as the server names it. `None` for an in-process job. + /// The job's type: `"query"` for SQL-backed handles, otherwise the cached + /// server-reported kind of work. `None` for an in-process job. pub fn job_type(&self) -> Option { + if let JobInner::Handle { handle, .. } = &self.inner + && let Some(job_type) = handle.job_type() + { + return Some(job_type.to_string()); + } self.with_description(|description| description.job_type.clone()) } @@ -536,6 +550,7 @@ struct SpawnedJob { outcome: watch::Receiver>, abort: AbortHandle, cancellation: Option, + id: Option, } type CancellationFuture = Pin> + Send>>; @@ -549,11 +564,17 @@ impl SpawnedJob { Self::new_inner(task, None) } - fn new_with_cancellation(task: JoinHandle>, cancellation: Cancellation) -> Self + fn new_with_cancellation( + query_id: Uuid, + task: JoinHandle>, + cancellation: Cancellation, + ) -> Self where T: Serialize + Send + 'static, { - Self::new_inner(task, Some(cancellation)) + let mut handle = Self::new_inner(task, Some(cancellation)); + handle.id = Some(query_id.to_string()); + handle } fn new_inner(task: JoinHandle>, cancellation: Option) -> Self @@ -582,12 +603,21 @@ impl SpawnedJob { outcome, abort, cancellation, + id: None, } } } #[async_trait] impl JobHandle for SpawnedJob { + fn id(&self) -> Option<&str> { + self.id.as_deref() + } + + fn job_type(&self) -> Option<&str> { + self.id.as_ref().map(|_| "query") + } + async fn status(&self) -> Result { let label = match &*self.outcome.borrow() { None => "running", @@ -648,11 +678,31 @@ mod tests { async fn mapped_spawned_job_reuses_outcome() { let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1); + assert_eq!(job.job_type(), None); + assert_eq!(job.id(), None); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.status().await.unwrap(), "finished"); } + #[tokio::test] + async fn mapped_query_job_preserves_identity_and_type() { + let query_id = Uuid::new_v4(); + let job = Job::from_query_task(query_id, tokio::spawn(async { Ok(41_u64) }), || async { + Ok(()) + }) + .map(|value| value + 1); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + let completed = Job::new_done(); + assert_eq!(completed.job_type(), None); + assert_eq!(completed.id(), None); + } + #[tokio::test] async fn mapped_spawned_job_preserves_cancellation() { let job = Job::spawned(tokio::spawn(async { pending::>().await })) diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index ecdd1783b..7ee76a160 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -1716,7 +1716,8 @@ impl CreateMaterializedViewBuilder { } /// Submit creation and initial population, returning a [`Job`] that - /// settles when the view is ready. + /// settles when the view is ready. Remote creation submits SQL; the handle's + /// [`Job::job_type`] is `"query"` and its ID identifies that query. pub async fn execute_async(self) -> Result { if self.connection.uri().starts_with("db://") { return self diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 768de79d2..d33565653 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -63,15 +63,6 @@ struct RemoteCloneTableRequest { is_shallow: Option, } -#[derive(serde::Serialize)] -struct RemoteCreateMaterializedViewRequest { - query: String, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_enabled: Option, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_interval_seconds: Option, -} - // the versions of the server that we support // for any new feature that we need to change the SDK behavior, we should bump the server version, // and add a feature flag as method of `ServerVersion` here. @@ -594,6 +585,11 @@ mod test_utils { } } + pub(in crate::remote) fn with_sql_client(mut self, client: SqlClient) -> Self { + self.sql_client = Some(client); + self + } + pub fn new_mock_with_config(handler: F, config: ClientConfig) -> Self where F: Fn(reqwest::Request) -> http::Response + Send + Sync + 'static, @@ -885,8 +881,7 @@ impl Database for RemoteDatabase { ) -> Result { if request.with_no_data { return Err(Error::NotSupported { - message: "remote materialized views are always populated by their SQL job" - .to_string(), + message: "remote materialized views must be populated during creation".to_string(), }); } if let Some(seconds) = request.auto_refresh_interval_seconds @@ -896,36 +891,37 @@ impl Database for RemoteDatabase { message: "auto_refresh_interval_seconds must be positive, finite seconds".into(), }); } - let identifier = build_table_identifier(&request.name, &request.namespace_path)?; - let body = RemoteCreateMaterializedViewRequest { - query: request.query, - auto_refresh_enabled: request.auto_refresh_enabled, - auto_refresh_interval_seconds: request.auto_refresh_interval_seconds, - }; - let request = self - .client - .post(&format!("/v1/materialized_view/{identifier}/create")) - .json(&body); - let (request_id, response) = self.client.send(request).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - if status != StatusCode::ACCEPTED { - return Err(Error::Http { - source: "materialized-view creation must return 202 Accepted".into(), - request_id, - status_code: Some(status), - }); + validate_table_name(&request.name)?; + validate_namespace(&request.namespace_path)?; + let client = self + .sql_client + .as_ref() + .ok_or_else(|| Error::NotSupported { + message: "SQL is unavailable for this remote database client".into(), + })?; + let mut options = Vec::new(); + if let Some(enabled) = request.auto_refresh_enabled { + options.push(format!("auto_refresh_enabled = {enabled}")); } - let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { - source: "materialized-view create response did not contain a valid job_id".into(), - request_id, - status_code: Some(status), - })?; - Ok(Job::new(Box::new(RemoteJob::new( - self.client.clone(), - job_id, - )))) + if let Some(seconds) = request.auto_refresh_interval_seconds { + options.push(format!("auto_refresh_interval_seconds = {seconds}")); + } + let name = request.name.replace('"', "\"\""); + let options = if options.is_empty() { + String::new() + } else { + format!(" WITH ({})", options.join(", ")) + }; + client + .submit_as_job( + format!( + "CREATE MATERIALIZED VIEW \"{name}\"{options} AS {}", + request.query + ), + request.namespace_path, + || async { Ok(()) }, + ) + .await } async fn drop_materialized_view_async( @@ -2016,72 +2012,22 @@ mod tests { } #[tokio::test] - async fn test_create_materialized_view_uses_rest_and_preserves_optional_settings() { - for (enabled, interval) in [(None, None), (Some(true), None), (Some(false), Some(0.5))] { - let db = super::RemoteDatabase::new_mock(move |request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/materialized_view/analytics$adults/create" - ); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(body["query"], "SELECT * FROM people"); - assert_eq!( - body.get("auto_refresh_enabled"), - enabled.map(serde_json::Value::from).as_ref() - ); - assert_eq!( - body.get("auto_refresh_interval_seconds"), - interval.map(serde_json::Value::from).as_ref() - ); - http::Response::builder() - .status(202) - .body(serde_json::json!({"job_id": "j1-mv-create"}).to_string()) - .unwrap() - }); - let job = db - .create_materialized_view_async(CreateMaterializedViewRequest { - name: "adults".into(), - namespace_path: vec!["analytics".into()], - query: "SELECT * FROM people".into(), - with_no_data: false, - auto_refresh_enabled: enabled, - auto_refresh_interval_seconds: interval, - }) - .await - .unwrap(); - assert_eq!(job.id(), Some("j1-mv-create")); - } - } - - #[tokio::test] - async fn test_create_materialized_view_rejects_invalid_acceptance() { - for (status, body) in [ - (200, r#"{"job_id":"j1"}"#), - (202, "{}"), - (202, r#"{"job_id":""}"#), - ] { - let db = super::RemoteDatabase::new_mock(move |_| { - http::Response::builder().status(status).body(body).unwrap() - }); - let error = db + async fn test_create_materialized_view_validates_before_submission() { + let db = super::RemoteDatabase::new_mock(|_| -> http::Response { + panic!("materialized-view creation must not send REST requests") + }); + for seconds in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let result = db .create_materialized_view_async(CreateMaterializedViewRequest { name: "adults".into(), namespace_path: vec![], query: "SELECT * FROM people".into(), with_no_data: false, auto_refresh_enabled: None, - auto_refresh_interval_seconds: None, + auto_refresh_interval_seconds: Some(seconds), }) - .await - .err() - .unwrap(); - assert!(error.to_string().contains(if status == 200 { - "202 Accepted" - } else { - "job_id" - })); + .await; + assert!(matches!(result, Err(Error::InvalidInput { .. }))); } } diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index d8f04119f..846e852f5 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::fs; -use std::future::{Future, pending}; +use std::future::Future; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::{Duration, Instant}; @@ -12,6 +12,7 @@ use arrow_array::RecordBatch; use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::error::FlightError; use arrow_flight::flight_service_client::FlightServiceClient; +#[cfg(test)] use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; use arrow_flight::{ Action, CancelFlightInfoRequest, CancelFlightInfoResult, CancelStatus, FlightClient, @@ -22,7 +23,7 @@ use futures::TryStreamExt; use http::header::{HeaderMap, HeaderName, HeaderValue}; use prost::Message; use serde::{Serialize, de::DeserializeOwned}; -use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch}; +use tokio::sync::{Mutex, Notify, OnceCell, mpsc}; use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}; use uuid::Uuid; @@ -226,6 +227,7 @@ impl SqlClient { .await } + #[cfg(test)] pub(super) async fn submit( &self, query: &str, @@ -254,40 +256,36 @@ impl SqlClient { .await } - pub(super) fn submit_as_job( + pub(super) async fn submit_as_job( &self, statement: String, namespace: Vec, finish: F, - ) -> Job + ) -> Result> where T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, F: FnOnce() -> Fut + Send + 'static, Fut: Future> + Send + 'static, { - let client = self.clone(); - let (query_tx, query_rx) = watch::channel(None::>); + let query = Arc::new( + self.execute(ExecuteQueryRequest { + query: statement, + default_namespace_path: namespace, + parameters: None, + }) + .await?, + ); + let query_id = query.id(); + let execution = query.clone(); let task = tokio::spawn(async move { - let query = Arc::new(client.submit(&statement, &namespace).await?); - query_tx.send_replace(Some(query.clone())); - let mut reader = query.reader().await?; + let mut reader = execution.reader().await?; while reader.try_next().await?.is_some() {} finish().await }); - Job::spawned_with_cancellation(task, move || { - let mut query_rx = query_rx.clone(); - async move { - loop { - let query = query_rx.borrow().clone(); - if let Some(query) = query { - return query.cancel().await; - } - if query_rx.changed().await.is_err() { - pending::<()>().await; - } - } - } - }) + Ok(Job::from_query_task(query_id, task, move || { + let query = query.clone(); + async move { query.cancel().await } + })) } pub(super) async fn describe(&self, query_id: Uuid) -> Result { @@ -315,6 +313,7 @@ impl SqlClientInner { ) } + #[cfg(test)] async fn poll( &self, descriptor: FlightDescriptor, diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 4837baab7..1368dd753 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -22,6 +22,7 @@ use tonic::{Request, Response, Status, Streaming}; use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse}; use super::*; use crate::database::Database; +use crate::materialized_view::CreateMaterializedViewRequest; use crate::remote::RemoteCatalogOptions; use crate::remote::client::HeaderProvider; use crate::remote::client::test_utils::client_with_handler; @@ -75,6 +76,7 @@ struct TestSqlService { defer: bool, invalid_decision: bool, query_count: Arc, + statements: Arc>>, do_get_count: Arc, cancel_count: Arc, cancel_denied_count: Arc, @@ -207,6 +209,7 @@ impl Default for TestSqlService { defer: false, invalid_decision: false, query_count: Arc::new(AtomicUsize::new(0)), + statements: Arc::new(std::sync::Mutex::new(Vec::new())), do_get_count: Arc::new(AtomicUsize::new(0)), cancel_count: Arc::new(AtomicUsize::new(0)), cancel_denied_count: Arc::new(AtomicUsize::new(0)), @@ -269,6 +272,7 @@ impl FlightService for TestSqlService { .and_then(|any| any.unpack::().ok().flatten()); let (query, stage) = if let Some(command) = command { self.query_count.fetch_add(1, Ordering::SeqCst); + self.statements.lock().unwrap().push(command.query.clone()); (command.query, 0_u8) } else { let continuation = std::str::from_utf8(request.get_ref().cmd.as_ref()) @@ -289,6 +293,9 @@ impl FlightService for TestSqlService { .ok_or_else(|| Status::invalid_argument("invalid continuation"))?; (query.to_string(), stage) }; + if query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } if (query == "SELECT slow" || query == "SELECT cancelling") && stage > 0 { tokio::time::sleep(Duration::from_millis(250)).await; } @@ -364,7 +371,11 @@ impl FlightService for TestSqlService { .ok() .and_then(|envelope| Any::decode(envelope.value).ok()) .and_then(|command| command.unpack::().ok().flatten()); - if initial.is_some() { + if let Some(command) = &initial { + self.statements.lock().unwrap().push(command.query.clone()); + if command.query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } self.capture_headers(request.metadata()); assert_eq!( request.metadata().get(PROTOCOL_HEADER).unwrap(), @@ -376,7 +387,8 @@ impl FlightService for TestSqlService { .map(|command| command.query.as_bytes()) .unwrap_or(request.get_ref().ticket.as_ref()); let ticket = query; - let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; + let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"" + || ticket == b"SELECT job cancel denied"; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; let large = ticket == b"SELECT large message"; @@ -458,7 +470,9 @@ impl FlightService for TestSqlService { { return Err(Status::not_found("query cancellation completed")); } - let status = if query == "SELECT cancel unspecified" { + let status = if query == "SELECT job cancel denied" { + CancelStatus::NotCancellable + } else if query == "SELECT cancel unspecified" { CancelStatus::Unspecified } else if query == "SELECT cancelling" { if self @@ -880,6 +894,132 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { server.await.unwrap().unwrap(); } +#[tokio::test] +async fn remote_create_uses_sql_and_preserves_options_and_namespace() { + for defer in [false, true] { + let service = TestSqlService { + adaptive: true, + defer, + ..Default::default() + }; + let statements = service.statements.clone(); + let headers = service.headers.clone(); + let count = service.query_count.clone(); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + panic!("CREATE must use SQL, not REST: {}", request.url().path()); + }) + .with_sql_client(sql_client.clone()); + for (enabled, interval, expected) in [ + (None, None, ""), + (Some(true), None, " WITH (auto_refresh_enabled = true)"), + ( + Some(false), + Some(0.5), + " WITH (auto_refresh_enabled = false, auto_refresh_interval_seconds = 0.5)", + ), + ] { + let job = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["team".into(), "reports".into()], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: enabled, + auto_refresh_interval_seconds: interval, + }) + .await + .unwrap(); + assert_eq!(job.job_type().as_deref(), Some("query")); + let query_id = + Uuid::parse_str(job.id().expect("SQL-backed job must expose its ID")).unwrap(); + job.wait().await.unwrap(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + assert!(job.events(Default::default()).await.is_err()); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!( + sql_client.describe(query_id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + statements.lock().unwrap().last().unwrap(), + &format!("CREATE MATERIALIZED VIEW \"adults\"{expected} AS SELECT * FROM people") + ); + } + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(statements.lock().unwrap().len(), 3); + assert!( + headers + .lock() + .unwrap() + .iter() + .all(|header| header.namespace_path == "team$reports") + ); + shutdown.send(()).unwrap(); + } +} + +#[tokio::test] +async fn query_backed_jobs_propagate_submission_and_execution_failures() { + let (client, shutdown) = start_sql_service( + TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }, + ClientConfig::default(), + ) + .await; + let rejected = client + .submit_as_job("SELECT rejected".into(), vec![], || async { Ok(()) }) + .await; + assert!(rejected.is_err()); + let finished = Arc::new(AtomicBool::new(false)); + let callback_finished = finished.clone(); + let job = client + .submit_as_job("SELECT fail".into(), vec![], move || async move { + callback_finished.store(true, Ordering::SeqCst); + Ok(()) + }) + .await + .unwrap(); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert_eq!(job.status().await.unwrap(), "failed"); + assert!(!finished.load(Ordering::SeqCst)); + shutdown.send(()).unwrap(); +} + +#[tokio::test] +async fn rejected_job_cancellation_does_not_abort_sql_completion() { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let do_get_count = service.do_get_count.clone(); + let release = service.gate_refresh_release.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let job = client + .submit_as_job("SELECT job cancel denied".into(), vec![], || async { + Ok(()) + }) + .await + .unwrap(); + wait_for(&do_get_count, 1, "job did not begin reading its SQL result").await; + assert!(matches!( + job.cancel().await, + Err(Error::NotSupported { .. }) + )); + assert_eq!(job.status().await.unwrap(), "running"); + release.notify_one(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + job.cancel().await.unwrap(); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn remote_refresh_twice_keeps_latest_selector() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -969,7 +1109,11 @@ async fn cancelling_remote_refresh_cancels_sql_query() { let (socket, _) = listener.accept().await?; Ok::<_, std::io::Error>(Some((socket, listener))) }); - let service = TestSqlService::default(); + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; let cancel_count = service.cancel_count.clone(); let do_get_count = service.do_get_count.clone(); let gate_refresh_release = service.gate_refresh_release.clone(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index b722dc4a1..4cbba04df 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2247,8 +2247,8 @@ impl BaseTable for RemoteTable { let freshness = self.freshness.clone(); let freshness_request = self.snapshot_freshness_headers(); let job_sql_client = sql_client.clone(); - Ok( - sql_client.submit_as_job(statement, namespace.clone(), move || async move { + sql_client + .submit_as_job(statement, namespace.clone(), move || async move { let table = Self::new_with_sql_client( client, name, @@ -2270,8 +2270,8 @@ impl BaseTable for RemoteTable { source_version: 0, version, }) - }), - ) + }) + .await } async fn query_snapshot(&self) -> Result> { let description = self.describe().await?; From dcfef7c5ab259d2ce36003b6d6bf10c6de22cd9e Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 25 Oct 2026 02:31:21 -0700 Subject: [PATCH 64/64] feat(python): open lazy blob files from SQL descriptors (#4668) SQL blob descriptors need pinned source identity to reopen a value after the query finishes. Adds `open_blob_file(descriptor)` and `open_blob_files(descriptors)` so Python can read the original blob incrementally without materializing it in the SQL result. ### Python experience ```python rows = db.execute_query(""" SELECT f.timestamp, blob_descriptor(v.video) AS video FROM videos v JOIN frames f ON v.id = f.video_id """).read_all() videos = db.open_table("videos") for row in rows.to_pylist(): if row["video"] is not None: with videos.open_blob_file(row["video"]) as video: frame = extract_frame(video, row["timestamp"]) ``` The singular API accepts the descriptor dictionary or an Arrow struct scalar directly and returns a seekable `BlobFile`. `videos.open_blob_files(rows["video"])` accepts an Arrow array/chunked array or another iterable, opens handles as iteration advances, and preserves order, duplicates, and nulls as `None`. Async tables provide `await videos.open_blob_file(...)` and `async for video in videos.open_blob_files(...)`. The iterable API can be used directly with the Arrow result column. Zip it with the rows to keep each frame timestamp: ```python for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): if video is not None: with video: frame = extract_frame(video, row["timestamp"]) ``` Each iteration opens one handle. Null descriptors yield `None`; handles preserve input order and duplicates, and the `with` block closes each non-null file. For an async table: ```python async for video in async_videos.open_blob_files(rows["video"]): if video is not None: with video: payload = await video.aread() ``` ### Execution The descriptor retains the canonical `kind`, `position`, `size`, `blob_id`, and `blob_uri` fields and carries a `source` struct with its dataset URI, version, commit timestamp, branch, column, row reference, and whether that reference is a physical address. Native opening resolves that exact snapshot without changing the table handle's checkout and rejects a descriptor belonging to another table. Remote handles send authenticated byte-range requests with the captured version/branch, commit timestamp, and source URI, so an old physical address is never interpreted against the latest table. Known-size remote descriptors open without a payload request. External blobs with an unspecified size use a metadata-only HEAD request. Payload bytes load only when the file is read; full reads still allocate the full value. Existing row-ID-based `fetch_blob_files` remains available when a descriptor is not already present. Coverage includes native snapshot preservation after overwrite, both row IDs and addresses, dictionary/Arrow inputs, null/empty values, reordered and duplicate descriptors, lazy iterable behavior, async iteration, and remote range pinning for all six storage kinds. Descriptor reads isolate their pinned snapshot from the live table handle's freshness state. Both native and remote reads verify the source commit timestamp to reject a replaced snapshot after table recreation. --- docs/src/python/python.md | 54 +++++++- python/README.md | 49 +++++++ python/python/lancedb/_blob.py | 14 +- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/remote/table.py | 5 + python/python/lancedb/table.py | 76 ++++++++++ python/python/tests/test_blob.py | 132 ++++++++++++++++++ python/src/table.rs | 19 ++- rust/lancedb/src/blob.rs | 111 ++++++++++++++- rust/lancedb/src/remote/client.rs | 5 + rust/lancedb/src/remote/table.rs | 4 + rust/lancedb/src/remote/table/blobs.rs | 183 ++++++++++++++++++++++++- rust/lancedb/src/table.rs | 33 ++++- 13 files changed, 674 insertions(+), 12 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 41d3e16d7..02ccced2c 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -538,12 +538,64 @@ instead of being materialized with the rest of the row. `lancedb.BlobType` is `lance.blob.BlobType` when pylance is installed. Without pylance, LanceDB uses a matching `lance.blob.v2` extension type so blob columns -still work. Queries return descriptors. Call +still work. Native table queries return descriptors. Call [`fetch_blob_files`][lancedb.table.Table.fetch_blob_files] for lazy reads or [`fetch_blobs`][lancedb.table.Table.fetch_blobs] for eager bytes. One missing or deleted row ID rejects the entire blob-read batch; the APIs do not return partial results. +Remote SQL returns blob bytes by default. Select `blob_descriptor(video)` to +return an openable reference instead, then use +[`open_blob_file`][lancedb.table.Table.open_blob_file] or the lazy iterable +[`open_blob_files`][lancedb.table.Table.open_blob_files]: + +```python +rows = db.execute_query(""" + SELECT f.timestamp, blob_descriptor(v.video) AS video + FROM videos v JOIN frames f ON v.id = f.video_id +""").read_all() +videos = db.open_table("videos") +for row in rows.to_pylist(): + if row["video"] is not None: + with videos.open_blob_file(row["video"]) as video: + frame = extract_frame(video, row["timestamp"]) +``` + +This example uses a remote SQL connection that produces source-bearing +descriptors. The descriptor contains `kind`, `position`, `size`, `blob_id`, +`blob_uri`, and a `source` struct with `uri`, `version`, `committed_at_micros`, +`branch`, `column`, `row_reference`, and `row_is_address`. The opener verifies +the source table and commit and reads the captured version without changing +the table handle's checkout. Raw native storage descriptors omit `source`; +use `fetch_blob_files(column, row_ids)` for those query results. + +The singular API accepts a dictionary or Arrow struct scalar. The iterable +accepts Arrow arrays/chunked arrays or another iterable, opens handles as +iteration advances, and preserves order, duplicates, and nulls as `None`. +Close each non-null handle after use. Async tables provide +`await table.open_blob_file(...)` and +`async for file in table.open_blob_files(...)`. + +```python +for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): + if video is not None: + with video: + frame = extract_frame(video, row["timestamp"]) +``` + +For an async table, the iterable is consumed with `async for`: + +```python +async for video in async_videos.open_blob_files(rows["video"]): + if video is not None: + with video: + payload = await video.aread() +``` + +Payload bytes load on reads. Opening an external blob with unspecified size +can issue a metadata-only HEAD request. Remote reads use authenticated range +requests that validate the captured URI and commit timestamp. + Blobs inside lists can be written and queried as descriptors. They are omitted from `blob_columns()` because row ids do not identify individual list elements, so the fetch APIs cannot read their bytes yet. diff --git a/python/README.md b/python/README.md index 907f7c099..0b640d27c 100644 --- a/python/README.md +++ b/python/README.md @@ -101,3 +101,52 @@ print(results) ### Development See [CONTRIBUTING.md](./CONTRIBUTING.md) for information on how to contribute to LanceDB. + +### Opening SQL blob descriptors + +A SQL `blob_descriptor(video)` result keeps the Blob V2 storage metadata and a +`source` struct containing its table URI, version, commit timestamp, branch, +column and row identity. This example uses a remote SQL connection whose SQL +service produces those source-bearing descriptors. Open the reference on the +source table to read incrementally: + +```python +rows = db.execute_query(""" + SELECT f.timestamp, blob_descriptor(v.video) AS video + FROM videos v JOIN frames f ON v.id = f.video_id +""").read_all() + +videos = db.open_table("videos") +for row in rows.to_pylist(): + if row["video"] is not None: + with videos.open_blob_file(row["video"]) as video: + frame = extract_frame(video, row["timestamp"]) +``` + +`open_blob_file` accepts the dictionary or an Arrow struct scalar directly. It +keeps the query's version and branch without changing the table handle's checkout. +Payload bytes load on reads. An external blob with an unspecified size may need a +metadata-only HEAD request to open. + +`videos.open_blob_files(rows["video"])` returns an iterator over the same handles, +preserving order, duplicates and nulls (`None`). It accepts Arrow arrays/chunked +arrays or another iterable and opens each handle as iteration advances. Close each +non-null handle after use. Async tables offer `await videos.open_blob_file(...)` +and `async for video in videos.open_blob_files(...)`. + +```python +for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + header = video.read(4096) +``` + +To extract a frame while retaining each row's timestamp, iterate over the +handles and rows together: + +```python +for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): + if video is not None: + with video: + frame = extract_frame(video, row["timestamp"]) +``` diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 1e9110dd0..46462553c 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -7,7 +7,7 @@ from __future__ import annotations import io import operator -from collections.abc import Awaitable, Callable, Iterable +from collections.abc import Awaitable, Callable, Iterable, Mapping from typing import TYPE_CHECKING, Optional, Union import pyarrow as pa @@ -31,6 +31,18 @@ FetchBlobsSync = Callable[[str, list[int]], pa.Array | pa.ChunkedArray] FetchBlobsAsync = Callable[[str, list[int]], Awaitable[pa.Array | pa.ChunkedArray]] +def _normalize_blob_descriptor(value): + if isinstance(value, pa.Scalar): + value = value.as_py() + if value is None: + return None + if not isinstance(value, Mapping): + raise TypeError( + "blob descriptor must be a mapping, an Arrow struct scalar, or None" + ) + return dict(value) + + class BlobFile(io.RawIOBase): """Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`. diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 9e602cce8..0b4fc82f1 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -580,6 +580,7 @@ class Table: async def fetch_blob_files( self, column: str, row_ids: list[int] ) -> list[Optional[BlobFile]]: ... + async def open_blob_file(self, descriptor: str) -> BlobFile: ... def vector_search(self) -> VectorQuery: ... class Tags: diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 53591182e..dd22dea48 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1221,6 +1221,11 @@ class RemoteTable(Table): ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + return LOOP.run(self._table.open_blob_file(descriptor)) + def head(self, n=5) -> pa.Table: """ Return the first `n` rows of the table. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index c08c6390e..665520479 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -15,9 +15,11 @@ from functools import cached_property from typing import ( TYPE_CHECKING, Any, + AsyncIterator, Callable, Dict, Iterable, + Iterator, List, Literal, Optional, @@ -33,6 +35,7 @@ from lancedb.scannable import _hf_batches, _register_optional_converters, to_sca from . import __version__ from ._blob import ( BlobFile, + _normalize_blob_descriptor, _normalize_blob_row_ids, _wrap_blob_files, strip_auto_row_ids, @@ -2289,6 +2292,44 @@ class Table(ABC): handles are returned. """ + @abstractmethod + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + """Open a seekable file from a SQL ``blob_descriptor`` result. + + Accepts the descriptor dictionary or an Arrow struct scalar. Its source + pins the table version, branch, column and row identity, so opening does + not change this table's checkout. Null descriptors return ``None``. + Payload bytes load only when the returned handle is read. An external + blob with an unspecified size may require a metadata request on opening. + + ```python + rows = db.execute_query( + "SELECT blob_descriptor(video) AS video FROM videos" + ).read_all() + with videos.open_blob_file(rows["video"][0]) as video: + header = video.read(4096) + ``` + """ + + def open_blob_files(self, descriptors: Iterable) -> Iterator[Optional[BlobFile]]: + """Iterate over lazy files from descriptor dictionaries or Arrow scalars. + + Accepts an Arrow array/chunked array or any iterable. Opens each handle + as iteration advances, preserving input order and nulls as ``None``. + Close each non-null handle after use, normally with a ``with`` block. + + ```python + for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + header = video.read(4096) + ``` + """ + for descriptor in descriptors: + yield self.open_blob_file(descriptor) + @abstractmethod def _execute_query( self, @@ -3239,6 +3280,11 @@ class LanceTable(Table): ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + return LOOP.run(self._table.open_blob_file(descriptor)) + @property def tags(self) -> Tags: """Tag management for the table. @@ -7451,6 +7497,36 @@ class AsyncTable: ) return _wrap_blob_files(handles) + async def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + """Open a snapshot-bound descriptor as a seekable file. + + See [Table.open_blob_file][lancedb.table.Table.open_blob_file]. + """ + descriptor = _normalize_blob_descriptor(descriptor) + if descriptor is None: + return None + return BlobFile(await self._inner.open_blob_file(json.dumps(descriptor))) + + async def open_blob_files( + self, descriptors: Iterable + ) -> AsyncIterator[Optional[BlobFile]]: + """Asynchronously iterate over files in descriptor order, preserving nulls. + + Use ``async for file in table.open_blob_files(descriptors)``. Each + non-null file must be closed after use. + + ```python + async for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + payload = await video.aread() + ``` + """ + for descriptor in descriptors: + yield await self.open_blob_file(descriptor) + @property def tags(self) -> AsyncTags: """Tag management for the dataset. diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index c2dc2a265..da362f053 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -2696,3 +2696,135 @@ def test_create_table_with_python_blob_dict_and_object(): b"from-dict", b"from-blob", ] + + +def _openable_descriptors(table, *, by_address): + hits = ( + table.search() + .select(["id", "image", "_rowaddr"]) + .with_row_id(True) + .limit(100) + .to_arrow(blob_mode="descriptions") + ) + version = table.version + committed_at = next( + item["timestamp"] + for item in table.list_versions() + if item["version"] == version + ) + descriptors = {} + for row in hits.to_pylist(): + descriptor = row["image"] + if descriptor is not None and descriptor.get("kind") is not None: + descriptor = dict(descriptor) + descriptor["source"] = { + "uri": table.to_lance().uri, + "version": version, + "committed_at_micros": round(committed_at.timestamp() * 1_000_000), + "branch": None, + "column": "image", + "row_reference": row["_rowaddr" if by_address else "_rowid"], + "row_is_address": by_address, + } + else: + descriptor = None + descriptors[row["id"]] = descriptor + return descriptors + + +@pytest.mark.parametrize("by_address", [False, True]) +def test_open_blob_files_preserve_snapshot_order_nulls_and_empty_values( + tmp_path, by_address +): + db = lancedb.connect(str(tmp_path)) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("videos", schema=schema) + table.add( + [ + {"id": 1, "image": b"alpha"}, + {"id": 2, "image": None}, + {"id": 3, "image": b"beta"}, + {"id": 4, "image": b""}, + ] + ) + descriptors = _openable_descriptors(table, by_address=by_address) + table.add([{"id": 1, "image": b"replacement"}], mode="overwrite") + latest_version = table.version + values = pa.array( + [descriptors[3], descriptors[2], descriptors[1], descriptors[3], descriptors[4]] + ) + files = table.open_blob_files( + pa.chunked_array([values.slice(0, 2), values.slice(2)]) + ) + assert iter(files) is files + actual = [] + for file in files: + if file is None: + actual.append(None) + else: + with file: + actual.append(file.read()) + assert file.closed + assert actual == [b"beta", None, b"alpha", b"beta", b""] + with table.open_blob_file(descriptors[1]) as file: + file.seek(2) + assert file.read_range(0, 2) == b"al" + assert file.tell() == 2 + assert file.read(2) == b"ph" + assert table.version == latest_version + assert table.open_blob_file(pa.scalar(None, type=values.type)) is None + other = db.create_table("other", schema=schema) + with pytest.raises(ValueError, match="different table"): + other.open_blob_file(descriptors[1]) + stale = dict(descriptors[1], source=dict(descriptors[1]["source"])) + stale["source"]["committed_at_micros"] += 1 + with pytest.raises(ValueError, match="snapshot has been replaced"): + table.open_blob_file(stale) + raw = {key: value for key, value in descriptors[1].items() if key != "source"} + with pytest.raises(ValueError, match="source"): + table.open_blob_file(raw) + with pytest.raises(TypeError, match="descriptor"): + table.open_blob_file("not a descriptor") + + +def test_open_blob_files_does_not_open_ahead(monkeypatch): + table = _blob_table("lazy_descriptor_iteration", [{"id": 1, "image": b"x"}]) + opened = [] + + def open_file(descriptor): + opened.append(descriptor) + return descriptor + + monkeypatch.setattr(table, "open_blob_file", open_file) + files = table.open_blob_files(iter([1, None, 2])) + assert opened == [] + assert next(files) == 1 + assert opened == [1] + assert list(files) == [None, 2] + assert opened == [1, None, 2] + + +@pytest.mark.asyncio +async def test_open_blob_files_async_iterates_snapshot_handles(tmp_path): + db = lancedb.connect(str(tmp_path)) + table = db.create_table( + "videos", + schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), + ) + table.add([{"id": 1, "image": b"alpha"}, {"id": 2, "image": None}]) + descriptors = _openable_descriptors(table, by_address=True) + async_db = await lancedb.connect_async(str(tmp_path)) + videos = await async_db.open_table("videos") + file = await videos.open_blob_file(pa.scalar(descriptors[1])) + with file: + assert file.read(2) == b"al" + actual = [] + async for file in videos.open_blob_files( + [descriptors[2], descriptors[1], descriptors[1]] + ): + if file is None: + actual.append(None) + else: + with file: + actual.append(await file.aread()) + assert actual == [None, b"alpha", b"alpha"] diff --git a/python/src/table.rs b/python/src/table.rs index e77c58578..e58b8a139 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -24,7 +24,7 @@ use arrow::{ ffi_stream::ArrowArrayStreamReader, pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; -use lancedb::blob::{BlobFile, BlobRangeRequest}; +use lancedb::blob::{BlobDescriptor, BlobFile, BlobRangeRequest}; use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, @@ -1551,6 +1551,23 @@ impl Table { }) } + #[pyo3(signature = (descriptor))] + pub fn open_blob_file<'py>( + self_: PyRef<'py, Self>, + descriptor: &str, + ) -> PyResult> { + let descriptor: BlobDescriptor = serde_json::from_str(descriptor) + .map_err(|error| PyValueError::new_err(format!("invalid blob descriptor: {error}")))?; + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let file = inner.open_blob_file(&descriptor).await.infer_error()?; + Ok(PyBlobFile { + inner: Arc::new(file), + closed: AtomicBool::new(false), + }) + }) + } + /// Optimize the on-disk data by compacting and pruning old data, for better performance. #[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))] pub fn optimize( diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index 05650d7c2..142b2264e 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -20,11 +20,120 @@ use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams}; use lance_arrow::FieldExt; use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; -use lance_io::object_store::ObjectStore; +use lance_io::object_store::{ObjectStore, uri_to_url}; use object_store::path::Path; +use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; +/// A stored Blob V2 descriptor with the source needed to reopen its snapshot. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct BlobDescriptor { + /// Canonical Blob V2 storage kind (0 through 5). + pub kind: u8, + /// Stored offset within the backing object. + pub position: u64, + /// Stored length; zero on an external blob can mean unspecified. + pub size: u64, + /// Stored sidecar or base identifier. + pub blob_id: u32, + /// Stored external URI or managed path. + pub blob_uri: String, + /// Source identity captured by the query. + pub source: BlobDescriptorSource, +} + +/// The pinned source of a descriptor returned by SQL. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct BlobDescriptorSource { + /// URI of the dataset on the source branch. + pub uri: String, + /// Exact table version read by the query. + pub version: u64, + /// Commit timestamp in microseconds, distinguishing recreated tables. + pub committed_at_micros: u64, + /// Source branch; `None` identifies main. + pub branch: Option, + /// Blob column path in the source table. + pub column: String, + /// Row ID or row address, as specified by `row_is_address`. + pub row_reference: u64, + /// Whether the row reference is a physical address. + pub row_is_address: bool, +} + +impl BlobDescriptor { + pub(crate) fn validate(&self) -> Result<()> { + if self.kind > 5 + || self.source.version == 0 + || self.source.uri.is_empty() + || self.source.column.is_empty() + { + return Err(Error::InvalidInput { + message: "invalid stored blob descriptor or missing pinned source".into(), + }); + } + Ok(()) + } +} + +pub(crate) async fn open_native_blob_file( + dataset: &Dataset, + descriptor: &BlobDescriptor, +) -> Result { + descriptor.validate()?; + let source = &descriptor.source; + let expected_uri = dataset + .branch_location() + .find_branch(source.branch.as_deref())? + .uri; + if uri_to_url(&expected_uri)? != uri_to_url(&source.uri)? { + return Err(Error::InvalidInput { + message: "blob descriptor belongs to a different table".into(), + }); + } + let dataset = Arc::new( + dataset + .checkout_version(lance::dataset::refs::Ref::Version( + source.branch.clone(), + Some(source.version), + )) + .await?, + ); + if uri_to_url(dataset.uri())? != uri_to_url(&source.uri)? { + return Err(Error::InvalidInput { + message: "blob descriptor belongs to a different table".into(), + }); + } + if dataset.manifest().timestamp_nanos / 1_000 != u128::from(source.committed_at_micros) { + return Err(Error::InvalidInput { + message: "blob descriptor source snapshot has been replaced".into(), + }); + } + ensure_blob_v2_column(dataset.schema(), &source.column)?; + let mut files = if source.row_is_address { + dataset + .take_blobs_by_addresses(&[source.row_reference], &source.column) + .await? + } else { + dataset + .take_blobs(&[source.row_reference], &source.column) + .await? + }; + if files.len() != 1 { + return Err(Error::InvalidInput { + message: "blob descriptor does not resolve to one stored value".into(), + }); + } + files + .pop() + .flatten() + .map(Into::into) + .ok_or_else(|| Error::InvalidInput { + message: "non-null blob descriptor resolved to a null value".into(), + }) +} + /// Seekable handle for one blob value, backed by local storage or a remote /// HTTP byte-range endpoint. #[derive(Debug)] diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index fe4a2072e..f54e2fe3e 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -809,6 +809,11 @@ impl RestfulLanceDbClient { self.client.get(full_uri) } + pub fn head(&self, uri: &str) -> RequestBuilder { + let full_uri = format!("{}{}", self.host, uri); + self.client.head(full_uri) + } + pub fn post(&self, uri: &str) -> RequestBuilder { let full_uri = format!("{}{}", self.host, uri); self.client.post(full_uri) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 4cbba04df..181ce61f9 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2927,6 +2927,10 @@ impl BaseTable for RemoteTable { self.fetch_blob_files_impl(column, row_ids).await } + async fn open_blob_file(&self, descriptor: &crate::blob::BlobDescriptor) -> Result { + self.open_blob_file_impl(descriptor).await + } + async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result { if let AnyQuery::Query(request) = query && let Some(offsets) = &request.take_offsets diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index bcaf47765..e3cdfc6a6 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -21,7 +21,7 @@ use reqwest::{Response, StatusCode, header}; use tokio::sync::{Mutex, OwnedSemaphorePermit}; use crate::Error; -use crate::blob::{BlobFile, BlobRangeRequest}; +use crate::blob::{BlobDescriptor, BlobDescriptorSource, BlobFile, BlobRangeRequest}; use crate::error::Result; use crate::query::{QueryFilter, QueryRequest, Select}; use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; @@ -49,6 +49,7 @@ fn is_fetch_blobs_byte_limit_error(error: &Error) -> bool { enum RangeRequestMode { SizeProbe, DataRead, + MetadataRead, } #[async_trait::async_trait] @@ -86,6 +87,9 @@ struct TableBlobRangeRequester { path: String, version: Option, branch: Option, + source_uri: Option, + source_commit_micros: Option, + row_is_address: bool, freshness: Arc>, parent_freshness: Arc>, parent_freshness_request: FreshnessHeaders, @@ -109,15 +113,27 @@ impl BlobRangeRequester for TableBlobRangeRequester { .expect("blob request semaphore is never closed"); let freshness_request = freshness_headers_snapshot(&self.freshness, self.read_consistency_interval); - let mut request = freshness_request - .apply(self.client.get(&self.path)) - .header(header::RANGE, range_header); + let mut request = if matches!(mode, RangeRequestMode::MetadataRead) { + freshness_request.apply(self.client.head(&self.path)) + } else { + freshness_request + .apply(self.client.get(&self.path)) + .header(header::RANGE, range_header) + }; if let Some(version) = version.or(self.version) { request = request.query(&[("version", version)]); } if let Some(branch) = &self.branch { request = request.query(&[("branch", branch)]); } + if let Some(source_uri) = &self.source_uri { + request = request + .query(&[("source_uri", source_uri)]) + .query(&[("row_is_address", self.row_is_address)]); + } + if let Some(timestamp) = self.source_commit_micros { + request = request.query(&[("source_commit_micros", timestamp)]); + } let (request_id, response) = self.client.send_with_retry(request, None, true).await?; // Preserve 416 size-probe responses so the caller can detect empty blobs. if response.status() == StatusCode::RANGE_NOT_SATISFIABLE @@ -131,8 +147,10 @@ impl BlobRangeRequester for TableBlobRangeRequester { } let response = self.client.check_response(&request_id, response).await?; freshness_request.observe_headers(&self.freshness, response.headers()); - self.parent_freshness_request - .observe_headers(&self.parent_freshness, response.headers()); + if self.source_uri.is_none() { + self.parent_freshness_request + .observe_headers(&self.parent_freshness, response.headers()); + } Ok(BlobRangeResponse { request_id, response, @@ -630,6 +648,51 @@ impl RemoteTable { Ok(files) } + pub(super) async fn open_blob_file_impl( + &self, + descriptor: &BlobDescriptor, + ) -> Result { + descriptor.validate()?; + self.ensure_blob_files_supported()?; + let source = &descriptor.source; + let snapshot = ReadSnapshot { + version: Some(source.version), + freshness_state: FreshnessState { + pinned: true, + ..Default::default() + }, + ..Default::default() + }; + let requester = self.blob_range_requester_with_source( + &source.column, + source.row_reference, + &snapshot, + Some(source), + ); + let size = if descriptor.kind == 3 && descriptor.size == 0 { + let metadata = requester + .request_range("", RangeRequestMode::MetadataRead, Some(source.version)) + .await?; + if metadata.response.status() != StatusCode::OK { + return Err(Error::InvalidInput { + message: "non-null blob descriptor did not resolve to a stored value".into(), + }); + } + metadata + .response + .headers() + .get(header::CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .ok_or_else(|| Error::InvalidInput { + message: "blob metadata response has no valid Content-Length".into(), + })? + } else { + descriptor.size + }; + Ok(RemoteBlobFile::new(requester, size, Some(source.version)).into()) + } + /// Probe each row once for its size. /// /// Serves callers that must not issue a query, such as a single blob that @@ -793,6 +856,16 @@ impl RemoteTable { column: &str, row_id: u64, read_snapshot: &ReadSnapshot, + ) -> Arc { + self.blob_range_requester_with_source(column, row_id, read_snapshot, None) + } + + fn blob_range_requester_with_source( + &self, + column: &str, + row_id: u64, + read_snapshot: &ReadSnapshot, + source: Option<&BlobDescriptorSource>, ) -> Arc { let path = format!( "/v1/table/{}/blob/{}/{row_id}/bytes", @@ -803,7 +876,10 @@ impl RemoteTable { client: self.client.clone(), path, version: read_snapshot.version, - branch: self.branch.clone(), + branch: source.map_or_else(|| self.branch.clone(), |source| source.branch.clone()), + source_uri: source.map(|source| source.uri.clone()), + source_commit_micros: source.map(|source| source.committed_at_micros), + row_is_address: source.is_some_and(|source| source.row_is_address), freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)), parent_freshness: self.freshness.clone(), parent_freshness_request: read_snapshot.freshness, @@ -1024,6 +1100,99 @@ mod tests { .unwrap() } + #[tokio::test] + async fn open_descriptor_files_are_lazy_and_pin_ranges_for_every_storage_kind() { + for kind in 0..=5 { + let requests = Arc::new(StdMutex::new(Vec::new())); + let captured = requests.clone(); + let table = RemoteTable::new_mock( + "videos".to_string(), + move |request| { + for name in [ + "x-lancedb-min-version", + "x-lancedb-min-read-version", + "x-lancedb-min-timestamp", + ] { + assert!(request.headers().get(name).is_none()); + } + assert_eq!( + request.url().path(), + "/v1/table/videos/blob/video/12884901889/bytes" + ); + let parameters = request + .url() + .query_pairs() + .into_owned() + .collect::>(); + assert_eq!(parameters["version"], "7"); + assert_eq!(parameters["source_commit_micros"], "1234"); + assert_eq!(parameters["branch"], "dev"); + assert_eq!(parameters["row_is_address"], "true"); + assert_eq!( + parameters["source_uri"], + "s3://bucket/videos.lance/tree/dev" + ); + captured.lock().unwrap().push(request.method().clone()); + if request.method() == reqwest::Method::HEAD { + assert!(request.headers().get(header::RANGE).is_none()); + http::Response::builder() + .status(200) + .header(header::CONTENT_LENGTH, PAYLOAD.len().to_string()) + .body(Vec::new()) + .unwrap() + } else { + range_response(&request, PAYLOAD) + } + }, + Some(Version::new(0, 5, 0)), + ); + { + let mut freshness = table.freshness.lock().unwrap(); + freshness.min_version = Some(99); + freshness.min_read_version = Some(2); + } + let descriptor = BlobDescriptor { + kind, + position: 0, + size: if kind == 3 { 0 } else { PAYLOAD.len() as u64 }, + blob_id: 0, + blob_uri: String::new(), + source: BlobDescriptorSource { + uri: "s3://bucket/videos.lance/tree/dev".into(), + version: 7, + committed_at_micros: 1234, + branch: Some("dev".into()), + column: "video".into(), + row_reference: 12884901889, + row_is_address: true, + }, + }; + let file = table.open_blob_file(&descriptor).await.unwrap(); + let expected = if kind == 3 { + vec![reqwest::Method::HEAD] + } else { + vec![] + }; + assert_eq!(*requests.lock().unwrap(), expected); + assert_eq!(file.size(), PAYLOAD.len() as u64); + file.seek(4).await.unwrap(); + assert_eq!(file.read_up_to(5).await.unwrap().as_ref(), &PAYLOAD[4..9]); + assert_eq!(file.tell().await.unwrap(), 9); + assert_eq!(table.freshness.lock().unwrap().min_read_version, Some(2)); + assert_eq!(file.read_range(0..3).await.unwrap().as_ref(), &PAYLOAD[..3]); + assert_eq!(file.tell().await.unwrap(), 9); + assert_eq!( + requests + .lock() + .unwrap() + .iter() + .filter(|method| **method == reqwest::Method::GET) + .count(), + 2 + ); + } + } + fn mock_remote_blob_table( requests: Arc>>, ) -> RemoteTable { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 3123eb776..3fcf4bdb7 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3,7 +3,7 @@ //! LanceDB Table APIs -use crate::blob::BlobFile; +use crate::blob::{BlobDescriptor, BlobFile}; use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; use arrow_schema::{Schema, SchemaRef}; use async_trait::async_trait; @@ -796,6 +796,12 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "fetch_blob_files is not supported on this table type".into(), }) } + /// Open a descriptor at its pinned source version. See [`Table::open_blob_file`]. + async fn open_blob_file(&self, _descriptor: &BlobDescriptor) -> Result { + Err(Error::NotSupported { + message: "open_blob_file is not supported on this table type".into(), + }) + } /// Gets the table tag manager. async fn tags(&self) -> Result>; /// Optimize the dataset. @@ -1377,6 +1383,26 @@ impl Table { self.inner.fetch_blob_files(column.as_ref(), row_ids).await } + /// Open a seekable file from a SQL `blob_descriptor` result. + /// + /// The descriptor carries its table version, branch, column and row identity. + /// Opening does not materialize the payload or change this table's checkout. + /// Remote reads use the table's authenticated HTTP byte-range endpoint. + /// + /// ``` + /// use lancedb::{Table, blob::BlobDescriptor}; + /// + /// # async fn read_header(videos: &Table, descriptor: &BlobDescriptor) -> lancedb::Result<()> { + /// let video = videos.open_blob_file(descriptor).await?; + /// let header = video.read_range(0..video.size().min(4096)).await?; + /// video.close().await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn open_blob_file(&self, descriptor: &BlobDescriptor) -> Result { + self.inner.open_blob_file(descriptor).await + } + /// Insert new records into this Table /// /// # Arguments @@ -3764,6 +3790,11 @@ impl BaseTable for NativeTable { crate::blob::take_blob_files_aligned(&dataset, column, row_ids).await } + async fn open_blob_file(&self, descriptor: &BlobDescriptor) -> Result { + let dataset = self.dataset.get().await?; + crate::blob::open_native_blob_file(&dataset, descriptor).await + } + /// Delete rows from the table async fn delete(&self, predicate: Predicate<'_>) -> Result { let result = delete::execute_delete(self, predicate).await?;