work

Author: Kazantsev Maksim

native/core/src/execution/jni_api.rs

@@ -67,6 +67,7 @@ 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::{
@@ -339,6 +340,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) {
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkConcat::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkHex::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(MapFromArrays::default()));
 }
 
 /// Prepares arrow arrays for output.

native/core/src/execution/planner.rs

@@ -134,7 +134,6 @@ use num::{BigInt, ToPrimitive};
 use object_store::path::Path;
 use std::cmp::max;
 use std::{collections::HashMap, sync::Arc};
-use datafusion_functions_nested::map::map;
 use url::Url;
 
 // For clippy error on type_complexity.
@@ -677,14 +676,6 @@ impl PhysicalPlanner {
             ExprStruct::MonotonicallyIncreasingId(_) => Ok(Arc::new(
                 MonotonicallyIncreasingId::from_partition_id(self.partition),
             )),
-            ExprStruct::CreateMap(expr) => {
-                let keys = expr.keys.iter().map(|expr| self.create_expr(expr, Arc::clone(&input_schema)))
-                    .collect::<Vec<_>>();
-                let values = expr.values.iter().map(|expr| self.create_expr(expr, Arc::clone(&input_schema)))
-                    .collect::<Result<Vec<_>, _>>()?;
-                let create_map = map(keys, values);
-                Ok(Arc::new(create_map))
-            },
             expr => Err(GeneralError(format!("Not implemented: {expr:?}"))),
         }
     }

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

@@ -22,10 +22,11 @@ package org.apache.comet.serde
 import scala.jdk.CollectionConverters._
 
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.types.{ArrayType, MapType}
+import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, MapType, StructType}
+import org.apache.spark.sql.types.DataTypes.{BinaryType, BooleanType, ByteType, DateType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampNTZType, TimestampType}
 
 import org.apache.comet.CometSparkSessionExtensions.withInfo
-import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType, serializeDataType}
 
 object CometMapKeys extends CometExpressionSerde[MapKeys] {
 
@@ -94,26 +95,38 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] {
 }
 
 object CometCreateMap extends CometExpressionSerde[CreateMap] {
+
+  override def getSupportLevel(expr: CreateMap): SupportLevel = {
+    Compatible(None)
+  }
+
   override def convert(
       expr: CreateMap,
       inputs: Seq[Attribute],
       binding: Boolean): Option[ExprOuterClass.Expr] = {
-    val keysProtoExpr = expr.keys.map(exprToProtoInternal(_, inputs, binding))
-    val valuesProtoExpr = expr.values.map(exprToProtoInternal(_, inputs, binding))
-    if (keysProtoExpr.forall(_.isDefined) && valuesProtoExpr.forall(_.isDefined)) {
-      val createMapProtoExpr = ExprOuterClass.CreateMap
-        .newBuilder()
-        .addAllValues(keysProtoExpr.map(_.get).asJava)
-        .addAllValues(valuesProtoExpr.map(_.get).asJava)
-        .build()
-      Some(
-        ExprOuterClass.Expr
-          .newBuilder()
-          .setCreateMap(createMapProtoExpr)
-          .build())
-    } else {
-      withInfo(expr, expr.children: _*)
-      None
+    val keysArray = CreateArray(expr.keys)
+    val valuesArray = CreateArray(expr.values)
+    val keysExprProto = exprToProtoInternal(keysArray, inputs, binding)
+    val valuesExprProto = exprToProtoInternal(valuesArray, inputs, binding)
+    val createMapExprProto =
+      scalarFunctionExprToProtoWithReturnType(
+        "map_from_arrays",
+        expr.dataType,
+        false,
+        keysExprProto,
+        valuesExprProto)
+    optExprWithInfo(createMapExprProto, expr, expr.children: _*)
+  }
+}
+
+sealed trait MapBase {
+
+  def containsBinary(dataType: DataType): Boolean = {
+    dataType match {
+      case BinaryType => true
+      case StructType(fields) => fields.exists(field => containsBinary(field.dataType))
+      case ArrayType(elementType, _) => containsBinary(elementType)
+      case _ => false
     }
   }
 }

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

@@ -25,6 +25,7 @@ import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.CometTestBase
 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}
 
@@ -125,4 +126,29 @@ class CometMapExpressionSuite extends CometTestBase {
     }
   }
 
+  test("create_map") {
+    withTempDir { dir =>
+      val path = new Path(dir.toURI.toString, "test.parquet")
+      val filename = path.toString
+      val random = new Random(42)
+      withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+        val schemaGenOptions =
+          SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false)
+        val dataGenOptions = DataGenOptions(allowNull = false, generateNegativeZero = false)
+        ParquetGenerator.makeParquetFile(
+          random,
+          spark,
+          filename,
+          100,
+          schemaGenOptions,
+          dataGenOptions)
+      }
+      val df = spark.read.parquet(filename)
+      df.createOrReplaceTempView("t1")
+      for (fieldName <- df.schema.filter(_.dataType != BinaryType).map(_.name)) {
+        checkSparkAnswerAndOperator(spark.sql(s"SELECT map($fieldName, $fieldName) FROM t1"))
+      }
+    }
+  }
+
 }