(spark) branch master updated: [SPARK-48944][CONNECT] Unify the JSON-format schema handling in Connect Server

[gu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 163256851ac0 [SPARK-48944][CONNECT] Unify the JSON-format schema handling in Connect Server
163256851ac0 is described below

commit 163256851ac08fb3cdde4c81e8db7863f3fe1272
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Sat Jul 20 09:53:25 2024 +0900

    [SPARK-48944][CONNECT] Unify the JSON-format schema handling in Connect Server
    
    ### What changes were proposed in this pull request?
    Simplify the JSON-format schema handling in Connect Server, by introducing a helper function `extractDataTypeFromJSON`
    
    ### Why are the changes needed?
    to unify the schema handling
    
    ### Does this PR introduce _any_ user-facing change?
    No, minor refactoring
    
    ### How was this patch tested?
    CI
    
    ### Was this patch authored or co-authored using generative AI tooling?
    No
    
    Closes #47415 from zhengruifeng/simplfy_from_json.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
 .../sql/connect/planner/SparkConnectPlanner.scala  | 76 +++++++++-------------
 1 file changed, 30 insertions(+), 46 deletions(-)

diff --git a/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
index 9eeec306cda8..013f0d83391c 100644
--- a/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
+++ b/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -1946,65 +1946,32 @@ class SparkConnectPlanner(
 
       case "from_json" if Seq(2, 3).contains(fun.getArgumentsCount) =>
         // JsonToStructs constructor doesn't accept JSON-formatted schema.
-        val children = fun.getArgumentsList.asScala.map(transformExpression)
-
-        var schema: DataType = null
-        children(1) match {
-          case Literal(s, StringType) if s != null =>
-            try {
-              schema = DataType.fromJson(s.toString)
-            } catch {
-              case _: Exception =>
-            }
-          case _ =>
-        }
-
-        if (schema != null) {
+        extractDataTypeFromJSON(fun.getArguments(1)).map { dataType =>
+          val children = fun.getArgumentsList.asScala.map(transformExpression)
+          val schema = CharVarcharUtils.failIfHasCharVarchar(dataType)
           var options = Map.empty[String, String]
           if (children.length == 3) {
             options = extractMapData(children(2), "Options")
           }
-          Some(
-            JsonToStructs(
-              schema = CharVarcharUtils.failIfHasCharVarchar(schema),
-              options = options,
-              child = children.head))
-        } else {
-          None
+          JsonToStructs(schema = schema, options = options, child = children.head)
         }
 
       case "from_xml" if Seq(2, 3).contains(fun.getArgumentsCount) =>
         // XmlToStructs constructor doesn't accept JSON-formatted schema.
-        val children = fun.getArgumentsList.asScala.map(transformExpression)
-
-        var schema: DataType = null
-        children(1) match {
-          case Literal(s, StringType) if s != null =>
-            try {
-              schema = DataType.fromJson(s.toString)
-            } catch {
-              case _: Exception =>
-            }
-          case _ =>
-        }
-
-        if (schema != null) {
-          schema match {
-            case t: StructType => t
-            case _ => throw DataTypeErrors.failedParsingStructTypeError(schema.sql)
+        extractDataTypeFromJSON(fun.getArguments(1)).map { dataType =>
+          val children = fun.getArgumentsList.asScala.map(transformExpression)
+          val schema = dataType match {
+            case t: StructType =>
+              CharVarcharUtils
+                .failIfHasCharVarchar(t)
+                .asInstanceOf[StructType]
+            case _ => throw DataTypeErrors.failedParsingStructTypeError(dataType.sql)
           }
-
           var options = Map.empty[String, String]
           if (children.length == 3) {
             options = extractMapData(children(2), "Options")
           }
-          Some(
-            XmlToStructs(
-              schema = CharVarcharUtils.failIfHasCharVarchar(schema).asInstanceOf[StructType],
-              options = options,
-              child = children.head))
-        } else {
-          None
+          XmlToStructs(schema = schema, options = options, child = children.head)
         }
 
       // Avro-specific functions
@@ -2189,6 +2156,23 @@ class SparkConnectPlanner(
     case other => throw InvalidPlanInput(s"$field should be created by map, but got $other")
   }
 
+  // Extract the schema from a literal string representing a JSON-formatted schema
+  private def extractDataTypeFromJSON(exp: proto.Expression): Option[DataType] = {
+    exp.getExprTypeCase match {
+      case proto.Expression.ExprTypeCase.LITERAL =>
+        exp.getLiteral.getLiteralTypeCase match {
+          case proto.Expression.Literal.LiteralTypeCase.STRING =>
+            try {
+              Some(DataType.fromJson(exp.getLiteral.getString))
+            } catch {
+              case _: Exception => None
+            }
+          case _ => None
+        }
+      case _ => None
+    }
+  }
+
   private def transformAlias(alias: proto.Expression.Alias): NamedExpression = {
     if (alias.getNameCount == 1) {
       val metadata = if (alias.hasMetadata() && alias.getMetadata.nonEmpty) {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org