[SPARK-59981][SQL] Derive the partition-key ordering when a V2 scan's reported ordering keeps no sort order

docs/sql-migration-guide.md

@@ -45,7 +45,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 but no explicit ordering now also reports itself sorted by its partition key expressions, because every row of such a partition evaluates them to the same value. 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 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.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, 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`.
 - Since Spark 4.4, when `spark.sql.groupByOrdinal` is true (the default), the pipe `AGGREGATE` operator rejects an aliased integer literal in its `GROUP BY` clause, such as `TABLE t |> AGGREGATE sum(x) AS s GROUP BY 2 AS grp`, with `UNSUPPORTED_FEATURE.PIPE_OPERATOR_AGGREGATE_UNSUPPORTED_CASE`, because a bare integer literal in the same position is an ordinal. Previously the aliased literal was grouped as a constant. Bare ordinals and aliases on other grouping expressions are not affected. To restore the previous behavior, set `spark.sql.legacy.allowAliasedPipeAggregateOrdinals` 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 but no explicit ordering, or an ordering that Spark ignores because it references a column that cannot be resolved. 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 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.
       </td>
       <td>4.2.0</td>
     </tr>

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsReportOrdering.java

@@ -44,6 +44,11 @@ public interface SupportsReportOrdering extends Scan {
    * avoid a sort, whether or not it is a partition key. Nor does it rely on the sort orders after
    * it. The exception is a sort order on a partition key that is not a transform, when Spark uses
    * the reported {@code KeyGroupedPartitioning}.
+   * <p>
+   * If Spark relies on none of the reported sort orders, for example when this returns an empty
+   * array, it may treat each partition as sorted by the partition keys that are not transforms.
+   * This needs a reported {@code KeyGroupedPartitioning}, whose partitions each hold a single key
+   * value. {@code spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled} controls it.
    */
   SortOrder[] outputOrdering();
 }

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

@@ -2680,11 +2680,13 @@ object SQLConf {
   val V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED =
     buildConf("spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled")
       .doc("When enabled, Spark derives output ordering from the partition key expressions of " +
-        "a V2 data source that reports a KeyedPartitioning but does not report explicit ordering " +
-        "via SupportsReportOrdering, or reports one that Spark ignores because it references a " +
-        "column that cannot be resolved. 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 " +
+        "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 " +
+        "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 " +
         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

@@ -139,38 +139,57 @@ trait DataSourceV2ScanExecBase
    * reference columns pruned out of the scan output, while a consumer binds it against the
    * output, e.g. the k-way merge of `GroupPartitionsExec`. Ordering is prefix-based, so the
    * leading run of sort orders over the output is kept and the rest is dropped, except the sort
-   * orders on a partition key when the output partitioning is a `KeyedPartitioning`: each
-   * partition holds a single key, so those still hold. Otherwise, when the output partitioning
-   * is a `KeyedPartitioning` and `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled`
-   * is on, each partition contains rows where the key expressions evaluate to a single constant
-   * value, so the data is trivially sorted by those expressions within the partition.
+   * orders on a partition key when the output partitioning is a `KeyedPartitioning`. Each
+   * partition holds a single key, so those still hold.
    *
-   * Either way, 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 a sort order over a pruned column.
-   * Dropping it loses nothing today when Spark can call the transform's function, since no
-   * operator then requires an ordering over the transform. The write path sorts by the function
-   * call instead. Keeping such a sort order would add comparisons nobody uses. Dropping it also
-   * handles a transform Spark cannot evaluate. Revisit this if an ordering over a transform
-   * becomes a real requirement.
+   * If nothing of the reported ordering is kept, the ordering is derived from the partition keys
+   * instead, when the output partitioning is a `KeyedPartitioning` and
+   * `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` is on. Each partition contains
+   * rows where the key expressions evaluate to a single constant value, so the data is trivially
+   * sorted by those expressions within the partition. Nothing is kept, for example, when the
+   * source reports no ordering, one Spark ignores, or an empty one.
+   *
+   * 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.
+   *
+   * 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
+   * a sort order over a pruned column. Dropping it loses nothing today when Spark can call the
+   * transform's function, since no operator then requires an ordering over the transform. The
+   * write path sorts by the function call instead. Keeping such a sort order would add
+   * comparisons nobody uses. Dropping it also handles a transform Spark cannot evaluate. Revisit
+   * this if an ordering over a transform becomes a real requirement.
    */
-  override def outputOrdering: Seq[SortOrder] = {
-    def holdsTransform(e: Expression): Boolean = e.exists(_.isInstanceOf[TransformExpression])
-    (ordering, outputPartitioning) match {
-      case (Some(o), p) =>
-        val (prefix, rest) =
-          o.span(order => order.references.subsetOf(outputSet) && !holdsTransform(order.child))
-        p match {
-          case k: KeyedPartitioning if rest.nonEmpty =>
-            val keyExprs = ExpressionSet(k.expressions.filterNot(holdsTransform))
-            prefix ++ rest.filter(order => keyExprs.contains(order.child))
-          case _ => prefix
-        }
-      case (_, k: KeyedPartitioning) if conf.v2BucketingPartitionKeyOrderingEnabled =>
-        k.expressions.filterNot(holdsTransform).map(SortOrder(_, Ascending))
-      case _ => Seq.empty
+  override def outputOrdering: Seq[SortOrder] =
+    if (keptOrdering.isEmpty && conf.v2BucketingPartitionKeyOrderingEnabled) {
+      derivedOrdering
+    } else {
+      keptOrdering
     }
+
+  /**
+   * What `outputOrdering` keeps of the reported ordering. A lazy val because `ordering`,
+   * `outputSet` and `outputPartitioning` do not change for an instance, so only the config is
+   * read on each call of `outputOrdering`.
+   */
+  @transient private lazy val keptOrdering: Seq[SortOrder] = {
+    val (prefix, rest) = ordering.getOrElse(Nil)
+      .span(order => order.references.subsetOf(outputSet) && !holdsTransform(order.child))
+    val keyExprs = ExpressionSet(derivedOrdering.map(_.child))
+    prefix ++ rest.filter(order => keyExprs.contains(order.child))
+  }
+
+  /** The ordering over the partition keys that `outputOrdering` can derive. */
+  @transient private lazy val derivedOrdering: Seq[SortOrder] = outputPartitioning match {
+    case k: KeyedPartitioning =>
+      k.expressions.filterNot(holdsTransform).map(SortOrder(_, Ascending))
+    case _ => Nil
   }
 
+  private def holdsTransform(e: Expression): Boolean = e.exists(_.isInstanceOf[TransformExpression])
+
   override def supportsColumnar: Boolean = {
     scan.columnarSupportMode() match {
       case Scan.ColumnarSupportMode.PARTITION_DEFINED =>

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

@@ -419,8 +419,8 @@ case class GroupPartitionsExec(
   override def outputOrdering: Seq[SortOrder] = {
     if (!hasCoalescing) {
       // No coalescing: each output partition is exactly one input partition. The child's
-      // within-partition ordering is fully preserved (including any key-derived ordering that
-      // `DataSourceV2ScanExecBase` already prepended).
+      // within-partition ordering is fully preserved, including an ordering that
+      // `DataSourceV2ScanExecBase` derives from the partition keys.
       child.outputOrdering
     } else if (usesSortedMerge) {
       // Coalescing with sorted merge: SortedMergeCoalescedRDD performs a k-way merge by the

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

@@ -95,15 +95,16 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with AliasHelper
           log"Spark ignores the ordering reported by ${MDC(RELATION_NAME, relation.name)} " +
             log"(scan ${MDC(CLASS_NAME, scan.getClass.getName)}) because the ordering columns " +
             log"cannot be resolved: ${MDC(COLUMN_NAMES, unresolvedColumns.mkString(", "))}.")
-        // Use None, not Some(Nil): only None lets DataSourceV2ScanExecBase.outputOrdering derive an
-        // ordering from a kept key-grouped partitioning, which does not depend on the dropped
-        // report.
+        // `None` and `Some(Nil)` mean the same to DataSourceV2ScanExecBase.outputOrdering, which
+        // can still derive an ordering from a kept key-grouped partitioning.
         d.copy(ordering = None)
       } else {
         // The ordering is kept as reported, even where it references columns pruned out of the
         // scan output: truncating it here would also drop the sort orders on a partition key past a
         // pruned column, which still hold. `DataSourceV2ScanRelation.doCanonicalize` and
-        // `DataSourceV2ScanExecBase.outputOrdering` restrict it to the scan output instead.
+        // `DataSourceV2ScanExecBase.outputOrdering` restrict it to the scan output instead, by
+        // 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