Kimahriman commented on a change in pull request #33040:
URL: https://github.com/apache/spark/pull/33040#discussion_r657293379
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveUnion.scala
##########
@@ -20,137 +20,66 @@ package org.apache.spark.sql.catalyst.analysis
import scala.collection.mutable
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.optimizer.{CombineUnions,
OptimizeUpdateFields}
+import org.apache.spark.sql.catalyst.optimizer.{CombineUnions}
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project,
Union}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.UNION
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.SchemaUtils
-import org.apache.spark.unsafe.types.UTF8String
/**
* Resolves different children of Union to a common set of columns.
*/
object ResolveUnion extends Rule[LogicalPlan] {
/**
- * This method sorts columns recursively in a struct expression based on
column names.
+ * Adds missing fields recursively into given `col` expression, based on the
expected struct
+ * fields from merging the two schemas. This is called by
`compareAndAddFields` when we find two
+ * struct columns with same name but different nested fields. This method
will recursively
+ * return a new struct with all of the expected fields, adding null values
when `col` doesn't
+ * already contain them. Currently we don't support merging structs nested
inside of arrays
+ * or maps.
*/
- private def sortStructFields(expr: Expression): Expression = {
- val existingExprs =
expr.dataType.asInstanceOf[StructType].fieldNames.zipWithIndex.map {
- case (name, i) =>
- val fieldExpr = GetStructField(KnownNotNull(expr), i)
- if (fieldExpr.dataType.isInstanceOf[StructType]) {
- (name, sortStructFields(fieldExpr))
- } else {
- (name, fieldExpr)
- }
- }.sortBy(_._1).flatMap(pair => Seq(Literal(pair._1), pair._2))
-
- val newExpr = CreateNamedStruct(existingExprs)
- if (expr.nullable) {
- If(IsNull(expr), Literal(null, newExpr.dataType), newExpr)
- } else {
- newExpr
- }
- }
-
- /**
- * Assumes input expressions are field expression of `CreateNamedStruct`.
This method
- * sorts the expressions based on field names.
- */
- private def sortFieldExprs(fieldExprs: Seq[Expression]): Seq[Expression] = {
- fieldExprs.grouped(2).map { e =>
- Seq(e.head, e.last)
- }.toSeq.sortBy { pair =>
- assert(pair.head.isInstanceOf[Literal])
- pair.head.eval().asInstanceOf[UTF8String].toString
- }.flatten
- }
-
- /**
- * This helper method sorts fields in a `UpdateFields` expression by field
name.
- */
- private def sortStructFieldsInWithFields(expr: Expression): Expression =
expr transformUp {
- case u: UpdateFields if u.resolved =>
- u.evalExpr match {
- case i @ If(IsNull(_), _, CreateNamedStruct(fieldExprs)) =>
- val sorted = sortFieldExprs(fieldExprs)
- val newStruct = CreateNamedStruct(sorted)
- i.copy(trueValue = Literal(null, newStruct.dataType), falseValue =
newStruct)
- case CreateNamedStruct(fieldExprs) =>
- val sorted = sortFieldExprs(fieldExprs)
- val newStruct = CreateNamedStruct(sorted)
- newStruct
- case other =>
- throw new IllegalStateException(s"`UpdateFields` has incorrect
expression: $other. " +
- "Please file a bug report with this error message, stack trace,
and the query.")
- }
- }
-
- /**
- * Adds missing fields recursively into given `col` expression, based on the
target `StructType`.
- * This is called by `compareAndAddFields` when we find two struct columns
with same name but
- * different nested fields. This method will find out the missing nested
fields from `col` to
- * `target` struct and add these missing nested fields. Currently we don't
support finding out
- * missing nested fields of struct nested in array or struct nested in map.
- */
- private def addFields(col: NamedExpression, target: StructType): Expression
= {
+ private def addFields(col: Expression, targetType: StructType): Expression =
{
Review comment:
Yep that's exactly what I had in mind as the solution for
https://github.com/apache/spark/pull/32972 once this got merged
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
