comphead commented on code in PR #5806:
URL: https://github.com/apache/datafusion-comet/pull/5806#discussion_r3981087399
##########
native/spark-expr/src/comet_scalar_funcs.rs:
##########
@@ -321,6 +321,9 @@ fn all_scalar_functions() -> Vec<Arc<ScalarUDF>> {
)),
Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::default())),
Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())),
+ // Overrides datafusion-functions-nested' `map_extract` with a
vectorized lookup that
+ // returns the value itself rather than a one-element list (#5795).
+ Arc::new(ScalarUDF::new_from_impl(SparkMapExtract::default())),
Review Comment:
DF's `MapExtract` declares `aliases: ["element_at"]`, and
`SessionState::register_udf` inserts one entry per alias. After this override,
`udf("map_extract")` returns the Comet kernel but `udf("element_at")` still
returns DF's list-returning one.
Nothing serializes that name today, so it is latent rather than live, but
the registry is now inconsistent and the next `element_at` serde would silently
get the wrong shape. Either add `aliases()` returning `["element_at"]`, or note
here why the alias is deliberately left alone.
##########
native/spark-expr/src/map_funcs/map_extract.rs:
##########
@@ -0,0 +1,632 @@
+// 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 arrow::array::{
+ new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray,
NullBufferBuilder, Scalar,
+ UInt32Array,
+};
+use arrow::buffer::BooleanBuffer;
+use arrow::compute::kernels::cmp::eq;
+use arrow::compute::take;
+use arrow::datatypes::{DataType, FieldRef};
+use datafusion::common::utils::take_function_args;
+use datafusion::common::{exec_err, Result as DataFusionResult};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use std::sync::Arc;
+
+/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(<map>, k)`.
+///
+/// Overrides DataFusion's `map_extract` under the same name, and differs from
it in two ways:
+///
+/// - it returns the matched **value** rather than a one-element list, so
the planner does not
+/// have to unwrap the list with a second `ListExtract` pass (see
`planner.rs`);
+/// - the lookup is vectorized. DataFusion's `general_map_extract_inner`
re-slices the query key
+/// and every candidate key into a fresh `ArrayRef` per comparison and
compares them through
+/// `dyn Array` equality, which made a constant-key lookup roughly 35x
more expensive than any
+/// other Comet map kernel and slower than Spark itself
+/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)).
Here a single Arrow
+/// `eq` covers the whole batch of entries at once, the per-row work is a
bit scan over the
+/// resulting mask, and the values are gathered with one `take`.
+///
+/// Spark's own lookup returns the first entry whose key compares equal, so
the mask scan stops at
+/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup
key all produce
+/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map
overload.
+///
+/// Key types whose Spark equality this cannot reproduce (floating point,
non-default collations,
+/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala`
declines them so the
+/// expression falls back to Spark.
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct SparkMapExtract {
+ signature: Signature,
+}
+
+impl Default for SparkMapExtract {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkMapExtract {
+ pub fn new() -> Self {
+ Self {
+ // `user_defined` so `coerce_types` runs and casts the lookup key
to the map's key
+ // type; Comet's planner applies that coercion to the argument
expression.
+ signature: Signature::user_defined(Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkMapExtract {
+ fn name(&self) -> &str {
+ "map_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) ->
DataFusionResult<DataType> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(map_entry_fields(map_type)?.1.data_type().clone())
+ }
+
+ fn coerce_types(&self, arg_types: &[DataType]) ->
DataFusionResult<Vec<DataType>> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(vec![
+ map_type.clone(),
+ map_entry_fields(map_type)?.0.data_type().clone(),
+ ])
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
DataFusionResult<ColumnarValue> {
+ let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?;
+ spark_map_extract(map_arg, key_arg, args.number_rows)
+ }
+}
+
+/// The `(key, value)` fields of a `Map`'s entry struct.
+fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef,
&FieldRef)> {
+ match map_type {
+ DataType::Map(entries, _) => match entries.data_type() {
+ DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0],
&fields[1])),
+ other => exec_err!("map_extract: map entries must be a two-field
struct, got {other}"),
+ },
+ other => exec_err!("map_extract: the first argument must be a map, got
{other}"),
+ }
+}
+
+/// Look up `key_arg` in each row of `map_arg`, returning the matched value or
`NULL`.
+pub fn spark_map_extract(
+ map_arg: &ColumnarValue,
+ key_arg: &ColumnarValue,
+ number_rows: usize,
+) -> DataFusionResult<ColumnarValue> {
+ let map_ref: ArrayRef = match map_arg {
+ ColumnarValue::Array(array) => Arc::clone(array),
+ ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?,
+ };
+ let Some(map_array) = map_ref.as_any().downcast_ref::<MapArray>() else {
+ return exec_err!(
+ "map_extract: the first argument must be a map, got {}",
+ map_ref.data_type()
+ );
+ };
+
+ let num_rows = map_array.len();
+ let value_type = map_array.value_type();
+
+ // Arrow keeps a sliced `MapArray`'s entries child intact and slices only
the offsets, so the
+ // offsets index the *unsliced* keys/values and the visible entries are
the half-open range
+ // [entries_start, entries_end). Comparing only that window keeps a native
OFFSET from paying
+ // for the entries it skipped.
+ let offsets = map_array.offsets();
+ let entries_start = offsets[0] as usize;
+ let entries_end = offsets[num_rows] as usize;
+ if entries_start == entries_end {
Review Comment:
Argument validation is now data-dependent: this early return runs before the
key-length check (L158) and the key-type check (L223). DF 55.0.0's
`map_extract_inner` validates the key type before touching data, so a mismatch
failed identically on every batch. Now a batch whose maps are all empty or all
NULL returns NULLs, while a later batch with entries errors. Same query,
failure depends on which partition holds data.
Both checks are cheap and neither needs the entries window, so they can be
hoisted above this return.
##########
native/spark-expr/src/map_funcs/mod.rs:
##########
@@ -15,5 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+mod map_extract;
mod map_sort;
+pub use map_extract::{spark_map_extract, SparkMapExtract};
Review Comment:
`spark_map_extract` has no caller outside this module.
`comet_scalar_funcs.rs` and the new bench both import only `SparkMapExtract`.
(`spark_map_sort` is exported because `comet_scalar_funcs.rs` calls it
directly.)
The tests are in the same file, so the free function can stay private and
only the UDF needs exporting.
##########
spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala:
##########
@@ -460,6 +460,34 @@ class CometMapExpressionSuite extends CometTestBase {
}
}
+ // The native lookup compares a whole batch of map entries in one pass and
then reads each row's
+ // window out of the resulting mask. A native OFFSET slices the batch, and
Arrow keeps a sliced
+ // MapArray's original entry offsets, so the visible entries start part way
into the keys child --
+ // the same trap `mapsort` hit below. Reading the mask from index 0 would
answer every row with
+ // some other row's entries.
+ test("element_at on a sliced map reads the visible entries") {
Review Comment:
This does exercise `entries_start != 0` today, but incidentally.
`_2` is a nullable Parquet map column, so under the Spark 4.1 ANSI default
`needsNullGuard` wraps the lookup in `CASE WHEN _2 IS NOT NULL`, and DF's
`CaseExpr` evaluates the THEN branch through `filter_record_batch`. It
preserves the slice only because the predicate is all-true, so arrow picks
`IterationStrategy::All` and returns `values.slice(0, count)`. Add one NULL map
row and the filter compacts the entries, `entries_start` becomes 0, and the
test quietly stops testing what its comment says.
An ANSI-off variant, or `_2['a3']` (`GetMapValue` has no guard) so the
sliced map always reaches the kernel, would pin the intent. This is analysis,
not an observed failure.
##########
native/spark-expr/benches/map_extract.rs:
##########
@@ -0,0 +1,131 @@
+// 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.
+
+//! Benchmarks for the map lookup behind `GetMapValue` and `element_at(<map>,
key)`.
+//!
+//! Each shape is run against both Comet's `SparkMapExtract` and the
+//! `datafusion-functions-nested` `map_extract` it overrides, so the gap that
motivated
Review Comment:
Worth a line here: DF main has already rewritten `general_map_extract_inner`
to a single `make_comparator` over the batch, so the per-comparison slicing is
55.0.0-specific. This gap will narrow a lot at the next DF bump and the
comparison will start measuring something different.
That does not change the case for the Comet kernel (one `eq` plus one
`take`, plus the removed `ListExtract` pass), but readers should not take these
ratios as permanent.
##########
native/spark-expr/src/map_funcs/map_extract.rs:
##########
@@ -0,0 +1,632 @@
+// 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 arrow::array::{
+ new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray,
NullBufferBuilder, Scalar,
+ UInt32Array,
+};
+use arrow::buffer::BooleanBuffer;
+use arrow::compute::kernels::cmp::eq;
+use arrow::compute::take;
+use arrow::datatypes::{DataType, FieldRef};
+use datafusion::common::utils::take_function_args;
+use datafusion::common::{exec_err, Result as DataFusionResult};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use std::sync::Arc;
+
+/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(<map>, k)`.
+///
+/// Overrides DataFusion's `map_extract` under the same name, and differs from
it in two ways:
+///
+/// - it returns the matched **value** rather than a one-element list, so
the planner does not
+/// have to unwrap the list with a second `ListExtract` pass (see
`planner.rs`);
+/// - the lookup is vectorized. DataFusion's `general_map_extract_inner`
re-slices the query key
+/// and every candidate key into a fresh `ArrayRef` per comparison and
compares them through
+/// `dyn Array` equality, which made a constant-key lookup roughly 35x
more expensive than any
+/// other Comet map kernel and slower than Spark itself
+/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)).
Here a single Arrow
+/// `eq` covers the whole batch of entries at once, the per-row work is a
bit scan over the
+/// resulting mask, and the values are gathered with one `take`.
+///
+/// Spark's own lookup returns the first entry whose key compares equal, so
the mask scan stops at
+/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup
key all produce
+/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map
overload.
+///
+/// Key types whose Spark equality this cannot reproduce (floating point,
non-default collations,
+/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala`
declines them so the
+/// expression falls back to Spark.
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct SparkMapExtract {
+ signature: Signature,
+}
+
+impl Default for SparkMapExtract {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkMapExtract {
+ pub fn new() -> Self {
+ Self {
+ // `user_defined` so `coerce_types` runs and casts the lookup key
to the map's key
+ // type; Comet's planner applies that coercion to the argument
expression.
+ signature: Signature::user_defined(Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkMapExtract {
+ fn name(&self) -> &str {
+ "map_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) ->
DataFusionResult<DataType> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(map_entry_fields(map_type)?.1.data_type().clone())
+ }
+
+ fn coerce_types(&self, arg_types: &[DataType]) ->
DataFusionResult<Vec<DataType>> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(vec![
+ map_type.clone(),
+ map_entry_fields(map_type)?.0.data_type().clone(),
+ ])
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
DataFusionResult<ColumnarValue> {
+ let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?;
+ spark_map_extract(map_arg, key_arg, args.number_rows)
+ }
+}
+
+/// The `(key, value)` fields of a `Map`'s entry struct.
+fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef,
&FieldRef)> {
+ match map_type {
+ DataType::Map(entries, _) => match entries.data_type() {
+ DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0],
&fields[1])),
+ other => exec_err!("map_extract: map entries must be a two-field
struct, got {other}"),
+ },
+ other => exec_err!("map_extract: the first argument must be a map, got
{other}"),
+ }
+}
+
+/// Look up `key_arg` in each row of `map_arg`, returning the matched value or
`NULL`.
+pub fn spark_map_extract(
+ map_arg: &ColumnarValue,
+ key_arg: &ColumnarValue,
+ number_rows: usize,
+) -> DataFusionResult<ColumnarValue> {
+ let map_ref: ArrayRef = match map_arg {
+ ColumnarValue::Array(array) => Arc::clone(array),
+ ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?,
+ };
+ let Some(map_array) = map_ref.as_any().downcast_ref::<MapArray>() else {
+ return exec_err!(
+ "map_extract: the first argument must be a map, got {}",
+ map_ref.data_type()
+ );
+ };
+
+ let num_rows = map_array.len();
+ let value_type = map_array.value_type();
+
+ // Arrow keeps a sliced `MapArray`'s entries child intact and slices only
the offsets, so the
+ // offsets index the *unsliced* keys/values and the visible entries are
the half-open range
+ // [entries_start, entries_end). Comparing only that window keeps a native
OFFSET from paying
+ // for the entries it skipped.
+ let offsets = map_array.offsets();
+ let entries_start = offsets[0] as usize;
+ let entries_end = offsets[num_rows] as usize;
+ if entries_start == entries_end {
+ // Every row is empty or NULL, so nothing can match.
+ return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows)));
+ }
+ let window_len = entries_end - entries_start;
+ let keys = map_array.keys().slice(entries_start, window_len);
+
+ let matched = match key_arg {
+ ColumnarValue::Scalar(scalar) => {
+ if scalar.is_null() {
+ // Spark map keys are never NULL, so a NULL lookup key matches
nothing.
+ return Ok(ColumnarValue::Array(new_null_array(value_type,
num_rows)));
+ }
+ let key = scalar.to_array_of_size(1)?;
+ key_match_mask(&keys, &key, true)?
+ }
+ ColumnarValue::Array(key_array) => {
+ if key_array.len() != num_rows {
+ return exec_err!(
+ "map_extract: expected {num_rows} lookup keys, got {}",
+ key_array.len()
+ );
+ }
+ // One vectorized compare needs a lookup key per *entry*, not per
row, so gather each
+ // row's key across that row's entries. Entries in a gap between
two rows (offsets are
+ // only required to be monotonic) keep index 0; the per-row scan
below never reads
+ // those positions.
+ let mut gather = vec![0u32; window_len];
+ for row in 0..num_rows {
+ let start = offsets[row] as usize - entries_start;
+ let end = offsets[row + 1] as usize - entries_start;
+ gather[start..end].fill(row as u32);
+ }
+ let per_entry_key = take(key_array, &UInt32Array::from(gather),
None)?;
+ key_match_mask(&keys, &per_entry_key, false)?
+ }
+ };
+
+ // Gather the first matching entry of each row. Map offsets are `i32`, so
an entry index always
+ // fits in `u32`.
+ //
+ // A NULL map row reads NULL whatever its entries hold. Arrow does not
require a null row's
+ // offset range to be empty, and neither Comet's struct-field helper
(which adds a parent null
+ // mask while preserving the child buffers) nor the UDF execution layer
clears those entries, so
+ // a null row can carry a live `a -> 7` that would otherwise match. Spark
returns NULL for a
+ // NULL map under both ANSI modes, for `element_at` and for `GetMapValue`
alike, and only
+ // `element_at` has a nullable-input guard upstream of this kernel.
+ let map_nulls = map_array.nulls();
+ let mut indices = vec![0u32; num_rows];
+ let mut nulls = NullBufferBuilder::new(num_rows);
+ for row in 0..num_rows {
+ if map_nulls.is_some_and(|n| n.is_null(row)) {
+ nulls.append(false);
+ continue;
+ }
+ let start = offsets[row] as usize - entries_start;
+ let end = offsets[row + 1] as usize - entries_start;
+ let found = (start..end).find(|&i| matched.value(i));
+ if let Some(i) = found {
+ indices[row] = (i + entries_start) as u32;
+ }
+ nulls.append(found.is_some());
+ }
+ let indices = UInt32Array::new(indices.into(), nulls.finish());
+
+ Ok(ColumnarValue::Array(take(
+ map_array.values(),
+ &indices,
+ None,
+ )?))
+}
+
+/// A bit per map entry: set where the stored key equals the lookup key.
`lookup` is either a
+/// length-1 array broadcast over every entry (constant key) or one key per
entry.
+fn key_match_mask(
+ keys: &ArrayRef,
+ lookup: &ArrayRef,
+ lookup_is_scalar: bool,
+) -> DataFusionResult<BooleanBuffer> {
+ // The planner casts the lookup key to the map's declared key type, so a
mismatch here means
+ // the runtime encoding is not the declared one (a dictionary-encoded key
column, say). Reject
+ // it rather than comparing incomparable encodings and reporting every row
as a miss.
+ if keys.data_type() != lookup.data_type() {
+ return exec_err!(
+ "map_extract: lookup key type {} does not match the map key type
{}",
+ lookup.data_type(),
+ keys.data_type()
+ );
+ }
+ let compared = if lookup_is_scalar {
+ eq(keys, &Scalar::new(Arc::clone(lookup)))
+ } else {
+ eq(keys, lookup)
+ };
+ match compared {
+ Ok(mask) => {
+ // A NULL on either side compares as NULL, which is not a match.
+ let (values, nulls) = mask.into_parts();
+ Ok(match nulls {
+ Some(nulls) if nulls.null_count() > 0 => &values &
nulls.inner(),
+ _ => values,
+ })
+ }
+ // `eq` rejects nested key types. `MapKeySupport` declines those
before they reach the
+ // native lookup, but keep DataFusion's element-wise comparison as a
backstop so this
+ // kernel is never less capable than the one it replaces.
+ Err(_) => Ok(elementwise_match_mask(keys, lookup, lookup_is_scalar)),
Review Comment:
This catches any `ArrowError`, not just the nested-type rejection the
comment describes. A future or unrelated `eq` failure silently drops into
`elementwise_match_mask`, which is exactly the per-row slice-and-compare
pattern this PR exists to remove. That is a silent ~50x throughput cliff with
no signal, and an untestable branch.
Deciding from the type up front would keep genuine errors visible:
```rust
if keys.data_type().is_nested() { elementwise_match_mask(..) } else {
eq(..)? }
```
##########
native/spark-expr/src/map_funcs/map_extract.rs:
##########
@@ -0,0 +1,632 @@
+// 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 arrow::array::{
+ new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray,
NullBufferBuilder, Scalar,
+ UInt32Array,
+};
+use arrow::buffer::BooleanBuffer;
+use arrow::compute::kernels::cmp::eq;
+use arrow::compute::take;
+use arrow::datatypes::{DataType, FieldRef};
+use datafusion::common::utils::take_function_args;
+use datafusion::common::{exec_err, Result as DataFusionResult};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use std::sync::Arc;
+
+/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(<map>, k)`.
+///
+/// Overrides DataFusion's `map_extract` under the same name, and differs from
it in two ways:
+///
+/// - it returns the matched **value** rather than a one-element list, so
the planner does not
+/// have to unwrap the list with a second `ListExtract` pass (see
`planner.rs`);
+/// - the lookup is vectorized. DataFusion's `general_map_extract_inner`
re-slices the query key
+/// and every candidate key into a fresh `ArrayRef` per comparison and
compares them through
+/// `dyn Array` equality, which made a constant-key lookup roughly 35x
more expensive than any
+/// other Comet map kernel and slower than Spark itself
+/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)).
Here a single Arrow
+/// `eq` covers the whole batch of entries at once, the per-row work is a
bit scan over the
+/// resulting mask, and the values are gathered with one `take`.
+///
+/// Spark's own lookup returns the first entry whose key compares equal, so
the mask scan stops at
+/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup
key all produce
+/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map
overload.
+///
+/// Key types whose Spark equality this cannot reproduce (floating point,
non-default collations,
+/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala`
declines them so the
+/// expression falls back to Spark.
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct SparkMapExtract {
+ signature: Signature,
+}
+
+impl Default for SparkMapExtract {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkMapExtract {
+ pub fn new() -> Self {
+ Self {
+ // `user_defined` so `coerce_types` runs and casts the lookup key
to the map's key
+ // type; Comet's planner applies that coercion to the argument
expression.
+ signature: Signature::user_defined(Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkMapExtract {
+ fn name(&self) -> &str {
+ "map_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) ->
DataFusionResult<DataType> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(map_entry_fields(map_type)?.1.data_type().clone())
+ }
+
+ fn coerce_types(&self, arg_types: &[DataType]) ->
DataFusionResult<Vec<DataType>> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(vec![
+ map_type.clone(),
+ map_entry_fields(map_type)?.0.data_type().clone(),
+ ])
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
DataFusionResult<ColumnarValue> {
+ let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?;
+ spark_map_extract(map_arg, key_arg, args.number_rows)
+ }
+}
+
+/// The `(key, value)` fields of a `Map`'s entry struct.
+fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef,
&FieldRef)> {
+ match map_type {
+ DataType::Map(entries, _) => match entries.data_type() {
+ DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0],
&fields[1])),
+ other => exec_err!("map_extract: map entries must be a two-field
struct, got {other}"),
+ },
+ other => exec_err!("map_extract: the first argument must be a map, got
{other}"),
+ }
+}
+
+/// Look up `key_arg` in each row of `map_arg`, returning the matched value or
`NULL`.
+pub fn spark_map_extract(
+ map_arg: &ColumnarValue,
+ key_arg: &ColumnarValue,
+ number_rows: usize,
+) -> DataFusionResult<ColumnarValue> {
+ let map_ref: ArrayRef = match map_arg {
+ ColumnarValue::Array(array) => Arc::clone(array),
+ ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?,
+ };
+ let Some(map_array) = map_ref.as_any().downcast_ref::<MapArray>() else {
+ return exec_err!(
+ "map_extract: the first argument must be a map, got {}",
+ map_ref.data_type()
+ );
+ };
+
+ let num_rows = map_array.len();
+ let value_type = map_array.value_type();
+
+ // Arrow keeps a sliced `MapArray`'s entries child intact and slices only
the offsets, so the
+ // offsets index the *unsliced* keys/values and the visible entries are
the half-open range
+ // [entries_start, entries_end). Comparing only that window keeps a native
OFFSET from paying
+ // for the entries it skipped.
+ let offsets = map_array.offsets();
+ let entries_start = offsets[0] as usize;
+ let entries_end = offsets[num_rows] as usize;
+ if entries_start == entries_end {
+ // Every row is empty or NULL, so nothing can match.
+ return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows)));
+ }
+ let window_len = entries_end - entries_start;
+ let keys = map_array.keys().slice(entries_start, window_len);
+
+ let matched = match key_arg {
+ ColumnarValue::Scalar(scalar) => {
+ if scalar.is_null() {
+ // Spark map keys are never NULL, so a NULL lookup key matches
nothing.
+ return Ok(ColumnarValue::Array(new_null_array(value_type,
num_rows)));
+ }
+ let key = scalar.to_array_of_size(1)?;
+ key_match_mask(&keys, &key, true)?
+ }
+ ColumnarValue::Array(key_array) => {
+ if key_array.len() != num_rows {
+ return exec_err!(
+ "map_extract: expected {num_rows} lookup keys, got {}",
+ key_array.len()
+ );
+ }
+ // One vectorized compare needs a lookup key per *entry*, not per
row, so gather each
+ // row's key across that row's entries. Entries in a gap between
two rows (offsets are
+ // only required to be monotonic) keep index 0; the per-row scan
below never reads
+ // those positions.
+ let mut gather = vec![0u32; window_len];
+ for row in 0..num_rows {
+ let start = offsets[row] as usize - entries_start;
+ let end = offsets[row + 1] as usize - entries_start;
+ gather[start..end].fill(row as u32);
+ }
+ let per_entry_key = take(key_array, &UInt32Array::from(gather),
None)?;
+ key_match_mask(&keys, &per_entry_key, false)?
+ }
+ };
+
+ // Gather the first matching entry of each row. Map offsets are `i32`, so
an entry index always
+ // fits in `u32`.
+ //
+ // A NULL map row reads NULL whatever its entries hold. Arrow does not
require a null row's
+ // offset range to be empty, and neither Comet's struct-field helper
(which adds a parent null
+ // mask while preserving the child buffers) nor the UDF execution layer
clears those entries, so
+ // a null row can carry a live `a -> 7` that would otherwise match. Spark
returns NULL for a
+ // NULL map under both ANSI modes, for `element_at` and for `GetMapValue`
alike, and only
+ // `element_at` has a nullable-input guard upstream of this kernel.
+ let map_nulls = map_array.nulls();
Review Comment:
Question about reachability, not about the fix.
`child_with_parent_nulls` only *adds* a mask, it does not create entries, so
this needs a producer that hands Comet a null struct row whose map child
already spans a non-empty range. Everything I checked leaves it empty: the
Parquet readers emit length 0 for a null parent, Spark's
`ArrowWriter.StructWriter.setNull` recurses into children, and arrow's `filter`
/ `take` / `concat` compact.
Keep the guard regardless. DF 55.0.0 has the same hole
(`general_map_extract_inner` never consults `map_array.is_valid` and builds its
`ListArray` with `None` nulls), so this closes a pre-existing latent bug
cheaply, and both tests do fail without it.
The ask is narrower: if you know the producing path, an end-to-end case
(`SELECT s.m['a']` over a nullable `struct<m: map<..>>`) would be much stronger
than a hand-built array. If you do not, consider softening "Comet's
struct-field helper ... leaves the entries in place", which reads as a live
path.
##########
native/spark-expr/src/map_funcs/map_extract.rs:
##########
@@ -0,0 +1,632 @@
+// 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 arrow::array::{
+ new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray,
NullBufferBuilder, Scalar,
+ UInt32Array,
+};
+use arrow::buffer::BooleanBuffer;
+use arrow::compute::kernels::cmp::eq;
+use arrow::compute::take;
+use arrow::datatypes::{DataType, FieldRef};
+use datafusion::common::utils::take_function_args;
+use datafusion::common::{exec_err, Result as DataFusionResult};
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use std::sync::Arc;
+
+/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(<map>, k)`.
+///
+/// Overrides DataFusion's `map_extract` under the same name, and differs from
it in two ways:
+///
+/// - it returns the matched **value** rather than a one-element list, so
the planner does not
+/// have to unwrap the list with a second `ListExtract` pass (see
`planner.rs`);
+/// - the lookup is vectorized. DataFusion's `general_map_extract_inner`
re-slices the query key
+/// and every candidate key into a fresh `ArrayRef` per comparison and
compares them through
+/// `dyn Array` equality, which made a constant-key lookup roughly 35x
more expensive than any
+/// other Comet map kernel and slower than Spark itself
+/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)).
Here a single Arrow
+/// `eq` covers the whole batch of entries at once, the per-row work is a
bit scan over the
+/// resulting mask, and the values are gathered with one `take`.
+///
+/// Spark's own lookup returns the first entry whose key compares equal, so
the mask scan stops at
+/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup
key all produce
+/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map
overload.
+///
+/// Key types whose Spark equality this cannot reproduce (floating point,
non-default collations,
+/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala`
declines them so the
+/// expression falls back to Spark.
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct SparkMapExtract {
+ signature: Signature,
+}
+
+impl Default for SparkMapExtract {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkMapExtract {
+ pub fn new() -> Self {
+ Self {
+ // `user_defined` so `coerce_types` runs and casts the lookup key
to the map's key
+ // type; Comet's planner applies that coercion to the argument
expression.
+ signature: Signature::user_defined(Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkMapExtract {
+ fn name(&self) -> &str {
+ "map_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) ->
DataFusionResult<DataType> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(map_entry_fields(map_type)?.1.data_type().clone())
+ }
+
+ fn coerce_types(&self, arg_types: &[DataType]) ->
DataFusionResult<Vec<DataType>> {
+ let [map_type, _] = take_function_args(self.name(), arg_types)?;
+ Ok(vec![
+ map_type.clone(),
+ map_entry_fields(map_type)?.0.data_type().clone(),
+ ])
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
DataFusionResult<ColumnarValue> {
+ let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?;
+ spark_map_extract(map_arg, key_arg, args.number_rows)
+ }
+}
+
+/// The `(key, value)` fields of a `Map`'s entry struct.
+fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef,
&FieldRef)> {
+ match map_type {
+ DataType::Map(entries, _) => match entries.data_type() {
+ DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0],
&fields[1])),
+ other => exec_err!("map_extract: map entries must be a two-field
struct, got {other}"),
+ },
+ other => exec_err!("map_extract: the first argument must be a map, got
{other}"),
+ }
+}
+
+/// Look up `key_arg` in each row of `map_arg`, returning the matched value or
`NULL`.
+pub fn spark_map_extract(
+ map_arg: &ColumnarValue,
+ key_arg: &ColumnarValue,
+ number_rows: usize,
+) -> DataFusionResult<ColumnarValue> {
+ let map_ref: ArrayRef = match map_arg {
+ ColumnarValue::Array(array) => Arc::clone(array),
+ ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?,
+ };
+ let Some(map_array) = map_ref.as_any().downcast_ref::<MapArray>() else {
+ return exec_err!(
+ "map_extract: the first argument must be a map, got {}",
+ map_ref.data_type()
+ );
+ };
+
+ let num_rows = map_array.len();
+ let value_type = map_array.value_type();
+
+ // Arrow keeps a sliced `MapArray`'s entries child intact and slices only
the offsets, so the
+ // offsets index the *unsliced* keys/values and the visible entries are
the half-open range
+ // [entries_start, entries_end). Comparing only that window keeps a native
OFFSET from paying
+ // for the entries it skipped.
+ let offsets = map_array.offsets();
+ let entries_start = offsets[0] as usize;
+ let entries_end = offsets[num_rows] as usize;
+ if entries_start == entries_end {
+ // Every row is empty or NULL, so nothing can match.
+ return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows)));
+ }
+ let window_len = entries_end - entries_start;
+ let keys = map_array.keys().slice(entries_start, window_len);
+
+ let matched = match key_arg {
+ ColumnarValue::Scalar(scalar) => {
+ if scalar.is_null() {
+ // Spark map keys are never NULL, so a NULL lookup key matches
nothing.
+ return Ok(ColumnarValue::Array(new_null_array(value_type,
num_rows)));
+ }
+ let key = scalar.to_array_of_size(1)?;
+ key_match_mask(&keys, &key, true)?
+ }
+ ColumnarValue::Array(key_array) => {
+ if key_array.len() != num_rows {
+ return exec_err!(
+ "map_extract: expected {num_rows} lookup keys, got {}",
+ key_array.len()
+ );
+ }
+ // One vectorized compare needs a lookup key per *entry*, not per
row, so gather each
+ // row's key across that row's entries. Entries in a gap between
two rows (offsets are
+ // only required to be monotonic) keep index 0; the per-row scan
below never reads
+ // those positions.
+ let mut gather = vec![0u32; window_len];
+ for row in 0..num_rows {
+ let start = offsets[row] as usize - entries_start;
+ let end = offsets[row + 1] as usize - entries_start;
+ gather[start..end].fill(row as u32);
+ }
+ let per_entry_key = take(key_array, &UInt32Array::from(gather),
None)?;
+ key_match_mask(&keys, &per_entry_key, false)?
+ }
+ };
+
+ // Gather the first matching entry of each row. Map offsets are `i32`, so
an entry index always
+ // fits in `u32`.
+ //
+ // A NULL map row reads NULL whatever its entries hold. Arrow does not
require a null row's
+ // offset range to be empty, and neither Comet's struct-field helper
(which adds a parent null
+ // mask while preserving the child buffers) nor the UDF execution layer
clears those entries, so
+ // a null row can carry a live `a -> 7` that would otherwise match. Spark
returns NULL for a
+ // NULL map under both ANSI modes, for `element_at` and for `GetMapValue`
alike, and only
+ // `element_at` has a nullable-input guard upstream of this kernel.
+ let map_nulls = map_array.nulls();
+ let mut indices = vec![0u32; num_rows];
+ let mut nulls = NullBufferBuilder::new(num_rows);
+ for row in 0..num_rows {
+ if map_nulls.is_some_and(|n| n.is_null(row)) {
+ nulls.append(false);
+ continue;
+ }
+ let start = offsets[row] as usize - entries_start;
+ let end = offsets[row + 1] as usize - entries_start;
+ let found = (start..end).find(|&i| matched.value(i));
+ if let Some(i) = found {
+ indices[row] = (i + entries_start) as u32;
+ }
+ nulls.append(found.is_some());
+ }
+ let indices = UInt32Array::new(indices.into(), nulls.finish());
+
+ Ok(ColumnarValue::Array(take(
+ map_array.values(),
+ &indices,
+ None,
+ )?))
+}
+
+/// A bit per map entry: set where the stored key equals the lookup key.
`lookup` is either a
+/// length-1 array broadcast over every entry (constant key) or one key per
entry.
+fn key_match_mask(
+ keys: &ArrayRef,
+ lookup: &ArrayRef,
+ lookup_is_scalar: bool,
+) -> DataFusionResult<BooleanBuffer> {
+ // The planner casts the lookup key to the map's declared key type, so a
mismatch here means
+ // the runtime encoding is not the declared one (a dictionary-encoded key
column, say). Reject
+ // it rather than comparing incomparable encodings and reporting every row
as a miss.
+ if keys.data_type() != lookup.data_type() {
+ return exec_err!(
+ "map_extract: lookup key type {} does not match the map key type
{}",
+ lookup.data_type(),
+ keys.data_type()
+ );
+ }
+ let compared = if lookup_is_scalar {
+ eq(keys, &Scalar::new(Arc::clone(lookup)))
+ } else {
+ eq(keys, lookup)
+ };
+ match compared {
+ Ok(mask) => {
+ // A NULL on either side compares as NULL, which is not a match.
+ let (values, nulls) = mask.into_parts();
+ Ok(match nulls {
+ Some(nulls) if nulls.null_count() > 0 => &values &
nulls.inner(),
+ _ => values,
+ })
+ }
+ // `eq` rejects nested key types. `MapKeySupport` declines those
before they reach the
+ // native lookup, but keep DataFusion's element-wise comparison as a
backstop so this
+ // kernel is never less capable than the one it replaces.
+ Err(_) => Ok(elementwise_match_mask(keys, lookup, lookup_is_scalar)),
+ }
+}
+
+fn elementwise_match_mask(
+ keys: &ArrayRef,
+ lookup: &ArrayRef,
+ lookup_is_scalar: bool,
+) -> BooleanBuffer {
+ let mut builder = BooleanBufferBuilder::new(keys.len());
+ for i in 0..keys.len() {
+ let lookup_row = if lookup_is_scalar { 0 } else { i };
+ builder.append(
+ !lookup.is_null(lookup_row)
+ && keys.slice(i, 1).as_ref() == lookup.slice(lookup_row,
1).as_ref(),
+ );
+ }
+ builder.finish()
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{Int32Array, StringArray, StructArray};
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ use arrow::datatypes::{Field, Fields};
+ use datafusion::common::ScalarValue;
+
+ /// One row of a test map: `None` for a NULL map, otherwise its `(key,
value)` entries.
+ type MapRow<'a> = Option<Vec<(&'a str, Option<i32>)>>;
+
+ /// `{"a": 1, "b": 2}`, `{}`, `{"c": 3, "a": 30}`, NULL, `{"b": NULL}`
+ fn test_map() -> MapArray {
+ map_from(vec![
+ Some(vec![("a", Some(1)), ("b", Some(2))]),
+ Some(vec![]),
+ Some(vec![("c", Some(3)), ("a", Some(30))]),
+ None,
+ Some(vec![("b", None)]),
+ ])
+ }
+
+ fn map_from(rows: Vec<MapRow>) -> MapArray {
+ let mut keys = Vec::new();
+ let mut values = Vec::new();
+ let mut offsets = vec![0i32];
+ let mut nulls = NullBufferBuilder::new(rows.len());
+ for row in &rows {
+ match row {
+ Some(entries) => {
+ for (k, v) in entries {
+ keys.push(*k);
+ values.push(*v);
+ }
+ nulls.append(true);
+ }
+ None => nulls.append(false),
+ }
+ offsets.push(keys.len() as i32);
+ }
+
+ let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
+ let value_field = Arc::new(Field::new("value", DataType::Int32, true));
+ let entries = StructArray::new(
+ Fields::from(vec![Arc::clone(&key_field),
Arc::clone(&value_field)]),
+ vec![
+ Arc::new(StringArray::from(keys)) as ArrayRef,
+ Arc::new(Int32Array::from(values)) as ArrayRef,
+ ],
+ None,
+ );
+ let entries_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(Fields::from(vec![key_field, value_field])),
+ false,
+ ));
+ MapArray::try_new(
+ entries_field,
+ OffsetBuffer::new(offsets.into()),
+ entries,
+ nulls.finish(),
+ false,
+ )
+ .unwrap()
+ }
+
+ fn extract(map: MapArray, key: ColumnarValue) -> Vec<Option<i32>> {
+ let num_rows = map.len();
+ let result = spark_map_extract(&ColumnarValue::Array(Arc::new(map)),
&key, num_rows)
+ .unwrap()
+ .into_array(num_rows)
+ .unwrap();
+ let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
+ (0..result.len())
+ .map(|i| (!result.is_null(i)).then(|| result.value(i)))
+ .collect()
+ }
+
+ fn key(value: &str) -> ColumnarValue {
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some(value.to_string())))
+ }
+
+ #[test]
+ fn constant_key_hit_and_miss() {
+ // A found key returns its value; an empty row, a NULL row, and a row
without the key all
+ // return NULL, as does a row whose stored value is NULL.
+ assert_eq!(
+ extract(test_map(), key("a")),
+ vec![Some(1), None, Some(30), None, None]
+ );
+ assert_eq!(
+ extract(test_map(), key("b")),
+ vec![Some(2), None, None, None, None]
+ );
+ assert_eq!(extract(test_map(), key("zz")), vec![None; 5]);
+ }
+
+ #[test]
+ fn duplicate_keys_return_the_first_match() {
+ // Spark's `GetMapValueUtil` scans entries in order and stops at the
first equal key, so a
+ // map that kept duplicates (EXCEPTION dedup is a write-side check)
resolves to the first.
+ let map = map_from(vec![Some(vec![("a", Some(1)), ("a", Some(2))])]);
+ assert_eq!(extract(map, key("a")), vec![Some(1)]);
+ }
+
+ #[test]
+ fn null_lookup_key_matches_nothing() {
+ assert_eq!(
+ extract(test_map(),
ColumnarValue::Scalar(ScalarValue::Utf8(None))),
+ vec![None; 5]
+ );
+ }
+
+ /// A NULL row whose entries were retained rather than dropped. `map_from`
gives a NULL row an
+ /// empty offset range, which is the shape Arrow's builders produce, but
nothing in the format
+ /// requires it: adding a parent null mask over intact child buffers
leaves the entries in
+ /// place. Such a row must still read NULL, not the value its live entry
holds.
+ fn null_row_with_retained_entries() -> MapArray {
Review Comment:
Minor: the `key_field` / `value_field` / `StructArray` / `entries_field` /
`try_new` scaffolding is repeated four times (here, `map_from`,
`non_string_keys`, `nested_key_falls_back_to_elementwise_comparison`). One
`map_of(keys, values, offsets, nulls) -> MapArray` helper would cut roughly 60
lines and make the three variants read as the variations they are.
--
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]
