[SPARK-59248][SQL] Keep storage-partitioned join when partition keys are pruned from the scan output

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala

@@ -343,11 +343,21 @@ case class DataSourceV2ScanRelation(
         output = this.relation.output.map(QueryPlan.normalizeExpressions(_, this.relation.output))
       ),
       output = this.output.map(QueryPlan.normalizeExpressions(_, this.output)),
+      // keyGroupedPartitioning may reference columns pruned out of `output`, which is kept as long
+      // as any key survives. A pruned key carries no information for plan comparison, since the
+      // physical outputPartitioning projects it away, so drop it before normalizing; otherwise the
+      // dangling attribute's exprId would keep otherwise-equivalent scans unequal and defeat
+      // subplan merging.
       keyGroupedPartitioning = keyGroupedPartitioning.map(
-        _.map(QueryPlan.normalizeExpressions(_, output))
+        _.filter(_.references.subsetOf(outputSet))
+          .map(QueryPlan.normalizeExpressions(_, output))
       ),
+      // ordering may likewise reference columns pruned out of `output`. Ordering is prefix-based,
+      // so keep only the leading run of sort orders that reference output columns and drop the rest
+      // before normalizing, for the same reason as keyGroupedPartitioning above.
       ordering = ordering.map(
-        _.map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, output)))
+        _.takeWhile(_.references.subsetOf(outputSet))
+          .map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, output)))
       ),
       // pushedFilters may reference columns pruned out of `output` (see the field doc), so they are
       // normalized against the relation's full output rather than `output`.

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

@@ -66,12 +66,22 @@ case class BatchScanExec(
     case other: BatchScanExec =>
       this.batch != null && this.batch == other.batch &&
           this.runtimeFilters == other.runtimeFilters &&
-          this.keyGroupedPartitioning == other.keyGroupedPartitioning
+          this.prunedKeyGroupedPartitioning == other.prunedKeyGroupedPartitioning
     case _ =>
       false
   }
 
-  override def hashCode(): Int = Objects.hash(batch, runtimeFilters, keyGroupedPartitioning)
+  override def hashCode(): Int = Objects.hash(batch, runtimeFilters, prunedKeyGroupedPartitioning)
+
+  /**
+   * The reported partitioning keys restricted to those still present in `output`. A key may
+   * reference a column pruned out of the scan, which is kept as long as any key survives (see
+   * V2ScanPartitioningAndOrdering). Such a dangling key carries no information, since the physical
+   * outputPartitioning projects it away, so `doCanonicalize`, `equals` and `hashCode` all use
+   * this view to stay consistent about ignoring it.
+   */
+  @transient lazy val prunedKeyGroupedPartitioning: Option[Seq[Expression]] =
+    keyGroupedPartitioning.map(_.filter(_.references.subsetOf(outputSet)))
 
   @transient override lazy val inputPartitions: Seq[InputPartition] =
     batch.planInputPartitions().toImmutableArraySeq
@@ -83,7 +93,7 @@ case class BatchScanExec(
       runtimeFilters,
       table,
       output,
-      outputPartitioning,
+      reportedKeyedPartitioning,
       inputPartitions)
 
   override lazy val readerFactory: PartitionReaderFactory = batch.createReaderFactory()
@@ -106,8 +116,8 @@ case class BatchScanExec(
       runtimeFilters = QueryPlan.normalizePredicates(
         runtimeFilters.filterNot(_ == DynamicPruningExpression(Literal.TrueLiteral)),
         output),
-      keyGroupedPartitioning = keyGroupedPartitioning.map(_.map(
-        QueryPlan.normalizeExpressions(_, output))))
+      keyGroupedPartitioning = prunedKeyGroupedPartitioning.map(
+        _.map(QueryPlan.normalizeExpressions(_, output))))
   }
 
   override def simpleString(maxFields: Int): String = {

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

@@ -89,22 +89,48 @@ trait DataSourceV2ScanExecBase
        |""".stripMargin
   }
 
-  // A `lazy val` because the planner asks a node for its partitioning many times, and the
-  // key-grouped arm sorts every partition key, wraps each one and runs a `distinct` over them.
-  @transient override lazy val outputPartitioning: physical.Partitioning = {
+  /**
+   * The partitioning as the source reported it: one key per input partition, holding every
+   * reported key position, with the partitions in the order those full keys sort into. It is built
+   * from the raw, full-width `HasPartitionKey.partitionKey()` rows, so a consumer of those rows
+   * (`filteredPartitions`) must take the keys, the key types and the order from here, not from the
+   * possibly-projected `outputPartitioning`.
+   *
+   * A `lazy val` because this is the expensive half: it sorts every partition key, wraps each one
+   * and runs a `distinct` over them, and both `outputPartitioning` and `filteredPartitions` ask
+   * for it.
+   */
+  @transient protected lazy val reportedKeyedPartitioning: Option[KeyedPartitioning] = {
     keyGroupedPartitioning match {
       case Some(exprs) if conf.v2BucketingEnabled && KeyedPartitioning.supportsExpressions(exprs) &&
           inputPartitions.nonEmpty && inputPartitions.forall(_.isInstanceOf[HasPartitionKey]) =>
         val dataTypes = exprs.map(_.dataType)
         val rowOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes)
         val partitionKeys =
           inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering)
-        KeyedPartitioning(exprs, partitionKeys)
-      case _ =>
-        super.outputPartitioning
+        Some(KeyedPartitioning(exprs, partitionKeys))
+      case _ => None
     }
   }
 
+  // A `lazy val` because the planner asks a node for its partitioning many times, and each ask
+  // would otherwise re-project the reported keys.
+  @transient override lazy val outputPartitioning: physical.Partitioning =
+    reportedKeyedPartitioning match {
+      case Some(partitioning) =>
+        // A partition key may reference a column that was pruned out of the scan output (see
+        // V2ScanPartitioningAndOrdering). Project such unresolvable key positions away so the
+        // reported partitioning only references output columns.
+        val exprs = partitioning.expressions
+        val resolvablePositions = exprs.indices.filter(i => exprs(i).references.subsetOf(outputSet))
+        if (resolvablePositions.isEmpty) {
+          super.outputPartitioning
+        } else {
+          partitioning.project(resolvablePositions)
+        }
+      case _ => super.outputPartitioning
+    }
+
   /**
    * Returns the output ordering for this scan. When the source reports ordering via
    * `SupportsReportOrdering`, that ordering is returned as-is. Otherwise, when the output

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

@@ -24,7 +24,7 @@ import org.apache.spark.internal.{Logging, LogKeys}
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, Literal, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils}
 import org.apache.spark.sql.catalyst.plans.logical.SampleMethod
-import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning}
+import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, InternalRowComparableWrapper}
@@ -273,23 +273,25 @@ object PushDownUtils extends Logging {
   }
 
   /**
-   * Pushes runtime filters into `scan` and re-plans its input partitions. For scans whose
-   * `outputPartitioning` is a [[KeyedPartitioning]] (SPJ-active), validates that the data source
-   * preserved the original partitioning and pads with `None` to preserve key alignment with the
-   * pre-filter partition set.
+   * Pushes runtime filters into `scan` and re-plans its input partitions. When the scan reported a
+   * [[KeyedPartitioning]] (SPJ-active), validates that the data source preserved the original
+   * partitioning and pads with `None` to preserve key alignment with the pre-filter partition set.
    *
    * Notes:
    *  - `filter` is mutating, and Spark may call this more than once for the same `scan` instance
    *    (see [[pushRuntimeFilters]]); successive calls are additive.
-   *  - When `outputPartitioning` is a [[KeyedPartitioning]], every split from
-   *    `planInputPartitions()` used on this path must implement [[HasPartitionKey]].
+   *  - With a [[KeyedPartitioning]], every split from `planInputPartitions()` used on this path
+   *    must implement [[HasPartitionKey]].
    *
    * @param scan                the V2 scan to push filters into
    * @param runtimeFilters      runtime filters to translate and push
    * @param table               the table backing the scan, used to derive the partition-predicate
    *                            schema for iterative [[PartitionPredicate]] pushdown
    * @param output              scan output attributes
-   * @param outputPartitioning  Spark-side output partitioning (used for SPJ validation)
+   * @param keyedPartitioning   the partitioning as the source reported it, at its full key width.
+   *                            The raw [[HasPartitionKey]] rows are read and ordered against it, so
+   *                            a projected partitioning must not be passed here: it would read each
+   *                            key row at the wrong positions and types
    * @param originalPartitions  unfiltered partitions, consulted only when no runtime filters fire
    * @return one entry per original input partition: `Some(part)` for surviving partitions and
    *         `None` for partition keys whose splits were entirely pruned (SPJ alignment)
@@ -299,15 +301,15 @@ object PushDownUtils extends Logging {
       runtimeFilters: Seq[Expression],
       table: Table,
       output: Seq[AttributeReference],
-      outputPartitioning: Partitioning,
+      keyedPartitioning: Option[KeyedPartitioning],
       originalPartitions: => Seq[InputPartition]): Seq[Option[InputPartition]] = {
     val filtered = pushRuntimeFilters(scan, runtimeFilters, table, output)
     if (filtered) {
       // call toBatch again to get filtered partitions
       val newPartitions = scan.toBatch.planInputPartitions()
 
-      outputPartitioning match {
-        case k: KeyedPartitioning =>
+      keyedPartitioning match {
+        case Some(k) =>
           if (newPartitions.exists(!_.isInstanceOf[HasPartitionKey])) {
             throw new SparkException("Data source must have preserved the original partitioning " +
                 "during runtime filtering: not all partitions implement HasPartitionKey after " +
@@ -344,22 +346,22 @@ object PushDownUtils extends Logging {
               fps.map(Some).padTo(size, None)
             }
 
-        case _ =>
+        case None =>
           // no validation is needed as the data source did not report any specific partitioning
           newPartitions.toSeq.map(Some)
       }
 
     } else {
       val parts = originalPartitions
-      (outputPartitioning match {
-        case k: KeyedPartitioning =>
+      (keyedPartitioning match {
+        case Some(k) =>
           if (parts.exists(!_.isInstanceOf[HasPartitionKey])) {
             throw new SparkException("Original partitions must implement HasPartitionKey when " +
-                "outputPartitioning is KeyedPartitioning.")
+                "the scan reported a KeyedPartitioning.")
           }
           parts.sortBy(_.asInstanceOf[HasPartitionKey].partitionKey())(k.keyRowOrdering)
 
-        case _ => parts
+        case None => parts
       }).map(Some)
     }
   }

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

@@ -50,15 +50,17 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging {
           val partitioning = sequenceToOption(
             kgp.keys().map(V2ExpressionUtils.toCatalystOpt(_, relation, relation.funCatalog))
               .toImmutableArraySeq)
-          if (partitioning.isEmpty) {
-            None
-          } else {
-            if (partitioning.get.forall(p => p.references.subsetOf(d.outputSet))) {
-              partitioning
-            } else {
-              None
-            }
-          }
+          // Keep the partitioning when at least one of its keys is still in the scan output: the
+          // scan projects the pruned key positions away when reporting its physical output
+          // partitioning (see DataSourceV2ScanExecBase.outputPartitioning). When no key survives,
+          // and likewise when the source reported no key at all, there is nothing to report.
+          // Grouping a projection that collapsed distinct keys onto the same key stays gated on
+          // allowKeysSubsetOfPartitionKeys one layer down, in KeyedPartitioning.mayGroupToSatisfy.
+          //
+          // A kept pruned key leaves a dangling attribute on the relation. What keeps that off
+          // `missingInput`, and so past the optimizer's plan-change validation, is the
+          // `DataSourceV2ScanRelation.references` override; see the comment there.
+          partitioning.filter(_.exists(_.references.subsetOf(d.outputSet)))
         case _: UnknownPartitioning => None
         case p =>
           logWarning(

sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala

@@ -981,9 +981,9 @@ class PlanMerger(
   // the merged scan's split count and partition values can still differ from an input's, since it
   // may push a different best-effort filter and so prune differently. And a report the merged scan
   // GAINS is not a degradation either. For partitioning that is because an input dropped its own
-  // only where a pruned column left the expressions inexpressible over that scan's output
-  // (V2ScanPartitioningAndOrdering's partitioning pass is reference-subset guarded), not because
-  // the source stopped reporting; the ordering pass has no such guard, so an ordering report is
+  // only where none of its keys survived in that scan's output (V2ScanPartitioningAndOrdering's
+  // partitioning pass keeps the report whenever any key survives), not because the source stopped
+  // reporting; the ordering pass has no such guard, so an ordering report is
   // never dropped by pruning and a gained one can only come from the source. Either way, keeping it
   // is exactly the win this merge is after.
   private def mergeDegradesReporting(

sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala

@@ -670,7 +670,7 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession {
     def replanAfterFiltering(afterFilter: Seq[InputPartition]): Unit = {
       val scan = new PartitioningBreakingScan(Seq(KeyedInputPartition(1)), afterFilter)
       PushDownUtils.replanWithRuntimeFilters(scan, Seq(EqualTo(partAttr, Literal(1))), table,
-        Seq(partAttr), partitioning, originalPartitions = Seq.empty)
+        Seq(partAttr), Some(partitioning), originalPartitions = Seq.empty)
     }
 
     val keyDropped = intercept[SparkException](replanAfterFiltering(Seq(new InputPartition {})))