sunchao commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4363411119
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala:
##########
@@ -133,6 +134,53 @@ object SchemaPruning extends SQLConfHelper {
} ++ rootFields
}
+ /**
+ * Gets the root field that the array `input` of an `explode` or
`posexplode` needs, when the
+ * generator's element output `element` is only read by `exprs` through
nested field accesses.
+ * For example, for `explode(items) AS item` read as `item.a` and `item.b`,
with
+ * `items: array<struct<a, b, c>>`, it returns `items: array<struct<a, b>>`.
+ *
+ * `elementFields` are fields of the element read by other means, e.g. as
the input of a
+ * generator above this one, named like the element.
+ *
+ * If nothing of the element is read, only the array sizes are needed, so a
single smallest
+ * field is kept, as in
[[org.apache.spark.sql.catalyst.optimizer.GenerateOptimization]].
+ *
+ * If the element is used whole, `input` is needed whole, e.g. `items.parts`
for
+ * `explode(items.parts)`, which still narrows `items` to `parts`.
+ *
+ * `input` must be expressed over the scan's output. Returns None if `input`
is not an access
+ * path [[SelectedField]] can narrow.
+ */
+ def getGeneratorInputRootField(
+ input: Expression,
+ element: Attribute,
+ exprs: Seq[Expression],
+ elementFields: Seq[StructField] = Nil): Option[RootField] =
input.dataType match {
+ case ArrayType(elementType, containsNull) =>
+ // Rename the element so that its fields can be told apart from other
columns by name.
+ val name = s"${element.name}#${element.exprId.id}"
+ val renamed = exprs.map(_.transform {
+ case a: Attribute if a.exprId == element.exprId => a.withName(name)
+ })
+ val fields = identifyRootFields(Nil, renamed,
+ elementFields.map(f => RootField(f.copy(name = name), derivedFromAtt =
false)))
+ .filter(_.field.name == name)
+ val prunedElementType = (fields, elementType) match {
+ // not read: only the array sizes are needed
+ case (Seq(), StructType(structFields)) if structFields.length > 1 =>
+ StructType(Array(structFields.minBy(_.dataType.defaultSize)))
+ case (Seq(), _) => elementType
+ // used whole: the input is needed whole, which still narrows what it
is extracted from
+ case _ if fields.exists(_.derivedFromAtt) => elementType
+ case _ =>
+ pruneSchema(StructType(Array(StructField(name, elementType))),
fields).head.dataType
+ }
+ SelectedField.withDataType(input, ArrayType(prunedElementType,
containsNull))
Review Comment:
Fixed in
[2ef75d92c01](https://github.com/apache/spark/pull/59706/commits/2ef75d92c0165492b176053b101485c3a9dc6f4a).
Generator input requirements now include the existing filter/comparator
requirement collection before propagation through the chain. Added coverage for
default struct sorting, whole-element filtering, useful field-specific pruning,
captures (including a sibling of a lower generator element), and strict
comparator error values.
The original wrong-results example was reproduced against the compiled
`11a1fbbd` runtime and passes with the fix. All 728 focused Parquet/ORC V1/V2
and query-shape cases passed with the changed production objects and test suite
recompiled against that CI build; this was incremental validation rather than a
full rebuild.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -990,7 +991,26 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
}
def pruneColumns(plan: LogicalPlan): LogicalPlan = plan.transform {
+ case op @ SchemaPruning.GeneratorsOverScan(_, _, sHolder:
ScanBuilderHolder) =>
+ // Fields read from the generators' elements above them narrow the scan
below them. If the
+ // plan above cannot be re-resolved, the scan is pruned again on its own
below, which calls
+ // `pruneColumns` on the scan builder a second time.
+ SchemaPruning.pruneThroughGenerators(op, sHolder) { (scan,
otherProjects, inputRootFields) =>
+ Some(pruneScan(scan.projects, scan.filtersStayUp,
scan.filtersPushDown, sHolder,
Review Comment:
Fixed in
[2ef75d92c01](https://github.com/apache/spark/pull/59706/commits/2ef75d92c0165492b176053b101485c3a9dc6f4a).
Direct array-field projections under `Size` are removed after scan pruning and
successful re-resolution. The array already has the required schema and
preserves length/nullness, so `size()` can read it directly, including when the
relevant container remains whole or only an unrelated column was pruned.
The baseline probe confirmed the unnecessary projection and the fixed probe
confirms its removal. Regression tests check the optimized expression, exact
scan schema, results, empty/null arrays, and ANSI behavior. All 728 focused
cases passed; no elapsed-time benchmark is claimed. The PR description now
explains that the projection is temporary.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
