[SPARK-57176][SQL][4.X] Extend nested column pruning through array-returning functions

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 =>

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])) {

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
     }