From ecf87f63710c4cc37d7810469845b2f19f522b26 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 12 Aug 2026 09:49:51 +0800 Subject: [PATCH] feat: add atomic generated column binding snapshots --- rust/lancedb/src/function.rs | 2 + rust/lancedb/src/function/binding_snapshot.rs | 186 +++++++ rust/lancedb/src/remote/table.rs | 496 +++++++++++++++++- rust/lancedb/src/table.rs | 134 +++++ ...erated_column_binding_snapshot_contract.rs | 117 +++++ 5 files changed, 930 insertions(+), 5 deletions(-) create mode 100644 rust/lancedb/src/function/binding_snapshot.rs create mode 100644 rust/lancedb/tests/first_class_generated_column_binding_snapshot_contract.rs diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index e06f8b9ce..80ed9c0d4 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -6,12 +6,14 @@ //! This module defines transport and metadata value types only. It does not //! provide catalogs, job execution, query planning, or generated-column runtime. +mod binding_snapshot; mod change_generated_column; mod create_generated_column; mod definition; mod refresh_generated_column; mod registration; +pub use binding_snapshot::{GeneratedColumnBindingEntry, GeneratedColumnBindingSnapshot}; pub use change_generated_column::ChangeGeneratedColumnJobSpec; pub use create_generated_column::CreateGeneratedColumnJobSpec; pub use definition::{FunctionCapability, FunctionDefinition, PythonFunctionDefinition}; diff --git a/rust/lancedb/src/function/binding_snapshot.rs b/rust/lancedb/src/function/binding_snapshot.rs new file mode 100644 index 000000000..6fcabe895 --- /dev/null +++ b/rust/lancedb/src/function/binding_snapshot.rs @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Atomic generated-column binding snapshot projection (FF-029). +//! +//! This is an implementation projection for table call binding. It is not a +//! catalog resource, Job, persistent model, wire payload, or table-version +//! replacement. + +use std::collections::HashSet; + +use arrow_schema::FieldRef; + +use super::invalid_input; +use crate::Result; + +/// One top-level field identity from a single table snapshot. +/// +/// Pairs a non-negative Lance stable field ID with the exact Arrow field from +/// that same snapshot. IDs are carried only here; they are never injected into +/// Arrow field metadata. +#[doc(hidden)] +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GeneratedColumnBindingEntry { + field_id: i32, + field: FieldRef, +} + +impl GeneratedColumnBindingEntry { + /// Stable Lance field ID for this top-level entry. + pub fn field_id(&self) -> i32 { + self.field_id + } + + /// Exact Arrow field from the same snapshot. + pub fn field(&self) -> &FieldRef { + &self.field + } +} + +/// Atomic table snapshot projection for generated-column call binding. +/// +/// Contains one table version and immutable top-level field entries in schema +/// order. Construction validates field/ID count equality, non-negative unique +/// IDs, and unique top-level names. +#[doc(hidden)] +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GeneratedColumnBindingSnapshot { + version: u64, + entries: Vec, +} + +impl GeneratedColumnBindingSnapshot { + /// Build a binding snapshot from one version and ordered field/ID pairs. + /// + /// `fields` and `field_ids` must have the same length. Every ID must be + /// non-negative and unique. Top-level field names must be unique. Order is + /// preserved exactly as provided. + pub fn try_new( + version: u64, + fields: impl IntoIterator, + field_ids: impl IntoIterator, + ) -> Result { + let fields: Vec = fields.into_iter().collect(); + let field_ids: Vec = field_ids.into_iter().collect(); + if fields.len() != field_ids.len() { + return Err(invalid_input( + "generated-column binding snapshot field count must equal field_ids count", + )); + } + + let mut seen_ids = HashSet::with_capacity(field_ids.len()); + let mut seen_names = HashSet::with_capacity(fields.len()); + let mut entries = Vec::with_capacity(fields.len()); + + for (field, field_id) in fields.into_iter().zip(field_ids) { + if field_id < 0 { + return Err(invalid_input( + "generated-column binding snapshot field IDs must be non-negative", + )); + } + if !seen_ids.insert(field_id) { + return Err(invalid_input( + "generated-column binding snapshot field IDs must be unique", + )); + } + if !seen_names.insert(field.name().clone()) { + return Err(invalid_input( + "generated-column binding snapshot top-level field names must be unique", + )); + } + entries.push(GeneratedColumnBindingEntry { field_id, field }); + } + + Ok(Self { version, entries }) + } + + /// Table version for this snapshot. + pub fn version(&self) -> u64 { + self.version + } + + /// Top-level entries in schema order. + pub fn entries(&self) -> &[GeneratedColumnBindingEntry] { + &self.entries + } + + /// Exact case-sensitive top-level field name lookup. + /// + /// A name containing `.` is a literal top-level field name, not a nested + /// path. Lookup does not fold case or interpret dotted selectors. + pub fn field(&self, name: &str) -> Option<&GeneratedColumnBindingEntry> { + self.entries + .iter() + .find(|entry| entry.field().name() == name) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::{DataType, Field}; + + use super::*; + use crate::Error; + + fn fields() -> Vec { + vec![ + Arc::new(Field::new("text", DataType::Utf8, true)), + Arc::new(Field::new("Score", DataType::Int32, false)), + Arc::new(Field::new("a.b", DataType::Utf8, true)), + ] + } + + #[test] + fn try_new_preserves_version_order_and_entry_data() { + let snapshot = + GeneratedColumnBindingSnapshot::try_new(11, fields(), vec![2, 4, 8]).unwrap(); + assert_eq!(snapshot.version(), 11); + assert_eq!(snapshot.entries().len(), 3); + assert_eq!(snapshot.entries()[0].field_id(), 2); + assert_eq!(snapshot.entries()[0].field().name(), "text"); + assert_eq!(snapshot.entries()[0].field().data_type(), &DataType::Utf8); + assert_eq!(snapshot.entries()[1].field_id(), 4); + assert_eq!(snapshot.entries()[1].field().name(), "Score"); + assert_eq!(snapshot.entries()[2].field_id(), 8); + assert_eq!(snapshot.entries()[2].field().name(), "a.b"); + } + + #[test] + fn lookup_is_exact_case_sensitive_and_treats_dot_literally() { + let snapshot = GeneratedColumnBindingSnapshot::try_new(1, fields(), vec![2, 4, 8]).unwrap(); + assert_eq!(snapshot.field("Score").unwrap().field_id(), 4); + assert!(snapshot.field("score").is_none()); + assert!(snapshot.field("TEXT").is_none()); + assert!(snapshot.field("a").is_none()); + assert!(snapshot.field("b").is_none()); + assert_eq!(snapshot.field("a.b").unwrap().field_id(), 8); + } + + #[test] + fn try_new_rejects_count_mismatch_negative_duplicate_ids_and_names() { + let base = fields(); + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, base.clone(), vec![1, 2]), + Err(Error::InvalidInput { .. }) + )); + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, base.clone(), vec![1, 2, -3]), + Err(Error::InvalidInput { .. }) + )); + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, base.clone(), vec![1, 2, 1]), + Err(Error::InvalidInput { .. }) + )); + let duplicate_names = vec![ + Arc::new(Field::new("text", DataType::Utf8, true)), + Arc::new(Field::new("text", DataType::Int32, false)), + ]; + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, duplicate_names, vec![1, 2]), + Err(Error::InvalidInput { .. }) + )); + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 0d843dd54..131ba3514 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -12,6 +12,7 @@ use super::{ARROW_FILE_CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE}; use crate::blob::BlobFile; use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions}; use crate::expr::expr_to_sql_string; +use crate::function::GeneratedColumnBindingSnapshot; use crate::index::Index; use crate::index::IndexStatistics; use crate::index::waiter::wait_for_index; @@ -563,6 +564,17 @@ impl RemoteTable { request: RequestBuilder, version: Option, ) -> Result { + let (_request_id, description) = self.describe_with_request_id(request, version).await?; + Ok(description) + } + + /// Same as [`Self::describe_with_request`], but preserves the real request ID + /// for sanitized protocol-error mapping on binding projections. + async fn describe_with_request_id( + &self, + request: RequestBuilder, + version: Option, + ) -> Result<(String, TableDescription)> { let mut body = serde_json::json!({ "version": version }); self.apply_branch_body(&mut body); let request = request.json(&body); @@ -572,11 +584,14 @@ impl RemoteTable { let response = self.check_table_response(&request_id, response).await?; match response.text().await { - Ok(body) => serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse table description: {}", e).into(), - request_id, - status_code: None, - }), + Ok(body) => { + let description = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse table description: {}", e).into(), + request_id: request_id.clone(), + status_code: None, + })?; + Ok((request_id, description)) + } Err(err) => { let status_code = err.status(); Err(Error::Http { @@ -1135,6 +1150,38 @@ struct TableDescription { version: u64, schema: JsonSchema, location: Option, + /// Optional top-level Lance stable field IDs matching `schema` order. + /// + /// Absent on old servers. Ordinary schema/version/seed paths ignore it; + /// generated-column binding requires it and fail-closes when missing or + /// invalid. + field_ids: Option>, +} + +/// Stable, payload-free Http error for binding-facing local protocol validation. +fn binding_protocol_http(request_id: String, message: &'static str) -> Error { + Error::Http { + source: message.into(), + request_id, + status_code: None, + } +} + +/// Sanitize describe JSON-parse failures for the binding path only. +/// +/// Transport/`reqwest` errors and HTTP status failures are left unchanged so +/// ordinary schema/version behavior and network diagnostics stay intact. +fn sanitize_binding_describe_protocol_error(err: Error) -> Error { + match err { + Error::Http { + request_id, + status_code: None, + source, + } if source.downcast_ref::().is_none() => { + binding_protocol_http(request_id, "invalid table description in describe response") + } + other => other, + } } /// How a response body frames its Arrow IPC payload. `/query` answers with file framing @@ -1853,6 +1900,43 @@ impl BaseTable for RemoteTable { .map_err(unwrap_shared_error) } + async fn generated_column_binding_snapshot(&self) -> Result { + // Bypass the Arrow-only schema cache. Cached schema must never invent + // or substitute Lance field IDs. One describe request carries version, + // schema, and optional field_ids together. + let version = self.current_version().await; + let request = self.post_read(&format!("/v1/table/{}/describe/", self.identifier)); + let (request_id, description) = self + .describe_with_request_id(request, version) + .await + .map_err(sanitize_binding_describe_protocol_error)?; + + let Some(field_ids) = description.field_ids else { + return Err(Error::NotSupported { + message: "generated-column binding snapshot requires field_ids in table describe response" + .into(), + }); + }; + + let Ok(arrow_schema) = arrow_schema::Schema::try_from(description.schema) else { + return Err(binding_protocol_http( + request_id, + "invalid table schema in describe response", + )); + }; + let fields = arrow_schema.fields().iter().cloned().collect::>(); + + GeneratedColumnBindingSnapshot::try_new(description.version, fields, field_ids).map_err( + |err| match err { + Error::InvalidInput { .. } => binding_protocol_http( + request_id, + "invalid generated-column binding projection in describe response", + ), + other => other, + }, + ) + } + async fn create_branch( &self, name: &str, @@ -3315,6 +3399,21 @@ mod tests { .unwrap() } + /// Describe response with additive top-level `field_ids` for binding tests. + fn describe_response_with_field_ids( + version: u64, + schema: &Schema, + field_ids: &[i32], + ) -> String { + let json_schema = JsonSchema::try_from(schema).unwrap(); + serde_json::to_string(&json!({ + "version": version, + "schema": json_schema, + "field_ids": field_ids, + })) + .unwrap() + } + fn nested_index_schema() -> Schema { let vector_type = DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 8); @@ -7622,6 +7721,393 @@ mod tests { assert_eq!(call_count.load(Ordering::SeqCst), 1); // Still 1, no new call } + #[tokio::test] + async fn test_generated_column_binding_snapshot_one_describe_with_ids() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let schema = Schema::new(vec![ + Field::new("text", DataType::Utf8, true), + Field::new("a.b", DataType::Int32, false), + ]); + let body = describe_response_with_field_ids(42, &schema, &[3, 9]); + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + let req_body = request_body_json(&request); + assert_eq!(req_body["version"], serde_json::Value::Null); + assert!(req_body.get("branch").is_none()); + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + assert_eq!(snapshot.version(), 42); + assert_eq!(snapshot.entries().len(), 2); + assert_eq!(snapshot.entries()[0].field_id(), 3); + assert_eq!(snapshot.entries()[0].field().name(), "text"); + assert_eq!(snapshot.entries()[0].field().data_type(), &DataType::Utf8); + assert_eq!(snapshot.entries()[1].field_id(), 9); + assert_eq!(snapshot.field("a.b").unwrap().field_id(), 9); + assert!( + !snapshot + .entries() + .iter() + .any(|e| e.field().metadata().contains_key("lance:field_id")) + ); + } + + #[tokio::test] + async fn test_generated_column_binding_snapshot_missing_ids_not_supported_schema_ok() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let body = describe_response(&schema); + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + + let err = table.generated_column_binding_snapshot().await.unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + let schema = table.schema().await.unwrap(); + assert_eq!(schema.field(0).name(), "a"); + // Ordinary schema may use cache / a second describe; binding must not + // invent IDs from that path. + assert!(call_count.load(Ordering::SeqCst) >= 1); + } + + #[tokio::test] + async fn test_generated_column_binding_snapshot_invalid_ids_are_protocol_errors() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ]); + + for field_ids in [vec![1], vec![1, -2], vec![1, 1]] { + let body = describe_response_with_field_ids(1, &schema, &field_ids); + let body_for_assert = body.clone(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + let err = table.generated_column_binding_snapshot().await.unwrap_err(); + match err { + Error::Http { + request_id, + status_code, + source, + } => { + assert!(!request_id.is_empty()); + assert!(status_code.is_none()); + let message = source.to_string(); + assert!(!message.contains(&body_for_assert)); + } + other => panic!("expected Http protocol error, got {other:?}"), + } + } + } + + #[tokio::test] + async fn test_generated_column_binding_snapshot_bypasses_seeded_schema_cache() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let seeded = describe_response(&schema); + let with_ids = describe_response_with_field_ids(7, &schema, &[11]); + + let remote = RemoteTable::new_mock( + "my_table".into(), + move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body(with_ids.clone()) + .unwrap() + }, + None, + ); + remote.seed_schema(&seeded); + let table = Table::from(Arc::new(remote) as Arc); + + // Cached/seeded Arrow schema satisfies ordinary schema reads. + let schema = table.schema().await.unwrap(); + assert_eq!(schema.field(0).name(), "a"); + assert_eq!(call_count.load(Ordering::SeqCst), 0); + + // Binding must issue a fresh describe that carries field_ids. + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + assert_eq!(snapshot.version(), 7); + assert_eq!(snapshot.field("a").unwrap().field_id(), 11); + + // No second fallback/lookup after a successful binding describe. + let _ = table.generated_column_binding_snapshot().await.unwrap(); + assert_eq!(call_count.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn test_generated_column_binding_snapshot_no_submit_or_mutation_endpoints() { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let body = describe_response_with_field_ids(1, &schema, &[0]); + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + assert_eq!( + path, "/v1/table/my_table/describe/", + "binding must not call submit/mutation endpoint {path}" + ); + assert!( + !path.contains("job") + && !path.contains("generated") + && !path.contains("add_columns") + && !path.contains("alter_columns") + && !path.contains("drop_columns") + && !path.contains("insert") + && !path.contains("update") + && !path.contains("delete") + && !path.contains("merge"), + "unexpected mutating path: {path}" + ); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + assert_eq!(snapshot.version(), 1); + assert_eq!(snapshot.entries()[0].field_id(), 0); + } + + #[tokio::test] + async fn test_generated_column_binding_snapshot_respects_checkout_version_and_branch() { + use lance::dataset::refs::Ref; + + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let describe_latest = describe_response_with_field_ids(5, &schema, &[1]); + let describe_pinned = describe_response_with_field_ids(3, &schema, &[1]); + let describe_branch = describe_response_with_field_ids(9, &schema, &[2]); + + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let body = request_body_json(&request); + if body.get("branch").and_then(|v| v.as_str()) == Some("exp") { + assert_eq!(body["version"], serde_json::Value::Null); + return http::Response::builder() + .status(200) + .body(describe_branch.clone()) + .unwrap(); + } + match body["version"].as_u64() { + Some(3) => http::Response::builder() + .status(200) + .body(describe_pinned.clone()) + .unwrap(), + None => http::Response::builder() + .status(200) + .body(describe_latest.clone()) + .unwrap(), + other => panic!("unexpected describe version: {other:?}"), + } + } + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + path => panic!("unexpected path: {path}"), + }); + + table.checkout(3).await.unwrap(); + let pinned = table.generated_column_binding_snapshot().await.unwrap(); + assert_eq!(pinned.version(), 3); + + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + let branched = branch.generated_column_binding_snapshot().await.unwrap(); + assert_eq!(branched.version(), 9); + assert_eq!(branched.field("a").unwrap().field_id(), 2); + } + + /// Binding describe must reuse the same read-freshness headers as ordinary + /// `post_read` describes: write watermark, read watermark, and consistency + /// interval timestamp — on the single describe request. + #[tokio::test] + async fn test_generated_column_binding_snapshot_preserves_freshness_headers() { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let describe_body = describe_response_with_field_ids(1, &schema, &[0]); + let describe_calls = Arc::new(AtomicUsize::new(0)); + let describe_calls_c = describe_calls.clone(); + let describe_headers = Arc::new(std::sync::Mutex::new(None::)); + let describe_headers_c = describe_headers.clone(); + + let table = Table::new_with_handler_and_interval( + "my_table", + move |request| match request.url().path() { + "/v1/table/my_table/update/" => http::Response::builder() + .status(200) + .body(r#"{"rows_updated":1,"version":7}"#.to_string()) + .unwrap(), + "/v1/table/my_table/count_rows/" => http::Response::builder() + .status(200) + .header("x-lancedb-version", "100") + .body("42".to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => { + describe_calls_c.fetch_add(1, Ordering::SeqCst); + *describe_headers_c.lock().unwrap() = Some(request.headers().clone()); + http::Response::builder() + .status(200) + .body(describe_body.clone()) + .unwrap() + } + path => panic!("unexpected path: {path}"), + }, + Some(Duration::ZERO), + ); + + // Establish the same observable freshness state ordinary reads use. + table.update().column("a", "a + 1").execute().await.unwrap(); + table.count_rows(None).await.unwrap(); + + let before = SystemTime::now(); + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + let after = SystemTime::now(); + + assert_eq!(snapshot.version(), 1); + assert_eq!( + describe_calls.load(Ordering::SeqCst), + 1, + "binding must issue exactly one describe" + ); + + let headers = describe_headers + .lock() + .unwrap() + .clone() + .expect("binding describe headers"); + assert_eq!( + headers + .get("x-lancedb-min-version") + .expect("binding describe must send x-lancedb-min-version") + .to_str() + .unwrap(), + "7" + ); + assert_eq!( + headers + .get("x-lancedb-min-read-version") + .expect("binding describe must send x-lancedb-min-read-version") + .to_str() + .unwrap(), + "100" + ); + let sent = parse_min_timestamp(&headers); + assert!( + sent >= before - FRESHNESS_TOLERANCE && sent <= after + FRESHNESS_TOLERANCE, + "binding describe must send x-lancedb-min-timestamp from consistency interval" + ); + } + + /// Malformed schema/type (and description parse) failures for binding must + /// be payload-free Http protocol errors with the real request ID. + #[tokio::test] + async fn test_generated_column_binding_snapshot_malformed_schema_is_sanitized_protocol_error() { + const MARKER: &str = "SENSITIVE_BINDING_SCHEMA_MARKER_7f3c9e2a"; + + let cases = [ + // Deserializes, but Arrow/JsonSchema conversion rejects the type. + format!( + r#"{{"version":1,"schema":{{"fields":[{{"name":"a","type":{{"type":"{MARKER}"}},"nullable":false}}]}},"field_ids":[0]}}"# + ), + // Completely malformed response body. + format!("not-json {MARKER}"), + // Valid JSON, but schema is not a JsonSchema object. + format!(r#"{{"version":1,"schema":"{MARKER}","field_ids":[0]}}"#), + ]; + + for body in cases { + let expected_request_id = Arc::new(std::sync::Mutex::new(None::)); + let expected_request_id_c = expected_request_id.clone(); + let body_for_handler = body.clone(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + let request_id = request + .headers() + .get("x-request-id") + .expect("client must set x-request-id") + .to_str() + .unwrap() + .to_string(); + *expected_request_id_c.lock().unwrap() = Some(request_id); + http::Response::builder() + .status(200) + .body(body_for_handler.clone()) + .unwrap() + }); + + let err = table + .generated_column_binding_snapshot() + .await + .expect_err("malformed binding describe must fail"); + match &err { + Error::Http { + request_id, + status_code, + source, + } => { + let expected = expected_request_id + .lock() + .unwrap() + .clone() + .expect("handler must observe request id"); + assert_eq!( + request_id, &expected, + "binding protocol errors must preserve the real request id" + ); + assert!(!request_id.is_empty()); + assert!( + status_code.is_none(), + "local protocol validation must not invent a status code" + ); + + let mut text = format!("{err:?}\n{err}\n{source}"); + let mut current = source.source(); + while let Some(inner) = current { + text.push('\n'); + text.push_str(&inner.to_string()); + current = inner.source(); + } + assert!( + !text.contains(MARKER), + "unique raw response marker must be absent from error chain: {text}" + ); + assert!( + !text.contains(&body), + "raw response body/schema payload must be absent from error chain: {text}" + ); + } + other => panic!("expected Http protocol error, got {other:?}"), + } + } + } + /// Test that schema cache expires after 30 seconds TTL #[tokio::test] async fn test_schema_cache_invalidation_after_ttl() { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 5a98b3ec3..46933f69d 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -54,6 +54,7 @@ use crate::database::listing::LANCE_FILE_EXTENSION; use crate::database::read_freshness::TableFreshness; use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; +use crate::function::GeneratedColumnBindingSnapshot; use crate::index::IndexStatistics; use crate::index::{Index, IndexBuilder}; use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType}; @@ -604,6 +605,17 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { fn id(&self) -> &str; /// Get the arrow [Schema] of the table. async fn schema(&self) -> Result; + /// Atomic generated-column binding snapshot for one table version. + /// + /// Default returns [`Error::NotSupported`] so existing third-party + /// [`BaseTable`] implementations keep compiling. Native and remote tables + /// override this with a single-snapshot projection of version, Arrow + /// fields, and Lance stable field IDs. + async fn generated_column_binding_snapshot(&self) -> Result { + Err(Error::NotSupported { + message: "generated_column_binding_snapshot is not supported on this table type".into(), + }) + } /// Count the number of rows in this table. async fn count_rows(&self, filter: Option) -> Result; /// Create a physical plan for the query. @@ -1116,6 +1128,17 @@ impl Table { self.inner.schema().await } + /// Atomic generated-column binding snapshot for one table version. + /// + /// Hidden implementation projection for generated-column call binding. Not + /// a catalog resource, Job, or replacement for [`Self::version`]. + #[doc(hidden)] + pub async fn generated_column_binding_snapshot( + &self, + ) -> Result { + self.inner.generated_column_binding_snapshot().await + } + /// Count the number of rows in this dataset. /// /// # Arguments @@ -3076,6 +3099,18 @@ impl BaseTable for NativeTable { Ok(Arc::new(Schema::from(&lance_schema))) } + async fn generated_column_binding_snapshot(&self) -> Result { + // One consistency-wrapper get(): version and Lance field IDs must come + // from the same Dataset snapshot. Do not compose schema() + version(). + let dataset = self.dataset.get().await?; + let version = dataset.version().version; + let lance_schema = dataset.schema(); + let field_ids: Vec = lance_schema.fields.iter().map(|field| field.id).collect(); + let arrow_schema = Schema::from(lance_schema); + let fields = arrow_schema.fields().iter().cloned().collect::>(); + GeneratedColumnBindingSnapshot::try_new(version, fields, field_ids) + } + async fn table_definition(&self) -> Result { let schema = self.schema().await?; TableDefinition::try_from_rich_schema(schema) @@ -5525,4 +5560,103 @@ mod tests { stats.fragment_stats.num_fragments ); } + + #[tokio::test] + async fn test_generated_column_binding_snapshot_matches_manifest() { + let tmp_dir = tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + let conn = ConnectBuilder::new(uri).execute().await.unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("text", DataType::Utf8, true), + Field::new("score", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![Some("a")])), + Arc::new(Int32Array::from(vec![1])), + ], + ) + .unwrap(); + let table = conn + .create_table("binding_native", batch) + .execute() + .await + .unwrap(); + + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + let public_schema = table.schema().await.unwrap(); + let version = table.version().await.unwrap(); + let native = table.as_native().expect("native table"); + let dataset = native.dataset.get().await.unwrap(); + let lance_schema = dataset.schema(); + + assert_eq!(snapshot.version(), version); + assert_eq!(snapshot.version(), dataset.version().version); + assert_eq!(snapshot.entries().len(), lance_schema.fields.len()); + assert_eq!(snapshot.entries().len(), public_schema.fields().len()); + + for (idx, entry) in snapshot.entries().iter().enumerate() { + let lance_field = &lance_schema.fields[idx]; + let public_field = public_schema.field(idx); + assert_eq!(entry.field_id(), lance_field.id); + assert_eq!(entry.field().name(), lance_field.name.as_str()); + assert_eq!(entry.field().name(), public_field.name()); + assert_eq!(entry.field().data_type(), public_field.data_type()); + assert!(!public_field.metadata().contains_key("lance:field_id")); + assert!(!entry.field().metadata().contains_key("lance:field_id")); + } + } + + #[tokio::test] + async fn test_generated_column_binding_snapshot_rename_preserves_id() { + let tmp_dir = tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + let conn = ConnectBuilder::new(uri).execute().await.unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("old_name", DataType::Int32, false), + Field::new("keep", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![7])), + Arc::new(StringArray::from(vec![Some("x")])), + ], + ) + .unwrap(); + let table = conn + .create_table("binding_rename", batch) + .execute() + .await + .unwrap(); + + let before = table.generated_column_binding_snapshot().await.unwrap(); + let old_entry = before.field("old_name").expect("old_name"); + let old_id = old_entry.field_id(); + let keep_id = before.field("keep").expect("keep").field_id(); + + table + .alter_columns(&[ColumnAlteration::new("old_name".into()).rename("new_name".into())]) + .await + .unwrap(); + + let after = table.generated_column_binding_snapshot().await.unwrap(); + assert!(after.field("old_name").is_none()); + let renamed = after.field("new_name").expect("new_name"); + assert_eq!(renamed.field_id(), old_id); + assert_eq!(after.field("keep").expect("keep").field_id(), keep_id); + assert!( + !table + .schema() + .await + .unwrap() + .field_with_name("new_name") + .unwrap() + .metadata() + .contains_key("lance:field_id") + ); + } } diff --git a/rust/lancedb/tests/first_class_generated_column_binding_snapshot_contract.rs b/rust/lancedb/tests/first_class_generated_column_binding_snapshot_contract.rs new file mode 100644 index 000000000..e3eca78ca --- /dev/null +++ b/rust/lancedb/tests/first_class_generated_column_binding_snapshot_contract.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Contract tests for GeneratedColumnBindingSnapshot (FF-029). +//! +//! Pins the hidden value projection and Table seam used by generated-column +//! call binding. These tests intentionally fail to compile until that API +//! exists. They do not submit Jobs, mutate generated-column state, or resolve +//! authored Function calls. + +use std::sync::Arc; + +use arrow_array::{Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use lancedb::connect; +use lancedb::function::{GeneratedColumnBindingEntry, GeneratedColumnBindingSnapshot}; +use lancedb::{Error, Result}; +use tempfile::tempdir; + +fn sample_fields() -> Vec { + vec![ + Arc::new(Field::new("text", DataType::Utf8, true)), + Arc::new(Field::new("score", DataType::Int32, false)), + Arc::new(Field::new("a.b", DataType::Utf8, true)), + ] +} + +#[test] +fn try_new_preserves_version_order_and_exact_lookup() -> Result<()> { + let fields = sample_fields(); + let snapshot = GeneratedColumnBindingSnapshot::try_new(7, fields.clone(), vec![3, 5, 9])?; + + assert_eq!(snapshot.version(), 7); + let entries = snapshot.entries(); + assert_eq!(entries.len(), 3); + assert_eq!(entries[0].field_id(), 3); + assert_eq!(entries[0].field().name(), "text"); + assert_eq!(entries[0].field().data_type(), &DataType::Utf8); + assert_eq!(entries[1].field_id(), 5); + assert_eq!(entries[1].field().name(), "score"); + assert_eq!(entries[2].field_id(), 9); + assert_eq!(entries[2].field().name(), "a.b"); + + let by_name = snapshot.field("score").expect("exact name"); + assert_eq!(by_name.field_id(), 5); + assert!(snapshot.field("Score").is_none()); + assert!(snapshot.field("a").is_none()); + let dotted = snapshot + .field("a.b") + .expect("literal dotted top-level name"); + assert_eq!(dotted.field_id(), 9); + assert_eq!(dotted.field().as_ref(), fields[2].as_ref()); + + // Type existence pin for the entry surface used by the next binding slice. + let _: &GeneratedColumnBindingEntry = by_name; + Ok(()) +} + +#[test] +fn try_new_rejects_invalid_projections() { + let fields = sample_fields(); + + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, fields.clone(), vec![1, 2]), + Err(Error::InvalidInput { .. }) + )); + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, fields.clone(), vec![1, 2, -1]), + Err(Error::InvalidInput { .. }) + )); + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, fields.clone(), vec![1, 2, 1]), + Err(Error::InvalidInput { .. }) + )); + + let duplicate_names = vec![ + Arc::new(Field::new("text", DataType::Utf8, true)), + Arc::new(Field::new("text", DataType::Int32, false)), + ]; + assert!(matches!( + GeneratedColumnBindingSnapshot::try_new(1, duplicate_names, vec![1, 2]), + Err(Error::InvalidInput { .. }) + )); +} + +#[tokio::test] +async fn table_seam_returns_atomic_native_snapshot() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let schema = Arc::new(Schema::new(vec![ + Field::new("text", DataType::Utf8, true), + Field::new("score", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec![Some("a")])), + Arc::new(Int32Array::from(vec![1])), + ], + )?; + let table = db.create_table("binding", batch).execute().await?; + + let snapshot = table.generated_column_binding_snapshot().await?; + let public_schema = table.schema().await?; + let version = table.version().await?; + + assert_eq!(snapshot.version(), version); + assert_eq!(snapshot.entries().len(), public_schema.fields().len()); + for (entry, field) in snapshot.entries().iter().zip(public_schema.fields()) { + assert_eq!(entry.field().name(), field.name()); + assert_eq!(entry.field().data_type(), field.data_type()); + assert!(entry.field_id() >= 0); + assert!(!field.metadata().contains_key("lance:field_id")); + assert!(!entry.field().metadata().contains_key("lance:field_id")); + } + Ok(()) +}