ohad-definity commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4365976379
##########
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:
Confirmed and fixed in abb931f4: the generator input now also keeps what a
comparator or predicate inside it reads, through the same
`getArrayReturningHigherOrderFunctionRootFields` the ordinary path uses
(exposed as `getGeneratorInputLambdaRootFields`). New tests in
`SchemaPruningSuite` on your data shape: `posexplode(array_sort(items))` with
the default comparator reads all of `a, b, c` and returns `(0,20,200),
(1,10,100)`; a comparator on `b` only still prunes to `<b, c>`; a whole-element
`filter` reads everything; `filter(items, x -> x.a > 1)` prunes to `<a, c>`.
All outer, on Parquet/ORC V1/V2.
##########
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:
Agreed, fixed in abb931f4. The length-only field access only decides what
the scan reads; once the scan is narrowed, `size()` in the matched chain reads
the narrowed array directly (`size(item.parts.pid)` back to
`size(item.parts)`), which covers both the no-pruning V2 case and a container
that is narrowed but still read. `RewriteSizeOfArrayStruct` runs only in the
operator optimization batches before early scan push-down, so nothing re-adds
the projection. Test: your `SELECT item, size(item.parts)` shape and a
partially pruned one, asserting no `size()` over a field projection and the
flag-off rows.
##########
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.
Review Comment:
@sunchao on the builder lifecycle you also raised: would a javadoc note on
`SupportsPushDownRequiredColumns` stating that the last `pruneColumns` call
before `build()` wins be acceptable, or would you prefer the V2 path to build
the scan only after the rewrite above it is validated? The latter needs the
types the connector will actually read before building, so it means
restructuring that path; happy to do either.
--
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]
