feat: add MAP type support for first level (#1603)

* feat: add MAP type support for first level

Author: comphead

common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala

@@ -277,7 +277,8 @@ object Utils {
       case v @ (_: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector |
           _: BigIntVector | _: Float4Vector | _: Float8Vector | _: VarCharVector |
           _: DecimalVector | _: DateDayVector | _: TimeStampMicroTZVector | _: VarBinaryVector |
-          _: FixedSizeBinaryVector | _: TimeStampMicroVector | _: StructVector | _: ListVector) =>
+          _: FixedSizeBinaryVector | _: TimeStampMicroVector | _: StructVector | _: ListVector |
+          _: MapVector) =>
         v.asInstanceOf[FieldVector]
       case _ =>
         throw new SparkException(s"Unsupported Arrow Vector for $reason: ${valueVector.getClass}")

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -70,6 +70,8 @@ object QueryPlanSerde extends Logging with CometExprShim {
       s.fields.map(_.dataType).forall(supportedDataType(_, allowComplex))
     case a: ArrayType if allowComplex =>
       supportedDataType(a.elementType, allowComplex)
+    case m: MapType if allowComplex =>
+      supportedDataType(m.keyType, allowComplex) && supportedDataType(m.valueType, allowComplex)
     case dt =>
       emitWarning(s"unsupported Spark data type: $dt")
       false

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

@@ -230,10 +230,10 @@ object CometNativeScanExec extends DataTypeSupport {
   }
 
   override def isAdditionallySupported(dt: DataType): Boolean = {
-    // TODO add map
     dt match {
       case s: StructType => s.fields.map(_.dataType).forall(isTypeSupported)
       case a: ArrayType => isTypeSupported(a.elementType)
+      case m: MapType => isTypeSupported(m.keyType) && isTypeSupported(m.valueType)
       case _ => false
     }
   }

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

@@ -480,10 +480,10 @@ object CometScanExec extends DataTypeSupport {
 
   override def isAdditionallySupported(dt: DataType): Boolean = {
     if (CometConf.COMET_NATIVE_SCAN_IMPL.get() == CometConf.SCAN_NATIVE_ICEBERG_COMPAT) {
-      // TODO add map
       dt match {
         case s: StructType => s.fields.map(_.dataType).forall(isTypeSupported)
         case a: ArrayType => isTypeSupported(a.elementType)
+        case m: MapType => isTypeSupported(m.keyType) && isTypeSupported(m.valueType)
         case _ => false
       }
     } else {

spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala

@@ -134,4 +134,13 @@ class CometNativeReaderSuite extends CometTestBase with AdaptiveSparkPlanHelper
         |""".stripMargin,
       "select c0 from tbl")
   }
+
+  test("native reader - read simple MAP fields") {
+    testSingleLineQuery(
+      """
+        |select map('a', 1) as c0 union all
+        |select map('b', 2)
+        |""".stripMargin,
+      "select c0 from tbl")
+  }
 }

spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala

@@ -107,19 +107,19 @@ abstract class ParquetReadSuite extends CometTestBase {
         Seq(
           StructField("f1", DecimalType.SYSTEM_DEFAULT),
           StructField("f2", StringType))) -> usingNativeIcebergCompat,
-      MapType(keyType = LongType, valueType = DateType) -> false,
+      MapType(keyType = LongType, valueType = DateType) -> usingNativeIcebergCompat,
       StructType(
         Seq(
           StructField("f1", ByteType),
           StructField("f2", StringType))) -> usingNativeIcebergCompat,
-      MapType(keyType = IntegerType, valueType = BinaryType) -> false).foreach {
-      case (dt, expected) =>
+      MapType(keyType = IntegerType, valueType = BinaryType) -> usingNativeIcebergCompat)
+      .foreach { case (dt, expected) =>
         assert(CometScanExec.isTypeSupported(dt) == expected)
         // usingDataFusionParquetExec does not support CometBatchScanExec yet
         if (!usingDataFusionParquetExec(conf)) {
           assert(CometBatchScanExec.isTypeSupported(dt) == expected)
         }
-    }
+      }
   }
 
   test("unsupported Spark schema") {
@@ -131,7 +131,7 @@ abstract class ParquetReadSuite extends CometTestBase {
     val cometScanExecSupported =
       if (sys.env.get("COMET_PARQUET_SCAN_IMPL").contains(SCAN_NATIVE_ICEBERG_COMPAT) && this
           .isInstanceOf[ParquetReadV1Suite])
-        Seq(true, true, false)
+        Seq(true, true, true)
       else Seq(true, false, false)
 
     val cometBatchScanExecSupported = Seq(true, false, false)