diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 2b95cb34c..0042963d7 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,31 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result(), + schema.metadata().clone(), + )); + fields.push(field); } Ok(fields) } +/// Check `(name, expression)` pairs against `schema` exactly as +/// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) will +/// admit them, without committing. For callers that stage declarations +/// behind other work and need the rejection before any of it lands. +pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<()> { + 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 @@ -1582,6 +1602,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..9f9124b2e 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -7,6 +7,12 @@ //! 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, each by its own refresh and +//! commit, so the expression never reads an input's placeholder null as a +//! value. 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 +//! already 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 @@ -41,7 +47,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. @@ -74,7 +81,35 @@ async fn execute_refresh_column_with_source( 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 bound = Arc::new(super::computed_columns::bind( + schema.clone(), + column, + &expression, + )?); + + // Inputs that are themselves computed are filled first, so their + // placeholder nulls are never read as values. Declarations are acyclic by + // construction: a column can only read what existed when it was declared. + for input in &bound.roots { + let Some(declaration) = schema + .field_with_name(input) + .ok() + .and_then(computed_column_from_field) + else { + continue; + }; + if !matches!(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}'" + ), + }); + } + Box::pin(execute_refresh_column_with_source(table, input)).await?; + } + // Re-read: the input fills above committed on this handle. + let dataset = table.dataset.get().await?; let field = dataset .schema() .field(column) @@ -414,6 +449,37 @@ 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 + ); + } + #[tokio::test] async fn test_refresh_fills_a_declared_column() { let table = table_with("refresh_fills", vec![1, 2, 3]).await;