sunchao commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4363269349
##########
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:
[P1] Preserve lambda requirements inside the generator input
`SelectedField.withDataType` also walks through `ArrayFilter` and
`ArraySort`, but this helper only collects reads of the generator's output.
Unlike the ordinary `getRootFields` path, it never adds
`getArrayReturningHigherOrderFunctionRootFields(input)`. Whole-element
predicates/comparators can therefore silently change meaning when
`ProjectionOverSchema` rebinds their lambda variables to the narrower element
type.
For example, use a Parquet/ORC table `t` with `items:
array<struct<a:int,b:int,c:int>>` containing one row with `[{a:2,b:10,c:100},
{a:1,b:20,c:200}]`:
```sql
SELECT pos, item.b, item.c
FROM t
LATERAL VIEW OUTER posexplode(array_sort(items)) v AS pos, item
```
The default comparator must sort by `a` first, producing `(0,20,200)` and
`(1,10,100)`. This path requests only `items<b,c>`; both comparator variables
are rebuilt as `struct<b:int,c:int>`, so it sorts by `b` instead, producing
`(0,10,100)` and `(1,20,200)`. The position column makes this a changed result,
independently of final row ordering. Selecting two fields avoids the old
single-field generator rewrite. There is no inferred nonempty-array filter in
this example to retain the missing requirements.
The rebuilt plan remains resolved with the same output schema, so the
fallback does not catch this. The same omission affects whole-element `filter`
predicates and the full-value error parameters of strict nullable sort
comparators, which the existing collector explicitly protects.
Please include all reads performed by accepted input wrappers, or
conservatively decline those wrappers until supported. Add V1/V2
outer-generator regressions for default struct sorting and whole-element
filtering, plus a field-specific lambda that should still prune. This
counterexample is established by source tracing at `11a1fbbd`; I did not
execute it locally.
##########
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:
[P2] Avoid retaining the size projection when its array stays unpruned
This callback always returns `Some`, even when `pruneScan` returns the
original nested schema. Consequently the speculative `lengthOnly` rewrites are
retained on a successful no-op V2 path.
For a V2 Parquet/ORC table with `items:
array<struct<id:int,parts:array<struct<a:int,b:int>>>>`, consider:
```sql
SELECT item, size(item.parts)
FROM t
LATERAL VIEW OUTER explode(items) v AS item
```
Selecting `item` whole requires every field, so there is no nested scan
saving. However, the rewrite leaves `size(item.parts.a)`. `Size(item.parts)`
just reads `numElements()`; `GetArrayStructFields(item.parts, a, ...)` first
allocates an array and loops over every part, in both interpreted and generated
code. This adds O(number of parts) work and temporary storage per item to an
otherwise constant-time length read. The existing `RewriteSizeOfArrayStruct`
does not introduce this work because it excludes roots produced by generators.
Please preserve the original size expression when the relevant container
remains required whole or the connector does not narrow it, and add this V2
no-pruning control. Checking only whether *some* scan field changed would miss
the same unnecessary work when an unrelated field is pruned. This is a
source-established allocation/complexity regression at `11a1fbbd`, not a
measured elapsed-time slowdown.
--
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]
