rich7420 commented on code in PR #5854:
URL: https://github.com/apache/datafusion-comet/pull/5854#discussion_r4016898751
##########
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:
Please preserve short-circuit evaluation. With ANSI enabled and a Parquet
row `(k = NULL, v = 'bad')`:
```sql
SELECT map_from_arrays(k, array(CAST(v AS INT))) FROM t;
```
Spark returns NULL; Comet raises `CAST_INVALID_INPUT` because the cast runs
before the NULL check. Reproduced on Spark 4.1.3; restoring the CASE guard
makes it pass. Please add this regression test.
##########
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 fix passes for both builders,
including `LAST_WIN` and different key/value offsets.
##########
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:
Verified at `3085702`: duplicate/NULL error ordering is fixed. Please also
update the two map-constructor sections in `map_funcs.md` that still say "ahead
of any duplicate-key check".
--
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]
