[SPARK-60011][SQL] Keep a DSv2 nested-field ordering after schema pruning

docs/sql-migration-guide.md

@@ -46,7 +46,7 @@ license: |
 - Since Spark 4.4, when `array_repeat` or `array_insert` is asked to build an array larger than the maximum supported array length, generated code raises the same error as interpreted evaluation. `array_repeat` now fails with `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER` instead of the internal error `_LEGACY_ERROR_TEMP_2176`, and `array_insert` fails with `COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION` instead of `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER`, which named a `count` parameter that `array_insert` does not have. Both functions raise an error under exactly the same conditions as before; only the reported error condition changes.
 - Since Spark 4.4, `/*+ ... */` inside a bracketed comment is parsed as a nested comment, so its closing `*/` no longer closes the outer comment. For example, `/* note /*+ x */ SELECT 1` previously returned `1`, but now raises `UNCLOSED_BRACKETED_COMMENT`. Close the outer comment explicitly, for example `/* note /*+ x */ */ SELECT 1`.
 - Since Spark 4.4, `spark.sql.sources.v2.bucketing.partition.filter.enabled` defaults to `true`. In a storage-partitioned join, the partition key groups pushed down to both sides may now be narrowed to those that can produce output for the join type, instead of always taking the union of the two sides' groups: an inner or semi join may keep only the groups present on both sides, a join that keeps or tests every left row (left outer, left anti, left single, existence) keeps the left side's groups, a right outer join keeps the right side's, and a full outer join is unaffected; a CROSS join carrying an equality condition is treated like an inner join. The narrowing is skipped when either side's partitioning may contain unknown partition keys, as after a shuffle on one side, so such a join keeps the full union. A group that is dropped is not scanned at all, so a query may read fewer files. To restore the previous behavior, set `spark.sql.sources.v2.bucketing.partition.filter.enabled` to `false`.
-- Since Spark 4.4, `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` defaults to `true`. A V2 scan that reports a keyed partitioning and no ordering that Spark keeps for the scan output now also reports itself sorted by its partition key expressions, because every row of such a partition evaluates them to the same value. That includes a scan that reports an empty ordering, or one that starts with a sort order on a column pruned from the scan output or on a partition transform and has no later sort order on a partition key other than a transform. It relies on each partition of the source holding a single partition key value, as `KeyGroupedPartitioning` requires. Partition transforms such as `days(ts)` or `bucket(8, id)` are left out of that ordering. This lets Spark drop a `Sort` it would otherwise place above the scan. An aggregate that groups by a leading run of the other key expressions, in key order, may also be planned as a sort aggregate rather than a hash aggregate. This is because `spark.sql.execution.replaceHashWithSortAgg` keys off the reported ordering. To restore the previous behavior, set `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` to `false`.
+- Since Spark 4.4, `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` defaults to `true`. A V2 scan that reports a keyed partitioning and no ordering that Spark keeps for the scan output now also reports itself sorted by its partition key expressions, because every row of such a partition evaluates them to the same value. That includes a scan that reports an empty ordering, or one that starts with a sort order on a column pruned from the scan output, on a nested field that nested schema pruning drops or on a struct that it narrows, or on a partition transform and has no later sort order on a partition key other than a transform. It relies on each partition of the source holding a single partition key value, as `KeyGroupedPartitioning` requires. Partition transforms such as `days(ts)` or `bucket(8, id)` and a key on a struct that nested schema pruning narrows are left out of that ordering. This lets Spark drop a `Sort` it would otherwise place above the scan. An aggregate that groups by a leading run of the other key expressions, in key order, may also be planned as a sort aggregate rather than a hash aggregate. This is because `spark.sql.execution.replaceHashWithSortAgg` keys off the reported ordering. To restore the previous behavior, set `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` to `false`.
 - Since Spark 4.4, `spark.sql.sources.v2.bucketing.preserveKeyOrderingOnCoalesce.enabled` defaults to `true`. When `GroupPartitionsExec` merges several input partitions that share one partition key value into a single output partition, it now keeps sort orders over the partition key expressions in the ordering it reports, since those expressions are constant within the merged partition. Sort orders over partition transforms such as `days(ts)` or `bucket(8, id)` are left out. Orders over other columns are still dropped, as the concatenation invalidates them, and a join that reduced the partition keys onto a common key space (see `spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled`) reports no order, since the merged partitions then share only the reduced key. This can remove a `Sort` downstream of the merge. To restore the previous behavior, set `spark.sql.sources.v2.bucketing.preserveKeyOrderingOnCoalesce.enabled` to `false`.
 - Since Spark 4.4, when the `from_protobuf` option `unwrap.primitive.wrapper.types` is enabled, `from_protobuf` unwraps a present Protobuf primitive wrapper (such as `google.protobuf.Int32Value` or `StringValue`) whose inner `value` is unset or set to its default to the type's default value (`0`, `false`, `""`, or empty binary), regardless of the `emit.default.values` option. Previously such a wrapper unwrapped to `NULL` unless the `emit.default.values` option was enabled, which could also place a `NULL` into a non-nullable array element or map value. An absent wrapper field still unwraps to `NULL`.
 - Since Spark 4.4, the value columns of `UNPIVOT` and `Dataset.unpivot` merge the nested nullability (array element, map value, and struct field nullability) of all unpivoted values. Previously, a value column took the nested nullability of the first listed value, so a `NULL` array element, map value, or struct field of a later value could be returned as a non-`NULL` value such as `0`, and `IS NULL` on it returned `false`. A view created by an earlier Spark version over such an `UNPIVOT` may have recorded the non-nullable nested type in its schema and can fail to read; recreate it with `CREATE OR REPLACE VIEW`. To restore the previous behavior, set `spark.sql.legacy.unpivot.useFirstValueNestedNullability` to `true`.

docs/sql-performance-tuning.md

@@ -708,7 +708,7 @@ The following SQL properties enable Storage Partition Join in different join que
       <td><code>spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled</code></td>
       <td>true</td>
       <td>
-        When enabled, Spark derives the output ordering of a V2 scan from its partition key expressions, if the source reports a keyed partitioning and no ordering that Spark keeps for the scan output. For example, the source reports none, an empty one, or one that starts with a sort order on a column pruned from the scan output or on a partition transform and has no later sort order on a partition key other than a transform. All rows of such a partition share one key value, so the partition is trivially sorted by those expressions, and a sort Spark would otherwise add becomes unnecessary. Partition transforms such as <code>days(ts)</code> or <code>bucket(8, id)</code> are left out of the ordering. This config requires <code>spark.sql.sources.v2.bucketing.enabled</code> to be true.
+        When enabled, Spark derives the output ordering of a V2 scan from its partition key expressions, if the source reports a keyed partitioning and no ordering that Spark keeps for the scan output. For example, the source reports none, an empty one, or one that starts with a sort order on a column pruned from the scan output, on a nested field that nested schema pruning drops or on a struct that it narrows, or on a partition transform and has no later sort order on a partition key other than a transform. All rows of such a partition share one key value, so the partition is trivially sorted by those expressions, and a sort Spark would otherwise add becomes unnecessary. Partition transforms such as <code>days(ts)</code> or <code>bucket(8, id)</code> and a key on a struct that nested schema pruning narrows are left out of the ordering. This config requires <code>spark.sql.sources.v2.bucketing.enabled</code> to be true.
       </td>
       <td>4.2.0</td>
     </tr>

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -2682,11 +2682,13 @@ object SQLConf {
       .doc("When enabled, Spark derives output ordering from the partition key expressions of " +
         "a V2 data source that reports a KeyedPartitioning and no ordering that Spark keeps for " +
         "the scan output. For example, the source reports none, an empty one, or one that " +
-        "starts with a sort order on a column pruned from the scan output or on a partition " +
+        "starts with a sort order on a column pruned from the scan output, on a nested field " +
+        "that nested schema pruning drops or on a struct that it narrows, or on a partition " +
         "transform and has no later sort order on a partition key other than a transform. " +
         "Within a single partition all rows share the same key value, so the data is trivially " +
         "sorted by those expressions. Partition transforms such as `days(ts)` or `bucket(8, id)` " +
-        "are left out of the ordering. Requires " +
+        "and a key on a struct that nested schema pruning narrows are left out of the ordering. " +
+        "Requires " +
         s"${V2_BUCKETING_ENABLED.key} to be enabled.")
       .version("4.2.0")
       .withBindingPolicy(ConfigBindingPolicy.SESSION)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala

@@ -17,16 +17,20 @@
 
 package org.apache.spark.sql.execution.datasources.v2
 
+import scala.util.control.NonFatal
+
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Ascending, Expression, ExpressionSet, SortOrder, TransformExpression}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeMap, AttributeReference, Expression, ExpressionSet, GetArrayStructFields, ProjectionOverSchema, SortOrder, TransformExpression}
 import org.apache.spark.sql.catalyst.plans.physical
 import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
 import org.apache.spark.sql.catalyst.util.truncatedString
 import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, PartitionReaderFactory, Scan}
 import org.apache.spark.sql.execution.{ExplainUtils, LeafExecNode, SafeForKWayMerge}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.internal.connector.SupportsMetadata
+import org.apache.spark.sql.types.{ArrayType, StructType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.Utils
 
@@ -152,7 +156,7 @@ trait DataSourceV2ScanExecBase
    * A kept sort order wins over the derived ordering, since ordering is prefix-based. Putting the
    * keys in front would stop a kept `[data]` from satisfying a required `[data]`. So the result
    * can depend on the columns a query reads. A report `[s]` derives the key ordering only when `s`
-   * is pruned.
+   * is pruned, or narrowed by nested schema pruning.
    *
    * In both the reported and the derived ordering, a sort order that holds a partition transform
    * is dropped, even on a partition key. In a reported ordering it also ends the leading run, like
@@ -176,18 +180,60 @@ trait DataSourceV2ScanExecBase
    */
   @transient private lazy val keptOrdering: Seq[SortOrder] = {
     val (prefix, rest) = ordering.getOrElse(Nil)
-      .span(order => order.references.subsetOf(outputSet) && !holdsTransform(order.child))
+      .map(order => bindToOutput(order.child).map(child => order.copy(child = child)))
+      .span(_.exists(order => !holdsTransform(order.child)))
     val keyExprs = ExpressionSet(derivedOrdering.map(_.child))
-    prefix ++ rest.filter(order => keyExprs.contains(order.child))
+    prefix.flatten ++ rest.flatten.filter(order => keyExprs.contains(order.child))
   }
 
-  /** The ordering over the partition keys that `outputOrdering` can derive. */
+  /**
+   * The ordering over the partition keys that `outputOrdering` can derive. The keys are bound to
+   * the output like the reported ordering, see `bindToOutput`, while `outputPartitioning` keeps
+   * them as reported for the full-width `HasPartitionKey` rows.
+   */
   @transient private lazy val derivedOrdering: Seq[SortOrder] = outputPartitioning match {
     case k: KeyedPartitioning =>
-      k.expressions.filterNot(holdsTransform).map(SortOrder(_, Ascending))
+      k.expressions.filterNot(holdsTransform).flatMap(bindToOutput).map(SortOrder(_, Ascending))
     case _ => Nil
   }
 
+  /**
+   * `e` bound to this scan's output, or None when it reads a column or nested field the scan does
+   * not output, or when pruning narrows its value: rows sorted by a struct are not necessarily
+   * sorted by a narrower one. Nested schema pruning keeps an attribute's exprId but narrows its
+   * type, so field ordinals are recomputed. Attribute names are first normalized to the output's,
+   * since `ProjectionOverSchema` matches them by name and a source may report a key in another
+   * case. Like `DataSourceV2ScanRelation.outputBoundPushedFilters`, but it also normalizes the
+   * names, drops a narrowed value and drops any failure.
+   */
+  private def bindToOutput(e: Expression): Option[Expression] = {
+    try {
+      val bound = e
+        .transform { case a: AttributeReference => a.withName(outputNames.getOrElse(a, a.name)) }
+        .transformDown { case outputProjection(projected) => projected }
+        .transform {
+          // `ProjectionOverSchema` keeps the reported spelling of an array field, which, unlike
+          // that of a struct field, takes part in `semanticEquals`. Take the output's.
+          case g: GetArrayStructFields => g.child.dataType match {
+            case ArrayType(struct: StructType, _) =>
+              g.copy(field = g.field.copy(name = struct(g.ordinal).name))
+            case _ => g
+          }
+        }
+      Option.when(bound.references.subsetOf(outputSet) &&
+          DataTypeUtils.sameType(bound.dataType, e.dataType))(bound)
+    } catch {
+      // `FIELD_NOT_FOUND` for a nested field that pruning removed is the expected failure. Any
+      // other comes from a source whose `readSchema()` changes a column's type. An ordering is only
+      // an optimization, so it is dropped either way.
+      case NonFatal(_) => None
+    }
+  }
+
+  // Built only when a sort order or a partition key is bound, which most scans do not have.
+  @transient private lazy val outputNames = AttributeMap(output.map(a => a -> a.name))
+  @transient private lazy val outputProjection = ProjectionOverSchema(schema, outputSet)
+
   private def holdsTransform(e: Expression): Boolean = e.exists(_.isInstanceOf[TransformExpression])
 
   override def supportsColumnar: Boolean = {

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala

@@ -106,9 +106,9 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with AliasHelper
         // different rules. `doCanonicalize` only feeds plan equality, so a plain `takeWhile` that
         // drops the pruned columns is enough there.
         // A nested sort key such as `s.x` resolves to a new `Alias(GetStructField(...))` on every
-        // run of this rule. Strip the alias so repeated runs give equal plans and, when the scan
-        // reads the whole struct, the ordering can satisfy a required ordering on the same field
-        // (see SPARK-60011).
+        // run of this rule. Strip the alias so repeated runs give equal plans and the ordering can
+        // satisfy a required ordering on the same field. When nested schema pruning narrows the
+        // struct, `DataSourceV2ScanExecBase` rebinds the field to the pruned type.
         val ordering =
           V2ExpressionUtils.toCatalystOrdering(reportedOrdering, relation, relation.funCatalog)
             .map(o => o.copy(child = trimAliases(o.child)))