[SPARK-59710][SQL] Fix JDBC V2 join pushdown for partitioned reads and partial AVG

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

@@ -671,8 +671,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
           aggExprToOutputOrdinal.clear()
           val newAggregates =
             collectAggregates(newResultExpressions, aggExprToOutputOrdinal)
-          val newNormalizedAggExprs = DataSourceStrategy.normalizeExprs(
-            newAggregates, holder.relation.output).asInstanceOf[Seq[AggregateExpression]]
+          val newNormalizedAggExprs =
+            normalizeExpressions(newAggregates, holder).asInstanceOf[Seq[AggregateExpression]]
           val newTranslatedAggOpt = DataSourceStrategy.translateAggregation(
             newNormalizedAggExprs, normalizedGroupingExpr)
           if (newTranslatedAggOpt.isEmpty) {

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

@@ -186,6 +186,12 @@ case class JDBCScanBuilder(
       return false
     }
 
+    // Join pushdown uses the query option, which cannot be combined with partitionColumn.
+    if (jdbcOptions.partitionColumn.isDefined) {
+      logDebug("Skipping JDBC join pushdown because partitionColumn is defined.")
+      return false
+    }
+
     val joinTypeStringOption = joinType match {
       case JoinType.INNER_JOIN => Some("INNER JOIN")
       case JoinType.LEFT_OUTER_JOIN => Some("LEFT JOIN")

sql/core/src/test/scala/org/apache/spark/sql/jdbc/v2/JDBCV2JoinPushdownIntegrationSuiteBase.scala

@@ -301,6 +301,31 @@ trait JDBCV2JoinPushdownIntegrationSuiteBase
     }
   }
 
+  gridTest("Join pushdown preserves partitioned input reads")(Seq(1, 2)) { numPartitions =>
+    val tableOptions = s"""WITH (
+      |'partitionColumn' '${caseConvert("id")}',
+      |'lowerBound' '0',
+      |'upperBound' '11',
+      |'numPartitions' '$numPartitions')""".stripMargin
+    val sqlQuery = s"""
+      |SELECT a.id, b.id
+      |FROM $catalogAndNamespace.$casedJoinTableName1 $tableOptions a
+      |JOIN $catalogAndNamespace.$casedJoinTableName1 $tableOptions b ON a.id = b.id + 1
+      |""".stripMargin
+
+    val expectedRows = for {
+      (leftId, _, _) <- table1Data
+      (rightId, _, _) <- table1Data
+      if leftId == rightId + 1
+    } yield Row(leftId, rightId)
+
+    withSQLConf(SQLConf.DATA_SOURCE_V2_JOIN_PUSHDOWN.key -> "true") {
+      val df = sql(sqlQuery)
+      checkJoinNotPushed(df)
+      checkAnswer(df, expectedRows)
+    }
+  }
+
   test("Test multi-way self join with conditions") {
     val sqlQuery = s"""
       |SELECT * FROM
@@ -479,6 +504,42 @@ trait JDBCV2JoinPushdownIntegrationSuiteBase
     }
   }
 
+  test("Join pushdown preserves aliases in partially pushed averages") {
+    assume(supportsAggregatePushdown, "Aggregate pushdown is not supported")
+    def sqlPartitionedQuery(numPartitions: Int): String = {
+      val idCol = caseConvert("id")
+      // numPartitions alone is valid and selects the partial aggregate path when greater than 1.
+      val tableOptions = s"WITH ('numPartitions' '$numPartitions')"
+      // Some databases return an integer for AVG over integer input, so use decimal input.
+      s"""
+        |SELECT avg(CAST(b.$idCol AS DECIMAL(10, 2)))
+        |FROM $catalogAndNamespace.$casedJoinTableName1 $tableOptions a
+        |JOIN $catalogAndNamespace.$casedJoinTableName1 $tableOptions b ON a.$idCol = b.$idCol + 1
+        |""".stripMargin
+    }
+
+    val rowsWithJoinPushdown = withSQLConf(SQLConf.DATA_SOURCE_V2_JOIN_PUSHDOWN.key -> "true") {
+      val completeAgg = sql(sqlPartitionedQuery(numPartitions = 1))
+      checkJoinPushed(completeAgg)
+      checkAggregateRemoved(completeAgg, pushed = true)
+      val expectedRows = completeAgg.collect().toSeq
+      assert(expectedRows.head.get(0) != null)
+
+      val partialAgg = sql(sqlPartitionedQuery(numPartitions = 2))
+      checkJoinPushed(partialAgg)
+      checkAggregateRemoved(partialAgg, pushed = false)
+      checkAnswer(partialAgg, expectedRows)
+      expectedRows
+    }
+
+    withSQLConf(SQLConf.DATA_SOURCE_V2_JOIN_PUSHDOWN.key -> "false") {
+      val df = sql(sqlPartitionedQuery(numPartitions = 2))
+      checkJoinNotPushed(df)
+      checkAggregateRemoved(df, pushed = false)
+      checkAnswer(df, rowsWithJoinPushdown)
+    }
+  }
+
   test("Test aggregate on top of multi-way self join") {
     val sqlQuery = s"""
       |SELECT min(a.id + b.id), min(a.id), min(c.id - 2)