Integrating map_sort with serdes

Author: rishvin

native/spark-expr/src/comet_scalar_funcs.rs

@@ -34,6 +34,7 @@ use datafusion::physical_plan::ColumnarValue;
 use std::any::Any;
 use std::fmt::Debug;
 use std::sync::Arc;
+use crate::map_funcs::spark_map_sort;
 
 macro_rules! make_comet_scalar_udf {
     ($name:expr, $func:ident, $data_type:ident) => {{
@@ -144,6 +145,10 @@ pub fn create_comet_physical_fun(
             let fail_on_error = fail_on_error.unwrap_or(false);
             make_comet_scalar_udf!("spark_modulo", func, without data_type, fail_on_error)
         }
+        "map_sort" => {
+            let func = Arc::new(spark_map_sort);
+            make_comet_scalar_udf!("spark_map_sort", func, without data_type)
+        }
         _ => registry.udf(fun_name).map_err(|e| {
             DataFusionError::Execution(format!(
                 "Function {fun_name} not found in the registry: {e}",

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

@@ -48,7 +48,7 @@ import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
 import org.apache.comet.CometConf
-import org.apache.comet.CometSparkSessionExtensions.{isCometScan, withInfo}
+import org.apache.comet.CometSparkSessionExtensions.{isCometScan, isSpark40Plus, withInfo}
 import org.apache.comet.expressions._
 import org.apache.comet.objectstore.NativeConfig
 import org.apache.comet.serde.ExprOuterClass.{AggExpr, DataType => ProtoDataType, Expr, ScalarFunc}
@@ -128,6 +128,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
     classOf[MapValues] -> CometMapValues,
     classOf[MapFromArrays] -> CometMapFromArrays,
     classOf[GetMapValue] -> CometMapExtract,
+    classOf[MapSort] -> CometMapSort,
     classOf[GreaterThan] -> CometGreaterThan,
     classOf[GreaterThanOrEqual] -> CometGreaterThanOrEqual,
     classOf[LessThan] -> CometLessThan,
@@ -1953,10 +1954,10 @@ object QueryPlanSerde extends Logging with CometExprShim {
 
         if (groupingExpressions.exists(expr =>
             expr.dataType match {
-              case _: MapType => true
+              case _: MapType if !isSpark40Plus => true
               case _ => false
             })) {
-          withInfo(op, "Grouping on map types is not supported")
+          withInfo(op, "Grouping on map types is not supported below Spark 4.0")
           return None
         }
 

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

@@ -94,3 +94,21 @@ object CometMapFromArrays extends CometExpressionSerde {
     optExprWithInfo(mapFromArraysExpr, expr, expr.children: _*)
   }
 }
+
+object CometMapSort extends CometExpressionSerde {
+
+  override def convert(
+      expr: Expression,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[ExprOuterClass.Expr] = {
+    // scalastyle:off println
+    println("Calling CometMapSort.convert")
+    val mapSortExpr = expr.asInstanceOf[MapSort]
+    val childExpr = exprToProtoInternal(mapSortExpr.child, inputs, binding)
+    val returnType = mapSortExpr.child.dataType
+
+    val mapSortScalarExpr =
+      scalarFunctionExprToProtoWithReturnType("map_sort", returnType, childExpr)
+    optExprWithInfo(mapSortScalarExpr, expr, expr.children: _*)
+  }
+}

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

@@ -31,6 +31,7 @@ import org.apache.spark.sql.functions.{avg, count_distinct, sum}
 import org.apache.spark.sql.internal.SQLConf
 
 import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 import org.apache.comet.testing.{DataGenOptions, ParquetGenerator}
 
 /**
@@ -1515,4 +1516,36 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     sparkPlan.collect { case s: CometHashAggregateExec => s }.size
   }
 
+  test("groupby with map column") {
+    assume(isSpark40Plus, "Groupby on map type is supported in Spark 4.0 and beyond")
+    //    withSQLConf(
+    //      CometConf.COMET_ENABLED.key -> "false",
+    //      CometConf.COMET_EXEC_ENABLED.key -> "false",
+    //      CometConf.COMET_EXPLAIN_FALLBACK_ENABLED.key -> "false",
+    //      CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION
+    //      //  CometConf.COMET_SHUFFLE_MODE.key -> "native"
+    //    )
+    //    {
+    withParquetTable(
+      Seq(
+        (1, Map("a" -> 1, "b" -> 2)),
+        (2, Map("b" -> 2, "a" -> 1)),
+        (3, Map("a" -> 5, "b" -> 6))),
+      "tbl") {
+      withSQLConf(
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_EXPLAIN_FALLBACK_ENABLED.key -> "true",
+        CometConf.COMET_SHUFFLE_MODE.key -> "auto",
+        CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION) {
+        val query = sql("SELECT count(*) AS testing FROM tbl group by _2")
+        println(query.queryExecution.executedPlan)
+        query.show()
+        // checkSparkAnswer(query)
+      }
+      // checkSparkAnswerAndOperator("SELECT _1, SUM(_2['b']) FROM tbl GROUP BY _1")
+    }
+    //   }
+  }
+
 }