diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala index d8817b9aa7582..060092cc91944 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala +++ b/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) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala index bcc332f4bd3cc..e81f25c656c08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala +++ b/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") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/v2/JDBCV2JoinPushdownIntegrationSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/v2/JDBCV2JoinPushdownIntegrationSuiteBase.scala index 4c46f5587a3f1..af390b707f09a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/v2/JDBCV2JoinPushdownIntegrationSuiteBase.scala +++ b/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)