From 4c9133a34f8c146e10b9b65b44901b8ee242e8a0 Mon Sep 17 00:00:00 2001 From: Hemanth Boyina Date: Tue, 1 Sep 2026 23:20:34 +0530 Subject: [PATCH 1/2] [SPARK-58735][SQL] Prune nested fields when computing size of an array of structs --- .../sql/catalyst/optimizer/Optimizer.scala | 1 + .../optimizer/RewriteSizeOfArrayStruct.scala | 92 +++++++++++++++++++ .../optimizer/NestedColumnAliasingSuite.scala | 41 +++++++++ .../datasources/SchemaPruningSuite.scala | 6 ++ 4 files changed, 140 insertions(+) create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 58b6479ddee30..760a9957f046d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/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 diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala new file mode 100644 index 0000000000000..cb8e0fc545875 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala @@ -0,0 +1,92 @@ +/* + * 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.{AttributeReference, 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>` 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.)`, 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.) + * }}} + */ +object RewriteSizeOfArrayStruct extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!SQLConf.get.nestedSchemaPruningEnabled) { + plan + } else { + plan.transformAllExpressions { + case s @ Size(child, _) if canRewrite(child) => + 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) that is rooted at a column reference (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): Boolean = { + !child.isInstanceOf[GetArrayStructFields] && isColumnReference(child) && (child.dataType match { + case ArrayType(st: StructType, _) => st.length > 1 + case _ => false + }) + } + + /** + * Returns true if the expression is built solely from a base column reference and value + * extractors, i.e. it reads from a scan column that nested column pruning can prune. + */ + private def isColumnReference(e: Expression): Boolean = e match { + case _: AttributeReference => true + case g: GetStructField => isColumnReference(g.child) + case g: GetArrayStructFields => isColumnReference(g.child) + case g: GetArrayItem => isColumnReference(g.child) + case g: GetMapValue => isColumnReference(g.child) + case m: MapValues => isColumnReference(m.child) + case m: MapKeys => isColumnReference(m.child) + case _ => false + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/NestedColumnAliasingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/NestedColumnAliasingSuite.scala index 38cd25cf491a1..93b85470d35d1 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/NestedColumnAliasingSuite.scala +++ b/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) 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 / already-extracted field is not rewritten") { + // `friends.first` is already a single-field extraction (array). 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"), 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 e37f004bf1fde..578c82929d1a4 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 @@ -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>>") + 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") From d97627a657fe3cf206a69030159cebbf05d41ffa Mon Sep 17 00:00:00 2001 From: Hemanth Boyina Date: Wed, 2 Sep 2026 18:15:05 +0530 Subject: [PATCH 2/2] fix build issues --- .../optimizer/RewriteSizeOfArrayStruct.scala | 44 ++++++++++--------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala index cb8e0fc545875..e5c27cac41319 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSizeOfArrayStruct.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.catalyst.optimizer -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, GetArrayItem, GetArrayStructFields, GetMapValue, GetStructField, MapKeys, MapValues, Size} +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 @@ -47,8 +47,11 @@ object RewriteSizeOfArrayStruct extends Rule[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) => + 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 @@ -64,29 +67,30 @@ object RewriteSizeOfArrayStruct extends Rule[LogicalPlan] { /** * 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) that is rooted at a column reference (so nested column pruning - * can actually prune it), and is not already a field extraction on an array of structs (which + * 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): Boolean = { - !child.isInstanceOf[GetArrayStructFields] && isColumnReference(child) && (child.dataType match { - case ArrayType(st: StructType, _) => st.length > 1 - case _ => false - }) + 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 true if the expression is built solely from a base column reference and value - * extractors, i.e. it reads from a scan column that nested column pruning can prune. + * 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 isColumnReference(e: Expression): Boolean = e match { - case _: AttributeReference => true - case g: GetStructField => isColumnReference(g.child) - case g: GetArrayStructFields => isColumnReference(g.child) - case g: GetArrayItem => isColumnReference(g.child) - case g: GetMapValue => isColumnReference(g.child) - case m: MapValues => isColumnReference(m.child) - case m: MapKeys => isColumnReference(m.child) - case _ => false + 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 } }