perf: vectorize the native map lookup behind element_at and GetMapValue

native/core/src/execution/planner.rs

@@ -707,17 +707,6 @@ impl PhysicalPlanner {
             ExprStruct::ScalarFunc(expr) => {
                 let func = self.create_scalar_function_expr(expr, input_schema);
                 match expr.func.as_ref() {
-                    // DataFusion map_extract returns array of struct entries even if lookup by key
-                    // Apache Spark wants a single value, so wrap the result into additional list extraction
-                    "map_extract" => Ok(Arc::new(ListExtract::new(
-                        func?,
-                        Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
-                        None,
-                        true,
-                        false,
-                        None, // No expr_id for internal map_extract wrapper
-                        Arc::clone(&self.query_context_registry),
-                    ))),
                     // DataFusion 49 hardcodes return type for MD5 built in function as UTF8View
                     // which is not yet supported in Comet
                     // Converting forcibly to UTF8. To be removed after UTF8View supported

native/spark-expr/Cargo.toml

@@ -219,6 +219,10 @@ harness = false
 name = "map_sort"
 harness = false
 
+[[bench]]
+name = "map_extract"
+harness = false
+
 [[bench]]
 name = "to_time"
 harness = false

native/spark-expr/benches/map_extract.rs

@@ -0,0 +1,138 @@
+// 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
+//! <https://github.com/apache/datafusion-comet/issues/5795> stays visible.
+//!
+//! Read the ratio as a measurement of the pinned DataFusion 55.0.0, not as a permanent gap.
+//! DataFusion main has since rewritten `general_map_extract_inner` around a single
+//! `make_comparator` over the batch, so the per-comparison `ArrayRef` slicing that dominates the
+//! baseline here is specific to the version Comet ships today, and the baseline arm will get much
+//! faster at the next DataFusion bump. What survives that bump is the rest of the case for this
+//! kernel: one `eq` plus one `take`, and the `ListExtract` unwrapping pass this removes.
+
+use arrow::array::builder::{MapBuilder, StringBuilder};
+use arrow::array::{ArrayRef, MapFieldNames, StringArray};
+use arrow::datatypes::{DataType, Field};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::ScalarValue;
+use datafusion::functions_nested::map_extract::map_extract_udf;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_comet_spark_expr::SparkMapExtract;
+use std::hint::black_box;
+use std::sync::Arc;
+
+const BATCH_SIZE: usize = 8192;
+/// Distinct keys per map column, as in the issue's `attrs map<string, string>` dataset.
+const DISTINCT_KEYS: usize = 60;
+
+/// `BATCH_SIZE` rows of `map<string, string>`, every tenth row NULL, each non-null row holding
+/// `entries_per_map` entries drawn from `DISTINCT_KEYS` keys. The stride is coprime with
+/// `DISTINCT_KEYS` so a given lookup key lands at a different entry position in every row rather
+/// than always being found (or missed) at the same depth.
+fn string_map(entries_per_map: usize) -> ArrayRef {
+    let mut builder = MapBuilder::new(
+        Some(MapFieldNames {
+            entry: "entries".into(),
+            key: "key".into(),
+            value: "value".into(),
+        }),
+        StringBuilder::new(),
+        StringBuilder::new(),
+    );
+    for row in 0..BATCH_SIZE {
+        if row % 10 == 0 {
+            builder.append(false).unwrap();
+            continue;
+        }
+        for entry in 0..entries_per_map {
+            builder
+                .keys()
+                .append_value(format!("a{}", (row * 13 + entry * 7) % DISTINCT_KEYS));
+            builder.values().append_value(format!("v{}", row % 400));
+        }
+        builder.append(true).unwrap();
+    }
+    Arc::new(builder.finish())
+}
+
+/// One lookup key per row, so the key cannot be hoisted out of the comparison.
+fn per_row_keys() -> ArrayRef {
+    Arc::new(StringArray::from_iter_values(
+        (0..BATCH_SIZE).map(|row| format!("a{}", (row * 29 + 11) % DISTINCT_KEYS)),
+    ))
+}
+
+fn call(udf: &dyn ScalarUDFImpl, args: &[ColumnarValue]) {
+    black_box(
+        udf.invoke_with_args(ScalarFunctionArgs {
+            args: args.to_vec(),
+            arg_fields: vec![],
+            number_rows: BATCH_SIZE,
+            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
+            config_options: Arc::new(ConfigOptions::default()),
+        })
+        .unwrap(),
+    );
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let comet = SparkMapExtract::new();
+    let datafusion = map_extract_udf();
+    let mut group = c.benchmark_group("map_extract");
+
+    for entries in [2usize, 8, 32] {
+        let map = string_map(entries);
+        let cases: [(&str, Vec<ColumnarValue>); 2] = [
+            (
+                "constant_key",
+                vec![
+                    ColumnarValue::Array(Arc::clone(&map)),
+                    ColumnarValue::Scalar(ScalarValue::Utf8(Some("a1".to_string()))),
+                ],
+            ),
+            (
+                "per_row_key",
+                vec![
+                    ColumnarValue::Array(Arc::clone(&map)),
+                    ColumnarValue::Array(per_row_keys()),
+                ],
+            ),
+        ];
+        for (case, args) in cases {
+            group.bench_with_input(
+                BenchmarkId::new(format!("comet/{case}"), entries),
+                &args,
+                |b, args| b.iter(|| call(&comet, args)),
+            );
+            group.bench_with_input(
+                BenchmarkId::new(format!("datafusion/{case}"), entries),
+                &args,
+                |b, args| b.iter(|| call(datafusion.inner().as_ref(), args)),
+            );
+        }
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

native/spark-expr/src/comet_scalar_funcs.rs

@@ -31,7 +31,7 @@ use crate::{
     EvalMode, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains,
     SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkIcebergBucket,
     SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval,
-    SparkMakeTime, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc,
+    SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc,
 };
 use arrow::datatypes::DataType;
 use datafusion::common::{DataFusionError, Result as DataFusionResult};
@@ -321,6 +321,10 @@ 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). It carries the same
+        // `element_at` alias so both registry entries the override replaces point here.
+        Arc::new(ScalarUDF::new_from_impl(SparkMapExtract::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkNextDay::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkSecondsToTimestamp::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkSizeFunc::default())),

native/spark-expr/src/lib.rs

@@ -61,7 +61,7 @@ pub mod jvm_udf;
 mod conditional_funcs;
 mod conversion_funcs;
 mod map_funcs;
-pub use map_funcs::spark_map_sort;
+pub use map_funcs::{spark_map_sort, SparkMapExtract};
 mod math_funcs;
 mod nondetermenistic_funcs;
 pub mod url_funcs;