You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2021/08/24 15:33:41 UTC
[spark] branch branch-3.1 updated: [SPARK-35876][SQL][3.1]
ArraysZip should retain field names to avoid being re-written by
analyzer/optimizer
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch branch-3.1
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-3.1 by this push:
new 644cad2 [SPARK-35876][SQL][3.1] ArraysZip should retain field names to avoid being re-written by analyzer/optimizer
644cad2 is described below
commit 644cad26f489b3c1e8210245f1d958d05c208e53
Author: Angerszhuuuu <an...@gmail.com>
AuthorDate: Tue Aug 24 23:33:00 2021 +0800
[SPARK-35876][SQL][3.1] ArraysZip should retain field names to avoid being re-written by analyzer/optimizer
#### What changes were proposed in this pull request?
This PR fixes an issue that field names of structs generated by arrays_zip function could be unexpectedly re-written by analyzer/optimizer.
Here is an example.
```
val df = sc.parallelize(Seq((Array(1, 2), Array(3, 4)))).toDF("a1", "b1").selectExpr("arrays_zip(a1, b1) as zipped")
df.printSchema
root
|-- zipped: array (nullable = true)
| |-- element: struct (containsNull = false)
| | |-- a1: integer (nullable = true) // OK. a1 is expected name
| | |-- b1: integer (nullable = true) // OK. b1 is expected name
df.explain
== Physical Plan ==
*(1) Project [arrays_zip(_1#3, _2#4) AS zipped#12] // Not OK. field names are re-written as _1 and _2 respectively
df.write.parquet("/tmp/test.parquet")
val df2 = spark.read.parquet("/tmp/test.parquet")
df2.printSchema
root
|-- zipped: array (nullable = true)
| |-- element: struct (containsNull = true)
| | |-- _1: integer (nullable = true) // Not OK. a1 is expected but got _1
| | |-- _2: integer (nullable = true) // Not OK. b1 is expected but got _2
```
This issue happens when aliases are eliminated by AliasHelper.replaceAliasButKeepName or AliasHelper.trimNonTopLevelAliases called via analyzer/optimizer
spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
Line 883 in b89cd8d
upper.map(replaceAliasButKeepName(_, aliases))
spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
Line 3759 in b89cd8d
val cleanedProjectList = projectList.map(trimNonTopLevelAliases)
I investigated functions which can be affected this issue but I found only arrays_zip so far.
To fix this issue, this PR changes the definition of ArraysZip to retain field names to avoid being re-written by analyzer/optimizer.
### Why are the changes needed?
This is apparently a bug.
### Does this PR introduce any user-facing change?
No. After this change, the field names are no longer re-written but it should be expected behavior for users.
#### How was this patch tested?
New tests.
Closes #33810 from AngersZhuuuu/SPARK-35876-3.1.
Lead-authored-by: Angerszhuuuu <an...@gmail.com>
Co-authored-by: Kousuke Saruta <sa...@oss.nttdata.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../spark/sql/catalyst/analysis/Analyzer.scala | 14 ++++-
.../expressions/collectionOperations.scala | 40 +++++++++---
.../apache/spark/sql/DataFrameFunctionsSuite.scala | 72 +++++++++++++++++++++-
3 files changed, 115 insertions(+), 11 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index a9ff6ee..b2af21d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -248,7 +248,7 @@ class Analyzer(override val catalogManager: CatalogManager)
ResolvePartitionSpec ::
AddMetadataColumns ::
ResolveReferences ::
- ResolveCreateNamedStruct ::
+ ResolveExpressionsWithNamePlaceholders ::
ResolveDeserializer ::
ResolveNewInstance ::
ResolveUpCast ::
@@ -3832,10 +3832,18 @@ object TimeWindowing extends Rule[LogicalPlan] {
}
/**
- * Resolve a [[CreateNamedStruct]] if it contains [[NamePlaceholder]]s.
+ * Resolve expressions if they contains [[NamePlaceholder]]s.
*/
-object ResolveCreateNamedStruct extends Rule[LogicalPlan] {
+object ResolveExpressionsWithNamePlaceholders extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveExpressions {
+ case e: ArraysZip if !e.resolved =>
+ val names = e.children.zip(e.names).map {
+ case (e: NamedExpression, NamePlaceholder) if e.resolved =>
+ Literal(e.name)
+ case (_, other) => other
+ }
+ ArraysZip(e.children, names)
+
case e: CreateNamedStruct if !e.resolved =>
val children = e.children.grouped(2).flatMap {
case Seq(NamePlaceholder, e: NamedExpression) if e.resolved =>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
index 0765bfd..b341895 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
@@ -23,7 +23,7 @@ import scala.collection.mutable
import scala.reflect.ClassTag
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion}
+import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion, UnresolvedAttribute}
import org.apache.spark.sql.catalyst.expressions.ArraySortLike.NullOrder
import org.apache.spark.sql.catalyst.expressions.codegen._
import org.apache.spark.sql.catalyst.expressions.codegen.Block._
@@ -174,16 +174,36 @@ case class MapKeys(child: Expression)
""",
group = "array_funcs",
since = "2.4.0")
-case class ArraysZip(children: Seq[Expression]) extends Expression with ExpectsInputTypes {
+case class ArraysZip(children: Seq[Expression], names: Seq[Expression])
+ extends Expression with ExpectsInputTypes {
+
+ def this(children: Seq[Expression]) = {
+ this(
+ children,
+ children.zipWithIndex.map {
+ case (u: UnresolvedAttribute, _) => Literal(u.nameParts.last)
+ case (e: NamedExpression, _) if e.resolved => Literal(e.name)
+ case (e: NamedExpression, _) => NamePlaceholder
+ case (_, idx) => Literal(idx.toString)
+ })
+ }
+
+ if (children.size != names.size) {
+ throw new IllegalArgumentException(
+ "The numbers of zipped arrays and field names should be the same")
+ }
+
+ override lazy val resolved: Boolean =
+ childrenResolved && checkInputDataTypes().isSuccess && names.forall(_.resolved)
override def inputTypes: Seq[AbstractDataType] = Seq.fill(children.length)(ArrayType)
@transient override lazy val dataType: DataType = {
- val fields = children.zip(arrayElementTypes).zipWithIndex.map {
- case ((expr: NamedExpression, elementType), _) =>
- StructField(expr.name, elementType, nullable = true)
- case ((_, elementType), idx) =>
- StructField(idx.toString, elementType, nullable = true)
+ val fields = arrayElementTypes.zip(names).map {
+ case (elementType, Literal(name, StringType)) =>
+ StructField(name.toString, elementType, nullable = true)
+ case _ =>
+ throw new IllegalStateException("Schema name of arrays_zip should be string literal.")
}
ArrayType(StructType(fields), containsNull = false)
}
@@ -322,6 +342,12 @@ case class ArraysZip(children: Seq[Expression]) extends Expression with ExpectsI
override def prettyName: String = "arrays_zip"
}
+object ArraysZip {
+ def apply(children: Seq[Expression]): ArraysZip = {
+ new ArraysZip(children)
+ }
+}
+
/**
* Returns an unordered array containing the values of the map.
*/
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
index 9baad7c..1c40215 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -17,6 +17,7 @@
package org.apache.spark.sql
+import java.io.File
import java.nio.charset.StandardCharsets
import java.sql.{Date, Timestamp}
@@ -24,7 +25,8 @@ import scala.util.Random
import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
+import org.apache.spark.sql.catalyst.expressions.{Alias, ArraysZip, AttributeReference, Expression, NamedExpression}
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
import org.apache.spark.sql.catalyst.plans.logical.OneRowRelation
import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.{withDefaultTimeZone, UTC}
@@ -552,6 +554,74 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession {
Row(Seq(Row(0, 1, 2, 3, 4, 5))))
}
+ test("SPARK-35876: arrays_zip should retain field names") {
+ withTempDir { dir =>
+ val df = spark.sparkContext.parallelize(
+ Seq((Seq(9001, 9002, 9003), Seq(4, 5, 6)))).toDF("val1", "val2")
+ val qualifiedDF = df.as("foo")
+
+ // Fields are UnresolvedAttribute
+ val zippedDF1 = qualifiedDF.select(arrays_zip($"foo.val1", $"foo.val2") as "zipped")
+ val maybeAlias1 = zippedDF1.queryExecution.logical.expressions.head
+ assert(maybeAlias1.isInstanceOf[Alias])
+ val maybeArraysZip1 = maybeAlias1.children.head
+ assert(maybeArraysZip1.isInstanceOf[ArraysZip])
+ assert(maybeArraysZip1.children.forall(_.isInstanceOf[UnresolvedAttribute]))
+ val file1 = new File(dir, "arrays_zip1")
+ zippedDF1.write.parquet(file1.getAbsolutePath)
+ val restoredDF1 = spark.read.parquet(file1.getAbsolutePath)
+ val fieldNames1 = restoredDF1.schema.head.dataType.asInstanceOf[ArrayType]
+ .elementType.asInstanceOf[StructType].fieldNames
+ assert(fieldNames1.toSeq === Seq("val1", "val2"))
+
+ // Fields are resolved NamedExpression
+ val zippedDF2 = df.select(arrays_zip(df("val1"), df("val2")) as "zipped")
+ val maybeAlias2 = zippedDF2.queryExecution.logical.expressions.head
+ assert(maybeAlias2.isInstanceOf[Alias])
+ val maybeArraysZip2 = maybeAlias2.children.head
+ assert(maybeArraysZip2.isInstanceOf[ArraysZip])
+ assert(maybeArraysZip2.children.forall(
+ e => e.isInstanceOf[AttributeReference] && e.resolved))
+ val file2 = new File(dir, "arrays_zip2")
+ zippedDF2.write.parquet(file2.getAbsolutePath)
+ val restoredDF2 = spark.read.parquet(file2.getAbsolutePath)
+ val fieldNames2 = restoredDF2.schema.head.dataType.asInstanceOf[ArrayType]
+ .elementType.asInstanceOf[StructType].fieldNames
+ assert(fieldNames2.toSeq === Seq("val1", "val2"))
+
+ // Fields are unresolved NamedExpression
+ val zippedDF3 = df.select(arrays_zip($"val1" as "val3", $"val2" as "val4") as "zipped")
+ val maybeAlias3 = zippedDF3.queryExecution.logical.expressions.head
+ assert(maybeAlias3.isInstanceOf[Alias])
+ val maybeArraysZip3 = maybeAlias3.children.head
+ assert(maybeArraysZip3.isInstanceOf[ArraysZip])
+ assert(maybeArraysZip3.children.forall(e => e.isInstanceOf[Alias] && !e.resolved))
+ val file3 = new File(dir, "arrays_zip3")
+ zippedDF3.write.parquet(file3.getAbsolutePath)
+ val restoredDF3 = spark.read.parquet(file3.getAbsolutePath)
+ val fieldNames3 = restoredDF3.schema.head.dataType.asInstanceOf[ArrayType]
+ .elementType.asInstanceOf[StructType].fieldNames
+ assert(fieldNames3.toSeq === Seq("val3", "val4"))
+
+ // Fields are neither UnresolvedAttribute nor NamedExpression
+ val zippedDF4 = df.select(arrays_zip(array_sort($"val1"), array_sort($"val2")) as "zipped")
+ val maybeAlias4 = zippedDF4.queryExecution.logical.expressions.head
+ assert(maybeAlias4.isInstanceOf[Alias])
+ val maybeArraysZip4 = maybeAlias4.children.head
+ assert(maybeArraysZip4.isInstanceOf[ArraysZip])
+ assert(maybeArraysZip4.children.forall {
+ case _: UnresolvedAttribute | _: NamedExpression => false
+ case _ => true
+ })
+ val file4 = new File(dir, "arrays_zip4")
+ zippedDF4.write.parquet(file4.getAbsolutePath)
+ val restoredDF4 = spark.read.parquet(file4.getAbsolutePath)
+ val fieldNames4 = restoredDF4.schema.head.dataType.asInstanceOf[ArrayType]
+ .elementType.asInstanceOf[StructType].fieldNames
+ assert(fieldNames4.toSeq === Seq("0", "1"))
+ }
+ }
+
def testSizeOfMap(sizeOfNull: Any): Unit = {
val df = Seq(
(Map[Int, Int](1 -> 1, 2 -> 2), "x"),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org