[SPARK-58735][SQL] Prune nested fields when computing size of an array of structs

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -111,6 +111,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
         OptimizeJoinCondition,
         LimitPushDown,
         LimitPushDownThroughWindow,
+        RewriteSizeOfArrayStruct,
         ColumnPruning,
         GenerateOptimization,
         // Operator combine

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala

@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, GetArrayItem, GetArrayStructFields, GetMapValue, GetStructField, MapKeys, MapValues, Size}
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, StructType}
+
+/**
+ * Computing the length of an array of structs (`size(arr)` / `array_size(arr)`) only requires the
+ * array's structural information (offsets / repetition levels), not the values of the element
+ * struct's fields. However, when `size` is applied to a whole `ARRAY<STRUCT<...>>` column, the
+ * column is referenced as a whole, which prevents [[NestedColumnAliasing]] (and the subsequent
+ * schema pruning at the file format reader) from pruning the unused nested fields. As a result,
+ * all nested fields are read from Parquet/ORC, causing large and unnecessary I/O (SPARK-58735).
+ *
+ * This rule rewrites `size(arr)` into `size(arr.<field>)`, picking a single (cheapest) field of the
+ * element struct. Because extracting a field from an array of structs preserves the array's length
+ * and null-ness, the result of `size` is unchanged, while the extra [[GetArrayStructFields]] lets
+ * the existing nested column pruning read only that one field.
+ *
+ * Example:
+ * {{{
+ *   size(events)  =>  size(events.<smallest_field>)
+ * }}}
+ */
+object RewriteSizeOfArrayStruct extends Rule[LogicalPlan] {
+
+  override def apply(plan: LogicalPlan): LogicalPlan = {
+    if (!SQLConf.get.nestedSchemaPruningEnabled) {
+      plan
+    } else {
+      // Only base (leaf) columns benefit: nested schema pruning happens at the scan, so rewriting
+      // `size` over an array produced by another operator (e.g. an aggregate) would be pure churn.
+      val baseAttrs = AttributeSet(plan.collectLeaves().flatMap(_.output))
+      plan.transformAllExpressions {
+        case s @ Size(child, _) if canRewrite(child, baseAttrs) =>
+          val array = child.dataType.asInstanceOf[ArrayType]
+          val struct = array.elementType.asInstanceOf[StructType]
+          // Pick the smallest field by default size, mirroring [[GenerateOptimization]]. Extracting
+          // any field preserves the array length, so the result of `size` is unchanged.
+          val (field, ordinal) =
+            struct.fields.zipWithIndex.minBy { case (f, _) => f.dataType.defaultSize }
+          val extractor = GetArrayStructFields(
+            child, field, ordinal, struct.length, array.containsNull || field.nullable)
+          s.withNewChildren(Seq(extractor))
+      }
+    }
+  }
+
+  /**
+   * We only rewrite when the child is an array of a struct with more than one field (with a single
+   * field there is nothing to prune), is rooted at a base (leaf) column so nested column pruning
+   * can actually prune it, and is not already a field extraction on an array of structs (which
+   * keeps this rule idempotent).
+   */
+  private def canRewrite(child: Expression, baseAttrs: AttributeSet): Boolean = {
+    !child.isInstanceOf[GetArrayStructFields] &&
+      rootAttribute(child).exists(baseAttrs.contains) && (child.dataType match {
+        case ArrayType(st: StructType, _) => st.length > 1
+        case _ => false
+      })
+  }
+
+  /**
+   * Returns the base attribute if the expression is built solely from an attribute and value
+   * extractors, i.e. it reads from a column that nested column pruning can prune; otherwise None.
+   */
+  private def rootAttribute(e: Expression): Option[Attribute] = e match {
+    case a: AttributeReference => Some(a)
+    case g: GetStructField => rootAttribute(g.child)
+    case g: GetArrayStructFields => rootAttribute(g.child)
+    case g: GetArrayItem => rootAttribute(g.child)
+    case g: GetMapValue => rootAttribute(g.child)
+    case m: MapValues => rootAttribute(m.child)
+    case m: MapKeys => rootAttribute(m.child)
+    case _ => None
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/NestedColumnAliasingSuite.scala

@@ -35,6 +35,7 @@ class NestedColumnAliasingSuite extends SchemaPruningTest {
 
   object Optimize extends RuleExecutor[LogicalPlan] {
     val batches = Batch("Nested column pruning", FixedPoint(100),
+      RewriteSizeOfArrayStruct,
       ColumnPruning,
       CollapseProject,
       RemoveNoopOperators) :: Nil
@@ -239,6 +240,46 @@ class NestedColumnAliasingSuite extends SchemaPruningTest {
     comparePlans(optimized, expected)
   }
 
+  test("SPARK-58735: size(array<struct>) prunes to a single element field") {
+    def collectArrayStructFields(plan: LogicalPlan): Seq[GetArrayStructFields] =
+      plan.flatMap(_.expressions.flatMap(_.collect { case g: GetArrayStructFields => g })).distinct
+
+    // `size(friends)` should be rewritten to read only one field of the element struct,
+    // so nested column pruning reads a single column instead of the whole `friends` struct.
+    val query = contact.select(Size($"friends", legacySizeOfNull = false)).analyze
+    val optimized = Optimize.execute(query)
+
+    // The rewrite preserves the user-visible output name (`size(friends)`); only the read
+    // schema changes.
+    val expected = contact
+      .select(Size(
+        GetArrayStructFields($"friends",
+          field = StructField("first", StringType),
+          ordinal = 0,
+          numFields = 3,
+          containsNull = true),
+        legacySizeOfNull = false).as("size(friends)"))
+      .analyze
+    comparePlans(optimized, expected)
+
+    // Only the single (first, atomic) field is extracted.
+    val extracted = collectArrayStructFields(optimized)
+    assert(extracted.map(_.field.name) == Seq("first"),
+      s"expected only `first` to be read, but got:\n$optimized")
+  }
+
+  test("SPARK-58735: size over array<primitive> / already-extracted field is not rewritten") {
+    // `friends.first` is already a single-field extraction (array<string>). The rule must not
+    // wrap it again (idempotence) and must not add any further extraction.
+    val alreadyExtracted = GetArrayStructFields($"friends",
+      field = StructField("first", StringType), ordinal = 0, numFields = 3, containsNull = true)
+    val query = contact.select(Size(alreadyExtracted, legacySizeOfNull = false)).analyze
+    val optimized = Optimize.execute(query)
+
+    val expected = contact.select(Size(alreadyExtracted, legacySizeOfNull = false)).analyze
+    comparePlans(optimized, expected)
+  }
+
   test("nested field pruning for getting struct field in map") {
     val field1 = GetStructField(GetMapValue($"relatives", Literal("key")), 0, Some("first"))
     val field2 = GetArrayStructFields(child = MapValues($"relatives"),

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala

@@ -177,6 +177,12 @@ abstract class SchemaPruningSuite
       Nil)
   }
 
+  testSchemaPruning("SPARK-58735: size of an array of structs reads a single nested field") {
+    val query = sql("select id, size(friends) from contacts where p=1")
+    checkScan(query, "struct<id:int,friends:array<struct<first:string>>>")
+    checkAnswer(query.orderBy("id"), Row(0, 1) :: Row(1, 0) :: Nil)
+  }
+
   testSchemaPruning("select a single complex field from a map entry and its parent map entry") {
     val query =
       sql("select relatives[\"brother\"].middle, relatives[\"brother\"] from contacts where p=1")