From 0968c21f4cee62443a4e9096ed301b4da9cb36f0 Mon Sep 17 00:00:00 2001 From: Dan Rammer Date: Fri, 16 Oct 2026 11:41:27 -0500 Subject: [PATCH] feat(remote): support set_unenforced_primary_key (#4158) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary `RemoteTable::set_unenforced_primary_key` returned `NotSupported`, so the call failed against LanceDB Cloud and enterprise from every SDK — Python and TypeScript both forward to it. That also blocked sharded LSM writes, since the server rejects bucket/identity sharding on a table that declares no unenforced primary key, and the SDK had no way to declare one. No server-side support was missing. The unenforced primary key is Lance schema field metadata, and the existing `update_field_metadata` endpoint writes exactly that, so the remote table now installs the key through it. The commit layer behind that endpoint installs the position, enforces immutability and runs `verify_primary_key()` — the same code a native table reaches, so both paths agree on semantics and not just on error messages. ## Changes - `rust/lancedb/src/table/primary_key.rs`: the request validation (single column, column exists, supported dtype, key not already set) and the metadata edit (set the 1-indexed position, clear the legacy boolean key) move into `validate()` and `install_edit()`, shared by both table types so they cannot drift. Native behaviour is unchanged. - `rust/lancedb/src/remote/table.rs`: real implementation — `check_mutable()`, validate against the described schema, then one `update_field_metadata` request. - `python/python/lancedb/remote/table.py`: drop the "Not supported on LanceDB Cloud." docstring. No binding changes were needed; Python and TypeScript already forward to the Rust `BaseTable` method. ## Tests Three mocked-server tests in `rust/lancedb/src/remote/table.rs`: - the happy path, asserting the request body carries `lance-schema:unenforced-primary-key:position: "1"` and a null legacy key on the right field path; - the invalid requests a native table also rejects (empty, compound, unknown column, unsupported dtype), each failing before any write reaches the server; - the already-set case, rejected from the schema the server already reports. `cargo test --features remote -p lancedb --lib`: 1186 passed, 0 failed. `cargo clippy --features remote -p lancedb --lib --tests` is clean. ## Note for enterprise Sophon pins the lancedb submodule at `v0.39.0-beta.6`, so this reaches enterprise on the next bump. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_011xV8EebEZDmDDV8hVBm6MD --------- Co-authored-by: Claude Opus 5 (1M context) --- nodejs/__test__/remote.test.ts | 4 - python/python/lancedb/remote/table.py | 2 +- python/python/tests/test_remote_db.py | 1 - rust/lancedb/src/remote/table.rs | 145 ++++++++++++++++++++++++-- rust/lancedb/src/table.rs | 2 +- rust/lancedb/src/table/primary_key.rs | 120 +++++++++++++-------- 6 files changed, 214 insertions(+), 60 deletions(-) diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 4f70ac1aa..c06cb737e 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -125,10 +125,6 @@ describe("remote connection", () => { const cases: [() => Promise, string][] = [ [() => table.usesV2ManifestPaths(), "uses_v2_manifest_paths"], [() => table.migrateManifestPathsV2(), "migrate_manifest_paths_v2"], - [ - () => table.setUnenforcedPrimaryKey("id"), - "set_unenforced_primary_key", - ], [() => table.optimize(), "optimize"], [ () => table.createIndex("vector", { config: Index.hnswPq() }), diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 53b6c402a..5996fdeb3 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1034,7 +1034,7 @@ class RemoteTable(Table): return LOOP.run(self._table.drop_columns(columns)) def set_unenforced_primary_key(self, columns: Union[str, Iterable[str]]) -> None: - """Not supported for remote tables.""" + """Set the unenforced primary key for this table to a single column.""" return LOOP.run(self._table.set_unenforced_primary_key(columns)) def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None: diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 717490dcd..c22aa78cd 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1623,7 +1623,6 @@ def remote_unsupported_operations(): return [ ("uses_v2_manifest_paths", (), {}, "uses_v2_manifest_paths"), ("migrate_v2_manifest_paths", (), {}, "migrate_manifest_paths_v2"), - ("set_unenforced_primary_key", (["id"],), {}, "set_unenforced_primary_key"), ("fetch_blob_ranges", ("image", [(1, 0, 1)]), {}, "fetch_blob_ranges"), ("optimize", (), {}, "optimize"), ("create_index", ("vector",), {"config": HnswPq()}, "IVF_HNSW_PQ"), diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 7d7a6301a..d9da2c1a1 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -37,6 +37,7 @@ use crate::table::Tags; use crate::table::UpdateResult; use crate::table::lsm_stats::GetLsmStatsResponse; use crate::table::merge::MergeFilter; +use crate::table::primary_key; use crate::table::query::create_multi_vector_plan; use crate::table::write_progress::FinishOnDrop; use crate::table::{ @@ -74,6 +75,7 @@ use lance::arrow::json::{JsonDataType, JsonSchema}; use lance::dataset::refs::TagContents; use lance::dataset::scanner::DatasetRecordBatchStream; use lance::dataset::{ColumnAlteration, NewColumnTransform, Version}; +use lance_core::datatypes::Schema as LanceSchema; use lance_datafusion::exec::{OneShotExec, execute_plan}; use reqwest::{RequestBuilder, Response}; use serde::{Deserialize, Serialize}; @@ -3202,10 +3204,27 @@ impl BaseTable for RemoteTable { } } - async fn set_unenforced_primary_key(&self, _columns: &[&str]) -> Result<()> { - Err(Error::NotSupported { - message: "set_unenforced_primary_key is not supported for remote tables.".into(), - }) + /// The unenforced primary key is Lance schema field metadata, so this + /// installs it through the `update_field_metadata` endpoint. The commit + /// layer behind that endpoint is what actually installs and validates the + /// key, exactly as on a native table; the checks here only fail fast with + /// the same messages a native table gives. + async fn set_unenforced_primary_key(&self, columns: &[&str]) -> Result<()> { + self.check_mutable().await?; + + let arrow_schema = self.schema().await?; + let schema = LanceSchema::try_from(arrow_schema.as_ref()).map_err(|e| Error::Schema { + message: format!("Invalid schema: {}", e), + })?; + primary_key::validate(&schema, columns)?; + + self.update_field_metadata(&[FieldMetadataUpdate { + path: columns[0].to_string(), + metadata: primary_key::install_edit(), + replace: false, + }]) + .await?; + Ok(()) } async fn flush_lsm(&self) -> Result<()> { @@ -5936,10 +5955,6 @@ mod tests { let mut fts = FullTextSearchQuery::new("test".into()); fts.wand_factor = Some(1.0); let results = [ - ( - "set_unenforced_primary_key", - table.set_unenforced_primary_key(["id"]).await, - ), ( "optimize", table.optimize(OptimizeAction::All).await.map(|_| ()), @@ -13365,6 +13380,120 @@ mod tests { assert_eq!(result.version, 7); } + /// The unenforced primary key is field metadata, so the remote table + /// installs it through the `update_field_metadata` endpoint. + #[tokio::test] + async fn test_set_unenforced_primary_key() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + ]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/update_field_metadata/" => { + let body = request_body_json(&request); + assert_eq!(body["updates"].as_array().unwrap().len(), 1); + let update = &body["updates"][0]; + assert_eq!(update["path"], "id"); + assert_eq!(update["replace"], json!(false)); + assert_eq!( + update["metadata"]["lance-schema:unenforced-primary-key:position"], + "1" + ); + assert_eq!( + update["metadata"]["lance-schema:unenforced-primary-key"], + json!(null) + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 3, "fields": {}}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + table.set_unenforced_primary_key(["id"]).await.unwrap(); + } + + /// Requests the native table rejects are rejected here too, before any + /// write reaches the server. + #[tokio::test] + async fn test_set_unenforced_primary_key_rejects_invalid_requests() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("score", DataType::Float32, true), + ]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + }); + + for columns in [ + vec![], + vec!["id", "score"], + vec!["nonexistent"], + vec!["score"], + ] { + let err = table + .set_unenforced_primary_key(columns.clone()) + .await + .unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { .. }), + "unexpected error for {:?}: {:?}", + columns, + err + ); + } + } + + /// The key is immutable once set, and the schema the server already + /// reports is enough to say so. + #[tokio::test] + async fn test_set_unenforced_primary_key_already_set() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false).with_metadata(HashMap::from([( + "lance-schema:unenforced-primary-key:position".to_string(), + "1".to_string(), + )])), + Field::new("name", DataType::Utf8, false), + ]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + }); + + for column in ["name", "id"] { + let err = table + .set_unenforced_primary_key([column]) + .await + .unwrap_err(); + assert!( + err.to_string().contains("already set"), + "unexpected error: {:?}", + err + ); + } + } + // ----- Branch support ----- /// Parse a request's in-memory JSON body. Only valid for JSON-body ops diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index ca58b54de..74b99255a 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -83,7 +83,7 @@ pub mod freshness; pub mod lsm_stats; pub mod merge; pub mod optimize; -mod primary_key; +pub(crate) mod primary_key; pub mod query; pub mod refresh; pub mod schema_evolution; diff --git a/rust/lancedb/src/table/primary_key.rs b/rust/lancedb/src/table/primary_key.rs index 8a7b48efa..d399155fd 100644 --- a/rust/lancedb/src/table/primary_key.rs +++ b/rust/lancedb/src/table/primary_key.rs @@ -11,24 +11,26 @@ //! Only a single-column primary key is supported, and the key cannot be //! changed once set. +use std::collections::HashMap; + use arrow_schema::DataType; -use lance_core::datatypes::{LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION}; +use lance_core::datatypes::{ + Field as LanceField, LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION, + Schema as LanceSchema, +}; use crate::error::{Error, Result}; use crate::table::NativeTable; -/// Set the unenforced primary key on `table` to the single column in `columns`. +/// Validate a `set_unenforced_primary_key` request against `schema`, returning +/// the field the key would be installed on. /// -/// Fails if `columns` is not exactly one column (compound primary keys are not -/// supported), if the column does not exist or has an unsupported dtype, or if -/// the table already has an unenforced primary key (changing the primary key -/// is not supported). -pub(super) async fn set_unenforced_primary_key( - table: &NativeTable, - columns: &[&str], -) -> Result<()> { - table.dataset.ensure_mutable()?; - +/// Shared by [`NativeTable`] and the remote table so both reject the same +/// requests with the same messages. Fails if `columns` is not exactly one +/// column (compound primary keys are not supported), if the column does not +/// exist or has an unsupported dtype, or if the table already has an +/// unenforced primary key (changing the primary key is not supported). +pub fn validate<'a>(schema: &'a LanceSchema, columns: &[&str]) -> Result<&'a LanceField> { if columns.is_empty() { return Err(Error::InvalidInput { message: "set_unenforced_primary_key: a column is required".into(), @@ -44,43 +46,71 @@ pub(super) async fn set_unenforced_primary_key( } let column = columns[0]; + // The primary key is immutable once set. The Lance commit layer is the + // source of truth for this (it also covers the concurrent-writer race); + // this check just fails fast with a clear message. + if !schema.unenforced_primary_key().is_empty() { + return Err(Error::InvalidInput { + message: "set_unenforced_primary_key: an unenforced primary key is already set on this table; changing it is not supported".into(), + }); + } + + let field = schema.field(column).ok_or_else(|| Error::InvalidInput { + message: format!( + "set_unenforced_primary_key: column '{}' not found on table", + column + ), + })?; + if !is_supported_pk_dtype(&field.data_type()) { + return Err(Error::InvalidInput { + message: format!( + "set_unenforced_primary_key: column '{}' has dtype {:?} which is not supported as a primary key. Supported: Int32, Int64, Utf8, LargeUtf8, Binary, LargeBinary, FixedSizeBinary", + column, + field.data_type() + ), + }); + } + Ok(field) +} + +/// The field metadata edit that installs the primary key on a field: keys to +/// set (`Some`) or delete (`None`). +/// +/// Position metadata is 1-indexed; `Schema::unenforced_primary_key` treats +/// position 0 as a legacy "no specific position" fallback, so the legacy +/// boolean key is cleared and only the position governs. +pub fn install_edit() -> HashMap> { + HashMap::from([ + (LANCE_UNENFORCED_PRIMARY_KEY.to_string(), None), + ( + LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(), + Some("1".to_string()), + ), + ]) +} + +/// Set the unenforced primary key on `table` to the single column in `columns`. +pub(super) async fn set_unenforced_primary_key( + table: &NativeTable, + columns: &[&str], +) -> Result<()> { + table.dataset.ensure_mutable()?; + let updates = { let dataset = table.dataset.get().await?; - let schema = dataset.schema(); + let field = validate(dataset.schema(), columns)?; - // The primary key is immutable once set. The Lance commit layer is the - // source of truth for this (it also covers the concurrent-writer race); - // this check just fails fast with a clear message. - if !schema.unenforced_primary_key().is_empty() { - return Err(Error::InvalidInput { - message: "set_unenforced_primary_key: an unenforced primary key is already set on this table; changing it is not supported".into(), - }); - } - - let field = schema.field(column).ok_or_else(|| Error::InvalidInput { - message: format!( - "set_unenforced_primary_key: column '{}' not found on table", - column - ), - })?; - if !is_supported_pk_dtype(&field.data_type()) { - return Err(Error::InvalidInput { - message: format!( - "set_unenforced_primary_key: column '{}' has dtype {:?} which is not supported as a primary key. Supported: Int32, Int64, Utf8, LargeUtf8, Binary, LargeBinary, FixedSizeBinary", - column, - field.data_type() - ), - }); - } - - // Position metadata is 1-indexed; `Schema::unenforced_primary_key` - // treats position 0 as a legacy "no specific position" fallback. let mut metadata = field.metadata.clone(); - metadata.remove(LANCE_UNENFORCED_PRIMARY_KEY); - metadata.insert( - LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(), - "1".to_string(), - ); + for (key, value) in install_edit() { + match value { + Some(value) => { + metadata.insert(key, value); + } + None => { + metadata.remove(&key); + } + } + } vec![(field_id_to_u32(field.id, &field.name)?, metadata)] };