mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-10 23:32:35 +00:00
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) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011xV8EebEZDmDDV8hVBm6MD
This commit is contained in:
co-authored by
Claude Opus 5
parent
bc4497b21a
commit
f9dab6c3c8
@@ -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:
|
||||
|
||||
@@ -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<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<String, Option<String>> {
|
||||
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)]
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user