This is an automated email from the ASF dual-hosted git repository.
csy pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new 7f2742b3 [AURON #1956] Add initial compatibility support for Spark 4.1
(UT/CI Pass) (#1958)
7f2742b3 is described below
commit 7f2742b33f27b391f27f0521e9162a31c9cf8326
Author: yew1eb <[email protected]>
AuthorDate: Thu Jan 29 16:45:21 2026 +0800
[AURON #1956] Add initial compatibility support for Spark 4.1 (UT/CI Pass)
(#1958)
# Which issue does this PR close?
Closes #1956
# Rationale for this change
This PR prioritizes Spark 4.1 as the first supported Spark 4.x version
to accelerate the Spark 4 compatibility initiative. The implementation
is designed for extensibility, enabling easy addition of Spark 4.0
support later if needed. This balances rapid adoption of the latest
stable Spark 4.x release with flexibility for other 4.x versions without
major rework.
# What changes are included in this PR?
## Spark 4 API Compatibility
### Servlet API Migration
Updated `AuronAllExecutionsPage.scala` to support both
`javax.servlet.http.HttpServletRequest` (Spark 3.x) and
`jakarta.servlet.http.HttpServletRequest` (Spark 4.x) via
version-specific `@sparkver` annotations, adapting to Spark 4's
migration to Jakarta EE Servlet API.
### Shuffle API Changes
Adapted shuffle components to address Spark 4.x's
`ShuffleWriteProcessor.write` API refinement
([SPARK-44605](https://github.com/apache/spark/pull/42234/changes)),
which triggers early execution of shuffle writers and breaks alignment
with Spark 3.x execution logic:
- Enhanced `AuronShuffleDependency` with a version-specific
`getInputRdd` method (returns `null` for Spark 3.x, returns `_rdd` for
Spark 4.x) — the exposed `inputRdd` field serializes the transient
`_rdd`, allowing the `_rdd` to be retrieved on `Executor` in Spark 4.1's
`ShuffleWriteProcessor.write` method.
- Returned `Iterator.empty` in `NativeRDD.compute()` for
`NativeRDD.ShuffleWrite` to defer execution to the
`ShuffleWriteProcessor.write()` method, aligning with Spark 3.x
execution logic.
- Added a Spark 4.1-specific override of `ShuffleWriteProcessor.write`
(which now takes `Iterator[_]` as its first parameter in Spark 4.x) in
`NativeShuffleExchangeExec`: it asserts the input iterator is empty
(validating adaptation logic), retrieves the RDD via
`AuronShuffleDependency.inputRdd`, and reuses core shuffle logic through
`internalWrite` to maintain consistency across Spark 3.x/4.x.
### SparkSession Package Path Change
Addressed Spark 4.x's SparkSession package restructure:
- Spark 3.x: org.apache.spark.sql.SparkSession → Spark 4.x:
org.apache.spark.sql.classic.SparkSession
- Updated references in NativeParquetInsertIntoHiveTableExec.scala and
NativeBroadcastExchangeBase.scala
### New Data Types
Added stubs for Spark 4.x's new
`GeographyVal`/`GeometryVal`/`VariantVal` data types in columnar data
structures (`AuronColumnarArray.scala`, `AuronColumnarStruct.scala`,
`AuronColumnarBatchRow.scala`). These stubs throw
`UnsupportedOperationException` to resolve compilation errors.
# Are there any user-facing changes?
# How was this patch tested?
- [x] Enabled Spark 4.1 in CI pipeline
- [x] Passed all existing Unit Tests (UT)
- [x] Passed all TPC-DS Integration Tests (IT)
---
.github/workflows/tpcds-reusable.yml | 2 +-
.github/workflows/tpcds.yml | 9 +++
auron-build.sh | 2 +-
auron-spark-ui/pom.xml | 13 ++++
.../sql/execution/ui/AuronAllExecutionsPage.scala | 32 ++++++++-
dev/auron-it/pom.xml | 41 +++++++++++-
.../scala/org/apache/auron/integration/Main.scala | 4 +-
.../apache/auron/integration/SessionManager.scala | 1 +
pom.xml | 64 ++++++++++++++++--
.../sql/auron/InterceptedValidateSparkPlan.scala | 4 +-
.../org/apache/spark/sql/auron/ShimsImpl.scala | 44 ++++++------
.../execution/auron/plan/ConvertToNativeExec.scala | 2 +-
.../sql/execution/auron/plan/NativeAggExec.scala | 10 +--
.../auron/plan/NativeBroadcastExchangeExec.scala | 2 +-
.../auron/plan/NativeCollectLimitExec.scala | 2 +-
.../execution/auron/plan/NativeExpandExec.scala | 2 +-
.../execution/auron/plan/NativeFilterExec.scala | 2 +-
.../execution/auron/plan/NativeGenerateExec.scala | 2 +-
.../auron/plan/NativeGlobalLimitExec.scala | 2 +-
.../auron/plan/NativeLocalLimitExec.scala | 2 +-
.../NativeParquetInsertIntoHiveTableExec.scala | 78 ++++++++++++++++++++--
.../auron/plan/NativeParquetSinkExec.scala | 2 +-
.../auron/plan/NativePartialTakeOrderedExec.scala | 2 +-
.../auron/plan/NativeProjectExecProvider.scala | 2 +-
.../plan/NativeRenameColumnsExecProvider.scala | 2 +-
.../auron/plan/NativeShuffleExchangeExec.scala | 52 ++++++++++++---
.../sql/execution/auron/plan/NativeSortExec.scala | 2 +-
.../auron/plan/NativeTakeOrderedExec.scala | 2 +-
.../sql/execution/auron/plan/NativeUnionExec.scala | 2 +-
.../execution/auron/plan/NativeWindowExec.scala | 2 +-
.../shuffle/AuronBlockStoreShuffleReader.scala | 2 +-
.../auron/shuffle/AuronRssShuffleManagerBase.scala | 2 +-
.../auron/shuffle/AuronShuffleManager.scala | 4 +-
.../auron/shuffle/AuronShuffleWriter.scala | 2 +-
.../joins/auron/plan/NativeBroadcastJoinExec.scala | 12 ++--
.../plan/NativeShuffledHashJoinExecProvider.scala | 2 +-
.../plan/NativeSortMergeJoinExecProvider.scala | 2 +-
.../org/apache/auron/AuronFunctionSuite.scala | 12 +++-
.../scala/org/apache/auron/AuronQuerySuite.scala | 11 ++-
.../scala/org/apache/auron/BaseAuronSQLSuite.scala | 3 +
.../execution/AuronAdaptiveQueryExecSuite.scala | 2 +-
.../org/apache/auron/util/SparkVersionUtil.scala | 14 ++--
.../org/apache/spark/sql/auron/NativeRDD.scala | 11 ++-
.../auron/columnar/AuronColumnarArray.scala | 20 +++++-
.../auron/columnar/AuronColumnarBatchRow.scala | 17 +++++
.../auron/columnar/AuronColumnarStruct.scala | 20 +++++-
.../auron/plan/NativeBroadcastExchangeBase.scala | 51 +++++++++-----
.../auron/shuffle/AuronShuffleDependency.scala | 16 ++++-
.../auron/plan/NativeHiveTableScanBase.scala | 2 +-
49 files changed, 471 insertions(+), 122 deletions(-)
diff --git a/.github/workflows/tpcds-reusable.yml
b/.github/workflows/tpcds-reusable.yml
index 40f81e54..75d84418 100644
--- a/.github/workflows/tpcds-reusable.yml
+++ b/.github/workflows/tpcds-reusable.yml
@@ -226,7 +226,7 @@ jobs:
if: steps.cache-spark-bin.outputs.cache-hit != 'true'
run: |
SPARK_PATH="spark/spark-${{
steps.get-dependency-version.outputs.sparkversion }}"
- if [ ${{ inputs.scalaver }} = "2.13" ]; then
+ if [ ${{ inputs.scalaver }} = "2.13" && "${{ inputs.sparkver }}" !=
"spark-4.1" ]; then
SPARK_FILE="spark-${{
steps.get-dependency-version.outputs.sparkversion }}-bin-${{
inputs.hadoop-profile }}-scala${{ inputs.scalaver }}.tgz"
else
SPARK_FILE="spark-${{
steps.get-dependency-version.outputs.sparkversion }}-bin-${{
inputs.hadoop-profile }}.tgz"
diff --git a/.github/workflows/tpcds.yml b/.github/workflows/tpcds.yml
index 87ac607f..e68e79db 100644
--- a/.github/workflows/tpcds.yml
+++ b/.github/workflows/tpcds.yml
@@ -87,3 +87,12 @@ jobs:
javaver: '21'
scalaver: '2.13'
hadoop-profile: 'hadoop3'
+
+ test-spark-41-jdk21-scala-2-13:
+ name: Test spark-4.1 JDK21 Scala-2.13
+ uses: ./.github/workflows/tpcds-reusable.yml
+ with:
+ sparkver: spark-4.1
+ javaver: '21'
+ scalaver: '2.13'
+ hadoop-profile: 'hadoop3'
\ No newline at end of file
diff --git a/auron-build.sh b/auron-build.sh
index d947b79e..e949021b 100755
--- a/auron-build.sh
+++ b/auron-build.sh
@@ -30,7 +30,7 @@
# Define constants for supported component versions
# -----------------------------------------------------------------------------
SUPPORTED_OS_IMAGES=("centos7" "ubuntu24" "rockylinux8" "debian11"
"azurelinux3")
-SUPPORTED_SPARK_VERSIONS=("3.0" "3.1" "3.2" "3.3" "3.4" "3.5")
+SUPPORTED_SPARK_VERSIONS=("3.0" "3.1" "3.2" "3.3" "3.4" "3.5" "4.1")
SUPPORTED_SCALA_VERSIONS=("2.12" "2.13")
SUPPORTED_CELEBORN_VERSIONS=("0.5" "0.6")
# Currently only one supported version, but kept plural for consistency
diff --git a/auron-spark-ui/pom.xml b/auron-spark-ui/pom.xml
index 08ba2039..876c4494 100644
--- a/auron-spark-ui/pom.xml
+++ b/auron-spark-ui/pom.xml
@@ -36,6 +36,19 @@
<artifactId>spark-sql_${scalaVersion}</artifactId>
<scope>provided</scope>
</dependency>
+ <!-- Required for XML processing (Scala 2.12+ split module, compatible
with Spark 4.x) -->
+ <dependency>
+ <groupId>org.scala-lang.modules</groupId>
+ <artifactId>scala-xml_${scalaVersion}</artifactId>
+ <version>${scala-xml.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.auron</groupId>
+ <artifactId>spark-version-annotation-macros_${scalaVersion}</artifactId>
+ <version>${project.version}</version>
+ <scope>compile</scope>
+ </dependency>
</dependencies>
<build>
<plugins>
diff --git
a/auron-spark-ui/src/main/scala/org/apache/spark/sql/execution/ui/AuronAllExecutionsPage.scala
b/auron-spark-ui/src/main/scala/org/apache/spark/sql/execution/ui/AuronAllExecutionsPage.scala
index c237557f..9b4630af 100644
---
a/auron-spark-ui/src/main/scala/org/apache/spark/sql/execution/ui/AuronAllExecutionsPage.scala
+++
b/auron-spark-ui/src/main/scala/org/apache/spark/sql/execution/ui/AuronAllExecutionsPage.scala
@@ -16,18 +16,44 @@
*/
package org.apache.spark.sql.execution.ui
-import javax.servlet.http.HttpServletRequest
-
import scala.xml.{Node, NodeSeq}
import org.apache.spark.internal.Logging
import org.apache.spark.ui.{UIUtils, WebUIPage}
+import org.apache.auron.sparkver
+
private[ui] class AuronAllExecutionsPage(parent: AuronSQLTab) extends
WebUIPage("") with Logging {
private val sqlStore = parent.sqlStore
- override def render(request: HttpServletRequest): Seq[Node] = {
+ @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ override def render(request: javax.servlet.http.HttpServletRequest):
Seq[Node] = {
+ val buildInfo = sqlStore.buildInfo()
+ val infos =
+ UIUtils.listingTable(propertyHeader, propertyRow, buildInfo.info,
fixedWidth = true)
+ val summary: NodeSeq =
+ <div>
+ <div>
+ <span class="collapse-sql-properties collapse-table"
+ onClick="collapseTable('collapse-sql-properties',
'sql-properties')">
+ <h4>
+ <span class="collapse-table-arrow arrow-open"></span>
+ <a>Auron Build Information</a>
+ </h4>
+ </span>
+ <div class="sql-properties collapsible-table">
+ {infos}
+ </div>
+ </div>
+ <br/>
+ </div>
+
+ UIUtils.headerSparkPage(request, "Auron", summary, parent)
+ }
+
+ @sparkver("4.1")
+ override def render(request: jakarta.servlet.http.HttpServletRequest):
Seq[Node] = {
val buildInfo = sqlStore.buildInfo()
val infos =
UIUtils.listingTable(propertyHeader, propertyRow, buildInfo.info,
fixedWidth = true)
diff --git a/dev/auron-it/pom.xml b/dev/auron-it/pom.xml
index 5cdd1710..37d4f3bc 100644
--- a/dev/auron-it/pom.xml
+++ b/dev/auron-it/pom.xml
@@ -331,6 +331,45 @@
</properties>
</profile>
+ <profile>
+ <id>spark-4.1</id>
+ <properties>
+ <shimName>spark-4.1</shimName>
+ <sparkVersion>4.1.1</sparkVersion>
+ </properties>
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-enforcer-plugin</artifactId>
+ <version>${maven-enforcer-plugin.version}</version>
+ <executions>
+ <execution>
+ <id>spark41-enforce-java-scala-version</id>
+ <goals>
+ <goal>enforce</goal>
+ </goals>
+ <configuration>
+ <rules>
+ <!-- Spark 4.1 requires JDK 17+ and Scala 2.13.x -->
+ <requireJavaVersion>
+ <version>[17,)</version>
+ <message>Spark 4.1 requires JDK 17 or higher. Current:
${java.version}</message>
+ </requireJavaVersion>
+ <requireProperty>
+ <property>scalaLongVersion</property>
+ <regex>2\.13\.\d+</regex>
+ <regexMessage>Spark 4.1 requires Scala 2.13.x. Current:
${scalaLongVersion}</regexMessage>
+ </requireProperty>
+ </rules>
+ </configuration>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </build>
+ </profile>
+
<profile>
<id>scala-2.12</id>
<properties>
@@ -343,7 +382,7 @@
<id>scala-2.13</id>
<properties>
<scalaVersion>2.13</scalaVersion>
- <scalaLongVersion>2.13.13</scalaLongVersion>
+ <scalaLongVersion>2.13.17</scalaLongVersion>
</properties>
</profile>
</profiles>
diff --git
a/dev/auron-it/src/main/scala/org/apache/auron/integration/Main.scala
b/dev/auron-it/src/main/scala/org/apache/auron/integration/Main.scala
index 48a7c4cc..b944b6c2 100644
--- a/dev/auron-it/src/main/scala/org/apache/auron/integration/Main.scala
+++ b/dev/auron-it/src/main/scala/org/apache/auron/integration/Main.scala
@@ -114,8 +114,8 @@ object Main {
|Spark Version: ${Shims.get.shimVersion}
|Data: ${args.dataLocation}
|Queries: [${args.queryFilter.mkString(", ")}] (${if
(args.queryFilter.isEmpty)
- "all"
- else args.queryFilter.length} queries)
+ "all"
+ else args.queryFilter.length} queries)
|Extra Spark Conf: ${args.extraSparkConf}""".stripMargin)
if (args.auronOnly) println("Mode: Auron-only (skip baseline)")
diff --git
a/dev/auron-it/src/main/scala/org/apache/auron/integration/SessionManager.scala
b/dev/auron-it/src/main/scala/org/apache/auron/integration/SessionManager.scala
index 3e65e2e3..a51c50f0 100644
---
a/dev/auron-it/src/main/scala/org/apache/auron/integration/SessionManager.scala
+++
b/dev/auron-it/src/main/scala/org/apache/auron/integration/SessionManager.scala
@@ -35,6 +35,7 @@ class SessionManager(val extraSparkConf: Map[String, String])
{
private lazy val commonConf: Map[String, String] = Map(
"spark.master" -> resolveMaster(),
"spark.sql.shuffle.partitions" -> "100",
+ "spark.sql.unionOutputPartitioning" -> "false",
"spark.ui.enabled" -> "false",
"spark.sql.sources.useV1SourceList" -> "parquet",
"spark.sql.autoBroadcastJoinThreshold" -> "-1")
diff --git a/pom.xml b/pom.xml
index 405b3474..3419a9e8 100644
--- a/pom.xml
+++ b/pom.xml
@@ -54,8 +54,11 @@
<protobufVersion>3.25.5</protobufVersion>
<nettyVersion>4.2.7.Final</nettyVersion>
<javaVersion>8</javaVersion>
+ <maven.compiler.source>${javaVersion}</maven.compiler.source>
+ <maven.compiler.target>${javaVersion}</maven.compiler.target>
<scalaVersion>2.12</scalaVersion>
<scalaLongVersion>2.12.18</scalaLongVersion>
+ <scala-xml.version>2.1.0</scala-xml.version>
<scalaJava8CompatVersion>1.0.2</scalaJava8CompatVersion>
<maven.version>3.9.12</maven.version>
<maven.plugin.scala.version>4.9.2</maven.plugin.scala.version>
@@ -383,6 +386,9 @@
<arg>-feature</arg>
<arg>-Ywarn-unused</arg>
<arg>-Xfatal-warnings</arg>
+
+ <arg>-Wconf:msg=method newInstance in class Class is
deprecated:s</arg>
+ <arg>-Wconf:msg=class ThreadDeath in package lang is
deprecated:s</arg>
</args>
</configuration>
<dependencies>
@@ -465,6 +471,8 @@
<version>${maven.plugin.surefire.version}</version>
<!-- Note config is repeated in scalatest config -->
<configuration>
+ <forkCount>1</forkCount>
+ <reuseForks>false</reuseForks>
<skipTests>false</skipTests>
<failIfNoSpecifiedTests>false</failIfNoSpecifiedTests>
<argLine>${extraJavaTestArgs}</argLine>
@@ -486,6 +494,7 @@
<reportsDirectory>${project.build.directory}/surefire-reports</reportsDirectory>
<junitxml>.</junitxml>
<filereports>TestSuite.txt</filereports>
+ <forkMode>once</forkMode>
<argLine>${extraJavaTestArgs}</argLine>
<environmentVariables />
<systemProperties>
@@ -827,6 +836,48 @@
</properties>
</profile>
+ <profile>
+ <id>spark-4.1</id>
+ <properties>
+ <shimName>spark-4.1</shimName>
+ <scalaTestVersion>3.2.9</scalaTestVersion>
+ <sparkVersion>4.1.1</sparkVersion>
+ <shortSparkVersion>4.1</shortSparkVersion>
+ <nettyVersion>4.1.118.Final</nettyVersion>
+ </properties>
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-enforcer-plugin</artifactId>
+ <version>${maven-enforcer-plugin.version}</version>
+ <executions>
+ <execution>
+ <id>spark41-enforce-java-scala-version</id>
+ <goals>
+ <goal>enforce</goal>
+ </goals>
+ <configuration>
+ <rules>
+ <!-- Spark 4.1 requires JDK 17+ and Scala 2.13.x -->
+ <requireJavaVersion>
+ <version>[17,)</version>
+ <message>Spark 4.1 requires JDK 17 or higher. Current:
${java.version}</message>
+ </requireJavaVersion>
+ <requireProperty>
+ <property>scalaLongVersion</property>
+ <regex>2\.13\.\d+</regex>
+ <regexMessage>Spark 4.1 requires Scala 2.13.x. Current:
${scalaLongVersion}</regexMessage>
+ </requireProperty>
+ </rules>
+ </configuration>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </build>
+ </profile>
+
<profile>
<id>jdk-8</id>
<activation>
@@ -835,7 +886,7 @@
<properties>
<javaVersion>8</javaVersion>
<spotless.plugin.version>2.30.0</spotless.plugin.version>
- <semanticdb.version>4.8.8</semanticdb.version>
+ <semanticdb.version>4.14.5</semanticdb.version>
<scalafmtVersion>3.0.0</scalafmtVersion>
</properties>
</profile>
@@ -848,7 +899,7 @@
<properties>
<javaVersion>11</javaVersion>
<spotless.plugin.version>2.30.0</spotless.plugin.version>
- <semanticdb.version>4.8.8</semanticdb.version>
+ <semanticdb.version>4.14.5</semanticdb.version>
<scalafmtVersion>3.0.0</scalafmtVersion>
</properties>
</profile>
@@ -861,7 +912,7 @@
<properties>
<javaVersion>17</javaVersion>
<spotless.plugin.version>2.45.0</spotless.plugin.version>
- <semanticdb.version>4.9.9</semanticdb.version>
+ <semanticdb.version>4.14.5</semanticdb.version>
<scalafmtVersion>3.9.9</scalafmtVersion>
</properties>
</profile>
@@ -874,7 +925,7 @@
<properties>
<javaVersion>21</javaVersion>
<spotless.plugin.version>2.45.0</spotless.plugin.version>
- <semanticdb.version>4.9.9</semanticdb.version>
+ <semanticdb.version>4.14.5</semanticdb.version>
<scalafmtVersion>3.9.9</scalafmtVersion>
</properties>
</profile>
@@ -921,7 +972,7 @@
</activation>
<properties>
<scalaVersion>2.13</scalaVersion>
- <scalaLongVersion>2.13.13</scalaLongVersion>
+ <scalaLongVersion>2.13.17</scalaLongVersion>
</properties>
<build>
<plugins>
@@ -944,11 +995,14 @@
<arg>-Wconf:msg=^(?=.*?method|value|type|object|trait|inheritance)(?=.*?deprecated)(?=.*?since
2.13).+$:s</arg>
<arg>-Wconf:msg=Auto-application to \`\(\)\` is
deprecated:s</arg>
<arg>-Wconf:msg=object JavaConverters in package collection is
deprecated:s</arg>
+ <arg>-Wconf:msg=method newInstance in class Class is
deprecated:s</arg>
+ <arg>-Wconf:msg=class ThreadDeath in package lang is
deprecated:s</arg>
<arg>-Wconf:cat=unchecked&msg=outer reference:s</arg>
<arg>-Wconf:cat=unchecked&msg=eliminated by erasure:s</arg>
<arg>-Wconf:cat=unused-nowarn:s</arg>
<arg>-Wconf:msg=early initializers are deprecated:s</arg>
<arg>-Wconf:cat=other-match-analysis:s</arg>
+ <arg>-Wconf:cat=feature-existentials:s</arg>
</args>
<compilerPlugins>
<compilerPlugin>
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/InterceptedValidateSparkPlan.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/InterceptedValidateSparkPlan.scala
index f90e4865..d8b574b3 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/InterceptedValidateSparkPlan.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/InterceptedValidateSparkPlan.scala
@@ -25,7 +25,7 @@ import org.apache.auron.sparkver
object InterceptedValidateSparkPlan extends Logging {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
def validate(plan: SparkPlan): Unit = {
import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec
import org.apache.spark.sql.execution.auron.plan.NativeRenameColumnsBase
@@ -79,7 +79,7 @@ object InterceptedValidateSparkPlan extends Logging {
throw new UnsupportedOperationException("validate is not supported in
spark 3.0.3 or 3.1.3")
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
private def errorOnInvalidBroadcastQueryStage(plan: SparkPlan): Unit = {
import org.apache.spark.sql.execution.adaptive.InvalidAQEPlanException
throw InvalidAQEPlanException("Invalid broadcast query stage", plan)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala
index 6353d7cb..e48acefb 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala
@@ -124,8 +124,10 @@ class ShimsImpl extends Shims with Logging {
override def shimVersion: String = "spark-3.4"
@sparkver("3.5")
override def shimVersion: String = "spark-3.5"
+ @sparkver("4.1")
+ override def shimVersion: String = "spark-4.1"
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def initExtension(): Unit = {
ValidateSparkPlanInjector.inject()
@@ -285,16 +287,16 @@ class ShimsImpl extends Shims with Logging {
child: SparkPlan): NativeGenerateBase =
NativeGenerateExec(generator, requiredChildOutput, outer, generatorOutput,
child)
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
private def effectiveLimit(rawLimit: Int): Int =
if (rawLimit == -1) Int.MaxValue else rawLimit
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
override def getLimitAndOffset(plan: GlobalLimitExec): (Int, Int) = {
(effectiveLimit(plan.limit), plan.offset)
}
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
override def getLimitAndOffset(plan: TakeOrderedAndProjectExec): (Int, Int)
= {
(effectiveLimit(plan.limit), plan.offset)
}
@@ -308,7 +310,7 @@ class ShimsImpl extends Shims with Logging {
override def createNativeLocalLimitExec(limit: Int, child: SparkPlan):
NativeLocalLimitBase =
NativeLocalLimitExec(limit, child)
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
override def getLimitAndOffset(plan: CollectLimitExec): (Int, Int) = {
(effectiveLimit(plan.limit), plan.offset)
}
@@ -456,7 +458,7 @@ class ShimsImpl extends Shims with Logging {
length: Long,
numRecords: Long): FileSegment = new FileSegment(file, offset, length)
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def commit(
dep: ShuffleDependency[_, _, _],
shuffleBlockResolver: IndexShuffleBlockResolver,
@@ -624,7 +626,7 @@ class ShimsImpl extends Shims with Logging {
expr.asInstanceOf[AggregateExpression].filter
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
private def isAQEShuffleRead(exec: SparkPlan): Boolean = {
import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
exec.isInstanceOf[AQEShuffleReadExec]
@@ -636,7 +638,7 @@ class ShimsImpl extends Shims with Logging {
exec.isInstanceOf[CustomShuffleReaderExec]
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
private def executeNativeAQEShuffleReader(exec: SparkPlan): NativeRDD = {
import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
import org.apache.spark.sql.execution.CoalescedMapperPartitionSpec
@@ -936,7 +938,7 @@ class ShimsImpl extends Shims with Logging {
}
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getSqlContext(sparkPlan: SparkPlan): SQLContext =
sparkPlan.session.sqlContext
@@ -958,7 +960,7 @@ class ShimsImpl extends Shims with Logging {
size: Long): PartitionedFile =
PartitionedFile(partitionValues, filePath, offset, size)
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
override def getPartitionedFile(
partitionValues: InternalRow,
filePath: String,
@@ -969,7 +971,7 @@ class ShimsImpl extends Shims with Logging {
PartitionedFile(partitionValues, SparkPath.fromPath(new Path(filePath)),
offset, size)
}
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getMinPartitionNum(sparkSession: SparkSession): Int =
sparkSession.sessionState.conf.filesMinPartitionNum
.getOrElse(sparkSession.sparkContext.defaultParallelism)
@@ -992,13 +994,13 @@ class ShimsImpl extends Shims with Logging {
}
@nowarn("cat=unused") // Some params temporarily unused
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
private def convertPromotePrecision(
e: Expression,
isPruningExpr: Boolean,
fallback: Expression => pb.PhysicalExprNode):
Option[pb.PhysicalExprNode] = None
- @sparkver("3.3 / 3.4 / 3.5")
+ @sparkver("3.3 / 3.4 / 3.5 / 4.1")
private def convertBloomFilterAgg(agg: AggregateFunction):
Option[pb.PhysicalAggExprNode] = {
import
org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate
agg match {
@@ -1028,7 +1030,7 @@ class ShimsImpl extends Shims with Logging {
@sparkver("3.0 / 3.1 / 3.2")
private def convertBloomFilterAgg(agg: AggregateFunction):
Option[pb.PhysicalAggExprNode] = None
- @sparkver("3.3 / 3.4 / 3.5")
+ @sparkver("3.3 / 3.4 / 3.5 / 4.1")
private def convertBloomFilterMightContain(
e: Expression,
isPruningExpr: Boolean,
@@ -1063,7 +1065,7 @@ class ShimsImpl extends Shims with Logging {
exec.initialPlan
}
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getAdaptiveInputPlan(exec: AdaptiveSparkPlanExec): SparkPlan = {
exec.inputPlan
}
@@ -1093,7 +1095,7 @@ class ShimsImpl extends Shims with Logging {
})
}
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getJoinBuildSide(exec: SparkPlan): JoinBuildSide = {
import org.apache.spark.sql.catalyst.optimizer.BuildLeft
convertJoinBuildSide(
@@ -1104,19 +1106,19 @@ class ShimsImpl extends Shims with Logging {
})
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean =
exec.isSkewJoin
@sparkver("3.0 / 3.1")
override def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean =
false
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] =
Some(exec.shuffleOrigin)
@sparkver("3.0")
override def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] = None
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def isNullAwareAntiJoin(exec: BroadcastHashJoinExec): Boolean =
exec.isNullAwareAntiJoin
@@ -1127,7 +1129,7 @@ class ShimsImpl extends Shims with Logging {
case class ForceNativeExecutionWrapper(override val child: SparkPlan)
extends ForceNativeExecutionWrapperBase(child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
@@ -1142,6 +1144,6 @@ case class NativeExprWrapper(
override val nullable: Boolean)
extends NativeExprWrapperBase(nativeExpr, dataType, nullable) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def withNewChildrenInternal(newChildren: IndexedSeq[Expression]):
Expression = copy()
}
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/ConvertToNativeExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/ConvertToNativeExec.scala
index 5231bd11..2028ac1f 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/ConvertToNativeExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/ConvertToNativeExec.scala
@@ -22,7 +22,7 @@ import org.apache.auron.sparkver
case class ConvertToNativeExec(override val child: SparkPlan) extends
ConvertToNativeBase(child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeAggExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeAggExec.scala
index 67b806b6..f9623bd2 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeAggExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeAggExec.scala
@@ -44,22 +44,22 @@ case class NativeAggExec(
child)
with BaseAggregateExec {
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override val requiredChildDistributionExpressions: Option[Seq[Expression]] =
theRequiredChildDistributionExpressions
- @sparkver("3.3 / 3.4 / 3.5")
+ @sparkver("3.3 / 3.4 / 3.5 / 4.1")
override val initialInputBufferOffset: Int = theInitialInputBufferOffset
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def isStreaming: Boolean = false
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def numShufflePartitions: Option[Int] = None
override def resultExpressions: Seq[NamedExpression] = outputAttributes
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeExec.scala
index f2f43f32..3a9ad1a9 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeExec.scala
@@ -43,7 +43,7 @@ case class NativeBroadcastExchangeExec(mode: BroadcastMode,
override val child:
relationFuturePromise.future
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitExec.scala
index 4ff7d804..4af2597f 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitExec.scala
@@ -23,7 +23,7 @@ import org.apache.auron.sparkver
case class NativeCollectLimitExec(limit: Int, offset: Int, override val child:
SparkPlan)
extends NativeCollectLimitBase(limit, offset, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeExpandExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeExpandExec.scala
index ca4aba5e..d83a1b1f 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeExpandExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeExpandExec.scala
@@ -28,7 +28,7 @@ case class NativeExpandExec(
override val child: SparkPlan)
extends NativeExpandBase(projections, output, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeFilterExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeFilterExec.scala
index 367ad7cc..0b51523f 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeFilterExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeFilterExec.scala
@@ -24,7 +24,7 @@ import org.apache.auron.sparkver
case class NativeFilterExec(condition: Expression, override val child:
SparkPlan)
extends NativeFilterBase(condition, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGenerateExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGenerateExec.scala
index f0019c66..3d2a1510 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGenerateExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGenerateExec.scala
@@ -30,7 +30,7 @@ case class NativeGenerateExec(
override val child: SparkPlan)
extends NativeGenerateBase(generator, requiredChildOutput, outer,
generatorOutput, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGlobalLimitExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGlobalLimitExec.scala
index 1b493f43..4b077812 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGlobalLimitExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeGlobalLimitExec.scala
@@ -23,7 +23,7 @@ import org.apache.auron.sparkver
case class NativeGlobalLimitExec(limit: Int, offset: Int, override val child:
SparkPlan)
extends NativeGlobalLimitBase(limit, offset, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeLocalLimitExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeLocalLimitExec.scala
index 805408c8..4b44aca9 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeLocalLimitExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeLocalLimitExec.scala
@@ -23,7 +23,7 @@ import org.apache.auron.sparkver
case class NativeLocalLimitExec(limit: Int, override val child: SparkPlan)
extends NativeLocalLimitBase(limit, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetInsertIntoHiveTableExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetInsertIntoHiveTableExec.scala
index e19c86e5..e13f1346 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetInsertIntoHiveTableExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetInsertIntoHiveTableExec.scala
@@ -17,7 +17,6 @@
package org.apache.spark.sql.execution.auron.plan
import org.apache.spark.sql.Row
-import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.auron.Shims
import org.apache.spark.sql.catalyst.catalog.CatalogTable
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
@@ -70,7 +69,26 @@ case class NativeParquetInsertIntoHiveTableExec(
metrics)
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("4.1")
+ override protected def getInsertIntoHiveTableCommand(
+ table: CatalogTable,
+ partition: Map[String, Option[String]],
+ query: LogicalPlan,
+ overwrite: Boolean,
+ ifPartitionNotExists: Boolean,
+ outputColumnNames: Seq[String],
+ metrics: Map[String, SQLMetric]): InsertIntoHiveTable = {
+ new AuronInsertIntoHiveTable41(
+ table,
+ partition,
+ query,
+ overwrite,
+ ifPartitionNotExists,
+ outputColumnNames,
+ metrics)
+ }
+
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
@@ -97,7 +115,9 @@ case class NativeParquetInsertIntoHiveTableExec(
override lazy val metrics: Map[String, SQLMetric] = outerMetrics
- override def run(sparkSession: SparkSession, child: SparkPlan): Seq[Row] =
{
+ override def run(
+ sparkSession: org.apache.spark.sql.SparkSession,
+ child: SparkPlan): Seq[Row] = {
val nativeParquetSink =
Shims.get.createNativeParquetSinkExec(sparkSession, table, partition,
child, metrics)
super.run(sparkSession, nativeParquetSink)
@@ -266,7 +286,57 @@ case class NativeParquetInsertIntoHiveTableExec(
override lazy val metrics: Map[String, SQLMetric] = outerMetrics
- override def run(sparkSession: SparkSession, child: SparkPlan): Seq[Row] =
{
+ override def run(
+ sparkSession: org.apache.spark.sql.SparkSession,
+ child: SparkPlan): Seq[Row] = {
+ val nativeParquetSink =
+ Shims.get.createNativeParquetSinkExec(sparkSession, table, partition,
child, metrics)
+ super.run(sparkSession, nativeParquetSink)
+ }
+ }
+
+ @sparkver("4.1")
+ class AuronInsertIntoHiveTable41(
+ table: CatalogTable,
+ partition: Map[String, Option[String]],
+ query: LogicalPlan,
+ overwrite: Boolean,
+ ifPartitionNotExists: Boolean,
+ outputColumnNames: Seq[String],
+ outerMetrics: Map[String, SQLMetric])
+ extends {
+ private val insertIntoHiveTable = InsertIntoHiveTable(
+ table,
+ partition,
+ query,
+ overwrite,
+ ifPartitionNotExists,
+ outputColumnNames)
+ private val initPartitionColumns = insertIntoHiveTable.partitionColumns
+ private val initBucketSpec = insertIntoHiveTable.bucketSpec
+ private val initOptions = insertIntoHiveTable.options
+ private val initFileFormat = insertIntoHiveTable.fileFormat
+ private val initHiveTmpPath = insertIntoHiveTable.hiveTmpPath
+
+ }
+ with InsertIntoHiveTable(
+ table,
+ partition,
+ query,
+ overwrite,
+ ifPartitionNotExists,
+ outputColumnNames,
+ initPartitionColumns,
+ initBucketSpec,
+ initOptions,
+ initFileFormat,
+ initHiveTmpPath) {
+
+ override lazy val metrics: Map[String, SQLMetric] = outerMetrics
+
+ override def run(
+ sparkSession: org.apache.spark.sql.classic.SparkSession,
+ child: SparkPlan): Seq[Row] = {
val nativeParquetSink =
Shims.get.createNativeParquetSinkExec(sparkSession, table, partition,
child, metrics)
super.run(sparkSession, nativeParquetSink)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetSinkExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetSinkExec.scala
index 5cbb61ef..5b548eda 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetSinkExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeParquetSinkExec.scala
@@ -31,7 +31,7 @@ case class NativeParquetSinkExec(
override val metrics: Map[String, SQLMetric])
extends NativeParquetSinkBase(sparkSession, table, partition, child,
metrics) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativePartialTakeOrderedExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativePartialTakeOrderedExec.scala
index faf541a6..eafb355c 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativePartialTakeOrderedExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativePartialTakeOrderedExec.scala
@@ -29,7 +29,7 @@ case class NativePartialTakeOrderedExec(
override val metrics: Map[String, SQLMetric])
extends NativePartialTakeOrderedBase(limit, sortOrder, child, metrics) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeProjectExecProvider.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeProjectExecProvider.scala
index 51845400..6902c9f5 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeProjectExecProvider.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeProjectExecProvider.scala
@@ -22,7 +22,7 @@ import org.apache.spark.sql.execution.SparkPlan
import org.apache.auron.sparkver
case object NativeProjectExecProvider {
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
def provide(projectList: Seq[NamedExpression], child: SparkPlan):
NativeProjectBase = {
import org.apache.spark.sql.execution.OrderPreservingUnaryExecNode
import org.apache.spark.sql.execution.PartitioningPreservingUnaryExecNode
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeRenameColumnsExecProvider.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeRenameColumnsExecProvider.scala
index 3ba34ba0..b3278248 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeRenameColumnsExecProvider.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeRenameColumnsExecProvider.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.execution.SparkPlan
import org.apache.auron.sparkver
case object NativeRenameColumnsExecProvider {
- @sparkver("3.4 / 3.5")
+ @sparkver("3.4 / 3.5 / 4.1")
def provide(child: SparkPlan, renamedColumnNames: Seq[String]):
NativeRenameColumnsBase = {
import org.apache.spark.sql.catalyst.expressions.NamedExpression
import org.apache.spark.sql.catalyst.expressions.SortOrder
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeShuffleExchangeExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeShuffleExchangeExec.scala
index 9d688ad1..d23e07c3 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeShuffleExchangeExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeShuffleExchangeExec.scala
@@ -20,19 +20,14 @@ import scala.collection.mutable
import scala.concurrent.Future
import org.apache.spark._
-import org.apache.spark.rdd.MapPartitionsRDD
import org.apache.spark.rdd.RDD
-import org.apache.spark.scheduler.MapStatus
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
import org.apache.spark.shuffle.ShuffleWriteProcessor
import org.apache.spark.sql.auron.NativeHelper
-import org.apache.spark.sql.auron.NativeRDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.plans.logical.Statistics
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.execution._
-import org.apache.spark.sql.execution.auron.shuffle.AuronRssShuffleWriterBase
-import org.apache.spark.sql.execution.auron.shuffle.AuronShuffleWriterBase
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.metric.SQLMetrics
import org.apache.spark.sql.execution.metric.SQLShuffleReadMetricsReporter
@@ -126,12 +121,46 @@ case class NativeShuffleExchangeExec(
new
SQLShuffleWriteMetricsReporter(context.taskMetrics().shuffleWriteMetrics,
metrics)
}
+ @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5")
override def write(
rdd: RDD[_],
dep: ShuffleDependency[_, _, _],
mapId: Long,
context: TaskContext,
- partition: Partition): MapStatus = {
+ partition: Partition): org.apache.spark.scheduler.MapStatus = {
+ internalWrite(rdd, dep, mapId, context, partition)
+ }
+
+ @sparkver("4.1")
+ override def write(
+ inputs: Iterator[_],
+ dep: ShuffleDependency[_, _, _],
+ mapId: Long,
+ mapIndex: Int,
+ context: TaskContext): org.apache.spark.scheduler.MapStatus = {
+ import
org.apache.spark.sql.execution.auron.shuffle.AuronShuffleDependency
+
+ // SPARK-44605: Spark 4+ refines ShuffleWriteProcessor API, leading to
early execution of NativeRDD.ShuffleWrite iterator
+ // Adaptation: Return empty iterator in NativeRDD.compute() to defer
execution to ShuffleWriteProcessor.write() (align with Spark 3.x logic)
+ assert(
+ inputs.isEmpty,
+ "Input iterator must be empty (SPARK-44605: adapt to Spark 4+
ShuffleWriteProcessor API changes)")
+
+ val rdd = dep.asInstanceOf[AuronShuffleDependency[_, _, _]].inputRdd
+ val partition = rdd.partitions(mapIndex)
+ internalWrite(rdd, dep, mapId, context, partition)
+ }
+
+ private def internalWrite(
+ rdd: RDD[_],
+ dep: ShuffleDependency[_, _, _],
+ mapId: Long,
+ context: TaskContext,
+ partition: Partition): org.apache.spark.scheduler.MapStatus = {
+ import org.apache.spark.rdd.MapPartitionsRDD
+ import org.apache.spark.sql.auron.NativeRDD
+ import
org.apache.spark.sql.execution.auron.shuffle.AuronRssShuffleWriterBase
+ import
org.apache.spark.sql.execution.auron.shuffle.AuronShuffleWriterBase
val writer = SparkEnv.get.shuffleManager.getWriter(
dep.shuffleHandle,
@@ -165,7 +194,7 @@ case class NativeShuffleExchangeExec(
// for databricks testing
val causedBroadcastJoinBuildOOM = false
- @sparkver("3.5")
+ @sparkver("3.5 / 4.1")
override def advisoryPartitionSize: Option[Long] = None
// If users specify the num partitions via APIs like `repartition`, we
shouldn't change it.
@@ -174,17 +203,22 @@ case class NativeShuffleExchangeExec(
override def canChangeNumPartitions: Boolean =
outputPartitioning != SinglePartition
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def shuffleOrigin:
org.apache.spark.sql.execution.exchange.ShuffleOrigin = {
import org.apache.spark.sql.execution.exchange.ShuffleOrigin;
_shuffleOrigin.get.asInstanceOf[ShuffleOrigin]
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
@sparkver("3.0 / 3.1")
override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
copy(child = newChildren.head)
+
+ @sparkver("4.1")
+ override def shuffleId: Int = {
+ shuffleDependency.shuffleId
+ }
}
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeSortExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeSortExec.scala
index 05cc7236..1e1896f2 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeSortExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeSortExec.scala
@@ -27,7 +27,7 @@ case class NativeSortExec(
override val child: SparkPlan)
extends NativeSortBase(sortOrder, global, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedExec.scala
index 16548155..310d22f9 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedExec.scala
@@ -28,7 +28,7 @@ case class NativeTakeOrderedExec(
override val child: SparkPlan)
extends NativeTakeOrderedBase(limit, offset, sortOrder, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeUnionExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeUnionExec.scala
index 665346e3..8406f0fc 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeUnionExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeUnionExec.scala
@@ -26,7 +26,7 @@ case class NativeUnionExec(
override val output: Seq[Attribute])
extends NativeUnionBase(children, output) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildrenInternal(newChildren:
IndexedSeq[SparkPlan]): SparkPlan =
copy(children = newChildren)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeWindowExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeWindowExec.scala
index ac5ff045..7f2b2ff6 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeWindowExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeWindowExec.scala
@@ -31,7 +31,7 @@ case class NativeWindowExec(
override val child: SparkPlan)
extends NativeWindowBase(windowExpression, partitionSpec, orderSpec,
groupLimit, child) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
copy(child = newChild)
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronBlockStoreShuffleReader.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronBlockStoreShuffleReader.scala
index 14ebba48..64d239bd 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronBlockStoreShuffleReader.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronBlockStoreShuffleReader.scala
@@ -41,7 +41,7 @@ class AuronBlockStoreShuffleReader[K, C](
private val _ = mapOutputTracker
override def readBlocks(): Iterator[InputStream] = {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
def fetchIterator = new ShuffleBlockFetcherIterator(
context,
blockManager.blockStoreClient,
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronRssShuffleManagerBase.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronRssShuffleManagerBase.scala
index af92b0e8..e9bf42e5 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronRssShuffleManagerBase.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronRssShuffleManagerBase.scala
@@ -77,7 +77,7 @@ abstract class AuronRssShuffleManagerBase(_conf: SparkConf)
extends ShuffleManag
context: TaskContext,
metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V]
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getReader[K, C](
handle: ShuffleHandle,
startMapIndex: Int,
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleManager.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleManager.scala
index 4b84de3d..5a73cd77 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleManager.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleManager.scala
@@ -52,7 +52,7 @@ class AuronShuffleManager(conf: SparkConf) extends
ShuffleManager with Logging {
sortShuffleManager.registerShuffle(shuffleId, dependency)
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getReader[K, C](
handle: ShuffleHandle,
startMapIndex: Int,
@@ -67,7 +67,7 @@ class AuronShuffleManager(conf: SparkConf) extends
ShuffleManager with Logging {
@sparkver("3.2")
def shuffleMergeFinalized =
baseShuffleHandle.dependency.shuffleMergeFinalized
- @sparkver("3.3 / 3.4 / 3.5")
+ @sparkver("3.3 / 3.4 / 3.5 / 4.1")
def shuffleMergeFinalized =
baseShuffleHandle.dependency.isShuffleMergeFinalizedMarked
val (blocksByAddress, canEnableBatchFetch) =
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleWriter.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleWriter.scala
index 2ba99710..a6d57df1 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleWriter.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleWriter.scala
@@ -23,6 +23,6 @@ import org.apache.auron.sparkver
class AuronShuffleWriter[K, V](metrics: ShuffleWriteMetricsReporter)
extends AuronShuffleWriterBase[K, V](metrics) {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def getPartitionLengths(): Array[Long] = partitionLengths
}
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala
index 2f04829c..4244642e 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala
@@ -48,7 +48,7 @@ case class NativeBroadcastJoinExec(
override val condition: Option[Expression] = None
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def buildSide: org.apache.spark.sql.catalyst.optimizer.BuildSide =
broadcastSide match {
case JoinBuildLeft => org.apache.spark.sql.catalyst.optimizer.BuildLeft
@@ -61,7 +61,7 @@ case class NativeBroadcastJoinExec(
case JoinBuildRight => org.apache.spark.sql.execution.joins.BuildRight
}
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def requiredChildDistribution
: List[org.apache.spark.sql.catalyst.plans.physical.Distribution] = {
import org.apache.spark.sql.catalyst.plans.physical.BroadcastDistribution
@@ -80,22 +80,22 @@ case class NativeBroadcastJoinExec(
override def rewriteKeyExprToLong(exprs: Seq[Expression]): Seq[Expression] =
HashJoin.rewriteKeyExpr(exprs)
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def supportCodegen: Boolean = false
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def inputRDDs(): Nothing = {
throw new NotImplementedError("NativeBroadcastJoin dose not support
codegen")
}
- @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def prepareRelation(
ctx: org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext)
: org.apache.spark.sql.execution.joins.HashedRelationInfo = {
throw new NotImplementedError("NativeBroadcastJoin dose not support
codegen")
}
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override protected def withNewChildrenInternal(
newLeft: SparkPlan,
newRight: SparkPlan): SparkPlan =
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeShuffledHashJoinExecProvider.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeShuffledHashJoinExecProvider.scala
index 5f763eec..44db7391 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeShuffledHashJoinExecProvider.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeShuffledHashJoinExecProvider.scala
@@ -29,7 +29,7 @@ import org.apache.auron.sparkver
case object NativeShuffledHashJoinExecProvider {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
def provide(
left: SparkPlan,
right: SparkPlan,
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeSortMergeJoinExecProvider.scala
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeSortMergeJoinExecProvider.scala
index 067ae071..e77ec6a3 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeSortMergeJoinExecProvider.scala
+++
b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeSortMergeJoinExecProvider.scala
@@ -25,7 +25,7 @@ import org.apache.auron.sparkver
case object NativeSortMergeJoinExecProvider {
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.2 / 3.3 / 3.4 / 3.5 / 4.1")
def provide(
left: SparkPlan,
right: SparkPlan,
diff --git
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
index ae1d2432..5fc20370 100644
---
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
+++
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
@@ -20,7 +20,7 @@ import java.text.SimpleDateFormat
import org.apache.spark.sql.{AuronQueryTest, Row}
-import org.apache.auron.util.AuronTestUtils
+import org.apache.auron.util.{AuronTestUtils, SparkVersionUtil}
class AuronFunctionSuite extends AuronQueryTest with BaseAuronSQLSuite {
@@ -83,6 +83,9 @@ class AuronFunctionSuite extends AuronQueryTest with
BaseAuronSQLSuite {
}
test("spark hash function") {
+ // TODO: Fix flaky codegen cache failures in SPARK-4.x,
https://github.com/apache/auron/issues/1961
+ assume(!SparkVersionUtil.isSparkV40OrGreater)
+
withTable("t1") {
sql("create table t1 using parquet as select array(1, 2) as arr")
val functions =
@@ -94,6 +97,9 @@ class AuronFunctionSuite extends AuronQueryTest with
BaseAuronSQLSuite {
}
test("expm1 function") {
+ // TODO: Fix flaky codegen cache failures in SPARK-4.x,
https://github.com/apache/auron/issues/1961
+ assume(!SparkVersionUtil.isSparkV40OrGreater)
+
withTable("t1") {
sql("create table t1(c1 double) using parquet")
sql("insert into t1 values(0.0), (1.1), (2.2)")
@@ -416,8 +422,8 @@ class AuronFunctionSuite extends AuronQueryTest with
BaseAuronSQLSuite {
val sqlStr = s"""SELECT
|nvl2(null_int, int_val, 999) AS int_only,
- |nvl2(1, str_val, int_val) AS has_str,
- |nvl2(null_int, int_val, str_val) AS str_in_false,
+ |nvl2(1, str_val, cast(int_val AS STRING))
AS has_str,
+ |nvl2(null_int, cast(int_val AS STRING), str_val)
AS str_in_false,
|nvl2(1, arr_val, array(888)) AS has_array,
|nvl2(null_int, null_str, null_str) AS all_null
|FROM t1""".stripMargin
diff --git
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala
index e82eb78f..a73d17f2 100644
---
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala
+++
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala
@@ -20,7 +20,7 @@ import org.apache.spark.sql.{AuronQueryTest, Row}
import org.apache.spark.sql.execution.joins.auron.plan.NativeBroadcastJoinExec
import org.apache.auron.spark.configuration.SparkAuronConfiguration
-import org.apache.auron.util.AuronTestUtils
+import org.apache.auron.util.{AuronTestUtils, SparkVersionUtil}
class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with
AuronSQLTestHelper {
import testImplicits._
@@ -42,6 +42,9 @@ class AuronQuerySuite extends AuronQueryTest with
BaseAuronSQLSuite with AuronSQ
}
test("test filter with year function") {
+ // TODO: Fix flaky codegen cache failures in SPARK-4.x,
https://github.com/apache/auron/issues/1961
+ assume(!SparkVersionUtil.isSparkV40OrGreater)
+
withTable("t1") {
sql("create table t1 using parquet as select '2024-12-18' as event_time")
checkSparkAnswerAndOperator(s"""
@@ -54,6 +57,9 @@ class AuronQuerySuite extends AuronQueryTest with
BaseAuronSQLSuite with AuronSQ
}
test("test select multiple spark ext functions with the same signature") {
+ // TODO: Fix flaky codegen cache failures in SPARK-4.x,
https://github.com/apache/auron/issues/1961
+ assume(!SparkVersionUtil.isSparkV40OrGreater)
+
withTable("t1") {
sql("create table t1 using parquet as select '2024-12-18' as event_time")
checkSparkAnswerAndOperator("select year(event_time), month(event_time)
from t1")
@@ -171,6 +177,9 @@ class AuronQuerySuite extends AuronQueryTest with
BaseAuronSQLSuite with AuronSQ
}
test("floor function with long input") {
+ // TODO: Fix flaky codegen cache failures in SPARK-4.x,
https://github.com/apache/auron/issues/1961
+ assume(!SparkVersionUtil.isSparkV40OrGreater)
+
withTable("t1") {
sql("create table t1 using parquet as select 1L as c1, 2.2 as c2")
checkSparkAnswerAndOperator("select floor(c1), floor(c2) from t1")
diff --git
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/BaseAuronSQLSuite.scala
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/BaseAuronSQLSuite.scala
index 315d4aab..587d8f96 100644
---
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/BaseAuronSQLSuite.scala
+++
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/BaseAuronSQLSuite.scala
@@ -58,5 +58,8 @@ trait BaseAuronSQLSuite extends SharedSparkSession {
.set("spark.auron.enable", "true")
.set("spark.ui.enabled", "false")
.set("spark.sql.warehouse.dir", warehouseDir)
+ // Avoid the code size overflow error in Spark code generation.
+ .set("spark.sql.codegen.wholeStage", "false")
+ .set("spark.sql.codegen.factoryMode", "NO_CODEGEN")
}
}
diff --git
a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/execution/AuronAdaptiveQueryExecSuite.scala
b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/execution/AuronAdaptiveQueryExecSuite.scala
index fc4757b8..dc4f9ff5 100644
---
a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/execution/AuronAdaptiveQueryExecSuite.scala
+++
b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/execution/AuronAdaptiveQueryExecSuite.scala
@@ -20,7 +20,7 @@ import org.apache.spark.sql.AuronQueryTest
import org.apache.auron.{sparkverEnableMembers, BaseAuronSQLSuite}
-@sparkverEnableMembers("3.5")
+@sparkverEnableMembers("3.5 / 4.1")
class AuronAdaptiveQueryExecSuite extends AuronQueryTest with
BaseAuronSQLSuite {
import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent}
diff --git
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleWriter.scala
b/spark-extension/src/main/scala/org/apache/auron/util/SparkVersionUtil.scala
similarity index 68%
copy from
spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleWriter.scala
copy to
spark-extension/src/main/scala/org/apache/auron/util/SparkVersionUtil.scala
index 2ba99710..9b95e723 100644
---
a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleWriter.scala
+++
b/spark-extension/src/main/scala/org/apache/auron/util/SparkVersionUtil.scala
@@ -14,15 +14,11 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.sql.execution.auron.shuffle
+package org.apache.auron.util
-import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
+import org.apache.spark.SPARK_VERSION
-import org.apache.auron.sparkver
-
-class AuronShuffleWriter[K, V](metrics: ShuffleWriteMetricsReporter)
- extends AuronShuffleWriterBase[K, V](metrics) {
-
- @sparkver("3.2 / 3.3 / 3.4 / 3.5")
- override def getPartitionLengths(): Array[Long] = partitionLengths
+object SparkVersionUtil {
+ lazy val SPARK_RUNTIME_VERSION: SemanticVersion =
SemanticVersion(SPARK_VERSION)
+ lazy val isSparkV40OrGreater: Boolean = SPARK_RUNTIME_VERSION >= "4.0"
}
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeRDD.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeRDD.scala
index c6205d17..11c62185 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeRDD.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeRDD.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.auron.metric.SparkMetricNode
import org.apache.auron.protobuf.PhysicalPlanNode
+import org.apache.auron.util.SparkVersionUtil
class NativeRDD(
@transient private val rddSparkContext: SparkContext,
@@ -65,7 +66,15 @@ class NativeRDD(
override def compute(split: Partition, context: TaskContext):
Iterator[InternalRow] = {
val computingNativePlan = nativePlanWrapper.plan(split, context)
- NativeHelper.executeNativePlan(computingNativePlan, metrics, split,
Some(context))
+
+ // SPARK-44605: Spark 4+ refines ShuffleWriteProcessor API (early
execution of NativeRDD.ShuffleWrite iterator)
+ // Adaptation for Spark 4.x: Defer NativeRDD.ShuffleWrite execution to
ShuffleWriteProcessor.write() to align with Spark 3.x logic
+ if (SparkVersionUtil.isSparkV40OrGreater &&
+ computingNativePlan.getPhysicalPlanTypeCase ==
PhysicalPlanNode.PhysicalPlanTypeCase.SHUFFLE_WRITER) {
+ Iterator.empty
+ } else {
+ NativeHelper.executeNativePlan(computingNativePlan, metrics, split,
Some(context))
+ }
}
}
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarArray.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarArray.scala
index 988875ad..4930f991 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarArray.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarArray.scala
@@ -31,8 +31,9 @@ import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.types.LongType
import org.apache.spark.sql.types.ShortType
import org.apache.spark.sql.types.TimestampType
-import org.apache.spark.unsafe.types.CalendarInterval
-import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+
+import org.apache.auron.sparkver
class AuronColumnarArray(data: AuronColumnVector, offset: Int, length: Int)
extends ArrayData {
override def numElements: Int = length
@@ -154,4 +155,19 @@ class AuronColumnarArray(data: AuronColumnVector, offset:
Int, length: Int) exte
override def setNullAt(ordinal: Int): Unit = {
throw new UnsupportedOperationException
}
+
+ @sparkver("4.1")
+ override def getGeography(i: Int):
org.apache.spark.unsafe.types.GeographyVal = {
+ throw new UnsupportedOperationException
+ }
+
+ @sparkver("4.1")
+ override def getGeometry(i: Int): org.apache.spark.unsafe.types.GeometryVal
= {
+ throw new UnsupportedOperationException
+ }
+
+ @sparkver("4.1")
+ override def getVariant(i: Int): org.apache.spark.unsafe.types.VariantVal = {
+ throw new UnsupportedOperationException
+ }
}
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarBatchRow.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarBatchRow.scala
index 6b24e0f5..62c6ed96 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarBatchRow.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarBatchRow.scala
@@ -37,6 +37,8 @@ import org.apache.spark.sql.types.StructType
import org.apache.spark.unsafe.types.CalendarInterval
import org.apache.spark.unsafe.types.UTF8String
+import org.apache.auron.sparkver
+
class AuronColumnarBatchRow(columns: Array[AuronColumnVector], var rowId: Int
= 0)
extends InternalRow {
override def numFields: Int = columns.length
@@ -133,4 +135,19 @@ class AuronColumnarBatchRow(columns:
Array[AuronColumnVector], var rowId: Int =
override def setNullAt(ordinal: Int): Unit = {
throw new UnsupportedOperationException
}
+
+ @sparkver("4.1")
+ override def getGeography(i: Int):
org.apache.spark.unsafe.types.GeographyVal = {
+ throw new UnsupportedOperationException
+ }
+
+ @sparkver("4.1")
+ override def getGeometry(i: Int): org.apache.spark.unsafe.types.GeometryVal
= {
+ throw new UnsupportedOperationException
+ }
+
+ @sparkver("4.1")
+ override def getVariant(i: Int): org.apache.spark.unsafe.types.VariantVal = {
+ throw new UnsupportedOperationException
+ }
}
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarStruct.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarStruct.scala
index 34e7a717..75842e6e 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarStruct.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/columnar/AuronColumnarStruct.scala
@@ -34,8 +34,9 @@ import org.apache.spark.sql.types.MapType
import org.apache.spark.sql.types.ShortType
import org.apache.spark.sql.types.StringType
import org.apache.spark.sql.types.StructType
-import org.apache.spark.unsafe.types.CalendarInterval
-import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+
+import org.apache.auron.sparkver
class AuronColumnarStruct(data: AuronColumnVector, rowId: Int) extends
InternalRow {
override def numFields: Int = data.dataType.asInstanceOf[StructType].size
@@ -143,4 +144,19 @@ class AuronColumnarStruct(data: AuronColumnVector, rowId:
Int) extends InternalR
override def setNullAt(ordinal: Int): Unit = {
throw new UnsupportedOperationException
}
+
+ @sparkver("4.1")
+ override def getGeography(i: Int):
org.apache.spark.unsafe.types.GeographyVal = {
+ throw new UnsupportedOperationException
+ }
+
+ @sparkver("4.1")
+ override def getGeometry(i: Int): org.apache.spark.unsafe.types.GeometryVal
= {
+ throw new UnsupportedOperationException
+ }
+
+ @sparkver("4.1")
+ override def getVariant(i: Int): org.apache.spark.unsafe.types.VariantVal = {
+ throw new UnsupportedOperationException
+ }
}
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeBase.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeBase.scala
index b37dae19..d51e0830 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeBase.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeBase.scala
@@ -35,7 +35,6 @@ import org.apache.spark.TaskContext
import org.apache.spark.broadcast
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.auron.NativeConverters
import org.apache.spark.sql.auron.NativeHelper
import org.apache.spark.sql.auron.NativeRDD
@@ -61,9 +60,10 @@ import
org.apache.spark.sql.execution.exchange.BroadcastExchangeLike
import org.apache.spark.sql.execution.joins.HashedRelationBroadcastMode
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.BinaryType
-import org.apache.auron.{protobuf => pb}
+import org.apache.auron.{protobuf => pb, sparkver}
import org.apache.auron.jni.JniBridge
import org.apache.auron.metric.SparkMetricNode
@@ -138,8 +138,7 @@ abstract class NativeBroadcastExchangeBase(mode:
BroadcastMode, override val chi
}
def doExecuteBroadcastNative[T](): broadcast.Broadcast[T] = {
- val conf = SparkSession.getActiveSession.map(_.sqlContext.conf).orNull
- val timeout: Long = conf.broadcastTimeout
+ val timeout: Long = SQLConf.get.broadcastTimeout
try {
relationFuture.get(timeout,
TimeUnit.SECONDS).asInstanceOf[broadcast.Broadcast[T]]
} catch {
@@ -258,23 +257,39 @@ abstract class NativeBroadcastExchangeBase(mode:
BroadcastMode, override val chi
lazy val relationFuturePromise: Promise[Broadcast[Any]] =
Promise[Broadcast[Any]]()
@transient
- lazy val relationFuture: Future[Broadcast[Any]] = {
+ lazy val relationFuture: Future[Broadcast[Any]] = getRelationFuture
+
+ @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ private def getRelationFuture = {
SQLExecution.withThreadLocalCaptured[Broadcast[Any]](
Shims.get.getSqlContext(this).sparkSession,
BroadcastExchangeExec.executionContext) {
- try {
- sparkContext.setJobGroup(
- getRunId.toString,
- s"native broadcast exchange (runId $getRunId)",
- interruptOnCancel = true)
- val broadcasted =
sparkContext.broadcast(collectNative().asInstanceOf[Any])
- relationFuturePromise.trySuccess(broadcasted)
- broadcasted
- } catch {
- case e: Throwable =>
- relationFuturePromise.tryFailure(e)
- throw e
- }
+ executeBroadcastJob()
+ }
+ }
+
+ @sparkver("4.1")
+ private def getRelationFuture = {
+ SQLExecution.withThreadLocalCaptured[Broadcast[Any]](
+ this.session.sqlContext.sparkSession,
+ BroadcastExchangeExec.executionContext) {
+ executeBroadcastJob()
+ }
+ }
+
+ private def executeBroadcastJob(): Broadcast[Any] = {
+ try {
+ sparkContext.setJobGroup(
+ getRunId.toString,
+ s"native broadcast exchange (runId $getRunId)",
+ interruptOnCancel = true)
+ val broadcasted =
sparkContext.broadcast(collectNative().asInstanceOf[Any])
+ relationFuturePromise.trySuccess(broadcasted)
+ broadcasted
+ } catch {
+ case e: Throwable =>
+ relationFuturePromise.tryFailure(e)
+ throw e
}
}
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleDependency.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleDependency.scala
index f4b3070c..5243958d 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleDependency.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/shuffle/AuronShuffleDependency.scala
@@ -25,6 +25,8 @@ import org.apache.spark.serializer.Serializer
import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleHandle,
ShuffleWriteProcessor}
import org.apache.spark.sql.types.StructType
+import org.apache.auron.sparkver
+
class AuronShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag](
@transient private val _rdd: RDD[_ <: Product2[K, V]],
override val partitioner: Partitioner,
@@ -41,7 +43,19 @@ class AuronShuffleDependency[K: ClassTag, V: ClassTag, C:
ClassTag](
keyOrdering,
aggregator,
mapSideCombine,
- shuffleWriterProcessor) {}
+ shuffleWriterProcessor) {
+
+ // Serialize _rdd
+ val inputRdd: RDD[_ <: Product2[K, V]] = getInputRdd
+
+ // Spark 3+: Not required
+ @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ def getInputRdd: RDD[_ <: Product2[K, V]] = null
+
+ // For Spark 4+ compatibility: _rdd is required to create
NativeRDD.ShuffleWrite in ShuffleWriteProcessor.write
+ @sparkver("4.1")
+ def getInputRdd: RDD[_ <: Product2[K, V]] = _rdd
+}
object AuronShuffleDependency extends Logging {
def isArrowShuffle(handle: ShuffleHandle): Boolean = {
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/auron/plan/NativeHiveTableScanBase.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/auron/plan/NativeHiveTableScanBase.scala
index 24db3a7d..2e57ba6e 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/auron/plan/NativeHiveTableScanBase.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/auron/plan/NativeHiveTableScanBase.scala
@@ -143,7 +143,7 @@ abstract class NativeHiveTableScanBase(basedHiveScan:
HiveTableScanExec)
override protected def doCanonicalize(): SparkPlan =
basedHiveScan.canonicalized
- @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.1")
override def simpleString(maxFields: Int): String =
s"$nodeName (${basedHiveScan.simpleString(maxFields)})"
}
