feat: support lookup map by key

Author: comphead

native/core/src/execution/planner.rs

@@ -555,7 +555,21 @@ impl PhysicalPlanner {
                     fail_on_error,
                 )))
             }
-            ExprStruct::ScalarFunc(expr) => self.create_scalar_function_expr(expr, input_schema),
+            ExprStruct::ScalarFunc(expr) => {
+                let func = self.create_scalar_function_expr(expr, input_schema)?;
+                match expr.func.as_ref() {
+                    // DataFusion map_extract returns array of struct entries even if lookup by key
+                    // Apache Spark waits a single value, so wrap the result into additional list extraction
+                    "map_extract" => Ok(Arc::new(ListExtract::new(
+                        func,
+                        Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
+                        None,
+                        true,
+                        false,
+                    ))),
+                    _ => Ok(func),
+                }
+            }
             ExprStruct::EqNullSafe(expr) => {
                 let left =
                     self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;

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

@@ -1973,6 +1973,10 @@ object QueryPlanSerde extends Logging with CometExprShim {
       case mv: MapValues =>
         val childExpr = exprToProtoInternal(mv.child, inputs, binding)
         scalarFunctionExprToProto("map_values", childExpr)
+      case gmv: GetMapValue =>
+        val mapExpr = exprToProtoInternal(gmv.child, inputs, binding)
+        val keyExpr = exprToProtoInternal(gmv.key, inputs, binding)
+        scalarFunctionExprToProto("map_extract", mapExpr, keyExpr)
       case _ =>
         withInfo(expr, s"${expr.prettyName} is not supported", expr.children: _*)
         None

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

@@ -394,4 +394,46 @@ class CometNativeReaderSuite extends CometTestBase with AdaptiveSparkPlanHelper
       "select * from tbl",
       readSchema = Some(readSchema))
   }
+
+  test("native reader - extract map by key") {
+    // existing key
+    testSingleLineQuery(
+      """
+        | select map(str0, str1) c0 from
+        | (
+        |    select 'key0' str0, named_struct('a', 1, 'b', 'str') str1
+        | )
+        |""".stripMargin,
+      "select c0['key0'] from tbl")
+
+    // existing key, existing struct subfield
+    testSingleLineQuery(
+      """
+        | select map(str0, str1) c0 from
+        | (
+        |    select 'key0' str0, named_struct('a', 1, 'b', 'str') str1
+        | )
+        |""".stripMargin,
+      "select c0['key0'].b from tbl")
+
+    // nonexisting key
+    testSingleLineQuery(
+      """
+        | select map(str0, str1) c0 from
+        | (
+        |    select 'key0' str0, named_struct('a', 1, 'b', 'str') str1
+        | )
+        |""".stripMargin,
+      "select c0['key1'] from tbl")
+
+    // nonexisting key, existing struct subfield
+    testSingleLineQuery(
+      """
+        | select map(str0, str1) c0 from
+        | (
+        |    select 'key0' str0, named_struct('a', 1, 'b', 'str') str1
+        | )
+        |""".stripMargin,
+      "select c0['key1'].b from tbl")
+  }
 }