zhangfengcdt commented on code in PR #6682: URL: https://github.com/apache/datafusion-comet/pull/6682#discussion_r4244579508
########## native/spark-expr/src/map_funcs/map_builders.rs: ########## @@ -0,0 +1,1526 @@ +// 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. + +//! Spark-compatible `map_from_arrays`, `map_from_entries` and `str_to_map`. +//! +//! The `datafusion-spark` kernels build the `MapArray` the way Spark's `ArrayBasedMapBuilder` +//! does: row by row, checking a row's key and value lengths and then its keys in order, under the +//! duplicate-key policy in `datafusion.spark.map_key_dedup_policy`. These wrappers add the one +//! check the kernels skip, a `NULL` key, and restate the kernels' errors as the `SparkError`s the +//! JVM side turns back into Spark's own: +//! +//! - a row whose key and value arrays differ in length raises `SparkError::MapKeyValueDiffSizes`, +//! which reaches the user as Spark's `_LEGACY_ERROR_TEMP_2128`; +//! - a `NULL` key raises `NULL_MAP_KEY` and, under `EXCEPTION`, a duplicate key raises +//! `DUPLICATED_MAP_KEY` naming the key. Spark inserts entries one at a time, so whichever comes +//! first in the row decides which of the two it reports. +//! +//! `str_to_map` builds its keys by splitting a string, so it needs only the duplicate-key +//! restatement. +//! +//! The kernels find duplicate keys by comparing `ScalarValue`s, which compare a `FLOAT` or +//! `DOUBLE` by its bits. Spark compares those keys as boxed values, where every NaN is one key, +//! and from Spark 4.0 it normalizes them first, so `-0.0` and `0.0` are one key too. For such a +//! key the wrappers hand the kernel the keys as Spark compares them, put back the keys that Spark +//! stores in the map it returns, and name a duplicate key as Spark does. [`MapFloatKeys`] says +//! which rule a call follows. + +use crate::conversion_funcs::java_float_string; +use crate::float_semantics::{canonicalize_nans, normalize_floats}; +use crate::SparkError; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, ListArray, MapArray, StructArray, +}; +use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer}; +use arrow::compute::filter; +use arrow::compute::kernels::zip::zip; +use arrow::datatypes::{DataType, FieldRef, Float32Type, Float64Type}; +use datafusion::common::config::MapKeyDedupPolicy; +use datafusion::common::{exec_err, internal_err, DataFusionError, HashSet, Result, ScalarValue}; +use datafusion::logical_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, +}; +use datafusion_spark::function::map::map_from_arrays::MapFromArrays as DataFusionMapFromArrays; +use datafusion_spark::function::map::map_from_entries::MapFromEntries as DataFusionMapFromEntries; +use datafusion_spark::function::map::str_to_map::SparkStrToMap as DataFusionStrToMap; +use std::sync::Arc; + +/// How a map builder compares and stores a `FLOAT` or `DOUBLE` key. Spark's `ArrayBasedMapBuilder` +/// finds duplicates in a `HashMap` of boxed keys and, from Spark 4.0, normalizes a float key +/// before it looks it up, unless `spark.sql.legacy.disableMapKeyNormalization` is set. The serde +/// picks the rule for the session, and each rule has its own function name. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum MapFloatKeys { + /// The equality of a boxed `java.lang.Double` or `java.lang.Float`, which compares + /// `doubleToLongBits` or `floatToIntBits`: every NaN is one key, while `-0.0` and `0.0` are + /// two. A map keeps each key as it first occurred. Spark 3.4 and 3.5, and Spark 4.0+ with + /// `spark.sql.legacy.disableMapKeyNormalization`. + #[default] + Boxed, + /// Spark 4.0+ compares and stores the key `NormalizeFloatingNumbers` gives: `-0.0` becomes + /// `0.0`, and every NaN the canonical NaN. `map_from_arrays` stores the keys it was given + /// when none of them repeats, because `ArrayBasedMapBuilder.from` then returns its input. + Normalized, +} + +/// Spark-compatible `map_from_arrays(keys, values)`. +#[derive(Debug, Default, PartialEq, Eq, Hash)] +pub struct SparkMapFromArrays { + inner: DataFusionMapFromArrays, + float_keys: MapFloatKeys, +} + +impl SparkMapFromArrays { + pub fn new(float_keys: MapFloatKeys) -> Self { + Self { + inner: DataFusionMapFromArrays::default(), + float_keys, + } + } +} + +impl ScalarUDFImpl for SparkMapFromArrays { + fn name(&self) -> &str { + match self.float_keys { + MapFloatKeys::Boxed => self.inner.name(), + MapFloatKeys::Normalized => "map_from_arrays_normalized_keys", + } + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + self.inner.return_type(arg_types) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { + self.inner.return_field_from_args(args) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + invoke_list_builder( + &self.inner, + args, + KeyLayout::Arrays, + self.float_keys, + |args, last_value_wins| match args { + [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { + validate_map_from_arrays(keys, values, last_value_wins) + } + other => exec_err!("map_from_arrays expects 2 arguments, got {}", other.len()), + }, + ) + } +} + +/// Spark-compatible `map_from_entries(entries)`. +#[derive(Debug, Default, PartialEq, Eq, Hash)] +pub struct SparkMapFromEntries { + inner: DataFusionMapFromEntries, + float_keys: MapFloatKeys, +} + +impl SparkMapFromEntries { + pub fn new(float_keys: MapFloatKeys) -> Self { + Self { + inner: DataFusionMapFromEntries::default(), + float_keys, + } + } +} + +impl ScalarUDFImpl for SparkMapFromEntries { + fn name(&self) -> &str { + match self.float_keys { + MapFloatKeys::Boxed => self.inner.name(), + MapFloatKeys::Normalized => "map_from_entries_normalized_keys", + } + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + self.inner.return_type(arg_types) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { + self.inner.return_field_from_args(args) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + invoke_list_builder( + &self.inner, + args, + KeyLayout::Entries, + self.float_keys, + |args, last_value_wins| match args { + [ColumnarValue::Array(entries)] => { + validate_map_from_entries(entries, last_value_wins) + } + other => exec_err!("map_from_entries expects 1 argument, got {}", other.len()), + }, + ) + } +} + +/// Spark-compatible `str_to_map(text[, pair_delim[, key_value_delim]])`. +#[derive(Debug, Default, PartialEq, Eq, Hash)] +pub struct SparkStrToMap { + inner: DataFusionStrToMap, +} + +impl ScalarUDFImpl for SparkStrToMap { + fn name(&self) -> &str { + self.inner.name() + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + self.inner.return_type(arg_types) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { + self.inner.return_field_from_args(args) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + // Splitting a string cannot produce a NULL key, so only the duplicate-key error needs + // restating here. + self.inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Quoted)) + } +} + +/// Runs `map_from_arrays` or `map_from_entries`: `validate` sees the arguments as arrays whose +/// entries start at offset zero and rejects a `NULL` key, which the kernel would store without a +/// word, then the kernel builds the maps and its own errors are restated. +fn invoke_list_builder( + inner: &dyn ScalarUDFImpl, + mut args: ScalarFunctionArgs, + layout: KeyLayout, + float_keys: MapFloatKeys, + validate: impl FnOnce(&[ColumnarValue], bool) -> Result<()>, +) -> Result<ColumnarValue> { + // The kernel evaluates an all-scalar call once and returns a scalar, which DataFusion then + // broadcasts to the batch; build that one row here too rather than `number_rows` copies. + let all_scalar = args + .args + .iter() + .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); + if all_scalar { + args.number_rows = 1; + } + expand_scalars(&mut args)?; + rebase_sliced_lists(&mut args)?; + let last_value_wins = + args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin; + let float_keys = FloatKeyRewrite::prepare(&mut args.args, layout, float_keys)?; + let result = validate(&args.args, last_value_wins).and_then(|()| { + inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Bare)) + }); + let result = match &float_keys { + None => result?, + Some(keys) => keys.restore(result.map_err(|error| keys.restate_duplicate(error))?)?, + }; + match (all_scalar, result) { + (true, ColumnarValue::Array(array)) => Ok(ColumnarValue::Scalar( + ScalarValue::try_from_array(&array, 0)?, + )), + (_, result) => Ok(result), + } +} + +/// Where a builder finds its keys among its arguments. +#[derive(Clone, Copy, PartialEq, Eq)] +enum KeyLayout { + /// `map_from_arrays(keys, values)`: the elements of the first list. + Arrays, + /// `map_from_entries(entries)`: the first field of the list's structs. + Entries, +} + +/// The `FLOAT` or `DOUBLE` keys of one call. The kernel compares keys by their bits, so it is +/// handed them as Spark compares them, with every NaN made canonical and, under +/// [`MapFloatKeys::Normalized`], `-0.0` made `0.0`. Its result then gets back the keys Spark +/// stores. +struct FloatKeyRewrite { + layout: KeyLayout, + rule: MapFloatKeys, + /// The flat keys as the call passed them. + original: ArrayRef, + /// The flat keys as Spark compares them, when that changes the bits of any of them. + compared: Option<ArrayRef>, + /// Each row's range of flat keys. + offsets: OffsetBuffer<i32>, + /// The rows Spark builds a map for. Every other row is a NULL map, whose keys it never sees. + built: BooleanBuffer, +} + +impl FloatKeyRewrite { + /// Hands `args` the keys as Spark compares them, when the keys are `FLOAT` or `DOUBLE`. Returns + /// `None` for any other key type. + fn prepare( + args: &mut [ColumnarValue], + layout: KeyLayout, + rule: MapFloatKeys, + ) -> Result<Option<Self>> { + let Some((original, offsets, built)) = float_keys_of(args, layout) else { + return Ok(None); + }; + let compared = match rule { + MapFloatKeys::Boxed => canonicalize_nans(&original), + MapFloatKeys::Normalized => normalize_floats(&original), + }; + // Keys that are not NaN, or not `-0.0` under normalization, compare the same either way, + // and then the kernel's own result is already Spark's. `ArrayData` compares primitive + // values byte for byte, so this sees their bits. + let compared = if original.to_data() == compared.to_data() { Review Comment: Agreed! I would be good to see Float64 benchmark numbers here next to the Int32. -- 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]
