fix: support map_keys

Author: comphead

native/core/src/execution/planner.rs

@@ -3082,9 +3082,6 @@ mod tests {
         let (mut scans, datafusion_plan) =
             planner.create_plan(&projection, &mut vec![], 1).unwrap();
 
-        // Feed the data into plan
-        //scans[0].set_input_batch(input_batch);
-
         // Start executing the plan in a separate thread
         // The plan waits for incoming batches and emitting result as input comes
         let mut stream = datafusion_plan.native_plan.execute(0, task_ctx).unwrap();
@@ -3240,11 +3237,11 @@ mod tests {
         select map(named_struct('a', 1, 'b', 'n', 'c', 'x'), named_struct('a', 1, 'b', 'n', 'c', 'x')) m
      */
     #[tokio::test]
-    async fn test_nested_types_map_keys_by_index() -> Result<(), DataFusionError> {
+    async fn test_nested_types_map_keys() -> Result<(), DataFusionError> {
         let session_ctx = SessionContext::new();
 
         // generate test data in the temp folder
-        let test_data = "select map([named_struct('a', 1, 'b', 'n', 'c', 'x')], [named_struct('a', 1, 'b', 'n', 'c', 'x')]) c0";
+        let test_data = "select map([named_struct('a', 1, 'b', 'n', 'c', 'x')], [named_struct('a', 2, 'b', 'm', 'c', 'y')]) c0";
         let tmp_dir = TempDir::new()?;
         let test_path = tmp_dir.path().to_str().unwrap().to_string();
 
@@ -3254,7 +3251,7 @@ mod tests {
             .create_physical_plan()
             .await?;
 
-        // Write parquet file into temp folder
+        // Write a parquet file into temp folder
         session_ctx
             .write_parquet(plan, test_path.clone(), None)
             .await?;
@@ -3334,7 +3331,7 @@ mod tests {
             "+------------------------------+",
             "| c0                           |",
             "+------------------------------+",
-            "| {{b: n}: {a: 1, b: n, c: x}} |",
+            "| {{b: n}: {a: 2, b: m, c: y}} |",
             "+------------------------------+",
         ];
         assert_batches_eq!(expected, &[actual.clone()]);

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

@@ -2047,6 +2047,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
       case mk: MapKeys =>
         val childExpr = exprToProtoInternal(mk.child, inputs, binding)
         scalarFunctionExprToProto("map_keys", childExpr)
+//  commented out because of correctness issue https://github.com/apache/datafusion-comet/issues/1789
 //      case mv: MapValues =>
 //        val childExpr = exprToProtoInternal(mv.child, inputs, binding)
 //        scalarFunctionExprToProto("map_values", childExpr)

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

@@ -304,25 +304,30 @@ class CometNativeReaderSuite extends CometTestBase with AdaptiveSparkPlanHelper
       "select c0[0].a, c0[0].a from tbl")
   }
 
-  test("native reader - select nested field from a complex map key using map_keys") {
+  test("native reader - select nested field from a complex map[struct, struct] using map_keys") {
     testSingleLineQuery(
       """
-        | select map(str0, str0) c0 from
+        | select map(str0, str1) c0 from
         | (
-        |   select named_struct('a', cast(1 as long), 'b', cast(2 as long), 'c', cast(3 as long)) str0 union all
-        |   select named_struct('a', cast(3 as long), 'b', cast(4 as long), 'c', cast(5 as long)) str0
+        |   select named_struct('a', cast(1 as long), 'b', cast(2 as long), 'c', cast(3 as long)) str0,
+        |          named_struct('x', cast(8 as long), 'y', cast(9 as long), 'z', cast(0 as long)) str1 union all
+        |   select named_struct('a', cast(3 as long), 'b', cast(4 as long), 'c', cast(5 as long)) str0,
+        |          named_struct('x', cast(6 as long), 'y', cast(7 as long), 'z', cast(8 as long)) str1
         | )
         |""".stripMargin,
       "select map_keys(c0).b from tbl")
   }
 
-//  test("native reader - select nested field from a complex map key using map_values") {
+//  commented out because of correctness issue https://github.com/apache/datafusion-comet/issues/1789
+//  test("native reader - select nested field from a complex map[struct, struct] using map_values") {
 //    testSingleLineQuery(
 //      """
-//        | select map(str0, str0) c0 from
+//        | select map(str0, str1) c0 from
 //        | (
-//        |   select named_struct('a', cast(1 as long), 'b', cast(2 as long), 'c', cast(3 as long)) str0 union all
-//        |   select named_struct('a', cast(3 as long), 'b', cast(4 as long), 'c', cast(5 as long)) str0
+//        |   select named_struct('a', cast(1 as long), 'b', cast(2 as long), 'c', cast(3 as long)) str0,
+//        |          named_struct('x', cast(8 as long), 'y', cast(9 as long), 'z', cast(0 as long)) str1 union all
+//        |   select named_struct('a', cast(3 as long), 'b', cast(4 as long), 'c', cast(5 as long)) str0,
+//        |          named_struct('x', cast(6 as long), 'y', cast(7 as long), 'z', cast(8 as long)) str1
 //        | )
 //        |""".stripMargin,
 //      "select map_values(c0).b from tbl")