[SPARK-33007][SQL] Simplify named_struct + get struct field + from_json expression chain

### What changes were proposed in this pull request?

This proposes to simplify named_struct + get struct field + from_json expression chain from `struct(from_json.col1, from_json.col2, from_json.col3...)` to `struct(from_json)`.

### Why are the changes needed?

Simplify complex expression tree that could be produced by query optimization or user.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Unit test.

Closes #29942 from viirya/SPARK-33007.

Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>

Author: viirya

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprs.scala

@@ -28,10 +28,46 @@ import org.apache.spark.sql.types.{ArrayType, StructType}
  * The optimization includes:
  * 1. JsonToStructs(StructsToJson(child)) => child.
  * 2. Prune unnecessary columns from GetStructField/GetArrayStructFields + JsonToStructs.
+ * 3. CreateNamedStruct(JsonToStructs(json).col1, JsonToStructs(json).col2, ...) =>
+ *      If(IsNull(json), nullStruct, KnownNotNull(JsonToStructs(prunedSchema, ..., json)))
+ *      if JsonToStructs(json) is shared among all fields of CreateNamedStruct. `prunedSchema`
+ *      contains all accessed fields in original CreateNamedStruct.
  */
 object OptimizeJsonExprs extends Rule[LogicalPlan] {
   override def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case p => p.transformExpressions {
+
+      case c: CreateNamedStruct
+          // If we create struct from various fields of the same `JsonToStructs`.
+          if c.valExprs.forall { v =>
+            v.isInstanceOf[GetStructField] &&
+              v.asInstanceOf[GetStructField].child.isInstanceOf[JsonToStructs] &&
+              v.children.head.semanticEquals(c.valExprs.head.children.head)
+          } =>
+        val jsonToStructs = c.valExprs.map(_.children.head)
+        val sameFieldName = c.names.zip(c.valExprs).forall {
+          case (name, valExpr: GetStructField) =>
+            name.toString == valExpr.childSchema(valExpr.ordinal).name
+          case _ => false
+        }
+
+        // Although `CreateNamedStruct` allows duplicated field names, e.g. "a int, a int",
+        // `JsonToStructs` does not support parsing json with duplicated field names.
+        val duplicateFields = c.names.map(_.toString).distinct.length != c.names.length
+
+        // If we create struct from various fields of the same `JsonToStructs` and we don't
+        // alias field names and there is no duplicated field in the struct.
+        if (sameFieldName && !duplicateFields) {
+          val fromJson = jsonToStructs.head.asInstanceOf[JsonToStructs].copy(schema = c.dataType)
+          val nullFields = c.children.grouped(2).flatMap {
+            case Seq(name, value) => Seq(name, Literal(null, value.dataType))
+          }.toSeq
+
+          If(IsNull(fromJson.child), c.copy(children = nullFields), KnownNotNull(fromJson))
+        } else {
+          c
+        }
+
       case jsonToStructs @ JsonToStructs(_, options1,
         StructsToJson(options2, child, timeZoneId2), timeZoneId1)
           if options1.isEmpty && options2.isEmpty && timeZoneId1 == timeZoneId2 &&

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala

@@ -199,4 +199,71 @@ class OptimizeJsonExprsSuite extends PlanTest with ExpressionEvalHelper {
         JsonToStructs(prunedSchema2, options, 'json), field2, 0, 1, false).as("b")).analyze
     comparePlans(optimized2, expected2)
   }
+
+  test("SPARK-33007: simplify named_struct + from_json") {
+    val options = Map.empty[String, String]
+    val schema = StructType.fromDDL("a int, b int, c long, d string")
+
+    val prunedSchema1 = StructType.fromDDL("a int, b int")
+    val nullStruct = namedStruct("a", Literal(null, IntegerType), "b", Literal(null, IntegerType))
+
+    val UTC_OPT = Option("UTC")
+    val json: BoundReference = 'json.string.canBeNull.at(0)
+
+    assertEquivalent(
+      testRelation2,
+      namedStruct(
+        "a", GetStructField(JsonToStructs(schema, options, json, UTC_OPT), 0),
+        "b", GetStructField(JsonToStructs(schema, options, json, UTC_OPT), 1)).as("struct"),
+      If(IsNull(json),
+        nullStruct,
+        KnownNotNull(JsonToStructs(prunedSchema1, options, json, UTC_OPT))).as("struct"))
+
+    val field1 = StructType.fromDDL("a int")
+    val field2 = StructType.fromDDL("b int")
+
+    // Skip optimization if `namedStruct` aliases field name.
+    assertEquivalent(
+      testRelation2,
+      namedStruct(
+        "a1", GetStructField(JsonToStructs(schema, options, json, UTC_OPT), 0),
+        "b", GetStructField(JsonToStructs(schema, options, json, UTC_OPT), 1)).as("struct"),
+      namedStruct(
+        "a1", GetStructField(JsonToStructs(field1, options, json, UTC_OPT), 0),
+        "b", GetStructField(JsonToStructs(field2, options, json, UTC_OPT), 0)).as("struct"))
+
+    assertEquivalent(
+      testRelation2,
+      namedStruct(
+        "a", GetStructField(JsonToStructs(schema, options, json, UTC_OPT), 0),
+        "a", GetStructField(JsonToStructs(schema, options, json, UTC_OPT), 0)).as("struct"),
+      namedStruct(
+        "a", GetStructField(JsonToStructs(field1, options, json, UTC_OPT), 0),
+        "a", GetStructField(JsonToStructs(field1, options, json, UTC_OPT), 0)).as("struct"))
+
+    val PST = getZoneId("-08:00")
+    // Skip optimization if `JsonToStructs`s are not the same.
+    assertEquivalent(
+      testRelation2,
+      namedStruct(
+        "a", GetStructField(JsonToStructs(schema, options, json, UTC_OPT), 0),
+        "b", GetStructField(JsonToStructs(schema, options, json, Option(PST.getId)), 1))
+        .as("struct"),
+      namedStruct(
+        "a", GetStructField(JsonToStructs(field1, options, json, UTC_OPT), 0),
+        "b", GetStructField(JsonToStructs(field2, options, json, Option(PST.getId)), 0))
+        .as("struct"))
+  }
+
+  private def assertEquivalent(relation: LocalRelation, e1: Expression, e2: Expression): Unit = {
+    val plan = relation.select(e1).analyze
+    val actual = Optimizer.execute(plan)
+    val expected = relation.select(e2).analyze
+    comparePlans(actual, expected)
+
+    Seq("""{"a":1, "b":2, "c": 123, "d": "test"}""", null).foreach(v => {
+      val row = create_row(v)
+      checkEvaluation(e1, e2.eval(row), row)
+    })
+  }
 }