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 via the `ListEmptyToNullExpr` planner bridge (works around [datafusion#19053](https://github.com/apache/datafusion/issues/19053)). 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 via the `ListEmptyToNullExpr` planner bridge (works around [datafusion#19053](https://github.com/apache/datafusion/issues/19053)). 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 via the `ListEmptyToNullExpr` planner bridge (works around [datafusion#19053](https://github.com/apache/datafusion/issues/19053)).
+- Same `CometExplodeExec` path as `posexplode`. Compatible for array and map inputs; empty and NULL arrays or maps both emit one row with null `pos` and null generated columns per Spark's `outer` semantics via the `ListEmptyToNullExpr` planner bridge (works around [datafusion#19053](https://github.com/apache/datafusion/issues/19053)).
 
 [Spark Expression Support]: ../../user-guide/latest/expressions.md

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

@@ -329,7 +329,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). 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

@@ -108,12 +108,12 @@ omitted from the tables below and may be reconsidered based on demand:
 
 ## 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` | ✅     |                                                                                                                           |
 
 ## 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

@@ -23,6 +23,7 @@ pub mod comparison;
 pub mod list_empty_to_null;
 pub mod list_positions;
 pub mod logical;
+pub mod map_entries;
 pub mod nullcheck;
 pub mod partition;
 pub mod random;

native/core/src/execution/planner.rs

@@ -41,6 +41,7 @@ use crate::execution::operators::{PartitionedRankLimitExec, WindowFnKind};
 use crate::execution::{
     expressions::list_empty_to_null::ListEmptyToNullExpr,
     expressions::list_positions::ListPositionsExpr,
+    expressions::map_entries::MapEntriesExpr,
     expressions::subquery::Subquery,
     operators::{
         CometFilterExec, ExecutionError, ExpandExec, ExplodeExec, ParquetCompression,
@@ -2057,7 +2058,7 @@ impl PhysicalPlanner {
                 let (scans, shuffle_scans, child) =
                     self.create_plan(&children[0], inputs, partition_count)?;
 
-                // Create the expression for the array to explode
+                // Create the expression for the collection to explode
                 let raw_child_expr = if let Some(child_expr) = &explode.child {
                     self.create_expr(child_expr, child.schema())?
                 } else {
@@ -2073,6 +2074,22 @@ impl PhysicalPlanner {
                     .name()
                     .to_string();
 
+                // Expose maps as List<Struct<key, value>>, sharing their entries buffers.
+                // Reuse the list path for outer rows, positions and bounded output batches,
+                // then flatten the entry struct into Spark's two output columns.
+                let map_fields = match raw_child_expr.data_type(&child_schema)? {
+                    DataType::Map(entries, _) => match entries.data_type() {
+                        DataType::Struct(fields) => Some(fields.clone()),
+                        _ => unreachable!("Map entries must be a struct"),
+                    },
+                    _ => None,
+                };
+                let raw_child_expr: Arc<dyn PhysicalExpr> = if map_fields.is_some() {
+                    Arc::new(MapEntriesExpr::new(raw_child_expr))
+                } else {
+                    raw_child_expr
+                };
+
                 // Bridge Spark's outer semantics: DataFusion's `UnnestExec` with
                 // `preserve_nulls = true` emits one null row for a NULL list but drops rows
                 // whose list is empty. Spark's `explode_outer`/`posexplode_outer` must emit
@@ -2182,11 +2199,22 @@ impl PhysicalPlanner {
                     }
                 };
 
-                output_fields.push(Field::new(
-                    array_field.name(),
-                    element_type,
-                    true, // Element is nullable after unnesting
-                ));
+                let struct_unnests = if let Some(fields) = map_fields {
+                    output_fields.extend(fields.iter().map(|field| {
+                        field
+                            .as_ref()
+                            .clone()
+                            .with_nullable(explode.outer || field.is_nullable())
+                    }));
+                    vec![array_input_index]
+                } else {
+                    output_fields.push(Field::new(
+                        array_field.name(),
+                        element_type,
+                        true, // Element is nullable after unnesting
+                    ));
+                    vec![]
+                };
 
                 let output_schema = Arc::new(Schema::new(output_fields));
 
@@ -2208,7 +2236,7 @@ impl PhysicalPlanner {
                 let unnest_exec = Arc::new(ExplodeExec::new(
                     project_exec,
                     list_unnests,
-                    vec![], // No struct columns to unnest
+                    struct_unnests,
                     output_schema,
                     unnest_options,
                 )?);
@@ -6277,11 +6305,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,
@@ -6300,6 +6340,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] {
@@ -6387,7 +6476,7 @@ mod tests {
                     );
                     assert_eq!(
                         projections,
-                        1 + usize::from(position && (outer || computed)),
+                        1 + usize::from(position && (outer || computed || map)),
                         "{context}"
                     );
                     let expected_values = if outer {
@@ -6407,6 +6496,23 @@ mod tests {
                         })
                         .collect();
                     assert_eq!(values, expected_values.repeat(2), "{context}");
+                    if map {
+                        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)]

native/proto/src/proto/operator.proto

@@ -942,9 +942,9 @@ message Expand {
 }
 
 message Explode {
-  // The array expression to explode into multiple rows
+  // The array or map expression to explode into multiple rows
   spark.spark_expression.Expr child = 1;
-  // Whether this is explode_outer (produces null row for empty/null arrays)
+  // Whether this is explode_outer (produces null row for empty/null collections)
   bool outer = 2;
   // Expressions for other columns to project alongside the exploded values
   repeated spark.spark_expression.Expr project_list = 3;

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -1682,12 +1682,8 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] {
       return Unsupported(Some(s"Unsupported generator: ${op.generator.nodeName}"))
     }
     op.generator.children.head.dataType match {
-      case _: ArrayType =>
+      case _: ArrayType | _: MapType =>
         Compatible()
-      case _: MapType =>
-        // TODO add support for map types
-        // https://github.com/apache/datafusion-comet/issues/2837
-        Unsupported(Some("Comet only supports explode/explode_outer for arrays, not maps"))
       case other =>
         Unsupported(Some(s"Unsupported data type: $other"))
     }