rich7420 commented on code in PR #6563: URL: https://github.com/apache/datafusion-comet/pull/6563#discussion_r4243520159
########## native/spark-expr/src/array_funcs/array_set_ops.rs: ########## @@ -0,0 +1,239 @@ +// 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::sync::Arc; + +use crate::float_semantics::normalize_nested_floats; +use arrow::datatypes::DataType; +use datafusion::common::{Result, ScalarValue}; +use datafusion::functions_nested::set_ops::{array_distinct_udf, array_union_udf}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, +}; + +/// Spark's `array_distinct` and `array_union` for elements that hold a float at any depth. +/// +/// Spark 4.2.0 normalizes the arguments of these functions in the plan (SPARK-54918), and 4.0.5, +/// 4.1.4 and 4.2.1 normalize while evaluating them (SPARK-59602). Either way, Spark treats `-0.0` +/// and `0.0`, and every NaN representation, as one value at any depth, and returns the normalized +/// value. DataFusion folds `-0.0` into `0.0` only in a flat float array and compares NaNs by their +/// bits, so normalize the input before delegating. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkArraySetOp { + name: &'static str, + datafusion_udf: Arc<ScalarUDF>, +} + +impl SparkArraySetOp { + pub fn distinct() -> Self { + Self { + name: "spark_array_distinct", + datafusion_udf: array_distinct_udf(), + } + } + + pub fn union() -> Self { + Self { + name: "spark_array_union", + datafusion_udf: array_union_udf(), + } + } +} + +impl ScalarUDFImpl for SparkArraySetOp { + fn name(&self) -> &str { + self.name + } + + fn signature(&self) -> &Signature { + self.datafusion_udf.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + self.datafusion_udf.return_type(arg_types) + } + + fn invoke_with_args(&self, mut args: ScalarFunctionArgs) -> Result<ColumnarValue> { + args.args = args + .args + .into_iter() + .map(|arg| match arg { + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(normalize_nested_floats(&array))) Review Comment: Could we normalize only the child ranges referenced by each slice? Native `explode` retains the full backing array, so each output batch can copy and normalize it again. A local release-mode component benchmark with 4,096 visible rows and 64× backing took 1.21 ms here versus 0.61 ms with recursive trimming. Could we update the shared helper and cover sliced inputs in a benchmark? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
