mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 03:58:26 +00:00
feat: expose generated column status
This commit is contained in:
@@ -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}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user