diff --git a/native-engine/auron-planner/proto/auron.proto b/native-engine/auron-planner/proto/auron.proto index a905c8a36..ef5c060a7 100644 --- a/native-engine/auron-planner/proto/auron.proto +++ b/native-engine/auron-planner/proto/auron.proto @@ -312,6 +312,7 @@ message PhysicalCastNode { message PhysicalNegativeNode { PhysicalExprNode expr = 1; + bool ansi_enabled = 2; } message PhysicalLikeExprNode { diff --git a/native-engine/auron-planner/src/planner.rs b/native-engine/auron-planner/src/planner.rs index b1ee15843..72a89c86e 100644 --- a/native-engine/auron-planner/src/planner.rs +++ b/native-engine/auron-planner/src/planner.rs @@ -42,8 +42,8 @@ use datafusion::{ physical_plan::{ ColumnStatistics, ExecutionPlan, Statistics, expressions as phys_expr, expressions::{ - BinaryExpr, CaseExpr, CastExpr, Column, IsNotNullExpr, IsNullExpr, Literal, - NegativeExpr, NotExpr, PhysicalSortExpr, + BinaryExpr, CaseExpr, CastExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, + PhysicalSortExpr, }, metrics::ExecutionPlanMetricsSet, }, @@ -54,7 +54,7 @@ use datafusion_ext_exprs::{ get_indexed_field::GetIndexedFieldExpr, get_map_value::GetMapValueExpr, named_struct::NamedStructExpr, row_num::RowNumExpr, spark_monotonically_increasing_id::SparkMonotonicallyIncreasingIdExpr, - spark_partition_id::SparkPartitionIdExpr, + spark_negative::SparkNegativeExpr, spark_partition_id::SparkPartitionIdExpr, spark_scalar_subquery_wrapper::SparkScalarSubqueryWrapperExpr, spark_udf_wrapper::SparkUDFWrapperExpr, string_contains::StringContainsExpr, string_ends_with::StringEndsWithExpr, string_starts_with::StringStartsWithExpr, @@ -962,8 +962,9 @@ impl PhysicalPlanner { ExprType::NotExpr(e) => Arc::new(NotExpr::new( self.try_parse_physical_expr_box_required(&e.expr, input_schema)?, )), - ExprType::Negative(e) => Arc::new(NegativeExpr::new( + ExprType::Negative(e) => Arc::new(SparkNegativeExpr::new( self.try_parse_physical_expr_box_required(&e.expr, input_schema)?, + e.ansi_enabled, )), ExprType::InList(e) => { let expr = self.try_parse_physical_expr_box_required(&e.expr, input_schema)?; diff --git a/native-engine/datafusion-ext-exprs/src/lib.rs b/native-engine/datafusion-ext-exprs/src/lib.rs index 6400f7d21..36389ae46 100644 --- a/native-engine/datafusion-ext-exprs/src/lib.rs +++ b/native-engine/datafusion-ext-exprs/src/lib.rs @@ -24,6 +24,7 @@ pub mod get_map_value; pub mod named_struct; pub mod row_num; pub mod spark_monotonically_increasing_id; +pub mod spark_negative; pub mod spark_partition_id; pub mod spark_scalar_subquery_wrapper; pub mod spark_udf_wrapper; diff --git a/native-engine/datafusion-ext-exprs/src/spark_negative.rs b/native-engine/datafusion-ext-exprs/src/spark_negative.rs new file mode 100644 index 000000000..15e395111 --- /dev/null +++ b/native-engine/datafusion-ext-exprs/src/spark_negative.rs @@ -0,0 +1,277 @@ +// 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::{ + any::Any, + fmt::{Debug, Display, Formatter}, + hash::{Hash, Hasher}, + sync::Arc, +}; + +use arrow::{ + array::{ArrayRef, Int8Array, Int16Array, Int32Array, Int64Array}, + datatypes::{DataType, Schema}, + record_batch::RecordBatch, +}; +use datafusion::{ + common::{Result, ScalarValue}, + logical_expr::ColumnarValue, + physical_expr::{PhysicalExpr, PhysicalExprRef}, + physical_plan::expressions::NegativeExpr, +}; +use datafusion_ext_commons::{df_execution_err, downcast_any}; + +pub struct SparkNegativeExpr { + expr: PhysicalExprRef, + ansi_enabled: bool, +} + +impl SparkNegativeExpr { + pub fn new(expr: PhysicalExprRef, ansi_enabled: bool) -> Self { + Self { expr, ansi_enabled } + } +} + +impl Display for SparkNegativeExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "negative({})", self.expr) + } +} + +impl Debug for SparkNegativeExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "negative({})", self.expr) + } +} + +impl PartialEq for SparkNegativeExpr { + fn eq(&self, other: &Self) -> bool { + self.expr.eq(&other.expr) && self.ansi_enabled == other.ansi_enabled + } +} + +impl Eq for SparkNegativeExpr {} + +impl Hash for SparkNegativeExpr { + fn hash(&self, state: &mut H) { + self.expr.hash(state); + self.ansi_enabled.hash(state); + } +} + +impl PhysicalExpr for SparkNegativeExpr { + fn as_any(&self) -> &dyn Any { + self + } + + fn data_type(&self, input_schema: &Schema) -> Result { + self.expr.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.expr.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + if self.ansi_enabled + && matches!( + self.expr.data_type(batch.schema().as_ref())?, + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 + ) + { + return checked_negate(self.expr.evaluate(batch)?); + } + + NegativeExpr::new(self.expr.clone()).evaluate(batch) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.expr] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + Ok(Arc::new(Self::new(children[0].clone(), self.ansi_enabled))) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "fmt_sql not used") + } +} + +fn checked_negate(value: ColumnarValue) -> Result { + Ok(match value { + ColumnarValue::Scalar(scalar) => ColumnarValue::Scalar(checked_negate_scalar(scalar)?), + ColumnarValue::Array(array) => ColumnarValue::Array(checked_negate_array(array.as_ref())?), + }) +} + +fn checked_negate_scalar(scalar: ScalarValue) -> Result { + Ok(match scalar { + ScalarValue::Int8(Some(v)) => ScalarValue::Int8(Some(checked_negate_value(v)?)), + ScalarValue::Int16(Some(v)) => ScalarValue::Int16(Some(checked_negate_value(v)?)), + ScalarValue::Int32(Some(v)) => ScalarValue::Int32(Some(checked_negate_value(v)?)), + ScalarValue::Int64(Some(v)) => ScalarValue::Int64(Some(checked_negate_value(v)?)), + ScalarValue::Int8(None) => ScalarValue::Int8(None), + ScalarValue::Int16(None) => ScalarValue::Int16(None), + ScalarValue::Int32(None) => ScalarValue::Int32(None), + ScalarValue::Int64(None) => ScalarValue::Int64(None), + other => return df_execution_err!("unsupported ANSI negative data type: {other}"), + }) +} + +macro_rules! checked_negate_primitive_array { + ($array:expr, $array_ty:ty) => {{ + let array = downcast_any!($array, $array_ty)?; + let mut values = Vec::with_capacity(array.len()); + for value in array.iter() { + values.push(match value { + Some(v) => Some(checked_negate_value(v)?), + None => None, + }); + } + Ok(Arc::new(<$array_ty>::from(values)) as ArrayRef) + }}; +} + +fn checked_negate_array(array: &dyn arrow::array::Array) -> Result { + match array.data_type() { + DataType::Int8 => checked_negate_primitive_array!(array, Int8Array), + DataType::Int16 => checked_negate_primitive_array!(array, Int16Array), + DataType::Int32 => checked_negate_primitive_array!(array, Int32Array), + DataType::Int64 => checked_negate_primitive_array!(array, Int64Array), + other => df_execution_err!("unsupported ANSI negative data type: {other}"), + } +} + +fn checked_negate_value(value: T) -> Result +where + T: CheckedNeg, +{ + value.checked_neg().ok_or_else(|| { + datafusion::common::DataFusionError::Execution( + "[ARITHMETIC_OVERFLOW] arithmetic overflow in unary minus".to_string(), + ) + }) +} + +trait CheckedNeg { + fn checked_neg(self) -> Option + where + Self: Sized; +} + +macro_rules! impl_checked_neg { + ($($ty:ty),+) => { + $( + impl CheckedNeg for $ty { + fn checked_neg(self) -> Option { + <$ty>::checked_neg(self) + } + } + )+ + }; +} + +impl_checked_neg!(i8, i16, i32, i64); + +#[cfg(test)] +mod test { + use std::{error::Error, sync::Arc}; + + use arrow::{ + array::{ArrayRef, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array}, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, + }; + use datafusion::physical_expr::{PhysicalExpr, expressions::Column}; + + use super::SparkNegativeExpr; + + #[test] + fn test_ansi_checked_negation() -> Result<(), Box> { + let cases: Vec<(DataType, ArrayRef)> = vec![ + ( + DataType::Int8, + Arc::new(Int8Array::from(vec![Some(i8::MIN)])), + ), + ( + DataType::Int16, + Arc::new(Int16Array::from(vec![Some(i16::MIN)])), + ), + ( + DataType::Int32, + Arc::new(Int32Array::from(vec![Some(i32::MIN)])), + ), + ( + DataType::Int64, + Arc::new(Int64Array::from(vec![Some(i64::MIN)])), + ), + ]; + + for (data_type, array) in cases { + let batch = batch(data_type, array)?; + let err = expression(true) + .evaluate(&batch) + .expect_err("expected overflow"); + assert!(err.to_string().contains("[ARITHMETIC_OVERFLOW]")); + } + Ok(()) + } + + #[test] + fn test_non_ansi_wrapping_negation() -> Result<(), Box> { + let batch = batch( + DataType::Int32, + Arc::new(Int32Array::from(vec![Some(i32::MIN), Some(1), None])), + )?; + let output = expression(false) + .evaluate(&batch)? + .into_array(batch.num_rows())?; + let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(i32::MIN), Some(-1), None])); + assert_eq!(&output, &expected); + Ok(()) + } + + #[test] + fn test_float_delegates_to_datafusion() -> Result<(), Box> { + let batch = batch( + DataType::Float64, + Arc::new(Float64Array::from(vec![Some(1.5), None])), + )?; + let output = expression(true) + .evaluate(&batch)? + .into_array(batch.num_rows())?; + let expected: ArrayRef = Arc::new(Float64Array::from(vec![Some(-1.5), None])); + assert_eq!(&output, &expected); + Ok(()) + } + + fn expression(ansi_enabled: bool) -> Arc { + Arc::new(SparkNegativeExpr::new( + Arc::new(Column::new("c", 0)), + ansi_enabled, + )) + } + + fn batch(data_type: DataType, array: ArrayRef) -> Result> { + Ok(RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("c", data_type, true)])), + vec![array], + )?) + } +} diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronExpressionSuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronExpressionSuite.scala index b0006d363..29b04d897 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronExpressionSuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronExpressionSuite.scala @@ -33,9 +33,52 @@ class AuronExpressionSuite extends AuronQueryTest with BaseAuronSQLSuite { } test("UnaryMinus") { - // Negating Int.MinValue overflows. Under ANSI mode (default in Spark 4.x) vanilla Spark - // throws while the native engine wraps, so the comparison diverges. Disable ANSI so both - // engines wrap consistently and the boundary value can still be exercised. + withSQLConf("spark.sql.ansi.enabled" -> "true") { + withTable("t1") { + sql("create table t1(col1 int) using parquet") + sql(""" + |insert into t1 values + | (1), + | (0), + | (-2147483648) + |""".stripMargin) + + withSQLConf("spark.auron.enable" -> "false") { + assertArithmeticOverflow(sql("SELECT negative(col1), -(col1) FROM t1"), "overflow") + } + withSQLConf("spark.auron.enable" -> "true") { + val df = sql("SELECT negative(col1), -(col1) FROM t1") + assertArithmeticOverflow(df, "[ARITHMETIC_OVERFLOW]") + assertNativePlan(df) + } + } + } + } + + test("UnaryMinusLong") { + withSQLConf("spark.sql.ansi.enabled" -> "true") { + withTable("t1") { + sql("create table t1(col1 bigint) using parquet") + sql(""" + |insert into t1 values + | (1), + | (0), + | (cast(-9223372036854775808 as bigint)) + |""".stripMargin) + + withSQLConf("spark.auron.enable" -> "false") { + assertArithmeticOverflow(sql("SELECT negative(col1), -(col1) FROM t1"), "overflow") + } + withSQLConf("spark.auron.enable" -> "true") { + val df = sql("SELECT negative(col1), -(col1) FROM t1") + assertArithmeticOverflow(df, "[ARITHMETIC_OVERFLOW]") + assertNativePlan(df) + } + } + } + } + + test("UnaryMinus without ANSI") { withSQLConf("spark.sql.ansi.enabled" -> "false") { withTable("t1") { sql("create table t1(col1 int) using parquet") @@ -45,4 +88,50 @@ class AuronExpressionSuite extends AuronQueryTest with BaseAuronSQLSuite { } } } + + test("UnaryMinus honors Spark's default ANSI setting") { + withTable("t1") { + sql("create table t1(col1 int) using parquet") + sql("insert into t1 values(-2147483648)") + + if (spark.conf.get("spark.sql.ansi.enabled").toBoolean) { + val df = sql("SELECT negative(col1), -(col1) FROM t1") + assertArithmeticOverflow(df, "[ARITHMETIC_OVERFLOW]") + assertNativePlan(df) + } else { + checkSparkAnswerAndOperator("SELECT negative(col1), -(col1) FROM t1") + } + } + } + + private def assertArithmeticOverflow( + df: => org.apache.spark.sql.DataFrame, + expectedMessage: String): Unit = { + val err = intercept[Exception] { + df.collect() + } + assert(allCauseMessages(err).toLowerCase.contains(expectedMessage.toLowerCase)) + } + + private def allCauseMessages(err: Throwable): String = { + val messages = scala.collection.mutable.ArrayBuffer.empty[String] + var current = err + while (current != null) { + Option(current.getMessage).foreach(messages += _) + current = current.getCause + } + messages.mkString(" | caused by: ") + } + + private def assertNativePlan(df: org.apache.spark.sql.DataFrame): Unit = { + val plan = stripAQEPlan(df.queryExecution.executedPlan) + plan + .collectFirst { case op if !isNativeOrPassThrough(op) => op } + .foreach { op => + fail(s""" + |Found non-native operator: ${op.nodeName} + |plan: + |${plan}""".stripMargin) + } + } } diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala index 378a8d662..adecbc59e 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala @@ -569,6 +569,7 @@ object NativeConverters extends Logging { pb.PhysicalNegativeNode .newBuilder() .setExpr(convertExprWithFallback(unaryMinus.child, isPruningExpr, fallback)) + .setAnsiEnabled(SQLConf.get.ansiEnabled) .build()) }