feat: Nested map support for columnar shuffle

common/src/main/java/org/apache/comet/vector/CometMapVector.java

@@ -0,0 +1,77 @@
+/*
+ * 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.
+ */
+
+package org.apache.comet.vector;
+
+import org.apache.arrow.vector.*;
+import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.util.TransferPair;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarMap;
+
+/** A Comet column vector for map type. */
+public class CometMapVector extends CometDecodedVector {
+  final MapVector mapVector;
+  final ValueVector dataVector;
+  final CometStructVector dataColumnVector;
+
+  final ColumnVector keys;
+  final ColumnVector values;
+
+  public CometMapVector(ValueVector vector, boolean useDecimal128) {
+    super(vector, vector.getField(), useDecimal128);
+
+    this.mapVector = ((MapVector) vector);
+    this.dataVector = mapVector.getDataVector();
+
+    if (dataVector instanceof StructVector) {
+      this.dataColumnVector = new CometStructVector(dataVector, useDecimal128);
+
+      if (dataColumnVector.children.size() != 2) {
+        throw new RuntimeException(
+            "MapVector's dataVector should have 2 children, but got: "
+                + dataColumnVector.children.size());
+      }
+
+      this.keys = dataColumnVector.getChild(0);
+      this.values = dataColumnVector.getChild(1);
+    } else {
+      throw new RuntimeException(
+          "MapVector's dataVector should be StructVector, but got: "
+              + dataVector.getClass().getSimpleName());
+    }
+  }
+
+  @Override
+  public ColumnarMap getMap(int i) {
+    int start = mapVector.getOffsetBuffer().getInt(i * MapVector.OFFSET_WIDTH);
+    int end = mapVector.getOffsetBuffer().getInt((i + 1) * MapVector.OFFSET_WIDTH);
+
+    return new ColumnarMap(keys, values, start, end - start);
+  }
+
+  @Override
+  public CometVector slice(int offset, int length) {
+    TransferPair tp = this.valueVector.getTransferPair(this.valueVector.getAllocator());
+    tp.splitAndTransfer(offset, length);
+
+    return new CometMapVector(tp.getTo(), useDecimal128);
+  }
+}

common/src/main/java/org/apache/comet/vector/CometVector.java

@@ -25,6 +25,7 @@
 import org.apache.arrow.vector.FixedWidthVector;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.MapVector;
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.dictionary.Dictionary;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
@@ -206,6 +207,8 @@ protected static CometVector getVector(
       ValueVector vector, boolean useDecimal128, DictionaryProvider dictionaryProvider) {
     if (vector instanceof StructVector) {
       return new CometStructVector(vector, useDecimal128);
+    } else if (vector instanceof MapVector) {
+      return new CometMapVector(vector, useDecimal128);
     } else if (vector instanceof ListVector) {
       return new CometListVector(vector, useDecimal128);
     } else {

core/src/execution/proto/expr.proto

@@ -421,15 +421,17 @@ message DataType {
     DATE = 12;
     NULL = 13;
     LIST = 14;
-    STRUCT = 15;
+    MAP = 15;
+    STRUCT = 16;
   }
   DataTypeId type_id = 1;
 
   message DataTypeInfo {
     oneof datatype_struct {
       DecimalInfo decimal = 2;
       ListInfo list = 3;
-      StructInfo struct = 4;
+      MapInfo map = 4;
+      StructInfo struct = 5;
     }
   }
 
@@ -443,6 +445,12 @@ message DataType {
     bool contains_null = 2;
   }
 
+  message MapInfo {
+    DataType key_type = 1;
+    DataType value_type = 2;
+    bool value_contains_null = 3;
+  }
+
   message StructInfo {
     repeated string field_names = 1;
     repeated DataType field_datatypes = 2;

core/src/execution/serde.rs

@@ -29,7 +29,7 @@ use crate::{
     },
 };
 use arrow::datatypes::{DataType as ArrowDataType, TimeUnit};
-use arrow_schema::Field;
+use arrow_schema::{Field, Fields};
 use prost::Message;
 use std::{io::Cursor, sync::Arc};
 
@@ -118,6 +118,34 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
             }
             _ => unreachable!(),
         },
+        DataTypeId::Map => match dt_value
+            .type_info
+            .as_ref()
+            .unwrap()
+            .datatype_struct
+            .as_ref()
+            .unwrap()
+        {
+            DatatypeStruct::Map(info) => {
+                let key_field = Field::new(
+                    "key",
+                    to_arrow_datatype(info.key_type.as_ref().unwrap()),
+                    false,
+                );
+                let value_field = Field::new(
+                    "value",
+                    to_arrow_datatype(info.value_type.as_ref().unwrap()),
+                    info.value_contains_null,
+                );
+                let struct_field = Field::new(
+                    "entries",
+                    ArrowDataType::Struct(Fields::from(vec![key_field, value_field])),
+                    false,
+                );
+                ArrowDataType::Map(Arc::new(struct_field), false)
+            }
+            _ => unreachable!(),
+        },
         DataTypeId::Struct => match dt_value
             .type_info
             .as_ref()