feat: create map

Author: Kazantsev Maksim

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

@@ -232,6 +232,7 @@ These settings can be used to determine which parts of the plan are accelerated
 | `spark.comet.expression.Cosh.enabled` | Enable Comet acceleration for `Cosh` | true |
 | `spark.comet.expression.Cot.enabled` | Enable Comet acceleration for `Cot` | true |
 | `spark.comet.expression.CreateArray.enabled` | Enable Comet acceleration for `CreateArray` | true |
+| `spark.comet.expression.CreateMap.enabled` | Enable Comet acceleration for `CreateMap` | true |
 | `spark.comet.expression.CreateNamedStruct.enabled` | Enable Comet acceleration for `CreateNamedStruct` | true |
 | `spark.comet.expression.DateAdd.enabled` | Enable Comet acceleration for `DateAdd` | true |
 | `spark.comet.expression.DateFormatClass.enabled` | Enable Comet acceleration for `DateFormatClass` | true |

native/core/src/execution/jni_api.rs

@@ -46,6 +46,7 @@ use datafusion_spark::function::datetime::date_add::SparkDateAdd;
 use datafusion_spark::function::datetime::date_sub::SparkDateSub;
 use datafusion_spark::function::hash::sha1::SparkSha1;
 use datafusion_spark::function::hash::sha2::SparkSha2;
+use datafusion_spark::function::map::map_from_arrays::MapFromArrays;
 use datafusion_spark::function::math::expm1::SparkExpm1;
 use datafusion_spark::function::math::hex::SparkHex;
 use datafusion_spark::function::string::char::CharFunc;
@@ -67,7 +68,6 @@ use std::collections::HashMap;
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
 use std::{sync::Arc, task::Poll};
-use datafusion_spark::function::map::map_from_arrays::MapFromArrays;
 use tokio::runtime::Runtime;
 
 use crate::execution::memory_pools::{

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

@@ -94,9 +94,18 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] {
   }
 }
 
-object CometCreateMap extends CometExpressionSerde[CreateMap] {
+object CometCreateMap extends CometExpressionSerde[CreateMap] with MapBase {
+  val keyUnsupportedReason = "Using BinaryType as Map keys is not allowed in create map function"
+  val valueUnsupportedReason =
+    "Using BinaryType as Map values is not allowed in create map function"
 
   override def getSupportLevel(expr: CreateMap): SupportLevel = {
+    if (containsBinary(expr.dataType.keyType)) {
+      return Incompatible(Some(keyUnsupportedReason))
+    }
+    if (containsBinary(expr.dataType.valueType)) {
+      return Incompatible(Some(valueUnsupportedReason))
+    }
     Compatible(None)
   }
 
@@ -121,7 +130,7 @@ object CometCreateMap extends CometExpressionSerde[CreateMap] {
 
 sealed trait MapBase {
 
-  def containsBinary(dataType: DataType): Boolean = {
+  protected def containsBinary(dataType: DataType): Boolean = {
     dataType match {
       case BinaryType => true
       case StructType(fields) => fields.exists(field => containsBinary(field.dataType))

spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala

@@ -27,7 +27,8 @@ import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.BinaryType
 
-import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions}
+import org.apache.comet.serde.CometCreateMap
+import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions}
 
 class CometMapExpressionSuite extends CometTestBase {
 
@@ -165,7 +166,10 @@ class CometMapExpressionSuite extends CometTestBase {
       val random = new Random(42)
       withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
         val schemaGenOptions =
-          SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false)
+          SchemaGenOptions(
+            generateArray = true,
+            generateStruct = true,
+            primitiveTypes = SchemaGenOptions.defaultPrimitiveTypes.filterNot(_ == BinaryType))
         val dataGenOptions = DataGenOptions(allowNull = false, generateNegativeZero = false)
         ParquetGenerator.makeParquetFile(
           random,
@@ -177,10 +181,23 @@ class CometMapExpressionSuite extends CometTestBase {
       }
       val df = spark.read.parquet(filename)
       df.createOrReplaceTempView("t1")
-      for (fieldName <- df.schema.filter(_.dataType != BinaryType).map(_.name)) {
+      for (fieldName <- df.schema.fieldNames) {
         checkSparkAnswerAndOperator(spark.sql(s"SELECT map($fieldName, $fieldName) FROM t1"))
       }
     }
   }
 
+  test("create_map - fallback for binary type") {
+    val table = "t2"
+    withTable(table) {
+      sql(
+        s"create table $table using parquet as select cast('abc' as binary) as c1 from range(10)")
+      checkSparkAnswerAndFallbackReason(
+        sql(s"select map(c1, 1) from $table"),
+        CometCreateMap.keyUnsupportedReason)
+      checkSparkAnswerAndFallbackReason(
+        sql(s"select map(1, c1) from $table"),
+        CometCreateMap.valueUnsupportedReason)
+    }
+  }
 }

spark/src/test/scala/org/apache/spark/sql/benchmark/CometMapExpressionBenchmark.scala

@@ -0,0 +1,79 @@
+/*
+ * 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.spark.sql.benchmark
+
+import org.apache.comet.CometConf
+
+/**
+ * Configuration for a map expression benchmark.
+ *
+ * @param name
+ *   Name for the benchmark
+ * @param query
+ *   SQL query to benchmark
+ * @param extraCometConfigs
+ *   Additional Comet configurations for the scan+exec case
+ */
+case class MapExprConfig(
+    name: String,
+    query: String,
+    extraCometConfigs: Map[String, String] = Map.empty)
+
+/**
+ * Benchmark to measure performance of Comet map expressions. To run this benchmark:
+ * {{{
+ *   SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometMapExpressionBenchmark
+ * }}}
+ * Results will be written to "spark/benchmarks/CometMapExpressionBenchmark-**results.txt".
+ */
+object CometMapExpressionBenchmark extends CometBenchmarkBase {
+
+  private val mapExpressions = List(
+    MapExprConfig(
+      "create_map",
+      "select map(c1, c1, c2, c2, c3, c3, c4, c4, c5, c5) from parquetV1Table"))
+
+  override def runCometBenchmark(args: Array[String]): Unit = {
+    runBenchmarkWithTable("Map expressions", 1024) { v =>
+      withTempPath { dir =>
+        withTempTable("parquetV1Table") {
+          prepareTable(
+            dir,
+            spark.sql(
+              s"SELECT " +
+                s"(value + 0) AS C1, " +
+                s"(value + 10) AS C2, " +
+                s"(value + 20) AS C3, " +
+                s"(value + 30) AS C4, " +
+                s"(value + 40) AS C5 FROM $tbl"))
+
+          val extraConfigs = Map(CometConf.COMET_CASE_CONVERSION_ENABLED.key -> "true")
+
+          mapExpressions.foreach { config =>
+            val allConfigs = extraConfigs ++ config.extraCometConfigs
+            runBenchmark(config.name) {
+              runExpressionBenchmark(config.name, v, config.query, allConfigs)
+            }
+          }
+        }
+      }
+    }
+  }
+}