You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by su...@apache.org on 2026/06/08 17:45:31 UTC
(spark) branch branch-4.x updated: [SPARK-57176][SQL][4.X] Extend nested column pruning through array-returning functions
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new b34acedc43e3 [SPARK-57176][SQL][4.X] Extend nested column pruning through array-returning functions
b34acedc43e3 is described below
commit b34acedc43e3e64da49f8282b413d28fcbc4527f
Author: Chao Sun <ch...@openai.com>
AuthorDate: Mon Jun 8 10:45:17 2026 -0700
[SPARK-57176][SQL][4.X] Extend nested column pruning through array-returning functions
### Why are the changes needed?
[SPARK-57176](https://issues.apache.org/jira/browse/SPARK-57176) follows [SPARK-57022](https://issues.apache.org/jira/browse/SPARK-57022), which added nested column pruning for `transform` over `array<struct>` inputs.
Array-returning functions still retain the complete input element struct even when downstream expressions and lambdas only require a subset of nested fields. For example:
```sql
SELECT filter(friends, friend -> friend.last = 'Smith').first
FROM contacts
```
If `friends` contains `first`, `middle`, and `last`, Spark currently reads all three fields even though the query only requires `first` and `last`.
### What changes were proposed in this PR?
- Merge downstream result-field requirements with lambda requirements for `filter` and comparator-based `array_sort`.
- Propagate projected element schemas through `reverse`, `shuffle`, `slice`, and `array_compact`.
- Rewrite bound lambda variable types and nested field ordinals after pruning.
- Retain the complete element schema when the whole result is used, when a lambda consumes the whole element, or when default `array_sort` natural ordering requires the full struct.
Functions that inspect full element equality or natural ordering remain out of scope because dropping nested fields could change results.
### Does this PR introduce _any_ user-facing change?
Yes. Eligible queries using array-returning functions over arrays of structs can read a narrower input schema. Query results and SQL APIs are unchanged.
### How was this patch tested?
- `JAVA_HOME=/opt/homebrew/opt/openjdk17/libexec/openjdk.jdk/Contents/Home PATH=/opt/homebrew/opt/openjdk17/bin:$PATH build/sbt "catalyst/testOnly org.apache.spark.sql.catalyst.expressions.SchemaPruningSuite" "sql/testOnly org.apache.spark.sql.execution.datasources.parquet.ParquetV1SchemaPruningSuite org.apache.spark.sql.execution.datasources.parquet.ParquetV2SchemaPruningSuite org.apache.spark.sql.execution.datasources.orc.OrcV1SchemaPruningSuite org.apache.spark.sql.execution.dataso [...]
- `JAVA_HOME=/opt/homebrew/opt/openjdk17/libexec/openjdk.jdk/Contents/Home PATH=/opt/homebrew/opt/openjdk17/bin:$PATH build/sbt catalyst/scalastyle sql/scalastyle`
- `git diff --check`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #56345 from sunchao/dev/chao/codex/spark-array-returning-function-pruning-4.x.
Authored-by: Chao Sun <ch...@openai.com>
Signed-off-by: Chao Sun <ch...@openai.com>
---
.../expressions/ProjectionOverSchema.scala | 80 +++++---
.../sql/catalyst/expressions/SchemaPruning.scala | 106 ++++++++--
.../sql/catalyst/expressions/SelectedField.scala | 20 ++
.../catalyst/expressions/SchemaPruningSuite.scala | 219 +++++++++++++++++++++
.../execution/datasources/SchemaPruningSuite.scala | 143 ++++++++++++++
5 files changed, 526 insertions(+), 42 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ProjectionOverSchema.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ProjectionOverSchema.scala
index 362643016d83..d69da52ce4c6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ProjectionOverSchema.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ProjectionOverSchema.scala
@@ -69,30 +69,28 @@ case class ProjectionOverSchema(schema: StructType, output: AttributeSet) {
case GetMapValue(child, key) =>
getProjection(child).map { projection => GetMapValue(projection, key) }
case transform @ ArrayTransform(argument, lambda: LambdaFunction) =>
- getProjection(argument).map {
- case projection @ ArrayTypeProjection(projectedElementSchema) =>
- lambda.arguments.headOption match {
- case Some(elementVar: NamedLambdaVariable) =>
- // Pruning fields changes the physical ordinal layout of the element struct.
- // For example, pruning struct<a, b, c> to struct<a, c> moves c from ordinal 2
- // to ordinal 1, so rewrite both the variable type and its field accesses.
- val projectedElementVar = elementVar.copy(dataType = projectedElementSchema)
- val lambdaProjection =
- ProjectionOverLambdaVariable(elementVar, projectedElementVar)
- val projectedBody = lambda.function.transformDown {
- case lambdaProjection(expr) => expr
- }
- transform.copy(
- argument = projection,
- function = lambda.copy(
- function = projectedBody,
- arguments = projectedElementVar +: lambda.arguments.tail))
- case _ =>
- transform.copy(argument = projection)
- }
- case projection =>
- transform.copy(argument = projection)
+ projectArrayHigherOrderFunction(argument, lambda, numElementVariables = 1) {
+ (projection, projectedLambda) =>
+ transform.copy(argument = projection, function = projectedLambda)
}
+ case filter @ ArrayFilter(argument, lambda: LambdaFunction) =>
+ projectArrayHigherOrderFunction(argument, lambda, numElementVariables = 1) {
+ (projection, projectedLambda) =>
+ filter.copy(argument = projection, function = projectedLambda)
+ }
+ case sort @ ArraySort(argument, lambda: LambdaFunction, _) =>
+ projectArrayHigherOrderFunction(argument, lambda, numElementVariables = 2) {
+ (projection, projectedLambda) =>
+ sort.copy(argument = projection, function = projectedLambda)
+ }
+ case reverse @ Reverse(child) =>
+ getProjection(child).map(projection => reverse.copy(child = projection))
+ case shuffle @ Shuffle(child, _) =>
+ getProjection(child).map(projection => shuffle.copy(child = projection))
+ case slice @ Slice(x, start, length) if start.foldable && length.foldable =>
+ getProjection(x).map(projection => slice.copy(x = projection))
+ case knownNotContainsNull @ KnownNotContainsNull(child) =>
+ getProjection(child).map(projection => knownNotContainsNull.copy(child = projection))
case GetStructFieldObject(child, field: StructField) =>
getProjection(child).map(p => (p, p.dataType)).map {
case (projection, projSchema: StructType) =>
@@ -108,6 +106,39 @@ case class ProjectionOverSchema(schema: StructType, output: AttributeSet) {
None
}
+ private def projectArrayHigherOrderFunction(
+ argument: Expression,
+ lambda: LambdaFunction,
+ numElementVariables: Int)(
+ rebuild: (Expression, LambdaFunction) => Expression): Option[Expression] = {
+ getProjection(argument).map {
+ case projection @ ArrayTypeProjection(projectedElementSchema) =>
+ val projectedArguments = lambda.arguments.zipWithIndex.map {
+ case (elementVar: NamedLambdaVariable, index) if index < numElementVariables =>
+ elementVar.copy(dataType = projectedElementSchema)
+ case (argument, _) =>
+ argument
+ }
+ val projectedBody =
+ lambda.arguments.zip(projectedArguments).foldLeft(lambda.function) {
+ case (body, (elementVar: NamedLambdaVariable, projectedElementVar:
+ NamedLambdaVariable)) if elementVar ne projectedElementVar =>
+ val lambdaProjection =
+ ProjectionOverLambdaVariable(elementVar, projectedElementVar)
+ body.transformDown {
+ case lambdaProjection(expr) => expr
+ }
+ case (body, _) =>
+ body
+ }
+ rebuild(
+ projection,
+ lambda.copy(function = projectedBody, arguments = projectedArguments))
+ case projection =>
+ rebuild(projection, lambda)
+ }
+ }
+
private object ArrayTypeProjection {
def unapply(expr: Expression): Option[StructType] = expr.dataType match {
case ArrayType(projectedElementSchema: StructType, _) => Some(projectedElementSchema)
@@ -118,6 +149,7 @@ case class ProjectionOverSchema(schema: StructType, output: AttributeSet) {
/**
* Rewrites references rooted at one bound lambda element to use its projected type and
* recomputes nested field ordinals against each projected struct in the access path.
+ * Bound lambda references are matched by exprId because they may be instantiated separately.
* This must support the same access paths collected by `SchemaPruning` for lambda variables;
* currently both sides support only `GetStructField` chains.
*/
@@ -127,7 +159,7 @@ case class ProjectionOverSchema(schema: StructType, output: AttributeSet) {
def unapply(expr: Expression): Option[Expression] = project(expr)
private def project(expr: Expression): Option[Expression] = expr match {
- case variable: NamedLambdaVariable if variable.semanticEquals(original) =>
+ case variable: NamedLambdaVariable if variable.exprId == original.exprId =>
Some(projected)
case GetStructFieldObject(child, field: StructField) =>
project(child).map { projection =>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala
index 2f99dd54f77a..59ff86071647 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala
@@ -141,29 +141,22 @@ object SchemaPruning extends SQLConfHelper {
private[catalyst] def getRootFields(expr: Expression): Seq[RootField] = {
expr match {
case ArrayTransform(argument, lambda: LambdaFunction) =>
- // Field accesses through the lambda variable are not directly rooted at the input
- // attribute. Convert them into a projected type for the transform argument so that
- // physical nested column pruning can see them.
- val nestedRootFields = lambda.arguments.headOption.collect {
- case elementVar: NamedLambdaVariable =>
- getArrayTransformRootField(argument, lambda.function, elementVar)
- }.flatten.toSeq.map(field => RootField(field, derivedFromAtt = false))
- if (nestedRootFields.nonEmpty) {
- nestedRootFields ++ getRootFields(lambda.function)
- } else {
- expr.children.flatMap(getRootFields)
- }
+ getArrayHigherOrderFunctionRootFields(expr, argument, lambda)
case att: Attribute =>
RootField(StructField(att.name, att.dataType, att.nullable, att.metadata),
derivedFromAtt = true) :: Nil
- case SelectedField(field) => RootField(field, derivedFromAtt = false) :: Nil
+ case SelectedField(field) =>
+ RootField(field, derivedFromAtt = false) +:
+ getArrayReturningHigherOrderFunctionRootFields(expr)
// Root field accesses by `IsNotNull` and `IsNull` are special cases as the expressions
// don't actually use any nested fields. These root field accesses might be excluded later
// if there are any nested fields accesses in the query plan.
case IsNotNull(SelectedField(field)) =>
- RootField(field, derivedFromAtt = false, prunedIfAnyChildAccessed = true) :: Nil
+ RootField(field, derivedFromAtt = false, prunedIfAnyChildAccessed = true) +:
+ getArrayReturningHigherOrderFunctionRootFields(expr)
case IsNull(SelectedField(field)) =>
- RootField(field, derivedFromAtt = false, prunedIfAnyChildAccessed = true) :: Nil
+ RootField(field, derivedFromAtt = false, prunedIfAnyChildAccessed = true) +:
+ getArrayReturningHigherOrderFunctionRootFields(expr)
case IsNotNull(_: Attribute) | IsNull(_: Attribute) =>
expr.children.flatMap(getRootFields).map(_.copy(prunedIfAnyChildAccessed = true))
case s: SubqueryExpression =>
@@ -175,7 +168,26 @@ object SchemaPruning extends SQLConfHelper {
}
}
- private def getArrayTransformRootField(
+ private def getArrayHigherOrderFunctionRootFields(
+ expr: Expression,
+ argument: Expression,
+ lambda: LambdaFunction): Seq[RootField] = {
+ // Field accesses through the lambda variable are not directly rooted at the input
+ // attribute. Convert them into a projected type for the array argument so that
+ // physical nested column pruning can see them.
+ val nestedRootFields = lambda.arguments.headOption.collect {
+ case elementVar: NamedLambdaVariable =>
+ getArrayHigherOrderFunctionRootField(argument, lambda.function, elementVar)
+ }.flatten.toSeq.map(field => RootField(field, derivedFromAtt = false))
+ if (nestedRootFields.nonEmpty) {
+ nestedRootFields ++ getRootFields(lambda.function) ++
+ getArrayReturningHigherOrderFunctionRootFields(argument)
+ } else {
+ expr.children.flatMap(getRootFields)
+ }
+ }
+
+ private def getArrayHigherOrderFunctionRootField(
argument: Expression,
function: Expression,
elementVar: NamedLambdaVariable): Option[StructField] = {
@@ -197,6 +209,58 @@ object SchemaPruning extends SQLConfHelper {
}
}
+ private def getArrayReturningHigherOrderFunctionRootFields(expr: Expression): Seq[RootField] = {
+ expr match {
+ case ArrayFilter(argument, lambda: LambdaFunction) =>
+ getArrayReturningHigherOrderFunctionRootFields(argument, lambda, numElementVariables = 1)
+ case ArraySort(argument, lambda: LambdaFunction, allowNullComparisonResult)
+ if !allowNullComparisonResult && lambda.function.nullable =>
+ // The strict null-comparator error includes the compared values, so pruning their
+ // element fields would change the observable error parameters.
+ getRootFields(argument) ++ getRootFields(lambda.function)
+ case ArraySort(argument, lambda: LambdaFunction, _) =>
+ getArrayReturningHigherOrderFunctionRootFields(argument, lambda, numElementVariables = 2)
+ case _ =>
+ expr.children.flatMap(getArrayReturningHigherOrderFunctionRootFields)
+ }
+ }
+
+ private def getArrayReturningHigherOrderFunctionRootFields(
+ argument: Expression,
+ lambda: LambdaFunction,
+ numElementVariables: Int): Seq[RootField] = {
+ val nestedRootFields = argument.dataType match {
+ case ArrayType(_: StructType, containsNull) =>
+ val elementVariables = lambda.arguments.take(numElementVariables).collect {
+ case elementVar: NamedLambdaVariable => elementVar
+ }
+ val selectedFields = elementVariables.map(collectLambdaVariableFields(lambda.function, _))
+ if (elementVariables.length == numElementVariables && selectedFields.forall(_.isDefined)) {
+ val fields = selectedFields.flatten.flatten
+ if (fields.nonEmpty) {
+ val mergedElementSchema = fields
+ .map(field => StructType(Array(field)))
+ .reduceLeft(_ merge _)
+ SelectedField.withDataType(
+ argument,
+ ArrayType(mergedElementSchema, containsNull)).toSeq match {
+ case Seq() => getRootFields(argument).map(_.field)
+ case fields => fields
+ }
+ } else {
+ Seq.empty
+ }
+ } else {
+ getRootFields(argument).map(_.field)
+ }
+ case _ =>
+ getRootFields(argument).map(_.field)
+ }
+ nestedRootFields.map(field => RootField(field, derivedFromAtt = false)) ++
+ getRootFields(lambda.function) ++
+ getArrayReturningHigherOrderFunctionRootFields(argument)
+ }
+
/**
* Collects statically identifiable nested fields read from `elementVar`.
*
@@ -205,6 +269,8 @@ object SchemaPruning extends SQLConfHelper {
* means the full element is required somewhere (for example, `x => struct(x.a, x)`), so it is
* not safe to prune the element struct.
*
+ * Bound lambda references are matched by exprId because they may be instantiated separately.
+ *
* Currently only `GetStructField` chains rooted at `elementVar` are collected; array or map
* traversal within the lambda conservatively requires the full element. Keep this set of
* supported paths in sync with `ProjectionOverLambdaVariable` in `ProjectionOverSchema`.
@@ -213,9 +279,13 @@ object SchemaPruning extends SQLConfHelper {
expr: Expression,
elementVar: NamedLambdaVariable): Option[Seq[StructField]] = {
expr match {
- case LambdaVariableField(field, variable) if variable.semanticEquals(elementVar) =>
+ case LambdaVariableField(field, variable) if variable.exprId == elementVar.exprId =>
Some(field :: Nil)
- case variable: NamedLambdaVariable if variable.semanticEquals(elementVar) =>
+ case IsNotNull(variable: NamedLambdaVariable) if variable.exprId == elementVar.exprId =>
+ Some(Seq.empty)
+ case IsNull(variable: NamedLambdaVariable) if variable.exprId == elementVar.exprId =>
+ Some(Seq.empty)
+ case variable: NamedLambdaVariable if variable.exprId == elementVar.exprId =>
None
case _ =>
expr.children.foldLeft(Option(Seq.empty[StructField])) {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala
index e36224e7d5c1..69dfdbfc9a08 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala
@@ -163,6 +163,26 @@ object SelectedField {
val opt = dataTypeOpt.map(dt => MapType(keyType, dt, valueContainsNull))
selectField(left, opt)
}
+ case ArrayFilter(argument, _: LambdaFunction) =>
+ selectField(argument, dataTypeOpt)
+ case ArraySort(argument, _: LambdaFunction, _) =>
+ selectField(argument, dataTypeOpt)
+ case Reverse(child) =>
+ selectField(child, dataTypeOpt)
+ case Shuffle(child, _) =>
+ selectField(child, dataTypeOpt)
+ case Slice(x, start, length) if start.foldable && length.foldable =>
+ selectField(x, dataTypeOpt)
+ case KnownNotContainsNull(child) =>
+ val ArrayType(_, containsNull) = child.dataType
+ val opt = dataTypeOpt.map {
+ case ArrayType(dataType, _) => ArrayType(dataType, containsNull)
+ case x =>
+ // This should not happen.
+ throw QueryCompilationErrors.dataTypeUnsupportedByClassError(
+ x, "KnownNotContainsNull")
+ }
+ selectField(child, opt)
case _ =>
None
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruningSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruningSuite.scala
index 9426ef91349e..8fae15316fc9 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruningSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruningSuite.scala
@@ -17,8 +17,12 @@
package org.apache.spark.sql.catalyst.expressions
+import java.util.concurrent.atomic.AtomicReference
+
import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.plans.SQLHelper
+import org.apache.spark.sql.catalyst.util.{ArrayData, GenericArrayData}
import org.apache.spark.sql.internal.SQLConf.CASE_SENSITIVE
import org.apache.spark.sql.types._
@@ -186,4 +190,219 @@ class SchemaPruningSuite extends SparkFunSuite with SQLHelper {
StructField("event", eventType, nullable = true),
derivedFromAtt = false)))
}
+
+ test("merge returned and lambda fields for array higher-order functions") {
+ val elementType = StructType.fromDDL("a int, b int, c int")
+ val eventType = StructType(Seq(
+ StructField("rules", ArrayType(elementType, containsNull = true))))
+ val event = AttributeReference("event", eventType)()
+ val argument = GetStructField(event, 0, Some("rules"))
+ val element = NamedLambdaVariable("x", elementType, nullable = true)
+ val predicate = LambdaFunction(
+ GreaterThan(GetStructField(element, 2, Some("c")), Literal(0)),
+ Seq(element))
+ val left = NamedLambdaVariable("left", elementType, nullable = true)
+ val right = NamedLambdaVariable("right", elementType, nullable = true)
+ val comparator = LambdaFunction(
+ Coalesce(Seq(
+ Subtract(
+ GetStructField(left, 2, Some("c")),
+ GetStructField(right, 2, Some("c"))),
+ Literal(0))),
+ Seq(left, right))
+
+ Seq(ArrayFilter(argument, predicate), ArraySort(argument, comparator)).foreach { function =>
+ val selected = GetArrayStructFields(
+ function,
+ elementType(0),
+ ordinal = 0,
+ numFields = elementType.length,
+ containsNull = true)
+ val rootFields = SchemaPruning.getRootFields(selected)
+ val prunedSchema = SchemaPruning.pruneSchema(
+ StructType(Seq(StructField("event", eventType))),
+ rootFields)
+
+ assert(prunedSchema === StructType.fromDDL(
+ "event struct<rules:array<struct<a:int,c:int>>>"))
+ }
+ }
+
+ test("match separately instantiated array lambda variables by exprId") {
+ val elementType = StructType.fromDDL("a int, b int, c int")
+ val sourceSchema = StructType(Seq(
+ StructField("rules", ArrayType(elementType, containsNull = true))))
+ val rules = AttributeReference("rules", ArrayType(elementType, containsNull = true))()
+
+ def selectFirstField(function: Expression): GetArrayStructFields = {
+ GetArrayStructFields(
+ function,
+ elementType(0),
+ ordinal = 0,
+ numFields = elementType.length,
+ containsNull = true)
+ }
+
+ def evaluateSelectedFields(selected: GetArrayStructFields): Seq[Int] = {
+ val rootFields = SchemaPruning.identifyRootFields(Seq(Alias(selected, "out")()), Seq.empty)
+ val prunedSchema = SchemaPruning.pruneSchema(sourceSchema, rootFields)
+ assert(prunedSchema === StructType.fromDDL("rules array<struct<a:int,c:int>>"))
+
+ val projected = ProjectionOverSchema(prunedSchema, AttributeSet(Seq(rules)))
+ .unapply(selected).get
+ val bound = BindReferences.bindReference(projected, Seq(rules))
+ val array = new GenericArrayData(Array[Any](InternalRow(1, 3), InternalRow(2, -1)))
+ val result = bound.eval(InternalRow(array)).asInstanceOf[ArrayData]
+ (0 until result.numElements()).map(result.getInt)
+ }
+
+ val filterArgument = NamedLambdaVariable("x", elementType, nullable = true)
+ val filterReference = filterArgument.copy(value = new AtomicReference[Any]())
+ val filter = ArrayFilter(
+ rules,
+ LambdaFunction(
+ GreaterThan(GetStructField(filterReference, 2, Some("c")), Literal(0)),
+ Seq(filterArgument)))
+ assert(evaluateSelectedFields(selectFirstField(filter)) === Seq(1))
+
+ val leftArgument = NamedLambdaVariable("left", elementType, nullable = true)
+ val rightArgument = NamedLambdaVariable("right", elementType, nullable = true)
+ val leftReference = leftArgument.copy(value = new AtomicReference[Any]())
+ val rightReference = rightArgument.copy(value = new AtomicReference[Any]())
+ val sort = ArraySort(
+ rules,
+ LambdaFunction(
+ Coalesce(Seq(
+ Subtract(
+ GetStructField(leftReference, 2, Some("c")),
+ GetStructField(rightReference, 2, Some("c"))),
+ Literal(0))),
+ Seq(leftArgument, rightArgument)),
+ allowNullComparisonResult = false)
+ assert(evaluateSelectedFields(selectFirstField(sort)) === Seq(2, 1))
+ }
+
+ test("retain full nested array elements for array-returning higher-order functions") {
+ val structType = StructType.fromDDL("a int, c int")
+ val arrayType = ArrayType(structType, containsNull = true)
+ val nestedArrayType = ArrayType(arrayType, containsNull = true)
+ val sourceSchema = StructType(Seq(StructField("rules", nestedArrayType)))
+ val rules = AttributeReference("rules", nestedArrayType)()
+
+ def selectFirstField(function: Expression): GetArrayStructFields = {
+ GetArrayStructFields(
+ GetArrayItem(function, Literal(0)),
+ structType(0),
+ ordinal = 0,
+ numFields = structType.length,
+ containsNull = true)
+ }
+
+ val element = NamedLambdaVariable("x", arrayType, nullable = true)
+ val elementC = GetArrayStructFields(
+ element,
+ structType(1),
+ ordinal = 1,
+ numFields = structType.length,
+ containsNull = true)
+ val filter = ArrayFilter(
+ rules,
+ LambdaFunction(GreaterThan(GetArrayItem(elementC, Literal(0)), Literal(0)), Seq(element)))
+
+ val left = NamedLambdaVariable("left", arrayType, nullable = true)
+ val right = NamedLambdaVariable("right", arrayType, nullable = true)
+ def firstC(variable: NamedLambdaVariable): Expression = {
+ GetArrayItem(
+ GetArrayStructFields(
+ variable,
+ structType(1),
+ ordinal = 1,
+ numFields = structType.length,
+ containsNull = true),
+ Literal(0))
+ }
+ val sort = ArraySort(
+ rules,
+ LambdaFunction(
+ Coalesce(Seq(Subtract(firstC(left), firstC(right)), Literal(0))),
+ Seq(left, right)),
+ allowNullComparisonResult = false)
+
+ Seq(filter, sort).foreach { function =>
+ val selected = Alias(selectFirstField(function), "out")()
+ val rootFields = SchemaPruning.identifyRootFields(Seq(selected), Seq.empty)
+ assert(SchemaPruning.pruneSchema(sourceSchema, rootFields) === sourceSchema)
+ }
+ }
+
+ test("do not prune ArrayFilter when the whole result is used") {
+ val elementType = StructType.fromDDL("a int, b int")
+ val eventType = StructType(Seq(
+ StructField("rules", ArrayType(elementType, containsNull = true))))
+ val event = AttributeReference("event", eventType)()
+ val element = NamedLambdaVariable("x", elementType, nullable = true)
+ val filtered = ArrayFilter(
+ GetStructField(event, 0, Some("rules")),
+ LambdaFunction(
+ GreaterThan(GetStructField(element, 0, Some("a")), Literal(0)),
+ Seq(element)))
+
+ val rootFields = SchemaPruning.getRootFields(filtered)
+
+ assert(rootFields.contains(
+ SchemaPruning.RootField(
+ StructField("event", eventType, nullable = true),
+ derivedFromAtt = false)))
+ }
+
+ test("do not prune strict ArraySort when the comparator can return null") {
+ val elementType = StructType.fromDDL("a int, b int")
+ val eventType = StructType(Seq(
+ StructField("rules", ArrayType(elementType, containsNull = true))))
+ val event = AttributeReference("event", eventType)()
+ val argument = GetStructField(event, 0, Some("rules"))
+ val left = NamedLambdaVariable("left", elementType, nullable = true)
+ val right = NamedLambdaVariable("right", elementType, nullable = true)
+ val comparator = LambdaFunction(Literal.create(null, IntegerType), Seq(left, right))
+ val sorted = ArraySort(argument, comparator, allowNullComparisonResult = false)
+ val selected = GetArrayStructFields(
+ sorted,
+ elementType(0),
+ ordinal = 0,
+ numFields = elementType.length,
+ containsNull = true)
+
+ val rootFields = SchemaPruning.getRootFields(selected)
+ val prunedSchema = SchemaPruning.pruneSchema(
+ StructType(Seq(StructField("event", eventType))),
+ rootFields)
+
+ assert(prunedSchema === StructType(Seq(StructField("event", eventType))))
+ }
+
+ test("retain input array nullability when pruning through KnownNotContainsNull") {
+ val elementType = StructType.fromDDL("a int, b int")
+ val eventType = StructType(Seq(
+ StructField("rules", ArrayType(elementType, containsNull = true))))
+ val event = AttributeReference("event", eventType)()
+ val element = NamedLambdaVariable("x", elementType, nullable = true)
+ val compacted = KnownNotContainsNull(ArrayFilter(
+ GetStructField(event, 0, Some("rules")),
+ LambdaFunction(IsNotNull(element), Seq(element))))
+ val selected = GetArrayStructFields(
+ compacted,
+ elementType(0),
+ ordinal = 0,
+ numFields = elementType.length,
+ containsNull = false)
+
+ val rootFields = SchemaPruning.getRootFields(selected)
+ val prunedSchema = SchemaPruning.pruneSchema(
+ StructType(Seq(StructField("event", eventType))),
+ rootFields)
+ val prunedEventType = prunedSchema("event").dataType.asInstanceOf[StructType]
+
+ assert(prunedEventType("rules").dataType ===
+ ArrayType(StructType.fromDDL("a int"), containsNull = true))
+ }
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
index 7f4b83ee342e..2b498cdf3c15 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
@@ -470,6 +470,149 @@ abstract class SchemaPruningSuite
Nil)
}
+ testSchemaPruning("select nested field returned by ArrayFilter") {
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(org.apache.spark.sql.functions.filter(
+ col("friends"), friend => friend.getField("last") === "Smith").getField("first"))
+
+ checkScan(query, "struct<friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query,
+ Row(Array("Susan")) ::
+ Row(Array.empty[String]) ::
+ Nil)
+ }
+
+ testSchemaPruning("do not prune ArrayFilter when the whole result is used") {
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(org.apache.spark.sql.functions.filter(
+ col("friends"), friend => friend.getField("last") === "Smith"))
+
+ checkScan(query, "struct<friends:array<struct<first:string,middle:string,last:string>>>")
+ checkAnswer(query,
+ Row(Array(Row("Susan", "Z.", "Smith"))) ::
+ Row(Array.empty[Row]) ::
+ Nil)
+ }
+
+ testSchemaPruning("select nested field returned by array functions under null checks") {
+ val expressions = Seq(
+ org.apache.spark.sql.functions.filter(
+ col("friends"), friend => friend.getField("last") === "Smith")
+ .getField("first").isNotNull,
+ array_sort(col("friends"), (left, right) =>
+ when(left.getField("last") < right.getField("last"), -1)
+ .when(left.getField("last") > right.getField("last"), 1)
+ .otherwise(0)).getField("first").isNull)
+
+ expressions.foreach { expression =>
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(expression)
+
+ checkScan(query, "struct<friends:array<struct<first:string,last:string>>>")
+ }
+ }
+
+ testSchemaPruning("select nested field returned by ArraySort") {
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(array_sort(col("friends"), (left, right) =>
+ when(left.getField("last") < right.getField("last"), -1)
+ .when(left.getField("last") > right.getField("last"), 1)
+ .otherwise(0)).getField("first"))
+
+ checkScan(query, "struct<friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query,
+ Row(Array("Susan")) ::
+ Row(Array.empty[String]) ::
+ Nil)
+ }
+
+ testSchemaPruning("do not prune default ArraySort when selecting a nested field") {
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(array_sort(col("friends")).getField("first"))
+
+ checkScan(query, "struct<friends:array<struct<first:string,middle:string,last:string>>>")
+ checkAnswer(query,
+ Row(Array("Susan")) ::
+ Row(Array.empty[String]) ::
+ Nil)
+ }
+
+ testSchemaPruning("select nested field returned by Array wrappers") {
+ val queries = Seq(
+ reverse(col("friends")).getField("first"),
+ shuffle(col("friends")).getField("first"),
+ slice(col("friends"), 1, 1).getField("first"))
+
+ queries.foreach { expression =>
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(expression)
+
+ checkScan(query, "struct<friends:array<struct<first:string>>>")
+ checkAnswer(query,
+ Row(Array("Susan")) ::
+ Row(Array.empty[String]) ::
+ Nil)
+ }
+ }
+
+ testSchemaPruning("do not prune through Slice with non-foldable bounds") {
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(
+ slice(col("friends"), col("employer.id") + 1, lit(1)).getField("first"),
+ col("employer.company.name"))
+
+ checkScan(query,
+ "struct<friends:array<struct<first:string,middle:string,last:string>>," +
+ "employer:struct<id:int,company:struct<name:string>>>")
+ checkAnswer(query,
+ Row(Array("Susan"), "abc") ::
+ Row(Array.empty[String], null) ::
+ Nil)
+ }
+
+ testSchemaPruning("select ArrayTransform over array-returning higher-order functions") {
+ val expressions = Seq(
+ transform(
+ org.apache.spark.sql.functions.filter(
+ col("friends"), friend => friend.getField("last") === "Smith"),
+ friend => friend.getField("first")),
+ transform(
+ array_sort(col("friends"), (left, right) =>
+ when(left.getField("last") < right.getField("last"), -1)
+ .when(left.getField("last") > right.getField("last"), 1)
+ .otherwise(0)),
+ friend => friend.getField("first")))
+
+ expressions.foreach { expression =>
+ val query = spark.table("contacts")
+ .where("p = 1")
+ .select(expression)
+
+ checkScan(query, "struct<friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query,
+ Row(Array("Susan")) ::
+ Row(Array.empty[String]) ::
+ Nil)
+ }
+ }
+
+ testSchemaPruning("select nested field returned by ArrayCompact with null elements") {
+ withDataSourceTable(organizations, "organizations") {
+ val query = spark.table("organizations")
+ .select(array_compact(col("team.members")).getField("id"))
+
+ checkScan(query, "struct<team:struct<members:array<struct<id:int>>>>")
+ checkAnswer(query, Row(Array(1, 0)) :: Nil)
+ }
+ }
+
testSchemaPruning("SPARK-34638: nested column prune on generator output") {
val query1 = spark.table("contacts")
.select(explode(col("friends")).as("friend"))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org