Skip to content

[SPARK-35213][SQL] Keep the correct ordering of nested structs in chained withField operations #32338

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.

Sign up for GitHub

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

Closed
Kimahriman wants to merge 2 commits into from
Closed

[SPARK-35213][SQL] Keep the correct ordering of nested structs in chained withField operations #32338

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1af0dae
Keep the correct ordering of nested structs in chained withField oper…
Kimahriman Apr 25, 2021
23ee428
Use the last casing instead of the first for case insensitivity
Kimahriman Apr 25, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 15 additions & 10 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,28 +49,33 @@ object OptimizeUpdateFields extends Rule[LogicalPlan] {
val values = withFields.map(_.valExpr)

val newNames = mutable.ArrayBuffer.empty[String]
val newValues = mutable.ArrayBuffer.empty[Expression]
val newValues = mutable.HashMap.empty[String, Expression]

if (caseSensitive) {
names.zip(values).reverse.foreach { case (name, value) =>
if (!newNames.contains(name)) {
names.zip(values).foreach { case (name, value) =>
if (newValues.contains(name)) {
newValues += name -> value
} else {
newNames += name
newValues += value
newValues += name -> value
}
}
} else {
val nameSet = mutable.HashSet.empty[String]
names.zip(values).reverse.foreach { case (name, value) =>
names.zip(values).foreach { case (name, value) =>
val lowercaseName = name.toLowerCase(Locale.ROOT)
if (!nameSet.contains(lowercaseName)) {
if (newValues.contains(lowercaseName)) {
// Add both versions so we pull out the latest value based on whichever casing was
// provided first
newValues += name -> value
newValues += lowercaseName -> value
} else {
newNames += name
newValues += value
nameSet += lowercaseName
newValues += name -> value
}
}
}

val newWithFields = newNames.reverse.zip(newValues.reverse).map(p => WithField(p._1, p._2))
val newWithFields = newNames.map(n => WithField(n, newValues(n)))
UpdateFields(structExpr, newWithFields.toSeq)

case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2) =>
Expand Down
23 changes: 22 additions & 1 deletion ...lyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeWithFieldsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ class OptimizeWithFieldsSuite extends PlanTest {
val correctAnswer = testRelation
.select(
Alias(UpdateFields('a, WithField("b1", Literal(5)) :: Nil), "out1")(),
Alias(UpdateFields('a, WithField("B1", Literal(5)) :: Nil), "out2")())
Alias(UpdateFields('a, WithField("b1", Literal(5)) :: Nil), "out2")())

@Kimahriman Kimahriman Apr 25, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One result is that for case-insensitive cases, the first casing seen for a field is maintained, rather than the last one. If this isn't what we want, I can update it to keep the last casing seen

@viirya viirya Apr 25, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As this is for case-insensitive, seems no big deal. Although for the semantics, the "B1" is specified later, so I guess it is more reasonable to keep later one.

@Kimahriman Kimahriman Apr 25, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed it to keep the last casing instead

.analyze

comparePlans(optimized, correctAnswer)
Expand All @@ -126,4 +126,25 @@ class OptimizeWithFieldsSuite extends PlanTest {
comparePlans(optimized, correctAnswer)
}
}

test("SPARK-35213: ensure optimize WithFields maintains correct struct ordering") {

@viirya viirya Apr 25, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

struct ordering -> withfield ordering

@Kimahriman Kimahriman Apr 25, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

val originalQuery = testRelation
.select(
Alias(UpdateFields('a,
WithField("a1", Literal(3)) ::
WithField("b1", Literal(4)) ::
WithField("a1", Literal(5)) ::
Nil), "out")())

val optimized = Optimize.execute(originalQuery.analyze)
val correctAnswer = testRelation
.select(
Alias(UpdateFields('a,
WithField("a1", Literal(5)) ::
WithField("b1", Literal(4)) ::
Nil), "out")())
.analyze

comparePlans(optimized, correctAnswer)
}

@viirya viirya Apr 25, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you check the output data type, you can see the struct type is not different:

optimized: ArrayBuffer(StructType(StructField(a1,IntegerType,false), StructField(b1,IntegerType,false)))
correctAnswer: ArrayBuffer(StructType(StructField(a1,IntegerType,false), StructField(b1,IntegerType,false)))

By design, UpdateFields will keep the order of fields in struct expression.

@viirya viirya Apr 25, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But yea, it looks better to keep original WithField order.

@Kimahriman Kimahriman Apr 25, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was just to sanity check the WithField order does actually stay the same, the tests on the Column Suite show how it can actually give you an incorrect schema. I don't fully know how a schema is determined (what part of the planning phase)

}
55 changes: 55 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1686,6 +1686,61 @@ class ColumnExpressionSuite extends QueryTest with SharedSparkSession {
StructType(Seq(StructField("a", IntegerType, nullable = true))))
}

test("SPARK-35213: chained withField operations should have correct schema for new columns") {
val df = spark.createDataFrame(
sparkContext.parallelize(Row(null) :: Nil),
StructType(Seq(StructField("data", NullType))))

@Kimahriman Kimahriman Apr 25, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to just create an empty dataframe with no columns in Scala? I mostly operate and python and can just do spark.createDataFrame([[]])


checkAnswer(
df.withColumn("data", struct()
.withField("a", struct())
.withField("b", struct())
.withField("a.aa", lit("aa1"))
.withField("b.ba", lit("ba1"))
.withField("a.ab", lit("ab1"))),
Row(Row(Row("aa1", "ab1"), Row("ba1"))) :: Nil,
StructType(Seq(
StructField("data", StructType(Seq(
StructField("a", StructType(Seq(
StructField("aa", StringType, nullable = false),
StructField("ab", StringType, nullable = false)
)), nullable = false),
StructField("b", StructType(Seq(
StructField("ba", StringType, nullable = false)
)), nullable = false)
)), nullable = false)
Comment on lines +1702 to +1711

@viirya viirya Apr 25, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Using ddl might be more readable?

@Kimahriman Kimahriman Apr 25, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it's kinda verbose, but I feel like for complicated things the objects are easier to understand than the DDL strings, especially with structs. Wasn't sure if there was an easier way to not have to explicitly mark everything as not nullable at least

))
)
}

test("SPARK-35213: optimized withField operations should maintain correct nested struct " +
"ordering") {
val df = spark.createDataFrame(
sparkContext.parallelize(Row(null) :: Nil),
StructType(Seq(StructField("data", NullType))))

checkAnswer(
df.withColumn("data", struct()
.withField("a", struct().withField("aa", lit("aa1")))
.withField("b", struct().withField("ba", lit("ba1")))
)
.withColumn("data", col("data").withField("b.bb", lit("bb1")))
.withColumn("data", col("data").withField("a.ab", lit("ab1"))),
Row(Row(Row("aa1", "ab1"), Row("ba1", "bb1"))) :: Nil,
StructType(Seq(
StructField("data", StructType(Seq(
StructField("a", StructType(Seq(
StructField("aa", StringType, nullable = false),
StructField("ab", StringType, nullable = false)
)), nullable = false),
StructField("b", StructType(Seq(
StructField("ba", StringType, nullable = false),
StructField("bb", StringType, nullable = false)
)), nullable = false)
)), nullable = false)
))
)
}

test("dropFields should throw an exception if called on a non-StructType column") {
intercept[AnalysisException] {
Expand Down