refactor: introduce CometInvokeExpressionSerde trait for generic Invoke dispatch

Author: rafafrdz

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

@@ -0,0 +1,81 @@
+/*
+ * 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.serde
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
+
+/**
+ * Serializer contract for Spark `Invoke` expressions that wrap a private evaluator object.
+ *
+ * In Spark 4.0 several built-in functions (e.g. `parse_url`) became `RuntimeReplaceable` and are
+ * rewritten by the analyser into an `Invoke(evaluator, arg, ...)` node whose first child is a
+ * `Literal` of `ObjectType` holding a private evaluator instance. The Spark expression class that
+ * Comet normally dispatches on (e.g. `ParseUrl`) is therefore never seen at serde time on Spark
+ * 4.0.
+ *
+ * Implementors expose:
+ *   - [[invokeTargetClassName]] - the fully-qualified name of the evaluator class embedded in the
+ *     first `Literal(_, ObjectType)` child of the `Invoke` node. This is the key used by
+ *     [[QueryPlanSerde]] to route the node to the correct handler.
+ *   - [[convertFromInvoke]] - the actual serialization logic, receiving the raw `Invoke`
+ *     expression (unerased, with all children including the evaluator literal).
+ *
+ * To register a new handler, add the object to [[QueryPlanSerde.invokeSerdeByTargetClassName]].
+ */
+trait CometInvokeExpressionSerde[T <: Expression] {
+
+  /**
+   * Fully-qualified class name of the private evaluator object held in the first child
+   * `Literal(_, ObjectType(...))` of the `Invoke` node.
+   *
+   * Example: `"org.apache.spark.sql.catalyst.expressions.url.ParseUrlEvaluator"`
+   */
+  def invokeTargetClassName: String
+
+  /**
+   * Serialize the `Invoke` expression into a Comet proto `Expr`.
+   *
+   * @param expr
+   *   The raw `Invoke` expression node (first child is the evaluator literal).
+   * @param inputs
+   *   Resolved input attributes for the enclosing operator.
+   * @param binding
+   *   Whether attributes are bound (relevant for aggregate expressions).
+   * @return
+   *   `Some(Expr)` on success, `None` if the expression cannot be handled (the implementor is
+   *   responsible for calling `withInfo` to record the reason).
+   */
+  def convertFromInvoke(
+      expr: T,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[ExprOuterClass.Expr]
+
+  /**
+   * Unchecked entry point used by [[QueryPlanSerde]] to dispatch through an existential
+   * `CometInvokeExpressionSerde[_]`. Casts `expr` to `T` before delegating to
+   * [[convertFromInvoke]]. Only call this when you have already verified (via
+   * [[invokeTargetClassName]]) that the handler owns this `Invoke` node.
+   */
+  final def convertFromInvokeUnchecked(
+      expr: Expression,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[ExprOuterClass.Expr] =
+    convertFromInvoke(expr.asInstanceOf[T], inputs, binding)
+}

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

@@ -44,14 +44,13 @@ import org.apache.comet.shims.CometExprShim
  */
 object QueryPlanSerde extends Logging with CometExprShim {
 
-  // Generic serializer contract for Spark Invoke expressions.
-  private type InvokeConverter = (Expression, Seq[Attribute], Boolean) => Option[Expr]
-
-  // Dispatch table keyed by the runtime class name stored in Invoke target ObjectType.
-  private val invokeConvertersByTargetClassName: Map[String, InvokeConverter] = Map(
-    CometParseUrl.invokeTargetClassName ->
-      ((expr: Expression, inputs: Seq[Attribute], binding: Boolean) =>
-        CometParseUrl.convertFromInvoke(expr, inputs, binding)))
+  // Registry of Invoke-expression handlers, keyed by the fully-qualified class name of the
+  // evaluator object embedded in the first Literal(_, ObjectType(...)) child of the Invoke node.
+  // To support a new RuntimeReplaceable expression rewritten to Invoke in Spark 4.0, implement
+  // CometInvokeExpressionSerde and add the object here.
+  private val invokeSerdeByTargetClassName
+      : Map[String, CometInvokeExpressionSerde[_ <: Expression]] =
+    Seq(CometParseUrl).map(s => s.invokeTargetClassName -> s).toMap
 
   // Extracts the target object class name from an Invoke-like expression.
   private def invokeTargetClassName(expr: Expression): Option[String] = {
@@ -63,17 +62,16 @@ object QueryPlanSerde extends Logging with CometExprShim {
     }
   }
 
-  // Routes Invoke expressions to a converter based on the target object class name.
+  // Routes Invoke expressions to a handler based on the target object class name.
   private def convertInvokeExpression(
       expr: Expression,
       inputs: Seq[Attribute],
       binding: Boolean): Option[Expr] =
     for {
-      innerObjectTypeExpression <- invokeTargetClassName(expr)
-      invokeExpr <- invokeConvertersByTargetClassName.get(innerObjectTypeExpression)
-      expression <- invokeExpr(expr, inputs, binding)
-
-    } yield expression
+      targetClassName <- invokeTargetClassName(expr)
+      handler <- invokeSerdeByTargetClassName.get(targetClassName)
+      result <- handler.convertFromInvokeUnchecked(expr, inputs, binding)
+    } yield result
 
   private val arrayExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
     classOf[ArrayAppend] -> CometArrayAppend,

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

@@ -383,55 +383,53 @@ object CometStringSplit extends CometExpressionSerde[StringSplit] {
   }
 }
 
-object CometParseUrl extends CometExpressionSerde[ParseUrl] {
-
-  // Class name of the Spark 4.0 internal evaluator embedded in the Invoke node that replaces
-  // ParseUrl at analysis time (RuntimeReplaceable). This constant is used as the dispatch key
-  // in QueryPlanSerde.invokeConvertersByTargetClassName.
-  val invokeTargetClassName: String =
+object CometParseUrl
+    extends CometExpressionSerde[ParseUrl]
+    with CometInvokeExpressionSerde[ParseUrl] {
+
+  // In Spark 4.0, ParseUrl became RuntimeReplaceable and the analyser rewrites it to
+  // Invoke(ParseUrlEvaluator.evaluate, url, part[, key]).  The first child is a
+  // Literal(evaluator, ObjectType(ParseUrlEvaluator)).  This class name is the key
+  // used by QueryPlanSerde to route the Invoke node to this handler.
+  override val invokeTargetClassName: String =
     "org.apache.spark.sql.catalyst.expressions.url.ParseUrlEvaluator"
 
-  // ---------------------------------------------------------------------------
-  // Spark 4.0 Invoke path helpers
-  // ---------------------------------------------------------------------------
-
-  // Extracts the failOnError flag from the ParseUrlEvaluator instance embedded in the
-  // Invoke literal. Uses reflection because ParseUrlEvaluator is a private class.
-  // Falls back to SQLConf.get.ansiEnabled when reflection fails (evaluator is null, or
-  // the method has been renamed in a future Spark version).
-  private[serde] def failOnErrorFromInvoke(expr: Expression): Boolean =
-    expr.children.headOption match {
-      case Some(Literal(evaluator, objectType: ObjectType))
-          if evaluator != null && objectType.cls.getName == invokeTargetClassName =>
-        try {
-          evaluator.getClass
-            .getMethod("failOnError")
-            .invoke(evaluator)
-            .asInstanceOf[Boolean]
-        } catch {
-          case _: ReflectiveOperationException => SQLConf.get.ansiEnabled
-        }
-      case _ =>
-        SQLConf.get.ansiEnabled
+  // Extracts the failOnError flag from the ParseUrlEvaluator instance via reflection.
+  // Falls back to SQLConf.get.ansiEnabled when reflection is unavailable (null evaluator
+  // or renamed accessor in a future Spark version).
+  private def failOnErrorFromEvaluator(evaluator: AnyRef): Boolean =
+    try {
+      evaluator.getClass.getMethod("failOnError").invoke(evaluator).asInstanceOf[Boolean]
+    } catch {
+      case _: ReflectiveOperationException => SQLConf.get.ansiEnabled
     }
 
-  // Drops the leading ParseUrlEvaluator literal from the Invoke children list, leaving
-  // only the actual URL/part/key arguments.
-  private def dropEvaluatorLiteral(children: Seq[Expression]): Seq[Expression] =
-    children.headOption match {
-      case Some(Literal(_, objectType: ObjectType))
+  override def convertFromInvoke(
+      expr: ParseUrl,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[Expr] = {
+    // The first child is Literal(evaluator, ObjectType(ParseUrlEvaluator)).
+    // Strip it and read failOnError from it; the remaining children are (url, part[, key]).
+    val (urlArgs, failOnError) = expr.children match {
+      case Literal(evaluator, objectType: ObjectType) +: rest
           if objectType.cls.getName == invokeTargetClassName =>
-        children.drop(1)
-      case _ =>
-        children
+        val foe =
+          if (evaluator != null) failOnErrorFromEvaluator(evaluator.asInstanceOf[AnyRef])
+          else SQLConf.get.ansiEnabled
+        (rest, foe)
+      case args =>
+        (args, SQLConf.get.ansiEnabled)
     }
+    toProto(expr, urlArgs, failOnError, inputs, binding)
+  }
 
-  // ---------------------------------------------------------------------------
-  // Core serialization
-  // ---------------------------------------------------------------------------
+  // In Spark 3.5, ParseUrl is a concrete expression node with a `failOnError` field
+  // that is directly accessible without reflection.
+  override def convert(expr: ParseUrl, inputs: Seq[Attribute], binding: Boolean): Option[Expr] =
+    toProto(expr, expr.children, expr.failOnError, inputs, binding)
 
-  // Converts the parse_url/try_parse_url arguments into a proto Expr.
-  // `urlArgs` must already be stripped of the evaluator literal and the failOnError flag literal.
+  // Serializes (url, part[, key]) arguments into the appropriate native function call.
+  // Uses parse_url (ANSI/strict) or try_parse_url (legacy/lenient) depending on failOnError.
   private def toProto(
       expr: Expression,
       urlArgs: Seq[Expression],
@@ -443,25 +441,6 @@ object CometParseUrl extends CometExpressionSerde[ParseUrl] {
     val optExpr = scalarFunctionExprToProto(functionName, childExprs: _*)
     optExprWithInfo(optExpr, expr, urlArgs: _*)
   }
-
-  // ---------------------------------------------------------------------------
-  // Public entry points
-  // ---------------------------------------------------------------------------
-
-  // Spark 3.5 path: ParseUrl is a concrete expression node with a `failOnError` field.
-  override def convert(expr: ParseUrl, inputs: Seq[Attribute], binding: Boolean): Option[Expr] =
-    toProto(expr, expr.children, expr.failOnError, inputs, binding)
-
-  // Spark 4.0 path: ParseUrl is replaced by Invoke(ParseUrlEvaluator, url, part[, key]).
-  // Called from QueryPlanSerde.convertInvokeExpression via invokeConvertersByTargetClassName.
-  def convertFromInvoke(
-      expr: Expression,
-      inputs: Seq[Attribute],
-      binding: Boolean): Option[Expr] = {
-    val failOnError = failOnErrorFromInvoke(expr)
-    val urlArgs = dropEvaluatorLiteral(expr.children)
-    toProto(expr, urlArgs, failOnError, inputs, binding)
-  }
 }
 
 trait CommonStringExprs {