[SPARK-50066][SQL] Codegen Support for SchemaOfXml (by Invoke & RuntimeReplaceable)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/XmlExpressionEvalUtils.scala

@@ -0,0 +1,42 @@
+/*
+ * 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.spark.sql.catalyst.expressions.xml
+
+import org.apache.spark.sql.catalyst.xml.XmlInferSchema
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, DataType, StructType}
+import org.apache.spark.unsafe.types.UTF8String
+
+object XmlExpressionEvalUtils {
+
+  def schemaOfXml(xmlInferSchema: XmlInferSchema, xml: UTF8String): UTF8String = {
+    val dataType = xmlInferSchema.infer(xml.toString).get match {
+      case st: StructType =>
+        xmlInferSchema.canonicalizeType(st).getOrElse(StructType(Nil))
+      case at: ArrayType if at.elementType.isInstanceOf[StructType] =>
+        xmlInferSchema
+          .canonicalizeType(at.elementType)
+          .map(ArrayType(_, containsNull = at.containsNull))
+          .getOrElse(ArrayType(StructType(Nil), containsNull = at.containsNull))
+      case other: DataType =>
+        xmlInferSchema.canonicalizeType(other).getOrElse(SQLConf.get.defaultStringType)
+    }
+
+    UTF8String.fromString(dataType.sql)
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala

@@ -21,7 +21,9 @@ import java.io.CharArrayWriter
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, TypeCheckSuccess}
-import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodegenFallback, ExprCode}
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
+import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
+import org.apache.spark.sql.catalyst.expressions.xml.XmlExpressionEvalUtils
 import org.apache.spark.sql.catalyst.util.{DropMalformedMode, FailFastMode, FailureSafeParser, PermissiveMode}
 import org.apache.spark.sql.catalyst.util.TypeUtils._
 import org.apache.spark.sql.catalyst.xml.{StaxXmlGenerator, StaxXmlParser, ValidatorUtil, XmlInferSchema, XmlOptions}
@@ -149,7 +151,9 @@ case class XmlToStructs(
 case class SchemaOfXml(
     child: Expression,
     options: Map[String, String])
-  extends UnaryExpression with CodegenFallback with QueryErrorsBase {
+  extends UnaryExpression
+  with RuntimeReplaceable
+  with QueryErrorsBase {
 
   def this(child: Expression) = this(child, Map.empty[String, String])
 
@@ -192,26 +196,20 @@ case class SchemaOfXml(
     }
   }
 
-  override def eval(v: InternalRow): Any = {
-    val dataType = xmlInferSchema.infer(xml.toString).get match {
-      case st: StructType =>
-        xmlInferSchema.canonicalizeType(st).getOrElse(StructType(Nil))
-      case at: ArrayType if at.elementType.isInstanceOf[StructType] =>
-        xmlInferSchema
-          .canonicalizeType(at.elementType)
-          .map(ArrayType(_, containsNull = at.containsNull))
-          .getOrElse(ArrayType(StructType(Nil), containsNull = at.containsNull))
-      case other: DataType =>
-        xmlInferSchema.canonicalizeType(other).getOrElse(SQLConf.get.defaultStringType)
-    }
-
-    UTF8String.fromString(dataType.sql)
-  }
-
   override def prettyName: String = "schema_of_xml"
 
   override protected def withNewChildInternal(newChild: Expression): SchemaOfXml =
     copy(child = newChild)
+
+  @transient private lazy val xmlInferSchemaObjectType = ObjectType(classOf[XmlInferSchema])
+
+  override def replacement: Expression = StaticInvoke(
+    XmlExpressionEvalUtils.getClass,
+    dataType,
+    "schemaOfXml",
+    Seq(Literal(xmlInferSchema, xmlInferSchemaObjectType), child),
+    Seq(xmlInferSchemaObjectType, child.dataType)
+  )
 }
 
 /**