Feat: Support map, map_keys & maps_values

native/core/src/execution/planner.rs

@@ -101,6 +101,8 @@ use datafusion_expr::{
     WindowFunctionDefinition,
 };
 use datafusion_functions_nested::array_has::ArrayHas;
+use datafusion_functions_nested::map_keys::map_keys_udf;
+use datafusion_functions_nested::map_values::map_values_udf;
 use datafusion_physical_expr::expressions::{Literal, StatsType};
 use datafusion_physical_expr::window::WindowExpr;
 use datafusion_physical_expr::LexOrdering;
@@ -765,6 +767,46 @@ impl PhysicalPlanner {
 
                 Ok(Arc::new(case_expr))
             }
+            ExprStruct::MapKeys(expr) => {
+                let map_expr =
+                    self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let return_type =
+                    if let DataType::Map(field, _) = map_expr.data_type(&input_schema)? {
+                        if let DataType::Struct(fields) = field.data_type() {
+                            fields[0].data_type().clone()
+                        } else {
+                            unreachable!()
+                        }
+                    } else {
+                        unreachable!()
+                    };
+                Ok(Arc::new(ScalarFunctionExpr::new(
+                    "map_keys",
+                    map_keys_udf(),
+                    vec![map_expr],
+                    return_type,
+                )))
+            }
+            ExprStruct::MapValues(expr) => {
+                let map_expr =
+                    self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let return_type =
+                    if let DataType::Map(field, _) = map_expr.data_type(&input_schema)? {
+                        if let DataType::Struct(fields) = field.data_type() {
+                            fields[0].data_type().clone()
+                        } else {
+                            unreachable!()
+                        }
+                    } else {
+                        unreachable!()
+                    };
+                Ok(Arc::new(ScalarFunctionExpr::new(
+                    "map_values",
+                    map_values_udf(),
+                    vec![map_expr],
+                    return_type,
+                )))
+            }
             expr => Err(ExecutionError::GeneralError(format!(
                 "Not implemented: {:?}",
                 expr
@@ -1114,7 +1156,6 @@ impl PhysicalPlanner {
                     codec,
                     writer.output_data_file.clone(),
                     writer.output_index_file.clone(),
-                    writer.enable_fast_encoding,
                 )?);
 
                 Ok((
@@ -2079,6 +2120,7 @@ impl PhysicalPlanner {
             .collect::<Result<Vec<_>, _>>()?;
 
         let fun_name = &expr.func;
+
         let input_expr_types = args
             .iter()
             .map(|x| x.data_type(input_schema.as_ref()))

native/proto/src/proto/expr.proto

@@ -86,6 +86,9 @@ message Expr {
     ArrayInsert array_insert = 59;
     BinaryExpr array_contains = 60;
     BinaryExpr array_remove = 61;
+    BinaryExpr make_map = 62;
+    UnaryExpr map_keys = 63;
+    UnaryExpr map_values = 64;
   }
 }
 

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

@@ -2300,6 +2300,36 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
             expr.children(1),
             inputs,
             (builder, binaryExpr) => builder.setArrayAppend(binaryExpr))
+
+        case CreateMap(children, _) =>
+          assert(children.size % 2 == 0)
+
+          val keys = children.indices.filter(_ % 2 == 0).map(children)
+          val values = children.indices.filter(_ % 2 != 0).map(children)
+          assert(keys.size == values.size)
+
+          val keysArray = CreateArray(keys)
+          val valuesArray = CreateArray(values)
+
+          val keysArrayExpr = exprToProtoInternal(keysArray, inputs);
+          val valuesArrayExpr = exprToProtoInternal(valuesArray, inputs);
+
+          if (keysArrayExpr.isDefined && valuesArrayExpr.isDefined) {
+            scalarExprToProto("map", keysArrayExpr, valuesArrayExpr)
+          } else {
+            withInfo(expr, "unsupported arguments for CreateMap", children: _*)
+            None
+          }
+        case _ if expr.prettyName == "map_keys" =>
+          createUnaryExpr(
+            expr.children.head,
+            inputs,
+            (builder, unaryExpr) => builder.setMapKeys(unaryExpr))
+        case _ if expr.prettyName == "map_values" =>
+          createUnaryExpr(
+            expr.children.head,
+            inputs,
+            (builder, unaryExpr) => builder.setMapValues(unaryExpr))
         case _ =>
           withInfo(expr, s"${expr.prettyName} is not supported", expr.children: _*)
           None

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

@@ -2551,4 +2551,19 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       }
     }
   }
+
+  test("CreateMap") {
+    withSQLConf("parquet.enable.dictionary" -> "true") {
+      val table = "test"
+      withTable(table) {
+        sql(s"create table $table(col1 long, col2 int) using parquet")
+        sql(s"insert into $table values(1, 2)")
+        sql(s"insert into $table values(2, 2)")
+        sql(s"insert into $table values(3, 4)")
+        sql(s"insert into $table values(4, 6)")
+
+        checkSparkAnswerAndOperator(s"SELECT map(col1, col2) FROM $table")
+      }
+    }
+  }
 }