diff --git a/native-engine/datafusion-ext-commons/src/lib.rs b/native-engine/datafusion-ext-commons/src/lib.rs index f63f1053a..f41031350 100644 --- a/native-engine/datafusion-ext-commons/src/lib.rs +++ b/native-engine/datafusion-ext-commons/src/lib.rs @@ -33,6 +33,7 @@ pub mod io; pub mod scalar_value; pub mod spark_bit_array; pub mod spark_bloom_filter; +pub mod spark_float; pub mod spark_hash; pub mod uda; diff --git a/native-engine/datafusion-ext-commons/src/spark_float.rs b/native-engine/datafusion-ext-commons/src/spark_float.rs new file mode 100644 index 000000000..698768759 --- /dev/null +++ b/native-engine/datafusion-ext-commons/src/spark_float.rs @@ -0,0 +1,68 @@ +// 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::cmp::Ordering; + +pub fn spark_float_cmp(left: T, right: T) -> Ordering { + // Spark places all NaNs last and treats signed zeros as equal. + if left == right || (left.is_nan() && right.is_nan()) { + Ordering::Equal + } else if left.is_nan() { + Ordering::Greater + } else if right.is_nan() || left < right { + Ordering::Less + } else { + Ordering::Greater + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn check_spark_float_cmp(nan: T) { + let values = [ + (T::neg_infinity(), 0), + (-T::one(), 1), + (-T::zero(), 2), + (T::zero(), 2), + (T::one(), 3), + (T::infinity(), 4), + (T::nan(), 5), + (-T::nan(), 5), + (nan, 5), + (-nan, 5), + ]; + for (left, left_rank) in values { + for (right, right_rank) in values { + assert_eq!( + spark_float_cmp(left, right), + left_rank.cmp(&right_rank), + "{left:?} vs {right:?}" + ); + } + } + } + + #[test] + fn test_spark_float_cmp_f32() { + check_spark_float_cmp(f32::from_bits(0x7fc0_0001)); + } + + #[test] + fn test_spark_float_cmp_f64() { + check_spark_float_cmp(f64::from_bits(0x7ff8_0000_0000_0001)); + } +} diff --git a/native-engine/datafusion-ext-functions/src/lib.rs b/native-engine/datafusion-ext-functions/src/lib.rs index a3f75a98a..eca828cd5 100644 --- a/native-engine/datafusion-ext-functions/src/lib.rs +++ b/native-engine/datafusion-ext-functions/src/lib.rs @@ -72,6 +72,7 @@ pub fn create_auron_ext_function( shared_function!(spark_get_json_object::spark_get_parsed_json_object) } "Spark_ParseJson" => shared_function!(spark_get_json_object::spark_parse_json), + "Spark_ArrayContains" => shared_function!(spark_array::array_contains), "Spark_ArrayReverse" => shared_function!(spark_array::array_reverse), "Spark_ArrayFlatten" => shared_function!(spark_array::array_flatten), "Spark_MakeArray" => shared_function!(spark_make_array::array), diff --git a/native-engine/datafusion-ext-functions/src/spark_array.rs b/native-engine/datafusion-ext-functions/src/spark_array.rs index 5137ee263..ca8e3611f 100644 --- a/native-engine/datafusion-ext-functions/src/spark_array.rs +++ b/native-engine/datafusion-ext-functions/src/spark_array.rs @@ -15,10 +15,13 @@ //! Array expressions -use std::sync::Arc; +use std::{cmp::Ordering, sync::Arc}; use arrow::{ - array::{Array, ArrayRef, ListArray, MutableArrayData, make_array}, + array::{ + Array, ArrayRef, BooleanArray, Float32Array, Float64Array, ListArray, MutableArrayData, + make_array, + }, buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}, datatypes::{DataType, Field}, }; @@ -27,9 +30,78 @@ use datafusion::{ logical_expr::ColumnarValue, }; use datafusion_ext_commons::{ - df_execution_err, downcast_any, scalar_value::compacted_scalar_value_from_array, + arrow::eq_comparator::{DynEqComparator, make_eq_comparator}, + df_execution_err, downcast_any, + scalar_value::compacted_scalar_value_from_array, + spark_float::spark_float_cmp, }; +pub fn array_contains(args: &[ColumnarValue]) -> Result { + if args.len() != 2 { + return df_execution_err!("array_contains requires exactly 2 arguments"); + } + let list_is_scalar = matches!(args[0], ColumnarValue::Scalar(_)); + let value_is_scalar = matches!(args[1], ColumnarValue::Scalar(_)); + // Keep scalar lists in one row to avoid copying their elements per input row. + let arrays = args + .iter() + .map(|arg| arg.to_array(1)) + .collect::>>()?; + let lists = downcast_any!(&arrays[0], ListArray)?; + let elements = lists.values(); + let values = &arrays[1]; + let num_rows = match (list_is_scalar, value_is_scalar) { + (true, _) => values.len(), + (_, true) => lists.len(), + _ if lists.len() == values.len() => lists.len(), + _ => return df_execution_err!("array_contains requires matching array lengths"), + }; + if elements.data_type() != values.data_type() { + return df_execution_err!("array_contains requires matching element and value types"); + } + + macro_rules! float_eq { + ($array_type:ty) => {{ + let left = downcast_any!(elements, $array_type)?.values().clone(); + let right = downcast_any!(values, $array_type)?.values().clone(); + Box::new(move |i: usize, j: usize| { + spark_float_cmp(left[i], right[j]) == Ordering::Equal + }) + }}; + } + let equals: DynEqComparator = match elements.data_type() { + DataType::Float32 => float_eq!(Float32Array), + DataType::Float64 => float_eq!(Float64Array), + _ => make_eq_comparator(elements.as_ref(), values.as_ref(), true)?, + }; + + let result: ArrayRef = Arc::new(BooleanArray::from_iter((0..num_rows).map(|row| { + let list_row = if list_is_scalar { 0 } else { row }; + let value_row = if value_is_scalar { 0 } else { row }; + if lists.is_null(list_row) || values.is_null(value_row) { + return None; + } + let mut result = Some(false); + let start = lists.value_offsets()[list_row] as usize; + let end = lists.value_offsets()[list_row + 1] as usize; + for index in start..end { + if elements.is_null(index) { + result = None; + } else if equals(index, value_row) { + return Some(true); + } + } + result + }))); + if list_is_scalar && value_is_scalar { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } else { + Ok(ColumnarValue::Array(result)) + } +} + pub fn array_reverse(args: &[ColumnarValue]) -> Result { if args.len() != 1 { return df_execution_err!("array_reverse requires exactly 1 argument"); @@ -170,14 +242,226 @@ mod test { use std::{error::Error, sync::Arc}; use arrow::{ - array::{ArrayRef, ListArray}, + array::{ + ArrayRef, BooleanArray, Float32Array, Float64Array, Int32Array, ListArray, StringArray, + }, buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}, - datatypes::{Field, Int32Type}, + datatypes::{Field, Float32Type, Float64Type, Int32Type}, }; use datafusion::common::ScalarValue; use super::*; + #[test] + fn test_array_contains_nulls_and_slices() -> Result<()> { + let lists = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(99)]), + Some(vec![Some(1), Some(2)]), + Some(vec![Some(1), Some(2)]), + Some(vec![None, Some(2)]), + Some(vec![Some(2), None]), + Some(vec![Some(1), None]), + Some(vec![]), + None, + Some(vec![Some(1)]), + Some(vec![None]), + Some(vec![]), + ]); + let values = Int32Array::from(vec![ + Some(99), + Some(2), + Some(3), + Some(2), + Some(2), + Some(3), + Some(1), + Some(1), + None, + Some(1), + None, + ]); + for (offset, len) in [(0, 11), (1, 10), (2, 0)] { + let result = array_contains(&[ + ColumnarValue::Array(Arc::new(lists.slice(offset, len))), + ColumnarValue::Array(Arc::new(values.slice(offset, len))), + ])? + .into_array(len)?; + let expected = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(false), + Some(true), + Some(true), + None, + Some(false), + None, + None, + None, + None, + ]); + let expected: ArrayRef = Arc::new(expected.slice(offset, len)); + assert_eq!(&result, &expected); + } + Ok(()) + } + + #[test] + fn test_array_contains_large_scalar_list() -> Result<()> { + let list = ScalarValue::List(Arc::new(ListArray::from_iter_primitive::( + [Some((0..10_000).map(Some))], + ))); + for len in [0, 10_000] { + let result = array_contains(&[ + ColumnarValue::Scalar(list.clone()), + ColumnarValue::Array(Arc::new(Int32Array::from(vec![0; len]))), + ])? + .into_array(len)?; + let expected: ArrayRef = Arc::new(BooleanArray::from(vec![true; len])); + assert_eq!(&result, &expected); + } + Ok(()) + } + + #[test] + fn test_array_contains_scalars() -> Result<()> { + let list = ScalarValue::List(ScalarValue::new_list( + &[ScalarValue::from("a"), ScalarValue::Utf8(None)], + &DataType::Utf8, + true, + )); + let result = array_contains(&[ + ColumnarValue::Scalar(list.clone()), + ColumnarValue::Array(Arc::new(StringArray::from(vec![ + Some("a"), + Some("b"), + None, + ]))), + ])? + .into_array(3)?; + let expected: ArrayRef = Arc::new(BooleanArray::from(vec![Some(true), None, None])); + assert_eq!(&result, &expected); + + for (value, expected) in [ + (ScalarValue::from("a"), Some(true)), + (ScalarValue::from("b"), None), + (ScalarValue::Utf8(None), None), + ] { + let result = array_contains(&[ + ColumnarValue::Scalar(list.clone()), + ColumnarValue::Scalar(value), + ])?; + assert!( + matches!(result, ColumnarValue::Scalar(ScalarValue::Boolean(value)) + if value == expected) + ); + } + + let lists = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1)]), + Some(vec![Some(2)]), + Some(vec![]), + None, + ]); + let result = array_contains(&[ + ColumnarValue::Array(Arc::new(lists)), + ColumnarValue::Scalar(ScalarValue::from(1)), + ])? + .into_array(4)?; + let expected: ArrayRef = Arc::new(BooleanArray::from(vec![ + Some(true), + Some(false), + Some(false), + None, + ])); + assert_eq!(&result, &expected); + + let null_list = ScalarValue::List(Arc::new(ListArray::new_null( + Arc::new(Field::new_list_field(DataType::Int32, true)), + 1, + ))); + assert!(matches!( + array_contains(&[ + ColumnarValue::Scalar(null_list), + ColumnarValue::Scalar(ScalarValue::from(1)), + ])?, + ColumnarValue::Scalar(ScalarValue::Boolean(None)) + )); + Ok(()) + } + + macro_rules! test_array_contains_float { + ($name:ident, $type:ty, $array:ty, $nan:expr) => { + #[test] + fn $name() -> Result<()> { + let nan = $nan; + let lists = ListArray::from_iter_primitive::<$type, _, _>(vec![ + Some(vec![Some(nan)]), + Some(vec![Some(-0.0)]), + Some(vec![Some(0.0)]), + Some(vec![Some(nan), None]), + Some(vec![Some(nan)]), + Some(vec![None, Some(nan)]), + ]); + let values = <$array>::from(vec![-nan, 0.0, -0.0, 1.0, 1.0, -nan]); + let result = array_contains(&[ + ColumnarValue::Array(Arc::new(lists)), + ColumnarValue::Array(Arc::new(values)), + ])? + .into_array(6)?; + let expected: ArrayRef = Arc::new(BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + None, + Some(false), + Some(true), + ])); + assert_eq!(&result, &expected); + Ok(()) + } + }; + } + + test_array_contains_float!( + test_array_contains_float32, + Float32Type, + Float32Array, + f32::from_bits(0x7fc00001) + ); + test_array_contains_float!( + test_array_contains_float64, + Float64Type, + Float64Array, + f64::from_bits(0x7ff8000000000001) + ); + + #[test] + fn test_array_contains_invalid_arguments() -> Result<()> { + assert!(array_contains(&[]).is_err()); + assert!( + array_contains(&[ + ColumnarValue::Scalar(ScalarValue::from(1)), + ColumnarValue::Scalar(ScalarValue::from(1)), + ]) + .is_err() + ); + let lists = ColumnarValue::Array(Arc::new( + ListArray::from_iter_primitive::(vec![Some(vec![Some(1)])]), + )); + assert!( + array_contains(&[lists.clone(), ColumnarValue::Scalar(ScalarValue::from("a")),]) + .is_err() + ); + assert!( + array_contains(&[ + lists, + ColumnarValue::Array(Arc::new(Int32Array::from(vec![1, 2]))), + ]) + .is_err() + ); + Ok(()) + } + #[test] fn test_array_reverse_int() -> Result<()> { let input: ArrayRef = Arc::new(ListArray::from_iter_primitive::(vec![ diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala index 63d2c3946..307833cc2 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala @@ -19,6 +19,10 @@ package org.apache.auron import java.text.SimpleDateFormat import org.apache.spark.sql.{AuronQueryTest, Row} +import org.apache.spark.sql.auron.NativeConverters +import org.apache.spark.sql.catalyst.expressions.Alias +import org.apache.spark.sql.execution.ProjectExec +import org.apache.spark.sql.execution.auron.plan.NativeProjectBase import org.apache.spark.sql.functions.{col, split} import org.apache.spark.sql.internal.SQLConf @@ -95,6 +99,142 @@ class AuronFunctionSuite extends AuronQueryTest with BaseAuronSQLSuite { } } + test("array_contains function") { + withTable("t1") { + sql("create table t1(id int, arr array, value int) using parquet") + sql("""insert into t1 values + | (0, array(1, 2), 2), + | (1, array(1, 2), 3), + | (2, array(null, 2), 2), + | (3, array(2, null), 2), + | (4, array(1, null), 3), + | (5, array(), 1), + | (6, null, 1), + | (7, array(1), null), + | (8, array(null), 1), + | (9, array(), null), + | (10, null, null) + |""".stripMargin) + + for (ansiEnabled <- Seq("true", "false")) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled) { + val query = "select id, array_contains(arr, value) from t1" + checkSparkAnswerAndOperator(query) + checkAnswer( + sql(query), + Seq( + Row(0, true), + Row(1, false), + Row(2, true), + Row(3, true), + Row(4, null), + Row(5, false), + Row(6, null), + Row(7, null), + Row(8, null), + Row(9, null), + Row(10, null))) + checkSparkAnswerAndOperator("""select + | array_contains(arr, 2), + | array_contains(array(1, 2), value) + |from t1 + |""".stripMargin) + checkSparkAnswerAndOperator("select id from t1 where array_contains(arr, value)") + } + } + } + } + + test("array_contains function with atomic element types") { + withTable("t1") { + sql("""create table t1( + | strings array, s string, + | bools array, b boolean, + | dates array, d date, + | timestamps array, ts timestamp, + | decimals array, n decimal(10, 2), + | ints array, l bigint) using parquet + |""".stripMargin) + sql("""insert into t1 values + | (array('a', null, 'b'), 'b', array(true, null), true, + | array(date'2024-02-29'), date'2024-02-29', + | array(timestamp'2024-02-29 12:34:56'), timestamp'2024-02-29 12:34:56', + | array(1.23, null, 4.56), 4.56, array(1, 2), 2), + | (array('a'), 'A', array(false), true, + | array(date'2024-02-29'), date'2024-03-01', + | array(timestamp'2024-02-29 12:34:56'), timestamp'2024-02-29 12:34:57', + | array(1.23), 4.56, array(1, 2), 3) + |""".stripMargin) + checkSparkAnswerAndOperator("""select + | array_contains(strings, s), array_contains(bools, b), + | array_contains(dates, d), array_contains(timestamps, ts), + | array_contains(decimals, n), array_contains(ints, l) + |from t1 + |""".stripMargin) + } + } + + test("array_contains function with floating point elements") { + withTable("t1") { + sql("create table t1(arr array, value double) using parquet") + sql("""insert into t1 values + | (array(cast('NaN' as double)), cast('NaN' as double)), + | (array(cast('-0.0' as double)), cast('0.0' as double)), + | (array(cast('0.0' as double)), cast('-0.0' as double)), + | (array(cast('Infinity' as double)), cast('Infinity' as double)), + | (array(cast('-Infinity' as double)), cast('Infinity' as double)), + | (array(cast('NaN' as double), null), 1.0), + | (array(cast('NaN' as double)), 1.0), + | (array(null, cast('NaN' as double)), cast('NaN' as double)) + |""".stripMargin) + val query = """select + | array_contains(arr, value), + | array_contains(cast(arr as array), cast(value as float)) + |from t1 + |""".stripMargin + checkSparkAnswerAndOperator(query) + checkAnswer( + sql(query), + Seq( + Row(true, true), + Row(true, true), + Row(true, true), + Row(true, true), + Row(false, false), + Row(null, null), + Row(false, false), + Row(true, true))) + } + } + + test("array_contains function with complex elements falls back") { + withTable("t1") { + sql("""create table t1 using parquet as select + | array(named_struct('a', cast(null as int))) as arr, + | named_struct('a', cast(null as int)) as value + |union all select + | array(named_struct('a', cast(null as int))), named_struct('a', 0) + |""".stripMargin) + val query = "select array_contains(arr, value) from t1" + for (fallbackEnabled <- Seq("false", "true")) { + withSQLConf("spark.auron.udf.singleChildFallback.enabled" -> fallbackEnabled) { + val df = sql(query) + checkAnswer(df, Seq(Row(true), Row(false))) + val plan = df.queryExecution.executedPlan + if (fallbackEnabled == "true") { + assertPlanIsNative(df) + val project = collectFirst(plan) { case p: NativeProjectBase => p }.get + val expression = project.expressions.collectFirst { case a: Alias => a.child }.get + val nativeExpr = NativeConverters.convertExpr(expression) + assert(nativeExpr.hasSparkUdfWrapperExpr) + } else { + assert(collectFirst(plan) { case _: ProjectExec => true }.contains(true)) + } + } + } + } + } + test("flatten function with mixed child array nullability") { withTable("t1") { sql(""" 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 e39c24cd9..c987fcec9 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 @@ -1168,6 +1168,9 @@ object NativeConverters extends Logging { buildExtScalarFunction("Spark_NormalizeNanAndZero", e.children, e.dataType) case e: CreateArray => buildExtScalarFunction("Spark_MakeArray", e.children, e.dataType) + case e: ArrayContains + if e.right.dataType.isInstanceOf[AtomicType] && isTypeSupported(e.right.dataType) => + buildExtScalarFunction("Spark_ArrayContains", e.children, e.dataType) case e: Flatten => buildExtScalarFunction("Spark_ArrayFlatten", e.children, e.dataType) case e: MapFromEntries => buildExtScalarFunction(