-
Notifications
You must be signed in to change notification settings - Fork 28.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[SPARK-32376][SQL] Make unionByName null-filling behavior work with struct columns #29587
Changes from all commits
95e8fd4
5db1e0f
8bec8a3
2515d78
4398e77
c787f66
3ea24af
ae14447
72800e6
337cea7
90fc4fc
7b0d65d
a77481e
b4270f4
61ff46f
9040c56
9b21d91
1829889
d16bf7d
8a9522e
bb8938f
9e73928
c07e30f
2ca1379
3d907d0
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -17,29 +17,188 @@ | |
|
||
package org.apache.spark.sql.catalyst.analysis | ||
|
||
import scala.collection.mutable | ||
|
||
import org.apache.spark.sql.AnalysisException | ||
import org.apache.spark.sql.catalyst.expressions.{Alias, Literal} | ||
import org.apache.spark.sql.catalyst.expressions._ | ||
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.internal.SQLConf | ||
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] { | ||
private def unionTwoSides( | ||
/** | ||
* This method sorts columns recursively in a struct expression based on column names. | ||
*/ | ||
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 => | ||
cloud-fan marked this conversation as resolved.
Show resolved
Hide resolved
|
||
throw new IllegalStateException(s"`UpdateFields` has incorrect expression: $other. " + | ||
"Please file a bug report with this error message, stack trace, and the query.") | ||
} | ||
} | ||
|
||
def simplifyWithFields(expr: Expression): Expression = { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yeah, actually there is #29812 for that, but is stuck by other PR that is refactoring |
||
expr.transformUp { | ||
case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2) => | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is it important to have this optimization inside this analyzer rule? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yea. Without optimizing the expressions, we cannot scale up well for deeply nested schema, e.g. the added test There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Actually I plan to move this optimization out of |
||
UpdateFields(struct, fieldOps1 ++ fieldOps2) | ||
} | ||
} | ||
|
||
/** | ||
* 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 = { | ||
assert(col.dataType.isInstanceOf[StructType], "Only support StructType.") | ||
|
||
val resolver = SQLConf.get.resolver | ||
val missingFieldsOpt = | ||
StructType.findMissingFields(col.dataType.asInstanceOf[StructType], target, resolver) | ||
|
||
// We need to sort columns in result, because we might add another column in other side. | ||
// E.g., we want to union two structs "a int, b long" and "a int, c string". | ||
// If we don't sort, we will have "a int, b long, c string" and | ||
// "a int, c string, b long", which are not compatible. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. is this behavior consistent with top-level columns? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think this is related to the comment: #29587 (comment) |
||
if (missingFieldsOpt.isEmpty) { | ||
sortStructFields(col) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We don't need to sort names recursively for nested struct cases? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Oh, we need this to make sure two sides have consistent schema. For example the test case from @fqaiser94 in #29587 (comment), when we add field to one side, another side still needs to sort its column, otherwise there is inconsistency. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @maropu I got you point when I was fixing the performance issue. Yeah, we should. I fixed it in latest commit. Thanks. |
||
} else { | ||
missingFieldsOpt.map { s => | ||
val struct = addFieldsInto(col, s.fields) | ||
// Combines `WithFields`s to reduce expression tree. | ||
val reducedStruct = simplifyWithFields(struct) | ||
val sorted = sortStructFieldsInWithFields(reducedStruct) | ||
sorted | ||
}.get | ||
} | ||
} | ||
|
||
/** | ||
* Adds missing fields recursively into given `col` expression. The missing fields are given | ||
* in `fields`. For example, given `col` as "z struct<z:int, y:int>, x int", and `fields` is | ||
* "z struct<w:long>, w string". This method will add a nested `z.w` field and a top-level | ||
* `w` field to `col` and fill null values for them. Note that because we might also add missing | ||
* fields at other side of Union, we must make sure corresponding attributes at two sides have | ||
* same field order in structs, so when we adding missing fields, we will sort the fields based on | ||
* field names. So the data type of returned expression will be | ||
* "w string, x int, z struct<w:long, y:int, z:int>". | ||
*/ | ||
private def addFieldsInto( | ||
col: Expression, | ||
fields: Seq[StructField]): Expression = { | ||
fields.foldLeft(col) { case (currCol, field) => | ||
field.dataType match { | ||
case st: StructType => | ||
val resolver = SQLConf.get.resolver | ||
val colField = currCol.dataType.asInstanceOf[StructType] | ||
.find(f => resolver(f.name, field.name)) | ||
if (colField.isEmpty) { | ||
// The whole struct is missing. Add a null. | ||
UpdateFields(currCol, field.name, Literal(null, st)) | ||
} else { | ||
UpdateFields(currCol, field.name, | ||
addFieldsInto(ExtractValue(currCol, Literal(field.name), resolver), st.fields)) | ||
} | ||
case dt => | ||
UpdateFields(currCol, field.name, Literal(null, dt)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. what if There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The top-level columns support There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good question. So it is actually the same issue as #29587 (comment), a.k.a adjusting the nested columns to have a more natural schema. As replied in the discussion, I plan to do it in followup. |
||
} | ||
} | ||
} | ||
|
||
/** | ||
* This method will compare right to left plan's outputs. If there is one struct attribute | ||
* at right side has same name with left side struct attribute, but two structs are not the | ||
* same data type, i.e., some missing (nested) fields at right struct attribute, then this | ||
* method will try to add missing (nested) fields into the right attribute with null values. | ||
*/ | ||
private def compareAndAddFields( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Although we have a rich comment in the function body, could you add a function description to give a general idea? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added. |
||
left: LogicalPlan, | ||
right: LogicalPlan, | ||
allowMissingCol: Boolean): LogicalPlan = { | ||
allowMissingCol: Boolean): (Seq[NamedExpression], Seq[NamedExpression]) = { | ||
val resolver = SQLConf.get.resolver | ||
val leftOutputAttrs = left.output | ||
val rightOutputAttrs = right.output | ||
|
||
// Builds a project list for `right` based on `left` output names | ||
val aliased = mutable.ArrayBuffer.empty[Attribute] | ||
|
||
val rightProjectList = leftOutputAttrs.map { lattr => | ||
rightOutputAttrs.find { rattr => resolver(lattr.name, rattr.name) }.getOrElse { | ||
val found = rightOutputAttrs.find { rattr => resolver(lattr.name, rattr.name) } | ||
if (found.isDefined) { | ||
dongjoon-hyun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
val foundAttr = found.get | ||
val foundDt = foundAttr.dataType | ||
(foundDt, lattr.dataType) match { | ||
case (source: StructType, target: StructType) | ||
if allowMissingCol && !source.sameType(target) => | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. To make the logic simpler, could we filter out all the unsupported case (e.g.,
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Hmm, I'm not sure where we can simplify the logic? By adding There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I read the comment and I thought first that all the unsupported cases are handled in the line 108-112. But, it also means unsupported cases if There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I see. I will add more comments explaining this. |
||
// Having an output with same name, but different struct type. | ||
// We need to add missing fields. Note that if there are deeply nested structs such as | ||
// nested struct of array in struct, we don't support to add missing deeply nested field | ||
// like that. We will sort columns in the struct expression to make sure two sides of | ||
// union have consistent schema. | ||
aliased += foundAttr | ||
Alias(addFields(foundAttr, target), foundAttr.name)() | ||
case _ => | ||
// We don't need/try to add missing fields if: | ||
// 1. The attributes of left and right side are the same struct type | ||
// 2. The attributes are not struct types. They might be primitive types, or array, map | ||
// types. We don't support adding missing fields of nested structs in array or map | ||
// types now. | ||
// 3. `allowMissingCol` is disabled. | ||
foundAttr | ||
} | ||
} else { | ||
if (allowMissingCol) { | ||
Alias(Literal(null, lattr.dataType), lattr.name)() | ||
} else { | ||
|
@@ -50,18 +209,29 @@ object ResolveUnion extends Rule[LogicalPlan] { | |
} | ||
} | ||
|
||
(rightProjectList, aliased.toSeq) | ||
} | ||
viirya marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
private def unionTwoSides( | ||
left: LogicalPlan, | ||
right: LogicalPlan, | ||
allowMissingCol: Boolean): LogicalPlan = { | ||
val rightOutputAttrs = right.output | ||
|
||
// Builds a project list for `right` based on `left` output names | ||
val (rightProjectList, aliased) = compareAndAddFields(left, right, allowMissingCol) | ||
|
||
// Delegates failure checks to `CheckAnalysis` | ||
val notFoundAttrs = rightOutputAttrs.diff(rightProjectList) | ||
val notFoundAttrs = rightOutputAttrs.diff(rightProjectList ++ aliased) | ||
val rightChild = Project(rightProjectList ++ notFoundAttrs, right) | ||
|
||
// Builds a project for `logicalPlan` based on `right` output names, if allowing | ||
// missing columns. | ||
val leftChild = if (allowMissingCol) { | ||
val missingAttrs = notFoundAttrs.map { attr => | ||
Alias(Literal(null, attr.dataType), attr.name)() | ||
} | ||
if (missingAttrs.nonEmpty) { | ||
Project(leftOutputAttrs ++ missingAttrs, left) | ||
// Add missing (nested) fields to left plan. | ||
val (leftProjectList, _) = compareAndAddFields(rightChild, left, allowMissingCol) | ||
if (leftProjectList.map(_.toAttribute) != left.output) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit:
? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Doesn't |
||
Project(leftProjectList, left) | ||
} else { | ||
left | ||
} | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -641,4 +641,39 @@ object StructType extends AbstractDataType { | |
fields.foreach(s => map.put(s.name, s)) | ||
map | ||
} | ||
|
||
/** | ||
* Returns a `StructType` that contains missing fields recursively from `source` to `target`. | ||
* Note that this doesn't support looking into array type and map type recursively. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I leverage |
||
*/ | ||
def findMissingFields( | ||
source: StructType, | ||
target: StructType, | ||
resolver: Resolver): Option[StructType] = { | ||
def bothStructType(dt1: DataType, dt2: DataType): Boolean = | ||
dt1.isInstanceOf[StructType] && dt2.isInstanceOf[StructType] | ||
|
||
val newFields = mutable.ArrayBuffer.empty[StructField] | ||
|
||
target.fields.foreach { field => | ||
val found = source.fields.find(f => resolver(field.name, f.name)) | ||
if (found.isEmpty) { | ||
// Found a missing field in `source`. | ||
newFields += field | ||
} else if (bothStructType(found.get.dataType, field.dataType) && | ||
!found.get.dataType.sameType(field.dataType)) { | ||
// Found a field with same name, but different data type. | ||
findMissingFields(found.get.dataType.asInstanceOf[StructType], | ||
field.dataType.asInstanceOf[StructType], resolver).map { missingType => | ||
newFields += found.get.copy(dataType = missingType) | ||
} | ||
} | ||
} | ||
|
||
if (newFields.isEmpty) { | ||
None | ||
} else { | ||
Some(StructType(newFields.toSeq)) | ||
} | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are functions having the same names, so could we assign different names? I think its a bit confusing.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
sure, let me think about better method names.