Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 10 additions & 19 deletions kernel/src/engine/arrow_expression/evaluate_expression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1185,13 +1185,10 @@ mod tests {
// Create coalesce expression with column that has no nulls, followed by
// a reference to a non-existent column. If short-circuit works, the
// non-existent column is never evaluated and no error occurs.
let expr = Expression::variadic(
VariadicExpressionOp::Coalesce,
vec![
Expression::column(["a"]),
Expression::column(["nonexistent"]), // Would fail if evaluated
],
);
let expr = Expression::coalesce([
Expression::column(["a"]),
Expression::column(["nonexistent"]), // Would fail if evaluated
]);

// Should return column "a" directly (short-circuit skips evaluating "nonexistent")
let result = evaluate_expression(&expr, &batch, Some(&DataType::INTEGER)).unwrap();
Expand All @@ -1216,14 +1213,11 @@ mod tests {

// Create coalesce expression: a has nulls, b has none, c doesn't exist.
// Short-circuit should stop after evaluating b.
let expr = Expression::variadic(
VariadicExpressionOp::Coalesce,
vec![
Expression::column(["a"]),
Expression::column(["b"]),
Expression::column(["nonexistent"]), // Would fail if evaluated
],
);
let expr = Expression::coalesce([
Expression::column(["a"]),
Expression::column(["b"]),
Expression::column(["nonexistent"]), // Would fail if evaluated
]);

// Should coalesce a and b, never evaluate "nonexistent"
let result = evaluate_expression(&expr, &batch, Some(&DataType::INTEGER)).unwrap();
Expand All @@ -1242,10 +1236,7 @@ mod tests {
let a_values = Int32Array::from(vec![1, 2, 3]); // No nulls - would short-circuit
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a_values)]).unwrap();

let expr = Expression::variadic(
VariadicExpressionOp::Coalesce,
vec![Expression::column(["a"])],
);
let expr = Expression::coalesce([Expression::column(["a"])]);

// Request STRING type but array is INT32 - should fail even with short-circuit
let result = evaluate_expression(&expr, &batch, Some(&DataType::STRING));
Expand Down
28 changes: 15 additions & 13 deletions kernel/src/expressions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -701,6 +701,14 @@ impl Expression {
Self::Variadic(VariadicExpression::new(op, exprs))
}

/// Creates a new COALESCE expression that returns the first non-null value.
///
/// COALESCE evaluates expressions in order and returns the first non-null result.
/// If all expressions evaluate to null, the result is null.
pub fn coalesce(exprs: impl IntoIterator<Item = impl Into<Expression>>) -> Self {
Self::variadic(VariadicExpressionOp::Coalesce, exprs)
}

/// Creates a new opaque expression
pub fn opaque(
op: impl OpaqueExpressionOp,
Expand Down Expand Up @@ -1158,7 +1166,7 @@ mod tests {
use crate::expressions::scalars::{ArrayData, DecimalData, MapData, StructData};
use crate::expressions::{
column_expr, column_name, BinaryExpressionOp, BinaryPredicateOp, ColumnName,
Expression, Predicate, Scalar, Transform, UnaryExpressionOp, VariadicExpressionOp,
Expression, Predicate, Scalar, Transform, UnaryExpressionOp,
};
use crate::schema::{ArrayType, DataType, DecimalType, MapType, StructField};
use crate::utils::test_utils::assert_result_error_with_message;
Expand Down Expand Up @@ -1297,14 +1305,11 @@ mod tests {

#[test]
fn test_variadic_expression_roundtrip() {
let expr = Expression::variadic(
VariadicExpressionOp::Coalesce,
[
column_expr!("a"),
column_expr!("b"),
Expression::literal("default"),
],
);
let expr = Expression::coalesce([
column_expr!("a"),
column_expr!("b"),
Expression::literal("default"),
]);
assert_roundtrip(&expr);
}

Expand Down Expand Up @@ -1505,10 +1510,7 @@ mod tests {
column_expr!("c"),
column_expr!("d"),
);
let coalesce = Expression::variadic(
VariadicExpressionOp::Coalesce,
[add, mul, Expression::literal(0)],
);
let coalesce = Expression::coalesce([add, mul, Expression::literal(0)]);
let pred = Predicate::gt(coalesce, Expression::literal(100));
assert_roundtrip(&pred);

Expand Down
21 changes: 9 additions & 12 deletions kernel/src/scan/log_replay.rs
Original file line number Diff line number Diff line change
Expand Up @@ -560,7 +560,7 @@ mod tests {

use crate::actions::get_commit_schema;
use crate::engine::sync::SyncEngine;
use crate::expressions::{BinaryExpressionOp, Scalar, VariadicExpressionOp};
use crate::expressions::{BinaryExpressionOp, Scalar};
use crate::log_replay::ActionsBatch;
use crate::scan::state::ScanFile;
use crate::scan::state_info::tests::{
Expand Down Expand Up @@ -761,17 +761,14 @@ mod tests {
assert!(row_id_transform.is_replace);
assert_eq!(row_id_transform.exprs.len(), 1);
let expr = &row_id_transform.exprs[0];
let expeceted_expr = Arc::new(Expr::variadic(
VariadicExpressionOp::Coalesce,
vec![
Expr::column(["row_id_col"]),
Expr::binary(
BinaryExpressionOp::Plus,
Expr::literal(42i64),
Expr::column(["row_indexes_for_row_id_0"]),
),
],
));
let expeceted_expr = Arc::new(Expr::coalesce([
Expr::column(["row_id_col"]),
Expr::binary(
BinaryExpressionOp::Plus,
Expr::literal(42i64),
Expr::column(["row_indexes_for_row_id_0"]),
),
]));
assert_eq!(expr, &expeceted_expr);
} else {
panic!("Should have been a transform expression");
Expand Down
42 changes: 17 additions & 25 deletions kernel/src/transforms.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,7 @@ use std::sync::Arc;
use itertools::Itertools;
use serde::{Deserialize, Serialize};

use crate::expressions::{
BinaryExpressionOp, Expression, ExpressionRef, Scalar, Transform, VariadicExpressionOp,
};
use crate::expressions::{BinaryExpressionOp, Expression, ExpressionRef, Scalar, Transform};
use crate::schema::{DataType, SchemaRef, StructType};
use crate::table_features::ColumnMappingMode;
use crate::{DeltaResult, Error};
Expand Down Expand Up @@ -137,17 +135,14 @@ pub(crate) fn get_transform_expr(
let base_row_id = base_row_id.ok_or_else(|| {
Error::generic("Asked to generate RowIds, but no baseRowId found.")
})?;
let expr = Arc::new(Expression::variadic(
VariadicExpressionOp::Coalesce,
vec![
Expression::column([field_name]),
Expression::binary(
BinaryExpressionOp::Plus,
Expression::literal(base_row_id),
Expression::column([row_index_field_name]),
),
],
));
let expr = Arc::new(Expression::coalesce([
Expression::column([field_name]),
Expression::binary(
BinaryExpressionOp::Plus,
Expression::literal(base_row_id),
Expression::column([row_index_field_name]),
),
]));
transform.with_replaced_field(field_name.clone(), expr)
}
MetadataDerivedColumn {
Expand Down Expand Up @@ -592,17 +587,14 @@ mod tests {
.expect("Should have row_id_col transform");
assert!(row_id_transform.is_replace);

let expeceted_expr = Arc::new(Expression::variadic(
VariadicExpressionOp::Coalesce,
vec![
Expression::column(["row_id_col"]),
Expression::binary(
BinaryExpressionOp::Plus,
Expression::literal(4i64),
Expression::column(["row_index_col"]),
),
],
));
let expeceted_expr = Arc::new(Expression::coalesce([
Expression::column(["row_id_col"]),
Expression::binary(
BinaryExpressionOp::Plus,
Expression::literal(4i64),
Expression::column(["row_index_col"]),
),
]));
assert_eq!(row_id_transform.exprs.len(), 1);
let expr = &row_id_transform.exprs[0];
assert_eq!(expr, &expeceted_expr);
Expand Down
Loading