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/23 08:58:46 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6180-a81a591e582ad631f0ad536360a45bab745c720a created (now 653c7954a4)
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-6180-a81a591e582ad631f0ad536360a45bab745c720a
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 653c7954a4 feat: support AtLeastNNonNulls natively (#6180)
This branch includes the following new commits:
new 653c7954a4 feat: support AtLeastNNonNulls natively (#6180)
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: feat: support AtLeastNNonNulls natively (#6180)
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-6180-a81a591e582ad631f0ad536360a45bab745c720a
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 653c7954a485666a4b7c5487616b8755bbcf643c
Author: KUAN-HAO HUANG <10...@users.noreply.github.com>
AuthorDate: Fri Oct 23 08:32:20 2026 +0000
feat: support AtLeastNNonNulls natively (#6180)
* feat: support AtLeastNNonNulls natively
* refactor: drop AtLeastNNonNulls small batch threshold
Keep only the bitmap counter. At 8192 rows, the row counter was slower
in 51 of 56 measured cases on a non-idle machine.
* fix: fall back when AtLeastNNonNulls would skip a shared subexpression
Spark's subexpression elimination evaluates a shared subexpression for
every row before the predicate, so an ANSI failure in it surfaces even
on rows where AtLeastNNonNulls reaches n first. Decline native
conversion when a child after the first n, which native evaluation may
skip, holds such a subexpression, reusing the check behind
canShortCircuitNulls.
* fix: preserve shared expressions across decimal promotion
---
.../expression-audits/predicate_funcs.md | 13 +
docs/source/user-guide/latest/expressions.md | 1 +
native/core/src/execution/planner.rs | 17 +-
native/proto/src/proto/expr.proto | 7 +
native/spark-expr/Cargo.toml | 4 +
native/spark-expr/benches/at_least_n_non_nulls.rs | 137 ++++++
native/spark-expr/src/lib.rs | 2 +-
.../src/predicate_funcs/at_least_n_non_nulls.rs | 523 +++++++++++++++++++++
native/spark-expr/src/predicate_funcs/mod.rs | 2 +
.../org/apache/comet/serde/QueryPlanSerde.scala | 23 +-
.../scala/org/apache/comet/serde/predicates.scala | 36 +-
.../org/apache/comet/CometExpressionSuite.scala | 173 ++++++-
.../apache/comet/serde/QueryPlanSerdeSuite.scala | 34 +-
.../CometPredicateExpressionBenchmark.scala | 98 +++-
14 files changed, 1053 insertions(+), 17 deletions(-)
diff --git a/docs/source/contributor-guide/expression-audits/predicate_funcs.md b/docs/source/contributor-guide/expression-audits/predicate_funcs.md
index 1396802126..1d77746366 100644
--- a/docs/source/contributor-guide/expression-audits/predicate_funcs.md
+++ b/docs/source/contributor-guide/expression-audits/predicate_funcs.md
@@ -84,6 +84,19 @@
- Spark 4.0.1 (audited 2026-05-27): semantics unchanged.
- Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1.
+## AtLeastNNonNulls (DataFrame API)
+
+- Spark 3.4.3 (audited 2026-09-24): non-nullable Boolean predicate used by `DataFrame.na.drop`. Counts non-null, non-NaN children and stops evaluating a row once the threshold is reached. Non-positive thresholds return true without evaluating children.
+- Spark 3.5.9 (audited 2026-09-24): identical to 3.4.3.
+- Spark 4.0.4 (audited 2026-09-24): identical to 3.4.3.
+- Spark 4.1.3 (audited 2026-09-24): identical to 3.4.3.
+
+Comet uses a native physical expression to preserve per-row short-circuit evaluation. Non-floating types use outer logical nullability; nested nulls and empty collections remain non-null values. Rust tests cover thresholds, sliced and empty batches, dictionary values, literals, and required and skipped child errors. DataFrame tests cover native routing, disabled-expression fallback, complex types and ANSI casts. This internal predicate has no SQL function name.
+
+Native evaluation, like Spark's, may skip the children after the first `n` on a row. When subexpression elimination shares a subexpression with one of those children, Spark evaluates it for every row, so Comet falls back to Spark rather than skip an error Spark would raise. Plain `DataFrame.na.drop` children are columns and stay native.
+
+- The [Criterion benchmark](../../../../native/spark-expr/benches/at_least_n_non_nulls.rs) covers batch sizes around the 64-row bitmap word boundary and NULL/NaN densities.
+
## between
- Spark 3.4.3 (audited 2026-05-27): the SQL form `expr BETWEEN low AND high` is rewritten at the parser level to `expr >= low AND expr <= high`. Comet sees only the resulting `And(GreaterThanOrEqual, LessThanOrEqual)` and routes via `CometAnd` + `CometGreaterThanOrEqual` + `CometLessThanOrEqual`.
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index 784b93f5eb..f9292b2734 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -705,6 +705,7 @@ Comet also accelerates a number of Catalyst expressions that have no Spark SQL f
- **Operator and optimizer-injected expressions:** runtime bloom-filter join probes (`BloomFilterMightContain`, `BloomFilterAggregate`), optimized `IN` sets (`InSet`), scalar subqueries (`ScalarSubquery`), and floating-point normalization (`KnownFloatingPointNormalized`).
- **Accessor expressions (subscript and field access, not functions):** struct field access (`col.field`), array element access (`arr[i]`), and map value access (`map[key]`).
- **Internal decimal arithmetic:** `CheckOverflow`, `MakeDecimal`, and `UnscaledValue`, which the analyzer inserts around decimal operations.
+- **Missing-data predicates:** `AtLeastNNonNulls`, used by `DataFrame.na.drop`, counts non-null and non-NaN values natively.
- **User-defined functions:** Scala UDFs registered through the DataFrame or SQL API.
- **DataSource V2 catalog functions:** Iceberg's system functions `bucket`, `truncate`, `years`, `months`, `days`, and `hours` (for example `system.bucket(16, id)`) run natively; see [Iceberg system functions](iceberg.md#iceberg-system-functions). Other DataSource V2 catalog functions run in Spark.
- **Lowered built-ins:** Spark lowers some built-in functions to `StaticInvoke` or `Invoke` calls. Those without a native mapping run through the JVM codegen dispatcher when their input and output types are supported.
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 6759dbd8e3..dfcc0072a6 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -116,10 +116,11 @@ use datafusion_comet_proto::{
};
use datafusion_comet_spark_expr::{
coerce_to_common_type, create_case_when, create_if_expr, decimal128_add_sub_width,
- jvm_udf::JvmScalarUdfExpr, spark_in_list, ArrayInsert, Cast, CheckOverflow, CreateNamedStruct,
- DecimalRescaleCheckOverflow, FloatOperands, GetArrayStructFields, GetStructField, IfExpr,
- ListExtract, NormalizeNaNAndZero, NormalizeNestedFloats, NullShortCircuit, SparkCastOptions,
- ToJson, UnboundColumn, WideDecimalBinaryExpr, WideDecimalOp,
+ jvm_udf::JvmScalarUdfExpr, spark_in_list, ArrayInsert, AtLeastNNonNulls, Cast, CheckOverflow,
+ CreateNamedStruct, DecimalRescaleCheckOverflow, FloatOperands, GetArrayStructFields,
+ GetStructField, IfExpr, ListExtract, NormalizeNaNAndZero, NormalizeNestedFloats,
+ NullShortCircuit, SparkCastOptions, ToJson, UnboundColumn, WideDecimalBinaryExpr,
+ WideDecimalOp,
};
use itertools::Itertools;
use jni::objects::{Global, JObject};
@@ -1080,6 +1081,14 @@ impl PhysicalPlanner {
csv_write_options,
)))
}
+ ExprStruct::AtLeastNNonNulls(expr) => {
+ let children = expr
+ .children
+ .iter()
+ .map(|child| self.create_expr(child, Arc::clone(&input_schema)))
+ .collect::<Result<Vec<_>, _>>()?;
+ Ok(Arc::new(AtLeastNNonNulls::new(expr.n, children)))
+ }
ExprStruct::ArraysZip(expr) => {
if expr.values.is_empty() {
return Err(GeneralError(
diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto
index 969b3aaa9c..72b1d03af7 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -95,6 +95,7 @@ message Expr {
RandStr rand_str = 73;
Uuid uuid = 74;
NativeScalarUdf native_scalar_udf = 75;
+ AtLeastNNonNulls at_least_n_non_nulls = 76;
}
reserved 20;
@@ -767,3 +768,9 @@ message NativeScalarUdf {
// does not need a wire change.
bool deterministic = 5;
}
+
+// Counts non-null and non-NaN children, with per-row short-circuit evaluation.
+message AtLeastNNonNulls {
+ int32 n = 1;
+ repeated Expr children = 2;
+}
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 27264ba5e8..563c4410bc 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -510,3 +510,7 @@ harness = false
[[bench]]
name = "greatest_least"
harness = false
+
+[[bench]]
+name = "at_least_n_non_nulls"
+harness = false
diff --git a/native/spark-expr/benches/at_least_n_non_nulls.rs b/native/spark-expr/benches/at_least_n_non_nulls.rs
new file mode 100644
index 0000000000..391453fe95
--- /dev/null
+++ b/native/spark-expr/benches/at_least_n_non_nulls.rs
@@ -0,0 +1,137 @@
+// 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 std::hint::black_box;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::{ArrayRef, AsArray, Float32Array, Float64Array, RecordBatch, StringArray};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::physical_expr::expressions::Column;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion_comet_spark_expr::AtLeastNNonNulls;
+use rand::{rngs::StdRng, RngExt, SeedableRng};
+
+fn benchmark(c: &mut Criterion) {
+ let mut group = c.benchmark_group("at_least_n_non_nulls");
+ group.sample_size(30);
+ group.warm_up_time(Duration::from_millis(200));
+ group.measurement_time(Duration::from_secs(1));
+ for rows in [1, 32, 63, 64, 65, 128, 129, 1024, 8192] {
+ for kind in ["f64", "f32", "utf8"] {
+ for width in [4, 32] {
+ for (nulls, nans) in [
+ (0, 0),
+ (0, 5),
+ (0, 9),
+ (0, 10),
+ (0, 11),
+ (0, 19),
+ (0, 50),
+ (0, 95),
+ (0, 100),
+ (5, 0),
+ (5, 5),
+ (9, 0),
+ (10, 0),
+ (11, 0),
+ (19, 0),
+ (50, 0),
+ (50, 5),
+ (95, 0),
+ (100, 0),
+ ] {
+ if kind == "utf8" && nans != 0 {
+ continue;
+ }
+ // Keep smaller batch controls focused; the default batch size covers
+ // the full NULL/NaN density matrix.
+ if rows != 8192 && (kind == "f32" || !matches!((nulls, nans), (0, 0) | (50, 0)))
+ {
+ continue;
+ }
+ let mut rng = StdRng::seed_from_u64(42);
+ let mut counts = vec![0; rows];
+ let arrays = (0..width)
+ .map(|column| {
+ let values = counts
+ .iter_mut()
+ .enumerate()
+ .map(|(row, count)| {
+ let bucket = rng.random_range(0..100);
+ *count += usize::from(bucket >= nulls + nans);
+ if bucket < nulls {
+ None
+ } else if bucket < nulls + nans {
+ Some(f64::NAN)
+ } else {
+ Some((row + column) as f64)
+ }
+ })
+ .collect::<Vec<_>>();
+ let array: ArrayRef = match kind {
+ "f64" => Arc::new(Float64Array::from(values)),
+ "f32" => Arc::new(Float32Array::from_iter(
+ values.into_iter().map(|v| v.map(|v| v as f32)),
+ )),
+ _ => Arc::new(StringArray::from_iter(
+ values.into_iter().map(|v| v.map(|_| "value")),
+ )),
+ };
+ (format!("c{column}"), array)
+ })
+ .collect::<Vec<_>>();
+ let batch = RecordBatch::try_from_iter(arrays).unwrap();
+ for n in [1, width / 2, width] {
+ if rows != 8192 && n != width / 2 {
+ continue;
+ }
+ let children: Vec<Arc<dyn PhysicalExpr>> = (0..width)
+ .map(|i| Arc::new(Column::new(&format!("c{i}"), i)) as _)
+ .collect();
+ let expr = AtLeastNNonNulls::new(n as i32, children);
+ let expected = counts.iter().map(|&count| count >= n).collect::<Vec<_>>();
+ let result = expr
+ .evaluate(&batch)
+ .unwrap()
+ .into_array(batch.num_rows())
+ .unwrap();
+ assert_eq!(result.null_count(), 0);
+ assert_eq!(
+ result.as_boolean().values().iter().collect::<Vec<_>>(),
+ expected
+ );
+ group.bench_with_input(
+ BenchmarkId::new(
+ format!("{rows}/{kind}_{width}_null{nulls}_nan{nans}"),
+ n,
+ ),
+ &batch,
+ |b, batch| {
+ b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap()))
+ },
+ );
+ }
+ }
+ }
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(benches, benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs
index 2ed90938f4..0e7cfadaea 100644
--- a/native/spark-expr/src/lib.rs
+++ b/native/spark-expr/src/lib.rs
@@ -37,7 +37,7 @@ mod unbound;
pub use unbound::UnboundColumn;
mod predicate_funcs;
pub mod utils;
-pub use predicate_funcs::{spark_isnan, RLike};
+pub use predicate_funcs::{spark_isnan, AtLeastNNonNulls, RLike};
mod agg_funcs;
mod array_funcs;
diff --git a/native/spark-expr/src/predicate_funcs/at_least_n_non_nulls.rs b/native/spark-expr/src/predicate_funcs/at_least_n_non_nulls.rs
new file mode 100644
index 0000000000..b3a2d652e2
--- /dev/null
+++ b/native/spark-expr/src/predicate_funcs/at_least_n_non_nulls.rs
@@ -0,0 +1,523 @@
+// 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 std::fmt::{Display, Formatter};
+use std::sync::Arc;
+
+use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, RecordBatch};
+use arrow::buffer::{BooleanBuffer, Buffer};
+use arrow::compute::cast;
+use arrow::datatypes::{DataType, Float32Type, Float64Type, Schema};
+use datafusion::common::{Result, ScalarValue};
+use datafusion::physical_expr::expressions::{Column, Literal};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_plan::ColumnarValue;
+
+/// Spark's per-row count of non-null, non-NaN children, stopping once `n` is reached.
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct AtLeastNNonNulls {
+ n: i32,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+}
+
+impl AtLeastNNonNulls {
+ pub fn new(n: i32, children: Vec<Arc<dyn PhysicalExpr>>) -> Self {
+ Self { n, children }
+ }
+}
+
+impl Display for AtLeastNNonNulls {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "atleastnnonnulls({}, {:?})", self.n, self.children)
+ }
+}
+
+impl PhysicalExpr for AtLeastNNonNulls {
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ Display::fmt(self, f)
+ }
+
+ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+ Ok(DataType::Boolean)
+ }
+
+ fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+ Ok(false)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ if self.n <= 0 {
+ return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))));
+ }
+ // The standard na.drop("all"/"any") cases reduce to bitmap OR/AND and do not
+ // need a per-row counter. For AND, still evaluate every child: Spark does not
+ // skip later expressions just because a row can no longer reach the threshold.
+ if self.n == 1 || self.n as usize == self.children.len() {
+ let all = self.n != 1;
+ let mut result = if all {
+ BooleanBuffer::new_set(batch.num_rows())
+ } else {
+ BooleanBuffer::new_unset(batch.num_rows())
+ };
+ for child in &self.children {
+ if batch.num_rows() == 0 || (!all && result.count_set_bits() == batch.num_rows()) {
+ break;
+ }
+ let value = if all || child.is::<Column>() || child.is::<Literal>() {
+ child.evaluate(batch)?
+ } else {
+ child.evaluate_selection(batch, &BooleanArray::new(!&result, None))?
+ };
+ let valid = valid_values(&value.into_array(batch.num_rows())?)?;
+ result = if all {
+ &result & &valid
+ } else {
+ &result | &valid
+ };
+ }
+ return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new(
+ result, None,
+ ))));
+ }
+ // One bit plane per binary digit holds 64 row counts in each chunk.
+ // Increment these counters with a half-adder, saturating at n by masking
+ // out rows that have already reached it.
+ let bits = (i32::BITS - self.n.leading_zeros()) as usize;
+ let chunks = batch.num_rows().div_ceil(64);
+ let mut counts = vec![0_u64; chunks * bits];
+ let mut result = vec![0_u64; chunks];
+ let mut remaining = batch.num_rows();
+ for child in &self.children {
+ if remaining == 0 {
+ break;
+ }
+ let value =
+ if remaining == batch.num_rows() || child.is::<Column>() || child.is::<Literal>() {
+ child.evaluate(batch)?
+ } else {
+ let selection = BooleanArray::new(
+ BooleanBuffer::collect_bool(batch.num_rows(), |row| {
+ result[row / 64] & (1 << (row % 64)) == 0
+ }),
+ None,
+ );
+ child.evaluate_selection(batch, &selection)?
+ };
+ let scalar = matches!(value, ColumnarValue::Scalar(_));
+ let array = value.into_array(if scalar { 1 } else { batch.num_rows() })?;
+ let valid = valid_values(&array)?;
+ let valid = if scalar {
+ if !valid.value(0) {
+ continue;
+ }
+ BooleanBuffer::new_set(batch.num_rows())
+ } else {
+ valid
+ };
+ for ((counts, result), valid) in counts
+ .chunks_exact_mut(bits)
+ .zip(&mut result)
+ .zip(valid.bit_chunks().iter_padded())
+ {
+ let mut carry = valid & !*result;
+ if carry == 0 {
+ continue;
+ }
+ for count in counts.iter_mut() {
+ let next = *count & carry;
+ *count ^= carry;
+ carry = next;
+ if carry == 0 {
+ break;
+ }
+ }
+ // Counts never exceed n, so intersecting its set bit planes
+ // identifies equality without comparing the zero bit planes.
+ let mut threshold = (self.n as u32) ^ (1 << (bits - 1));
+ let mut reached = counts[bits - 1];
+ while threshold != 0 && reached != 0 {
+ reached &= counts[threshold.trailing_zeros() as usize];
+ threshold &= threshold - 1;
+ }
+ remaining -= (reached & !*result).count_ones() as usize;
+ *result |= reached;
+ }
+ }
+ // Arrow stores bitmap bytes in little-endian order on every platform.
+ let result = result.into_iter().map(u64::to_le).collect::<Vec<_>>();
+ Ok(ColumnarValue::Array(Arc::new(BooleanArray::new(
+ BooleanBuffer::new(Buffer::from_vec(result), 0, batch.num_rows()),
+ None,
+ ))))
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ self.children.iter().collect()
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ Ok(Arc::new(Self::new(self.n, children)))
+ }
+}
+
+fn valid_values(array: &ArrayRef) -> Result<BooleanBuffer> {
+ // Dictionary validity also includes null dictionary values. Floating-point dictionaries
+ // need decoding to inspect NaNs; other types need only their logical validity bitmap.
+ let array = match array.data_type() {
+ DataType::Dictionary(_, value_type) if value_type.is_floating() => cast(array, value_type)?,
+ _ => Arc::clone(array),
+ };
+ let valid = array.logical_nulls().map_or_else(
+ || BooleanBuffer::new_set(array.len()),
+ |n| n.inner().clone(),
+ );
+ // Computing NaN bits independently of nulls lets Arrow scan the values without
+ // a per-element validity branch. The final AND removes values beneath nulls.
+ Ok(match array.data_type() {
+ DataType::Float32 => {
+ let not_nan =
+ BooleanArray::from_unary(array.as_primitive::<Float32Type>(), |v| !v.is_nan());
+ &valid & not_nan.values()
+ }
+ DataType::Float64 => {
+ let not_nan =
+ BooleanArray::from_unary(array.as_primitive::<Float64Type>(), |v| !v.is_nan());
+ &valid & not_nan.values()
+ }
+ _ => valid,
+ })
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{
+ DictionaryArray, Float32Array, Float64Array, Int32Array, NullArray, StringArray,
+ };
+ use arrow::compute::CastOptions;
+ use arrow::datatypes::{Field, Int32Type};
+ use datafusion::physical_expr::expressions::CastExpr;
+
+ fn batch(arrays: Vec<ArrayRef>) -> RecordBatch {
+ let schema = Schema::new(
+ arrays
+ .iter()
+ .enumerate()
+ .map(|(i, a)| Field::new(format!("c{i}"), a.data_type().clone(), true))
+ .collect::<Vec<_>>(),
+ );
+ RecordBatch::try_new(Arc::new(schema), arrays).unwrap()
+ }
+
+ fn columns(batch: &RecordBatch) -> Vec<Arc<dyn PhysicalExpr>> {
+ batch
+ .schema()
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(i, f)| Arc::new(Column::new(f.name(), i)) as _)
+ .collect()
+ }
+
+ fn result(expr: &AtLeastNNonNulls, batch: &RecordBatch) -> Vec<bool> {
+ let array = expr
+ .evaluate(batch)
+ .unwrap()
+ .into_array(batch.num_rows())
+ .unwrap();
+ assert_eq!(array.null_count(), 0);
+ array.as_boolean().values().iter().collect()
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_survives_rewrite() {
+ let input = batch(vec![Arc::new(Int32Array::from(vec![Some(1), None]))]);
+ let expr = Arc::new(AtLeastNNonNulls::new(2, vec![]));
+ let rewritten = expr.with_new_children(columns(&input)).unwrap();
+ let rewritten = rewritten.downcast_ref::<AtLeastNNonNulls>().unwrap();
+ assert_eq!(rewritten.n, 2);
+ assert_eq!(result(rewritten, &input), [false, false]);
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_mixed_types() {
+ let batch = batch(vec![
+ Arc::new(StringArray::from(vec![
+ Some(""),
+ None,
+ Some("NaN"),
+ None,
+ None,
+ None,
+ ])),
+ Arc::new(Float32Array::from(vec![
+ Some(f32::NAN),
+ Some(-0.0),
+ None,
+ None,
+ // Native arrays preserve NaN signs and payloads that Parquet canonicalizes.
+ Some(f32::from_bits(0xffc0_0001)),
+ Some(-0.0),
+ ])),
+ Arc::new(Float64Array::from(vec![
+ Some(1.0),
+ Some(f64::NAN),
+ Some(f64::INFINITY),
+ None,
+ Some(-0.0),
+ Some(f64::from_bits(0xfff8_0000_0000_0001)),
+ ])),
+ Arc::new(NullArray::new(6)),
+ ]);
+ for n in [-1, 0, 1, 2, 3, 4, 5, i32::MAX] {
+ let expr = AtLeastNNonNulls::new(n, columns(&batch));
+ assert_eq!(
+ result(&expr, &batch),
+ [2, 1, 2, 0, 1, 1].map(|count| count >= n)
+ );
+ assert!(!expr.nullable(batch.schema_ref()).unwrap());
+ }
+ let sliced = batch.slice(1, 2);
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(2, columns(&sliced)), &sliced),
+ [false, true]
+ );
+ assert!(result(
+ &AtLeastNNonNulls::new(1, columns(&batch)),
+ &batch.slice(0, 0)
+ )
+ .is_empty());
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_sliced_bitmaps() {
+ let batch = batch(
+ (0..5)
+ .map(|column| {
+ Arc::new(Float64Array::from_iter((0..133).map(|row| {
+ match (row + column) % 7 {
+ 0 => None,
+ 1 => Some(f64::NAN),
+ _ => Some(row as f64),
+ }
+ }))) as ArrayRef
+ })
+ .collect(),
+ )
+ .slice(3, 129);
+ for n in 0..=6 {
+ let expected = (3..132)
+ .map(|row| (0..5).filter(|column| (row + column) % 7 > 1).count() >= n)
+ .collect::<Vec<_>>();
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(n as i32, columns(&batch)), &batch),
+ expected
+ );
+ }
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_dictionary_values_and_scalars() {
+ let dict = DictionaryArray::<Int32Type>::try_new(
+ Int32Array::from(vec![Some(0), Some(1), Some(2), None]),
+ Arc::new(Float64Array::from(vec![Some(3.0), Some(f64::NAN), None])),
+ )
+ .unwrap();
+ let strings = DictionaryArray::<Int32Type>::try_new(
+ Int32Array::from(vec![Some(0), Some(1), Some(2), None]),
+ Arc::new(StringArray::from(vec![Some(""), Some("NaN"), None])),
+ )
+ .unwrap();
+ let batch = batch(vec![Arc::new(dict), Arc::new(strings)]);
+ let mut children = columns(&batch);
+ children.push(Arc::new(Literal::new(ScalarValue::Utf8(Some("".into())))));
+ children.push(Arc::new(Literal::new(ScalarValue::Float32(Some(f32::NAN)))));
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(3, children.clone()), &batch),
+ [true, false, false, false]
+ );
+ assert_eq!(
+ result(
+ &AtLeastNNonNulls::new(1, vec![Arc::new(Column::new("c1", 1))]),
+ &batch
+ ),
+ [true, true, false, false]
+ );
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(1, vec![]), &batch),
+ [false; 4]
+ );
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_short_circuits_per_row() {
+ let batch = batch(vec![
+ Arc::new(Int32Array::from(vec![Some(1), None])),
+ Arc::new(StringArray::from(vec!["bad", "7"])),
+ ]);
+ let children: Vec<Arc<dyn PhysicalExpr>> = vec![
+ Arc::new(Column::new("c0", 0)),
+ Arc::new(CastExpr::new(
+ Arc::new(Column::new("c1", 1)),
+ DataType::Int32,
+ Some(CastOptions {
+ safe: false,
+ ..Default::default()
+ }),
+ )),
+ ];
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(1, children.clone()), &batch),
+ [true, true]
+ );
+ // Also exercise the general counter path, not only the n=1 bitmap path.
+ let mut general = children.clone();
+ general.insert(0, Arc::new(Literal::new(ScalarValue::Boolean(Some(false)))));
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(2, general), &batch),
+ [true, true]
+ );
+ // Both non-positive thresholds and an empty batch must avoid evaluating the bad cast.
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(0, children.clone()), &batch),
+ [true, true]
+ );
+ assert!(result(
+ &AtLeastNNonNulls::new(2, children.clone()),
+ &batch.slice(0, 0)
+ )
+ .is_empty());
+ // An impossible threshold is not permission to skip evaluations Spark would perform.
+ let mut required = children;
+ required[0] = Arc::new(Literal::new(ScalarValue::Int32(None)));
+ for n in [2, 3] {
+ assert!(AtLeastNNonNulls::new(n, required.clone())
+ .evaluate(&batch)
+ .is_err());
+ }
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_counter_bit_boundaries() {
+ // Every possible count up to 65 occurs on both sides of word boundaries.
+ // Rotating the valid children also exercises different carry orders.
+ let input = batch(
+ (0..65)
+ .map(|column| {
+ Arc::new(Float64Array::from_iter((0..200).map(|row| {
+ if (column + row) % 65 < row % 66 {
+ Some(row as f64)
+ } else if column % 2 == 0 {
+ Some(f64::NAN)
+ } else {
+ None
+ }
+ }))) as ArrayRef
+ })
+ .collect(),
+ );
+ for offset in [0, 1, 7, 63] {
+ for len in [0, 1, 31, 32, 63, 64, 65, 127, 128, 129] {
+ let sliced = input.slice(offset, len);
+ for n in (0..=66).chain([i32::MAX as usize]) {
+ let expected = (offset..offset + len)
+ .map(|row| row % 66 >= n)
+ .collect::<Vec<_>>();
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(n as i32, columns(&sliced)), &sliced),
+ expected,
+ "offset={offset}, len={len}, n={n}"
+ );
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_bitmap_dictionary_and_scalars() {
+ let keys =
+ Int32Array::from_iter(
+ (0..137).map(|row| if row % 4 == 3 { None } else { Some(row % 4) }),
+ );
+ let floats = DictionaryArray::<Int32Type>::try_new(
+ keys.clone(),
+ Arc::new(Float64Array::from(vec![Some(3.0), Some(f64::NAN), None])),
+ )
+ .unwrap();
+ let strings = DictionaryArray::<Int32Type>::try_new(
+ keys,
+ Arc::new(StringArray::from(vec![Some(""), Some("NaN"), None])),
+ )
+ .unwrap();
+ let input = batch(vec![Arc::new(floats), Arc::new(strings)]).slice(3, 129);
+ let mut children = columns(&input);
+ children.push(Arc::new(Literal::new(ScalarValue::Utf8(Some("".into())))));
+ children.push(Arc::new(Literal::new(ScalarValue::Float32(Some(f32::NAN)))));
+ for n in 1..=5 {
+ let expected = (3..132)
+ .map(|row| [3, 2, 1, 1][row % 4] >= n)
+ .collect::<Vec<_>>();
+ assert_eq!(
+ result(&AtLeastNNonNulls::new(n, children.clone()), &input),
+ expected
+ );
+ }
+ }
+
+ #[test]
+ fn at_least_n_non_nulls_mixed_completion_before_cast() {
+ let input = batch(vec![
+ Arc::new(Int32Array::from_iter((0..137).map(|row| {
+ if row % 3 == 0 {
+ None
+ } else {
+ Some(1)
+ }
+ }))),
+ Arc::new(StringArray::from_iter((0..137).map(|row| {
+ if row % 3 == 0 {
+ Some("7")
+ } else {
+ Some("bad")
+ }
+ }))),
+ ]);
+ let children: Vec<Arc<dyn PhysicalExpr>> = vec![
+ Arc::new(Column::new("c0", 0)),
+ Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
+ Arc::new(CastExpr::new(
+ Arc::new(Column::new("c1", 1)),
+ DataType::Int32,
+ Some(CastOptions {
+ safe: false,
+ ..Default::default()
+ }),
+ )),
+ Arc::new(Literal::new(ScalarValue::Int32(None))),
+ ];
+ for len in [1, 63, 64, 65, 129] {
+ let sliced = input.slice(7, len);
+ let expr = |n| AtLeastNNonNulls::new(n, children.clone());
+ assert_eq!(result(&expr(2), &sliced), vec![true; len]);
+ // These rows still need the cast when n=3, so its error is required.
+ assert!(expr(3).evaluate(&sliced).is_err());
+ }
+ }
+}
diff --git a/native/spark-expr/src/predicate_funcs/mod.rs b/native/spark-expr/src/predicate_funcs/mod.rs
index 5f1f570c05..6f04549318 100644
--- a/native/spark-expr/src/predicate_funcs/mod.rs
+++ b/native/spark-expr/src/predicate_funcs/mod.rs
@@ -15,8 +15,10 @@
// specific language governing permissions and limitations
// under the License.
+mod at_least_n_non_nulls;
mod is_nan;
mod rlike;
+pub use at_least_n_non_nulls::AtLeastNNonNulls;
pub use is_nan::spark_isnan;
pub use rlike::RLike;
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 82e5efbe47..afa51ee6ab 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -100,6 +100,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
private val predicateExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
classOf[And] -> CometAnd,
+ classOf[AtLeastNNonNulls] -> CometAtLeastNNonNulls,
classOf[EqualTo] -> CometEqualTo,
classOf[EqualNullSafe] -> CometEqualNullSafe,
classOf[GreaterThan] -> CometGreaterThan,
@@ -1313,11 +1314,12 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
private val sharedSubexpressions = new DynamicVariable[Option[SharedSubexpressions]](None)
private final class SharedSubexpressions(expressions: Seq[Expression], alwaysShared: Boolean) {
+ // Find shared nodes in Spark's original trees, then compare them after decimal promotion.
lazy val common: Seq[Expression] =
if (alwaysShared || SQLConf.get.subexpressionEliminationEnabled) {
val equivalence = new EquivalentExpressions
expressions.foreach(equivalence.addExprTree(_))
- equivalence.getCommonSubexpressions
+ equivalence.getCommonSubexpressions.map(DecimalPrecision.promote)
} else {
Nil
}
@@ -1343,11 +1345,22 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
* 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
+ def canShortCircuitNulls(expr: Expression): Boolean = canSkipEvaluation(expr.children.tail)
+
+ /**
+ * Whether native execution may leave `skippable`, expressions that Spark evaluates only for
+ * some rows, unevaluated on the rows Spark skips. Not if one of them holds a subexpression that
+ * the operator's expressions share and Spark eliminates, nor outside
+ * `withSharedSubexpressions`; see `canShortCircuitNulls`. Normalize both sides with decimal
+ * promotion so callers may pass either original or already-promoted trees.
+ */
+ def canSkipEvaluation(skippable: Seq[Expression]): Boolean = {
// 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))))
+ skippable.forall(_.isInstanceOf[LeafExpression]) || sharedSubexpressions.value.exists { s =>
+ // Avoid decimal promotion when Spark shares nothing.
+ s.common.isEmpty || !skippable.exists { expr =>
+ DecimalPrecision.promote(expr).exists(e => s.common.exists(_.semanticEquals(e)))
+ }
}
}
diff --git a/spark/src/main/scala/org/apache/comet/serde/predicates.scala b/spark/src/main/scala/org/apache/comet/serde/predicates.scala
index 446568d44f..ae13a41926 100644
--- a/spark/src/main/scala/org/apache/comet/serde/predicates.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/predicates.scala
@@ -21,13 +21,13 @@ package org.apache.comet.serde
import scala.jdk.CollectionConverters._
-import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BinaryExpression, EqualNullSafe, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, InSet, IsNaN, IsNotNull, IsNull, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, Literal, Not, Or}
+import org.apache.spark.sql.catalyst.expressions.{And, AtLeastNNonNulls, Attribute, BinaryExpression, EqualNullSafe, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, InSet, IsNaN, IsNotNull, IsNull, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, Literal, Not, Or}
import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{BooleanType, DoubleType, FloatType}
import org.apache.comet.CometConf
-import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus}
+import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, withFallbackReason}
import org.apache.comet.serde.ExprOuterClass.Expr
import org.apache.comet.serde.QueryPlanSerde._
@@ -491,3 +491,35 @@ object ComparisonUtils {
}
}
}
+
+object CometAtLeastNNonNulls extends CometExpressionSerde[AtLeastNNonNulls] {
+ override def convert(
+ expr: AtLeastNNonNulls,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[Expr] = {
+ // Native evaluation, like Spark's, stops counting a row once it reaches n, which takes at
+ // least the first n children. So it may skip the children after those, which a shared
+ // subexpression that Spark eliminates would evaluate for every row.
+ val skippable = if (expr.n <= 0) expr.children else expr.children.drop(expr.n)
+ if (!canSkipEvaluation(skippable)) {
+ withFallbackReason(
+ expr,
+ "AtLeastNNonNulls would skip a subexpression that Spark evaluates for every row")
+ return None
+ }
+ val children = expr.children.map(exprToProtoInternal(_, inputs, binding))
+ if (children.forall(_.isDefined)) {
+ Some(
+ Expr
+ .newBuilder()
+ .setAtLeastNNonNulls(
+ ExprOuterClass.AtLeastNNonNulls
+ .newBuilder()
+ .setN(expr.n)
+ .addAllChildren(children.map(_.get).asJava))
+ .build())
+ } else {
+ None
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index a6992e2b7e..1d104351e2 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -24,7 +24,7 @@ import scala.util.Random
import org.apache.hadoop.fs.Path
import org.apache.spark.SparkThrowable
import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row}
-import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, Cast, Divide, FromUnixTime, In, InSet, IntegralDivide, Literal, Multiply, Remainder, StructsToJson, Subtract, TruncDate, TruncTimestamp}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AtLeastNNonNulls, AttributeReference, Cast, Divide, FromUnixTime, In, InSet, IntegralDivide, Literal, Multiply, Remainder, StructsToJson, Subtract, TruncDate, TruncTimestamp}
import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, OptimizeIn, SimplifyExtractValueOps}
import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec, CometSortExec, CometTakeOrderedAndProjectExec}
import org.apache.spark.sql.execution.{LocalTableScanExec, ProjectExec, SparkPlan}
@@ -62,6 +62,177 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
checkCometAnswer(cometDf, sparkAnswer)
}
+ test("AtLeastNNonNulls: na.drop counts non-null and non-NaN values") {
+ val input = sql("""SELECT * FROM VALUES
+ (1, '', cast('NaN' AS FLOAT), 1D),
+ (2, NULL, -0F, cast('NaN' AS DOUBLE)),
+ (3, 'NaN', cast(NULL AS FLOAT), cast('Infinity' AS DOUBLE)),
+ (4, NULL, cast(NULL AS FLOAT), cast(NULL AS DOUBLE)) AS t(id, s, f, d)""")
+ withTempPath { dir =>
+ input.write.parquet(dir.getCanonicalPath)
+ withParquetTable(dir.getCanonicalPath, "nonnulls") {
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "2") {
+ for (n <- Seq(-1, 0, 1, 2, 3, 4)) {
+ checkSparkAnswerAndImpl(
+ spark.table("nonnulls").na.drop(n, Seq("s", "f", "d")),
+ native = Seq("atleastnnonnulls"))
+ }
+ checkSparkAnswerAndImpl(
+ spark.table("nonnulls").na.drop("all", Seq("s", "f", "d")),
+ native = Seq("atleastnnonnulls"))
+ withSQLConf(CometConf.getExprEnabledConfigKey("AtLeastNNonNulls") -> "false") {
+ checkSparkAnswerAndFallbackReason(
+ spark.table("nonnulls").na.drop(2, Seq("s", "f", "d")),
+ "Expression support is disabled. Set spark.comet.expression.AtLeastNNonNulls.enabled=true to enable it.")
+ }
+ }
+ }
+ }
+ }
+
+ test("AtLeastNNonNulls: counter and batch boundaries") {
+ val columns = (0 until 33).map(i => s"c$i")
+ val fields = columns.zipWithIndex.map { case (name, i) =>
+ s"CASE WHEN (id + $i) % 33 < id % 34 THEN cast(id AS DOUBLE) " +
+ s"WHEN $i % 2 = 0 THEN cast('NaN' AS DOUBLE) ELSE NULL END AS $name"
+ }
+ withTempPath { dir =>
+ spark.range(137).selectExpr((Seq("id") ++ fields): _*).write.parquet(dir.getCanonicalPath)
+ withParquetTable(dir.getCanonicalPath, "nonnulls") {
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "65") {
+ for (n <- Seq(2, 3, 7, 8, 15, 16, 17, 31, 32, 33, 34)) {
+ checkSparkAnswerAndImpl(
+ spark.table("nonnulls").na.drop(n, columns),
+ native = Seq("atleastnnonnulls"))
+ }
+ }
+ }
+ }
+ }
+
+ test("AtLeastNNonNulls: complex values use only their outer nullability") {
+ val input = sql("""SELECT * FROM VALUES
+ (array(cast(NULL AS INT)), named_struct('a', cast(NULL AS INT)), map(1, cast(NULL AS INT))),
+ (array(), named_struct('a', 1), map()),
+ (NULL, NULL, NULL) AS t(a, s, m)""")
+ withTempPath { dir =>
+ input.write.parquet(dir.getCanonicalPath)
+ withParquetTable(dir.getCanonicalPath, "nonnulls") {
+ for (n <- Seq(1, 2, 3, 4)) {
+ checkSparkAnswerAndImpl(
+ spark.table("nonnulls").na.drop(n),
+ native = Seq("atleastnnonnulls"))
+ }
+ }
+ }
+ }
+
+ test("AtLeastNNonNulls: short-circuit child evaluation") {
+ withParquetTable(Seq((Some(1), "bad"), (None, "7")), "nonnulls") {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ def query(n: Int, prefix: Boolean = false): DataFrame = {
+ val df = spark.table("nonnulls")
+ val children = Seq(df.col("_1").expr, Cast(df.col("_2").expr, IntegerType))
+ val args = if (prefix) Literal(1) +: children else children
+ df.select(getColumnFromExpression(AtLeastNNonNulls(n, args)))
+ }
+ for (n <- Seq(-1, 0, 1)) {
+ checkSparkAnswerAndImpl(query(n), native = Seq("atleastnnonnulls"))
+ }
+ // Reach the threshold on different children for different rows in the counter path.
+ checkSparkAnswerAndImpl(query(2, prefix = true), native = Seq("atleastnnonnulls"))
+ for (n <- Seq(2, 3)) {
+ assertExpressionImpl(
+ query(n).queryExecution.executedPlan,
+ Seq("atleastnnonnulls"),
+ Seq.empty)
+ checkSparkError(query(n), "CAST_INVALID_INPUT")
+ }
+ }
+ }
+ }
+
+ test("AtLeastNNonNulls: falls back rather than skip a shared subexpression") {
+ withParquetTable(Seq((Some(1), "bad"), (None, "7")), "nonnulls") {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ // Every row reaches n=2 before its second cast, and the first row before either cast.
+ // Spark evaluates the repeated cast for every row only when it eliminates it as a
+ // shared subexpression, and then fails on "bad".
+ def query: DataFrame = {
+ val df = spark.table("nonnulls")
+ val cast = Cast(df.col("_2").expr, IntegerType)
+ df.select(
+ getColumnFromExpression(
+ AtLeastNNonNulls(2, Seq(Literal(1), df.col("_1").expr, cast, cast))))
+ }
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "false") {
+ checkSparkAnswerAndImpl(query, native = Seq("atleastnnonnulls"))
+ }
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+ // A shared subexpression among the first n children is evaluated for every row anyway.
+ val df = spark.table("nonnulls")
+ val plus = Add(df.col("_1").expr, Literal(1))
+ checkSparkAnswerAndImpl(
+ df.select(getColumnFromExpression(
+ AtLeastNNonNulls(2, Seq(plus, plus, Cast(df.col("_2").expr, IntegerType))))),
+ native = Seq("atleastnnonnulls"))
+ }
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+ val df = query
+ val reasons =
+ new ExtendedExplainInfo().getFallbackReasons(df.queryExecution.executedPlan)
+ assert(
+ reasons.exists(_.contains("would skip a subexpression that Spark evaluates")),
+ reasons)
+ checkSparkAnswerMaybeThrows(df) match {
+ case (Some(sparkException), Some(cometException)) =>
+ for (error <- Seq(sparkException, cometException)) {
+ assert(
+ causeChain(error).exists(e =>
+ String.valueOf(e.getMessage).contains("CAST_INVALID_INPUT")),
+ error)
+ }
+ case result => fail(s"Expected both Spark and Comet to fail, got $result")
+ }
+ }
+ }
+ }
+ }
+
+ test("AtLeastNNonNulls: shared decimal subexpressions preserve divide-by-zero errors") {
+ val rows = Seq((Some(1), BigDecimal(1), BigDecimal(0)), (None, BigDecimal(7), BigDecimal(1)))
+ withParquetTable(rows, "nonnulls") {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ def query: DataFrame = {
+ val df = spark.table("nonnulls")
+ val cast = Cast(Divide(df.col("_2").expr, df.col("_3").expr), IntegerType)
+ df.select(
+ getColumnFromExpression(
+ AtLeastNNonNulls(2, Seq(Literal(1), df.col("_1").expr, cast, cast))))
+ }
+
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "false") {
+ checkSparkAnswerAndImpl(query, native = Seq("atleastnnonnulls"))
+ }
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+ val df = query
+ checkSparkAnswerMaybeThrows(df) match {
+ case (Some(sparkException), Some(cometException)) =>
+ for (error <- Seq(sparkException, cometException)) {
+ assert(arithmeticError(error).getErrorClass == "DIVIDE_BY_ZERO")
+ }
+ case result => fail(s"Expected both Spark and Comet to fail, got $result")
+ }
+ val reasons =
+ new ExtendedExplainInfo().getFallbackReasons(df.queryExecution.executedPlan)
+ assert(
+ reasons.exists(_.contains("would skip a subexpression that Spark evaluates")),
+ reasons)
+ }
+ }
+ }
+ }
+
test("nested floating point membership uses native In and InSet") {
withTable("nested_in_plan") {
sql("CREATE TABLE nested_in_plan (a ARRAY<DOUBLE>) USING parquet")
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 71f329eb17..856cc8a7d0 100644
--- a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala
@@ -21,12 +21,13 @@ 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.catalyst.expressions.{ArrayContains, ArrayPosition, AttributeReference, Cast, Divide}
+import org.apache.spark.sql.comet.DecimalPrecision
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus}
-import org.apache.comet.serde.QueryPlanSerde.{canShortCircuitNulls, supportedDataType, withSharedSubexpressions}
+import org.apache.comet.serde.QueryPlanSerde.{canShortCircuitNulls, canSkipEvaluation, supportedDataType, withSharedSubexpressions}
class QueryPlanSerdeSuite extends AnyFunSuite {
@@ -149,4 +150,33 @@ class QueryPlanSerdeSuite extends AnyFunSuite {
assert(!canShortCircuitNulls(contains))
}
}
+
+ test("canSkipEvaluation recognizes shared subexpressions across decimal promotion") {
+ val left = AttributeReference("a", DecimalType(10, 0))()
+ val right = AttributeReference("b", DecimalType(10, 0))()
+ val cast = Cast(Divide(left, right), IntegerType)
+ val array = AttributeReference("ai", ArrayType(IntegerType))()
+ val contains = ArrayContains(array, cast)
+ val position = ArrayPosition(array, cast)
+ val promoted = DecimalPrecision.promote(cast)
+ assert(!cast.semanticEquals(promoted))
+
+ for (enabled <- Seq(false, true); alwaysShared <- Seq(false, true)) {
+ val conf = new SQLConf
+ conf.setConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED, enabled)
+ SQLConf.withExistingConf(conf) {
+ withSharedSubexpressions(Seq(contains, position), alwaysShared) {
+ val canSkip = !(enabled || alwaysShared)
+ assert(canSkipEvaluation(Seq(cast)) == canSkip)
+ assert(canSkipEvaluation(Seq(promoted)) == canSkip)
+ assert(canShortCircuitNulls(contains) == canSkip)
+ assert(canShortCircuitNulls(DecimalPrecision.promote(contains)) == canSkip)
+ }
+ withSharedSubexpressions(Seq(contains), alwaysShared) {
+ assert(canSkipEvaluation(Seq(cast)))
+ assert(canSkipEvaluation(Seq(promoted)))
+ }
+ }
+ }
+ }
}
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPredicateExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPredicateExpressionBenchmark.scala
index a7e22d8e90..d72b58210c 100644
--- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPredicateExpressionBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPredicateExpressionBenchmark.scala
@@ -19,11 +19,19 @@
package org.apache.spark.sql.benchmark
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.Row
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
/**
* Benchmark to measure Comet execution performance. To run this benchmark:
* `SPARK_GENERATE_BENCHMARK_FILES=1 make
* benchmark-org.apache.spark.sql.benchmark.CometPredicateExpressionBenchmark` Results will be
- * written to "spark/benchmarks/CometPredicateExpressionBenchmark -**results.txt".
+ * written to "spark/benchmarks/CometPredicateExpressionBenchmark -**results.txt". Pass `--
+ * atleastnnonnulls` to benchmark `DataFrame.na.drop`. Add `--focus-double` for the 32-column,
+ * half-threshold controls with count-only and full-column consumers, 0/50% NULL and 0/5% NaN. Add
+ * `--reverse-modes` to reverse mode order in an independent run.
*/
object CometPredicateExpressionBenchmark extends CometBenchmarkBase {
@@ -43,8 +51,94 @@ object CometPredicateExpressionBenchmark extends CometBenchmarkBase {
}
}
+ def atLeastNNonNullsBenchmark(rows: Int, focusDouble: Boolean, reverseModes: Boolean): Unit = {
+ for (kind <- (if (focusDouble) Seq("double") else Seq("double", "string"));
+ width <- (if (focusDouble) Seq(32) else Seq(4, 32));
+ missing <- Seq(0, 50);
+ nanPercent <- (if (kind == "double") Seq(0, 5) else Seq(0))) {
+ withTempPath { dir =>
+ val fields = (0 until width).map { i =>
+ val bucket = s"pmod(xxhash64(id, $i), 100)"
+ val nan = if (nanPercent > 0) {
+ s"WHEN $bucket < ${missing + nanPercent} THEN cast('NaN' AS DOUBLE) "
+ } else ""
+ s"CASE WHEN $bucket < $missing THEN NULL " + nan +
+ s"ELSE cast(id + $i AS $kind) END AS c$i"
+ }
+ spark.range(rows.toLong).selectExpr(fields: _*).write.parquet(dir.getCanonicalPath)
+ val key = CometConf.getExprEnabledConfigKey("AtLeastNNonNulls")
+ val defaultModes = Seq(
+ ("Comet native", true, true),
+ ("Comet fallback", true, false),
+ ("Spark", false, true))
+ val modes = if (reverseModes) defaultModes.reverse else defaultModes
+ for (threshold <- (if (focusDouble) Seq(width / 2) else Seq(1, width / 2, width));
+ consumed <-
+ (if (focusDouble) Seq(0, width)
+ else Seq(0, 1, math.min(4, width), width).distinct)) {
+ val benchmark = new Benchmark(
+ s"na.drop: $width $kind columns, $missing% NULL, $nanPercent% NaN, " +
+ s"n=$threshold, consumed=$consumed",
+ rows.toLong,
+ output = output)
+ var expected: Option[Row] = None
+ modes.foreach { case (name, comet, native) =>
+ def run(verify: Boolean = false): Row = {
+ var result = Row.empty
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> comet.toString,
+ CometConf.COMET_EXEC_ENABLED.key -> comet.toString,
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false",
+ key -> native.toString) {
+ // Compare count-only, narrow and full consumers of the filtered output.
+ val sums = (0 until consumed).map { i =>
+ val value =
+ if (kind == "string") s"length(c$i)" else s"if(isnan(c$i), 0D, c$i)"
+ s"sum($value)"
+ }
+ val df = spark.read
+ .parquet(dir.getCanonicalPath)
+ .na
+ .drop(threshold)
+ .selectExpr((Seq("count(*)") ++ sums): _*)
+ result = df.collect().head
+ if (verify) {
+ val plan = df.queryExecution.executedPlan
+ val info = new ExtendedExplainInfo()
+ require(
+ info
+ .getNativeExpressions(plan)
+ .contains("atleastnnonnulls") == (comet && native))
+ require(!info.getCodegenDispatchExpressions(plan).contains("atleastnnonnulls"))
+ if (kind == "double" && consumed > 0 && comet && native) {
+ // The cleaned sums must exercise native IF, not silently fall back.
+ require(info.getNativeExpressions(plan).contains("if"))
+ require(!info.getCodegenDispatchExpressions(plan).contains("if"))
+ }
+ }
+ }
+ result
+ }
+ val result = run(verify = true)
+ expected.foreach(value => require(value == result, s"$name: $result != $value"))
+ expected = Some(result)
+ benchmark.addCase(name)(_ => run())
+ }
+ benchmark.run()
+ }
+ }
+ }
+ }
+
override def runCometBenchmark(mainArgs: Array[String]): Unit = {
- val values = 1024 * 1024;
+ val values = 1024 * 1024
+ if (mainArgs.contains("atleastnnonnulls")) {
+ atLeastNNonNullsBenchmark(
+ values,
+ mainArgs.contains("--focus-double"),
+ mainArgs.contains("--reverse-modes"))
+ return
+ }
runBenchmarkWithTable("inExpr", values) { v =>
inExprBenchmark(v)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org