andygrove commented on code in PR #6563: URL: https://github.com/apache/datafusion-comet/pull/6563#discussion_r4248919469
########## 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: Agreed, and I changed the shared helper in 9f49d10f0a. `normalize_nested_floats` now rebuilds a list over only the part of its child array that the offsets reference, at every depth, and rebases the offsets onto it. A list whose offsets already span the whole child is rebuilt as before. `test_normalize_sliced_list_rebuilds_only_referenced_values` covers a sliced list and a list of lists sliced at the outer level, and it fails without the trimming. `sliced_input` runs both functions over sliced arguments. `float_arrays` now has `array_distinct` and `array_union` cases over a whole list and over 8,192 rows sliced from a list 16 times as long, each next to DataFusion's kernel on the same input. With 8 doubles per row and no nulls, the medians of two interleaved runs were: | Case | Before | After | DataFusion | | --- | ---: | ---: | ---: | | `array_distinct`, whole | 0.71 ms | 0.71 ms | 0.69 ms | | `array_distinct`, sliced | 1.19 ms | 0.72 ms | 0.96 ms | | `array_union`, whole | 1.16 ms | 1.16 ms | 1.14 ms | | `array_union`, sliced | 2.12 ms | 1.18 ms | 1.68 ms | With 50 doubles per row, the sliced cases went from 6.6 to 3.7 ms and from 12.7 to 6.8 ms. DataFusion's own kernels are slower on the sliced input because they scan the whole child for `-0.0` before slicing it. `arrays_overlap` and `array_position` pass `list.values()` to the helper themselves, so their top level still covers the whole child. I left those two alone here. -- 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]
