diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 2b95cb34c..ae55e655d 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -22,7 +22,7 @@ use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; -use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; +use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef}; use datafusion_common::tree_node::TreeNode; use datafusion_physical_plan::PhysicalExpr; use lance::dataset::NewColumnTransform; @@ -1273,6 +1273,11 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< /// refresh time: that the expression parses, that every column it reads /// exists, and that the target name is free. A declaration that survives this /// is one a refresh can always act on. +/// +/// Each accepted column joins the schema the next one resolves against, so a +/// batch may declare `a` and then `b = a + 1` in one commit. Refresh fills a +/// column's computed inputs before the column, so the order of refresh calls +/// does not matter. pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result> { if columns.is_empty() { return Err(Error::InvalidInput { @@ -1280,11 +1285,11 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result = Vec::with_capacity(columns.len()); for (name, expression) in columns { - if schema.field_with_name(name).is_ok() || declared.contains(&name.as_str()) { + if schema.field_with_name(name).is_ok() { return Err(Error::ColumnAlreadyExists { name: name.clone() }); } @@ -1292,16 +1297,50 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result(), + schema.metadata().clone(), + )); + fields.push(field); } Ok(fields) } +/// Run the schema-level checks of +/// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) against +/// `schema` without committing: the Function-binding guard and the planning of +/// every declaration. For callers that stage declarations behind other work +/// and need those rejections before any of it lands. +/// +/// Only the schema is consulted. Declaring also refuses a table with an LSM +/// write spec or retained SSTables; that is table state, checked at commit. +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow_schema::{DataType, Field, Schema}; +/// use lancedb::table::computed_columns::validate_declarations; +/// +/// let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); +/// let declarations = vec![ +/// ("a".to_string(), "x + 1".to_string()), +/// ("b".to_string(), "a * 2".to_string()), +/// ]; +/// assert!(validate_declarations(schema.clone(), &declarations).is_ok()); +/// assert!(validate_declarations(schema, &[("c".into(), "random()".into())]).is_err()); +/// ``` +pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<()> { + ensure_no_function_bindings_for_mutation(schema.as_ref(), "schema evolution")?; + plan(schema, columns).map(drop) +} + /// Build the transform that declares `columns` against `schema`. /// /// An all-null column is how a binding with no values yet is carried into a @@ -1340,6 +1379,22 @@ pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &st #[cfg(test)] mod tests { + /// The gate's reproducer: the validator applies the same schema-level + /// guard declaring does, so a staging caller is refused before it commits + /// anything else. + #[test] + fn test_validate_declarations_matches_schema_admission_barriers() { + let schema = Arc::new(ArrowSchema::new_with_metadata( + vec![ArrowField::new("x", DataType::Int32, true)], + HashMap::from([( + FUNCTION_BINDINGS_META_KEY.to_string(), + "not valid binding metadata".to_string(), + )]), + )); + let declarations = vec![("a".to_string(), "x + 1".to_string())]; + assert!(super::validate_declarations(schema, &declarations).is_err()); + } + #[test] fn output_arrow_type_grammar_matches_the_shared_golden() { let golden: serde_json::Value = serde_json::from_str(include_str!( @@ -1582,6 +1637,40 @@ mod tests { assert!(declared(&table).await.is_empty()); } + /// A batch may build on itself: one commit, and the later entry's inputs + /// name the earlier one. + #[tokio::test] + async fn test_a_declaration_may_read_one_declared_before_it() { + let table = table_with_ints("chain").await; + let before = table.version().await.unwrap(); + add_computed( + &table, + &[("a".into(), "x + 1".into()), ("b".into(), "a * 2".into())], + ) + .await + .unwrap(); + assert_eq!(table.version().await.unwrap(), before + 1); + let declared = declared(&table).await; + assert_eq!(declared[1].name, "b"); + assert_eq!(declared[1].inputs, vec!["a".to_string()]); + + // Order is the dependency order; reading ahead is still unknown. + let err = add_computed( + &table, + &[("c".into(), "d + 1".into()), ("d".into(), "x + 1".into())], + ) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "c")); + assert!( + validate_declarations( + table.schema().await.unwrap(), + &[("e".into(), "random()".into())] + ) + .is_err() + ); + } + /// A column added by an ordinary transform is materialized, not bound, so /// it carries no declaration to report. #[tokio::test] diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 35f883411..b02fdc136 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -7,6 +7,16 @@ //! therefore idempotent and does not observe input mutation -- once a row is //! filled, changing what the expression reads leaves the stored result alone. //! +//! A column's computed inputs are filled first -- the dependency graph is +//! walked once, each reachable column filled once in dependency order, each +//! fill its own commit. Every fill in the pass, the requested column's +//! included, covers only the fragments of the snapshot the pass started +//! from: a commit may rebase over a concurrent append, and the fragment that +//! admits carries placeholder nulls no earlier fill covered, so it waits for +//! a later refresh rather than being read as values. Two concurrent fills of +//! one input collide on its field in lance's conflict check, so a dependent +//! fill can only commit over inputs that were durable when it read them. +//! //! Two passes per fragment. The first scans only the unfilled live rows and //! evaluates the expression over them, which yields the exact fill count and //! decides whether the fragment is staged at all -- a fragment where nothing @@ -19,6 +29,7 @@ //! inputs masked to null first, so a poison value in a row nobody is filling //! cannot fail the refresh. +use std::collections::HashSet; use std::sync::Arc; use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions}; @@ -41,7 +52,8 @@ use crate::{Error, Result}; /// The result of refreshing a computed column. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct RefreshColumnResult { - /// Rows that had a value computed. + /// Rows that had a value computed, in the requested column only; inputs + /// filled on its behalf are not counted. #[serde(default)] pub rows_filled: u64, /// The commit version associated with the operation. @@ -51,7 +63,79 @@ pub struct RefreshColumnResult { struct RefreshExecution { result: RefreshColumnResult, + /// The snapshot the requested column was evaluated against, after its + /// inputs were filled. source_version: u64, + /// The last version any fill in the pass committed, inputs included. + published_version: Option, + /// Unfilled live rows of the requested column in fragments the pass did + /// not cover, counted on the snapshot the result reports -- the published + /// one, or the source when nothing was published. + rows_deferred: u64, +} + +/// One column's fill against one snapshot. +struct Fill { + rows_filled: u64, + /// The dataset the commit produced, if anything was filled. + committed: Option>, +} + +/// Test-only one-shot pauses before a named column's fill: +/// `(column, reached, resume)`, one per column, consumed when hit. +#[cfg(test)] +type Pause = (String, Arc, Arc); +#[cfg(test)] +static PAUSES_BEFORE_FILL: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); + +/// The SQL-computed columns `column` reads, transitively, each once, in an +/// order that fills every column after the columns it reads. +/// +/// Declarations are acyclic by construction: a column can only read what +/// existed when it was declared. +fn dependency_order(schema: &ArrowSchema, column: &str) -> Result> { + fn visit( + schema: &ArrowSchema, + column: &str, + visited: &mut HashSet, + order: &mut Vec, + ) -> Result<()> { + let Some(declaration) = schema + .field_with_name(column) + .ok() + .and_then(computed_column_from_field) + else { + return Ok(()); + }; + for input in &declaration.inputs { + let input = super::computed_columns::root(input); + if visited.contains(input) { + continue; + } + let Some(input_declaration) = schema + .field_with_name(input) + .ok() + .and_then(computed_column_from_field) + else { + continue; + }; + if !matches!(input_declaration.kind, ComputedColumnKind::Sql { .. }) { + return Err(Error::NotSupported { + message: format!( + "computed column '{column}' reads '{input}', which this refresh \ + cannot fill first; refresh '{input}' before '{column}'" + ), + }); + } + visit(schema, input, visited, order)?; + visited.insert(input.to_string()); + order.push(input.to_string()); + } + Ok(()) + } + let mut order = Vec::new(); + visit(schema, column, &mut HashSet::new(), &mut order)?; + Ok(order) } /// Internal implementation of the refresh logic. @@ -70,9 +154,96 @@ async fn execute_refresh_column_with_source( ) -> Result { table.dataset.ensure_mutable()?; ensure_no_lsm_write_spec(table).await?; - let dataset = table.dataset.get().await?; + let mut dataset = table.dataset.get().await?; + declared_expression(&dataset, column)?; - let expression = declared_expression(&dataset, column)?; + // The pass covers exactly these fragments. A commit below may rebase over + // a concurrent append, and the fragment that admits was scanned by no + // earlier fill, so it is excluded from every later one. + let eligible: HashSet = dataset + .get_fragments() + .iter() + .map(|f| f.id() as u64) + .collect(); + + let mut published_version = None; + for input in dependency_order(&ArrowSchema::from(dataset.schema()), column)? { + pause_before_fill(&input).await; + if let Some(committed) = fill_column(table, &dataset, &input, &eligible) + .await? + .committed + { + published_version = Some(committed.version().version); + dataset = committed; + } + } + + pause_before_fill(column).await; + let source_version = dataset.version().version; + let fill = fill_column(table, &dataset, column, &eligible).await?; + if let Some(committed) = &fill.committed { + published_version = Some(committed.version().version); + } + // Counted after the commit: it may have rebased over an append, and the + // fragment that admits is deferred but part of the published version. + let terminal = fill.committed.as_ref().unwrap_or(&dataset); + let rows_deferred = count_deferred(terminal, column, &eligible).await?; + Ok(RefreshExecution { + result: RefreshColumnResult { + rows_filled: fill.rows_filled, + version: published_version.unwrap_or(source_version), + }, + source_version, + published_version, + rows_deferred, + }) +} + +/// Unfilled live rows of `column` in `dataset`'s fragments outside `eligible`. +async fn count_deferred(dataset: &Dataset, column: &str, eligible: &HashSet) -> Result { + let deferred: Vec<_> = dataset + .get_fragments() + .into_iter() + .filter(|fragment| !eligible.contains(&(fragment.id() as u64))) + .map(|fragment| fragment.metadata().clone()) + .collect(); + if deferred.is_empty() { + return Ok(0); + } + let mut scanner = dataset.scan(); + scanner + .with_fragments(deferred) + .filter(&format!("{} IS NULL", quote_identifier(column)))?; + Ok(scanner.count_rows().await?) +} + +#[cfg(test)] +async fn pause_before_fill(column: &str) { + let pause = { + let mut pauses = PAUSES_BEFORE_FILL.lock().unwrap(); + pauses + .iter() + .position(|(paused, _, _)| paused == column) + .map(|index| pauses.remove(index)) + }; + if let Some((_, reached, resume)) = pause { + reached.notify_one(); + resume.notified().await; + } +} + +#[cfg(not(test))] +async fn pause_before_fill(_column: &str) {} + +/// Fill `column`'s unfilled live rows in the `eligible` fragments as they +/// stand in `dataset`, committing against that snapshot. +async fn fill_column( + table: &NativeTable, + dataset: &Arc, + column: &str, + eligible: &HashSet, +) -> Result { + let expression = declared_expression(dataset, column)?; let schema = Arc::new(ArrowSchema::from(dataset.schema())); let bound = Arc::new(super::computed_columns::bind(schema, column, &expression)?); let field = dataset @@ -91,23 +262,22 @@ async fn execute_refresh_column_with_source( let mut rows_filled = 0u64; let mut replacements = Vec::new(); for fragment in dataset.get_fragments() { - let gained = count_fragment_gains(&dataset, &fragment, &bound, column).await?; + if !eligible.contains(&(fragment.id() as u64)) { + continue; + } + let gained = count_fragment_gains(dataset, &fragment, &bound, column).await?; if gained == 0 { continue; } rows_filled += gained; - let values = fill_stream(&dataset, &fragment, bound.clone(), column).await?; + let values = fill_stream(dataset, &fragment, bound.clone(), column).await?; replacements.push(fragment.write_columns(values, &column_schema).await?); } if replacements.is_empty() { - let source_version = dataset.version().version; - return Ok(RefreshExecution { - result: RefreshColumnResult { - rows_filled: 0, - version: source_version, - }, - source_version, + return Ok(Fill { + rows_filled: 0, + committed: None, }); } @@ -125,15 +295,10 @@ async fn execute_refresh_column_with_source( false, ) .await?; - - let version = new_dataset.version().version; - table.dataset.update(new_dataset); - Ok(RefreshExecution { - result: RefreshColumnResult { - rows_filled, - version, - }, - source_version: read_version, + table.dataset.update(new_dataset.clone()); + Ok(Fill { + rows_filled, + committed: Some(Arc::new(new_dataset)), }) } @@ -158,10 +323,9 @@ pub(crate) async fn execute_refresh_column_async( Ok(crate::function::RefreshColumnResult { rows_assigned: execution.result.rows_filled, rows_failed: 0, - rows_remaining: 0, + rows_remaining: execution.rows_deferred, source_version: execution.source_version, - published_version: (execution.result.rows_filled > 0) - .then_some(execution.result.version), + published_version: execution.published_version, }) }))) } @@ -414,6 +578,181 @@ mod tests { table.add(batch).execute().await.unwrap(); } + /// The gate's reproducer: `b = coalesce(a, 0)` refreshed before `a` + /// must not bake zeros from `a`'s placeholder null. + #[tokio::test] + async fn test_dependent_refresh_cannot_fill_from_placeholder_null() { + let table = table_with("dependent_refresh_order", vec![1, 2, 3]).await; + table + .add_columns() + .computed("a", "x + 1") + .computed("b", "coalesce(a, 0)") + .execute() + .await + .unwrap(); + + let result = table.refresh_column("b").await.unwrap(); + assert_eq!(result.rows_filled, 3); + assert_eq!(read(&table, "a").await, vec![Some(2), Some(3), Some(4)]); + assert_eq!( + table.count_rows(Some("b = a".to_string())).await.unwrap(), + 3 + ); + assert_eq!(table.refresh_column("a").await.unwrap().rows_filled, 0); + + // Appended rows: the input is filled in the new fragment first too. + append(&table, vec![10]).await; + assert_eq!(table.refresh_column("b").await.unwrap().rows_filled, 1); + assert_eq!( + table.count_rows(Some("b = 0".to_string())).await.unwrap(), + 0 + ); + } + + /// The gate's reproducer: a fragment appended between the input fill and + /// the requested column's fill has inputs the pass never covered, so it + /// is left unfilled rather than read as null. + #[tokio::test] + async fn test_dependent_refresh_fences_an_append_after_input_fill() { + let table = table_with("dependent_refresh_append_gap", vec![1, 2, 3]).await; + table + .add_columns() + .computed("a_gap", "x + 1") + .computed("b_gap", "coalesce(a_gap, 0)") + .execute() + .await + .unwrap(); + let reached = Arc::new(tokio::sync::Notify::new()); + let resume = Arc::new(tokio::sync::Notify::new()); + super::PAUSES_BEFORE_FILL.lock().unwrap().push(( + "b_gap".to_string(), + reached.clone(), + resume.clone(), + )); + + let refresh = { + let table = table.clone(); + tokio::spawn(async move { + table + .refresh_column_async("b_gap") + .await + .unwrap() + .wait() + .await + }) + }; + reached.notified().await; + append(&table, vec![10]).await; + resume.notify_one(); + let result = refresh.await.unwrap().unwrap(); + assert_eq!(result.rows_assigned, 3); + // The target commit rebased over the append: the published version + // holds the deferred row, and the count is taken there. + assert_eq!(result.rows_remaining, 1); + + table.refresh_column("a_gap").await.unwrap(); + assert_eq!(table.count_rows(Some("b_gap = 0".into())).await.unwrap(), 0); + assert_eq!(table.refresh_column("b_gap").await.unwrap().rows_filled, 1); + } + + /// The gate's reproducer: an append landing between two input fills is + /// rebased into the second's snapshot, but the first never covered it. + #[tokio::test] + async fn test_dependent_refresh_fences_an_append_between_input_fills() { + let table = table_with("dependent_refresh_between_inputs", vec![1, 2, 3]).await; + table + .add_columns() + .computed("a_mid", "x + 1") + .computed("c_mid", "x + 2") + .computed("b_mid", "coalesce(a_mid, 0) + coalesce(c_mid, 0)") + .execute() + .await + .unwrap(); + let reached = Arc::new(tokio::sync::Notify::new()); + let resume = Arc::new(tokio::sync::Notify::new()); + super::PAUSES_BEFORE_FILL.lock().unwrap().push(( + "c_mid".to_string(), + reached.clone(), + resume.clone(), + )); + + let refresh = { + let table = table.clone(); + tokio::spawn(async move { + table + .refresh_column_async("b_mid") + .await + .unwrap() + .wait() + .await + }) + }; + reached.notified().await; + append(&table, vec![10]).await; + resume.notify_one(); + let result = refresh.await.unwrap().unwrap(); + assert_eq!(result.rows_assigned, 3); + // The appended row is in the reported source version but was deferred. + assert_eq!(result.rows_remaining, 1); + + table.refresh_column("a_mid").await.unwrap(); + table.refresh_column("c_mid").await.unwrap(); + assert_eq!(table.count_rows(Some("b_mid = 0".into())).await.unwrap(), 0); + assert_eq!(table.refresh_column("b_mid").await.unwrap().rows_filled, 1); + assert_eq!( + table.count_rows(Some("b_mid = 23".into())).await.unwrap(), + 1 + ); + } + + /// A dependency commit is a publication even when the requested column + /// itself fills nothing. + #[tokio::test] + async fn test_refresh_async_reports_a_dependency_publication() { + let table = table_with("refresh_async_dependency_publication", vec![1, 2, 3]).await; + table + .add_columns() + .computed("a", "x + 1") + .computed("b", "nullif(a, a)") + .execute() + .await + .unwrap(); + let result = table + .refresh_column_async("b") + .await + .unwrap() + .wait() + .await + .unwrap(); + assert_eq!(result.rows_assigned, 0); + assert_eq!(result.published_version, Some(result.source_version)); + } + + /// The gate's reproducer: a dense graph is walked once, not once per path. + #[tokio::test] + async fn test_dependency_refresh_visits_each_column_once() { + let table = table_with("dependency_refresh_deduplicates", vec![1]).await; + let mut declaration = table.add_columns(); + for i in 0..=6 { + let expression = if i == 0 { + "x + 1".to_string() + } else { + (0..i) + .map(|j| format!("a{j}")) + .collect::>() + .join(" + ") + }; + declaration = declaration.computed(format!("a{i}"), expression); + } + declaration.execute().await.unwrap(); + let order = super::dependency_order(table.schema().await.unwrap().as_ref(), "a6").unwrap(); + assert_eq!(order, (0..6).map(|i| format!("a{i}")).collect::>()); + // One commit per column, the requested one included. + let before = table.version().await.unwrap(); + table.refresh_column("a6").await.unwrap(); + assert_eq!(table.version().await.unwrap(), before + 7); + } + #[tokio::test] async fn test_refresh_fills_a_declared_column() { let table = table_with("refresh_fills", vec![1, 2, 3]).await;