@@ -210,15 +210,28 @@ fn criterion_benchmark(c: &mut Criterion) {
 }
 
 fn bench_map_extract(c: &mut Criterion) {
+    // DataFusion's default execution batch size.
+    const ROWS: usize = 8192;
+
     let udf = map_extract_udf();
     let config_options = Arc::new(ConfigOptions::default());
     let mut group = c.benchmark_group("map_extract");
 
-    for (rows, width) in [(1, 0), (1, 1), (1024, 1), (1024, 32)] {
-        for key_type in ["int32", "utf8_view", "struct"] {
+    // Cases are named `{key type}/{lookup}/{rows}x{entries}`. `shuffled` looks
+    // up a key that every row holds at a different position, and `varying`
+    // looks up a different key per row, mixing matches and misses.
+    let shapes: &[(usize, &[&str])] = &[
+        (4, &["last", "shuffled", "missing", "varying"]),
+        (32, &["first", "last", "shuffled", "missing", "varying"]),
+    ];
+    for &(width, lookups) in shapes {
+        for key_type in ["int32", "utf8", "utf8_view", "struct"] {
             let make_keys = |keys: Vec<i32>| -> ArrayRef {
                 match key_type {
                     "int32" => Arc::new(Int32Array::from(keys)),
+                    "utf8" => Arc::new(StringArray::from_iter_values(
+                        keys.iter().map(|key| format!("key_{key:016}")),
+                    )),
                     "utf8_view" => Arc::new(StringViewArray::from_iter_values(
                         keys.iter().map(|key| format!("key_{key:016}")),
                     )),
@@ -229,39 +242,52 @@ fn bench_map_extract(c: &mut Criterion) {
                     _ => unreachable!(),
                 }
             };
-            let keys = make_keys((0..rows).flat_map(|_| 0..width as i32).collect());
-            let entries = StructArray::from(vec![
-                (
-                    Arc::new(Field::new("key", keys.data_type().clone(), false)),
-                    keys,
-                ),
-                (
-                    Arc::new(Field::new("value", DataType::Int32, false)),
-                    Arc::new(Int32Array::from_iter_values(0..(rows * width) as i32))
-                        as ArrayRef,
-                ),
-            ]);
-            let map: ArrayRef = Arc::new(MapArray::new(
-                Arc::new(Field::new("entries", entries.data_type().clone(), false)),
-                OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows)),
-                entries,
-                None,
-                false,
-            ));
-            let lookups: &[&str] = if width <= 1 {
-                &["last"]
-            } else {
-                &["first", "last", "missing", "varying"]
+            // Every row holds the keys `0..width`. With `shuffled`, each
+            // row's entries are rotated by the row number.
+            let make_map = |shuffled: bool| -> ArrayRef {
+                let keys = (0..ROWS)
+                    .flat_map(|row| {
+                        (0..width).map(move |position| {
+                            if shuffled {
+                                ((position + row) % width) as i32
+                            } else {
+                                position as i32
+                            }
+                        })
+                    })
+                    .collect();
+                let keys = make_keys(keys);
+                let entries = StructArray::from(vec![
+                    (
+                        Arc::new(Field::new("key", keys.data_type().clone(), false)),
+                        keys,
+                    ),
+                    (
+                        Arc::new(Field::new("value", DataType::Int32, false)),
+                        Arc::new(Int32Array::from_iter_values(0..(ROWS * width) as i32))
+                            as ArrayRef,
+                    ),
+                ]);
+                Arc::new(MapArray::new(
+                    Arc::new(Field::new("entries", entries.data_type().clone(), false)),
+                    OffsetBuffer::from_lengths(std::iter::repeat_n(width, ROWS)),
+                    entries,
+                    None,
+                    false,
+                ))
             };
+            let map = make_map(false);
+            let shuffled_map = make_map(true);
             for &lookup in lookups {
-                let query_keys = match lookup {
-                    "first" => vec![0],
-                    "last" => vec![width.saturating_sub(1) as i32],
-                    "missing" => vec![width as i32],
-                    // Mix matches and misses with a different lookup key per row.
-                    "varying" => {
-                        (0..rows).map(|row| (row % (width + 1)) as i32).collect()
-                    }
+                let (map, query_keys) = match lookup {
+                    "first" => (&map, vec![0]),
+                    "last" => (&map, vec![width.saturating_sub(1) as i32]),
+                    "shuffled" => (&shuffled_map, vec![0]),
+                    "missing" => (&map, vec![width as i32]),
+                    "varying" => (
+                        &map,
+                        (0..ROWS).map(|row| (row % (width + 1)) as i32).collect(),
+                    ),
                     _ => unreachable!(),
                 };
                 let query_keys = make_keys(query_keys);
@@ -272,7 +298,7 @@ fn bench_map_extract(c: &mut Criterion) {
                         ScalarValue::try_from_array(&query_keys, 0).unwrap(),
                     )
                 };
-                let args = vec![ColumnarValue::Array(Arc::clone(&map)), query_keys];
+                let args = vec![ColumnarValue::Array(Arc::clone(map)), query_keys];
                 let arg_fields = args
                     .iter()
                     .map(|arg| Field::new("arg", arg.data_type(), true).into())
@@ -289,15 +315,15 @@ fn bench_map_extract(c: &mut Criterion) {
                 group.bench_function(
                     BenchmarkId::new(
                         format!("{key_type}/{lookup}"),
-                        format!("{rows}x{width}"),
+                        format!("{ROWS}x{width}"),
                     ),
                     |b| {
                         b.iter(|| {
                             black_box(
                                 udf.invoke_with_args(ScalarFunctionArgs {
                                     args: args.clone(),
                                     arg_fields: arg_fields.clone(),
-                                    number_rows: rows,
+                                    number_rows: ROWS,
                                     return_field: Arc::clone(&return_field),
                                     config_options: Arc::clone(&config_options),
                                 })
 
@@ -17,20 +17,19 @@
 
 //! [`ScalarUDFImpl`] definitions for map_extract functions.
 
-use crate::utils::{get_map_entry_field, make_scalar_function};
-use arrow::array::{
-    Array, ArrayRef, ListArray, MapArray, MutableArrayData, make_array, new_empty_array,
-};
+use crate::utils::get_map_entry_field;
+use arrow::array::{Array, ArrayRef, ListArray, MapArray, UInt32Array};
 use arrow::buffer::OffsetBuffer;
-use arrow::compute::SortOptions;
+use arrow::compute::take;
 use arrow::datatypes::{DataType, Field};
-use arrow_ord::ord::make_comparator;
 use datafusion_common::utils::take_function_args;
 use datafusion_common::{Result, cast::as_map_array, exec_err};
+use datafusion_expr::function::Hint;
 use datafusion_expr::{
     ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
     Volatility,
 };
+use datafusion_functions::utils::{make_scalar_function, map_lookup};
 use datafusion_macros::user_doc;
 use std::sync::Arc;
 
@@ -119,7 +118,11 @@ impl ScalarUDFImpl for MapExtract {
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
-        make_scalar_function(map_extract_inner)(&args.args)
+        // A scalar key is passed through as a single row rather than expanded
+        // to the batch size; the lookup applies it to every map row.
+        make_scalar_function(map_extract_inner, vec![Hint::Pad, Hint::AcceptsSingular])(
+            &args.args,
+        )
     }
 
     fn aliases(&self) -> &[String] {
@@ -149,85 +152,40 @@ fn general_map_extract_inner(
     map_array: &MapArray,
     query_keys_array: &dyn Array,
 ) -> Result<ArrayRef> {
-    let keys = map_array.keys();
-    let values = map_array.values();
-    let field = Arc::new(Field::new_list_field(map_array.value_type().clone(), true));
-    let map_offsets = map_array.value_offsets();
-    if map_offsets.first() == map_offsets.last() {
-        return Ok(Arc::new(ListArray::new(
-            field,
-            OffsetBuffer::new_zeroed(map_array.len()),
-            new_empty_array(values.data_type()),
-            map_array.nulls().cloned(),
-        )));
-    }
-
-    // Compare keys by index using a single comparator for the batch.
-    let compare =
-        make_comparator(keys.as_ref(), query_keys_array, SortOptions::default())?;
-    let mut offsets = Vec::with_capacity(map_array.len() + 1);
-    offsets.push(0_i32);
-
-    let original_data = values.to_data();
-    // There is at most one output value per map row.
-    let mut mutable = MutableArrayData::new(
-        vec![&original_data],
-        false,
-        map_array.len().min(values.len()),
-    );
-
-    for (row_index, offset_window) in map_offsets.windows(2).enumerate() {
-        let start = offset_window[0] as usize;
-        let end = offset_window[1] as usize;
-        let mut offset = offsets[row_index];
-
-        if map_array.is_valid(row_index)
-            && let Some(index) = (start..end).find(|&i| compare(i, row_index).is_eq())
-        {
-            mutable.try_extend(0, index, index + 1)?;
-            offset += 1;
-        }
-
-        // A missing key results in an empty list.
-        offsets.push(offset);
-    }
-
-    let data = mutable.freeze();
-
+    let indices = map_lookup(map_array, query_keys_array)?;
+    // Each matched row contributes one list element. Every other row is an
+    // empty list, or NULL when the map itself is NULL.
+    let lengths = indices.iter().map(|index| usize::from(index.is_some()));
+    let mut matched = Vec::with_capacity(indices.len() - indices.null_count());
+    matched.extend(indices.iter().flatten());
+    let values = take(
+        map_array.values().as_ref(),
+        &UInt32Array::from(matched),
+        None,
+    )?;
     Ok(Arc::new(ListArray::new(
-        field,
-        OffsetBuffer::<i32>::new(offsets.into()),
-        make_array(data),
+        Arc::new(Field::new_list_field(map_array.value_type().clone(), true)),
+        OffsetBuffer::from_lengths(lengths),
+        values,
         map_array.nulls().cloned(),
     )))
 }
 
 fn map_extract_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
     let [map_arg, key_arg] = take_function_args("map_extract", args)?;
-
-    let map_array = match map_arg.data_type() {
-        DataType::Map(_, _) => as_map_array(&map_arg)?,
-        DataType::Null => return Ok(Arc::clone(map_arg)),
-        _ => return exec_err!("The first argument in map_extract must be a map"),
-    };
-
-    let key_type = map_array.key_type();
-
-    if key_type != key_arg.data_type() {
-        return exec_err!(
-            "The key type {} does not match the map key type {}",
-            key_arg.data_type(),
-            key_type
-        );
+    match map_arg.data_type() {
+        DataType::Map(_, _) => {
+            general_map_extract_inner(as_map_array(map_arg.as_ref())?, key_arg.as_ref())
+        }
+        DataType::Null => Ok(Arc::clone(map_arg)),
+        _ => exec_err!("The first argument in map_extract must be a map"),
     }
-
-    general_map_extract_inner(map_array, key_arg)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{Float64Array, Int32Array, StructArray};
+    use arrow::array::{Float64Array, Int32Array, NullArray, StructArray};
     use arrow::buffer::NullBuffer;
     use arrow::datatypes::Int32Type;
 
@@ -299,6 +257,24 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn map_extract_untyped_null_keys() -> Result<()> {
+        let map = make_map(
+            Arc::new(Int32Array::from(vec![1, 2])),
+            vec![10, 20],
+            vec![0, 1, 2],
+            Some(NullBuffer::from(vec![true, false])),
+        );
+        let expected =
+            ListArray::from_iter_primitive::<Int32Type, _, _>([Some(vec![]), None]);
+        // Direct callers can pass untyped NULLs without SQL's key coercion.
+        for len in [1, map.len()] {
+            let result = general_map_extract_inner(&map, &NullArray::new(len))?;
+            assert_eq!(result.as_ref(), &expected);
+        }
+        Ok(())
+    }
+
     #[test]
     fn map_extract_float_keys() -> Result<()> {
         let nan = f64::NAN;