You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/12 19:23:15 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6716-23a9a7cfb9e0775b4f3876c19dd2e80b4242aa0d created (now daf33bd61b)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6716-23a9a7cfb9e0775b4f3876c19dd2e80b4242aa0d
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at daf33bd61b fix: skip an array or map function's later arguments where the array or map is NULL (#6716)
This branch includes the following new commits:
new daf33bd61b fix: skip an array or map function's later arguments where the array or map is NULL (#6716)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: skip an array or map function's later arguments where the array or map is NULL (#6716)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6716-23a9a7cfb9e0775b4f3876c19dd2e80b4242aa0d
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit daf33bd61b44d80e3be7f81f55367ed8c4d23645
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon Oct 12 19:22:55 2026 +0000
fix: skip an array or map function's later arguments where the array or map is NULL (#6716)
* fix: skip an array or map function's later arguments where the array or map is NULL
Spark's array_contains, array_position, array_remove, arrays_overlap,
array_union, slice, m[k] and a[i] are null-intolerant BinaryExpressions or
TernaryExpressions, which return NULL for a NULL array or map without
evaluating the other arguments. The native functions evaluated every argument
over the whole batch, so an argument that raises, such as an ANSI cast of a
malformed string, failed a query that Spark completes, and a nondeterministic
argument saw rows Spark never gives it.
Add a native NullShortCircuit wrapper that evaluates the first argument over
the whole batch and each later one only for the rows where no earlier argument
is NULL, as ArrayInsert already does for its own arguments. A
null_short_circuit flag on ScalarFunc and ListExtract enables it, set through
QueryPlanSerde.withNullShortCircuit. The planner leaves a function unwrapped
when every later argument can be evaluated for any row.
Closes #6613.
* Simplify NullShortCircuit and skip masking for infallible arguments
Track the skipped rows with NullBuffer::union and collect an argument's
columns with collect_columns. Keep each argument's batch field on the
argument itself instead of in a parallel list. Evaluate a later argument
that can be evaluated for any row over the whole batch, as literals
already were, and build a projection only for the arguments that are
masked. Leave a column whose index does not change unrewritten. Point
the guide and the proto comment at withNullShortCircuit instead of
repeating its conditions.
* Evaluate a shared subexpression for every row, as Spark's generated code does
Spark's subexpression elimination evaluates a subexpression that an
operator's expressions share for every row, before the expressions that
use it, in projections and aggregations. A shared ANSI cast therefore
raises in Spark even on rows whose array is NULL. Before serializing an
operator, ask Spark's EquivalentExpressions which subexpressions its
expressions share, and leave unmasked any call whose later argument
holds one, so it raises there as well.
* Evaluate a masked argument through one copy, so shuffle keeps its state
A masked argument was evaluated through its rebuilt copy, which reads its
columns from a batch of just those columns, only for a batch with a NULL
in an earlier argument, and through the original for every other batch.
Rebuilding an argument whose column indices change gives it new state, so
array_union(a, shuffle(b)) restarted shuffle's random sequence each time
evaluation moved between the two, and drew different permutations from
Spark's. The rebuilt copy now evaluates a masked argument for every batch.
A unit test evaluates array_union(a, shuffle(b)) over a batch without a
NULL array, one with, and one without again, against the unwrapped
function, which draws the same permutations as Spark when b is NULL
wherever a is. A Spark 4.0+ SQL fixture reads 8200 rows from one file,
so that only the scan's second batch has NULL arrays. Both fail with
the previous evaluation.
* fix: count LEGACY and TRY string-to-integer casts as infallible, and find shared subexpressions lazily
Cast::is_infallible now accepts a cast from Utf8 or LargeUtf8 to an integer
in LEGACY or TRY mode, which returns NULL for a string it cannot parse.
Comet's parsers for those modes never raise. With ANSI off, the native
planner then leaves array_contains(ids, CAST(s AS INT)) unwrapped, and CASE
keeps its eager path for such a cast.
withSharedSubexpressions now runs EquivalentExpressions on the operator's
first canShortCircuitNulls call instead of for every converted operator.
canShortCircuitNulls skips it entirely when every later argument is a leaf,
which subexpression elimination never shares.
array_null_short_circuit_non_ansi.sql runs the array shapes with ANSI off.
The casts return NULL there, while a slice start of 0 and a rand argument
still depend on the short-circuit. The shuffle fixture pins
spark.comet.batchSize, and the map fixture covers map_from_arrays, which its
CASE WHEN guards already protect.
* perf: keep CASE lazy for string-parsing casts, and give NullShortCircuit its own check
Counting a LEGACY or TRY string-to-integer cast as infallible also put CASE
and IF on their eager path, which parses every row for each branch that some
row chooses. With a selective predicate, that made IF about 8x slower and a
three-branch CASE about 3x slower.
Cast::is_infallible keeps its original set: casts that cannot fail and cost
little more than reading their input. Cast::cannot_fail adds the LEGACY and
TRY string-to-integer casts. case_when.rs exposes both checks. CASE uses
is_infallible, so it chooses the path it chose before. NullShortCircuit uses
cannot_fail, since it masks an argument only to keep it from failing on rows
that Spark skips.
The conditional bench gains IF and CASE shapes whose branches parse strings.
* fix: find shared subexpressions where shuffle keys and TakeOrderedAndProject are serialized
The native shuffle writer serializes its partitioning keys, and CometTakeOrderedAndProjectExec its
sort keys and projection, when the plan runs, outside the operator conversion that sets the shared
subexpressions. canShortCircuitNulls took the missing set as an empty one, so a cast that two
partitioning keys share was skipped on the rows whose array is NULL, where Spark's projection
evaluates it for every row and raises.
canShortCircuitNulls now skips no argument outside withSharedSubexpressions. The shuffle writer
calls it with the partitioning keys, which Spark evaluates together in one projection, and
CometTakeOrderedAndProjectExec with its own expressions.
* fix: scope a TakeOrderedAndProject's sort keys and projection apart, as Spark evaluates them
3090d859a gave CometTakeOrderedAndProjectExec's sort keys and projection one shared set. Spark's
ordering evaluates each sort key on its own, without subexpression elimination, and its projection
runs separately over the rows the sort keeps. So a cast in two sort keys, or in a sort key and the
projection, raised in Comet where Spark returns the rows. The top-K plans now use an empty set, and
the projection its projection list.
---
.../contributor-guide/adding_a_new_expression.md | 4 +
native/core/src/execution/jni_api.rs | 1 +
native/core/src/execution/planner.rs | 20 +-
native/proto/src/proto/expr.proto | 6 +
native/spark-expr/benches/conditional.rs | 60 +-
.../spark-expr/src/conditional_funcs/case_when.rs | 68 +-
native/spark-expr/src/conditional_funcs/mod.rs | 2 +
.../src/conditional_funcs/null_short_circuit.rs | 733 +++++++++++++++++++++
native/spark-expr/src/conversion_funcs/cast.rs | 19 +-
.../org/apache/comet/rules/CometExecRule.scala | 33 +-
.../org/apache/comet/serde/QueryPlanSerde.scala | 67 ++
.../main/scala/org/apache/comet/serde/arrays.scala | 24 +-
.../main/scala/org/apache/comet/serde/maps.scala | 4 +-
.../sql/comet/CometTakeOrderedAndProjectExec.scala | 21 +-
.../shuffle/CometNativeShuffleWriter.scala | 17 +-
.../expressions/array/array_null_short_circuit.sql | 125 ++++
.../array/array_null_short_circuit_non_ansi.sql | 66 ++
.../array/array_null_short_circuit_shuffle.sql | 38 ++
.../expressions/map/map_null_short_circuit.sql | 63 ++
.../apache/comet/serde/QueryPlanSerdeSuite.scala | 23 +-
20 files changed, 1329 insertions(+), 65 deletions(-)
diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md
index 20bcae06e6..3ce1b4e473 100644
--- a/docs/source/contributor-guide/adding_a_new_expression.md
+++ b/docs/source/contributor-guide/adding_a_new_expression.md
@@ -152,6 +152,10 @@ object CometLevenshtein extends CometExpressionSerde[Levenshtein] {
When the return type is set on the proto, the native planner skips the registry lookup entirely and routes straight to the Comet UDF registered in `create_comet_physical_fun_with_eval_mode`.
+#### Arguments that Spark skips after a NULL
+
+A null-intolerant `BinaryExpression` or `TernaryExpression` returns NULL as soon as one of its arguments is NULL, without evaluating the arguments after it. Native execution evaluates every argument of a `ScalarFunc` over the whole batch, so an argument that can fail, such as an ANSI cast of a malformed string, would fail on rows that Spark skips. Wrap the serialized call of such an expression in `withNullShortCircuit`, as `CometArrayContains` and `CometSlice` do, after checking the cond [...]
+
#### Registering the Expression Handler
Once you've created your `CometExpressionSerde` implementation, register it in `QueryPlanSerde.scala` by adding it to the appropriate expression map (e.g., `mathExpressions`, `stringExpressions`, `predicateExpressions`, etc.):
diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
index f8e9e497f6..262d8b8fd4 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -2784,6 +2784,7 @@ mod tests {
}],
return_type: None,
fail_on_error: false,
+ null_short_circuit: false,
})),
query_context: None,
expr_id: None,
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index ceca9e657a..15a5c48acb 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -146,8 +146,8 @@ use datafusion_comet_spark_expr::{
ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct,
DecimalRescaleCheckOverflow, FloatOperands, GetArrayStructFields, GetStructField, HllPlusPlus,
HllSketchAgg, HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero,
- NormalizeNestedFloats, Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev, SumDecimal,
- ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp,
+ NormalizeNestedFloats, NullShortCircuit, Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev,
+ SumDecimal, ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp,
};
use itertools::Itertools;
use jni::objects::{Global, JObject};
@@ -968,7 +968,7 @@ impl PhysicalPlanner {
.as_ref()
.map(|e| self.create_expr(e, Arc::clone(&input_schema)))
.transpose()?;
- Ok(Arc::new(ListExtract::new(
+ let list_extract: Arc<dyn PhysicalExpr> = Arc::new(ListExtract::new(
child,
ordinal,
default_value,
@@ -976,7 +976,12 @@ impl PhysicalPlanner {
expr.fail_on_error,
spark_expr.expr_id,
Arc::clone(&self.query_context_registry),
- )))
+ ));
+ if expr.null_short_circuit {
+ Ok(NullShortCircuit::wrap(list_extract, &input_schema)?)
+ } else {
+ Ok(list_extract)
+ }
}
ExprStruct::GetArrayStructFields(expr) => {
let child =
@@ -4055,6 +4060,11 @@ impl PhysicalPlanner {
Arc::new(Field::new(fun_name, data_type.clone(), true)),
Arc::new(ConfigOptions::default()),
));
+ let scalar_expr = if expr.null_short_circuit {
+ NullShortCircuit::wrap(scalar_expr, &input_schema)?
+ } else {
+ scalar_expr
+ };
// DF53 changed some UDFs (e.g. md5) to return StringViewArray at execution
// time (apache/datafusion#20045). Comet does not yet support view types, so
@@ -6866,6 +6876,7 @@ mod tests {
args: vec![array_col, array_col_1],
return_type: None,
fail_on_error: false,
+ null_short_circuit: false,
})),
query_context: None,
expr_id: None,
@@ -6992,6 +7003,7 @@ mod tests {
args: vec![array_col, array_col_1],
return_type: None,
fail_on_error: false,
+ null_short_circuit: false,
})),
query_context: None,
expr_id: None,
diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto
index b85f08f653..50d4716012 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -558,6 +558,10 @@ message ScalarFunc {
repeated Expr args = 2;
DataType return_type = 3;
bool fail_on_error = 4;
+ // Evaluate each argument after the first only for the rows where no earlier argument is NULL,
+ // as Spark does for a null-intolerant BinaryExpression or TernaryExpression. See
+ // QueryPlanSerde.withNullShortCircuit.
+ bool null_short_circuit = 5;
}
message CaseWhen {
@@ -631,6 +635,8 @@ message ListExtract {
Expr default_value = 3;
bool one_based = 4;
bool fail_on_error = 5;
+ // Evaluate the ordinal only for the rows where the child is not NULL, as in ScalarFunc
+ bool null_short_circuit = 6;
}
message GetArrayStructFields {
diff --git a/native/spark-expr/benches/conditional.rs b/native/spark-expr/benches/conditional.rs
index 2febcddee9..8ded476240 100644
--- a/native/spark-expr/benches/conditional.rs
+++ b/native/spark-expr/benches/conditional.rs
@@ -45,9 +45,10 @@ type Expr = Arc<dyn PhysicalExpr>;
type Shape<'a> = (&'a str, fn() -> Expr, &'a RecordBatch);
/// Columns shaped like the ones in `CometConditionalExpressionBenchmark`:
-/// `c1` a random long, `c2` an int in `0..100`, `c3` another random long, and `c4` / `c5` short
-/// and `c6` / `c7` long strings. `null_density` applies to every column. With `sorted`, `c1` and
-/// `c2` ascend through the batch, so every predicate over them selects one contiguous run of rows.
+/// `c1` a random long, `c2` an int in `0..100`, `c3` another random long, `c4` / `c5` short
+/// and `c6` / `c7` long strings, and `c8` / `c9` integers written as strings. `null_density`
+/// applies to every column. With `sorted`, `c1` and `c2` ascend through the batch, so every
+/// predicate over them selects one contiguous run of rows.
fn make_batch(null_density: f32, sorted: bool) -> RecordBatch {
let mut rng = StdRng::seed_from_u64(42);
let mut c1: Vec<i64> = (0..NUM_ROWS).map(|_| rng.random::<i64>()).collect();
@@ -57,7 +58,7 @@ fn make_batch(null_density: f32, sorted: bool) -> RecordBatch {
c2.sort_unstable();
}
let c3: Vec<i64> = (0..NUM_ROWS).map(|_| rng.random::<i64>()).collect();
- let valid: Vec<Vec<bool>> = (0..7)
+ let valid: Vec<Vec<bool>> = (0..9)
.map(|_| {
(0..NUM_ROWS)
.map(|_| rng.random::<f32>() >= null_density)
@@ -84,6 +85,11 @@ fn make_batch(null_density: f32, sorted: bool) -> RecordBatch {
.map(|i| valid[col][i].then(|| format!("{tag}{i}-").repeat(12)))
.collect()
};
+ let numeric = |col: usize, seed: usize| -> StringArray {
+ (0..NUM_ROWS)
+ .map(|i| valid[col][i].then(|| format!("{}", (i * 7919 + seed) % 2_000_000)))
+ .collect()
+ };
let columns: Vec<ArrayRef> = vec![
Arc::new(c1),
Arc::new(c2),
@@ -92,6 +98,8 @@ fn make_batch(null_density: f32, sorted: bool) -> RecordBatch {
Arc::new(short(4, "t")),
Arc::new(long(5, "long value ")),
Arc::new(long(6, "other value ")),
+ Arc::new(numeric(7, 17)),
+ Arc::new(numeric(8, 4242)),
];
RecordBatch::try_new(Arc::new(schema()), columns).unwrap()
}
@@ -105,6 +113,8 @@ fn schema() -> Schema {
Field::new("c5", DataType::Utf8, true),
Field::new("c6", DataType::Utf8, true),
Field::new("c7", DataType::Utf8, true),
+ Field::new("c8", DataType::Utf8, true),
+ Field::new("c9", DataType::Utf8, true),
])
}
@@ -309,13 +319,45 @@ fn case_fallible_branch() -> Expr {
)
}
+/// A string-to-integer cast with ANSI off, which returns NULL for a string it cannot parse.
+fn parse_int(column: &str) -> Expr {
+ spark_cast(col(column), DataType::Int32)
+}
+
+/// A selective predicate whose branch parses a string: 5% of the rows choose it.
+fn if_selective_parse() -> Expr {
+ if_expr(c2_below(5), parse_int("c8"), lit(ScalarValue::Int32(None)))
+}
+
+/// A predicate that most rows match, with a branch that parses a string.
+fn if_mostly_parse() -> Expr {
+ if_expr(c2_below(95), parse_int("c8"), lit(ScalarValue::Int32(None)))
+}
+
+/// Both branches parse a string, and half of the rows choose each.
+fn if_parse_either() -> Expr {
+ if_expr(c2_below(50), parse_int("c8"), parse_int("c9"))
+}
+
+/// Three selective branches, each parsing a string for 2% of the rows.
+fn case_selective_parse_3_branches() -> Expr {
+ case_when(
+ vec![
+ (c2_below(2), parse_int("c8")),
+ (c2_below(4), parse_int("c9")),
+ (c2_below(6), parse_int("c8")),
+ ],
+ None,
+ )
+}
+
fn criterion_benchmark(c: &mut Criterion) {
let no_nulls = make_batch(0.0, false);
let sparse_nulls = make_batch(0.1, false);
let dense_nulls = make_batch(0.9, false);
let sorted = make_batch(0.0, true);
- let shapes: [Shape; 26] = [
+ let shapes: [Shape; 30] = [
(
"case literal 3 branches",
case_literal_3_branches,
@@ -374,6 +416,14 @@ fn criterion_benchmark(c: &mut Criterion) {
case_column_10_branches,
&sorted,
),
+ ("if selective parse", if_selective_parse, &no_nulls),
+ ("if mostly parse", if_mostly_parse, &no_nulls),
+ ("if parse either", if_parse_either, &no_nulls),
+ (
+ "case selective parse 3 branches",
+ case_selective_parse_3_branches,
+ &no_nulls,
+ ),
];
let mut group = c.benchmark_group("conditional");
group.throughput(Throughput::Elements(NUM_ROWS as u64));
diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs
index 192e6046f6..c5478bb1d2 100644
--- a/native/spark-expr/src/conditional_funcs/case_when.rs
+++ b/native/spark-expr/src/conditional_funcs/case_when.rs
@@ -389,12 +389,27 @@ impl PhysicalExpr for CaseWhenExpr {
}
}
-/// Whether `expr` can be evaluated for rows that Spark would not evaluate it for.
+/// Whether `expr` can be evaluated for rows that Spark would not evaluate it for, at little cost.
///
-/// That holds when it can neither fail nor return something different for seeing more rows: a
-/// column, a literal, and comparisons, boolean logic, null checks, widening casts and wrapping
-/// arithmetic over them. Anything else, including every function, is assumed to be able to fail.
-fn is_infallible(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema) -> bool {
+/// That holds when it can neither fail nor return something different for seeing more rows, and is
+/// cheap: a column, a literal, and comparisons, boolean logic, null checks, widening casts and
+/// wrapping arithmetic over them. Anything else, including every function and a cast that parses a
+/// string, is treated as able to fail. `CASE` uses this to decide whether to evaluate every branch
+/// over the whole batch, which pays off only when that costs little more than selecting rows.
+pub(super) fn is_infallible(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema) -> bool {
+ infallible(expr, input_schema, false)
+}
+
+/// Whether `expr` can be evaluated for rows that Spark would not evaluate it for, at any cost: what
+/// [`is_infallible`] accepts, and casts that [`Cast::cannot_fail`] accepts beyond it, such as a
+/// LEGACY or TRY cast from a string to an integer, which returns NULL for a string it cannot parse.
+/// `NullShortCircuit` uses this, since it masks an argument only to keep it from failing on rows
+/// that Spark skips.
+pub(super) fn cannot_fail(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema) -> bool {
+ infallible(expr, input_schema, true)
+}
+
+fn infallible(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema, parsing: bool) -> bool {
if expr.is::<Column>() || expr.is::<Literal>() {
return true;
}
@@ -403,7 +418,11 @@ fn is_infallible(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema) -> bool {
} else if let Some(comparison) = expr.downcast_ref::<SparkComparison>() {
comparison.is_infallible(input_schema)
} else if let Some(cast) = expr.downcast_ref::<Cast>() {
- cast.is_infallible(input_schema)
+ if parsing {
+ cast.cannot_fail(input_schema)
+ } else {
+ cast.is_infallible(input_schema)
+ }
} else if let Some(normalize) = expr.downcast_ref::<NormalizeNaNAndZero>() {
normalize.is_infallible()
} else if let Some(not) = expr.downcast_ref::<NotExpr>() {
@@ -421,7 +440,7 @@ fn is_infallible(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema) -> bool {
&& expr
.children()
.into_iter()
- .all(|child| is_infallible(child, input_schema))
+ .all(|child| infallible(child, input_schema, parsing))
}
fn binary_is_infallible(binary: &BinaryExpr, input_schema: &Schema) -> bool {
@@ -1151,6 +1170,7 @@ mod tests {
Field::new("d", DataType::Float64, true),
Field::new("e", DataType::Float64, true),
Field::new("s", DataType::Utf8, true),
+ Field::new("ls", DataType::LargeUtf8, true),
Field::new("dec", DataType::Decimal128(10, 2), true),
Field::new("l", list_of(DataType::Float64), true),
Field::new("m", list_of(DataType::Float64), true),
@@ -1164,15 +1184,17 @@ mod tests {
),
]);
let c = |name: &str| col(name, &schema).unwrap();
- let cast = |e: Arc<dyn PhysicalExpr>, to: DataType| -> Arc<dyn PhysicalExpr> {
- Arc::new(Cast::new(
+ let cast_in = |mode: EvalMode, e: Arc<dyn PhysicalExpr>, to: DataType| {
+ let cast: Arc<dyn PhysicalExpr> = Arc::new(Cast::new(
e,
to,
- SparkCastOptions::new_without_timezone(EvalMode::Ansi),
+ SparkCastOptions::new_without_timezone(mode),
None,
None,
- ))
+ ));
+ cast
};
+ let cast = |e: Arc<dyn PhysicalExpr>, to: DataType| cast_in(EvalMode::Ansi, e, to);
let infallible = |e: Arc<dyn PhysicalExpr>| is_infallible(&e, &schema);
// A comparison as the planner builds it, which follows Spark's ordering for floats
let compare = |left: &str, op: Operator, right: Arc<dyn PhysicalExpr>| {
@@ -1229,9 +1251,31 @@ mod tests {
assert!(!infallible(checked));
// Decimal arithmetic can overflow
assert!(!infallible(binary(c("dec"), Operator::Plus, c("dec"))));
- // Narrowing and parsing casts can fail
+ // Narrowing and parsing casts can fail under ANSI
assert!(!infallible(cast(c("i"), DataType::Int32)));
assert!(!infallible(cast(c("s"), DataType::Int32)));
+ // A string that LEGACY or TRY mode cannot parse as an integer becomes NULL, so the cast
+ // cannot fail. Parsing costs too much for CASE to do it for rows it would skip, though.
+ for mode in [EvalMode::Legacy, EvalMode::Try] {
+ for to in [
+ DataType::Int8,
+ DataType::Int16,
+ DataType::Int32,
+ DataType::Int64,
+ ] {
+ for column in ["s", "ls"] {
+ let parse = cast_in(mode, c(column), to.clone());
+ assert!(!infallible(Arc::clone(&parse)), "{mode:?} {column} {to}");
+ assert!(cannot_fail(&parse, &schema), "{mode:?} {column} {to}");
+ }
+ }
+ // Other string casts are not claimed
+ assert!(
+ !cannot_fail(&cast_in(mode, c("s"), DataType::Float64), &schema),
+ "{mode:?}"
+ );
+ }
+ assert!(!cannot_fail(&cast(c("s"), DataType::Int32), &schema));
// So can anything built from a part that can fail
assert!(!infallible(binary(
binary(c("i"), Operator::Divide, c("i")),
diff --git a/native/spark-expr/src/conditional_funcs/mod.rs b/native/spark-expr/src/conditional_funcs/mod.rs
index dc13347b3a..36d854aecc 100644
--- a/native/spark-expr/src/conditional_funcs/mod.rs
+++ b/native/spark-expr/src/conditional_funcs/mod.rs
@@ -17,6 +17,8 @@
mod case_when;
mod if_expr;
+mod null_short_circuit;
pub use case_when::{create_case_when, create_if_expr, CaseWhenExpr};
pub use if_expr::IfExpr;
+pub use null_short_circuit::NullShortCircuit;
diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs
new file mode 100644
index 0000000000..232b9725cf
--- /dev/null
+++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs
@@ -0,0 +1,733 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::case_when::cannot_fail;
+use arrow::array::{new_null_array, Array, BooleanArray};
+use arrow::buffer::NullBuffer;
+use arrow::datatypes::{DataType, FieldRef, Schema};
+use arrow::record_batch::{RecordBatch, RecordBatchOptions};
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::common::{internal_datafusion_err, internal_err, Result};
+use datafusion::logical_expr::ColumnarValue;
+use datafusion::physical_expr::expressions::{Column, Literal};
+use datafusion::physical_expr::utils::collect_columns;
+use datafusion::physical_expr::PhysicalExpr;
+use std::fmt::{Display, Formatter};
+use std::hash::{Hash, Hasher};
+use std::sync::Arc;
+
+/// Evaluates a function's arguments the way Spark evaluates those of a null-intolerant
+/// `BinaryExpression` or `TernaryExpression`: the first for every row, and each later one only for
+/// the rows where no earlier argument is NULL.
+///
+/// Spark returns NULL as soon as an argument is NULL, without evaluating the ones after it.
+/// DataFusion evaluates every argument over the whole batch before it calls a function, so an
+/// argument that can fail, such as an ANSI cast of a malformed string, would fail on a row that
+/// Spark answers with NULL, and a nondeterministic argument would see rows that it does not see in
+/// Spark. An argument that can be evaluated for any row, such as a column, is evaluated for all of
+/// them.
+///
+/// The function still sees every row, with NULL for each argument skipped there, so it has to
+/// return NULL wherever an argument other than the last is NULL. Its arguments have to be its
+/// children, in the order Spark evaluates them, and it has to compute its result from them alone.
+/// `ArrayInsert` evaluates its own arguments the same way.
+#[derive(Debug)]
+pub struct NullShortCircuit {
+ /// The function over its arguments
+ inner: Arc<dyn PhysicalExpr>,
+ /// The function over the evaluated arguments, reading each one that is not a literal from a
+ /// column of the batch that holds them
+ body: Arc<dyn PhysicalExpr>,
+ args: Vec<Arg>,
+}
+
+impl NullShortCircuit {
+ /// Wraps `inner` if evaluating one of its arguments after the first for a row that Spark skips
+ /// could change the outcome, and returns it unchanged otherwise.
+ pub fn wrap(
+ inner: Arc<dyn PhysicalExpr>,
+ input_schema: &Schema,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let masked = inner
+ .children()
+ .iter()
+ .enumerate()
+ .map(|(i, arg)| i > 0 && !cannot_fail(arg, input_schema))
+ .collect::<Vec<_>>();
+ if !masked.contains(&true) {
+ return Ok(inner);
+ }
+ let plans = inner
+ .children()
+ .into_iter()
+ .zip(masked)
+ .map(|(arg, masked)| {
+ if arg.is::<Literal>() {
+ return Ok((None, masked));
+ }
+ // A masked argument is NULL on the rows it is skipped for
+ let field = arg.return_field(input_schema)?;
+ Ok((
+ Some(Arc::new(field.as_ref().clone().with_nullable(true))),
+ masked,
+ ))
+ })
+ .collect::<Result<Vec<_>>>()?;
+ Ok(Arc::new(Self::try_new(inner, plans)?))
+ }
+
+ /// `plans` holds each argument's [`Arg::field`] and whether it is masked.
+ fn try_new(inner: Arc<dyn PhysicalExpr>, plans: Vec<(Option<FieldRef>, bool)>) -> Result<Self> {
+ let children = inner.children();
+ if children.len() != plans.len() {
+ return internal_err!(
+ "NullShortCircuit has {} plans for {} arguments",
+ plans.len(),
+ children.len()
+ );
+ }
+ let args = children
+ .into_iter()
+ .zip(plans)
+ .map(|(expr, (field, masked))| Arg::try_new(expr, field, masked))
+ .collect::<Result<Vec<_>>>()?;
+ let mut body_children = Vec::with_capacity(args.len());
+ let mut column = 0;
+ for arg in &args {
+ if let Some(field) = &arg.field {
+ body_children.push(Arc::new(Column::new(field.name(), column)) as _);
+ column += 1;
+ } else {
+ body_children.push(Arc::clone(&arg.expr));
+ }
+ }
+ let body = Arc::clone(&inner).with_new_children(body_children)?;
+ Ok(Self { inner, body, args })
+ }
+}
+
+impl Hash for NullShortCircuit {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.inner.hash(state);
+ }
+}
+
+impl PartialEq for NullShortCircuit {
+ fn eq(&self, other: &Self) -> bool {
+ self.inner.eq(&other.inner)
+ }
+}
+
+impl Eq for NullShortCircuit {}
+
+impl Display for NullShortCircuit {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "NullShortCircuit({})", self.inner)
+ }
+}
+
+impl PhysicalExpr for NullShortCircuit {
+ fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+ self.inner.data_type(input_schema)
+ }
+
+ fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+ self.inner.nullable(input_schema)
+ }
+
+ fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> {
+ self.inner.return_field(input_schema)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ let num_rows = batch.num_rows();
+ // The rows where an argument evaluated so far is NULL, or `None` while there are none
+ let mut nulls: Option<NullBuffer> = None;
+ let mut fields = vec![];
+ let mut columns = vec![];
+ let last = self.args.len() - 1;
+ for (i, arg) in self.args.iter().enumerate() {
+ let value = arg.evaluate(batch, nulls.as_ref())?;
+ if i < last {
+ let value_nulls = match &value {
+ ColumnarValue::Array(array) if array.len() != num_rows => {
+ return internal_err!(
+ "{} returned {} rows for a batch of {num_rows}",
+ arg.expr,
+ array.len()
+ );
+ }
+ ColumnarValue::Array(array) => array.logical_nulls(),
+ ColumnarValue::Scalar(scalar) if scalar.is_null() => {
+ Some(NullBuffer::new_null(num_rows))
+ }
+ ColumnarValue::Scalar(_) => None,
+ };
+ nulls = NullBuffer::union(nulls.as_ref(), value_nulls.as_ref())
+ .filter(|n| n.null_count() > 0);
+ if nulls.as_ref().is_some_and(|n| n.null_count() == num_rows) {
+ // Every row is NULL, and no later argument is evaluated for any of them
+ let data_type = self.inner.data_type(batch.schema_ref())?;
+ return Ok(ColumnarValue::Array(new_null_array(&data_type, num_rows)));
+ }
+ }
+ if let Some(field) = &arg.field {
+ let column = value.into_array(num_rows)?;
+ // The batch is checked against the types the arguments have at run time, which
+ // can differ from their planned types in the names and nullability of nested fields
+ fields.push(if field.data_type() == column.data_type() {
+ Arc::clone(field)
+ } else {
+ Arc::new(
+ field
+ .as_ref()
+ .clone()
+ .with_data_type(column.data_type().clone()),
+ )
+ });
+ columns.push(column);
+ }
+ }
+ let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
+ let args =
+ RecordBatch::try_new_with_options(Arc::new(Schema::new(fields)), columns, &options)?;
+ self.body.evaluate(&args)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ self.inner.children()
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let inner = Arc::clone(&self.inner).with_new_children(children)?;
+ let plans = self
+ .args
+ .iter()
+ .map(|arg| (arg.field.clone(), arg.masked.is_some()))
+ .collect();
+ Ok(Arc::new(Self::try_new(inner, plans)?))
+ }
+
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ self.inner.fmt_sql(f)
+ }
+
+ fn is_volatile_node(&self) -> bool {
+ self.inner.is_volatile_node()
+ }
+}
+
+/// An argument, and how it is evaluated.
+#[derive(Debug)]
+struct Arg {
+ expr: Arc<dyn PhysicalExpr>,
+ /// The field of the column that `body` reads the argument from, or `None` for a literal,
+ /// which `body` keeps
+ field: Option<FieldRef>,
+ /// For an argument evaluated for only the rows where no earlier argument is NULL, the input
+ /// columns it reads, and the argument reading them from a batch of just those columns, so
+ /// that skipping rows filters no other column. It is evaluated in place of `expr` for every
+ /// batch, even one with no NULL, because rebuilding the argument gave it its own copy of any
+ /// state kept across batches, such as the random generator of `shuffle`
+ masked: Option<(Vec<usize>, Arc<dyn PhysicalExpr>)>,
+}
+
+impl Arg {
+ fn try_new(
+ expr: &Arc<dyn PhysicalExpr>,
+ field: Option<FieldRef>,
+ masked: bool,
+ ) -> Result<Self> {
+ let masked = masked.then(|| project(expr)).transpose()?;
+ Ok(Self {
+ expr: Arc::clone(expr),
+ field,
+ masked,
+ })
+ }
+
+ /// Evaluates the argument, for just the rows outside `nulls` if it is masked, returning a
+ /// value for every row of the batch.
+ fn evaluate(&self, batch: &RecordBatch, nulls: Option<&NullBuffer>) -> Result<ColumnarValue> {
+ let Some((projection, projected)) = &self.masked else {
+ return self.expr.evaluate(batch);
+ };
+ let batch = batch.project(projection)?;
+ match nulls {
+ Some(nulls) => {
+ let rows = BooleanArray::new(nulls.inner().clone(), None);
+ projected.evaluate_selection(&batch, &rows)
+ }
+ None => projected.evaluate(&batch),
+ }
+ }
+}
+
+/// The input columns that `expr` reads, and `expr` reading them from a batch of just those columns.
+fn project(expr: &Arc<dyn PhysicalExpr>) -> Result<(Vec<usize>, Arc<dyn PhysicalExpr>)> {
+ let mut projection = collect_columns(expr)
+ .iter()
+ .map(Column::index)
+ .collect::<Vec<_>>();
+ projection.sort_unstable();
+ projection.dedup();
+ let projected = Arc::clone(expr)
+ .transform_down(|e| {
+ let Some(column) = e.downcast_ref::<Column>() else {
+ return Ok(Transformed::no(e));
+ };
+ let index = projection
+ .binary_search(&column.index())
+ .map_err(|_| internal_datafusion_err!("{column} is not projected"))?;
+ if index == column.index() {
+ return Ok(Transformed::no(e));
+ }
+ Ok(Transformed::yes(
+ Arc::new(Column::new(column.name(), index)) as Arc<dyn PhysicalExpr>,
+ ))
+ })?
+ .data;
+ Ok((projection, projected))
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::{
+ create_query_context_map, Cast, EvalMode, ListExtract, ShuffleExpr, SparkArrayPositionFunc,
+ SparkArraySlice, SparkCastOptions,
+ };
+ use arrow::array::{ArrayRef, Int32Array, Int64Array, ListArray, StringArray};
+ use arrow::datatypes::{Field, Int32Type};
+ use datafusion::common::config::ConfigOptions;
+ use datafusion::common::ScalarValue;
+ use datafusion::functions_nested::set_ops::array_union_udf;
+ use datafusion::logical_expr::ScalarUDF;
+ use datafusion::physical_expr::expressions::{col, lit, CaseExpr, IsNotNullExpr};
+ use datafusion::physical_expr::ScalarFunctionExpr;
+ use std::sync::Mutex;
+
+ /// An array per row, a string that is not an integer on most rows whose array is NULL, and
+ /// the integer it is elsewhere.
+ fn test_batch() -> RecordBatch {
+ let arrays = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ None,
+ Some(vec![Some(1), Some(2)]),
+ Some(vec![Some(3)]),
+ None,
+ Some(vec![Some(5), None, Some(6)]),
+ None,
+ ]);
+ let ints = Int32Array::from(vec![7, 2, 1, 7, 6, 7]);
+ let strings = StringArray::from(vec!["bad", "2", "3", "worse", "3", "x"]);
+ let lengths = Int64Array::from(vec![Some(1), Some(1), None, Some(1), Some(2), Some(1)]);
+ let schema = Schema::new(vec![
+ Field::new("a", arrays.data_type().clone(), true),
+ Field::new("i", DataType::Int32, false),
+ Field::new("s", DataType::Utf8, false),
+ Field::new("len", DataType::Int64, true),
+ ]);
+ RecordBatch::try_new(
+ Arc::new(schema),
+ vec![
+ Arc::new(arrays),
+ Arc::new(ints),
+ Arc::new(strings),
+ Arc::new(lengths),
+ ],
+ )
+ .unwrap()
+ }
+
+ fn ansi_cast(expr: Arc<dyn PhysicalExpr>, to: DataType) -> Arc<dyn PhysicalExpr> {
+ Arc::new(Cast::new(
+ expr,
+ to,
+ SparkCastOptions::new_without_timezone(EvalMode::Ansi),
+ None,
+ None,
+ ))
+ }
+
+ fn function(
+ udf: ScalarUDF,
+ args: Vec<Arc<dyn PhysicalExpr>>,
+ return_type: DataType,
+ ) -> Arc<dyn PhysicalExpr> {
+ let name = udf.name().to_string();
+ Arc::new(ScalarFunctionExpr::new(
+ &name,
+ Arc::new(udf),
+ args,
+ Arc::new(Field::new(&name, return_type, true)),
+ Arc::new(ConfigOptions::default()),
+ ))
+ }
+
+ fn array_position(
+ array: Arc<dyn PhysicalExpr>,
+ element: Arc<dyn PhysicalExpr>,
+ ) -> Arc<dyn PhysicalExpr> {
+ let udf = ScalarUDF::new_from_impl(SparkArrayPositionFunc::new());
+ function(udf, vec![array, element], DataType::Int64)
+ }
+
+ fn slice(args: Vec<Arc<dyn PhysicalExpr>>, schema: &Schema) -> Arc<dyn PhysicalExpr> {
+ let udf = ScalarUDF::new_from_impl(SparkArraySlice::new());
+ function(udf, args, schema.field(0).data_type().clone())
+ }
+
+ fn evaluate(expr: &Arc<dyn PhysicalExpr>, batch: &RecordBatch) -> Result<ArrayRef> {
+ expr.evaluate(batch)?.into_array(batch.num_rows())
+ }
+
+ /// `array_position(a, CAST(s AS INT))` over `test_batch`
+ fn expected_positions() -> ArrayRef {
+ Arc::new(Int64Array::from(vec![
+ None,
+ Some(2),
+ Some(1),
+ None,
+ Some(0),
+ None,
+ ]))
+ }
+
+ /// `slice(a, CAST(s AS INT), len)` over `test_batch`
+ fn expected_slices() -> ArrayRef {
+ Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ None,
+ Some(vec![Some(2)]),
+ None,
+ None,
+ Some(vec![Some(6)]),
+ None,
+ ]))
+ }
+
+ #[test]
+ fn later_argument_is_not_evaluated_where_the_array_is_null() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ let array = Arc::new(RowRecorder::new(col("a", &schema).unwrap()));
+ let element = ansi_cast(col("s", &schema).unwrap(), DataType::Int32);
+ let unguarded = array_position(Arc::clone(&array) as _, element);
+ // The cast fails on "bad", whose array is NULL
+ let error = evaluate(&unguarded, &batch).unwrap_err().to_string();
+ assert!(error.contains("CAST_INVALID_INPUT"), "{error}");
+ array.seen.lock().unwrap().clear();
+
+ let guarded = NullShortCircuit::wrap(unguarded, &schema).unwrap();
+ assert!(guarded.is::<NullShortCircuit>());
+ let actual = evaluate(&guarded, &batch).unwrap();
+ assert_eq!(actual.as_ref(), expected_positions().as_ref());
+ // The array is evaluated once, over the whole batch
+ assert_eq!(*array.seen.lock().unwrap(), vec![(6, 4)]);
+ // An offset into the batch reaches the arrays' buffers and the filter
+ let actual = evaluate(&guarded, &batch.slice(1, 4)).unwrap();
+ assert_eq!(actual.as_ref(), expected_positions().slice(1, 4).as_ref());
+ }
+
+ #[test]
+ fn later_argument_still_fails_where_the_array_is_not_null() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ let element = ansi_cast(col("s", &schema).unwrap(), DataType::Int32);
+ let guarded =
+ NullShortCircuit::wrap(array_position(col("a", &schema).unwrap(), element), &schema)
+ .unwrap();
+ let arrays =
+ ListArray::from_iter_primitive::<Int32Type, _, _>(vec![None, Some(vec![Some(1)])]);
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(arrays),
+ Arc::new(Int32Array::from(vec![1, 1])),
+ Arc::new(StringArray::from(vec!["bad", "x"])),
+ Arc::new(Int64Array::from(vec![1, 1])),
+ ],
+ )
+ .unwrap();
+ // "bad" is skipped, but the array of "x" is not NULL
+ let error = evaluate(&guarded, &batch).unwrap_err().to_string();
+ assert!(error.contains("'x'"), "{error}");
+ }
+
+ #[test]
+ fn infallible_arguments_are_not_wrapped() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ let column = array_position(col("a", &schema).unwrap(), col("i", &schema).unwrap());
+ let wrapped = NullShortCircuit::wrap(Arc::clone(&column), &schema).unwrap();
+ assert!(Arc::ptr_eq(&column, &wrapped));
+ let widened = slice(
+ vec![
+ col("a", &schema).unwrap(),
+ ansi_cast(col("i", &schema).unwrap(), DataType::Int64),
+ lit(1i64),
+ ],
+ &schema,
+ );
+ let wrapped = NullShortCircuit::wrap(Arc::clone(&widened), &schema).unwrap();
+ assert!(Arc::ptr_eq(&widened, &wrapped));
+ // With ANSI off, a string that the cast cannot parse becomes NULL, so it cannot fail either
+ let parsed = array_position(
+ col("a", &schema).unwrap(),
+ Arc::new(Cast::new(
+ col("s", &schema).unwrap(),
+ DataType::Int32,
+ SparkCastOptions::new_without_timezone(EvalMode::Legacy),
+ None,
+ None,
+ )),
+ );
+ let wrapped = NullShortCircuit::wrap(Arc::clone(&parsed), &schema).unwrap();
+ assert!(Arc::ptr_eq(&parsed, &wrapped));
+ }
+
+ #[test]
+ fn each_argument_waits_for_every_earlier_one() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ // The length is only evaluated where neither the array nor the start is NULL
+ let start = ansi_cast(
+ ansi_cast(col("s", &schema).unwrap(), DataType::Int32),
+ DataType::Int64,
+ );
+ let length = Arc::new(RowRecorder::new(col("len", &schema).unwrap()));
+ let args = vec![col("a", &schema).unwrap(), start, Arc::clone(&length) as _];
+ let guarded = NullShortCircuit::wrap(slice(args, &schema), &schema).unwrap();
+ let actual = evaluate(&guarded, &batch).unwrap();
+ assert_eq!(actual.as_ref(), expected_slices().as_ref());
+ // The length saw the three rows that have an array, and only the column it reads
+ assert_eq!(*length.seen.lock().unwrap(), vec![(3, 1)]);
+ }
+
+ #[test]
+ fn infallible_later_argument_is_evaluated_for_every_row() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ let start = ansi_cast(
+ ansi_cast(col("s", &schema).unwrap(), DataType::Int32),
+ DataType::Int64,
+ );
+ let args = vec![
+ col("a", &schema).unwrap(),
+ start,
+ col("len", &schema).unwrap(),
+ ];
+ let guarded = NullShortCircuit::wrap(slice(args, &schema), &schema).unwrap();
+ // Only the start, which can fail, skips the rows whose array is NULL
+ let masked = guarded
+ .downcast_ref::<NullShortCircuit>()
+ .unwrap()
+ .args
+ .iter()
+ .map(|arg| arg.masked.is_some())
+ .collect::<Vec<_>>();
+ assert_eq!(masked, vec![false, true, false]);
+ let actual = evaluate(&guarded, &batch).unwrap();
+ assert_eq!(actual.as_ref(), expected_slices().as_ref());
+ }
+
+ #[test]
+ fn null_literal_array_evaluates_no_later_argument() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ let null_array = lit(ScalarValue::try_new_null(schema.field(0).data_type()).unwrap());
+ let element = Arc::new(RowRecorder::new(ansi_cast(
+ col("s", &schema).unwrap(),
+ DataType::Int32,
+ )));
+ let guarded = NullShortCircuit::wrap(
+ array_position(null_array, Arc::clone(&element) as _),
+ &schema,
+ )
+ .unwrap();
+ let actual = evaluate(&guarded, &batch).unwrap();
+ assert_eq!(actual.data_type(), &DataType::Int64);
+ assert_eq!(actual.null_count(), batch.num_rows());
+ assert!(element.seen.lock().unwrap().is_empty());
+ }
+
+ #[test]
+ fn list_extract_is_guarded() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ // element_at(a, CAST(s AS INT)), with a NULL for an index past the end
+ let extract = Arc::new(ListExtract::new(
+ col("a", &schema).unwrap(),
+ ansi_cast(col("s", &schema).unwrap(), DataType::Int32),
+ None,
+ true,
+ false,
+ None,
+ create_query_context_map(),
+ ));
+ let guarded = NullShortCircuit::wrap(extract, &schema).unwrap();
+ let actual = evaluate(&guarded, &batch).unwrap();
+ let expected: ArrayRef = Arc::new(Int32Array::from(vec![
+ None,
+ Some(2),
+ None,
+ None,
+ Some(6),
+ None,
+ ]));
+ assert_eq!(actual.as_ref(), expected.as_ref());
+ }
+
+ #[test]
+ fn column_indices_can_be_rewritten() {
+ let batch = test_batch();
+ let schema = batch.schema();
+ let element = ansi_cast(col("s", &schema).unwrap(), DataType::Int32);
+ let guarded =
+ NullShortCircuit::wrap(array_position(col("a", &schema).unwrap(), element), &schema)
+ .unwrap();
+ // DataFusion's CASE projects the batch down to the columns that it reads, and rewrites
+ // their indices in its branches through `with_new_children`
+ let when = Arc::new(IsNotNullExpr::new(col("len", &schema).unwrap()));
+ let case: Arc<dyn PhysicalExpr> =
+ Arc::new(CaseExpr::try_new(None, vec![(when, guarded)], None).unwrap());
+ let actual = evaluate(&case, &batch).unwrap();
+ let expected: ArrayRef = Arc::new(Int64Array::from(vec![
+ None,
+ Some(2),
+ None,
+ None,
+ Some(0),
+ None,
+ ]));
+ assert_eq!(actual.as_ref(), expected.as_ref());
+ }
+
+ #[test]
+ fn stateful_argument_keeps_its_state_across_batches() {
+ let list = DataType::new_list(DataType::Int32, true);
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", list.clone(), true),
+ Field::new("b", list.clone(), true),
+ ]));
+ // array_union(a, shuffle(b)). The masked shuffle reads `b` from the first column of a
+ // batch of just that column, so it is rebuilt, with a generator of its own.
+ let union = || {
+ let shuffle = Arc::new(ShuffleExpr::new(col("b", &schema).unwrap(), 42));
+ let udf = array_union_udf().as_ref().clone();
+ function(udf, vec![col("a", &schema).unwrap(), shuffle], list.clone())
+ };
+ // An empty array and [1, 2, 3, 4] on each row that is not NULL. A shuffle draws nothing
+ // for a NULL row, so the unwrapped function, which evaluates it for the rows whose array
+ // is NULL as well, draws the same permutations as Spark.
+ let batch = |valid: &[bool]| {
+ let four = || (1..=4).map(Some).collect::<Vec<_>>();
+ let a = valid.iter().map(|&v| v.then(Vec::<Option<i32>>::new));
+ let b = valid.iter().map(|&v| v.then(four));
+ let columns: Vec<ArrayRef> = vec![
+ Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(a)),
+ Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(b)),
+ ];
+ RecordBatch::try_new(Arc::clone(&schema), columns).unwrap()
+ };
+ let unguarded = union();
+ let guarded = NullShortCircuit::wrap(union(), &schema).unwrap();
+ assert!(guarded.is::<NullShortCircuit>());
+ // The permutations carry on from one batch to the next, whether or not it has a NULL
+ for valid in [vec![true, true], vec![false, true, true], vec![true, true]] {
+ let batch = batch(&valid);
+ let expected = evaluate(&unguarded, &batch).unwrap();
+ let actual = evaluate(&guarded, &batch).unwrap();
+ assert_eq!(actual.as_ref(), expected.as_ref());
+ }
+ }
+
+ /// Evaluates `child`, recording the rows and columns of each batch that it is evaluated for.
+ #[derive(Debug)]
+ struct RowRecorder {
+ child: Arc<dyn PhysicalExpr>,
+ /// Shared with the copies that `with_new_children` makes
+ seen: Arc<Mutex<Vec<(usize, usize)>>>,
+ }
+
+ impl RowRecorder {
+ fn new(child: Arc<dyn PhysicalExpr>) -> Self {
+ Self {
+ child,
+ seen: Arc::new(Mutex::new(vec![])),
+ }
+ }
+ }
+
+ impl Hash for RowRecorder {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.child.hash(state);
+ }
+ }
+
+ impl PartialEq for RowRecorder {
+ fn eq(&self, other: &Self) -> bool {
+ self.child.eq(&other.child)
+ }
+ }
+
+ impl Eq for RowRecorder {}
+
+ impl Display for RowRecorder {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "RowRecorder({})", self.child)
+ }
+ }
+
+ impl PhysicalExpr for RowRecorder {
+ fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+ self.child.data_type(input_schema)
+ }
+
+ fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+ self.child.nullable(input_schema)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ self.seen
+ .lock()
+ .unwrap()
+ .push((batch.num_rows(), batch.num_columns()));
+ self.child.evaluate(batch)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ Ok(Arc::new(Self {
+ child: Arc::clone(&children[0]),
+ seen: Arc::clone(&self.seen),
+ }))
+ }
+
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ self.child.fmt_sql(f)
+ }
+ }
+}
diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs
index 30b3fe997a..884573d01e 100644
--- a/native/spark-expr/src/conversion_funcs/cast.rs
+++ b/native/spark-expr/src/conversion_funcs/cast.rs
@@ -119,8 +119,9 @@ impl Cast {
}
}
- /// Whether the cast cannot fail for any value, in any eval mode: a cast to the child's own
- /// type, from NULL, or to a wider number. `cast_array` rejects a dictionary target first.
+ /// Whether the cast cannot fail for any value, in any eval mode, and costs little more than
+ /// reading its input: a cast to the child's own type, from NULL, or to a wider number.
+ /// `cast_array` rejects a dictionary target first.
pub(crate) fn is_infallible(&self, input_schema: &Schema) -> bool {
use DataType::*;
let Ok(from) = self.child.data_type(input_schema) else {
@@ -138,6 +139,20 @@ impl Cast {
| (Float32, Float64)
))
}
+
+ /// Whether the cast cannot fail for any value: what [`Self::is_infallible`] accepts, and a cast
+ /// from a string to an integer in LEGACY or TRY mode, which returns NULL for a string it cannot
+ /// parse. Parsing costs far more per row than those casts.
+ pub(crate) fn cannot_fail(&self, input_schema: &Schema) -> bool {
+ use DataType::*;
+ self.is_infallible(input_schema)
+ || (self
+ .child
+ .data_type(input_schema)
+ .is_ok_and(|from| matches!(from, Utf8 | LargeUtf8))
+ && matches!(self.data_type, Int8 | Int16 | Int32 | Int64)
+ && self.cast_options.eval_mode != EvalMode::Ansi)
+ }
}
/// Spark cast options
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 32faea60e6..5a6e761d44 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -1134,24 +1134,21 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
}
val builder = OperatorOuterClass.Operator.newBuilder().setPlanId(op.id)
- if (op.children.nonEmpty && op.children.forall(_.isInstanceOf[CometNativeExec])) {
- val childOp = op.children.map(_.asInstanceOf[CometNativeExec].nativeOp)
- childOp.foreach(builder.addChildren)
- return serde
- .convert(op, builder, childOp: _*)
- .map { nativeOp =>
- val exec = serde.createExec(nativeOp, op)
- rollUpInfoMessages(op, exec)
- exec
- }
- } else {
- return serde
- .convert(op, builder)
- .map { nativeOp =>
- val exec = serde.createExec(nativeOp, op)
- rollUpInfoMessages(op, exec)
- exec
- }
+ // Which subexpressions Spark evaluates for every row decides where native execution may
+ // skip a function's arguments (see `QueryPlanSerde.canShortCircuitNulls`)
+ val nativeOp = QueryPlanSerde.withSharedSubexpressions(op.expressions) {
+ if (op.children.nonEmpty && op.children.forall(_.isInstanceOf[CometNativeExec])) {
+ val childOp = op.children.map(_.asInstanceOf[CometNativeExec].nativeOp)
+ childOp.foreach(builder.addChildren)
+ serde.convert(op, builder, childOp: _*)
+ } else {
+ serde.convert(op, builder)
+ }
+ }
+ return nativeOp.map { nativeOp =>
+ val exec = serde.createExec(nativeOp, op)
+ rollUpInfoMessages(op, exec)
+ exec
}
}
None
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 43960ad560..56cb421ae8 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -24,6 +24,7 @@ import java.util.concurrent.atomic.AtomicLong
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
+import scala.util.DynamicVariable
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.expressions._
@@ -1297,6 +1298,72 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
Some(ExprOuterClass.Expr.newBuilder().setScalarFunc(builder).build())
}
+ /**
+ * The subexpressions that occur more than once in the expressions of the operator being
+ * serialized, as Spark's subexpression elimination finds them. Spark's generated code for a
+ * projection or an aggregation evaluates each of them for every row, before the expressions
+ * that use them. They are treated that way for every operator, and whatever
+ * `spark.sql.subexpressionElimination.enabled` says, which only keeps them evaluated for every
+ * row where Spark might not. They are found on the first `canShortCircuitNulls` call for the
+ * operator, so an operator without an argument to skip never runs `EquivalentExpressions`.
+ * `None` outside `withSharedSubexpressions`, where they are unknown.
+ */
+ private val sharedSubexpressions = new DynamicVariable[Option[SharedSubexpressions]](None)
+
+ private final class SharedSubexpressions(expressions: Seq[Expression]) {
+ lazy val common: Seq[Expression] = {
+ val equivalence = new EquivalentExpressions
+ expressions.foreach(equivalence.addExprTree(_))
+ equivalence.getCommonSubexpressions
+ }
+ }
+
+ /**
+ * Runs `f`, which serializes an operator whose expressions are `expressions`. Code that
+ * serializes expressions outside an operator's conversion, such as when the plan runs, calls it
+ * with the expressions Spark evaluates together, or no argument is skipped.
+ */
+ def withSharedSubexpressions[T](expressions: Seq[Expression])(f: => T): T =
+ sharedSubexpressions.withValue(Some(new SharedSubexpressions(expressions)))(f)
+
+ /**
+ * Whether native execution may skip the arguments of `expr` after the first on the rows where
+ * an earlier argument is NULL, as Spark's evaluation of a null-intolerant `BinaryExpression` or
+ * `TernaryExpression` does. Not if one of them holds a subexpression that the operator's
+ * expressions share: Spark evaluates that for every row, and raises if it fails on a row whose
+ * array is NULL. Nor outside `withSharedSubexpressions`, where what they share is unknown.
+ */
+ def canShortCircuitNulls(expr: Expression): Boolean = {
+ val later = expr.children.tail
+ // Subexpression elimination never shares a leaf, such as a column or a literal
+ later.forall(_.isInstanceOf[LeafExpression]) || sharedSubexpressions.value.exists { s =>
+ !later.exists(_.exists(e => s.common.exists(_.semanticEquals(e))))
+ }
+ }
+
+ /**
+ * Has native execution evaluate the arguments of `scalarFunc`, the serialized `expr`, the way
+ * Spark evaluates those of a null-intolerant `BinaryExpression` or `TernaryExpression`: each
+ * argument after the first only for the rows where no earlier argument is NULL, unless
+ * `canShortCircuitNulls` rules it out. Otherwise every argument is evaluated over the whole
+ * batch, so one that can fail, such as an ANSI cast of a malformed string, fails on a row that
+ * Spark returns NULL for without evaluating it (#6613). The arguments must be in the order
+ * Spark evaluates them, and the function must return NULL wherever an argument other than the
+ * last is NULL. The native planner drops the guard when every later argument can be evaluated
+ * for any row, such as a column or a literal.
+ */
+ def withNullShortCircuit(expr: Expression, scalarFunc: Option[Expr]): Option[Expr] =
+ scalarFunc.map { proto =>
+ require(proto.hasScalarFunc, s"not a scalar function: $proto")
+ if (canShortCircuitNulls(expr)) {
+ proto.toBuilder
+ .setScalarFunc(proto.getScalarFunc.toBuilder.setNullShortCircuit(true))
+ .build()
+ } else {
+ proto
+ }
+ }
+
/**
* If `handler` is a `CodegenDispatchFallback`, run `expr` through the JVM codegen dispatcher
* and return `Some((handler, proto))` on success; otherwise return `None`. Shared by the
diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
index a16d42c94d..aa17b6ef53 100644
--- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
@@ -56,7 +56,7 @@ object CometArrayRemove
} else {
"array_remove_all"
}
- scalarFunctionExprToProto(function, arrayExprProto, keyExprProto)
+ withNullShortCircuit(expr, scalarFunctionExprToProto(function, arrayExprProto, keyExprProto))
}
}
@@ -184,7 +184,7 @@ object CometArrayContains
} else {
"array_contains"
}
- scalarFunctionExprToProto(function, arrayExprProto, keyExprProto)
+ withNullShortCircuit(expr, scalarFunctionExprToProto(function, arrayExprProto, keyExprProto))
}
}
@@ -376,7 +376,7 @@ object CometArraysOverlap
false,
leftArrayExprProto,
rightArrayExprProto)
- arraysOverlapScalarExpr
+ withNullShortCircuit(expr, arraysOverlapScalarExpr)
}
}
@@ -583,11 +583,13 @@ object CometSlice extends CometExpressionSerde[Slice] {
// output, so only `return_field_from_args` is guaranteed to match. Spark's `expr.dataType` is
// not: `CometCreateArray` may have widened the input to a deeply-nullable element type, and
// DataFusion list elements are nullable where Spark's `containsNull` says otherwise.
- scalarFunctionExprToProto(
- "spark_array_slice",
- arrayExprProto,
- startExprProto,
- lengthExprProto)
+ withNullShortCircuit(
+ expr,
+ scalarFunctionExprToProto(
+ "spark_array_slice",
+ arrayExprProto,
+ startExprProto,
+ lengthExprProto))
}
}
@@ -674,7 +676,7 @@ object CometArrayUnion extends CometExpressionSerde[ArrayUnion] {
val arraysUnionScalarExpr =
scalarFunctionExprToProto("array_union", leftArrayExprProto, rightArrayExprProto)
- arraysUnionScalarExpr
+ withNullShortCircuit(expr, arraysUnionScalarExpr)
}
}
@@ -734,6 +736,8 @@ object CometGetArrayItem extends CometExpressionSerde[GetArrayItem] {
.setOrdinal(ordinalExpr.get)
.setOneBased(false)
.setFailOnError(expr.failOnError)
+ // Spark evaluates the ordinal only where the array is not NULL, see withNullShortCircuit
+ .setNullShortCircuit(canShortCircuitNulls(expr))
Some(
ExprOuterClass.Expr
@@ -1007,7 +1011,7 @@ object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with Array
// (matching Spark's behavior)
val optExpr =
scalarFunctionExprToProto("spark_array_position", arrayExprProto, elementExprProto)
- optExpr
+ withNullShortCircuit(expr, optExpr)
}
}
diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala
index f9d3e2dbc0..ae32794d7f 100644
--- a/spark/src/main/scala/org/apache/comet/serde/maps.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.comet.DataTypeSupport.isComplexType
-import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto, withNullShortCircuit}
import org.apache.comet.shims.CometTypeShim
/**
@@ -128,7 +128,7 @@ object CometMapExtract extends CometExpressionSerde[GetMapValue] {
val mapExpr = exprToProtoInternal(expr.child, inputs, binding)
val keyExpr = exprToProtoInternal(expr.key, inputs, binding)
val mapExtractExpr = scalarFunctionExprToProto("map_extract", mapExpr, keyExpr)
- mapExtractExpr
+ withNullShortCircuit(expr, mapExtractExpr)
}
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
index 88a797d123..af78e16e82 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala
@@ -34,7 +34,7 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.comet.{CometConf, ConfigEntry}
import org.apache.comet.CometSparkSessionExtensions.isCometShuffleEnabled
import org.apache.comet.serde.{Compatible, OperatorOuterClass, SupportLevel, Unsupported}
-import org.apache.comet.serde.QueryPlanSerde.{exprToProto, supportedSortType}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProto, supportedSortType, withSharedSubexpressions}
import org.apache.comet.serde.operator.CometSink
object CometTakeOrderedAndProjectExec extends CometSink[TakeOrderedAndProjectExec] {
@@ -135,11 +135,21 @@ case class CometTakeOrderedAndProjectExec(
// physical limit includes the offset; the native limit applies that offset once.
CometExecUtils.getLimitNativePlan(child.output, limit, offset)
} else {
- CometExecUtils.getTopKNativePlan(child.output, sortOrder, child, limit, offset)
+ topKNativePlan(offset)
+ }
+ // Spark's projection runs on its own over the rows the sort keeps, so only its list can share
+ // a subexpression (see QueryPlanSerde.canShortCircuitNulls)
+ withSharedSubexpressions(projectList) {
+ selection.flatMap(CometExecUtils.getProjectionNativePlan(projectList, child.output, _))
}
- selection.flatMap(CometExecUtils.getProjectionNativePlan(projectList, child.output, _))
}
+ // Spark's ordering evaluates each sort key on its own, without subexpression elimination
+ private def topKNativePlan(offset: Int): Option[OperatorOuterClass.Operator] =
+ withSharedSubexpressions(Seq.empty) {
+ CometExecUtils.getTopKNativePlan(child.output, sortOrder, child, limit, offset)
+ }
+
protected override def doExecuteColumnar(): RDD[ColumnarBatch] = {
val childRDD = child.executeColumnar()
if (childRDD.getNumPartitions == 0 || limit == 0) {
@@ -157,10 +167,7 @@ case class CometTakeOrderedAndProjectExec(
} else {
val numParts = childRDD.getNumPartitions
// Serialize the plan once before mapping to avoid repeated serialization per partition
- val topK =
- CometExecUtils
- .getTopKNativePlan(child.output, sortOrder, child, limit)
- .get
+ val topK = topKNativePlan(offset = 0).get
val serializedTopK = CometExec.serializeNativePlan(topK)
val numOutputCols = child.output.length
val inputSchema = CometUtils.fromAttributes(child.output)
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
index fcd85b6f6a..b7264d0704 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
@@ -343,8 +343,13 @@ class CometNativeShuffleWriter[K, V](
val partitioning = PartitioningOuterClass.HashPartition.newBuilder()
partitioning.setNumPartitions(effectivePartitionCount)
- val partitionExprs = hashPartitioning.expressions
- .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes))
+ // Spark evaluates the keys in one projection, whose subexpression elimination sees what
+ // they share (see QueryPlanSerde.canShortCircuitNulls)
+ val partitionExprs =
+ QueryPlanSerde.withSharedSubexpressions(hashPartitioning.expressions) {
+ hashPartitioning.expressions
+ .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes))
+ }
if (partitionExprs.length != hashPartitioning.expressions.length) {
throw new UnsupportedOperationException(
@@ -377,8 +382,12 @@ class CometNativeShuffleWriter[K, V](
}
{
- val orderingExprs = rangePartitioning.ordering
- .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes))
+ // As for hash partitioning, Spark evaluates the keys in one projection
+ val orderingExprs =
+ QueryPlanSerde.withSharedSubexpressions(rangePartitioning.ordering.map(_.child)) {
+ rangePartitioning.ordering
+ .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes))
+ }
if (orderingExprs.length != rangePartitioning.ordering.length) {
throw new UnsupportedOperationException(
s"Partitioning $rangePartitioning is not supported.")
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit.sql b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit.sql
new file mode 100644
index 0000000000..559574fc38
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit.sql
@@ -0,0 +1,125 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- These array functions are null-intolerant BinaryExpressions or TernaryExpressions in Spark,
+-- which return NULL for a NULL array without evaluating the other arguments. Comet evaluates each
+-- later argument only for the rows where every earlier one is not NULL, so an argument that raises,
+-- here an ANSI cast of a malformed string, cannot raise on a row Spark skips. Each query names the
+-- expression with expect_native, because the codegen dispatcher would pass by running Spark's code.
+-- https://github.com/apache/datafusion-comet/issues/6613
+
+-- Config: spark.sql.ansi.enabled=true
+
+statement
+CREATE TABLE array_null_short_circuit(ai array<int>, s string, i string) USING parquet
+
+-- `s` and `i` are malformed only where the array is NULL. COALESCE(1) writes a single file, so the
+-- NULL arrays share a batch with the others.
+statement
+INSERT INTO array_null_short_circuit SELECT /*+ COALESCE(1) */ * FROM VALUES
+ (NULL, 'bad', 'bad'),
+ (array(1, 2), '2', '1'),
+ (array(3), '7', '0'),
+ (NULL, 'worse', 'worse'),
+ (array(4, NULL, 5), '5', '2')
+ AS v(ai, s, i)
+
+query expect_native(array_contains)
+SELECT array_contains(ai, CAST(s AS INT)) FROM array_null_short_circuit
+
+query expect_native(array_position)
+SELECT array_position(ai, CAST(s AS INT)) FROM array_null_short_circuit
+
+query expect_native(array_remove)
+SELECT array_remove(ai, CAST(s AS INT)) FROM array_null_short_circuit
+
+query expect_native(arrays_overlap)
+SELECT arrays_overlap(ai, array(CAST(s AS INT))) FROM array_null_short_circuit
+
+query expect_native(array_union)
+SELECT array_union(ai, array(CAST(s AS INT))) FROM array_null_short_circuit
+
+-- slice is a TernaryExpression: the length is evaluated only where the start is not NULL either
+query expect_native(slice)
+SELECT slice(ai, CAST(s AS INT), 1), slice(ai, 1, CAST(i AS INT)) FROM array_null_short_circuit
+
+query expect_native(getarrayitem)
+SELECT ai[CAST(i AS INT)] FROM array_null_short_circuit
+
+-- A start of 0 makes slice raise in any mode, and only rows whose array is NULL have one
+query expect_native(array_union)
+SELECT array_union(ai, slice(array(1), IF(ai IS NULL, 0, 1), 1)) FROM array_null_short_circuit
+
+-- A nondeterministic argument sees the same rows as in Spark, so it returns the same values
+query expect_native(array_position)
+SELECT array_position(ai, CAST(rand(7L) * 3 AS INT)) FROM array_null_short_circuit
+
+-- Spark's subexpression elimination evaluates a subexpression that an operator's expressions share
+-- for every row, before them, so this cast raises on the rows whose array is NULL as well, in a
+-- projection and in an aggregation. Comet does not skip an argument that holds one.
+query expect_error(CAST_INVALID_INPUT)
+SELECT array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT))
+FROM array_null_short_circuit
+
+query expect_error(CAST_INVALID_INPUT)
+SELECT sum(array_position(ai, CAST(s AS INT))), max(array_contains(ai, CAST(s AS INT)))
+FROM array_null_short_circuit
+
+-- Here the whole call is shared, and Spark evaluates its cast only where the array is not NULL
+query expect_native(array_contains)
+SELECT array_contains(ai, CAST(s AS INT)), NOT array_contains(ai, CAST(s AS INT))
+FROM array_null_short_circuit
+
+-- A shuffle evaluates its partitioning keys together in one projection, and a TakeOrderedAndProject
+-- its projection list, so a cast that two of them share raises on the rows whose array is NULL too.
+-- Both are serialized when the plan runs, after the operator's conversion.
+query expect_error(CAST_INVALID_INPUT)
+SELECT ai, s FROM array_null_short_circuit
+DISTRIBUTE BY array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT))
+
+query expect_error(CAST_INVALID_INPUT)
+SELECT array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT))
+FROM (SELECT * FROM array_null_short_circuit ORDER BY i LIMIT 10)
+
+-- A cast that only one of them holds is skipped, for hash and range partitioning keys and for a
+-- TakeOrderedAndProject
+query
+SELECT ai, s FROM array_null_short_circuit DISTRIBUTE BY array_contains(ai, CAST(s AS INT))
+
+query
+SELECT ai, s FROM array_null_short_circuit ORDER BY array_contains(ai, CAST(s AS INT)), s
+
+query expect_native(array_contains)
+SELECT array_contains(ai, CAST(s AS INT))
+FROM (SELECT * FROM array_null_short_circuit ORDER BY i LIMIT 10)
+
+-- A TakeOrderedAndProject compares rows one sort key at a time and evaluates its projection apart
+-- from the keys, so a cast in two sort keys, or in a sort key and the projection, is skipped too
+query expect_native(array_contains, array_position)
+SELECT i FROM (SELECT * FROM array_null_short_circuit
+ ORDER BY array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT)), i LIMIT 10)
+
+query expect_native(array_contains, array_position)
+SELECT array_contains(ai, CAST(s AS INT))
+FROM (SELECT * FROM array_null_short_circuit ORDER BY array_position(ai, CAST(s AS INT)), i LIMIT 10)
+
+-- Where the array is not NULL, the argument is evaluated and raises as in Spark
+query expect_error(CAST_INVALID_INPUT)
+SELECT array_contains(ai, CAST(s || 'x' AS INT)) FROM array_null_short_circuit
+
+query expect_error(CAST_INVALID_INPUT)
+SELECT ai[CAST(i || 'x' AS INT)] FROM array_null_short_circuit
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_non_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_non_ansi.sql
new file mode 100644
index 0000000000..1f557326d4
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_non_ansi.sql
@@ -0,0 +1,66 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- The queries of array_null_short_circuit.sql with ANSI off. A cast of a malformed string then
+-- returns NULL instead of raising, so the cast queries match Spark however the arguments are
+-- evaluated, and the native planner drops the guard for them. A start of 0 still makes slice raise,
+-- and a nondeterministic argument still advances on every row it sees, so those two shapes depend on
+-- skipping the rows whose array is NULL in every mode.
+-- https://github.com/apache/datafusion-comet/issues/6613
+
+-- Config: spark.sql.ansi.enabled=false
+
+statement
+CREATE TABLE array_null_short_circuit_non_ansi(ai array<int>, s string, i string) USING parquet
+
+statement
+INSERT INTO array_null_short_circuit_non_ansi SELECT /*+ COALESCE(1) */ * FROM VALUES
+ (NULL, 'bad', 'bad'),
+ (array(1, 2), '2', '1'),
+ (array(3), '7', '0'),
+ (NULL, 'worse', 'worse'),
+ (array(4, NULL, 5), '5', '2')
+ AS v(ai, s, i)
+
+query expect_native(array_contains)
+SELECT array_contains(ai, CAST(s AS INT)) FROM array_null_short_circuit_non_ansi
+
+query expect_native(array_position)
+SELECT array_position(ai, CAST(s AS INT)) FROM array_null_short_circuit_non_ansi
+
+query expect_native(array_remove)
+SELECT array_remove(ai, CAST(s AS INT)) FROM array_null_short_circuit_non_ansi
+
+query expect_native(arrays_overlap)
+SELECT arrays_overlap(ai, array(CAST(s AS INT))) FROM array_null_short_circuit_non_ansi
+
+query expect_native(array_union)
+SELECT array_union(ai, array(CAST(s AS INT))) FROM array_null_short_circuit_non_ansi
+
+query expect_native(slice)
+SELECT slice(ai, CAST(s AS INT), 1), slice(ai, 1, CAST(i AS INT)) FROM array_null_short_circuit_non_ansi
+
+query expect_native(getarrayitem)
+SELECT ai[CAST(i AS INT)] FROM array_null_short_circuit_non_ansi
+
+-- A start of 0 makes slice raise in any mode, and only rows whose array is NULL have one
+query expect_native(array_union)
+SELECT array_union(ai, slice(array(1), IF(ai IS NULL, 0, 1), 1)) FROM array_null_short_circuit_non_ansi
+
+-- A nondeterministic argument sees the same rows as in Spark, so it returns the same values
+query expect_native(array_position)
+SELECT array_position(ai, CAST(rand(7L) * 3 AS INT)) FROM array_null_short_circuit_non_ansi
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_shuffle.sql b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_shuffle.sql
new file mode 100644
index 0000000000..e5fe87017f
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_shuffle.sql
@@ -0,0 +1,38 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- array_union evaluates shuffle(b) only for the rows whose array is not NULL, as in
+-- array_null_short_circuit.sql. shuffle keeps one random generator for every batch of a partition,
+-- so its permutations have to carry on from a batch with no NULL array to one with some. The two
+-- argument shuffle(array, seed) form only exists in Spark 4.0+.
+-- https://github.com/apache/datafusion-comet/issues/6613
+
+-- MinSparkVersion: 4.0
+-- Config: spark.comet.batchSize=8192
+
+statement
+CREATE TABLE array_null_short_circuit_shuffle(id bigint, a array<int>, b array<int>) USING parquet
+
+-- A single file, which the scan reads in batches of 8192 rows (the batch size is pinned above). Only
+-- the second batch has NULL arrays.
+statement
+INSERT INTO array_null_short_circuit_shuffle
+SELECT id, IF(id >= 8192 AND id % 2 = 0, NULL, array(0)), array(1, 2, 3, 4)
+FROM range(0, 8200, 1, 1)
+
+query expect_native(array_union, shuffle)
+SELECT id, array_union(a, shuffle(b, 42)) FROM array_null_short_circuit_shuffle
diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit.sql b/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit.sql
new file mode 100644
index 0000000000..6f295f275a
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit.sql
@@ -0,0 +1,63 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- Spark returns NULL for a NULL map without evaluating the key, so a key that raises, here an ANSI
+-- cast of a malformed string, cannot raise on a row whose map is NULL. map_contains_key is
+-- array_contains(map_keys(m), k) after analysis. The array counterparts, and why each query names
+-- its expression with expect_native, are in array/array_null_short_circuit.sql.
+-- https://github.com/apache/datafusion-comet/issues/6613
+
+-- Config: spark.sql.ansi.enabled=true
+
+statement
+CREATE TABLE map_null_short_circuit(m map<int, int>, s string) USING parquet
+
+-- `s` is malformed only where the map is NULL. COALESCE(1) writes a single file, so the NULL maps
+-- share a batch with the others.
+statement
+INSERT INTO map_null_short_circuit SELECT /*+ COALESCE(1) */ * FROM VALUES
+ (NULL, 'bad'),
+ (map(1, 10, 2, 20), '2'),
+ (map(3, 30), '4'),
+ (NULL, 'worse')
+ AS v(m, s)
+
+query expect_native(array_contains)
+SELECT map_contains_key(m, CAST(s AS INT)) FROM map_null_short_circuit
+
+query expect_native(getmapvalue)
+SELECT m[CAST(s AS INT)] FROM map_null_short_circuit
+
+-- Where the map is not NULL, the key is evaluated and raises as in Spark
+query expect_error(CAST_INVALID_INPUT)
+SELECT m[CAST(s || 'x' AS INT)] FROM map_null_short_circuit
+
+-- map_from_arrays builds a map rather than reading one, but it is null-intolerant too: Spark skips
+-- the values where the keys are NULL. Comet guards it with CASE WHEN keys IS NOT NULL instead.
+statement
+CREATE TABLE map_from_arrays_null_short_circuit(k array<int>, s string) USING parquet
+
+statement
+INSERT INTO map_from_arrays_null_short_circuit SELECT /*+ COALESCE(1) */ * FROM VALUES
+ (NULL, 'bad'),
+ (array(1), '10'),
+ (NULL, 'worse'),
+ (array(2), '20')
+ AS v(k, s)
+
+query
+SELECT map_from_arrays(k, array(CAST(s AS INT))) FROM map_from_arrays_null_short_circuit
diff --git a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala
index 19476f8e13..8d248669b0 100644
--- a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala
@@ -21,10 +21,11 @@ package org.apache.comet.serde
import org.scalatest.funsuite.AnyFunSuite
+import org.apache.spark.sql.catalyst.expressions.{ArrayContains, ArrayPosition, AttributeReference, Cast}
import org.apache.spark.sql.types._
import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus}
-import org.apache.comet.serde.QueryPlanSerde.supportedDataType
+import org.apache.comet.serde.QueryPlanSerde.{canShortCircuitNulls, supportedDataType, withSharedSubexpressions}
class QueryPlanSerdeSuite extends AnyFunSuite {
@@ -100,4 +101,24 @@ class QueryPlanSerdeSuite extends AnyFunSuite {
rejected.foreach(dt => assert(!supports(dt), s"$name should reject $dt"))
}
}
+
+ test("canShortCircuitNulls skips an argument only where the shared subexpressions are known") {
+ val array = AttributeReference("ai", ArrayType(IntegerType))()
+ val value = AttributeReference("i", IntegerType)()
+ val cast = Cast(AttributeReference("s", StringType)(), IntegerType)
+ val contains = ArrayContains(array, cast)
+ val position = ArrayPosition(array, cast)
+
+ // A column is never shared, so it can be skipped anywhere
+ assert(canShortCircuitNulls(ArrayContains(array, value)))
+ // Something serialized outside an operator's conversion might share the cast
+ assert(!canShortCircuitNulls(contains))
+ withSharedSubexpressions(Seq(contains)) {
+ assert(canShortCircuitNulls(contains))
+ }
+ withSharedSubexpressions(Seq(contains, position)) {
+ assert(!canShortCircuitNulls(contains))
+ assert(!canShortCircuitNulls(position))
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org