diff --git a/rust/lancedb/src/function/binding_snapshot.rs b/rust/lancedb/src/function/binding_snapshot.rs index b2868c483..e1f5a47ad 100644 --- a/rust/lancedb/src/function/binding_snapshot.rs +++ b/rust/lancedb/src/function/binding_snapshot.rs @@ -11,7 +11,9 @@ use std::collections::HashSet; use arrow_schema::FieldRef; -use super::{FunctionCall, invalid_input}; +use super::{ + FunctionCall, GENERATED_COLUMN_METADATA_KEY, GeneratedColumnDefinition, invalid_input, +}; use crate::Result; /// One top-level field identity from a single table snapshot. @@ -36,6 +38,32 @@ impl GeneratedColumnBindingEntry { pub fn field(&self) -> &FieldRef { &self.field } + + /// Strict generated-column definition from this entry's Arrow metadata. + /// + /// Reads only [`GENERATED_COLUMN_METADATA_KEY`] on the exact snapshot field + /// and decodes through + /// [`GeneratedColumnDefinition::from_metadata_json`] with + /// [`Self::field_id`] as the expected output identity. The same-snapshot + /// stable field ID is mandatory so decode rejects metadata whose embedded + /// `output_field_id` does not match this entry; name/ordinal/hash fallbacks + /// are not used. + /// + /// Returns [`Ok`]`(`[`None`]`)` when the key is absent. Present but invalid + /// metadata fails closed as [`crate::Error::InvalidInput`] with a short + /// field-ID diagnostic that does not echo the raw metadata payload. + pub(crate) fn generated_column_definition(&self) -> Result> { + let Some(raw) = self.field.metadata().get(GENERATED_COLUMN_METADATA_KEY) else { + return Ok(None); + }; + match GeneratedColumnDefinition::from_metadata_json(raw, self.field_id) { + Ok(definition) => Ok(Some(definition)), + Err(_) => Err(invalid_input(format!( + "invalid generated-column metadata for field id {}", + self.field_id + ))), + } + } } /// Atomic table snapshot projection for generated-column call binding. @@ -356,4 +384,177 @@ mod tests { GeneratedColumnBindingSnapshot::try_new(1, Vec::::new(), vec![]).unwrap(); empty.validate_field_arguments(&literal_only).unwrap(); } + + fn status_sample_function() -> crate::function::Function { + use crate::function::{ + Function, FunctionId, FunctionOutput, FunctionParameter, FunctionSignature, + }; + Function::new( + FunctionId::try_new("fn.exact.status.binding").unwrap(), + FunctionSignature::try_new( + vec![FunctionParameter::new("label", DataType::Utf8)], + FunctionOutput::new(DataType::Int32, true), + ) + .unwrap(), + ) + } + + fn status_sample_call() -> crate::function::FunctionCall { + use crate::function::{FunctionArgument, FunctionCall}; + use arrow_array::{ArrayRef, StringArray}; + let function = status_sample_function(); + FunctionCall::try_new( + &function, + vec![( + "label".to_string(), + FunctionArgument::try_literal( + Arc::new(StringArray::from(vec![Some("ok")])) as ArrayRef + ) + .unwrap(), + )], + ) + .unwrap() + } + + fn definition_json( + output_field_id: i32, + dependency_epoch: u64, + materialized_epoch: u64, + ) -> String { + use crate::function::GeneratedColumnDefinition; + GeneratedColumnDefinition::try_new( + output_field_id, + status_sample_call(), + dependency_epoch, + materialized_epoch, + ) + .unwrap() + .to_metadata_json() + .unwrap() + } + + fn entry_with_metadata( + name: &str, + field_id: i32, + metadata_json: Option<&str>, + ) -> GeneratedColumnBindingEntry { + use crate::function::GENERATED_COLUMN_METADATA_KEY; + let field = if let Some(json) = metadata_json { + Field::new(name, DataType::Int32, true).with_metadata( + [(GENERATED_COLUMN_METADATA_KEY.to_string(), json.to_string())].into(), + ) + } else { + Field::new(name, DataType::Int32, true) + }; + let snapshot = + GeneratedColumnBindingSnapshot::try_new(1, vec![Arc::new(field)], vec![field_id]) + .unwrap(); + snapshot.entries()[0].clone() + } + + #[test] + fn generated_column_definition_absent_returns_none() { + let entry = entry_with_metadata("ordinary", 3, None); + let got = entry.generated_column_definition().unwrap(); + assert!(got.is_none()); + } + + #[test] + fn generated_column_definition_decodes_complete_and_incomplete() { + use crate::function::{GeneratedColumnDefinition, GeneratedColumnStatus}; + + let complete_json = definition_json(5, 3, 3); + let complete_entry = entry_with_metadata("gen_complete", 5, Some(&complete_json)); + let complete = complete_entry + .generated_column_definition() + .unwrap() + .expect("complete metadata present"); + assert_eq!(complete.output_field_id(), 5); + assert_eq!(complete.dependency_epoch(), 3); + assert_eq!(complete.materialized_epoch(), 3); + assert_eq!(complete.status(), GeneratedColumnStatus::Complete); + assert_eq!( + complete, + GeneratedColumnDefinition::from_metadata_json(&complete_json, 5).unwrap() + ); + + let incomplete_json = definition_json(7, 4, 2); + let incomplete_entry = entry_with_metadata("gen_incomplete", 7, Some(&incomplete_json)); + let incomplete = incomplete_entry + .generated_column_definition() + .unwrap() + .expect("incomplete metadata present"); + assert_eq!(incomplete.output_field_id(), 7); + assert_eq!(incomplete.dependency_epoch(), 4); + assert_eq!(incomplete.materialized_epoch(), 2); + assert_eq!(incomplete.status(), GeneratedColumnStatus::Incomplete); + assert_eq!( + incomplete, + GeneratedColumnDefinition::from_metadata_json(&incomplete_json, 7).unwrap() + ); + } + + #[test] + fn generated_column_definition_fail_closed_for_invalid_metadata() { + use crate::function::GENERATED_COLUMN_METADATA_KEY; + + let field_id = 9i32; + let valid = definition_json(field_id, 2, 2); + let mut mismatched: serde_json::Value = serde_json::from_str(&valid).unwrap(); + mismatched["output_field_id"] = serde_json::json!(field_id + 1); + + let mut unsupported: serde_json::Value = serde_json::from_str(&valid).unwrap(); + unsupported["format_version"] = serde_json::json!(2); + + let mut reversed: serde_json::Value = serde_json::from_str(&valid).unwrap(); + reversed["dependency_epoch"] = serde_json::json!(1); + reversed["materialized_epoch"] = serde_json::json!(2); + + let malformed_json = "{not-json"; + let mut malformed_call: serde_json::Value = serde_json::from_str(&valid).unwrap(); + malformed_call["function_call"] = serde_json::json!("not-an-object"); + + for raw in [ + mismatched.to_string(), + unsupported.to_string(), + reversed.to_string(), + malformed_json.to_string(), + malformed_call.to_string(), + ] { + let entry = entry_with_metadata("gen_bad", field_id, Some(&raw)); + assert!( + entry + .field() + .metadata() + .contains_key(GENERATED_COLUMN_METADATA_KEY), + "fixture must carry generated-column metadata" + ); + let err = entry.generated_column_definition().unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + } + } + + #[test] + fn generated_column_definition_errors_omit_raw_metadata_marker() { + const MARKER: &str = "SENSITIVE_STATUS_METADATA_MARKER_b3d1_9f2e"; + let raw = format!( + r#"{{"format_version":1,"output_field_id":3,"function_call":{MARKER},"dependency_epoch":1,"materialized_epoch":1}}"# + ); + assert!(raw.contains(MARKER)); + let entry = entry_with_metadata("gen_redact", 3, Some(&raw)); + let err = entry.generated_column_definition().unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + let text = format!("{err}\n{err:?}"); + assert!( + !text.contains(MARKER), + "status definition diagnostics must not echo raw metadata marker: {text}" + ); + assert!( + !text.contains(&raw), + "status definition diagnostics must not echo raw metadata payload: {text}" + ); + } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 88a143e40..22effc81c 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -8174,6 +8174,498 @@ mod tests { } } + // ------------------------------------------------------------------------- + // Generated-column status projection (table semantic read) + // ------------------------------------------------------------------------- + + fn status_remote_definition_json( + output_field_id: i32, + dependency_epoch: u64, + materialized_epoch: u64, + ) -> String { + use crate::function::{ + Function, FunctionArgument, FunctionCall, FunctionId, FunctionOutput, + FunctionParameter, FunctionSignature, GeneratedColumnDefinition, + }; + use arrow_array::StringArray; + + let function = Function::new( + FunctionId::try_new("fn.exact.status.remote").unwrap(), + FunctionSignature::try_new( + vec![FunctionParameter::new("label", DataType::Utf8)], + FunctionOutput::new(DataType::Int32, true), + ) + .unwrap(), + ); + let call = FunctionCall::try_new( + &function, + vec![( + "label".to_string(), + FunctionArgument::try_literal( + Arc::new(StringArray::from(vec![Some("ok")])) as arrow_array::ArrayRef + ) + .unwrap(), + )], + ) + .unwrap(); + GeneratedColumnDefinition::try_new( + output_field_id, + call, + dependency_epoch, + materialized_epoch, + ) + .unwrap() + .to_metadata_json() + .unwrap() + } + + fn status_remote_schema_with_epochs() -> Schema { + use crate::function::GENERATED_COLUMN_METADATA_KEY; + + let complete = Field::new("complete_col", DataType::Int32, true).with_metadata( + [( + GENERATED_COLUMN_METADATA_KEY.to_string(), + status_remote_definition_json(5, 3, 3), + )] + .into(), + ); + let incomplete = Field::new("incomplete_col", DataType::Int32, true).with_metadata( + [( + GENERATED_COLUMN_METADATA_KEY.to_string(), + status_remote_definition_json(7, 4, 1), + )] + .into(), + ); + Schema::new(vec![ + Field::new("ordinary", DataType::Utf8, true), + complete, + incomplete, + ]) + } + + #[tokio::test] + async fn test_generated_column_status_one_describe_complete_and_incomplete() { + use crate::function::GeneratedColumnStatus; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let schema = status_remote_schema_with_epochs(); + let body = describe_response_with_field_ids(11, &schema, &[1, 5, 7]); + + 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() + }); + + assert_eq!( + table.generated_column_status("complete_col").await.unwrap(), + GeneratedColumnStatus::Complete + ); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + + assert_eq!( + table + .generated_column_status("incomplete_col") + .await + .unwrap(), + GeneratedColumnStatus::Incomplete + ); + assert_eq!(call_count.load(Ordering::SeqCst), 2); + + for name in ["missing", "Complete_Col", "ordinary"] { + let before = call_count.load(Ordering::SeqCst); + let err = table.generated_column_status(name).await.unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput for `{name}`, got {err:?}" + ); + assert_eq!( + call_count.load(Ordering::SeqCst), + before + 1, + "fail-closed status lookup still uses exactly one describe" + ); + } + } + + #[tokio::test] + async fn test_generated_column_status_respects_checkout_version_and_branch() { + use crate::function::GeneratedColumnStatus; + use lance::dataset::refs::Ref; + + let schema = status_remote_schema_with_epochs(); + let describe_pinned = describe_response_with_field_ids(3, &schema, &[1, 5, 7]); + let describe_branch = describe_response_with_field_ids(9, &schema, &[1, 5, 7]); + + 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(), + 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(); + assert_eq!( + table.generated_column_status("complete_col").await.unwrap(), + GeneratedColumnStatus::Complete + ); + + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + assert_eq!( + branch + .generated_column_status("incomplete_col") + .await + .unwrap(), + GeneratedColumnStatus::Incomplete + ); + } + + #[tokio::test] + async fn test_generated_column_status_preserves_freshness_headers() { + use crate::function::GeneratedColumnStatus; + + let schema = status_remote_schema_with_epochs(); + // Describe body version (1) is intentionally distinct from write (9) / read (7) + // watermarks so status must not confuse describe version with freshness state. + let describe_body = describe_response_with_field_ids(1, &schema, &[1, 5, 7]); + 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 count_calls = Arc::new(AtomicUsize::new(0)); + let count_calls_c = count_calls.clone(); + let post_status_headers = Arc::new(std::sync::Mutex::new(None::)); + let post_status_headers_c = post_status_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":9}"#.to_string()) + .unwrap(), + "/v1/table/my_table/count_rows/" => { + let n = count_calls_c.fetch_add(1, Ordering::SeqCst); + if n == 0 { + http::Response::builder() + .status(200) + .header("x-lancedb-version", "7") + .body("42".to_string()) + .unwrap() + } else { + *post_status_headers_c.lock().unwrap() = Some(request.headers().clone()); + http::Response::builder() + .status(200) + .header("x-lancedb-version", "7") + .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), + ); + + // Distinct preexisting write (9) and read (7) watermarks. + table.update().column("a", "a + 1").execute().await.unwrap(); + table.count_rows(None).await.unwrap(); + assert_eq!(count_calls.load(Ordering::SeqCst), 1); + + let before = SystemTime::now(); + assert_eq!( + table.generated_column_status("complete_col").await.unwrap(), + GeneratedColumnStatus::Complete + ); + let after = SystemTime::now(); + + assert_eq!( + describe_calls.load(Ordering::SeqCst), + 1, + "status must issue exactly one describe" + ); + let status_headers = describe_headers + .lock() + .unwrap() + .clone() + .expect("status describe headers"); + assert_eq!( + status_headers + .get("x-lancedb-min-version") + .expect("status describe must send x-lancedb-min-version") + .to_str() + .unwrap(), + "9" + ); + assert_eq!( + status_headers + .get("x-lancedb-min-read-version") + .expect("status describe must send x-lancedb-min-read-version") + .to_str() + .unwrap(), + "7" + ); + assert_ne!( + status_headers + .get("x-lancedb-min-version") + .unwrap() + .to_str() + .unwrap(), + "1", + "write watermark must not become the status describe body version" + ); + assert_ne!( + status_headers + .get("x-lancedb-min-read-version") + .unwrap() + .to_str() + .unwrap(), + "1", + "read watermark must not become the status describe body version" + ); + let sent = parse_min_timestamp(&status_headers); + assert!( + sent >= before - FRESHNESS_TOLERANCE && sent <= after + FRESHNESS_TOLERANCE, + "status describe must send x-lancedb-min-timestamp from consistency interval" + ); + + // Subsequent ordinary read must still carry the same preexisting watermarks. + table.count_rows(None).await.unwrap(); + assert_eq!(count_calls.load(Ordering::SeqCst), 2); + let read_headers = post_status_headers + .lock() + .unwrap() + .clone() + .expect("post-status count_rows headers"); + assert_eq!( + read_headers + .get("x-lancedb-min-version") + .expect("post-status read must send x-lancedb-min-version") + .to_str() + .unwrap(), + "9" + ); + assert_eq!( + read_headers + .get("x-lancedb-min-read-version") + .expect("post-status read must send x-lancedb-min-read-version") + .to_str() + .unwrap(), + "7" + ); + assert_eq!( + describe_calls.load(Ordering::SeqCst), + 1, + "ordinary post-status read must not issue another describe" + ); + } + + #[tokio::test] + async fn test_generated_column_status_bypasses_seeded_schema_cache() { + use crate::function::GeneratedColumnStatus; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let schema = status_remote_schema_with_epochs(); + let seeded = describe_response(&schema); + let with_ids = describe_response_with_field_ids(7, &schema, &[1, 5, 7]); + + 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); + + let schema_before = table.schema().await.unwrap(); + assert_eq!( + schema_before + .field_with_name("complete_col") + .unwrap() + .name(), + "complete_col" + ); + assert_eq!(call_count.load(Ordering::SeqCst), 0); + + assert_eq!( + table.generated_column_status("complete_col").await.unwrap(), + GeneratedColumnStatus::Complete + ); + assert_eq!( + call_count.load(Ordering::SeqCst), + 1, + "status must issue exactly one authoritative describe for IDs" + ); + + let schema_after = table.schema().await.unwrap(); + assert!( + Arc::ptr_eq(&schema_before, &schema_after), + "status must leave the preexisting schema cache Arc installed" + ); + assert_eq!( + call_count.load(Ordering::SeqCst), + 1, + "ordinary schema read after status must add zero requests" + ); + } + + #[tokio::test] + async fn test_generated_column_status_empty_name_zero_requests() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let table = + Table::new_with_handler("my_table", move |_request| -> http::Response { + call_count_clone.fetch_add(1, Ordering::SeqCst); + panic!("empty status name must not issue HTTP requests"); + }); + + let err = table.generated_column_status("").await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + assert_eq!(call_count.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn test_generated_column_status_missing_ids_not_supported() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let schema = status_remote_schema_with_epochs(); + 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_status("complete_col") + .await + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_generated_column_status_invalid_ids_are_protocol_errors() { + let schema = status_remote_schema_with_epochs(); + for field_ids in [vec![1, 5], vec![1, 5, -7], vec![1, 5, 5]] { + let body = describe_response_with_field_ids(1, &schema, &field_ids); + let body_for_assert = body.clone(); + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + 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_status("complete_col") + .await + .unwrap_err(); + match err { + Error::Http { + request_id, + status_code, + source, + } => { + assert!(!request_id.is_empty()); + assert!(status_code.is_none()); + assert!(!source.to_string().contains(&body_for_assert)); + } + other => panic!("expected Http protocol error, got {other:?}"), + } + assert_eq!(call_count.load(Ordering::SeqCst), 1); + } + } + + #[tokio::test] + async fn test_generated_column_status_malformed_metadata_fail_closed() { + use crate::function::GENERATED_COLUMN_METADATA_KEY; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let bad = Field::new("gen_bad", DataType::Int32, true).with_metadata( + [( + GENERATED_COLUMN_METADATA_KEY.to_string(), + r#"{"format_version":1,"output_field_id":99,"function_call":{},"dependency_epoch":1,"materialized_epoch":1}"# + .to_string(), + )] + .into(), + ); + let schema = Schema::new(vec![bad]); + let body = describe_response_with_field_ids(1, &schema, &[3]); + + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path(); + assert_eq!(path, "/v1/table/my_table/describe/"); + assert!( + !path.contains("job") + && !path.contains("generated") + && !path.contains("function") + && !path.contains("submit"), + "status must not hit create/Job/Function endpoints: {path}" + ); + call_count_clone.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + + let err = table.generated_column_status("gen_bad").await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + } + // ------------------------------------------------------------------------- // CreateGeneratedColumnJobSpec remote table submit transport (FF-031) // ------------------------------------------------------------------------- diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index ccf9814cc..ec240f2cc 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -54,7 +54,9 @@ 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::{CreateGeneratedColumnJobSpec, GeneratedColumnBindingSnapshot}; +use crate::function::{ + CreateGeneratedColumnJobSpec, GeneratedColumnBindingSnapshot, GeneratedColumnStatus, +}; use crate::index::IndexStatistics; use crate::index::{Index, IndexBuilder}; use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType}; @@ -1156,6 +1158,48 @@ impl Table { self.inner.generated_column_binding_snapshot().await } + /// Project generated-column completeness for one top-level column name. + /// + /// Projection-only: loads one authoritative + /// [`Self::generated_column_binding_snapshot`], looks up the exact + /// case-sensitive top-level name (`.` is literal, not a nested path), and + /// returns that field's [`GeneratedColumnStatus::Complete`] or + /// [`GeneratedColumnStatus::Incomplete`]. Does not submit or wait on a Job, + /// validate Function catalog identity, execute a UDF, or mutate + /// data/schema/cache state. + /// + /// Returns [`Error::InvalidInput`] for an empty name (before any table + /// access), a missing top-level field, or an ordinary field without a + /// valid generated-column definition. Tables whose binding snapshot is + /// unsupported surface that as [`Error::NotSupported`] (or the snapshot's + /// existing protocol error) without a separate status transport. + pub async fn generated_column_status( + &self, + column_name: impl AsRef, + ) -> Result { + let column_name = column_name.as_ref(); + if column_name.is_empty() { + return Err(Error::InvalidInput { + message: "generated column name must not be empty".into(), + }); + } + + let snapshot = self.inner.generated_column_binding_snapshot().await?; + let Some(entry) = snapshot.field(column_name) else { + return Err(Error::InvalidInput { + message: format!( + "generated column '{column_name}' was not found in the table schema" + ), + }); + }; + let Some(definition) = entry.generated_column_definition()? else { + return Err(Error::InvalidInput { + message: format!("column '{column_name}' is not a generated column"), + }); + }; + Ok(definition.status()) + } + /// Submit a create-generated-column Job from an already-bound snapshot pair. /// /// Hidden table submit seam for enterprise generated-column create. Input @@ -5693,4 +5737,226 @@ mod tests { .contains_key("lance:field_id") ); } + + fn status_projection_function() -> crate::function::Function { + use crate::function::{ + Function, FunctionId, FunctionOutput, FunctionParameter, FunctionSignature, + }; + Function::new( + FunctionId::try_new("fn.exact.status.native").unwrap(), + FunctionSignature::try_new( + vec![FunctionParameter::new("label", DataType::Utf8)], + FunctionOutput::new(DataType::Int32, true), + ) + .unwrap(), + ) + } + + fn status_projection_definition( + output_field_id: i32, + dependency_epoch: u64, + materialized_epoch: u64, + ) -> crate::function::GeneratedColumnDefinition { + use crate::function::{FunctionArgument, FunctionCall, GeneratedColumnDefinition}; + let function = status_projection_function(); + let call = FunctionCall::try_new( + &function, + vec![( + "label".to_string(), + FunctionArgument::try_literal( + Arc::new(StringArray::from(vec![Some("ok")])) as arrow_array::ArrayRef + ) + .unwrap(), + )], + ) + .unwrap(); + GeneratedColumnDefinition::try_new( + output_field_id, + call, + dependency_epoch, + materialized_epoch, + ) + .unwrap() + } + + async fn plant_generated_column_metadata( + table: &Table, + column: &str, + dependency_epoch: u64, + materialized_epoch: u64, + ) -> i32 { + use crate::function::GENERATED_COLUMN_METADATA_KEY; + + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + let field_id = snapshot.field(column).expect(column).field_id(); + let json = status_projection_definition(field_id, dependency_epoch, materialized_epoch) + .to_metadata_json() + .unwrap(); + table + .update_field_metadata(&[ + FieldMetadataUpdate::new(column).set(GENERATED_COLUMN_METADATA_KEY, json) + ]) + .await + .unwrap(); + field_id + } + + async fn create_status_projection_table(name: &str) -> (tempfile::TempDir, Table) { + 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("gen_out", DataType::Int32, true), + Field::new("ordinary", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some("x")])), + ], + ) + .unwrap(); + let table = conn.create_table(name, batch).execute().await.unwrap(); + (tmp_dir, table) + } + + #[tokio::test] + async fn test_generated_column_status_complete_and_incomplete() { + use crate::function::GeneratedColumnStatus; + + let (_tmp, table) = create_status_projection_table("status_epochs").await; + plant_generated_column_metadata(&table, "gen_out", 3, 3).await; + assert_eq!( + table.generated_column_status("gen_out").await.unwrap(), + GeneratedColumnStatus::Complete + ); + + plant_generated_column_metadata(&table, "gen_out", 5, 2).await; + assert_eq!( + table.generated_column_status("gen_out").await.unwrap(), + GeneratedColumnStatus::Incomplete + ); + } + + #[tokio::test] + async fn test_generated_column_status_reject_empty_missing_case_ordinary() { + let (_tmp, table) = create_status_projection_table("status_reject").await; + plant_generated_column_metadata(&table, "gen_out", 1, 1).await; + + for name in ["", "missing", "Gen_Out", "GEN_OUT", "ordinary"] { + let err = table.generated_column_status(name).await.unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput for `{name}`, got {err:?}" + ); + } + } + + #[tokio::test] + async fn test_generated_column_status_rename_preserves_status_and_id() { + use crate::function::GeneratedColumnStatus; + + let (_tmp, table) = create_status_projection_table("status_rename").await; + let old_id = plant_generated_column_metadata(&table, "gen_out", 4, 4).await; + assert_eq!( + table.generated_column_status("gen_out").await.unwrap(), + GeneratedColumnStatus::Complete + ); + + table + .alter_columns(&[ColumnAlteration::new("gen_out".into()).rename("gen_renamed".into())]) + .await + .unwrap(); + + assert_eq!( + table.generated_column_status("gen_renamed").await.unwrap(), + GeneratedColumnStatus::Complete + ); + let after = table.generated_column_binding_snapshot().await.unwrap(); + assert_eq!(after.field("gen_renamed").unwrap().field_id(), old_id); + let err = table.generated_column_status("gen_out").await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + } + + #[tokio::test] + async fn test_generated_column_status_drop_recreate_does_not_inherit() { + use crate::function::GeneratedColumnStatus; + + let (_tmp, table) = create_status_projection_table("status_drop_recreate").await; + let old_id = plant_generated_column_metadata(&table, "gen_out", 2, 2).await; + assert_eq!( + table.generated_column_status("gen_out").await.unwrap(), + GeneratedColumnStatus::Complete + ); + + table.drop_columns(&["gen_out"]).await.unwrap(); + table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "gen_out".into(), + "cast(NULL as int)".into(), + )])) + .execute() + .await + .unwrap(); + + let after = table.generated_column_binding_snapshot().await.unwrap(); + let recreated = after.field("gen_out").expect("recreated"); + assert_ne!(recreated.field_id(), old_id); + assert!( + !recreated + .field() + .metadata() + .contains_key(crate::function::GENERATED_COLUMN_METADATA_KEY) + ); + let err = table.generated_column_status("gen_out").await.unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { .. }), + "drop/recreate ordinary field must not inherit status, got {err:?}" + ); + } + + #[tokio::test] + async fn test_generated_column_status_is_read_only() { + use crate::function::GeneratedColumnStatus; + use futures::TryStreamExt; + + let (_tmp, table) = create_status_projection_table("status_readonly").await; + // Incomplete on purpose: later query guards must reject references to gen_out, + // so this read-only proof only projects the ordinary non-generated column. + plant_generated_column_metadata(&table, "gen_out", 6, 3).await; + + let version_before = table.version().await.unwrap(); + let schema_before = table.schema().await.unwrap(); + let data_before = table + .query() + .select(Select::columns(&["ordinary"])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(data_before[0].num_columns(), 1); + assert_eq!(data_before[0].schema().field(0).name(), "ordinary"); + + assert_eq!( + table.generated_column_status("gen_out").await.unwrap(), + GeneratedColumnStatus::Incomplete + ); + + assert_eq!(table.version().await.unwrap(), version_before); + assert_eq!(table.schema().await.unwrap(), schema_before); + let data_after = table + .query() + .select(Select::columns(&["ordinary"])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(data_after, data_before); + } }