feat: route higher-order functions through codegen dispatcher

Register the array and map higher-order (lambda) functions that previously fell
back to Spark so they stay native via the codegen dispatcher:

- array: transform, exists, forall, aggregate, array_sort (comparator), zip_with
- map: map_filter, transform_keys, transform_values, map_zip_with

These have no native (rust) implementation and extend Spark's CodegenFallback,
which the dispatcher's canHandle already admits, so the projection stays native
and matches Spark exactly. When the dispatcher is disabled they fall back to
Spark. Adds CometHigherOrderFunctionSuite covering the dispatch path and the
fallback-when-disabled path for each family.

Author: andygrove

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

@@ -73,7 +73,13 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     classOf[Flatten] -> CometFlatten,
     classOf[GetArrayItem] -> CometGetArrayItem,
     classOf[Size] -> CometSize,
-    classOf[ArraysZip] -> CometArraysZip)
+    classOf[ArraysZip] -> CometArraysZip,
+    classOf[ArrayTransform] -> CometArrayTransform,
+    classOf[ArrayExists] -> CometArrayExists,
+    classOf[ArrayForAll] -> CometArrayForAll,
+    classOf[ArrayAggregate] -> CometArrayAggregate,
+    classOf[ArraySort] -> CometArraySort,
+    classOf[ZipWith] -> CometZipWith)
 
   private val conditionalExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
     Map(classOf[CaseWhen] -> CometCaseWhen, classOf[If] -> CometIf)
@@ -153,7 +159,11 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     classOf[MapFromArrays] -> CometMapFromArrays,
     classOf[MapContainsKey] -> CometMapContainsKey,
     classOf[MapFromEntries] -> CometMapFromEntries,
-    classOf[StringToMap] -> CometStrToMap)
+    classOf[StringToMap] -> CometStrToMap,
+    classOf[MapFilter] -> CometMapFilter,
+    classOf[TransformKeys] -> CometTransformKeys,
+    classOf[TransformValues] -> CometTransformValues,
+    classOf[MapZipWith] -> CometMapZipWith)
 
   private[comet] val structExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
     Map(

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

@@ -22,7 +22,7 @@ package org.apache.comet.serde
 import scala.annotation.tailrec
 import scala.jdk.CollectionConverters._
 
-import org.apache.spark.sql.catalyst.expressions.{And, ArrayAppend, ArrayContains, ArrayExcept, ArrayFilter, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArrayRepeat, ArraysOverlap, ArraysZip, ArrayUnion, Attribute, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, Literal, Reverse, Size, Slice, SortArray}
+import org.apache.spark.sql.catalyst.expressions.{And, ArrayAggregate, ArrayAppend, ArrayContains, ArrayExcept, ArrayExists, ArrayFilter, ArrayForAll, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArrayRepeat, ArraySort, ArraysOverlap, ArraysZip, ArrayTransform, ArrayUnion, Attribute, Cast, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, Literal, Reverse, Size, Slice, SortArray, ZipWith}
 import org.apache.spark.sql.catalyst.util.GenericArrayData
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
@@ -843,3 +843,19 @@ trait ArraysBase {
     }
   }
 }
+
+// Array higher-order (lambda) functions have no native (rust) implementation. They extend
+// Spark's `CodegenFallback`, which the codegen dispatcher admits: the lambda body evaluates
+// through the kernel's typed Arrow getters, so the projection stays native while matching Spark
+// exactly. See `CometCodegenHOFSuite`.
+object CometArrayTransform extends CometCodegenDispatch[ArrayTransform]
+
+object CometArrayExists extends CometCodegenDispatch[ArrayExists]
+
+object CometArrayForAll extends CometCodegenDispatch[ArrayForAll]
+
+object CometArrayAggregate extends CometCodegenDispatch[ArrayAggregate]
+
+object CometArraySort extends CometCodegenDispatch[ArraySort]
+
+object CometZipWith extends CometCodegenDispatch[ZipWith]

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

@@ -163,3 +163,14 @@ object CometMapFromEntries extends CometScalarFunction[MapFromEntries]("map_from
 }
 
 object CometStrToMap extends CometScalarFunction[StringToMap]("str_to_map")
+
+// Map higher-order (lambda) functions have no native (rust) implementation. They extend Spark's
+// `CodegenFallback`, which the codegen dispatcher admits, so the projection stays native while
+// matching Spark exactly. See `CometCodegenHOFSuite`.
+object CometMapFilter extends CometCodegenDispatch[MapFilter]
+
+object CometTransformKeys extends CometCodegenDispatch[TransformKeys]
+
+object CometTransformValues extends CometCodegenDispatch[TransformValues]
+
+object CometMapZipWith extends CometCodegenDispatch[MapZipWith]

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

@@ -0,0 +1,145 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+
+/**
+ * Higher-order (lambda) functions have no native (rust) implementation. They are wired into the
+ * codegen dispatcher via [[org.apache.comet.serde.CometArrayTransform]] and friends, so a
+ * top-level HOF projection stays native (running Spark's own `doGenCode`/interpreted-eval inside
+ * the Comet kernel) and matches Spark exactly. When the dispatcher is disabled, the HOF has no
+ * native path and the projection falls back to Spark.
+ *
+ * `CometCodegenHOFSuite` covers HOFs nested inside a registered `ScalaUDF`; this suite covers the
+ * HOF as the top-level expression, which only stays native once the serde is registered.
+ */
+class CometHigherOrderFunctionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
+
+  private def withArrayTable(thunk: => Unit): Unit = {
+    withTable("t") {
+      sql("CREATE TABLE t (a ARRAY<INT>, b ARRAY<INT>) USING parquet")
+      sql(
+        "INSERT INTO t VALUES " +
+          "(array(1, 2, 3), array(10, 20, 30)), " +
+          "(array(-5, 5), array(1, 2)), " +
+          "(array(), array()), " +
+          "(null, null)")
+      thunk
+    }
+  }
+
+  private def withMapTable(thunk: => Unit): Unit = {
+    withTable("t") {
+      sql("CREATE TABLE t (m MAP<STRING, INT>, n MAP<STRING, INT>) USING parquet")
+      sql(
+        "INSERT INTO t VALUES " +
+          "(map('a', 1, 'b', 2), map('a', 10, 'c', 30)), " +
+          "(map('x', -1), map('x', 5)), " +
+          "(map(), map()), " +
+          "(null, null)")
+      thunk
+    }
+  }
+
+  test("transform") {
+    withArrayTable {
+      checkSparkAnswerAndOperator("SELECT transform(a, x -> x + 1) FROM t")
+    }
+  }
+
+  test("exists") {
+    withArrayTable {
+      checkSparkAnswerAndOperator("SELECT exists(a, x -> x > 2) FROM t")
+    }
+  }
+
+  test("forall") {
+    withArrayTable {
+      checkSparkAnswerAndOperator("SELECT forall(a, x -> x > 0) FROM t")
+    }
+  }
+
+  test("aggregate") {
+    withArrayTable {
+      checkSparkAnswerAndOperator("SELECT aggregate(a, 0, (acc, x) -> acc + x) FROM t")
+    }
+  }
+
+  test("array_sort with comparator") {
+    withArrayTable {
+      checkSparkAnswerAndOperator(
+        "SELECT array_sort(a, (l, r) -> " +
+          "CASE WHEN l < r THEN 1 WHEN l > r THEN -1 ELSE 0 END) FROM t")
+    }
+  }
+
+  test("zip_with") {
+    withArrayTable {
+      checkSparkAnswerAndOperator("SELECT zip_with(a, b, (x, y) -> x + y) FROM t")
+    }
+  }
+
+  test("map_filter") {
+    withMapTable {
+      checkSparkAnswerAndOperator("SELECT map_filter(m, (k, v) -> v > 1) FROM t")
+    }
+  }
+
+  test("transform_keys") {
+    withMapTable {
+      checkSparkAnswerAndOperator("SELECT transform_keys(m, (k, v) -> concat(k, 'X')) FROM t")
+    }
+  }
+
+  test("transform_values") {
+    withMapTable {
+      checkSparkAnswerAndOperator("SELECT transform_values(m, (k, v) -> v + 1) FROM t")
+    }
+  }
+
+  test("map_zip_with") {
+    withMapTable {
+      checkSparkAnswerAndOperator(
+        "SELECT map_zip_with(m, n, (k, v1, v2) -> coalesce(v1, 0) + coalesce(v2, 0)) FROM t")
+    }
+  }
+
+  test("array HOF falls back to Spark when codegen dispatcher disabled") {
+    withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+      withArrayTable {
+        checkSparkAnswerAndFallbackReason(
+          "SELECT transform(a, x -> x + 1) FROM t",
+          CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key)
+      }
+    }
+  }
+
+  test("map HOF falls back to Spark when codegen dispatcher disabled") {
+    withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+      withMapTable {
+        checkSparkAnswerAndFallbackReason(
+          "SELECT transform_values(m, (k, v) -> v + 1) FROM t",
+          CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key)
+      }
+    }
+  }
+}