rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 14:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: plan computed columns against the session's DataFusion functions (#4924)

Browse Source 
A computed column's expression is planned with Lance's default planner,
so a function the embedder registers with DataFusion can't be declared,
and the declared field drops the return-field metadata a function
attaches, such as a logical type.

Every site that plans a declaration (declare, validate, refresh, schema
changes, view and table creation) now uses the DataFusion session the
table's lance session carries, when there is one. The declared field
keeps the expression's return-field metadata, without schema
declarations such as a primary key.
This commit is contained in:
Wyatt Alt
2026-11-09 21:55:58 -08:00
committed by GitHub
parent d8957238bf
commit 0a5c36c1e2
8 changed files with 291 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/database/namespace.rs
+3
View File
@@ -336,6 +336,9 @@ impl Database for LanceNamespaceDatabase {
// Refuse a bad declaration before the namespace records a table.
crate::table::computed_columns::ensure_declarations_are_planned(
&request.data.arrow_schema(),
self.session
.as_deref()
.and_then(lance::session::Session::datafusion_session),
)?;
let mut table_id = request.namespace_path.clone();
table_id.push(request.name.clone());
rust/lancedb/src/materialized_view.rs
+9 -4
View File
@@ -958,7 +958,7 @@ fn is_declaration(key: &str) -> bool {
|| crate::table::computed_columns::is_declaration_key(key)
}
fn without_declarations(field: &ArrowField) -> ArrowField {
pub(crate) fn without_declarations(field: &ArrowField) -> ArrowField {
let metadata: HashMap<String, String> = field
.metadata()
.iter()
@@ -1123,6 +1123,10 @@ impl PreparedDeclaration {
/// are validated over the assembled schema, and read only columns the
/// view holds (see [`PreparedDeclaration::input_column`]).
///
/// `session` holds the DataFusion functions a declaration may call, as
/// for [`AddColumnsBuilder::computed`](crate::table::AddColumnsBuilder::computed);
/// `None` plans against Lance's built-in functions only.
///
/// ```no_run
/// # #![recursion_limit = "256"]
/// # use std::collections::HashMap;
@@ -1148,7 +1152,7 @@ impl PreparedDeclaration {
/// (INPUTS_META_KEY.into(), format!("[\"{text}\"]")),
/// ]));
/// let view = prepared
/// .with_computed_columns(vec![(1, length)], &[])?
/// .with_computed_columns(vec![(1, length)], &[], None)?
/// .create("lengths")
/// .await?;
/// view.refresh().execute().await?; // rows land with `length` NULL
@@ -1160,6 +1164,7 @@ impl PreparedDeclaration {
mut self,
columns: Vec<(usize, ArrowField)>,
bindings: &[FunctionBinding],
session: Option<&datafusion::execution::context::SessionContext>,
) -> Result<Self> {
let invalid = |message: String| Error::InvalidInput { message };
if columns.is_empty() {
@@ -1223,7 +1228,7 @@ impl PreparedDeclaration {
}
rewrite_column_definitions(&mut metadata, self.schema.as_ref(), &fields)?;
let schema = ArrowSchema::new_with_metadata(fields, metadata);
ensure_declarations_are_planned(&schema)?;
ensure_declarations_are_planned(&schema, session)?;
self.schema = Arc::new(schema);
Ok(self)
}
@@ -2685,7 +2690,7 @@ mod tests {
(INPUTS_META_KEY.into(), "[\"x\"]".into()),
]));
let view = prepared
.with_computed_columns(vec![(1, doubled)], &[])
.with_computed_columns(vec![(1, doubled)], &[], None)
.unwrap()
.create("m")
.await
rust/lancedb/src/materialized_view/refresh.rs
+1 -1
View File
@@ -133,7 +133,7 @@ pub(crate) async fn execute_refresh(
})?;
let physical = Arc::new(ArrowSchema::from(view_ds.schema()));
ensure_declarations_are_planned(&physical)?;
ensure_declarations_are_planned(&physical, view_ds.session().datafusion_session())?;
if let Some(field) = physical
.fields()
.iter()
rust/lancedb/src/table.rs
+11 -2
View File
@@ -3407,7 +3407,11 @@ impl NativeTable {
namespace_client: Option<Arc<dyn LanceNamespace>>,
pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
) -> Result<Self> {
let batches = computed_columns::admit_create_source(batches)?;
let session = params.as_ref().and_then(|p| p.session.as_deref());
let batches = computed_columns::admit_create_source(
batches,
session.and_then(lance::session::Session::datafusion_session),
)?;
// Default params uses format v1.
let params = params.unwrap_or(WriteParams {
..Default::default()
@@ -3508,7 +3512,12 @@ impl NativeTable {
pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
session: Option<Arc<lance::session::Session>>,
) -> Result<Self> {
let batches = computed_columns::admit_create_source(batches)?;
let batches = computed_columns::admit_create_source(
batches,
session
.as_deref()
.and_then(lance::session::Session::datafusion_session),
)?;
// Build table_id from namespace + name for the storage options provider
let mut table_id = namespace.clone();
table_id.push(name.to_string());
rust/lancedb/src/table/add_columns.rs
+24
View File
@@ -89,6 +89,30 @@ impl AddColumnsBuilder {
/// # Ok(())
/// # }
/// ```
///
/// In LanceDB Embedded, an expression can also call functions registered
/// on a DataFusion `SessionContext` attached to the connection's Lance
/// session. They are not stored with the table, so attach them on every
/// connection that opens it.
///
/// ```no_run
/// # use std::sync::Arc;
/// # use datafusion::execution::context::SessionContext;
/// # use datafusion::logical_expr::ScalarUDF;
/// # async fn declare(my_udf: ScalarUDF) -> Result<(), Box<dyn std::error::Error>> {
/// let functions = SessionContext::new();
/// functions.register_udf(my_udf);
/// let session = lance::session::Session::default().with_datafusion_session(functions);
/// let db = lancedb::connect("data/db")
/// .session(Arc::new(session))
/// .execute()
/// .await?;
/// let table = db.open_table("t").execute().await?;
/// table.add_columns().computed("y", "my_udf(x)").execute().await?;
/// table.refresh_column("y").await?;
/// # Ok(())
/// # }
/// ```
pub fn computed(mut self, name: impl Into<String>, expression: impl Into<String>) -> Self {
self.computed.push((name.into(), expression.into()));
self
rust/lancedb/src/table/computed_columns.rs
+222 -38
View File
@@ -26,6 +26,7 @@ use std::collections::{BTreeSet, HashMap, HashSet};
use std::sync::Arc;
use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef};
use datafusion::execution::context::SessionContext;
use datafusion_common::{ScalarValue, tree_node::TreeNode};
use datafusion_expr::Expr;
use datafusion_physical_plan::PhysicalExpr;
@@ -40,6 +41,7 @@ use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::function::{FUNCTION_BLOB_V2_TYPE, FunctionApplication, FunctionBinding};
use crate::materialized_view::without_declarations;
use crate::utils::resolve_arrow_field_path;
use crate::{Error, Result};
@@ -1640,8 +1642,12 @@ pub(crate) fn plan_function_application(
///
/// Paths are compared at their root: a declaration reading `metadata` is
/// invalidated by a change to `metadata.age` just as surely.
pub(crate) fn ensure_not_an_input(schema: &SchemaRef, paths: &[&str]) -> Result<()> {
ensure_not_an_input_of(schema, paths, &[])
pub(crate) fn ensure_not_an_input(
schema: &SchemaRef,
paths: &[&str],
session: Option<&SessionContext>,
) -> Result<()> {
ensure_not_an_input_of(schema, paths, &[], session)
}
/// [`ensure_not_an_input`] where `removed` names columns whose own
@@ -1651,6 +1657,7 @@ pub(crate) fn ensure_not_an_input_of(
schema: &SchemaRef,
paths: &[&str],
removed: &[&str],
session: Option<&SessionContext>,
) -> Result<()> {
// By identity, not spelling: a quoted path names the same column, while
// a path *into* one removes no declaration at all.
@@ -1668,7 +1675,7 @@ pub(crate) fn ensure_not_an_input_of(
// The expression, not stored inputs, is the source of truth; an
// expression that no longer parses proves nothing, so refuse.
let inputs = match &declaration.kind {
ComputedColumnKind::Sql { expression } => Planner::new(schema.clone())
ComputedColumnKind::Sql { expression } => session_planner(schema.clone(), session)
.parse_expr(expression)
.map(|parsed| Planner::column_names_in_expr(&parsed))
.map_err(|e| Error::InvalidInput {
@@ -1760,7 +1767,10 @@ pub(crate) fn ensure_batch_writes_no_computed_values(
/// Function declaration satisfies the binding contract, and no SQL
/// declaration reads another computed column. What passes here is what
/// `refresh_column` can execute.
pub(crate) fn ensure_declarations_are_planned(schema: &ArrowSchema) -> Result<()> {
pub(crate) fn ensure_declarations_are_planned(
schema: &ArrowSchema,
session: Option<&SessionContext>,
) -> Result<()> {
let invalid = |message: String| Error::InvalidInput { message };
// A field with any declaration key is a declaration; a partial one is
// not "no declaration", it is a broken one.
@@ -1794,7 +1804,12 @@ pub(crate) fn ensure_declarations_are_planned(schema: &ArrowSchema) -> Result<()
.filter(|f| f.name() != &column.name)
.map(|f| f.as_ref().clone())
.collect();
let bound = bind(Arc::new(ArrowSchema::new(others)), &column.name, expression)?;
let bound = bind(
Arc::new(ArrowSchema::new(others)),
&column.name,
expression,
session,
)?;
if let Some(input) = bound.roots.iter().find(|r| declared.contains(*r)) {
return Err(invalid(format!(
"computed column '{}' reads computed column '{input}'",
@@ -1944,6 +1959,17 @@ pub(crate) struct BoundExpression {
pub blob_paths: Vec<String>,
/// A directly projected Blob v2 field whose semantics the output inherits.
projected_blob_field: Option<ArrowField>,
/// Field metadata of the expression's return field, e.g. a function's
/// logical type.
return_metadata: HashMap<String, String>,
}
/// Plans against `session`'s functions when given, else Lance's defaults.
fn session_planner(schema: SchemaRef, session: Option<&SessionContext>) -> Planner {
match session {
Some(session) => Planner::new_with_session(schema, session.clone()),
None => Planner::new(schema),
}
}
fn is_direct_field_projection(expr: &Expr) -> bool {
@@ -2089,7 +2115,12 @@ fn referenced_blob_paths(schema: &ArrowSchema, inputs: &[String]) -> Result<Vec<
/// simplifier can fold a referenced column out entirely (`true OR x > 0`),
/// and the guard protecting the stored SQL has to see every column the text
/// names, not just the ones the simplified form still reads.
pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<BoundExpression> {
pub(crate) fn bind(
schema: SchemaRef,
column: &str,
expression: &str,
session: Option<&SessionContext>,
) -> Result<BoundExpression> {
let invalid = |message: String| Error::InvalidExpression {
column: column.to_string(),
message,
@@ -2102,7 +2133,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
.into_iter()
.collect::<HashSet<_>>();
let parsing_schema = blob_runtime_schema(schema.as_ref(), &all_blob_paths);
let planner = Planner::new(parsing_schema);
let planner = session_planner(parsing_schema, session);
let parsed = planner
.parse_expr(expression)
.map_err(|e| invalid(e.to_string()))?;
@@ -2171,22 +2202,22 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
.map(|field| field.name().clone())
.collect();
let runtime_planner = Planner::new(runtime_schema);
let runtime_planner = session_planner(runtime_schema, session);
let optimized = runtime_planner
.optimize_expr(parsed)
.map_err(|e| invalid(e.to_string()))?;
let physical = Planner::new(read_schema.clone())
let physical = session_planner(read_schema.clone(), session)
.create_physical_expr(&optimized)
.map_err(|e| invalid(e.to_string()))?;
let data_type = physical
.data_type(read_schema.as_ref())
let return_field = physical
.return_field(read_schema.as_ref())
.map_err(|e| invalid(e.to_string()))?;
Ok(BoundExpression {
inputs,
roots,
physical,
data_type,
data_type: return_field.data_type().clone(),
blob_paths: blob_paths
.iter()
.map(|path| {
@@ -2195,6 +2226,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
})
.collect(),
projected_blob_field,
return_metadata: return_field.metadata().clone(),
})
}
@@ -2210,7 +2242,11 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
/// batch may declare `a` and then `b = a + 1` in one commit. Refresh order
/// then matters, and refresh enforces it: `b` is refused while `a` still has
/// unfilled rows.
fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<Vec<ArrowField>> {
fn plan_declarations(
schema: SchemaRef,
columns: &[(String, String)],
session: Option<&SessionContext>,
) -> Result<Vec<ArrowField>> {
if columns.is_empty() {
return Err(Error::InvalidInput {
message: "at least one computed column is required".into(),
@@ -2225,23 +2261,22 @@ fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<
return Err(Error::ColumnAlreadyExists { name: name.clone() });
}
let bound = bind(schema.clone(), name, expression)?;
let bound = bind(schema.clone(), name, expression, session)?;
// Declared columns start entirely null, so nullability is a property
// of the declaration rather than of what the expression yields.
let computed_metadata = computed_column_metadata(expression, &bound.inputs);
let field = match bound.projected_blob_field {
Some(source) => {
let mut metadata = source.metadata().clone();
metadata.retain(|key, _| !is_declaration_key(key));
metadata.extend(computed_metadata);
source
.with_name(name)
.with_nullable(true)
.with_metadata(metadata)
let projected = match bound.projected_blob_field {
Some(source) => source.with_name(name),
None => {
ArrowField::new(name, bound.data_type, true).with_metadata(bound.return_metadata)
}
None => ArrowField::new(name, bound.data_type, true).with_metadata(computed_metadata),
};
// Keep the source's descriptive metadata, not its schema declarations.
let projected = without_declarations(&projected);
let mut metadata = projected.metadata().clone();
metadata.extend(computed_metadata);
let field = projected.with_nullable(true).with_metadata(metadata);
schema = Arc::new(ArrowSchema::new_with_metadata(
schema
.fields()
@@ -2257,8 +2292,12 @@ fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<
Ok(fields)
}
pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Vec<ArrowField>> {
plan_declarations(schema, columns)
pub(crate) fn plan(
schema: SchemaRef,
columns: &[(String, String)],
session: Option<&SessionContext>,
) -> Result<Vec<ArrowField>> {
plan_declarations(schema, columns, session)
}
/// Run the schema-level checks of
@@ -2267,6 +2306,11 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
/// every declaration. For callers that stage declarations behind other work
/// and need those rejections before any of it lands.
///
/// Expressions may call the functions of `session`, the DataFusion session
/// the table's lance session carries (see
/// [`lance::session::Session::with_datafusion_session`]); `None` allows Lance's
/// defaults only.
///
/// Only the schema is consulted. Declaring also refuses a table with an RTE
/// write spec or retained SSTables; that is table state, checked at commit.
///
@@ -2280,16 +2324,20 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
/// ("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());
/// assert!(validate_declarations(schema.clone(), &declarations, None).is_ok());
/// assert!(validate_declarations(schema, &[("c".into(), "random()".into())], None).is_err());
/// ```
pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<()> {
pub fn validate_declarations(
schema: SchemaRef,
columns: &[(String, String)],
session: Option<&SessionContext>,
) -> Result<()> {
ensure_not_function_bound(
schema.as_ref(),
"schema evolution",
columns.iter().map(|(name, _)| name),
)?;
plan(schema, columns).map(drop)
plan(schema, columns, session).map(drop)
}
/// Build the transform that declares `columns` against `schema`.
@@ -2302,8 +2350,9 @@ pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) ->
pub(crate) fn declare(
schema: SchemaRef,
columns: &[(String, String)],
session: Option<&SessionContext>,
) -> Result<NewColumnTransform> {
let fields = plan_declarations(schema, columns)?;
let fields = plan_declarations(schema, columns, session)?;
Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(
fields,
))))
@@ -2334,9 +2383,10 @@ pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &st
/// created.
pub(crate) fn admit_create_source<S: lance_datafusion::utils::StreamingWriteSource>(
batches: S,
session: Option<&SessionContext>,
) -> Result<UnfilledDeclarations<S>> {
let schema = batches.arrow_schema();
ensure_declarations_are_planned(&schema)?;
ensure_declarations_are_planned(&schema, session)?;
let declared = computed_columns(&schema)
.into_iter()
.map(|c| c.name)
@@ -2496,7 +2546,7 @@ mod tests {
)]),
));
let declarations = vec![("a".to_string(), "x + 1".to_string())];
assert!(super::validate_declarations(schema, &declarations).is_err());
assert!(super::validate_declarations(schema, &declarations, None).is_err());
}
#[test]
@@ -2588,7 +2638,7 @@ mod tests {
}
}
use arrow_array::record_batch;
use arrow_array::{RecordBatch, record_batch};
use arrow_schema::{DataType, TimeUnit};
use futures::TryStreamExt;
use lance::dataset::ColumnAlteration;
@@ -2618,6 +2668,137 @@ mod tests {
computed_columns(table.schema().await.unwrap().as_ref())
}
/// Echoes an Int32 and tags its return field, like a function with a
/// logical type.
#[derive(Debug, PartialEq, Eq, Hash)]
struct Tagged(datafusion_expr::Signature);
impl datafusion_expr::ScalarUDFImpl for Tagged {
fn name(&self) -> &str {
"tagged"
}
fn signature(&self) -> &datafusion_expr::Signature {
&self.0
}
fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
Ok(DataType::Int32)
}
fn return_field_from_args(
&self,
_: datafusion_expr::ReturnFieldArgs,
) -> datafusion_common::Result<arrow_schema::FieldRef> {
Ok(Arc::new(
ArrowField::new("tagged", DataType::Int32, true)
.with_metadata(HashMap::from([("tag".into(), "yes".into())])),
))
}
fn invoke_with_args(
&self,
args: datafusion_expr::ScalarFunctionArgs,
) -> datafusion_common::Result<datafusion_expr::ColumnarValue> {
Ok(args.args[0].clone())
}
}
#[tokio::test]
#[cfg_attr(
not(feature = "experimental-functions"),
ignore = "requires experimental-functions"
)]
async fn test_declarations_plan_against_the_session_functions() {
let functions = SessionContext::new();
functions.register_udf(datafusion_expr::ScalarUDF::new_from_impl(Tagged(
datafusion_expr::Signature::exact(
vec![DataType::Int32],
datafusion_expr::Volatility::Immutable,
),
)));
let session =
Arc::new(lance::session::Session::default().with_datafusion_session(functions));
let conn = connect("memory://")
.session(session)
.execute()
.await
.unwrap();
let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [4, 5, 6])).unwrap();
let table = conn
.create_table("session_functions", batch)
.execute()
.await
.unwrap();
add_computed(&table, &[("t".into(), "tagged(x)".into())])
.await
.unwrap();
let schema = table.schema().await.unwrap();
assert_eq!(
schema.field_with_name("t").unwrap().metadata().get("tag"),
Some(&"yes".to_string())
);
table.refresh_column("t").await.unwrap();
// Dropping a column re-plans every declaration on the table.
table.drop_columns(&["y"]).await.unwrap();
let batches = table
.query()
.select(Select::columns(&["t"]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let expected = record_batch!(("t", Int32, [1, 2, 3])).unwrap();
assert_eq!(batches[0].column(0).as_ref(), expected.column(0).as_ref());
let plain = table_with_ints("no_session").await;
let err = add_computed(&plain, &[("t".into(), "tagged(x)".into())])
.await
.unwrap_err();
assert!(err.to_string().contains("tagged"), "{err}");
}
/// A projection keeps its source's descriptive metadata, not its schema
/// declarations: a copy of the primary key is not a primary key.
#[tokio::test]
#[cfg_attr(
not(feature = "experimental-functions"),
ignore = "requires experimental-functions"
)]
async fn test_projecting_the_primary_key_drops_its_declaration() {
let conn = connect("memory://").execute().await.unwrap();
let ints = record_batch!(("x", Int32, [1, 2, 3])).unwrap();
let schema = ArrowSchema::new(vec![
ArrowField::new("x", DataType::Int32, false)
.with_metadata(HashMap::from([("note".into(), "kept".into())])),
]);
let batch = RecordBatch::try_new(Arc::new(schema), ints.columns().to_vec()).unwrap();
let table = conn
.create_table("key_projection", batch)
.execute()
.await
.unwrap();
table.set_unenforced_primary_key(["x"]).await.unwrap();
add_computed(&table, &[("copy".into(), "x".into())])
.await
.unwrap();
let schema = table.schema().await.unwrap();
let copy = schema.field_with_name("copy").unwrap();
assert!(
!copy
.metadata()
.keys()
.any(|key| key.starts_with("lance-schema:")),
"{copy:?}"
);
assert_eq!(copy.metadata().get("note"), Some(&"kept".to_string()));
}
#[tokio::test]
#[cfg_attr(
not(feature = "experimental-functions"),
@@ -2658,6 +2839,7 @@ mod tests {
("first".to_string(), "image".to_string()),
("second".to_string(), "first".to_string()),
],
None,
)
.unwrap();
@@ -2680,6 +2862,7 @@ mod tests {
let fields = plan(
schema,
&[("payload".to_string(), "coalesce(image, image)".to_string())],
None,
)
.unwrap();
@@ -2914,7 +3097,8 @@ mod tests {
assert!(
validate_declarations(
table.schema().await.unwrap(),
&[("e".into(), "random()".into())]
&[("e".into(), "random()".into())],
None,
)
.is_err()
);
@@ -4803,7 +4987,7 @@ mod tests {
let schema = declare_plan(&schema, &downstream, "fb_downstream");
ensure_supported_function_metadata(&schema).unwrap();
ensure_declarations_are_planned(&schema).unwrap();
ensure_declarations_are_planned(&schema, None).unwrap();
let bindings = function_bindings(&schema).unwrap();
assert!(
computed_function_inputs(&schema, &bindings[0])
@@ -4868,7 +5052,7 @@ mod tests {
let base = ArrowSchema::new(vec![ArrowField::new("x", DataType::Int32, true)]);
let mut fields = base.fields().to_vec();
fields.extend(
plan(Arc::new(base), &[("doubled".into(), "x * 2".into())])
plan(Arc::new(base), &[("doubled".into(), "x * 2".into())], None)
.unwrap()
.into_iter()
.map(Arc::new),
@@ -4887,7 +5071,7 @@ mod tests {
.is_none_or(|metadata| !metadata.keys().any(|key| is_declaration_key(key)))
);
let schema = Arc::new(declare_plan(&schema, &downstream, "fb_downstream"));
ensure_declarations_are_planned(&schema).unwrap();
ensure_declarations_are_planned(&schema, None).unwrap();
let batch = RecordBatch::try_new(
schema,
vec![
rust/lancedb/src/table/refresh.rs
+7 -1
View File
@@ -308,6 +308,7 @@ async fn execute_refresh_column_with_source(
schema.clone(),
column,
&expression,
dataset.session().datafusion_session(),
)?);
ensure_inputs_filled(&dataset, &schema, column, &bound).await?;
let field = dataset
@@ -490,7 +491,12 @@ async fn ensure_inputs_filled(
),
});
};
let input_bound = super::computed_columns::bind(schema.clone(), input, expression)?;
let input_bound = super::computed_columns::bind(
schema.clone(),
input,
expression,
dataset.session().datafusion_session(),
)?;
let mut unfilled = 0u64;
for fragment in dataset.get_fragments() {
unfilled += count_fragment_gains(dataset, &fragment, &input_bound, input)
rust/lancedb/src/table/schema_evolution.rs
+14 -3
View File
@@ -166,7 +166,9 @@ pub(crate) async fn execute_declare(
.into(),
});
}
let transform = computed_columns::declare(table.schema().await?, columns)?;
let session = table.dataset.get().await?.session();
let transform =
computed_columns::declare(table.schema().await?, columns, session.datafusion_session())?;
let NewColumnTransform::AllNulls(schema) = transform else {
unreachable!("declarations are all-null")
};
@@ -307,7 +309,11 @@ pub(crate) async fn execute_alter_columns(
.filter(|alteration| alteration.rename.is_some() || alteration.data_type.is_some())
.map(|alteration| alteration.path.as_str())
.collect::<Vec<_>>();
computed_columns::ensure_not_an_input(&schema, &rebinding)?;
computed_columns::ensure_not_an_input(
&schema,
&rebinding,
dataset.session().datafusion_session(),
)?;
let retyped = alterations
.iter()
.filter(|alteration| alteration.data_type.is_some())
@@ -344,7 +350,12 @@ pub(crate) async fn execute_drop_columns(
let dropped = names.iter().map(String::as_str).collect::<Vec<_>>();
ensure_not_embedding_column(&schema, "drop", dropped.iter().copied())?;
computed_columns::ensure_not_bound_by(&unbinding.retained, "schema evolution", &dropped)?;
computed_columns::ensure_not_an_input_of(&schema, &dropped, &dropped)?;
computed_columns::ensure_not_an_input_of(
&schema,
&dropped,
&dropped,
dataset.session().datafusion_session(),
)?;
if !unbinding.is_noop() {
// The retirement commits before the projection, so the whole