You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ur...@apache.org on 2026/09/08 09:21:33 UTC
(spark) branch master updated: [SPARK-58735][SQL] Prune nested fields when computing size of an array of structs
This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 3de625e179ba [SPARK-58735][SQL] Prune nested fields when computing size of an array of structs
3de625e179ba is described below
commit 3de625e179ba316ac68239f03173869876ad1852
Author: hemanthboyina <he...@apache.org>
AuthorDate: Tue Sep 8 09:21:13 2026 +0000
[SPARK-58735][SQL] Prune nested fields when computing size of an array of structs
**What changes were proposed in this pull request?**
This adds a new optimizer rule, RewriteSizeOfArrayStruct.
When a query calls size() (or array_size()) on a column that is an array of structs, the rule rewrites it to take the size of a single field of that struct instead of the whole array. For example,
size(events) becomes size of events with only one inner field selected. It always picks the smallest field of the struct, following the same approach the existing GenerateOptimization rule already uses for
the explode + count case.
The rewrite runs in the operator optimization batch, just before column pruning, and only when nested schema pruning is enabled.
**Why are the changes needed?**
size() only needs to know how many elements are in the array. It never looks at the values inside the struct. But today, when you write size() over a whole array-of-struct column, Spark treats it as "the
entire column is needed," so nested column pruning and the Parquet/ORC schema pruning cannot kick in. As a result, every field of the struct is read from disk, even though none of them are used.
The nested pruning machinery only prunes when it sees a field being accessed in the plan. A plain column reference has no field access, so it is read in full. By rewriting the expression to access one small
field, we give the pruner the signal it needs, and it reads just that one field.
**Does this PR introduce any user-facing change?**
No. The query results are unchanged; only the amount of data read from disk is reduced.
**How was this patch tested?**
Added unit tests in NestedColumnAliasingSuite that check the rewritten plan reads only a single nested field, and that arrays of primitives and already-pruned expressions are left untouched.
**Was this patch authored or co-authored using generative AI tooling?**
Generated-by: Yes, used Claude Code
Closes #58456 from hemanthboyina/SPARK-58735-size-array-pruning.
Authored-by: hemanthboyina <he...@apache.org>
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
---
.../spark/sql/catalyst/optimizer/Optimizer.scala | 1 +
.../optimizer/RewriteSizeOfArrayStruct.scala | 96 ++++++++++++++++++++++
.../optimizer/NestedColumnAliasingSuite.scala | 41 +++++++++
.../execution/datasources/SchemaPruningSuite.scala | 6 ++
4 files changed, 144 insertions(+)
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 a1b650fd315d..f1484ff02e15 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 000000000000..e5c27cac4131
--- /dev/null
+++ b/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
+ }
+}
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 38cd25cf491a..93b85470d35d 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<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"),
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 e37f004bf1fd..578c82929d1a 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<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")
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org