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..e5c27cac41319 --- /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>` 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 { + // 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 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")