rich7420 commented on code in PR #5854: URL: https://github.com/apache/datafusion-comet/pull/5854#discussion_r4016897902
########## native/spark-expr/src/map_funcs/map_builders.rs: ########## @@ -0,0 +1,651 @@ +// 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` and already follow Spark's +//! `spark.sql.mapKeyDedupPolicy`, which Comet forwards as +//! `datafusion.spark.map_key_dedup_policy`. These wrappers add the checks Spark's +//! `ArrayBasedMapBuilder` performs before inserting an entry, and restate the upstream errors +//! as the Spark error classes `SparkErrorConverter` turns back into `QueryExecutionErrors`: +//! +//! - a `NULL` key element raises `[NULL_MAP_KEY]`, ahead of any duplicate-key check, because +//! Spark rejects the `NULL` before it reaches the dedup map; +//! - a key array and value array of different lengths raise `[MAP_KEY_VALUE_DIFF_SIZES]`; +//! - a duplicate key under `EXCEPTION` raises `[DUPLICATED_MAP_KEY]` naming the key. +//! +//! `str_to_map` builds its keys by splitting a string, so it needs only the duplicate-key +//! restatement. + +use crate::SparkError; +use arrow::array::{Array, ArrayRef, AsArray, StructArray}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion::common::{exec_err, DataFusionError, Result}; +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; + +/// Spark-compatible `map_from_arrays(keys, values)`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkMapFromArrays { + inner: DataFusionMapFromArrays, +} + +impl Default for SparkMapFromArrays { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapFromArrays { + pub fn new() -> Self { + Self { + inner: DataFusionMapFromArrays::new(), + } + } +} + +impl ScalarUDFImpl for SparkMapFromArrays { + 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> { + let args = expand_scalars(args)?; + match args.args.as_slice() { + [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { + validate_map_from_arrays(keys, values)? + } + other => return exec_err!("map_from_arrays expects 2 arguments, got {}", other.len()), + } + self.inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Bare)) Review Comment: Verified at `3085702`. The sliced-list regressions pass, and additional `LAST_WIN` probes pass for both builders, including different key/value offsets and unused trailing entries in `map_from_arrays`. ########## native/spark-expr/src/map_funcs/map_builders.rs: ########## @@ -0,0 +1,651 @@ +// 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` and already follow Spark's +//! `spark.sql.mapKeyDedupPolicy`, which Comet forwards as +//! `datafusion.spark.map_key_dedup_policy`. These wrappers add the checks Spark's +//! `ArrayBasedMapBuilder` performs before inserting an entry, and restate the upstream errors +//! as the Spark error classes `SparkErrorConverter` turns back into `QueryExecutionErrors`: +//! +//! - a `NULL` key element raises `[NULL_MAP_KEY]`, ahead of any duplicate-key check, because +//! Spark rejects the `NULL` before it reaches the dedup map; +//! - a key array and value array of different lengths raise `[MAP_KEY_VALUE_DIFF_SIZES]`; +//! - a duplicate key under `EXCEPTION` raises `[DUPLICATED_MAP_KEY]` naming the key. +//! +//! `str_to_map` builds its keys by splitting a string, so it needs only the duplicate-key +//! restatement. + +use crate::SparkError; +use arrow::array::{Array, ArrayRef, AsArray, StructArray}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion::common::{exec_err, DataFusionError, Result}; +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; + +/// Spark-compatible `map_from_arrays(keys, values)`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkMapFromArrays { + inner: DataFusionMapFromArrays, +} + +impl Default for SparkMapFromArrays { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapFromArrays { + pub fn new() -> Self { + Self { + inner: DataFusionMapFromArrays::new(), + } + } +} + +impl ScalarUDFImpl for SparkMapFromArrays { + 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> { + let args = expand_scalars(args)?; + match args.args.as_slice() { + [ColumnarValue::Array(keys), ColumnarValue::Array(values)] => { + validate_map_from_arrays(keys, values)? + } + other => return exec_err!("map_from_arrays expects 2 arguments, got {}", other.len()), + } + self.inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Bare)) + } +} + +/// Spark-compatible `map_from_entries(entries)`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkMapFromEntries { + inner: DataFusionMapFromEntries, +} + +impl Default for SparkMapFromEntries { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapFromEntries { + pub fn new() -> Self { + Self { + inner: DataFusionMapFromEntries::new(), + } + } +} + +impl ScalarUDFImpl for SparkMapFromEntries { + 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> { + let args = expand_scalars(args)?; + match args.args.as_slice() { + [ColumnarValue::Array(entries)] => validate_map_from_entries(entries)?, + other => return exec_err!("map_from_entries expects 1 argument, got {}", other.len()), + } + self.inner + .invoke_with_args(args) + .map_err(|error| as_spark_error(error, DuplicateKeyFormat::Bare)) + } +} + +/// Spark-compatible `str_to_map(text[, pair_delim[, key_value_delim]])`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkStrToMap { + inner: DataFusionStrToMap, +} + +impl Default for SparkStrToMap { + fn default() -> Self { + Self::new() + } +} + +impl SparkStrToMap { + pub fn new() -> Self { + Self { + inner: DataFusionStrToMap::new(), + } + } +} + +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)) + } +} + +/// Materializes scalar arguments so the validation below indexes rows the same way the kernel +/// does. `make_scalar_function` inside the kernel expands them anyway, so this only moves that +/// work earlier. +fn expand_scalars(mut args: ScalarFunctionArgs) -> Result<ScalarFunctionArgs> { + let number_rows = args.number_rows; + for arg in args.args.iter_mut() { + if let ColumnarValue::Scalar(scalar) = arg { + *arg = ColumnarValue::Array(scalar.to_array_of_size(number_rows)?); + } + } + Ok(args) +} + +/// Rejects the inputs Spark's `MapFromArrays` rejects before building the map: a row whose key +/// and value arrays differ in length, and a `NULL` key element. +fn validate_map_from_arrays(keys: &ArrayRef, values: &ArrayRef) -> Result<()> { + // A `NULL`-typed argument makes every row a NULL map, which never reaches the builder. + if matches!(keys.data_type(), DataType::Null) || matches!(values.data_type(), DataType::Null) { + return Ok(()); + } + let (flat_keys, key_offsets) = list_values_and_offsets(keys)?; + let (_, value_offsets) = list_values_and_offsets(values)?; + if key_offsets.len() != value_offsets.len() { + return exec_err!("map_from_arrays: keys and values must have the same number of rows"); + } + let key_nulls = element_validity(&flat_keys); + + for row in 0..key_offsets.len().saturating_sub(1) { + // `MapFromArrays` is null intolerant, so a NULL input array yields a NULL map without + // evaluating the builder. + if !keys.is_valid(row) || !values.is_valid(row) { + continue; + } + let (start, end) = (key_offsets[row], key_offsets[row + 1]); + if end - start != value_offsets[row + 1] - value_offsets[row] { + return Err(SparkError::MapKeyValueDiffSizes.into()); + } + if let Some(nulls) = &key_nulls { + if nulls.slice(start, end - start).null_count() > 0 { + return Err(SparkError::NullMapKey.into()); + } Review Comment: The runtime ordering is fixed at `3085702`: `[1, 1, NULL]` reports `DUPLICATED_MAP_KEY` in both builders, and the NULL-first and `LAST_WIN` regressions also pass. The two map-constructor sections in `expression-audits/map_funcs.md` still say "ahead of any duplicate-key check"; please update those to match the corrected behavior. ########## spark/src/main/scala/org/apache/comet/serde/maps.scala: ########## @@ -132,79 +133,86 @@ object CometMapExtract extends CometExpressionSerde[GetMapValue] { } } -private object MapKeyDedupPolicySupport { - val incompatibleReason: String = - s"`${SQLConf.MAP_KEY_DEDUP_POLICY.key}` is set to " + - s"`${SQLConf.MapKeyDedupPolicy.LAST_WIN}`; Comet's native map construction " + - "does not implement LAST_WIN dedup semantics." - - val nullKeyReason: String = - "Spark rejects a `NULL` element inside the keys array with a `RuntimeException`" + - " (`Cannot use null as map key`); Comet's native `map_from_arrays` / `map_from_entries`" + - " does not detect a per-element `NULL` key and produces a map with a `NULL` key instead" + - " ([#4680](https://github.com/apache/datafusion-comet/issues/4680))." - - def isLastWin: Boolean = - SQLConf.get - .getConf(SQLConf.MAP_KEY_DEDUP_POLICY) - .toString - .equalsIgnoreCase(SQLConf.MapKeyDedupPolicy.LAST_WIN.toString) +/** + * Shared gate for the native map constructors (`map_from_arrays`, `map_from_entries`), which + * reproduce Spark's `ArrayBasedMapBuilder`: they reject a `NULL` key with `NULL_MAP_KEY` and + * follow `spark.sql.mapKeyDedupPolicy`, whose value Comet forwards to the native session as + * `datafusion.spark.map_key_dedup_policy`. + */ +private object MapBuilderSupport { + + /** + * Floating-point keys differ from Spark only on 4.0 and later, and differently per function. + * `ArrayBasedMapBuilder` gained `keyNormalizer` in 4.0 (with + * `spark.sql.legacy.disableMapKeyNormalization` to turn it off); 3.4 and 3.5 do not normalize + * at all, so the native builders already match there. + * + * On 4.0+ the normalized key decides duplicates for both functions, so a map built from both + * `-0.0` and `+0.0` is one key in Spark and two natively. What each function stores then + * diverges: `MapFromArrays` calls `ArrayBasedMapBuilder.from`, which returns the input arrays + * untouched when no key repeated, so a lone `-0.0` key stays `-0.0` in Spark too; while + * `MapFromEntries` puts entries one at a time and always calls `build()`, which emits the + * normalized keys, so a lone `-0.0` key comes back as `+0.0` in Spark and as `-0.0` natively. + * + * A note rather than a decline, because a map keyed on `-0.0` or `NaN` is rare; + * `spark.comet.exec.strictFloatingPoint` declines it for anyone who wants the guarantee. That + * gate is not conditioned on the Spark version: declining on 3.4 and 3.5 costs those users + * nothing beyond a fallback they opted into. + */ + val floatingPointKeyNote: String = + "On Spark 4.0 and later, `ArrayBasedMapBuilder` normalizes a floating-point map key before " + + "comparing it, so `-0.0` counts as the same key as `+0.0` and all `NaN`s count as one " + + "key. Comet's native map construction compares the raw Arrow values, so a map built from " + + "both `-0.0` and `+0.0` keeps two entries where Spark reports a duplicate key. " + + "`map_from_entries` also stores the normalized key, so Spark returns `+0.0` for a `-0.0` " + + "key where Comet returns `-0.0`; `map_from_arrays` keeps the original keys in both " + + "engines when nothing repeated. Spark 3.4 and 3.5 do not normalize at all, so they match " + + s"Comet already. Set `${COMET_EXEC_STRICT_FLOATING_POINT.key}=true` to fall back to Spark " + + "for a floating-point map key." + + /** + * `ArrayBasedMapBuilder` keys its dedup map on `TypeUtils.getInterpretedOrdering` once the key + * type contains a string, so under `UTF8_LCASE` the keys `'a'` and `'A'` are one key. The + * native builders compare the raw Arrow bytes and would keep both, missing the duplicate that + * Spark reports (or, under `LAST_WIN`, the overwrite Spark performs). `MapKeySupport` declines + * a collated key for `map_extract` for the same reason. + */ + val collationKeyReason: String = + "Comet's native map construction compares string keys as `UTF8_BINARY`, so it cannot honour " + + "a non-default collation when it looks for a duplicate key." + + /** The support level for a map constructor whose result has key type `keyType`. */ + def keySupport(keyType: DataType): SupportLevel = + if (hasNonDefaultStringCollation(keyType)) { + Incompatible(Some(collationKeyReason)) + } else { + SupportLevel + .strictFloatingPointReason(keyType, "Map construction on a floating-point key") + .map(reason => Incompatible(Some(reason))) + .getOrElse(Compatible(None)) + } } object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] { override def getIncompatibleReasons(): Seq[String] = - Seq(MapKeyDedupPolicySupport.incompatibleReason) + Seq(MapBuilderSupport.collationKeyReason) override def getCompatibleNotes(): Seq[String] = - Seq(MapKeyDedupPolicySupport.nullKeyReason) + Seq(MapBuilderSupport.floatingPointKeyNote) - override def getSupportLevel(expr: MapFromArrays): SupportLevel = { - if (MapKeyDedupPolicySupport.isLastWin) { - Incompatible(Some(MapKeyDedupPolicySupport.incompatibleReason)) - } else { - Compatible(None) - } - } + override def getSupportLevel(expr: MapFromArrays): SupportLevel = + MapBuilderSupport.keySupport(expr.dataType.keyType) override def convert( expr: MapFromArrays, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val keysExpr = exprToProtoInternal(expr.left, inputs, binding) val valuesExpr = exprToProtoInternal(expr.right, inputs, binding) - val keyType = expr.left.dataType.asInstanceOf[ArrayType].elementType - val valueType = expr.right.dataType.asInstanceOf[ArrayType].elementType - val returnType = MapType(keyType = keyType, valueType = valueType) - for { - andBinaryExprProto <- createAndBinaryExpr(expr, inputs, binding) - mapFromArraysExprProto <- scalarFunctionExprToProto("map", keysExpr, valuesExpr) - nullLiteralExprProto <- exprToProtoInternal(Literal(null, returnType), inputs, binding) - } yield { - val caseWhenExprProto = ExprOuterClass.CaseWhen - .newBuilder() - .addWhen(andBinaryExprProto) - .addThen(mapFromArraysExprProto) - .setElseExpr(nullLiteralExprProto) - .build() - ExprOuterClass.Expr - .newBuilder() - .setCaseWhen(caseWhenExprProto) - .build() - } - } - - private def createAndBinaryExpr( - expr: MapFromArrays, - inputs: Seq[Attribute], - binding: Boolean): Option[ExprOuterClass.Expr] = { - createBinaryExpr( - expr, - IsNotNull(expr.left), - IsNotNull(expr.right), - inputs, - binding, - (builder, binaryExpr) => builder.setAnd(binaryExpr)) + // Native `map_from_arrays` is null intolerant like Spark's: a NULL keys or values array + // yields a NULL map for that row, so no CaseWhen guard is needed here. + scalarFunctionExprToProto("map_from_arrays", keysExpr, valuesExpr) Review Comment: Removing this guard makes a query that should return NULL fail. Reproduced on Spark 4.1.3 at `3085702` with a rebuilt native library: ```sql SET spark.sql.ansi.enabled=true; CREATE TABLE t(k ARRAY<INT>, v STRING) USING parquet; INSERT INTO t VALUES (NULL, 'bad'); SELECT map_from_arrays(k, array(CAST(v AS INT))) FROM t; ``` Spark returns NULL; Comet raises `CAST_INVALID_INPUT` from `CometProject`. Spark skips the right child when the keys array is NULL, but `ScalarFunctionExpr` evaluates both arguments before the wrapper can check them. Restoring the previous CASE guard while keeping the new `map_from_arrays` kernel makes this same test pass. Please preserve the short-circuit evaluation and add a column-based regression for a NULL keys array with a values expression that would otherwise throw. -- 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]
