From 219f41339df819459c1da0ff69de2554debc2a48 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Wed, 5 Aug 2026 19:31:39 -0700 Subject: [PATCH] refactor(rust): tag a computed column's definition by kind ComputedColumn carried a bare expression string, which asserts that every computed column is a SQL expression. That holds for the only kind there is, but it is the wrong shape for the next one: a column defined by a registered function cannot be typed by parsing its definition, so its type and its inputs have a different provenance than a SQL column's. A single string has nowhere to say which it is. The definition is now ComputedColumnKind, non-exhaustive so another kind is additive, and the metadata carries a matching computed_column.kind tag beside the payload. Tagging the persisted form is the point -- the Rust type stays cheap to change and field metadata does not, and a second kind distinguished only by which keys happen to be present would leave every reader sniffing the shape. An unrecognized kind reads back as Unrecognized rather than as absent. A newer version's declaration is a computed column this one cannot evaluate, not a plain column: reported as absent it would be redeclarable over and would fail refresh as "not a computed column". Co-Authored-By: Claude Opus 5 (1M context) --- rust/lancedb/src/table.rs | 4 +- rust/lancedb/src/table/computed_columns.rs | 156 ++++++++++++++++++--- rust/lancedb/src/table/refresh.rs | 26 +++- 3 files changed, 166 insertions(+), 20 deletions(-) diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 1cc511035..bb85df8a8 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -93,7 +93,9 @@ pub use branch_merge::{ MergeBranchResult, MergeBranchStatus, MergePreview, RowCountSummary, }; pub use chrono::Duration; -pub use computed_columns::{ComputedColumn, computed_column_from_field, computed_columns}; +pub use computed_columns::{ + ComputedColumn, ComputedColumnKind, computed_column_from_field, computed_columns, +}; pub use delete::DeleteResult; use futures::future::join_all; pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTags}; diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index eb80d2d1c..2463cc18d 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -1,15 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! Expression-backed computed columns. +//! Computed columns. //! -//! A computed column is defined by a SQL expression rather than by values -//! supplied at write time. Declaring one commits the column carrying its -//! expression in field metadata but no data, so the cost does not scale with -//! the table; a later refresh fills the rows. +//! A computed column is defined by a rule rather than by values supplied at +//! write time. Declaring one commits the column carrying that rule in field +//! metadata but no data, so the cost does not scale with the table; a later +//! refresh fills the rows. //! -//! The expression is the whole definition: both the result type and the input -//! columns are derived from it, so a caller writes neither. +//! The rule is tagged by kind ([`ComputedColumnKind`]) because kinds differ in +//! where the column's type and inputs come from. A SQL expression is +//! self-describing -- both are derived from the expression, so a caller writes +//! neither -- while a kind resolved through a registry cannot be typed without +//! consulting it. Only SQL exists today; the tag is what lets another kind be +//! added without a second reading of the same key. //! //! [`computed_columns`] and [`computed_column_from_field`] read declarations //! back off a schema. @@ -26,27 +30,63 @@ use crate::{Error, Result}; /// Field metadata key marking a column as computed. The value is `"true"`. pub const COMPUTED_COLUMN_META_KEY: &str = "computed_column"; +/// Field metadata key naming the kind of rule that defines the column. +pub const KIND_META_KEY: &str = "computed_column.kind"; + /// Field metadata key holding the SQL expression that defines the column. pub const EXPRESSION_META_KEY: &str = "computed_column.expression"; /// Field metadata key holding the column's inputs, as a JSON array of names. pub const INPUTS_META_KEY: &str = "computed_column.inputs"; +/// Value of [`KIND_META_KEY`] for a column defined by a SQL expression. +pub const SQL_KIND: &str = "sql"; + +/// The rule that defines a computed column's values. +/// +/// Non-exhaustive: a kind added later is an additive change, and a caller that +/// only handles the kinds it knows keeps compiling. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum ComputedColumnKind { + /// A SQL expression evaluated by DataFusion. It is the whole definition: + /// the column's type and its inputs are both derived from it. + Sql { + /// The expression. + expression: String, + }, + /// A kind this version does not understand, written by a newer one. + /// + /// Reported rather than hidden so a caller can tell a column it cannot + /// refresh apart from one that was never computed. Nothing produces this. + Unrecognized { + /// The kind as it was found in the metadata. + kind: String, + }, +} + /// A computed column's declaration, as read back from field metadata. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ComputedColumn { /// Name of the computed column. pub name: String, - /// The SQL expression that defines it. - pub expression: String, - /// Columns the expression reads, parsed from it at declaration time. + /// The rule that defines it. + pub kind: ComputedColumnKind, + /// Columns the rule reads, recorded at declaration time. + /// + /// Outside the kind because every kind has inputs and the consumers that + /// use them -- refresh planning, dependency ordering -- do not care which + /// kind produced them. Where they come from does differ, and that is + /// settled at declaration: derived from a SQL expression, supplied by the + /// caller for a kind that cannot be parsed. pub inputs: Vec, } -/// Build the field metadata recording a binding. +/// Build the field metadata recording a SQL binding. fn computed_column_metadata(expression: &str, inputs: &[String]) -> HashMap { HashMap::from([ (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), + (KIND_META_KEY.to_string(), SQL_KIND.to_string()), (EXPRESSION_META_KEY.to_string(), expression.to_string()), ( INPUTS_META_KEY.to_string(), @@ -57,22 +97,32 @@ fn computed_column_metadata(expression: &str, inputs: &[String]) -> HashMap Option { let metadata = field.metadata(); if metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str) != Some("true") { return None; } - let expression = metadata.get(EXPRESSION_META_KEY)?; + let kind = match metadata.get(KIND_META_KEY)?.as_str() { + SQL_KIND => ComputedColumnKind::Sql { + expression: metadata.get(EXPRESSION_META_KEY)?.clone(), + }, + other => ComputedColumnKind::Unrecognized { + kind: other.to_string(), + }, + }; let inputs = metadata .get(INPUTS_META_KEY) .and_then(|raw| serde_json::from_str::>(raw).ok()) .unwrap_or_default(); Some(ComputedColumn { name: field.name().clone(), - expression: expression.clone(), + kind, inputs, }) } @@ -193,6 +243,28 @@ pub(crate) fn declare( )))) } +/// Commit a declaration of a kind this version does not produce, the way a +/// newer lancedb would leave one behind. Shared with the refresh tests, which +/// need the same column to check that refresh refuses it. +#[cfg(test)] +pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &str) { + use arrow_schema::DataType; + + let field = ArrowField::new(name, DataType::Int32, true).with_metadata(HashMap::from([ + (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), + (KIND_META_KEY.to_string(), kind.to_string()), + (INPUTS_META_KEY.to_string(), r#"["x"]"#.to_string()), + ])); + table + .add_columns() + .transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( + vec![field], + )))) + .execute() + .await + .unwrap(); +} + #[cfg(test)] mod tests { use arrow_array::record_batch; @@ -243,7 +315,9 @@ mod tests { declared(&table).await, vec![ComputedColumn { name: "doubled".into(), - expression: "x * 2".into(), + kind: ComputedColumnKind::Sql { + expression: "x * 2".into() + }, inputs: vec!["x".into()], }] ); @@ -264,6 +338,7 @@ mod tests { metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str), Some("true") ); + assert_eq!(metadata.get(KIND_META_KEY).map(String::as_str), Some("sql")); assert_eq!( metadata.get(EXPRESSION_META_KEY).map(String::as_str), Some("x * 2") @@ -473,6 +548,53 @@ mod tests { assert_eq!(declared[2].inputs, vec!["n".to_string()]); } + /// The reason the kind is tagged: a declaration written by a newer version + /// has to read back as a computed column this one cannot evaluate, not as + /// an ordinary column. Reported as absent it would be refreshable by + /// nothing and redeclarable over, silently. + #[tokio::test] + async fn test_unrecognized_kind_is_reported_rather_than_hidden() { + let table = table_with_ints("foreign_kind").await; + super::add_foreign_kind(&table, "embedding", "udf").await; + + assert_eq!( + declared(&table).await, + vec![ComputedColumn { + name: "embedding".into(), + kind: ComputedColumnKind::Unrecognized { kind: "udf".into() }, + inputs: vec!["x".into()], + }] + ); + + let err = add_computed(&table, &[("embedding".into(), "x * 2".into())]) + .await + .unwrap_err(); + assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "embedding")); + } + + /// A kind is what makes a declaration readable at all, so the flag alone + /// is half-formed in the same way a missing expression is. + #[test] + fn test_flag_without_a_kind_is_not_a_declaration() { + let field = + ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([( + COMPUTED_COLUMN_META_KEY.to_string(), + "true".to_string(), + )])); + assert_eq!(computed_column_from_field(&field), None); + } + + /// A SQL declaration is its expression; without one there is nothing to + /// refresh from. + #[test] + fn test_sql_kind_without_an_expression_is_not_a_declaration() { + let field = ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([ + (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), + (KIND_META_KEY.to_string(), SQL_KIND.to_string()), + ])); + assert_eq!(computed_column_from_field(&field), None); + } + #[tokio::test] async fn test_inputs_are_deduplicated_and_sorted() { let conn = connect("memory://").execute().await.unwrap(); diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index fcaeb0917..c13d7dfaa 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -8,7 +8,7 @@ use lance::dataset::UpdateBuilder as LanceUpdateBuilder; use serde::{Deserialize, Serialize}; use super::NativeTable; -use super::computed_columns::computed_column_from_field; +use super::computed_columns::{ComputedColumnKind, computed_column_from_field}; use crate::{Error, Result}; /// The result of refreshing a computed column. @@ -40,13 +40,24 @@ pub(crate) async fn execute_refresh_column( computed_column_from_field(field).ok_or_else(|| Error::NotAComputedColumn { name: column.to_string(), })?; + let expression = match &declaration.kind { + ComputedColumnKind::Sql { expression } => expression, + ComputedColumnKind::Unrecognized { kind } => { + return Err(Error::NotSupported { + message: format!( + "computed column '{column}' is defined by '{kind}', which this version of \ + lancedb cannot evaluate" + ), + }); + } + }; // Rows still holding no value are the ones to fill. A row whose expression // evaluates to null is indistinguishable from an unfilled one and is // recomputed, which costs work but cannot change the result. let builder = LanceUpdateBuilder::new(dataset) .update_where(&format!("{column} IS NULL"))? - .set(column, &declaration.expression)?; + .set(column, expression)?; let result = builder.build()?.execute().await?; let version = result.new_dataset.version().version; @@ -202,4 +213,15 @@ mod tests { let err = table.refresh_column("nope").await.unwrap_err(); assert!(matches!(err, Error::ColumnNotFound { name } if name == "nope")); } + + /// A declaration of a kind this version cannot evaluate is refused by + /// name, rather than mistaken for a plain column or fed to the SQL path. + #[tokio::test] + async fn test_refresh_rejects_a_kind_it_cannot_evaluate() { + let table = table_with("refresh_foreign", vec![1, 2, 3]).await; + super::super::computed_columns::add_foreign_kind(&table, "embedding", "udf").await; + + let err = table.refresh_column("embedding").await.unwrap_err(); + assert!(matches!(err, Error::NotSupported { message } if message.contains("udf"))); + } }