This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 3ab17271ec [CORE] Fix nested field rebinding after schema pruning
(#13478)
3ab17271ec is described below
commit 3ab17271eca2036449f1c47a8d032c4fa09d48c9
Author: kevinwilfong <[email protected]>
AuthorDate: Thu Oct 29 02:43:00 2026 -0700
[CORE] Fix nested field rebinding after schema pruning (#13478)
---
.../gluten/expression/StructFieldBindSuite.scala | 67 ++++++-
.../gluten/expression/ExpressionConverter.scala | 139 +++++++-------
...utenExpressionConverterSchemaPruningSuite.scala | 202 +++++++++++++++++++++
3 files changed, 338 insertions(+), 70 deletions(-)
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/expression/StructFieldBindSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/expression/StructFieldBindSuite.scala
index 2603e9e89b..7cb84c3ad0 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/expression/StructFieldBindSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/expression/StructFieldBindSuite.scala
@@ -16,13 +16,14 @@
*/
package org.apache.gluten.expression
-import org.apache.gluten.execution.{ProjectExecTransformer,
VeloxWholeStageTransformerSuite}
+import org.apache.gluten.execution.{FileSourceScanExecTransformer,
ProjectExecTransformer, VeloxWholeStageTransformerSuite}
import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
Expression}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
Expression, GetArrayStructFields}
import org.apache.spark.sql.catalyst.expressions.GetStructField
-import org.apache.spark.sql.execution.ProjectExec
-import org.apache.spark.sql.types.{IntegerType, LongType, StringType,
StructField, StructType}
+import org.apache.spark.sql.execution.{FileSourceScanExec, ProjectExec,
SQLExecution}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, IntegerType, LongType, MapType,
StringType, StructField, StructType}
class StructFieldBindSuite extends VeloxWholeStageTransformerSuite {
override protected val resourcePath: String = ""
@@ -34,7 +35,10 @@ class StructFieldBindSuite extends
VeloxWholeStageTransformerSuite {
.asInstanceOf[VeloxGetStructFieldTransformer]
private def readsStructField(expressions: Seq[Expression]): Boolean =
- expressions.exists(_.find(_.isInstanceOf[GetStructField]).isDefined)
+ expressions.exists(_.find {
+ case _: GetStructField | _: GetArrayStructFields => true
+ case _ => false
+ }.isDefined)
// The struct fields must be read by a native projection, or the result
comparison would not
// exercise the binding.
@@ -113,6 +117,59 @@ class StructFieldBindSuite extends
VeloxWholeStageTransformerSuite {
assert(e.getMessage.contains("ambiguous field a"), e.getMessage)
}
+ for (container <- Seq("map", "array", "array_fields")) {
+ test(s"native $container filter binds retained predicates to the pruned
schema") {
+ withSQLConf(
+ SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key -> "true",
+ SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ withTempPath {
+ path =>
+ val struct = "named_struct('unused0', id, 'unused1', id,
'unused2', id, 'value', id)"
+ val (column, field) = container match {
+ case "map" => (s"map('key', $struct)", "nested['key'].value")
+ case "array" => (s"array($struct)", "nested[0].value")
+ case "array_fields" => (s"array($struct)", "nested.value[0]")
+ }
+ spark.range(1, 4).selectExpr(s"$column AS
nested").write.parquet(path.toString)
+ val data = spark.read.parquet(path.toString)
+ def query: DataFrame = data.where(s"$field =
1").selectExpr(s"$field AS value")
+ var originalFilters = Seq.empty[Expression]
+ withSQLConf(
+ "spark.gluten.enabled" -> "false",
+ SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key -> "false") {
+ originalFilters = query.queryExecution.executedPlan.collect {
+ case scan: FileSourceScanExec => scan.dataFilters
+ }.flatten
+ }
+ val originalOrdinals = originalFilters.flatMap(_.collect {
+ case field: GetStructField => field.ordinal
+ case fields: GetArrayStructFields => fields.ordinal
+ })
+ assert(originalOrdinals.nonEmpty && originalOrdinals.forall(_ ==
3))
+ val df = query
+ assertStructReadOffloaded(df)
+ val plan = df.queryExecution.executedPlan
+ val scans = collect(plan) { case scan:
FileSourceScanExecTransformer => scan }
+ assert(scans.size == 1, plan)
+ val prunedStruct = scans.head.requiredSchema("nested").dataType
match {
+ case MapType(_, struct: StructType, _) => struct
+ case ArrayType(struct: StructType, _) => struct
+ }
+ assert(prunedStruct.fieldNames.toSeq == Seq("value"))
+ // Model the retained pre-pruning scan predicates, then execute
the native query.
+ val retainedPlan = plan.transform {
+ case scan: FileSourceScanExecTransformer =>
scan.copy(dataFilters = originalFilters)
+ }
+ val rows = SQLExecution.withNewExecutionId(df.queryExecution) {
+ retainedPlan.executeCollect()
+ }
+ assert(rows.map(_.getLong(0)).toSeq == Seq(1L))
+ }
+ }
+ }
+ }
+
test("same-named struct columns with different layouts in a join") {
// Both join sides carry a struct column `s`, with the fields in a
different order. The
// query reads every field, so nested column pruning keeps the whole
structs above the join
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
index 0464734c32..b43f54b555 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala
@@ -33,7 +33,7 @@ import org.apache.spark.sql.hive.HiveUDFTransformer
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
-import scala.collection.mutable.ArrayBuffer
+import scala.util.control.NonFatal
trait Transformable {
def getTransformer(childrenTransformers: Seq[ExpressionTransformer]):
ExpressionTransformer
@@ -41,6 +41,67 @@ trait Transformable {
object ExpressionConverter extends SQLConfHelper with Logging {
+ /**
+ * Rebinds extraction paths handled by Spark's ProjectionOverSchema,
including maps and arrays.
+ * HiveTableScanNestedColumnPruning applies the same rewrite during plan
construction.
+ */
+ private def bindToInputSchema(
+ expression: Expression,
+ inputAttributes: Seq[Attribute]): Expression = {
+ val input = AttributeSeq(Option(inputAttributes).getOrElse(Seq.empty))
+ try {
+ // Expression.references also includes synthetic attributes for lambda
variables, which
+ // are bound by their LambdaFunction rather than the operator's input
schema.
+ val attributes = AttributeSet(expression.collect { case a:
AttributeReference => a })
+ attributes.foldLeft(expression) {
+ case (bound, attribute) =>
+ val index = input.indexOf(attribute.exprId)
+ if (index < 0) {
+ throw new IllegalStateException(s"no input attribute with the
exprId of $attribute")
+ }
+ if (attribute.dataType != input.attrs(index).dataType) {
+ // Project one root at a time, using exprId to select its input
type. Inputs such as
+ // joins can contain same-named attributes with different nested
layouts.
+ // ProjectionOverSchema reads only the root field's name and
dataType; the wrapper's
+ // nullability and metadata do not affect the rewrite.
+ val schema = StructType(Seq(StructField(attribute.name,
input.attrs(index).dataType)))
+ val projectionOverSchema = ProjectionOverSchema(schema,
AttributeSet(Seq(attribute)))
+ bound.transformDown {
+ case projectionOverSchema(projected) =>
+ val field = projected match {
+ case f: GetStructField => Some((f.childSchema, f.ordinal))
+ case f: GetArrayStructFields =>
+ val struct =
f.child.dataType.asInstanceOf[ArrayType].elementType
+ .asInstanceOf[StructType]
+ Some((struct, f.ordinal))
+ case _ => None
+ }
+ field.foreach {
+ case (struct, ordinal) =>
+ if (ordinal < 0 || ordinal >= struct.length) {
+ throw new IllegalStateException(s"invalid field ordinal
$ordinal in $struct")
+ }
+ val name = struct(ordinal).name
+ if (struct.fields.count(_.name == name) > 1) {
+ throw new IllegalStateException(s"ambiguous field $name")
+ }
+ }
+ projected
+ }
+ } else {
+ // Preserve resolved ordinals for unchanged layouts, including
duplicate field names.
+ bound
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ val message = s"Failed to bind reference for $expression in " +
+ s"${input.attrs.mkString("[", ",", "]")}: ${e.getMessage}"
+ logWarning(message, e)
+ throw new UnsupportedOperationException(message, e)
+ }
+ }
+
def replaceWithExpressionTransformer(
exprs: Seq[Expression],
attributeSeq: Seq[Attribute]): Seq[ExpressionTransformer] = {
@@ -526,33 +587,26 @@ object ExpressionConverter extends SQLConfHelper with
Logging {
expr => replaceWithExpressionTransformer0(expr, attributeSeq,
expressionsMap)),
m)
case getStructField: GetStructField =>
- try {
- val bindRef =
- bindGetStructField(getStructField, attributeSeq)
- // Different backends may have different result.
-
BackendsApiManager.getSparkPlanExecApiInstance.genGetStructFieldTransformer(
- substraitExprName,
- replaceWithExpressionTransformer0(getStructField.child,
attributeSeq, expressionsMap),
- bindRef.ordinal,
- getStructField)
- } catch {
- case e: IllegalStateException =>
- // This situation may need developers to fix, although we just
throw the below
- // exception to let the corresponding operator fall back.
- throw new UnsupportedOperationException(
- s"Failed to bind reference for $getStructField: ${e.getMessage}")
- }
+ val boundStructField =
+ bindToInputSchema(getStructField,
attributeSeq).asInstanceOf[GetStructField]
+
BackendsApiManager.getSparkPlanExecApiInstance.genGetStructFieldTransformer(
+ substraitExprName,
+ replaceWithExpressionTransformer0(boundStructField.child,
attributeSeq, expressionsMap),
+ boundStructField.ordinal,
+ boundStructField)
case getArrayStructFields: GetArrayStructFields =>
+ val boundArrayStructFields =
+ bindToInputSchema(getArrayStructFields,
attributeSeq).asInstanceOf[GetArrayStructFields]
GenericExpressionTransformer(
substraitExprName,
Seq(
replaceWithExpressionTransformer0(
- getArrayStructFields.child,
+ boundArrayStructFields.child,
attributeSeq,
expressionsMap),
- LiteralTransformer(getArrayStructFields.ordinal)),
- getArrayStructFields
+ LiteralTransformer(boundArrayStructFields.ordinal)),
+ boundArrayStructFields
)
case t: StringTranslate =>
BackendsApiManager.getSparkPlanExecApiInstance.genStringTranslateTransformer(
@@ -977,49 +1031,4 @@ object ExpressionConverter extends SQLConfHelper with
Logging {
}
}
- private def bindGetStructField(
- structField: GetStructField,
- input: AttributeSeq): BoundReference = {
- val names = new ArrayBuffer[String]
- var root: Expression = structField
- while (root.isInstanceOf[GetStructField]) {
- val curField = root.asInstanceOf[GetStructField]
- val name = curField.childSchema.fields(curField.ordinal).name
- names += name
- root = root.asInstanceOf[GetStructField].child
- }
- // For map/array type, the reference is correct no matter
NESTED_SCHEMA_PRUNING_ENABLED or not
- if (!root.isInstanceOf[AttributeReference]) {
- return BoundReference(structField.ordinal, structField.dataType,
structField.nullable)
- }
- def cannotBind(reason: String): Nothing =
- throw new IllegalStateException(
- s"Couldn't bind $structField ($reason) in ${input.attrs.mkString("[",
",", "]")}")
- // Bind the root by exprId like BindReferences: the input can hold several
same-named struct
- // attributes with different layouts, e.g. a struct column of the same
name on both join sides.
- val rootAttr = root.asInstanceOf[AttributeReference]
- val rootIndex = input.indexOf(rootAttr.exprId)
- if (rootIndex < 0) {
- cannotBind(s"no input attribute with the exprId of ${rootAttr.name}")
- }
- val attribute = input.attrs(rootIndex)
- // Same layout as the one the expression was resolved against, so its own
ordinal is right.
- // It is also the only right answer when the struct has duplicate field
names.
- if (attribute.dataType == rootAttr.dataType) {
- return BoundReference(structField.ordinal, structField.dataType,
structField.nullable)
- }
- // The input struct type differs (e.g. nested nullability merged by a
union): look each level
- // up by name from the root down, and fall back instead of guessing on a
duplicate name.
- val (_, ordinal) = names.reverseIterator.foldLeft[(DataType,
Int)]((attribute.dataType, -1)) {
- case ((struct: StructType, _), name) =>
- struct.fields.indexWhere(_.name == name) match {
- case -1 => cannotBind(s"no field $name")
- case idx if struct.fields.lastIndexWhere(_.name == name) != idx =>
- cannotBind(s"ambiguous field $name")
- case idx => (struct.fields(idx).dataType, idx)
- }
- case ((other, _), name) => cannotBind(s"field $name looked up in
non-struct $other")
- }
- BoundReference(ordinal, structField.dataType, structField.nullable)
- }
}
diff --git
a/gluten-ut/test/src/test/scala/org/apache/gluten/expressions/GlutenExpressionConverterSchemaPruningSuite.scala
b/gluten-ut/test/src/test/scala/org/apache/gluten/expressions/GlutenExpressionConverterSchemaPruningSuite.scala
new file mode 100644
index 0000000000..2a024d0fe9
--- /dev/null
+++
b/gluten-ut/test/src/test/scala/org/apache/gluten/expressions/GlutenExpressionConverterSchemaPruningSuite.scala
@@ -0,0 +1,202 @@
+/*
+ * 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.gluten.expressions
+
+import org.apache.gluten.expression.{ExpressionConverter,
ExpressionTransformer}
+import org.apache.gluten.substrait.SubstraitContext
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.expressions.{And, ArrayTransform,
AttributeReference, Expression, GetArrayItem, GetArrayStructFields,
GetMapValue, GetStructField, IsNotNull, LambdaFunction, Literal,
NamedLambdaVariable}
+import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, MapType,
StringType, StructField, StructType}
+
+class GlutenExpressionConverterSchemaPruningSuite extends SparkFunSuite {
+ private val originalSchema = StructType(
+ Seq(StructField("unused", IntegerType), StructField("value", IntegerType)))
+
+ private def nestedType(container: String, schema: StructType): DataType =
container match {
+ case "struct" => schema
+ case "map" => MapType(StringType, schema)
+ case _ => ArrayType(schema)
+ }
+
+ private def extractField(container: String, attribute: Expression):
Expression = {
+ val schema = attribute.dataType match {
+ case struct: StructType => struct
+ case MapType(_, struct: StructType, _) => struct
+ case ArrayType(struct: StructType, _) => struct
+ }
+ val ordinal = schema.fieldIndex("value")
+ container match {
+ case "struct" => GetStructField(attribute, ordinal, Some("value"))
+ case "map" =>
+ GetStructField(GetMapValue(attribute, Literal("key")), ordinal,
Some("value"))
+ case "array" =>
+ GetStructField(
+ GetArrayItem(attribute, Literal(0), failOnError = false),
+ ordinal,
+ Some("value"))
+ case "array_fields" =>
+ GetArrayStructFields(
+ attribute,
+ schema("value"),
+ ordinal,
+ schema.length,
+ containsNull = true)
+ }
+ }
+
+ private def structOrdinals(transformer: ExpressionTransformer): Seq[Int] = {
+ val ordinal = transformer.original match {
+ case field: GetStructField =>
+ Seq(field.ordinal)
+ case fields: GetArrayStructFields =>
+ Seq(fields.ordinal)
+ case _ => Seq.empty
+ }
+ ordinal ++ transformer.children.flatMap(structOrdinals)
+ }
+
+ for (container <- Seq("struct", "map", "array", "array_fields")) {
+ test(s"$container extraction preserves lambda variable binding") {
+ val variable =
+ NamedLambdaVariable("value", nestedType(container, originalSchema),
nullable = true)
+ val input = AttributeReference("nested", ArrayType(variable.dataType))()
+ val field = extractField(container, variable)
+ assert(field.references.exists(_.exprId == variable.exprId))
+ val expression = ArrayTransform(input, LambdaFunction(field,
Seq(variable)))
+ val transformer =
ExpressionConverter.replaceWithExpressionTransformer(expression, Seq(input))
+ assert(structOrdinals(transformer) == Seq(1))
+ transformer.doTransform(new SubstraitContext).toProtobuf
+ }
+
+ for (pruning <- Seq(false, true)) {
+ test(s"rebind $container extraction: nested pruning=$pruning") {
+ val original = AttributeReference("nested", nestedType(container,
originalSchema))()
+ val schema = if (pruning) StructType(Seq(originalSchema("value")))
else originalSchema
+ val input = original.withDataType(nestedType(container, schema))
+ val predicate = IsNotNull(extractField(container, original))
+ val rewritten = IsNotNull(extractField(container, input))
+ val expectedOrdinal = if (pruning) 0 else 1
+ val transformers =
ExpressionConverter.replaceWithExpressionTransformer(
+ Seq(predicate),
+ Seq(input))
+ assert(transformers.flatMap(structOrdinals) == Seq(expectedOrdinal))
+ if (container == "array_fields") {
+ val fields =
transformers.head.children.head.original.asInstanceOf[GetArrayStructFields]
+ assert(fields.numFields == schema.length)
+ }
+
+ // Both versions share an exprId. Check mixed conditions in both
orders.
+ for (condition <- Seq(And(predicate, rewritten), And(rewritten,
predicate))) {
+ val transformer =
+ ExpressionConverter.replaceWithExpressionTransformer(condition,
Seq(input))
+ assert(structOrdinals(transformer) == Seq(expectedOrdinal,
expectedOrdinal))
+ transformer.doTransform(new SubstraitContext).toProtobuf
+ }
+ }
+ }
+
+ for (pruning <- Seq(false, true)) {
+ test(
+ s"ambiguous $container field falls back: pruning=$pruning") {
+ val duplicates = StructType(
+ Seq(
+ StructField("value", IntegerType),
+ StructField("value", IntegerType),
+ StructField("unused", IntegerType)))
+ val schema = if (pruning) {
+ StructType(duplicates.fields.take(2))
+ } else {
+ StructType(duplicates.fields.map(_.copy(nullable = false)))
+ }
+ val original = AttributeReference("nested", nestedType(container,
duplicates))()
+ val expression = extractField(container, original)
+ val input = Seq(original.withDataType(nestedType(container, schema)))
+ val failure = intercept[UnsupportedOperationException] {
+ ExpressionConverter.replaceWithExpressionTransformer(expression,
input)
+ }
+ assert(failure.getMessage.contains("ambiguous field value"))
+
assert(!ExpressionConverter.canReplaceWithExpressionTransformer(expression,
input))
+ }
+ }
+
+ test(s"missing $container field triggers unsupported-expression fallback")
{
+ val original = AttributeReference("nested", nestedType(container,
originalSchema))()
+ val expression = extractField(container, original)
+ val input = Seq(original.withDataType(
+ nestedType(container, StructType(Seq(originalSchema("unused"))))))
+ val failure = intercept[UnsupportedOperationException] {
+ ExpressionConverter.replaceWithExpressionTransformer(expression, input)
+ }
+ assert(failure.getCause != null)
+ intercept[UnsupportedOperationException] {
+ ExpressionConverter.replaceWithExpressionTransformer(Seq(expression),
input)
+ }
+
assert(!ExpressionConverter.canReplaceWithExpressionTransformer(expression,
input))
+ }
+ }
+
+ test("expression binding preserves duplicate field ordinals when the layout
is unchanged") {
+ val schema =
+ StructType(Seq(StructField("duplicate", IntegerType),
StructField("duplicate", IntegerType)))
+ val attribute = AttributeReference("nested", schema)()
+ for (ordinal <- Seq(0, 1)) {
+ val expression = GetStructField(attribute, ordinal, Some("duplicate"))
+ val transformer = ExpressionConverter.replaceWithExpressionTransformer(
+ expression,
+ Seq(attribute.withNullability(false)))
+ assert(structOrdinals(transformer) == Seq(ordinal))
+ }
+ }
+
+ test("expression binding selects same-named input attributes by exprId") {
+ val schema =
+ StructType(Seq(StructField("unused", IntegerType), StructField("value",
IntegerType)))
+ val left = AttributeReference("nested", schema)()
+ val right = AttributeReference("nested", schema)()
+ val input = Seq(left.withDataType(StructType(Seq(schema("value")))), right)
+ val expressions = Seq(left, right).map(GetStructField(_, 1, Some("value")))
+ val transformers =
ExpressionConverter.replaceWithExpressionTransformer(expressions, input)
+ assert(transformers.map(structOrdinals) == Seq(Seq(0), Seq(1)))
+ }
+
+ test("expression binding reports a missing input attribute as unsupported") {
+ val schema = StructType(Seq(StructField("value", IntegerType)))
+ val attribute = AttributeReference("nested", schema)()
+ val expression = GetStructField(attribute, 0, Some("value"))
+ val failure = intercept[UnsupportedOperationException] {
+ ExpressionConverter.replaceWithExpressionTransformer(expression,
Seq.empty)
+ }
+ assert(failure.getMessage.contains("no input attribute"))
+ assert(failure.getMessage.contains("in []"))
+
assert(!ExpressionConverter.canReplaceWithExpressionTransformer(expression,
Seq.empty))
+ }
+
+ test("converter rejects null input for attribute references") {
+ val schema = StructType(Seq(StructField("value", IntegerType)))
+ val expression = GetStructField(AttributeReference("nested", schema)(), 0,
Some("value"))
+ intercept[UnsupportedOperationException] {
+ ExpressionConverter.replaceWithExpressionTransformer(expression, null)
+ }
+ intercept[UnsupportedOperationException] {
+ ExpressionConverter.replaceWithExpressionTransformer(Seq(expression),
null)
+ }
+ // Literal translation does not require input attributes.
+ ExpressionConverter.replaceWithExpressionTransformer(Literal(1), null)
+ ExpressionConverter.replaceWithExpressionTransformer(Seq(Literal(1)), null)
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
