feat: support map inputs for explode

docs/source/contributor-guide/expression-audits/generator_funcs.md

@@ -23,18 +23,18 @@
 
 ## explode
 
-- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), not via the expression serde maps, so it is not auto-detected by the function-registry checkbox logic. Compatible for array inputs; map inputs fall back ([#2837](https://github.com/apache/datafusion-comet/issues/2837)).
+- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), not via the expression serde maps, so it is not auto-detected by the function-registry checkbox logic. Compatible for array and map inputs. Maps emit key/value columns, with a position column for `posexplode`.
 
 ## explode_outer
 
-- Same `CometExplodeExec` path as `explode`. Compatible for array inputs; empty and NULL arrays both emit one null-valued row per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. Map inputs fall back.
+- Same `CometExplodeExec` path as `explode`. Compatible for array and map inputs; empty and NULL collections both emit one null-valued row per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. Map entries reuse the list unnest path and expand into key/value columns.
 
 ## posexplode
 
-- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), like `explode`. Compatible for array inputs; map inputs fall back ([#2837](https://github.com/apache/datafusion-comet/issues/2837)).
+- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), like `explode`. Compatible for array and map inputs. Maps emit key/value columns, with a position column for `posexplode`.
 
 ## posexplode_outer
 
-- Same `CometExplodeExec` path as `posexplode`. Compatible for array inputs; empty and NULL arrays both emit one row with null `pos` and null `value` per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`.
+- Same `CometExplodeExec` path as `posexplode`. Compatible for array and map inputs; empty and NULL collections both emit one row with null `pos` and null generated columns per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`.
 
 [Spark Expression Support]: ../../user-guide/latest/expressions.md

docs/source/user-guide/latest/expressions.md

@@ -337,8 +337,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
 ## generator_funcs
 
 `explode`, `explode_outer`, `posexplode`, and `posexplode_outer` are supported via
-`CometExplodeExec` (operator-level, not expression-level) for array input; map input falls back
-to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via
+`CometExplodeExec` for array and map inputs (operator-level, not expression-level). Enabled by default via
 `spark.comet.exec.explode.enabled`.
 
 | Function | Status | Implementation | Notes |

docs/source/user-guide/latest/operators.md

@@ -124,12 +124,15 @@ combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinEx
 
 ## Generators and set operations
 
-| Operator       | Status | Notes                                                                                                            |
-| -------------- | ------ | ---------------------------------------------------------------------------------------------------------------- |
-| `GenerateExec` | ✅     | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays. `inline` / `stack` fall back. |
-| `ExpandExec`   | ✅     |                                                                                                                  |
-| `UnionExec`    | ✅     |                                                                                                                  |
-| `CoalesceExec` | ✅     |                                                                                                                  |
+| Operator       | Status | Notes                                                                                                                     |
+| -------------- | ------ | ------------------------------------------------------------------------------------------------------------------------- |
+| `GenerateExec` | ✅     | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays and maps. `inline` / `stack` fall back. |
+| `ExpandExec`   | ✅     |                                                                                                                           |
+| `UnionExec`    | ✅     |                                                                                                                           |
+| `CoalesceExec` | ✅     |                                                                                                                           |
+
+Chained array and map generators run natively. They materialize carried columns between generators,
+so their latency can exceed Spark's fused generators on a single task with idle cores.
 
 ## Writes
 

native/core/src/execution/expressions/map_entries.rs

@@ -0,0 +1,119 @@
+// 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 std::fmt::{Display, Formatter};
+use std::hash::{Hash, Hasher};
+use std::sync::Arc;
+
+use arrow::array::{Array, ListArray, MapArray, RecordBatch};
+use arrow::datatypes::{DataType, FieldRef, Schema};
+use datafusion::common::{exec_err, Result as DataFusionResult};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_plan::ColumnarValue;
+
+/// Exposes a map's entries as a list for `ExplodeExec`, sharing the input buffers.
+/// Preserve the original entry fields, including nullability and metadata; the
+/// SQL `map_entries` function rebuilds those fields and loses their metadata.
+#[derive(Debug, Clone)]
+pub struct MapEntriesExpr {
+    child: Arc<dyn PhysicalExpr>,
+}
+
+impl MapEntriesExpr {
+    pub fn new(child: Arc<dyn PhysicalExpr>) -> Self {
+        Self { child }
+    }
+}
+
+impl Display for MapEntriesExpr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "map_entries({})", self.child)
+    }
+}
+
+impl PartialEq for MapEntriesExpr {
+    fn eq(&self, other: &Self) -> bool {
+        self.child.eq(&other.child)
+    }
+}
+
+impl Eq for MapEntriesExpr {}
+
+impl Hash for MapEntriesExpr {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.child.hash(state);
+    }
+}
+
+impl PhysicalExpr for MapEntriesExpr {
+    fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        Display::fmt(self, f)
+    }
+
+    fn return_field(&self, input_schema: &Schema) -> DataFusionResult<FieldRef> {
+        let field = self.child.return_field(input_schema)?;
+        let DataType::Map(entries, _) = field.data_type() else {
+            return exec_err!(
+                "MapEntriesExpr expected Map input, got {}",
+                field.data_type()
+            );
+        };
+        Ok(Arc::new(
+            field
+                .as_ref()
+                .clone()
+                .with_data_type(DataType::List(Arc::clone(entries))),
+        ))
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue> {
+        let array = self.child.evaluate(batch)?.into_array(batch.num_rows())?;
+        let Some(map) = array.as_any().downcast_ref::<MapArray>() else {
+            return exec_err!(
+                "MapEntriesExpr expected Map input, got {}",
+                array.data_type()
+            );
+        };
+        let DataType::Map(entries, _) = map.data_type() else {
+            unreachable!("MapArray downcast guarantees DataType::Map");
+        };
+        let list = ListArray::try_new(
+            Arc::clone(entries),
+            map.offsets().clone(),
+            Arc::new(map.entries().clone()),
+            map.nulls().cloned(),
+        )?;
+        Ok(ColumnarValue::Array(Arc::new(list)))
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
+        if children.len() != 1 {
+            return exec_err!(
+                "MapEntriesExpr expects exactly 1 child, got {}",
+                children.len()
+            );
+        }
+        Ok(Arc::new(Self::new(Arc::clone(&children[0]))))
+    }
+}

native/core/src/execution/expressions/mod.rs

@@ -21,6 +21,7 @@ pub mod arithmetic;
 pub mod bitwise;
 pub mod comparison;
 pub mod logical;
+pub mod map_entries;
 pub mod nullcheck;
 pub mod partition;
 pub mod random;

native/core/src/execution/planner.rs

@@ -5094,11 +5094,23 @@ mod tests {
 
     #[tokio::test]
     async fn explode_evaluates_array_once_per_batch() {
+        check_explode_evaluates_collection_once_per_batch(false).await;
+    }
+
+    #[tokio::test]
+    async fn explode_evaluates_map_once_per_batch() {
+        check_explode_evaluates_collection_once_per_batch(true).await;
+    }
+
+    async fn check_explode_evaluates_collection_once_per_batch(map: bool) {
+        use arrow::array::{AsArray, MapArray, StructArray};
         use arrow::datatypes::Int32Type;
         use datafusion::common::tree_node::{Transformed, TreeNode};
         use datafusion::logical_expr::{create_udf, Volatility};
         use datafusion::physical_plan::projection::ProjectionExec;
-        use spark_expression::data_type::{data_type_info::DatatypeStruct, DataTypeInfo, ListInfo};
+        use spark_expression::data_type::{
+            data_type_info::DatatypeStruct, DataTypeInfo, ListInfo, MapInfo,
+        };
 
         let array_type = spark_expression::DataType {
             type_id: 14,
@@ -5117,6 +5129,55 @@ mod tests {
             Some(vec![Some(20)]),
         ])) as ArrayRef;
 
+        let (array_type, arrays) = if map {
+            // Slice away a leading entry to exercise non-zero map offsets.
+            let values = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+                Some(vec![Some(99)]),
+                Some(vec![Some(10), None]),
+                Some(vec![]),
+                None,
+                Some(vec![Some(20)]),
+            ])
+            .slice(1, 4);
+            let fields: Fields = vec![
+                Field::new("key", DataType::Int32, false)
+                    .with_metadata([("PARQUET:field_id".to_string(), "10".to_string())].into()),
+                Field::new("value", DataType::Int32, true)
+                    .with_metadata([("PARQUET:field_id".to_string(), "11".to_string())].into()),
+            ]
+            .into();
+            let entries = StructArray::new(
+                fields.clone(),
+                vec![
+                    Arc::new(Int32Array::from(vec![99, 1, 2, 3])),
+                    Arc::clone(values.values()),
+                ],
+                None,
+            );
+            let maps = MapArray::new(
+                Arc::new(Field::new("entries", DataType::Struct(fields), false)),
+                values.offsets().clone(),
+                entries,
+                values.nulls().cloned(),
+                false,
+            );
+            let map_type = spark_expression::DataType {
+                type_id: 15,
+                type_info: Some(Box::new(DataTypeInfo {
+                    datatype_struct: Some(DatatypeStruct::Map(Box::new(MapInfo {
+                        key_type: Some(Box::new(create_proto_datatype())),
+                        value_type: Some(Box::new(create_proto_datatype())),
+                        value_contains_null: true,
+                        key_field_id: Some(10),
+                        value_field_id: Some(11),
+                    }))),
+                })),
+            };
+            (map_type, Arc::new(maps) as ArrayRef)
+        } else {
+            (array_type, arrays)
+        };
+
         for outer in [false, true] {
             for position in [false, true] {
                 for computed in [false, true] {
@@ -5195,20 +5256,21 @@ mod tests {
                     let results = collect(native_plan.execute(0, task_ctx).unwrap())
                         .await
                         .unwrap();
-                    let context =
-                        format!("outer={outer}, position={position}, computed={computed}");
+                    let context = format!(
+                        "outer={outer}, position={position}, computed={computed}, map={map}"
+                    );
                     assert_eq!(
                         calls.load(Ordering::Relaxed),
                         if computed { 2 } else { 0 },
                         "{context}"
                     );
                     // The array is pre-projected only to share one evaluation between the
-                    // `pos` and `value` references, so only a computed child needs it. `outer`
+                    // `pos` and `value` references, so a computed child or map wrapper needs it. `outer`
                     // does not, since it is now a `UnnestOptions` mode rather than a wrapper
                     // expression around the child.
                     assert_eq!(
                         projections,
-                        1 + usize::from(position && computed),
+                        1 + usize::from(position && (computed || map)),
                         "{context}"
                     );
                     let expected_values = if outer {
@@ -5228,6 +5290,38 @@ mod tests {
                         })
                         .collect();
                     assert_eq!(values, expected_values.repeat(2), "{context}");
+                    if map {
+                        let schema = native_plan.schema();
+                        for (index, id) in [
+                            (schema.fields().len() - 2, "10"),
+                            (schema.fields().len() - 1, "11"),
+                        ] {
+                            assert_eq!(
+                                schema
+                                    .field(index)
+                                    .metadata()
+                                    .get("PARQUET:field_id")
+                                    .map(String::as_str),
+                                Some(id)
+                            );
+                        }
+                        assert_eq!(schema.field(schema.fields().len() - 2).is_nullable(), outer);
+                        let expected_keys = if outer {
+                            vec![Some(1), Some(2), None, None, Some(3)]
+                        } else {
+                            vec![Some(1), Some(2), Some(3)]
+                        };
+                        let keys: Vec<_> = results
+                            .iter()
+                            .flat_map(|batch| {
+                                batch
+                                    .column(batch.num_columns() - 2)
+                                    .as_primitive::<Int32Type>()
+                                    .iter()
+                            })
+                            .collect();
+                        assert_eq!(keys, expected_keys.repeat(2), "{context}");
+                    }
                     if position {
                         let expected_positions = if outer {
                             vec![Some(0), Some(1), None, None, Some(0)]