feat: expose generated column status

This commit is contained in:
Xuanwo
2026-08-12 12:17:02 +08:00
parent 7a94ab7d6c
commit 93ba442ac2
3 changed files with 961 additions and 2 deletions
+202 -1
View File
@@ -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<Option<GeneratedColumnDefinition>> {
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::<FieldRef>::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}"
);
}
}
+492
View File
@@ -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::<http::HeaderMap>));
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::<http::HeaderMap>));
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<dyn BaseTable>);
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<String> {
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)
// -------------------------------------------------------------------------
+267 -1
View File
@@ -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<str>,
) -> Result<GeneratedColumnStatus> {
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::<Vec<_>>()
.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::<Vec<_>>()
.await
.unwrap();
assert_eq!(data_after, data_before);
}
}