From f9dab6c3c8fb942650da45247e0126641f97b470 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Thu, 10 Sep 2026 10:38:20 -0500 Subject: [PATCH] feat(remote): support set_unenforced_primary_key `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. 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 messages. The request validation and the metadata edit move into shared helpers so the native and remote paths cannot drift. Native behaviour is unchanged. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_011xV8EebEZDmDDV8hVBm6MD --- python/python/lancedb/remote/table.py | 2 +- rust/lancedb/src/remote/table.rs | 141 +++++++++++++++++++++++++- rust/lancedb/src/table.rs | 2 +- rust/lancedb/src/table/primary_key.rs | 120 ++++++++++++++-------- 4 files changed, 214 insertions(+), 51 deletions(-) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 89b958165..532c2d00d 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1008,7 +1008,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 on LanceDB Cloud.""" + """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/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 5c32e3cdd..ef5987227 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -32,6 +32,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::{ @@ -68,6 +69,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}; @@ -2992,10 +2994,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 on LanceDB cloud.".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<()> { @@ -11832,6 +11851,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 7f139c4cb..526282c86 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -78,7 +78,7 @@ pub mod delete; 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)] };