work

Author: Kazantsev Maksim

native/core/src/execution/planner.rs

@@ -134,6 +134,7 @@ 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.
@@ -676,6 +677,14 @@ 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:?}"))),
         }
     }

native/proto/src/proto/expr.proto

@@ -85,6 +85,7 @@ message Expr {
     Rand randn = 62;
     EmptyExpr spark_partition_id = 63;
     EmptyExpr monotonically_increasing_id = 64;
+    CreateMap create_map = 65;
   }
 }
 
@@ -410,3 +411,8 @@ message ArrayJoin {
 message Rand {
   int64 seed = 1;
 }
+
+message CreateMap {
+  repeated Expr keys = 1;
+  repeated Expr values = 2;
+}

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

@@ -19,10 +19,13 @@
 
 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.comet.serde.QueryPlanSerde.{exprToProto, exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType}
+import org.apache.comet.CometSparkSessionExtensions.withInfo
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType}
 
 object CometMapKeys extends CometExpressionSerde[MapKeys] {
 
@@ -95,23 +98,22 @@ object CometCreateMap extends CometExpressionSerde[CreateMap] {
       expr: CreateMap,
       inputs: Seq[Attribute],
       binding: Boolean): Option[ExprOuterClass.Expr] = {
-    val keys = CreateArray(expr.keys)
-    val values = CreateArray(expr.values)
-    val keysProtoExpr = exprToProtoInternal(keys, inputs, binding)
-    val valuesProtoExpr = exprToProtoInternal(values, inputs, binding)
-    // scalastyle:off println
-    println(keysProtoExpr)
-    // scalastyle:on println line=102 column=4
-    // scalastyle:off println
-    // println(valuesProtoExpr)
-    // scalastyle:on println line=103 column=4
-    val createMapScalarExpr =
-      scalarFunctionExprToProtoWithReturnType(
-        "map",
-        expr.dataType,
-        false,
-        keysProtoExpr,
-        valuesProtoExpr)
-    optExprWithInfo(createMapScalarExpr, expr, expr.children: _*)
+    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
+    }
   }
 }