You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@gluten.apache.org by we...@apache.org on 2026/09/29 09:16:26 UTC
(gluten) branch main updated: [GLUTEN-12569][VL] Add Spark 4.2 shim layer and build profile (#13126)
This is an automated email from the ASF dual-hosted git repository.
weiting-chen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 198c810200 [GLUTEN-12569][VL] Add Spark 4.2 shim layer and build profile (#13126)
198c810200 is described below
commit 198c810200b87b4e4ccd5eb10f3f9af5c5c2d734
Author: Akshay Tayal <50...@users.noreply.github.com>
AuthorDate: Tue Sep 29 02:16:19 2026 -0700
[GLUTEN-12569][VL] Add Spark 4.2 shim layer and build profile (#13126)
* Spark: Move 4.1 as 4.2
* Spark: Copy back 4.2 as 4.1
* Spark: Initial support for 4.2.0
Add the Spark 4.2 shim layer and build profile for the Velox backend. The
shims/spark42 module is created from shims/spark41 via a history-preserving
move + copy-back (apache/gluten discussion #11352) so blame/log follow the
spark42 files back through spark41; the changes folded into this commit are:
| Cause | Type | Category | Description | Affected Files |
|-------|------|----------|-------------|----------------|
| - | Feat | Feature | Introduce Spark42Shims and the spark-4.2 build configuration | pom.xml<br>shims/pom.xml<br>shims/spark42/pom.xml<br>shims/spark42/.../shims/spark42/Spark42Shims.scala<br>shims/spark42/.../shims/spark42/SparkShimProvider.scala<br>shims/spark42/.../META-INF/services/org.apache.gluten.sql.shims.SparkShimProvider |
| - | Feat | Build | Add the spark-4.2 profile to gluten-it | tools/gluten-it/pom.xml |
| - | Fix | Compatibility | Implement the new SparkShims methods for Spark 4.2 | shims/spark42/.../shims/spark42/Spark42Shims.scala<br>shims/spark42/.../v2/BatchScanExecShim.scala |
| - | Fix | Compatibility | Fix the MemoryStream shim for Spark 4.2 | shims/spark42/.../streaming/MemoryStream.scala |
| - | Fix | Compatibility | Fix the PythonUDFRunner.writeUDFs call for Spark 4.2 | shims/spark42/.../python/BasePythonRunnerShim.scala |
| - | Fix | Compatibility | Adapt the columnar vector shims to Spark 4.2 | shims/spark42/.../vectorized/WritableColumnVectorShim.java<br>shims/spark42/.../vectorized/ArrowColumnarArray.java |
| - | Fix | Compatibility | Port the BatchScanExec storage-partitioned-join shims to Spark 4.2 | shims/spark42/.../shims/spark42/Spark42Shims.scala<br>shims/spark42/.../v2/AbstractBatchScanExec.scala<br>shims/spark42/.../v2/BatchScanExecShim.scala |
| - | Fix | Compatibility | Remove the stale widerDecimalType override so shims/spark42 compiles under -Pspark-4.2 | shims/spark42/.../shims/spark42/Spark42Shims.scala |
| - | Fix | Compatibility | Adapt the Arrow/Python UDF command framing to Spark 4.2 (runnerConf/evalConf); pre-4.2 behavior unchanged | backends-velox/.../python/ColumnarArrowEvalPythonExec.scala<br>gluten-core/.../util/SparkVersionUtil.scala<br>shims/spark3[45]/.../python/BasePythonRunnerShim.scala<br>shims/spark4[012]/.../python/BasePythonRunnerShim.scala |
| - | Fix | CI | Add the Spark 4.2 CI resources, change detection and a spark-4.2 compile lane | .github/workflows/velox_backend_x86.yml<br>.github/workflows/util/install-spark-resources.sh<br>dev/docker/Dockerfile.centos8-dynamic-build<br>dev/docker/Dockerfile.centos9-dynamic-build<br>dev/format-scala-code.sh |
| - | Fix | CI | Honor INSTALL_DIR in the Spark 4.2 resource install | .github/workflows/util/install-spark-resources.sh |
| - | Fix | Build | Preserve the Spark 4.2 bundle across subsequent clean builds | package/pom.xml |
| - | Fix | Build | Register the Spark 4.2 shim in the developer test helpers | dev/bloop-test.sh<br>dev/run-scala-test.sh |
| - | Fix | Build | Use the canonical ASF license header in shims/spark42 | shims/spark42/pom.xml |
| - | Fix | Compatibility | Document the Spark 4.2 key-grouped partitioning (SPJ) limitation (tracked in apache/gluten#13139) | shims/spark42/.../v2/BatchScanExecShim.scala |
---------
Co-authored-by: Akshay Tayal <ak...@users.noreply.github.com>
---
.github/workflows/util/install-spark-resources.sh | 7 +
.github/workflows/velox_backend_x86.yml | 21 +-
.../api/python/ColumnarArrowEvalPythonExec.scala | 29 +-
dev/bloop-test.sh | 1 +
dev/docker/Dockerfile.centos8-dynamic-build | 1 +
dev/docker/Dockerfile.centos9-dynamic-build | 1 +
dev/format-scala-code.sh | 2 +-
dev/run-scala-test.sh | 1 +
.../org/apache/spark/util/SparkVersionUtil.scala | 1 +
package/pom.xml | 1 +
pom.xml | 90 +-
shims/pom.xml | 6 +
.../execution/python/BasePythonRunnerShim.scala | 5 +
.../execution/python/BasePythonRunnerShim.scala | 5 +
.../execution/python/BasePythonRunnerShim.scala | 5 +
.../execution/python/BasePythonRunnerShim.scala | 5 +
shims/spark42/pom.xml | 121 ++
.../vectorized/WritableColumnVectorShim.java | 237 +++
.../spark/sql/vectorized/ArrowColumnarArray.java | 30 +-
.../org.apache.gluten.sql.shims.SparkShimProvider | 1 +
.../gluten/execution/GenerateTreeStringShim.scala | 93 ++
.../gluten/execution/PartitionedFileUtilShim.scala | 174 +++
.../gluten/sql/shims/spark42/Spark42Shims.scala | 490 ++++++
.../sql/shims/spark42/SparkShimProvider.scala | 18 +-
.../org/apache/gluten/utils/InternalRowUtil.scala | 26 +-
.../main/scala/org/apache/spark/ShuffleUtils.scala | 48 +
.../scala/org/apache/spark/SparkContextUtils.scala | 29 +-
.../apache/spark/shuffle/GlutenMapStatusUtil.scala | 30 +-
.../spark/shuffle/SparkSortShuffleWriterUtil.scala | 25 +-
.../catalyst/expressions/PromotePrecision.scala | 29 +-
.../expressions/objects/InvokeExtractors.scala | 46 +
.../catalyst/optimizer/CollapseProjectShim.scala | 27 +-
.../apache/spark/sql/classic/ClassicColumn.scala | 24 +-
.../apache/spark/sql/classic/ClassicDataset.scala | 21 +-
.../apache/spark/sql/classic/ClassicTypes.scala | 23 +-
.../org/apache/spark/sql/classic/conversions.scala | 24 +-
.../sql/execution/AbstractFileSourceScanExec.scala | 307 ++++
.../execution/ExpandOutputPartitioningShim.scala | 93 ++
.../sql/execution/FileSourceScanExecShim.scala | 286 ++++
.../sql/execution/GlutenFileFormatWriter.scala | 49 +
...PartitioningAndOrderingPreservingNodeShim.scala | 19 +-
.../parquet/ParquetFooterReaderShim.scala | 43 +
.../datasources/v2/AbstractBatchScanExec.scala | 156 ++
.../datasources/v2/BatchScanExecShim.scala | 161 ++
.../datasources/v2/utils/CatalogUtil.scala | 22 +-
.../execution/python/BasePythonRunnerShim.scala | 25 +-
.../sql/execution/python/EvalPythonExecBase.scala | 26 +-
.../sql/execution/streaming/MemoryStream.scala | 25 +-
.../apache/spark/sql/execution/ui/TypeAlias.scala | 22 +-
.../hive/execution/AbstractHiveTableScanExec.scala | 277 ++++
.../spark/storage/GlutenPushBasedFetchHelper.scala | 399 +++++
.../GlutenShuffleBlockFetcherIterator.scala | 1601 ++++++++++++++++++++
.../spark42/Spark42LocalTableScanStreamSuite.scala | 59 +
tools/gluten-it/pom.xml | 12 +
54 files changed, 5039 insertions(+), 240 deletions(-)
diff --git a/.github/workflows/util/install-spark-resources.sh b/.github/workflows/util/install-spark-resources.sh
index 98b0d1bd00..84d01bdfbf 100755
--- a/.github/workflows/util/install-spark-resources.sh
+++ b/.github/workflows/util/install-spark-resources.sh
@@ -122,6 +122,13 @@ if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
install_spark "4.1.1" "3" "2.12"
mv /opt/shims/spark41/spark_home/assembly/target/scala-2.12 /opt/shims/spark41/spark_home/assembly/target/scala-2.13
;;
+ 4.2)
+ # Spark-4.x, scala 2.12 // using 2.12 as a hack as 4.2 does not have a 2.13 suffix
+ cd ${INSTALL_DIR} && \
+ install_spark "4.2.0" "3" "2.12"
+ mv "${INSTALL_DIR}/shims/spark42/spark_home/assembly/target/scala-2.12" \
+ "${INSTALL_DIR}/shims/spark42/spark_home/assembly/target/scala-2.13"
+ ;;
*)
echo "Spark version is expected to be specified."
exit 1
diff --git a/.github/workflows/velox_backend_x86.yml b/.github/workflows/velox_backend_x86.yml
index 349a234c3a..66da3898b4 100644
--- a/.github/workflows/velox_backend_x86.yml
+++ b/.github/workflows/velox_backend_x86.yml
@@ -97,6 +97,7 @@ jobs:
shims35: ${{ steps.filter.outputs.shims35 }}
shims40: ${{ steps.filter.outputs.shims40 }}
shims41: ${{ steps.filter.outputs.shims41 }}
+ shims42: ${{ steps.filter.outputs.shims42 }}
tools_it: ${{ steps.filter.outputs.tools_it }}
steps:
- uses: actions/checkout@v7
@@ -114,7 +115,7 @@ jobs:
recent=$(git log --oneline --after="24 hours ago" HEAD | wc -l)
if [ "$recent" -eq 0 ]; then
echo "No commits in the last 24 hours -- skipping nightly run."
- for flag in cpp java shims34 shims35 shims40 shims41 tools_it; do
+ for flag in cpp java shims34 shims35 shims40 shims41 shims42 tools_it; do
echo "$flag=false" >> $GITHUB_OUTPUT
done
echo "full_run=false" >> $GITHUB_OUTPUT
@@ -122,7 +123,7 @@ jobs:
fi
fi
# Run everything, including the extended matrices PR runs skip.
- for flag in cpp java shims34 shims35 shims40 shims41 tools_it; do
+ for flag in cpp java shims34 shims35 shims40 shims41 shims42 tools_it; do
echo "$flag=true" >> $GITHUB_OUTPUT
done
echo "full_run=true" >> $GITHUB_OUTPUT
@@ -142,6 +143,7 @@ jobs:
echo "shims35=$(match '^(shims/(spark35|common)/|gluten-ut/spark35/)')" >> $GITHUB_OUTPUT
echo "shims40=$(match '^(shims/(spark40|common)/|gluten-ut/spark40/)')" >> $GITHUB_OUTPUT
echo "shims41=$(match '^(shims/(spark41|common)/|gluten-ut/spark41/)')" >> $GITHUB_OUTPUT
+ echo "shims42=$(match '^(shims/(spark42|common)/|gluten-ut/spark42/)')" >> $GITHUB_OUTPUT
echo "tools_it=$(match '^tools/gluten-it/')" >> $GITHUB_OUTPUT
build-native-lib-centos-8:
@@ -1244,6 +1246,21 @@ jobs:
cd $GITHUB_WORKSPACE/
$MVN_CMD clean test-compile -Pspark-4.0 -Pscala-2.13 -Pbackends-velox -Pspark-ut -Piceberg,iceberg-test,delta,paimon -Pfast-build
+ build-test-spark42:
+ needs: detect-changes
+ if: >-
+ needs.detect-changes.outputs.java == 'true' ||
+ needs.detect-changes.outputs.shims42 == 'true' ||
+ needs.detect-changes.outputs.cpp == 'true'
+ runs-on: ubuntu-22.04
+ container: apache/gluten:centos-8-jdk17
+ steps:
+ - uses: actions/checkout@v7
+ - name: Build Spark 4.2 Velox (Java 17, Scala 2.13)
+ run: |
+ cd $GITHUB_WORKSPACE/
+ $MVN_CMD clean install -Pspark-4.2 -Pjava-17 -Pscala-2.13 -Pbackends-velox -DskipTests
+
spark-test-spark40:
needs: [detect-changes, build-native-lib-centos-8]
if: >-
diff --git a/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala b/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
index 831259b012..ef1bc4a2e7 100644
--- a/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
+++ b/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
@@ -63,6 +63,13 @@ class ColumnarArrowPythonRunner(
// Keep this source compatible with older Spark profiles where PythonEvalType differs.
private val SQL_ARROW_BATCHED_UDF = 101
+ // On Spark 4.2 the runner config and (for Arrow-batched UDFs) the input schema are delivered to
+ // the worker via runnerConf/evalConf instead of the writeCommand prefix (see BasePythonRunnerShim
+ // and Spark's ArrowPythonRunner/ArrowPythonWithNamedArgumentRunner). These are no-ops on <4.2.
+ override protected def pythonRunnerConfMap: Map[String, String] = conf
+
+ override protected def pythonInputSchema: StructType = schema
+
override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
require(
bufferSize >= 4,
@@ -148,14 +155,20 @@ class ColumnarArrowPythonRunner(
context: TaskContext): Writer = {
new Writer(env, worker, inputIterator, partitionIndex, context) {
override protected def writeCommand(dataOut: DataOutputStream): Unit = {
- // Write config for the worker as a number of key -> value pairs of strings
- dataOut.writeInt(conf.size)
- for ((k, v) <- conf) {
- PythonRDD.writeUTF(k, dataOut)
- PythonRDD.writeUTF(v, dataOut)
- }
- if (SparkVersionUtil.gteSpark41 && evalType == SQL_ARROW_BATCHED_UDF) {
- PythonRDD.writeUTF(schema.json, dataOut)
+ // Spark 4.2 (SPARK-51384) writes evalType -> runnerConf -> evalConf -> writeCommand, so the
+ // config and input schema flow through runnerConf/evalConf (see BasePythonRunnerShim) and
+ // writeCommand must emit only the UDF definitions. Older profiles carry the config prefix
+ // here, exactly as before.
+ if (!SparkVersionUtil.gteSpark42) {
+ // Write config for the worker as a number of key -> value pairs of strings
+ dataOut.writeInt(conf.size)
+ for ((k, v) <- conf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ if (SparkVersionUtil.gteSpark41 && evalType == SQL_ARROW_BATCHED_UDF) {
+ PythonRDD.writeUTF(schema.json, dataOut)
+ }
}
ColumnarArrowPythonRunner.this.writeUdf(dataOut, argMetas)
}
diff --git a/dev/bloop-test.sh b/dev/bloop-test.sh
index 1d4730e118..a341385059 100755
--- a/dev/bloop-test.sh
+++ b/dev/bloop-test.sh
@@ -77,6 +77,7 @@ declare -A MODULE_MAP=(
["shims/spark35"]="spark-sql-columnar-shims-spark35"
["shims/spark40"]="spark-sql-columnar-shims-spark40"
["shims/spark41"]="spark-sql-columnar-shims-spark41"
+ ["shims/spark42"]="spark-sql-columnar-shims-spark42"
# Unit test modules
["gluten-ut/common"]="gluten-ut-common"
diff --git a/dev/docker/Dockerfile.centos8-dynamic-build b/dev/docker/Dockerfile.centos8-dynamic-build
index bf960e61a7..388f71f922 100644
--- a/dev/docker/Dockerfile.centos8-dynamic-build
+++ b/dev/docker/Dockerfile.centos8-dynamic-build
@@ -54,6 +54,7 @@ RUN set -ex; \
./install-spark-resources.sh 3.5-scala2.13; \
./install-spark-resources.sh 4.0; \
./install-spark-resources.sh 4.1; \
+ ./install-spark-resources.sh 4.2; \
ARCH=$(uname -m); \
if [[ "$ARCH" == "aarch64" || "$ARCH" == "ppc64le" ]]; then \
export CPU_TARGET="$ARCH"; \
diff --git a/dev/docker/Dockerfile.centos9-dynamic-build b/dev/docker/Dockerfile.centos9-dynamic-build
index 92bac1d692..9a8a26378d 100644
--- a/dev/docker/Dockerfile.centos9-dynamic-build
+++ b/dev/docker/Dockerfile.centos9-dynamic-build
@@ -46,6 +46,7 @@ RUN set -ex; \
./install-spark-resources.sh 3.5-scala2.13; \
./install-spark-resources.sh 4.0; \
./install-spark-resources.sh 4.1; \
+ ./install-spark-resources.sh 4.2; \
ARCH=$(uname -m); \
if [[ "$ARCH" == "aarch64" || "$ARCH" == "ppc64le" ]]; then \
export CPU_TARGET="$ARCH"; \
diff --git a/dev/format-scala-code.sh b/dev/format-scala-code.sh
index d3ec8d4895..bb02aa89b6 100755
--- a/dev/format-scala-code.sh
+++ b/dev/format-scala-code.sh
@@ -22,7 +22,7 @@ MVN_CMD="${BASEDIR}/../build/mvn"
# If a new profile is introduced for new modules, please add it here to ensure
# the new modules are covered.
PROFILES="-Pbackends-velox,backends-clickhouse -Pceleborn,uniffle -Piceberg,delta,hudi,paimon \
- -Pspark-3.4,spark-3.5,spark-4.0,spark-4.1 -Pspark-ut -Pkafka"
+ -Pspark-3.4,spark-3.5,spark-4.0,spark-4.1,spark-4.2 -Pspark-ut -Pkafka"
COMMAND=$1
diff --git a/dev/run-scala-test.sh b/dev/run-scala-test.sh
index 8e98bca197..f702179ecf 100755
--- a/dev/run-scala-test.sh
+++ b/dev/run-scala-test.sh
@@ -164,6 +164,7 @@ declare -A MODULE_MAP=(
["spark-sql-columnar-shims-spark35"]="shims/spark35:java"
["spark-sql-columnar-shims-spark40"]="shims/spark40:java"
["spark-sql-columnar-shims-spark41"]="shims/spark41:java"
+ ["spark-sql-columnar-shims-spark42"]="shims/spark42:java"
# Unit test modules
["gluten-ut-common"]="gluten-ut/common:scala"
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
index f7326e98d6..11775c5b59 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
@@ -20,6 +20,7 @@ object SparkVersionUtil {
val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+ val gteSpark42: Boolean = compareMajorMinorVersion((4, 2)) >= 0
// Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
def compareMajorMinorVersion(other: (Int, Int)): Int = {
diff --git a/package/pom.xml b/package/pom.xml
index 1ad1c06b5d..8fdac49fe2 100644
--- a/package/pom.xml
+++ b/package/pom.xml
@@ -240,6 +240,7 @@
<exclude>*spark3.5*</exclude>
<exclude>*spark4.0*</exclude>
<exclude>*spark4.1*</exclude>
+ <exclude>*spark4.2*</exclude>
</excludes>
<followSymlinks>false</followSymlinks>
</fileset>
diff --git a/pom.xml b/pom.xml
index 68575cec4b..e01cc34d9a 100644
--- a/pom.xml
+++ b/pom.xml
@@ -979,7 +979,7 @@
<rules>
<requireActiveProfile>
<!-- Spark 4.0 requires Java 17+ -->
- <profiles>spark-3.4,spark-3.5,spark-4.0,spark-4.1</profiles>
+ <profiles>spark-3.4,spark-3.5,spark-4.0,spark-4.1,spark-4.2</profiles>
<all>false</all>
<message>Missing spark version profile: -Pspark-<version></message>
</requireActiveProfile>
@@ -1445,6 +1445,94 @@
</plugins>
</build>
</profile>
+ <profile>
+ <id>spark-4.2</id>
+ <properties>
+ <sparkbundle.version>4.2</sparkbundle.version>
+ <sparkshim.artifactId>spark-sql-columnar-shims-spark42</sparkshim.artifactId>
+ <spark.version>4.2.0</spark.version>
+ <!-- No Spark 4.2 artifacts published for Iceberg/Delta/Hudi yet. These values are
+ placeholders inherited from spark-4.1; the delta/iceberg/hudi profiles are
+ opt-in (activeByDefault=false) and are NOT enabled on the Spark 4.2 CI lane. -->
+ <iceberg.version>1.11.0</iceberg.version>
+ <delta.package.name>delta-spark_4.1</delta.package.name>
+ <delta.version>4.1.0</delta.version>
+ <delta.binary.version>40</delta.binary.version>
+ <hudi.version>1.2.0</hudi.version>
+ <fasterxml.version>2.18.2</fasterxml.version>
+ <fasterxml.jackson.version>2.18.2</fasterxml.jackson.version>
+ <fasterxml.jackson.databind.version>2.18.2</fasterxml.jackson.databind.version>
+ <hadoop.version>3.4.1</hadoop.version>
+ <antlr4.version>4.13.1</antlr4.version>
+ <guava.version>33.4.0-jre</guava.version>
+ <slf4j.version>2.0.16</slf4j.version>
+ <log4j.version>2.24.3</log4j.version>
+ <commons-lang3.version>3.17.0</commons-lang3.version>
+ <!-- Must track Spark 4.2's own Arrow (19.0.0): scope is provided and Gluten
+ deliberately lets Spark's Arrow through. -->
+ <arrow.version>19.0.0</arrow.version>
+ <arrow.deps.scope>provided</arrow.deps.scope>
+ <arrow-memory.scope>provided</arrow-memory.scope>
+ <!-- Matches no groupId: let Spark's own Arrow through. -->
+ <spark.arrow.exclusion.groupId>none</spark.arrow.exclusion.groupId>
+ <scala.compiler.version>4.9.5</scala.compiler.version>
+ </properties>
+ <dependencies>
+ <dependency>
+ <groupId>org.slf4j</groupId>
+ <artifactId>slf4j-api</artifactId>
+ <version>${slf4j.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-slf4j2-impl</artifactId>
+ <version>${log4j.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ </dependencies>
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-enforcer-plugin</artifactId>
+ <executions>
+ <execution>
+ <id>enforce-java-17+</id>
+ <goals>
+ <goal>enforce</goal>
+ </goals>
+ <configuration>
+ <rules>
+ <requireActiveProfile>
+ <!-- Spark 4.2 requires Java 17+ -->
+ <profiles>java-17,java-21,java-25</profiles>
+ <all>false</all>
+ <message>"-P spark-4.2" requires JDK 17+</message>
+ </requireActiveProfile>
+ </rules>
+ </configuration>
+ </execution>
+ <execution>
+ <id>enforce-scala-213</id>
+ <goals>
+ <goal>enforce</goal>
+ </goals>
+ <configuration>
+ <rules>
+ <requireActiveProfile>
+ <!-- Spark 4.2 requires Scala 2.13 -->
+ <profiles>scala-2.13</profiles>
+ <message>"-P spark-4.2" requires Scala 2.13</message>
+ </requireActiveProfile>
+ </rules>
+ </configuration>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </build>
+ </profile>
<profile>
<id>hadoop-2.7.4</id>
<properties>
diff --git a/shims/pom.xml b/shims/pom.xml
index 50ffa892b5..69e868fe0e 100644
--- a/shims/pom.xml
+++ b/shims/pom.xml
@@ -95,6 +95,12 @@
<module>spark41</module>
</modules>
</profile>
+ <profile>
+ <id>spark-4.2</id>
+ <modules>
+ <module>spark42</module>
+ </modules>
+ </profile>
<profile>
<id>default</id>
<activation>
diff --git a/shims/spark34/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala b/shims/spark34/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
index 82d971d5d6..a26c8b4b20 100644
--- a/shims/spark34/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
+++ b/shims/spark34/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
@@ -20,6 +20,7 @@ import org.apache.spark.SparkEnv
import org.apache.spark.TaskContext
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions}
import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.ColumnarBatch
import java.io.DataOutputStream
@@ -60,4 +61,8 @@ abstract class BasePythonRunnerShim(
context: TaskContext): Writer = {
createNewWriter(env, worker, inputIterator, partitionIndex, context)
}
+
+ protected def pythonRunnerConfMap: Map[String, String] = Map.empty
+
+ protected def pythonInputSchema: StructType = new StructType()
}
diff --git a/shims/spark35/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala b/shims/spark35/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
index ecabc9a93b..1280db2a9e 100644
--- a/shims/spark35/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
+++ b/shims/spark35/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
@@ -20,6 +20,7 @@ import org.apache.spark.SparkEnv
import org.apache.spark.TaskContext
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions}
import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.ColumnarBatch
import java.io.DataOutputStream
@@ -61,4 +62,8 @@ abstract class BasePythonRunnerShim(
context: TaskContext): Writer = {
createNewWriter(env, worker, inputIterator, partitionIndex, context)
}
+
+ protected def pythonRunnerConfMap: Map[String, String] = Map.empty
+
+ protected def pythonInputSchema: StructType = new StructType()
}
diff --git a/shims/spark40/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala b/shims/spark40/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
index 127a0fc3cf..fc2389dbff 100644
--- a/shims/spark40/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
+++ b/shims/spark40/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
@@ -21,6 +21,7 @@ import org.apache.spark.TaskContext
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions, PythonWorker}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata
+import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.ColumnarBatch
import java.io.DataOutputStream
@@ -63,4 +64,8 @@ abstract class BasePythonRunnerShim(
createNewWriter(env, worker, inputIterator, partitionIndex, context)
}
+ protected def pythonRunnerConfMap: Map[String, String] = Map.empty
+
+ protected def pythonInputSchema: StructType = new StructType()
+
}
diff --git a/shims/spark41/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala b/shims/spark41/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
index 127a0fc3cf..fc2389dbff 100644
--- a/shims/spark41/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
+++ b/shims/spark41/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
@@ -21,6 +21,7 @@ import org.apache.spark.TaskContext
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions, PythonWorker}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata
+import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.ColumnarBatch
import java.io.DataOutputStream
@@ -63,4 +64,8 @@ abstract class BasePythonRunnerShim(
createNewWriter(env, worker, inputIterator, partitionIndex, context)
}
+ protected def pythonRunnerConfMap: Map[String, String] = Map.empty
+
+ protected def pythonInputSchema: StructType = new StructType()
+
}
diff --git a/shims/spark42/pom.xml b/shims/spark42/pom.xml
new file mode 100644
index 0000000000..8d28104976
--- /dev/null
+++ b/shims/spark42/pom.xml
@@ -0,0 +1,121 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ ~ Licensed to the Apache Software Foundation (ASF) under one or more
+ ~ contributor license agreements. See the NOTICE file distributed with
+ ~ this work for additional information regarding copyright ownership.
+ ~ The ASF licenses this file to You under the Apache License, Version 2.0
+ ~ (the "License"); you may not use this file except in compliance with
+ ~ the License. You may obtain a copy of the License at
+ ~
+ ~ http://www.apache.org/licenses/LICENSE-2.0
+ ~
+ ~ Unless required by applicable law or agreed to in writing, software
+ ~ distributed under the License is distributed on an "AS IS" BASIS,
+ ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ ~ See the License for the specific language governing permissions and
+ ~ limitations under the License.
+ -->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+
+ <parent>
+ <groupId>org.apache.gluten</groupId>
+ <artifactId>spark-sql-columnar-shims</artifactId>
+ <version>1.8.0-SNAPSHOT</version>
+ <relativePath>../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>spark-sql-columnar-shims-spark42</artifactId>
+ <packaging>jar</packaging>
+ <name>Gluten Shims for Spark 4.2</name>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.gluten</groupId>
+ <artifactId>${project.prefix}-shims-common</artifactId>
+ <version>${project.version}</version>
+ <scope>compile</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ <scope>provided</scope>
+ <optional>true</optional>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-catalyst_${scala.binary.version}</artifactId>
+ <scope>provided</scope>
+ <optional>true</optional>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-core_${scala.binary.version}</artifactId>
+ <scope>provided</scope>
+ <optional>true</optional>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-common</artifactId>
+ <version>${hadoop.version}</version>
+ <scope>provided</scope>
+ </dependency>
+
+ <!--test-->
+ <dependency>
+ <groupId>org.scalatest</groupId>
+ <artifactId>scalatest_${scala.binary.version}</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-core_${scala.binary.version}</artifactId>
+ <type>test-jar</type>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ <type>test-jar</type>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-catalyst_${scala.binary.version}</artifactId>
+ <type>test-jar</type>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-hive_${scala.binary.version}</artifactId>
+ <scope>provided</scope>
+ </dependency>
+
+ </dependencies>
+
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-checkstyle-plugin</artifactId>
+ </plugin>
+ <plugin>
+ <groupId>org.scalastyle</groupId>
+ <artifactId>scalastyle-maven-plugin</artifactId>
+ </plugin>
+ <plugin>
+ <groupId>com.diffplug.spotless</groupId>
+ <artifactId>spotless-maven-plugin</artifactId>
+ </plugin>
+ <plugin>
+ <groupId>net.alchim31.maven</groupId>
+ <artifactId>scala-maven-plugin</artifactId>
+ </plugin>
+ <plugin>
+ <groupId>org.scalatest</groupId>
+ <artifactId>scalatest-maven-plugin</artifactId>
+ </plugin>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-compiler-plugin</artifactId>
+ </plugin>
+ </plugins>
+ </build>
+</project>
diff --git a/shims/spark42/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVectorShim.java b/shims/spark42/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVectorShim.java
new file mode 100644
index 0000000000..4933352125
--- /dev/null
+++ b/shims/spark42/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVectorShim.java
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.vectorized;
+
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.unsafe.types.BinaryView;
+import org.apache.spark.unsafe.types.UTF8String;
+
+import java.nio.ByteBuffer;
+
+/**
+ * because spark33 add new function abstract method 'putBooleans(int, byte)' in
+ * 'WritableColumnVector' And function getByteBuffer()
+ */
+public class WritableColumnVectorShim extends WritableColumnVector {
+ /**
+ * Sets up the common state and also handles creating the child columns if this is a nested type.
+ *
+ * @param capacity
+ * @param type
+ */
+ protected WritableColumnVectorShim(int capacity, DataType type) {
+ super(capacity, type);
+ }
+
+ protected void releaseMemory() {}
+
+ @Override
+ public int getDictId(int rowId) {
+ return 0;
+ }
+
+ @Override
+ protected void reserveInternal(int capacity) {}
+
+ @Override
+ public void putNotNull(int rowId) {}
+
+ @Override
+ public void putNull(int rowId) {}
+
+ @Override
+ public void putNulls(int rowId, int count) {}
+
+ @Override
+ public void putNotNulls(int rowId, int count) {}
+
+ @Override
+ public void putBoolean(int rowId, boolean value) {}
+
+ @Override
+ public void putBooleans(int rowId, int count, boolean value) {}
+
+ @Override
+ public void putBooleans(int rowId, byte src) {
+ throw new UnsupportedOperationException("Unsupported function");
+ }
+
+ @Override
+ public void putByte(int rowId, byte value) {}
+
+ @Override
+ public void putBytes(int rowId, int count, byte value) {}
+
+ @Override
+ public void putBytes(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putBytes(int rowId, int count, ByteBuffer src, int srcIndex) {}
+
+ @Override
+ public void putShort(int rowId, short value) {}
+
+ @Override
+ public void putShorts(int rowId, int count, short value) {}
+
+ @Override
+ public void putShorts(int rowId, int count, short[] src, int srcIndex) {}
+
+ @Override
+ public void putShorts(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putShortsFromIntsLittleEndian(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putInt(int rowId, int value) {}
+
+ @Override
+ public void putInts(int rowId, int count, int value) {}
+
+ @Override
+ public void putInts(int rowId, int count, int[] src, int srcIndex) {}
+
+ @Override
+ public void putInts(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putIntsLittleEndian(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putLong(int rowId, long value) {}
+
+ @Override
+ public void putLongs(int rowId, int count, long value) {}
+
+ @Override
+ public void putLongs(int rowId, int count, long[] src, int srcIndex) {}
+
+ @Override
+ public void putLongs(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putLongsLittleEndian(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putFloat(int rowId, float value) {}
+
+ @Override
+ public void putFloats(int rowId, int count, float value) {}
+
+ @Override
+ public void putFloats(int rowId, int count, float[] src, int srcIndex) {}
+
+ @Override
+ public void putFloats(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putFloatsLittleEndian(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putDouble(int rowId, double value) {}
+
+ @Override
+ public void putDoubles(int rowId, int count, double value) {}
+
+ @Override
+ public void putDoubles(int rowId, int count, double[] src, int srcIndex) {}
+
+ @Override
+ public void putDoubles(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putDoublesLittleEndian(int rowId, int count, byte[] src, int srcIndex) {}
+
+ @Override
+ public void putArray(int rowId, int offset, int length) {}
+
+ @Override
+ public int putByteArray(int rowId, byte[] value, int offset, int count) {
+ return 0;
+ }
+
+ @Override
+ protected UTF8String getBytesAsUTF8String(int rowId, int count) {
+ return null;
+ }
+
+ @Override
+ protected BinaryView getBytesAsBinaryView(int rowId, int count) {
+ return null;
+ }
+
+ @Override
+ public ByteBuffer getByteBuffer(int rowId, int count) {
+ throw new UnsupportedOperationException("Unsupported this function");
+ }
+
+ @Override
+ public int getArrayLength(int rowId) {
+ return 0;
+ }
+
+ @Override
+ public int getArrayOffset(int rowId) {
+ return 0;
+ }
+
+ @Override
+ public WritableColumnVector reserveNewColumn(int capacity, DataType type) {
+ return null;
+ }
+
+ @Override
+ public boolean isNullAt(int rowId) {
+ return false;
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return false;
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return 0;
+ }
+
+ @Override
+ public short getShort(int rowId) {
+ return 0;
+ }
+
+ @Override
+ public int getInt(int rowId) {
+ return 0;
+ }
+
+ @Override
+ public long getLong(int rowId) {
+ return 0;
+ }
+
+ @Override
+ public float getFloat(int rowId) {
+ return 0;
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return 0;
+ }
+}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnarArray.java
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnarArray.java
index f7326e98d6..939e0648d9 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnarArray.java
@@ -14,20 +14,24 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.vectorized;
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.unsafe.types.BinaryView;
+import org.apache.spark.unsafe.types.VariantVal;
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+public final class ArrowColumnarArray extends AbstractColumnarArray {
+
+ public ArrowColumnarArray(ColumnVector data, int offset, int length) {
+ super(data, offset, length);
+ }
+
+ @Override
+ public VariantVal getVariant(int ordinal) {
+ return data.getVariant(offset + ordinal);
+ }
+
+ @Override
+ public BinaryView getBinaryView(int ordinal) {
+ return data.getBinaryView(offset + ordinal);
}
}
diff --git a/shims/spark42/src/main/resources/META-INF/services/org.apache.gluten.sql.shims.SparkShimProvider b/shims/spark42/src/main/resources/META-INF/services/org.apache.gluten.sql.shims.SparkShimProvider
new file mode 100644
index 0000000000..3e56d5d48b
--- /dev/null
+++ b/shims/spark42/src/main/resources/META-INF/services/org.apache.gluten.sql.shims.SparkShimProvider
@@ -0,0 +1 @@
+org.apache.gluten.sql.shims.spark42.SparkShimProvider
\ No newline at end of file
diff --git a/shims/spark42/src/main/scala/org/apache/gluten/execution/GenerateTreeStringShim.scala b/shims/spark42/src/main/scala/org/apache/gluten/execution/GenerateTreeStringShim.scala
new file mode 100644
index 0000000000..3d329a8f31
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/gluten/execution/GenerateTreeStringShim.scala
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.execution
+
+import org.apache.spark.sql.execution.UnaryExecNode
+
+/**
+ * Spark 3.5 has changed the parameter type of the generateTreeString API in TreeNode. In order to
+ * support multiple versions of Spark, we cannot directly override the generateTreeString method in
+ * WhostageTransformer. Therefore, we have defined the GenerateTreeStringShim trait in the shim to
+ * allow different Spark versions to override their own generateTreeString.
+ */
+
+trait WholeStageTransformerGenerateTreeStringShim extends UnaryExecNode {
+
+ def stageId: Int
+
+ def substraitPlanJson: String
+
+ def wholeStageTransformerContextDefined: Boolean
+
+ override def generateTreeString(
+ depth: Int,
+ lastChildren: java.util.ArrayList[Boolean],
+ append: String => Unit,
+ verbose: Boolean,
+ prefix: String = "",
+ addSuffix: Boolean = false,
+ maxFields: Int,
+ printNodeId: Boolean,
+ printOutputColumns: Boolean,
+ indent: Int = 0): Unit = {
+ val prefix = if (printNodeId) "^ " else s"^($stageId) "
+ child.generateTreeString(
+ depth,
+ lastChildren,
+ append,
+ verbose,
+ prefix,
+ addSuffix = false,
+ maxFields,
+ printNodeId = printNodeId,
+ printOutputColumns = printOutputColumns,
+ indent = indent)
+
+ if (verbose && wholeStageTransformerContextDefined) {
+ append(prefix + "Substrait plan:\n")
+ append(substraitPlanJson)
+ append("\n")
+ }
+ }
+}
+
+trait InputAdapterGenerateTreeStringShim extends UnaryExecNode {
+
+ override def generateTreeString(
+ depth: Int,
+ lastChildren: java.util.ArrayList[Boolean],
+ append: String => Unit,
+ verbose: Boolean,
+ prefix: String = "",
+ addSuffix: Boolean = false,
+ maxFields: Int,
+ printNodeId: Boolean,
+ printOutputColumns: Boolean,
+ indent: Int = 0): Unit = {
+ child.generateTreeString(
+ depth,
+ lastChildren,
+ append,
+ verbose,
+ prefix = "",
+ addSuffix = false,
+ maxFields,
+ printNodeId,
+ printOutputColumns,
+ indent)
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/gluten/execution/PartitionedFileUtilShim.scala b/shims/spark42/src/main/scala/org/apache/gluten/execution/PartitionedFileUtilShim.scala
new file mode 100644
index 0000000000..2d3287afe4
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/gluten/execution/PartitionedFileUtilShim.scala
@@ -0,0 +1,174 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.execution
+
+import org.apache.spark.paths.SparkPath
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.PartitionedFileUtil
+import org.apache.spark.sql.execution.datasources.{FileStatusWithMetadata, PartitionedFile}
+
+import org.apache.hadoop.fs.Path
+
+import java.lang.reflect.Method
+
+object PartitionedFileUtilShim {
+
+ private val clz: Class[_] = PartitionedFileUtil.getClass
+ private val module = clz.getField("MODULE$").get(null)
+
+ private lazy val getPartitionedFileMethod: Method = {
+ try {
+ val m = clz.getDeclaredMethod(
+ "getPartitionedFile",
+ classOf[FileStatusWithMetadata],
+ classOf[InternalRow])
+ m.setAccessible(true)
+ m
+ } catch {
+ case _: NoSuchMethodException => null
+ }
+ }
+
+ private lazy val getPartitionedFileByPathMethod: Method = {
+ try {
+ val m = clz.getDeclaredMethod(
+ "getPartitionedFile",
+ classOf[FileStatusWithMetadata],
+ classOf[Path],
+ classOf[InternalRow])
+ m.setAccessible(true)
+ m
+ } catch {
+ case _: NoSuchMethodException => null
+ }
+ }
+
+ private lazy val getPartitionedFileByPathSizeMethod: Method = {
+ try {
+ val m = clz.getDeclaredMethod(
+ "getPartitionedFile",
+ classOf[FileStatusWithMetadata],
+ classOf[Path],
+ classOf[InternalRow],
+ classOf[Long],
+ classOf[Long])
+ m.setAccessible(true)
+ m
+ } catch {
+ case _: NoSuchMethodException => null
+ }
+ }
+
+ def getPartitionedFile(
+ file: FileStatusWithMetadata,
+ partitionValues: InternalRow): PartitionedFile = {
+ if (getPartitionedFileMethod != null) {
+ getPartitionedFileMethod
+ .invoke(module, file, partitionValues)
+ .asInstanceOf[PartitionedFile]
+ } else if (getPartitionedFileByPathMethod != null) {
+ getPartitionedFileByPathMethod
+ .invoke(module, file, file.getPath, partitionValues)
+ .asInstanceOf[PartitionedFile]
+ } else if (getPartitionedFileByPathSizeMethod != null) {
+ getPartitionedFileByPathSizeMethod
+ .invoke(module, file, file.getPath, partitionValues, 0, file.getLen)
+ .asInstanceOf[PartitionedFile]
+ } else {
+ val params = clz.getDeclaredMethods
+ .find(_.getName == "getPartitionedFile")
+ .map(_.getGenericParameterTypes.mkString(", "))
+ throw new RuntimeException(
+ s"getPartitionedFile with $params is not correctly shimmed " +
+ "in PartitionedFileUtilShim")
+ }
+ }
+
+ private lazy val splitFilesMethod: Method = {
+ try {
+ val m = clz.getDeclaredMethod(
+ "splitFiles",
+ classOf[SparkSession],
+ classOf[FileStatusWithMetadata],
+ classOf[Boolean],
+ classOf[Long],
+ classOf[InternalRow])
+ m.setAccessible(true)
+ m
+ } catch {
+ case _: NoSuchMethodException => null
+ }
+ }
+
+ private lazy val splitFilesByPathMethod: Method = {
+ try {
+ val m = clz.getDeclaredMethod(
+ "splitFiles",
+ classOf[FileStatusWithMetadata],
+ classOf[Path],
+ classOf[Boolean],
+ classOf[Long],
+ classOf[InternalRow])
+ m.setAccessible(true)
+ m
+ } catch {
+ case _: NoSuchMethodException => null
+ }
+ }
+
+ def splitFiles(
+ sparkSession: SparkSession,
+ file: FileStatusWithMetadata,
+ isSplitable: Boolean,
+ maxSplitBytes: Long,
+ partitionValues: InternalRow): Seq[PartitionedFile] = {
+ if (splitFilesMethod != null) {
+ splitFilesMethod
+ .invoke(
+ module,
+ sparkSession,
+ file,
+ java.lang.Boolean.valueOf(isSplitable),
+ java.lang.Long.valueOf(maxSplitBytes),
+ partitionValues)
+ .asInstanceOf[Seq[PartitionedFile]]
+ } else if (splitFilesByPathMethod != null) {
+ splitFilesByPathMethod
+ .invoke(
+ module,
+ file,
+ file.getPath,
+ java.lang.Boolean.valueOf(isSplitable),
+ java.lang.Long.valueOf(maxSplitBytes),
+ partitionValues)
+ .asInstanceOf[Seq[PartitionedFile]]
+ } else {
+ val params = clz.getDeclaredMethods
+ .find(_.getName == "splitFiles")
+ .map(_.getGenericParameterTypes.mkString(", "))
+ throw new RuntimeException(
+ s"splitFiles with $params is not correctly shimmed " +
+ "in PartitionedFileUtilShim")
+ }
+ }
+
+ // Helper method to create PartitionedFile from path and length.
+ def makePartitionedFileFromPath(path: String, length: Long): PartitionedFile = {
+ PartitionedFile(null, SparkPath.fromPathString(path), 0, length, Array.empty)
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/gluten/sql/shims/spark42/Spark42Shims.scala b/shims/spark42/src/main/scala/org/apache/gluten/sql/shims/spark42/Spark42Shims.scala
new file mode 100644
index 0000000000..8f664fca6b
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/gluten/sql/shims/spark42/Spark42Shims.scala
@@ -0,0 +1,490 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.sql.shims.spark42
+
+import org.apache.gluten.execution.PartitionedFileUtilShim
+import org.apache.gluten.expression.{ExpressionNames, Sig}
+import org.apache.gluten.sql.shims.SparkShims
+
+import org.apache.spark._
+import org.apache.spark.sql.{AnalysisException, SparkSession}
+import org.apache.spark.sql.catalyst.{ExtendedAnalysisException, InternalRow}
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
+import org.apache.spark.sql.catalyst.plans.{JoinType, LeftSingle}
+import org.apache.spark.sql.catalyst.plans.QueryPlan
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning}
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.catalyst.util.{CollationFactory, InternalRowComparableWrapper, MapData}
+import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec
+import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, Scan}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.execution._
+import org.apache.spark.sql.execution.datasources._
+import org.apache.spark.sql.execution.datasources.parquet.{ParquetFileFormat, ParquetFilters}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanExecBase}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike}
+import org.apache.spark.sql.execution.window.{Final, Partial, _}
+import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf}
+import org.apache.spark.sql.types._
+import org.apache.spark.storage.{GlutenShuffleBlockFetcherIterator, GlutenShuffleBlockFetcherIteratorBase, ShuffleBlockFetcherIteratorParams}
+
+import org.apache.hadoop.fs.{FileStatus, Path}
+import org.apache.parquet.hadoop.metadata.{CompressionCodecName, ParquetMetadata}
+import org.apache.parquet.hadoop.metadata.FileMetaData.EncryptionType
+import org.apache.parquet.schema.{GroupType, LogicalTypeAnnotation, MessageType}
+
+import java.util.{Map => JMap}
+
+import scala.jdk.CollectionConverters._
+
+class Spark42Shims extends SparkShims {
+
+ override def getSampleSeed(plan: SampleExec): Long = plan.resolvedSeed
+
+ override def isKeyGroupedPartitioning(partitioning: Partitioning): Boolean =
+ partitioning.isInstanceOf[KeyedPartitioning]
+
+ override def getLocalTableScanStream(plan: LocalTableScanExec): Option[SparkDataStream] =
+ plan.stream
+
+ override def scalarExpressionMappings: Seq[Sig] = {
+ Seq(
+ Sig[Empty2Null](ExpressionNames.EMPTY2NULL),
+ Sig[Mask](ExpressionNames.MASK),
+ Sig[ArrayInsert](ExpressionNames.ARRAY_INSERT),
+ Sig[CheckOverflowInTableInsert](ExpressionNames.CHECK_OVERFLOW_IN_TABLE_INSERT),
+ Sig[ArrayAppend](ExpressionNames.ARRAY_APPEND),
+ Sig[UrlEncode](ExpressionNames.URL_ENCODE),
+ Sig[KnownNotContainsNull](ExpressionNames.KNOWN_NOT_CONTAINS_NULL),
+ Sig[UrlDecode](ExpressionNames.URL_DECODE),
+ Sig[ToPrettyString](ExpressionNames.TO_PRETTY_STRING),
+ Sig[RandStr](ExpressionNames.RANDSTR),
+ Sig[RegExpInStr](ExpressionNames.REGEXP_INSTR),
+ Sig[DayName](ExpressionNames.DAY_NAME),
+ Sig[MonthName](ExpressionNames.MONTH_NAME)
+ )
+ }
+
+ override def aggregateExpressionMappings: Seq[Sig] = {
+ Seq(
+ Sig[RegrSlope](ExpressionNames.REGR_SLOPE),
+ Sig[RegrIntercept](ExpressionNames.REGR_INTERCEPT),
+ Sig[RegrSXY](ExpressionNames.REGR_SXY),
+ Sig[RegrReplacement](ExpressionNames.REGR_REPLACEMENT),
+ Sig[BitmapConstructAgg](ExpressionNames.BITMAP_CONSTRUCT_AGG)
+ )
+ }
+
+ override def runtimeReplaceableExpressionMappings: Seq[Sig] = {
+ Seq(
+ Sig[ArrayCompact](ExpressionNames.ARRAY_COMPACT),
+ Sig[ArrayPrepend](ExpressionNames.ARRAY_PREPEND),
+ Sig[EqualNull](ExpressionNames.EQUAL_NULL),
+ Sig[Get](ExpressionNames.GET),
+ Sig[Luhncheck](ExpressionNames.LUHN_CHECK)
+ )
+ }
+
+ override def isNullIntolerant(expr: Expression): Boolean = expr.nullIntolerant
+
+ override def filesGroupedToBuckets(
+ selectedPartitions: Array[PartitionDirectory]): Map[Int, Array[PartitionedFile]] = {
+ selectedPartitions
+ .flatMap(p => p.files.map(f => PartitionedFileUtilShim.getPartitionedFile(f, p.values)))
+ .groupBy {
+ f =>
+ BucketingUtils
+ .getBucketId(f.toPath.getName)
+ .getOrElse(throw invalidBucketFile(f.urlEncodedPath))
+ }
+ }
+
+ // https://issues.apache.org/jira/browse/SPARK-40400
+ private def invalidBucketFile(path: String): Throwable = {
+ new SparkException(
+ errorClass = "INVALID_BUCKET_FILE",
+ messageParameters = Map("path" -> path),
+ cause = null)
+ }
+
+ override def isWindowGroupLimitExec(plan: SparkPlan): Boolean = plan match {
+ case _: WindowGroupLimitExec => true
+ case _ => false
+ }
+
+ override def isEmptyRelationExec(plan: SparkPlan): Boolean = plan match {
+ case _: EmptyRelationExec => true
+ case _ => false
+ }
+
+ override def getWindowGroupLimitExecShim(plan: SparkPlan): WindowGroupLimitExecShim = {
+ val windowGroupLimitPlan = plan.asInstanceOf[WindowGroupLimitExec]
+ val mode = windowGroupLimitPlan.mode match {
+ case Partial => GlutenPartial
+ case Final => GlutenFinal
+ }
+ WindowGroupLimitExecShim(
+ windowGroupLimitPlan.partitionSpec,
+ windowGroupLimitPlan.orderSpec,
+ windowGroupLimitPlan.rankLikeFunction,
+ windowGroupLimitPlan.limit,
+ mode,
+ windowGroupLimitPlan.child
+ )
+ }
+
+ override def getWindowGroupLimitExec(
+ windowGroupLimitExecShim: WindowGroupLimitExecShim): SparkPlan = {
+ val mode = windowGroupLimitExecShim.mode match {
+ case GlutenPartial => Partial
+ case GlutenFinal => Final
+ }
+ WindowGroupLimitExec(
+ windowGroupLimitExecShim.partitionSpec,
+ windowGroupLimitExecShim.orderSpec,
+ windowGroupLimitExecShim.rankLikeFunction,
+ windowGroupLimitExecShim.limit,
+ mode,
+ windowGroupLimitExecShim.child
+ )
+ }
+
+ override def setJobDescriptionOrTagForBroadcastExchange(
+ sc: SparkContext,
+ broadcastExchange: BroadcastExchangeLike): Unit = {
+ // Setup a job tag here so later it may get cancelled by tag if necessary.
+ sc.addJobTag(broadcastExchange.jobTag)
+ sc.setInterruptOnCancel(true)
+ }
+
+ override def cancelJobGroupForBroadcastExchange(
+ sc: SparkContext,
+ broadcastExchange: BroadcastExchangeLike): Unit = {
+ sc.cancelJobsWithTag(broadcastExchange.jobTag)
+ }
+
+ override def getShuffleAdvisoryPartitionSize(shuffle: ShuffleExchangeLike): Option[Long] =
+ shuffle.advisoryPartitionSize
+
+ def getFileStatus(partition: PartitionDirectory): Seq[(FileStatus, Map[String, Any])] =
+ partition.files.map(f => (f.fileStatus, f.metadata))
+
+ def isFileSplittable(
+ relation: HadoopFsRelation,
+ filePath: Path,
+ sparkSchema: StructType): Boolean = {
+ relation.fileFormat
+ .isSplitable(relation.sparkSession, relation.options, filePath)
+ }
+
+ def isRowIndexMetadataColumn(name: String): Boolean =
+ name == ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME ||
+ name.equalsIgnoreCase("__delta_internal_is_row_deleted")
+
+ def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = {
+ sparkSchema.fields.zipWithIndex.find {
+ case (field: StructField, _: Int) =>
+ field.name == ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME
+ } match {
+ case Some((field: StructField, idx: Int)) =>
+ if (field.dataType != LongType && field.dataType != IntegerType) {
+ throw new RuntimeException(
+ s"${ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME} " +
+ "must be of LongType or IntegerType")
+ }
+ idx
+ case _ => -1
+ }
+ }
+
+ def splitFiles(
+ sparkSession: SparkSession,
+ file: FileStatus,
+ filePath: Path,
+ isSplitable: Boolean,
+ maxSplitBytes: Long,
+ partitionValues: InternalRow,
+ metadata: Map[String, Any] = Map.empty): Seq[PartitionedFile] = {
+ PartitionedFileUtilShim.splitFiles(
+ sparkSession,
+ FileStatusWithMetadata(file, metadata),
+ isSplitable,
+ maxSplitBytes,
+ partitionValues)
+ }
+
+ def structFromAttributes(attrs: Seq[Attribute]): StructType = {
+ DataTypeUtils.fromAttributes(attrs)
+ }
+
+ def attributesFromStruct(structType: StructType): Seq[Attribute] = {
+ DataTypeUtils.toAttributes(structType)
+ }
+
+ def getAnalysisExceptionPlan(ae: AnalysisException): Option[LogicalPlan] = {
+ ae match {
+ case eae: ExtendedAnalysisException =>
+ eae.plan
+ case _ =>
+ None
+ }
+ }
+ override def getCommonPartitionValues(
+ batchScan: BatchScanExec): Option[Seq[(InternalRow, Int)]] = {
+ // Spark 4.2 removed `StoragePartitionJoinParams` (and `BatchScanExec.spjParams`), so the
+ // "common partition values" that a partially-clustered storage-partitioned join used to expose
+ // on the scan node are no longer available here -- Spark 4.2 computes and applies them in
+ // `EnsureRequirements`/`GroupPartitionsExec` instead. There is no equivalent accessor on the
+ // 4.2 `BatchScanExec`, so we conservatively return `None`, which simply disables the
+ // partially-clustered-distribution refinement in Gluten's own scan planner (DEGRADED: see the
+ // note in `orderPartitions`). This does not affect the base (fully-clustered) SPJ path.
+ None
+ }
+
+ // please ref BatchScanExec::inputRDD
+ override def orderPartitions(
+ batchScan: DataSourceV2ScanExecBase,
+ scan: Scan,
+ keyGroupedPartitioning: Option[Seq[Expression]],
+ filteredPartitions: Seq[Seq[InputPartition]],
+ outputPartitioning: Partitioning,
+ commonPartitionValues: Option[Seq[(InternalRow, Int)]],
+ applyPartialClustering: Boolean,
+ replicatePartitions: Boolean,
+ joinKeyPositions: Option[Seq[Int]] = None): Seq[Seq[InputPartition]] = {
+ scan match {
+ case _ if keyGroupedPartitioning.isDefined =>
+ outputPartitioning match {
+ case p: KeyedPartitioning =>
+ val partExpressions = keyGroupedPartitioning.get
+
+ // DEGRADED (Spark 4.2 port): Spark 4.2 removed `KeyGroupedPartitioning` and
+ // `StoragePartitionJoinParams`, and moved the storage-partitioned-join refinements that
+ // used to run here into `EnsureRequirements`/`GroupPartitionsExec`:
+ // - subset-of-join-keys projection (`joinKeyPositions`),
+ // - compatible partition-expression reduction (`reducers`),
+ // - partially-clustered replication (`commonPartitionValues` /
+ // `applyPartialClustering` / `replicatePartitions`).
+ // Gluten never populates `joinKeyPositions`/`reducers`, and `getCommonPartitionValues`
+ // returns `None` on 4.2, so `commonPartitionValues` is always empty here. These
+ // parameters therefore have no 4.2 equivalent that can be reproduced on the scan node
+ // and are intentionally NOT applied; only the base key-grouped ordering is reproduced.
+ // The base (fully-clustered) SPJ path is unaffected.
+ val groupedPartitions = filteredPartitions.map {
+ splits =>
+ assert(splits.nonEmpty && splits.head.isInstanceOf[HasPartitionKey])
+ (splits.head.asInstanceOf[HasPartitionKey].partitionKey(), splits)
+ }
+
+ val partitionMapping = groupedPartitions.map {
+ case (partValue, splits) =>
+ InternalRowComparableWrapper(partValue, partExpressions) -> splits
+ }.toMap
+
+ // Use the unique, sorted partition keys as the canonical partition order (Spark 4.2's
+ // `KeyedPartitioning.toGrouped` returns distinct keys sorted ascending), filling absent
+ // keys with empty split groups so both sides of a storage-partitioned join stay
+ // aligned. This mirrors the old `KeyGroupedPartitioning.uniquePartitionValues` path.
+ p.toGrouped.partitionKeys.map {
+ keyWrapper =>
+ // Use empty partition for those partition values that are not present
+ partitionMapping.getOrElse(keyWrapper, Seq.empty)
+ }
+
+ case _ => filteredPartitions
+ }
+ case _ =>
+ filteredPartitions
+ }
+ }
+
+ override def createParquetFilters(
+ conf: SQLConf,
+ schema: MessageType,
+ caseSensitive: Option[Boolean] = None): ParquetFilters = {
+ new ParquetFilters(
+ schema,
+ conf.parquetFilterPushDownDate,
+ conf.parquetFilterPushDownTimestamp,
+ conf.parquetFilterPushDownDecimal,
+ conf.parquetFilterPushDownStringPredicate,
+ conf.parquetFilterPushDownInFilterThreshold,
+ caseSensitive.getOrElse(conf.caseSensitiveAnalysis),
+ RebaseSpec(LegacyBehaviorPolicy.CORRECTED)
+ )
+ }
+
+ override def withOperatorIdMap[T](idMap: java.util.Map[QueryPlan[_], Int])(body: => T): T = {
+ val prevIdMap = QueryPlan.localIdMap.get()
+ try {
+ QueryPlan.localIdMap.set(idMap)
+ body
+ } finally {
+ QueryPlan.localIdMap.set(prevIdMap)
+ }
+ }
+
+ override def getOperatorId(plan: QueryPlan[_]): Option[Int] = {
+ Option(QueryPlan.localIdMap.get().get(plan))
+ }
+
+ override def setOperatorId(plan: QueryPlan[_], opId: Int): Unit = {
+ val map = QueryPlan.localIdMap.get()
+ assert(!map.containsKey(plan))
+ map.put(plan, opId)
+ }
+
+ override def unsetOperatorId(plan: QueryPlan[_]): Unit = {
+ QueryPlan.localIdMap.get().remove(plan)
+ }
+
+ override def isParquetFileEncrypted(footer: ParquetMetadata): Boolean = {
+ footer.getFileMetaData.getEncryptionType match {
+ // UNENCRYPTED file has a plaintext footer and no file encryption,
+ // We can leverage file metadata for this check and return unencrypted.
+ case EncryptionType.UNENCRYPTED =>
+ false
+ // PLAINTEXT_FOOTER has a plaintext footer however the file is encrypted.
+ // In such cases, read the footer and use the metadata for encryption check.
+ case EncryptionType.PLAINTEXT_FOOTER =>
+ true
+ case _ =>
+ false
+ }
+ }
+
+ override def shouldFallbackForParquetVariantAnnotation(footer: ParquetMetadata): Boolean = {
+ if (SQLConf.get.getConf(SQLConf.PARQUET_IGNORE_VARIANT_ANNOTATION)) {
+ false
+ } else {
+ containsVariantAnnotation(footer.getFileMetaData.getSchema)
+ }
+ }
+
+ private def containsVariantAnnotation(groupType: GroupType): Boolean = {
+ groupType.getFields.asScala.exists {
+ field =>
+ Option(field.getLogicalTypeAnnotation)
+ .exists(_.isInstanceOf[LogicalTypeAnnotation.VariantLogicalTypeAnnotation]) ||
+ (!field.isPrimitive && containsVariantAnnotation(field.asGroupType()))
+ }
+ }
+
+ override def getOtherConstantMetadataColumnValues(file: PartitionedFile): JMap[String, Object] =
+ file.otherConstantMetadataColumnValues.asJava.asInstanceOf[JMap[String, Object]]
+
+ override def extractExpressionTimestampAddUnit(exp: Expression): Option[Seq[String]] = {
+ exp match {
+ // Velox does not support quantity larger than Int.MaxValue.
+ case TimestampAdd(_, LongLiteral(quantity), _, _) if quantity > Integer.MAX_VALUE =>
+ Option.empty
+ case timestampAdd: TimestampAdd =>
+ Option.apply(Seq(timestampAdd.unit, timestampAdd.timeZoneId.getOrElse("")))
+ case _ => Option.empty
+ }
+ }
+
+ override def decimalAllowPrecisionLoss(expr: BinaryArithmetic): Boolean = expr match {
+ case a: Add => a.evalContext.allowDecimalPrecisionLoss
+ case s: Subtract => s.evalContext.allowDecimalPrecisionLoss
+ case m: Multiply => m.evalContext.allowDecimalPrecisionLoss
+ case d: Divide => d.evalContext.allowDecimalPrecisionLoss
+ // Remainder and Pmod do not carry evalContext in Spark 4.1. They also throw
+ // GlutenNotSupportException in DecimalArithmeticUtil.getResultType, so they never
+ // reach Velox execution; SQLConf.get is a safe fallback for the name-lookup path.
+ case _ => SQLConf.get.decimalOperationsAllowPrecisionLoss
+ }
+
+ override def getErrorMessage(raiseError: RaiseError): Option[Expression] = {
+ raiseError.errorParms match {
+ case CreateMap(children, _)
+ if children.size == 2 && children.head.isInstanceOf[Literal]
+ && children.head.asInstanceOf[Literal].value.toString == "errorMessage" =>
+ Some(children(1))
+ case lit: Literal if lit.value.isInstanceOf[MapData] =>
+ // Constant-folded CreateMap: look up "errorMessage" in the MapData
+ val mapData = lit.value.asInstanceOf[MapData]
+ (0 until mapData.numElements())
+ .find(i => mapData.keyArray().getUTF8String(i).toString == "errorMessage")
+ .map(i => Literal(mapData.valueArray().getUTF8String(i), StringType))
+ case _ => None
+ }
+ }
+
+ override def throwExceptionInWrite(
+ t: Throwable,
+ writePath: String,
+ descriptionPath: String): Unit = {
+ throw t
+ }
+
+ override def enrichWriteException(cause: Throwable, path: String): Nothing = {
+ GlutenFileFormatWriter.wrapWriteError(cause, path)
+ }
+ override def getFileSourceScanStream(scan: FileSourceScanExec): Option[SparkDataStream] = {
+ scan.stream
+ }
+
+ override def unsupportedCodec: Seq[CompressionCodecName] = {
+ Seq(CompressionCodecName.LZO, CompressionCodecName.BROTLI, CompressionCodecName.LZ4_RAW)
+ }
+
+ /**
+ * Shim layer for QueryExecution to maintain compatibility across different Spark versions.
+ *
+ * @since Spark
+ * 4.1
+ */
+ override def createSparkPlan(
+ sparkSession: SparkSession,
+ planner: SparkPlanner,
+ plan: LogicalPlan): SparkPlan =
+ QueryExecution.createSparkPlan(planner, plan)
+
+ override def isLeftSingleJoinType(joinType: JoinType): Boolean = {
+ joinType == LeftSingle
+ }
+
+ override def isBinaryCollationString(dt: StringType): Boolean =
+ dt.collationId == CollationFactory.UTF8_BINARY_COLLATION_ID
+
+ override def getShuffleBlockFetcherIterator(params: ShuffleBlockFetcherIteratorParams)
+ : GlutenShuffleBlockFetcherIteratorBase = {
+ new GlutenShuffleBlockFetcherIterator(
+ params.context,
+ params.shuffleClient,
+ params.blockManager,
+ params.mapOutputTracker,
+ params.blocksByAddress,
+ params.streamWrapper,
+ params.maxBytesInFlight,
+ params.maxReqsInFlight,
+ params.maxBlocksInFlightPerAddress,
+ params.maxReqSizeShuffleToMem,
+ params.maxAttemptsOnNettyOOM,
+ params.detectCorrupt,
+ params.detectCorruptUseExtraMemory,
+ params.checksumEnabled,
+ params.checksumAlgorithm,
+ params.shuffleMetrics,
+ params.doBatchFetch,
+ params.clock
+ )
+ }
+}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/gluten/sql/shims/spark42/SparkShimProvider.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/gluten/sql/shims/spark42/SparkShimProvider.scala
index f7326e98d6..b831fd4ff7 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/gluten/sql/shims/spark42/SparkShimProvider.scala
@@ -14,20 +14,12 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.gluten.sql.shims.spark42
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.gluten.sql.shims.SparkShims
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+class SparkShimProvider extends org.apache.gluten.sql.shims.SparkShimProvider {
+ def createShim: SparkShims = {
+ new Spark42Shims()
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/gluten/utils/InternalRowUtil.scala
similarity index 56%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/gluten/utils/InternalRowUtil.scala
index f7326e98d6..0f7c13448f 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/gluten/utils/InternalRowUtil.scala
@@ -14,20 +14,20 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.gluten.utils
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.types.StructType
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+object InternalRowUtil {
+ def toString(struct: StructType, rows: Iterator[InternalRow]): String = {
+ val encoder = ExpressionEncoder(struct).resolveAndBind()
+ val deserializer = encoder.createDeserializer()
+ rows.map(deserializer).mkString(System.lineSeparator())
+ }
+
+ def toString(struct: StructType, rows: Iterator[InternalRow], start: Int, length: Int): String = {
+ toString(struct, rows.slice(start, start + length))
}
}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/ShuffleUtils.scala b/shims/spark42/src/main/scala/org/apache/spark/ShuffleUtils.scala
new file mode 100644
index 0000000000..c2a6cd5cff
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/ShuffleUtils.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark
+
+import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleHandle}
+import org.apache.spark.storage.{BlockId, BlockManagerId}
+
+object ShuffleUtils {
+ def getReaderParam[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int): Tuple2[Iterator[(BlockManagerId, Seq[(BlockId, Long, Int)])], Boolean] = {
+ val baseShuffleHandle = handle.asInstanceOf[BaseShuffleHandle[K, _, C]]
+ if (baseShuffleHandle.dependency.isShuffleMergeFinalizedMarked) {
+ val res = SparkEnv.get.mapOutputTracker.getPushBasedShuffleMapSizesByExecutorId(
+ handle.shuffleId,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition)
+ (res.iter.map(b => (b._1, b._2.toSeq)), res.enableBatchFetch)
+ } else {
+ val address = SparkEnv.get.mapOutputTracker.getMapSizesByExecutorId(
+ handle.shuffleId,
+ startMapIndex,
+ endMapIndex,
+ startPartition,
+ endPartition)
+ (address.map(b => (b._1, b._2.toSeq)), true)
+ }
+ }
+}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/SparkContextUtils.scala
similarity index 57%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/SparkContextUtils.scala
index f7326e98d6..41e09c6466 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/SparkContextUtils.scala
@@ -14,20 +14,23 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.rdd.{RDD, SQLPartitioningAwareUnionRDD}
+import org.apache.spark.sql.vectorized.ColumnarBatch
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+import scala.reflect.ClassTag
+
+object SparkContextUtils {
+ def broadcastInternal[T: ClassTag](sc: SparkContext, value: T): Broadcast[T] = {
+ sc.broadcastInternal(value, serializedOnly = true)
+ }
+
+ def createPartitioningAwareUnionRDD(
+ sc: SparkContext,
+ rdds: Seq[RDD[ColumnarBatch]],
+ numPartitions: Int): RDD[ColumnarBatch] = {
+ new SQLPartitioningAwareUnionRDD(sc, rdds, numPartitions)
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/shuffle/GlutenMapStatusUtil.scala
similarity index 57%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/shuffle/GlutenMapStatusUtil.scala
index f7326e98d6..bf8108c4b0 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/shuffle/GlutenMapStatusUtil.scala
@@ -14,20 +14,24 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.shuffle
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.storage.BlockManagerId
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+object GlutenMapStatusUtil {
+ def isRowBasedChecksumEnabled: Boolean = {
+ val sqlConf = SQLConf.get
+ sqlConf.shuffleOrderIndependentChecksumEnabled ||
+ sqlConf.shuffleChecksumMismatchFullRetryEnabled
+ }
+
+ def createMapStatus(
+ loc: BlockManagerId,
+ uncompressedSizes: Array[Long],
+ mapTaskId: Long,
+ checksumValue: Long): MapStatus = {
+ MapStatus(loc, uncompressedSizes, mapTaskId, checksumValue)
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/shuffle/SparkSortShuffleWriterUtil.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/shuffle/SparkSortShuffleWriterUtil.scala
index f7326e98d6..95b15f04e7 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/shuffle/SparkSortShuffleWriterUtil.scala
@@ -14,20 +14,19 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.shuffle
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.TaskContext
+import org.apache.spark.shuffle.api.ShuffleExecutorComponents
+import org.apache.spark.shuffle.sort.SortShuffleWriter
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+object SparkSortShuffleWriterUtil {
+ def create[K, V, C](
+ handle: BaseShuffleHandle[K, V, C],
+ mapId: Long,
+ context: TaskContext,
+ writeMetrics: ShuffleWriteMetricsReporter,
+ shuffleExecutorComponents: ShuffleExecutorComponents): ShuffleWriter[K, V] = {
+ new SortShuffleWriter(handle, mapId, context, writeMetrics, shuffleExecutorComponents)
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/expressions/PromotePrecision.scala
similarity index 50%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/expressions/PromotePrecision.scala
index f7326e98d6..b18a79b864 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/expressions/PromotePrecision.scala
@@ -14,20 +14,21 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.catalyst.expressions
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
+import org.apache.spark.sql.types._
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
- }
+case class PromotePrecision(child: Expression) extends UnaryExpression {
+ override def dataType: DataType = child.dataType
+ override def eval(input: InternalRow): Any = child.eval(input)
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+ child.genCode(ctx)
+ override def prettyName: String = "promote_precision"
+ override def sql: String = child.sql
+ override lazy val canonicalized: Expression = child.canonicalized
+
+ override protected def withNewChildInternal(newChild: Expression): Expression =
+ copy(child = newChild)
}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/InvokeExtractors.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/InvokeExtractors.scala
new file mode 100644
index 0000000000..8abebd0d2a
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/InvokeExtractors.scala
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.catalyst.expressions.objects
+
+import org.apache.spark.sql.catalyst.expressions.{Expression, Literal}
+import org.apache.spark.sql.catalyst.expressions.json.StructsToJsonEvaluator
+
+/**
+ * Extractors for Invoke expressions to ensure compatibility across different Spark versions.
+ *
+ * Since Spark 4.0, StructsToJson has been replaced with Invoke expressions using
+ * StructsToJsonEvaluator. This extractor provides a unified interface to extract evaluator options,
+ * child expression, and timeZoneId from the Invoke pattern.
+ */
+object StructsToJsonInvoke {
+ def unapply(expr: Expression): Option[(Map[String, String], Expression, Option[String])] = {
+ expr match {
+ case Invoke(
+ Literal(evaluator: StructsToJsonEvaluator, _),
+ "evaluate",
+ _,
+ Seq(child),
+ _,
+ _,
+ _,
+ _) =>
+ Some((evaluator.options, child, evaluator.timeZoneId))
+ case _ =>
+ None
+ }
+ }
+}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CollapseProjectShim.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CollapseProjectShim.scala
index f7326e98d6..1df1456f40 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CollapseProjectShim.scala
@@ -14,20 +14,21 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.catalyst.optimizer
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression}
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+object CollapseProjectShim {
+ def canCollapseExpressions(
+ consumers: Seq[Expression],
+ producers: Seq[NamedExpression],
+ alwaysInline: Boolean): Boolean = {
+ CollapseProject.canCollapseExpressions(consumers, producers, alwaysInline)
+ }
+
+ def buildCleanedProjectList(
+ upper: Seq[NamedExpression],
+ lower: Seq[NamedExpression]): Seq[NamedExpression] = {
+ CollapseProject.buildCleanedProjectList(upper, lower)
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicColumn.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicColumn.scala
index f7326e98d6..5b1f5a2a80 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicColumn.scala
@@ -14,20 +14,18 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.classic
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.expressions.Expression
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+/**
+ * Ensures compatibility with Spark 4.0. Since Spark 4.0, the Column class is private to the package
+ * org.apache.spark. This class provides a way to construct a Column from an Expression in code that
+ * is outside the org.apache.spark package by delegating to ExpressionUtils.column().
+ */
+object ClassicColumn {
+ def apply(e: Expression): Column = {
+ ExpressionUtils.column(e)
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicDataset.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicDataset.scala
index f7326e98d6..1a48a8e110 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicDataset.scala
@@ -14,20 +14,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.classic
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.classic
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+/** Since Spark 4.0, the method ofRows cannot be invoked directly from sql.Dataset. */
+object ClassicDataset {
+ def ofRows(sparkSession: classic.SparkSession, logicalPlan: LogicalPlan): DataFrame = {
+ // Redirect to the classic.Dataset companion method.
+ classic.Dataset.ofRows(sparkSession, logicalPlan)
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicTypes.scala
similarity index 57%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicTypes.scala
index f7326e98d6..ea3fb8893f 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/ClassicTypes.scala
@@ -14,20 +14,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.classic
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
- }
+/**
+ * Prior to Spark 4.0, `ClassicSparkSession` refers to `sql.SparkSession`. Since Spark 4.0, it
+ * refers to `sql.classic.SparkSession`.
+ */
+object ClassicTypes {
+
+ type ClassicSparkSession = sql.classic.SparkSession
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/conversions.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/classic/conversions.scala
index f7326e98d6..b168838f05 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/classic/conversions.scala
@@ -14,20 +14,20 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.classic
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
+/**
+ * Enables access to the methods in the companion object classic.SparkSession via sql.SparkSession.
+ * Since Spark 4.0, these methods have been moved from sql.SparkSession to sql.classic.SparkSession.
+ */
+object ExtendedClassicConversions {
+
+ implicit class RichSqlSparkSession(sqlSparkSession: sql.SparkSession.type) {
+ def cleanupAnyExistingSession(): Unit = {
+ // Redirect to the classic.SparkSession companion method.
+ sql.classic.SparkSession.cleanupAnyExistingSession()
}
}
}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/execution/AbstractFileSourceScanExec.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/AbstractFileSourceScanExec.scala
new file mode 100644
index 0000000000..ac17c259b8
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/AbstractFileSourceScanExec.scala
@@ -0,0 +1,307 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution
+
+import org.apache.gluten.execution.PartitionedFileUtilShim
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow, TableIdentifier}
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.execution.datasources._
+import org.apache.spark.sql.execution.datasources.parquet.{ParquetFileFormat => ParquetSource}
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.collection.BitSet
+
+import org.apache.hadoop.fs.Path
+
+import java.util.concurrent.TimeUnit._
+
+/**
+ * Physical plan node for scanning data from HadoopFsRelations.
+ *
+ * @param relation
+ * The file-based relation to scan.
+ * @param output
+ * Output attributes of the scan, including data attributes and partition attributes.
+ * @param requiredSchema
+ * Required schema of the underlying relation, excluding partition columns.
+ * @param partitionFilters
+ * Predicates to use for partition pruning.
+ * @param optionalBucketSet
+ * Bucket ids for bucket pruning.
+ * @param optionalNumCoalescedBuckets
+ * Number of coalesced buckets.
+ * @param dataFilters
+ * Filters on non-partition columns.
+ * @param tableIdentifier
+ * Identifier for the table in the metastore.
+ * @param disableBucketedScan
+ * Disable bucketed scan based on physical query plan, see rule [[DisableUnnecessaryBucketedScan]]
+ * for details.
+ */
+abstract class AbstractFileSourceScanExec(
+ @transient override val relation: HadoopFsRelation,
+ override val output: Seq[Attribute],
+ override val requiredSchema: StructType,
+ override val partitionFilters: Seq[Expression],
+ override val optionalBucketSet: Option[BitSet],
+ override val optionalNumCoalescedBuckets: Option[Int],
+ override val dataFilters: Seq[Expression],
+ override val tableIdentifier: Option[TableIdentifier],
+ override val disableBucketedScan: Boolean = false)
+ extends FileSourceScanLike {
+
+ override def supportsColumnar: Boolean = {
+ // The value should be defined in GlutenPlan.
+ throw new UnsupportedOperationException(
+ "Unreachable code from org.apache.spark.sql.execution.AbstractFileSourceScanExec" +
+ ".supportsColumnar")
+ }
+
+ private lazy val needsUnsafeRowConversion: Boolean = {
+ if (relation.fileFormat.isInstanceOf[ParquetSource]) {
+ conf.parquetVectorizedReaderEnabled
+ } else {
+ false
+ }
+ }
+
+ lazy val inputRDD: RDD[InternalRow] = {
+ val options = relation.options +
+ (FileFormat.OPTION_RETURNING_BATCH -> supportsColumnar.toString)
+ val readFile: (PartitionedFile) => Iterator[InternalRow] =
+ relation.fileFormat.buildReaderWithPartitionValues(
+ sparkSession = relation.sparkSession,
+ dataSchema = relation.dataSchema,
+ partitionSchema = relation.partitionSchema,
+ requiredSchema = requiredSchema,
+ filters = pushedDownFilters,
+ options = options,
+ hadoopConf = relation.sparkSession.sessionState.newHadoopConfWithOptions(relation.options)
+ )
+
+ val readRDD = if (bucketedScan) {
+ createBucketedReadRDD(relation.bucketSpec.get, readFile, dynamicallySelectedPartitions)
+ } else {
+ createReadRDD(readFile, dynamicallySelectedPartitions)
+ }
+ sendDriverMetrics()
+ readRDD
+ }
+
+ override def inputRDDs(): Seq[RDD[InternalRow]] = {
+ inputRDD :: Nil
+ }
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ val numOutputRows = longMetric("numOutputRows")
+ if (needsUnsafeRowConversion) {
+ inputRDD.mapPartitionsWithIndexInternal {
+ (index, iter) =>
+ val toUnsafe = UnsafeProjection.create(schema)
+ toUnsafe.initialize(index)
+ iter.map {
+ row =>
+ numOutputRows += 1
+ toUnsafe(row)
+ }
+ }
+ } else {
+ inputRDD.mapPartitionsInternal {
+ iter =>
+ iter.map {
+ row =>
+ numOutputRows += 1
+ row
+ }
+ }
+ }
+ }
+
+ override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ val numOutputRows = longMetric("numOutputRows")
+ val scanTime = longMetric("scanTime")
+ inputRDD.asInstanceOf[RDD[ColumnarBatch]].mapPartitionsInternal {
+ batches =>
+ new Iterator[ColumnarBatch] {
+
+ override def hasNext: Boolean = {
+ // The `FileScanRDD` returns an iterator which scans the file during the `hasNext` call.
+ val startNs = System.nanoTime()
+ val res = batches.hasNext
+ scanTime += NANOSECONDS.toMillis(System.nanoTime() - startNs)
+ res
+ }
+
+ override def next(): ColumnarBatch = {
+ val batch = batches.next()
+ numOutputRows += batch.numRows()
+ batch
+ }
+ }
+ }
+ }
+
+ override val nodeNamePrefix: String = "File"
+
+ /**
+ * Create an RDD for bucketed reads. The non-bucketed variant of this function is
+ * [[createReadRDD]].
+ *
+ * The algorithm is pretty simple: each RDD partition being returned should include all the files
+ * with the same bucket id from all the given Hive partitions.
+ *
+ * @param bucketSpec
+ * the bucketing spec.
+ * @param readFile
+ * a function to read each (part of a) file.
+ * @param selectedPartitions
+ * Hive-style partition that are part of the read.
+ */
+ private def createBucketedReadRDD(
+ bucketSpec: BucketSpec,
+ readFile: (PartitionedFile) => Iterator[InternalRow],
+ selectedPartitions: ScanFileListing): RDD[InternalRow] = {
+ logInfo(s"Planning with ${bucketSpec.numBuckets} buckets")
+ val partitionArray = selectedPartitions.toPartitionArray
+ val filesGroupedToBuckets = partitionArray.groupBy {
+ f =>
+ BucketingUtils
+ .getBucketId(f.toPath.getName)
+ .getOrElse(throw QueryExecutionErrors.invalidBucketFile(f.urlEncodedPath))
+ }
+
+ val prunedFilesGroupedToBuckets = if (optionalBucketSet.isDefined) {
+ val bucketSet = optionalBucketSet.get
+ filesGroupedToBuckets.filter(f => bucketSet.get(f._1))
+ } else {
+ filesGroupedToBuckets
+ }
+
+ val filePartitions = optionalNumCoalescedBuckets
+ .map {
+ numCoalescedBuckets =>
+ logInfo(s"Coalescing to $numCoalescedBuckets buckets")
+ val coalescedBuckets = prunedFilesGroupedToBuckets.groupBy(_._1 % numCoalescedBuckets)
+ Seq.tabulate(numCoalescedBuckets) {
+ bucketId =>
+ val partitionedFiles = coalescedBuckets
+ .get(bucketId)
+ .map {
+ _.values.flatten.toArray
+ }
+ .getOrElse(Array.empty)
+ FilePartition(bucketId, partitionedFiles)
+ }
+ }
+ .getOrElse {
+ Seq.tabulate(bucketSpec.numBuckets) {
+ bucketId =>
+ FilePartition(bucketId, prunedFilesGroupedToBuckets.getOrElse(bucketId, Array.empty))
+ }
+ }
+
+ new FileScanRDD(
+ relation.sparkSession,
+ readFile,
+ filePartitions,
+ new StructType(requiredSchema.fields ++ relation.partitionSchema.fields),
+ fileConstantMetadataColumns,
+ relation.fileFormat.fileConstantMetadataExtractors,
+ new FileSourceOptions(CaseInsensitiveMap(relation.options))
+ )
+ }
+
+ /**
+ * Create an RDD for non-bucketed reads. The bucketed variant of this function is
+ * [[createBucketedReadRDD]].
+ *
+ * @param readFile
+ * a function to read each (part of a) file.
+ * @param selectedPartitions
+ * Hive-style partition that are part of the read.
+ */
+ private def createReadRDD(
+ readFile: (PartitionedFile) => Iterator[InternalRow],
+ selectedPartitions: ScanFileListing): RDD[InternalRow] = {
+ val openCostInBytes = relation.sparkSession.sessionState.conf.filesOpenCostInBytes
+ val maxSplitBytes =
+ FilePartition.maxSplitBytes(relation.sparkSession, selectedPartitions)
+ logInfo(
+ s"Planning scan with bin packing, max size: $maxSplitBytes bytes, " +
+ s"open cost is considered as scanning $openCostInBytes bytes.")
+
+ // Filter files with bucket pruning if possible
+ val bucketingEnabled = relation.sparkSession.sessionState.conf.bucketingEnabled
+ val shouldProcess: Path => Boolean = optionalBucketSet match {
+ case Some(bucketSet) if bucketingEnabled =>
+ // Do not prune the file if bucket file name is invalid
+ filePath => BucketingUtils.getBucketId(filePath.getName).forall(bucketSet.get)
+ case _ =>
+ _ => true
+ }
+
+ val splitFiles = selectedPartitions.filePartitionIterator
+ .flatMap {
+ partition =>
+ partition.files.flatMap {
+ file =>
+ if (shouldProcess(file.getPath)) {
+ val isSplitable = relation.fileFormat.isSplitable(
+ relation.sparkSession,
+ relation.options,
+ file.getPath)
+ PartitionedFileUtilShim.splitFiles(
+ sparkSession = relation.sparkSession,
+ file = file,
+ isSplitable = isSplitable,
+ maxSplitBytes = maxSplitBytes,
+ partitionValues = partition.values
+ )
+ } else {
+ Seq.empty
+ }
+ }
+ }
+ .toArray
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+
+ val partitions =
+ FilePartition.getFilePartitions(relation.sparkSession, splitFiles, maxSplitBytes)
+
+ new FileScanRDD(
+ relation.sparkSession,
+ readFile,
+ partitions,
+ new StructType(requiredSchema.fields ++ relation.partitionSchema.fields),
+ fileConstantMetadataColumns,
+ relation.fileFormat.fileConstantMetadataExtractors,
+ new FileSourceOptions(CaseInsensitiveMap(relation.options))
+ )
+ }
+
+ // Filters unused DynamicPruningExpression expressions - one which has been replaced
+ // with DynamicPruningExpression(Literal.TrueLiteral) during Physical Planning
+ protected def filterUnusedDynamicPruningExpressions(
+ predicates: Seq[Expression]): Seq[Expression] = {
+ predicates.filterNot(_ == DynamicPruningExpression(Literal.TrueLiteral))
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/execution/ExpandOutputPartitioningShim.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/ExpandOutputPartitioningShim.scala
new file mode 100644
index 0000000000..791490064a
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/ExpandOutputPartitioningShim.scala
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioningLike, Partitioning, PartitioningCollection}
+
+import scala.collection.mutable
+
+// https://issues.apache.org/jira/browse/SPARK-31869
+class ExpandOutputPartitioningShim(
+ streamedKeyExprs: Seq[Expression],
+ buildKeyExprs: Seq[Expression],
+ expandLimit: Int) {
+ // An one-to-many mapping from a streamed key to build keys.
+ private lazy val streamedKeyToBuildKeyMapping = {
+ val mapping = mutable.Map.empty[Expression, Seq[Expression]]
+ streamedKeyExprs.zip(buildKeyExprs).foreach {
+ case (streamedKey, buildKey) =>
+ val key = streamedKey.canonicalized
+ mapping.get(key) match {
+ case Some(v) => mapping.put(key, v :+ buildKey)
+ case None => mapping.put(key, Seq(buildKey))
+ }
+ }
+ mapping.toMap
+ }
+
+ def expandPartitioning(partitioning: Partitioning): Partitioning = {
+ partitioning match {
+ case h: HashPartitioningLike => expandOutputPartitioning(h)
+ case c: PartitioningCollection => expandOutputPartitioning(c)
+ case _ => partitioning
+ }
+ }
+
+ // Expands the given partitioning collection recursively.
+ private def expandOutputPartitioning(
+ partitioning: PartitioningCollection): PartitioningCollection = {
+ PartitioningCollection(partitioning.partitionings.flatMap {
+ case h: HashPartitioningLike => expandOutputPartitioning(h).partitionings
+ case c: PartitioningCollection => Seq(expandOutputPartitioning(c))
+ case other => Seq(other)
+ })
+ }
+
+ // Expands the given hash partitioning by substituting streamed keys with build keys.
+ // For example, if the expressions for the given partitioning are Seq("a", "b", "c")
+ // where the streamed keys are Seq("b", "c") and the build keys are Seq("x", "y"),
+ // the expanded partitioning will have the following expressions:
+ // Seq("a", "b", "c"), Seq("a", "b", "y"), Seq("a", "x", "c"), Seq("a", "x", "y").
+ // The expanded expressions are returned as PartitioningCollection.
+ private def expandOutputPartitioning(
+ partitioning: HashPartitioningLike): PartitioningCollection = {
+ val maxNumCombinations = expandLimit
+ var currentNumCombinations = 0
+
+ def generateExprCombinations(
+ current: Seq[Expression],
+ accumulated: Seq[Expression]): Seq[Seq[Expression]] = {
+ if (currentNumCombinations >= maxNumCombinations) {
+ Nil
+ } else if (current.isEmpty) {
+ currentNumCombinations += 1
+ Seq(accumulated)
+ } else {
+ val buildKeysOpt = streamedKeyToBuildKeyMapping.get(current.head.canonicalized)
+ generateExprCombinations(current.tail, accumulated :+ current.head) ++
+ buildKeysOpt
+ .map(_.flatMap(b => generateExprCombinations(current.tail, accumulated :+ b)))
+ .getOrElse(Nil)
+ }
+ }
+
+ PartitioningCollection(
+ generateExprCombinations(partitioning.expressions, Nil)
+ .map(exprs => partitioning.withNewChildren(exprs).asInstanceOf[HashPartitioningLike]))
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/execution/FileSourceScanExecShim.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/FileSourceScanExecShim.scala
new file mode 100644
index 0000000000..0a0bbcae9c
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/FileSourceScanExecShim.scala
@@ -0,0 +1,286 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution
+
+import org.apache.gluten.metrics.GlutenTimeMetric
+import org.apache.gluten.sql.shims.SparkShimLoader
+
+import org.apache.spark.Partition
+import org.apache.spark.internal.LogKeys.{COUNT, MAX_SPLIT_BYTES, OPEN_COST_IN_BYTES}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier}
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, BoundReference, Expression, FileSourceConstantMetadataAttribute, FileSourceGeneratedMetadataAttribute, PlanExpression, Predicate}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.execution.datasources.{BucketingUtils, FilePartition, HadoopFsRelation, PartitionDirectory}
+import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.collection.BitSet
+
+import org.apache.hadoop.fs.Path
+
+abstract class FileSourceScanExecShim(
+ @transient relation: HadoopFsRelation,
+ output: Seq[Attribute],
+ requiredSchema: StructType,
+ partitionFilters: Seq[Expression],
+ optionalBucketSet: Option[BitSet],
+ optionalNumCoalescedBuckets: Option[Int],
+ dataFilters: Seq[Expression],
+ tableIdentifier: Option[TableIdentifier],
+ disableBucketedScan: Boolean = false)
+ extends AbstractFileSourceScanExec(
+ relation,
+ output,
+ requiredSchema,
+ partitionFilters,
+ optionalBucketSet,
+ optionalNumCoalescedBuckets,
+ dataFilters,
+ tableIdentifier,
+ disableBucketedScan) {
+
+ // Note: "metrics" is made transient to avoid sending driver-side metrics to tasks.
+ @transient override lazy val metrics: Map[String, SQLMetric] = Map()
+
+ lazy val metadataColumns: Seq[AttributeReference] = output.collect {
+ case FileSourceConstantMetadataAttribute(attr) => attr
+ case FileSourceGeneratedMetadataAttribute(attr, _) => attr
+ }
+
+ protected lazy val driverMetricsAlias = driverMetrics
+
+ def dataFiltersInScan: Seq[Expression] = dataFilters.filterNot(_.references.exists {
+ attr => SparkShimLoader.getSparkShims.isRowIndexMetadataColumn(attr.name)
+ })
+
+ def hasUnsupportedColumns: Boolean = {
+ // TODO, fallback if user define same name column due to we can't right now
+ // detect which column is metadata column which is user defined column.
+ val metadataColumnsNames = metadataColumns.map(_.name)
+ output
+ .filterNot(metadataColumns.toSet)
+ .exists(v => metadataColumnsNames.contains(v.name))
+ }
+
+ def isMetadataColumn(attr: Attribute): Boolean = metadataColumns.contains(attr)
+
+ def hasFieldIds: Boolean = ParquetUtils.hasFieldIds(requiredSchema)
+
+ protected def isDynamicPruningFilter(e: Expression): Boolean =
+ e.find(_.isInstanceOf[PlanExpression[_]]).isDefined
+
+ protected def setFilesNumAndSizeMetric(partitions: ScanFileListing, static: Boolean): Unit = {
+ val filesNum = partitions.totalNumberOfFiles
+ val filesSize = partitions.totalFileSize
+ if (!static || !partitionFilters.exists(isDynamicPruningFilter)) {
+ driverMetrics("numFiles").set(filesNum)
+ driverMetrics("filesSize").set(filesSize)
+ } else {
+ driverMetrics("staticFilesNum").set(filesNum)
+ driverMetrics("staticFilesSize").set(filesSize)
+ }
+ if (relation.partitionSchema.nonEmpty) {
+ driverMetrics("numPartitions").set(partitions.partitionCount)
+ }
+ }
+
+ @transient override lazy val dynamicallySelectedPartitions: ScanFileListing = {
+ val dynamicDataFilters = dataFilters.filter(isDynamicPruningFilter)
+ val dynamicPartitionFilters =
+ partitionFilters.filter(isDynamicPruningFilter)
+ val selected = if (dynamicPartitionFilters.nonEmpty) {
+ GlutenTimeMetric.withMillisTime {
+ // call the file index for the files matching all filters except dynamic partition filters
+ val boundedFilters = dynamicPartitionFilters.map {
+ dynamicPartitionFilter =>
+ dynamicPartitionFilter.transform {
+ case a: AttributeReference =>
+ val index = relation.partitionSchema.indexWhere(a.name == _.name)
+ BoundReference(index, relation.partitionSchema(index).dataType, nullable = true)
+ }
+ }
+ val boundPredicate = Predicate.create(boundedFilters.reduce(And), Nil)
+ val returnedFiles =
+ selectedPartitions.filterAndPruneFiles(boundPredicate, dynamicDataFilters)
+ setFilesNumAndSizeMetric(returnedFiles, false)
+ returnedFiles
+ }(t => driverMetrics("pruningTime").set(t))
+ } else {
+ selectedPartitions
+ }
+ sendDriverMetrics()
+ selected
+ }
+
+ def getPartitionArray: Array[PartitionDirectory] = {
+ // Use dynamicallySelectedPartitions which triggers selectedPartitions evaluation,
+ // ensuring driver metrics (numFiles, filesSize, etc.) are properly set.
+ // Convert ScanFileListing to Array[PartitionDirectory] via filePartitionIterator.
+ dynamicallySelectedPartitions.filePartitionIterator.map {
+ listing => PartitionDirectory(listing.values, listing.files.toSeq)
+ }.toArray
+ }
+
+ /**
+ * Create an RDD for bucketed reads. The non-bucketed variant of this function is
+ * [[createReadRDD]].
+ *
+ * The algorithm is pretty simple: each RDD partition being returned should include all the files
+ * with the same bucket id from all the given Hive partitions.
+ *
+ * @param bucketSpec
+ * the bucketing spec.
+ * @param selectedPartitions
+ * Hive-style partition that are part of the read.
+ */
+ private def createBucketedReadPartition(
+ bucketSpec: BucketSpec,
+ selectedPartitions: ScanFileListing): Seq[FilePartition] = {
+ logInfo(log"Planning with ${MDC(COUNT, bucketSpec.numBuckets)} buckets")
+ val partitionArray = selectedPartitions.toPartitionArray
+ val filesGroupedToBuckets = partitionArray.groupBy {
+ f =>
+ BucketingUtils
+ .getBucketId(f.toPath.getName)
+ .getOrElse(throw QueryExecutionErrors.invalidBucketFile(f.urlEncodedPath))
+ }
+
+ val prunedFilesGroupedToBuckets = if (optionalBucketSet.isDefined) {
+ val bucketSet = optionalBucketSet.get
+ filesGroupedToBuckets.filter(f => bucketSet.get(f._1))
+ } else {
+ filesGroupedToBuckets
+ }
+
+ val filePartitions = optionalNumCoalescedBuckets
+ .map {
+ numCoalescedBuckets =>
+ logInfo(log"Coalescing to ${MDC(COUNT, numCoalescedBuckets)} buckets")
+ val coalescedBuckets = prunedFilesGroupedToBuckets.groupBy(_._1 % numCoalescedBuckets)
+ Seq.tabulate(numCoalescedBuckets) {
+ bucketId =>
+ val partitionedFiles = coalescedBuckets
+ .get(bucketId)
+ .map {
+ _.values.flatten.toArray
+ }
+ .getOrElse(Array.empty)
+ FilePartition(bucketId, partitionedFiles)
+ }
+ }
+ .getOrElse {
+ Seq.tabulate(bucketSpec.numBuckets) {
+ bucketId =>
+ FilePartition(bucketId, prunedFilesGroupedToBuckets.getOrElse(bucketId, Array.empty))
+ }
+ }
+ filePartitions
+ }
+
+ /**
+ * Create an RDD for non-bucketed reads. The bucketed variant of this function is
+ * [[createBucketedReadRDD]].
+ *
+ * @param selectedPartitions
+ * Hive-style partition that are part of the read.
+ */
+ private def createReadPartitions(selectedPartitions: ScanFileListing): Seq[FilePartition] = {
+ val openCostInBytes = relation.sparkSession.sessionState.conf.filesOpenCostInBytes
+ val maxSplitBytes =
+ FilePartition.maxSplitBytes(relation.sparkSession, selectedPartitions)
+ logInfo(log"Planning scan with bin packing, max size: ${MDC(MAX_SPLIT_BYTES, maxSplitBytes)} " +
+ log"bytes, open cost is considered as scanning ${MDC(OPEN_COST_IN_BYTES, openCostInBytes)} " +
+ log"bytes.")
+
+ // Filter files with bucket pruning if possible
+ val bucketingEnabled = relation.sparkSession.sessionState.conf.bucketingEnabled
+ val shouldProcess: Path => Boolean = optionalBucketSet match {
+ case Some(bucketSet) if bucketingEnabled =>
+ // Do not prune the file if bucket file name is invalid
+ filePath => BucketingUtils.getBucketId(filePath.getName).forall(bucketSet.get)
+ case _ =>
+ _ => true
+ }
+
+ val splitFiles = selectedPartitions.filePartitionIterator
+ .flatMap {
+ partition =>
+ val ListingPartition(partitionVals, _, fileStatusIterator) = partition
+ fileStatusIterator.flatMap {
+ file =>
+ // getPath() is very expensive so we only want to call it once in this block:
+ val filePath = file.getPath
+ if (shouldProcess(filePath)) {
+ val isSplitable =
+ relation.fileFormat.isSplitable(relation.sparkSession, relation.options, filePath)
+ PartitionedFileUtil.splitFiles(
+ file = file,
+ filePath = filePath,
+ isSplitable = isSplitable,
+ maxSplitBytes = maxSplitBytes,
+ partitionValues = partitionVals
+ )
+ } else {
+ Seq.empty
+ }
+ }
+ }
+ .toArray
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+
+ val partitions = FilePartition
+ .getFilePartitions(relation.sparkSession, splitFiles.toImmutableArraySeq, maxSplitBytes)
+ partitions
+ }
+
+ def getPartitionsSeq(): Seq[Partition] = {
+ if (bucketedScan) {
+ createBucketedReadPartition(relation.bucketSpec.get, dynamicallySelectedPartitions)
+ } else {
+ createReadPartitions(dynamicallySelectedPartitions)
+ }
+ }
+}
+
+abstract class ArrowFileSourceScanLikeShim(original: FileSourceScanExec)
+ extends FileSourceScanLike {
+ override val nodeNamePrefix: String = "ArrowFile"
+
+ override def tableIdentifier: Option[TableIdentifier] = original.tableIdentifier
+
+ override def inputRDDs(): Seq[RDD[InternalRow]] = original.inputRDDs()
+
+ override def dataFilters: Seq[Expression] = original.dataFilters
+
+ override def disableBucketedScan: Boolean = original.disableBucketedScan
+
+ override def optionalBucketSet: Option[BitSet] = original.optionalBucketSet
+
+ override def optionalNumCoalescedBuckets: Option[Int] = original.optionalNumCoalescedBuckets
+
+ override def partitionFilters: Seq[Expression] = original.partitionFilters
+
+ override def relation: HadoopFsRelation = original.relation
+
+ override def requiredSchema: StructType = original.requiredSchema
+
+ override def getStream: Option[SparkDataStream] = original.stream
+}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/execution/GlutenFileFormatWriter.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/GlutenFileFormatWriter.scala
new file mode 100644
index 0000000000..8dc5fbc964
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/GlutenFileFormatWriter.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution
+
+import org.apache.spark.internal.io.FileCommitProtocol
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.execution.datasources.{FileFormatWriter, WriteJobDescription, WriteTaskResult}
+
+object GlutenFileFormatWriter {
+ def writeFilesExecuteTask(
+ description: WriteJobDescription,
+ jobTrackerID: String,
+ sparkStageId: Int,
+ sparkPartitionId: Int,
+ sparkAttemptNumber: Int,
+ committer: FileCommitProtocol,
+ iterator: Iterator[InternalRow]): WriteTaskResult = {
+ FileFormatWriter.executeTask(
+ description,
+ jobTrackerID,
+ sparkStageId,
+ sparkPartitionId,
+ sparkAttemptNumber,
+ committer,
+ iterator,
+ None
+ )
+ }
+
+ // Wrapper for throwing standardized write error using QueryExecutionErrors
+ def wrapWriteError(cause: Throwable, writePath: String): Nothing = {
+ throw QueryExecutionErrors.taskFailedWhileWritingRowsError(writePath, cause)
+ }
+}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/PartitioningAndOrderingPreservingNodeShim.scala
similarity index 57%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/execution/PartitioningAndOrderingPreservingNodeShim.scala
index f7326e98d6..f0ffcc8b69 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/PartitioningAndOrderingPreservingNodeShim.scala
@@ -14,20 +14,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.execution
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
-
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
- }
-}
+trait OrderPreservingNodeShim extends OrderPreservingUnaryExecNode
+trait PartitioningPreservingNodeShim extends PartitioningPreservingUnaryExecNode
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFooterReaderShim.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFooterReaderShim.scala
new file mode 100644
index 0000000000..f4cc013ad4
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFooterReaderShim.scala
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.parquet
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileStatus, Path}
+import org.apache.parquet.format.converter.ParquetMetadataConverter
+import org.apache.parquet.hadoop.metadata.ParquetMetadata
+import org.apache.parquet.hadoop.util.HadoopInputFile
+
+/** Shim layer for ParquetFooterReader to maintain compatibility across different Spark versions. */
+object ParquetFooterReaderShim {
+
+ /** @since Spark 4.1 */
+ def readFooter(
+ configuration: Configuration,
+ fileStatus: FileStatus,
+ filter: ParquetMetadataConverter.MetadataFilter): ParquetMetadata = {
+ ParquetFooterReader.readFooter(HadoopInputFile.fromStatus(fileStatus, configuration), filter)
+ }
+
+ /** @since Spark 4.1 */
+ def readFooter(
+ configuration: Configuration,
+ file: Path,
+ filter: ParquetMetadataConverter.MetadataFilter): ParquetMetadata = {
+ ParquetFooterReader.readFooter(HadoopInputFile.fromPath(file, configuration), filter)
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AbstractBatchScanExec.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AbstractBatchScanExec.scala
new file mode 100644
index 0000000000..7348e3ac46
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AbstractBatchScanExec.scala
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2
+
+import org.apache.spark.SparkException
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, SinglePartition}
+import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper}
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.util.ArrayImplicits._
+
+import com.google.common.base.Objects
+
+/**
+ * Physical plan node for scanning a batch of data from a data source v2. Please ref BatchScanExec
+ * in Spark.
+ *
+ * Ported to the Spark 4.2 storage-partitioned-join API: Spark 4.2 removed
+ * `StoragePartitionJoinParams` / `KeyGroupedPartitioning` and replaced them with the single
+ * `keyGroupedPartitioning` expression list plus `KeyedPartitioning`. The partition
+ * grouping/replication that BatchScanExec used to perform in `inputRDD` was moved out of the scan
+ * into `GroupPartitionsExec`, so this class now mirrors Spark 4.2's `BatchScanExec`.
+ */
+abstract class AbstractBatchScanExec(
+ output: Seq[AttributeReference],
+ @transient scan: Scan,
+ val runtimeFilters: Seq[Expression],
+ ordering: Option[Seq[SortOrder]] = None,
+ @transient table: Table,
+ val keyGroupedPartitioning: Option[Seq[Expression]] = None
+) extends DataSourceV2ScanExecBase {
+
+ @transient lazy val batch: Batch = if (scan == null) null else scan.toBatch
+
+ // TODO: unify the equal/hashCode implementation for all data source v2 query plans.
+ override def equals(other: Any): Boolean = other match {
+ case other: AbstractBatchScanExec =>
+ this.batch != null && this.batch == other.batch &&
+ this.runtimeFilters == other.runtimeFilters &&
+ this.keyGroupedPartitioning == other.keyGroupedPartitioning
+ case _ =>
+ false
+ }
+
+ override def hashCode(): Int = Objects.hashCode(batch, runtimeFilters)
+
+ @transient override lazy val inputPartitions: Seq[InputPartition] = inputPartitionsShim
+
+ @transient protected lazy val inputPartitionsShim: Seq[InputPartition] =
+ batch.planInputPartitions().toImmutableArraySeq
+
+ @transient private lazy val filteredPartitions: Seq[Option[InputPartition]] = {
+ val originalPartitioning = outputPartitioning
+
+ val filtered = PushDownUtils.pushRuntimeFilters(scan, runtimeFilters, table, output)
+ if (filtered) {
+ // call toBatch again to get filtered partitions
+ val newPartitions = scan.toBatch.planInputPartitions()
+
+ originalPartitioning match {
+ case k: KeyedPartitioning =>
+ 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 " +
+ "filtering")
+ }
+
+ val inputMap = k.partitionKeys.groupBy(identity).view.mapValues(_.size)
+ val comparableKeyWrapperFactory = InternalRowComparableWrapper
+ .getInternalRowComparableWrapperFactory(k.expressionDataTypes)
+ val filteredMap = newPartitions.groupBy(
+ p => comparableKeyWrapperFactory(p.asInstanceOf[HasPartitionKey].partitionKey()))
+
+ if (!filteredMap.keySet.subsetOf(inputMap.keySet)) {
+ throw new SparkException(
+ "During runtime filtering, data source must not report new " +
+ "partition keys that are not present in the original partitioning.")
+ }
+
+ inputMap.toSeq
+ .sortBy(_._1)(k.keyOrdering)
+ .flatMap {
+ case (key, size) =>
+ // We require the new number of partitions to be equal or less than the old number
+ // of partitions for a given key. In the case of less than, empty partitions are
+ // added.
+ val fps = filteredMap.getOrElse(key, Array.empty)
+
+ if (fps.size > size) {
+ throw new SparkException(
+ "During runtime filtering, data source must not report " +
+ s"new partitions for a given key. Before: $size partitions. " +
+ s"After: ${fps.size} partitions")
+ }
+
+ fps.map(Some(_)).padTo(size, None)
+ }
+
+ case _ =>
+ // no validation is needed as the data source did not report any specific partitioning
+ newPartitions.toSeq.map(Some(_))
+ }
+
+ } else {
+ (originalPartitioning match {
+ case k: KeyedPartitioning =>
+ inputPartitions.sortBy(_.asInstanceOf[HasPartitionKey].partitionKey())(k.keyRowOrdering)
+
+ case _ => inputPartitions
+ }).map(Some(_))
+ }
+ }
+
+ override lazy val readerFactory: PartitionReaderFactory = batch.createReaderFactory()
+
+ override lazy val inputRDD: RDD[InternalRow] = {
+ val rdd = if (filteredPartitions.isEmpty && outputPartitioning == SinglePartition) {
+ // return an empty RDD with 1 partition if dynamic filtering removed the only split
+ sparkContext.parallelize(Array.empty[InternalRow].toImmutableArraySeq, 1)
+ } else {
+ new DataSourceRDD(
+ sparkContext,
+ filteredPartitions,
+ readerFactory,
+ supportsColumnar,
+ customMetrics)
+ }
+ postDriverMetrics(scan.reportDriverMetrics())
+ rdd
+ }
+
+ override def simpleString(maxFields: Int): String = {
+ val truncatedOutputString = truncatedString(output, "[", ", ", "]", maxFields)
+ val runtimeFiltersString = s"RuntimeFilters: ${runtimeFilters.mkString("[", ",", "]")}"
+ val result = s"$nodeName$truncatedOutputString ${scan.description()} $runtimeFiltersString"
+ redact(result)
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExecShim.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExecShim.scala
new file mode 100644
index 0000000000..c5ad98e1a6
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExecShim.scala
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2
+
+import org.apache.spark.SparkException
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
+import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.connector.catalog.functions.Reducer
+import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
+import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, Scan}
+import org.apache.spark.sql.execution.datasources.v2.orc.OrcScan
+import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+// Spark 4.2 removed `StoragePartitionJoinParams` and no longer accepts the SPJ parameters
+// (`joinKeyPositions`, `commonPartitionValues`, `reducers`, `applyPartialClustering`,
+// `replicatePartitions`) on the scan node -- that grouping/replication now happens in
+// `GroupPartitionsExec`. To keep the public constructor identical to the other Spark shims
+// (Gluten's own planner reads these vals), they are kept here as shim-local fields and are simply
+// not forwarded into the Spark superclass, which now only takes `keyGroupedPartitioning`.
+abstract class BatchScanExecShim(
+ output: Seq[AttributeReference],
+ @transient scan: Scan,
+ runtimeFilters: Seq[Expression],
+ keyGroupedPartitioning: Option[Seq[Expression]] = None,
+ ordering: Option[Seq[SortOrder]] = None,
+ @transient val table: Table,
+ val joinKeyPositions: Option[Seq[Int]] = None,
+ val commonPartitionValues: Option[Seq[(InternalRow, Int)]] = None,
+ val reducers: Option[Seq[Option[Reducer[_, _]]]] = None,
+ val applyPartialClustering: Boolean = false,
+ val replicatePartitions: Boolean = false)
+ extends AbstractBatchScanExec(
+ output,
+ scan,
+ runtimeFilters,
+ ordering,
+ table,
+ keyGroupedPartitioning
+ ) {
+
+ // Note: "metrics" is made transient to avoid sending driver-side metrics to tasks.
+ @transient override lazy val metrics: Map[String, SQLMetric] = Map()
+
+ lazy val metadataColumns: Seq[AttributeReference] = output.collect {
+ case FileSourceConstantMetadataAttribute(attr) => attr
+ case FileSourceGeneratedMetadataAttribute(attr, _) => attr
+ }
+
+ def hasUnsupportedColumns: Boolean = {
+ // TODO, fallback if user define same name column due to we can't right now
+ // detect which column is metadata column which is user defined column.
+ val metadataColumnsNames = metadataColumns.map(_.name)
+ output
+ .filterNot(metadataColumns.toSet)
+ .exists(v => metadataColumnsNames.contains(v.name))
+ }
+
+ // Spark 4.2 moved `postDriverMetrics` to SupportsCustomDriverMetrics and made the reported
+ // task metrics an explicit argument (see BatchScanExec in Spark 4.2).
+ def doPostDriverMetrics(): Unit = {
+ postDriverMetrics(scan.reportDriverMetrics())
+ }
+
+ override def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ throw new UnsupportedOperationException("Need to implement this method")
+ }
+
+ @transient protected lazy val filteredPartitions: Seq[Seq[InputPartition]] = {
+ val originalPartitioning = outputPartitioning
+
+ val filtered = PushDownUtils.pushRuntimeFilters(scan, runtimeFilters, table, output)
+ // call toBatch again to get filtered partitions if any runtime filter was pushed
+ val newPartitions =
+ if (filtered) scan.toBatch.planInputPartitions().toSeq else inputPartitions
+
+ originalPartitioning match {
+ case k: KeyedPartitioning =>
+ 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 " +
+ "filtering")
+ }
+
+ if (filtered) {
+ // Validate that runtime filtering only removed partition keys, never introduced new ones.
+ val newPartitionKeys = newPartitions
+ .map(
+ partition =>
+ InternalRowComparableWrapper(
+ partition.asInstanceOf[HasPartitionKey].partitionKey(),
+ k.expressions))
+ .toSet
+ val oldPartitionKeys = k.partitionKeys.toSet
+ // We require the new number of partition keys to be equal or less than the old number.
+ if (oldPartitionKeys.size < newPartitionKeys.size) {
+ throw new SparkException(
+ "During runtime filtering, data source must either report " +
+ "the same number of partition values, or a subset of partition values from the " +
+ s"original. Before: ${oldPartitionKeys.size} partition values. " +
+ s"After: ${newPartitionKeys.size} partition values")
+ }
+ if (!newPartitionKeys.forall(oldPartitionKeys.contains)) {
+ throw new SparkException(
+ "During runtime filtering, data source must not report new " +
+ "partition values that are not present in the original partitioning.")
+ }
+ }
+
+ // Group the splits that share the same partition key and sort the groups by key ascending.
+ // This reproduces the key-grouped layout Spark 4.1's `BatchScanExec` produced and that
+ // Gluten's native `SparkShims.orderPartitions` still expects.
+ //
+ // FIXME(Spark 4.2, apache/gluten#13139): this collapsing is incorrect for Spark 4.2, which
+ // keeps splits ungrouped and defers grouping to `GroupPartitionsExec` -- e.g. keys
+ // [A, A, B] advertise 3 partitions but build only 2, so `GroupPartitionsExec` requests an
+ // index the reduced RDD lacks. It is currently unreachable (only `KeyGroupedPartitioning`
+ // /SPJ sources hit it, and no such source -- e.g. Iceberg -- is enabled on Spark 4.2), so
+ // it is dormant today. It must be reworked before native keyed DSv2 is enabled on 4.2; see
+ // the issue for the two fix options.
+ newPartitions
+ .map(part => (part.asInstanceOf[HasPartitionKey].partitionKey(), part))
+ .groupBy { case (key, _) => InternalRowComparableWrapper(key, k.expressions) }
+ .toSeq
+ .sortBy(_._1)(k.keyOrdering)
+ .map { case (_, keyedParts) => keyedParts.map(_._2) }
+
+ case _ =>
+ // no validation is needed as the data source did not report any specific partitioning
+ newPartitions.map(Seq(_))
+ }
+ }
+
+ @transient lazy val pushedAggregate: Option[Aggregation] = {
+ scan match {
+ case s: ParquetScan => s.pushedAggregate
+ case o: OrcScan => o.pushedAggregate
+ case _ => None
+ }
+ }
+}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/utils/CatalogUtil.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/utils/CatalogUtil.scala
index f7326e98d6..c6686ba552 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/utils/CatalogUtil.scala
@@ -14,20 +14,16 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.execution.datasources.v2.utils
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
+import org.apache.spark.sql.connector.expressions.Transform
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+object CatalogUtil {
+
+ def convertPartitionTransforms(partitions: Seq[Transform]): (Seq[String], Option[BucketSpec]) = {
+ import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.TransformHelper
+ val (identityCols, bucketSpec, _) = partitions.convertTransforms
+ (identityCols, bucketSpec)
}
}
diff --git a/shims/spark40/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
similarity index 68%
copy from shims/spark40/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
index 127a0fc3cf..8839fad099 100644
--- a/shims/spark40/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/python/BasePythonRunnerShim.scala
@@ -21,6 +21,7 @@ import org.apache.spark.TaskContext
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions, PythonWorker}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata
+import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.vectorized.ColumnarBatch
import java.io.DataOutputStream
@@ -50,8 +51,7 @@ abstract class BasePythonRunnerShim(
PythonUDFRunner.writeUDFs(
dataOut,
funcs,
- argOffsets.map(_.map(pair => ArgumentMetadata(pair._1, pair._2))),
- None)
+ argOffsets.map(_.map(pair => ArgumentMetadata(pair._1, pair._2))))
}
override protected def newWriter(
@@ -63,4 +63,25 @@ abstract class BasePythonRunnerShim(
createNewWriter(env, worker, inputIterator, partitionIndex, context)
}
+ // Spark 4.2 (SPARK-51384) frames the Python worker command as
+ // evalType -> runnerConf -> evalConf -> writeCommand, and expects writeCommand to emit only the
+ // UDF definitions. Older profiles hand-wrote the config prefix inside writeCommand; on Spark 4.2
+ // that config must instead flow through these hooks (mirroring Spark's own ArrowPythonRunner /
+ // ArrowPythonWithNamedArgumentRunner), otherwise the worker misreads the extra prefix.
+ // The concrete runner supplies the values via pythonRunnerConfMap / pythonInputSchema.
+ private val SQL_ARROW_BATCHED_UDF = 101
+
+ protected def pythonRunnerConfMap: Map[String, String] = Map.empty
+
+ protected def pythonInputSchema: StructType = new StructType()
+
+ override def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConfMap
+
+ override def evalConf: Map[String, String] =
+ if (evalType == SQL_ARROW_BATCHED_UDF) {
+ super.evalConf + ("input_type" -> pythonInputSchema.json)
+ } else {
+ super.evalConf
+ }
+
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonExecBase.scala
similarity index 57%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonExecBase.scala
index f7326e98d6..7ad7ca6b09 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonExecBase.scala
@@ -14,20 +14,22 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.execution.python
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.catalyst.expressions.{Expression, NamedArgumentExpression}
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
+abstract class EvalPythonExecBase extends EvalPythonExec {
+
+ override protected def evaluatorFactory: EvalPythonEvaluatorFactory = {
+ throw new IllegalStateException("EvalPythonExecTransformer doesn't support evaluate")
+ }
+}
+
+object EvalPythonExecBase {
+ object NamedArgumentExpressionShim {
+ def unapply(expr: Expression): Option[(String, Expression)] = expr match {
+ case NamedArgumentExpression(key, value) => Some((key, value))
+ case _ => None
}
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/streaming/MemoryStream.scala
similarity index 58%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/execution/streaming/MemoryStream.scala
index f7326e98d6..cd29f54fc6 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/streaming/MemoryStream.scala
@@ -14,20 +14,19 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.execution.streaming
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
+import org.apache.spark.sql.{Encoder, SQLContext}
+import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream => RuntimeMemoryStream}
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
+object MemoryStream {
+ def apply[A: Encoder](implicit sqlContext: SQLContext): RuntimeMemoryStream[A] = {
+ RuntimeMemoryStream[A](implicitly[Encoder[A]], sqlContext.sparkSession)
+ }
+
+ def apply[A: Encoder](
+ numPartitions: Int)(
+ implicit sqlContext: SQLContext): RuntimeMemoryStream[A] = {
+ RuntimeMemoryStream[A](sqlContext.sparkSession, numPartitions)
}
}
diff --git a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/ui/TypeAlias.scala
similarity index 57%
copy from gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
copy to shims/spark42/src/main/scala/org/apache/spark/sql/execution/ui/TypeAlias.scala
index f7326e98d6..40c66590e0 100644
--- a/gluten-core/src/main/scala/org/apache/spark/util/SparkVersionUtil.scala
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/execution/ui/TypeAlias.scala
@@ -14,20 +14,12 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.util
+package org.apache.spark.sql.execution.ui
-object SparkVersionUtil {
- val gteSpark35: Boolean = compareMajorMinorVersion((3, 5)) >= 0
- val gteSpark40: Boolean = compareMajorMinorVersion((4, 0)) >= 0
- val gteSpark41: Boolean = compareMajorMinorVersion((4, 1)) >= 0
-
- // Returns X. X < 0 if one < other, x == 0 if one == other, x > 0 if one > other.
- def compareMajorMinorVersion(other: (Int, Int)): Int = {
- val (major, minor) = VersionUtils.majorMinorVersion(org.apache.spark.SPARK_VERSION)
- if (major == other._1) {
- minor - other._2
- } else {
- major - other._1
- }
- }
+/**
+ * Ensures compatibility for the type HttpServletRequest across Spark 4.0 and earlier versions.
+ * Starting from Spark 4.0, `jakarta.servlet.http.HttpServletRequest` is used.
+ */
+object TypeAlias {
+ type HttpServletRequest = jakarta.servlet.http.HttpServletRequest
}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala b/shims/spark42/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala
new file mode 100644
index 0000000000..5e747e7df3
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala
@@ -0,0 +1,277 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.hive.execution
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.CastSupport
+import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.SchemaPruning.RootField
+import org.apache.spark.sql.execution._
+import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.sql.hive._
+import org.apache.spark.sql.hive.client.HiveClientImpl
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BooleanType, DataType, StructType}
+import org.apache.spark.util.Utils
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.hive.ql.io.{DelegateSymlinkTextInputFormat, SymlinkTextInputFormat}
+import org.apache.hadoop.hive.ql.metadata.{Partition => HivePartition}
+import org.apache.hadoop.hive.ql.plan.TableDesc
+import org.apache.hadoop.hive.serde.serdeConstants
+import org.apache.hadoop.hive.serde2.objectinspector._
+import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils.ObjectInspectorCopyOption
+import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoUtils
+import org.apache.hadoop.mapred.InputFormat
+
+import scala.collection.JavaConverters._
+
+/**
+ * The Hive table scan operator. Column and partition pruning are both handled.
+ *
+ * @param requestedAttributes
+ * Attributes to be fetched from the Hive table.
+ * @param relation
+ * The Hive table be scanned.
+ * @param partitionPruningPred
+ * An optional partition pruning predicate for partitioned table.
+ * @param prunedOutput
+ * The pruned output.
+ */
+abstract private[hive] class AbstractHiveTableScanExec(
+ requestedAttributes: Seq[Attribute],
+ relation: HiveTableRelation,
+ partitionPruningPred: Seq[Expression],
+ prunedOutput: Seq[Attribute] = Seq.empty[Attribute])(
+ @transient protected val sparkSession: SparkSession)
+ extends LeafExecNode
+ with CastSupport {
+
+ require(
+ partitionPruningPred.isEmpty || relation.isPartitioned,
+ "Partition pruning predicates only supported for partitioned tables.")
+
+ override def conf: SQLConf = sparkSession.sessionState.conf
+
+ override def nodeName: String = s"Scan hive ${relation.tableMeta.qualifiedName}"
+
+ override lazy val metrics = Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"))
+
+ override def producedAttributes: AttributeSet = outputSet ++
+ AttributeSet(partitionPruningPred.flatMap(_.references))
+
+ private val originalAttributes = AttributeMap(relation.output.map(a => a -> a))
+
+ override def output: Seq[Attribute] = {
+ if (prunedOutput.nonEmpty) {
+ prunedOutput
+ } else {
+ // Retrieve the original attributes based on expression ID so that capitalization matches.
+ requestedAttributes.map(attr => originalAttributes.getOrElse(attr, attr)).distinct
+ }
+ }
+
+ // Bind all partition key attribute references in the partition pruning predicate for later
+ // evaluation.
+ private lazy val boundPruningPred = partitionPruningPred.reduceLeftOption(And).map {
+ pred =>
+ require(
+ pred.dataType == BooleanType,
+ s"Data type of predicate $pred must be ${BooleanType.catalogString} rather than " +
+ s"${pred.dataType.catalogString}.")
+
+ BindReferences.bindReference(pred, relation.partitionCols)
+ }
+
+ @transient private lazy val hiveQlTable = HiveClientImpl.toHiveTable(relation.tableMeta)
+ @transient private lazy val tableDesc = new TableDesc(
+ getInputFormat(hiveQlTable.getInputFormatClass, conf),
+ hiveQlTable.getOutputFormatClass,
+ hiveQlTable.getMetadata)
+
+ // Create a local copy of hadoopConf,so that scan specific modifications should not impact
+ // other queries
+ @transient private lazy val hadoopConf = {
+ val c = sparkSession.sessionState.newHadoopConf()
+ // append columns ids and names before broadcast
+ addColumnMetadataToConf(c)
+ c
+ }
+
+ @transient private lazy val hadoopReader =
+ new HadoopTableReader(output, relation.partitionCols, tableDesc, sparkSession, hadoopConf)
+
+ private def castFromString(value: String, dataType: DataType) = {
+ cast(Literal(value), dataType).eval(null)
+ }
+
+ private def addColumnMetadataToConf(hiveConf: Configuration): Unit = {
+ // Specifies needed column IDs for those non-partitioning columns.
+ val columnOrdinals = AttributeMap(relation.dataCols.zipWithIndex)
+ val neededColumnIDs = output.flatMap(columnOrdinals.get).map(o => o: Integer)
+ val neededColumnNames = output.filter(columnOrdinals.contains).map(_.name)
+
+ HiveShim.appendReadColumns(hiveConf, neededColumnIDs, neededColumnNames)
+
+ val deserializer = tableDesc.getDeserializerClass.getConstructor().newInstance()
+ deserializer.initialize(hiveConf, tableDesc.getProperties)
+
+ // Specifies types and object inspectors of columns to be scanned.
+ val structOI = ObjectInspectorUtils
+ .getStandardObjectInspector(deserializer.getObjectInspector, ObjectInspectorCopyOption.JAVA)
+ .asInstanceOf[StructObjectInspector]
+
+ val columnTypeNames = structOI.getAllStructFieldRefs.asScala
+ .map(_.getFieldObjectInspector)
+ .map(TypeInfoUtils.getTypeInfoFromObjectInspector(_).getTypeName)
+ .mkString(",")
+
+ hiveConf.set(serdeConstants.LIST_COLUMN_TYPES, columnTypeNames)
+ hiveConf.set(serdeConstants.LIST_COLUMNS, relation.dataCols.map(_.name).mkString(","))
+ }
+
+ /**
+ * Prunes partitions not involve the query plan.
+ *
+ * @param partitions
+ * All partitions of the relation.
+ * @return
+ * Partitions that are involved in the query plan.
+ */
+ private[hive] def prunePartitions(partitions: Seq[HivePartition]): Seq[HivePartition] = {
+ boundPruningPred match {
+ case None => partitions
+ case Some(shouldKeep) =>
+ partitions.filter {
+ part =>
+ val dataTypes = relation.partitionCols.map(_.dataType)
+ val castedValues = part.getValues.asScala
+ .zip(dataTypes)
+ .map { case (value, dataType) => castFromString(value, dataType) }
+
+ // Only partitioned values are needed here, since the predicate has
+ // already been bound to partition key attribute references.
+ val row = InternalRow.fromSeq(castedValues.toSeq)
+ shouldKeep.eval(row).asInstanceOf[Boolean]
+ }
+ }
+ }
+
+ // This is used on the driver side, so it is important to avoid executing subqueries
+ @transient lazy val basePrunedPartitions: Seq[HivePartition] = {
+ if (relation.prunedPartitions.nonEmpty) {
+ relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
+ } else {
+ rawPartitions
+ }
+ }
+
+ @transient lazy val prunedPartitions: Seq[HivePartition] =
+ if (relation.prunedPartitions.nonEmpty) {
+ if (partitionPruningPred.forall(!ExecSubqueryExpression.hasSubquery(_))) {
+ basePrunedPartitions
+ } else {
+ prunePartitions(basePrunedPartitions)
+ }
+ } else if (
+ sparkSession.sessionState.conf.metastorePartitionPruning &&
+ partitionPruningPred.nonEmpty
+ ) {
+ basePrunedPartitions
+ } else {
+ prunePartitions(basePrunedPartitions)
+ }
+
+ // exposed for tests
+ @transient lazy val rawPartitions: Seq[HivePartition] = {
+ val prunedPartitions =
+ if (
+ sparkSession.sessionState.conf.metastorePartitionPruning &&
+ partitionPruningPred.nonEmpty
+ ) {
+ // Retrieve the original attributes based on expression ID so that capitalization matches.
+ val normalizedFilters = partitionPruningPred.map(_.transform {
+ case a: AttributeReference => originalAttributes(a)
+ })
+ sparkSession.sessionState.catalog
+ .listPartitionsByFilter(relation.tableMeta.identifier, normalizedFilters)
+ } else {
+ sparkSession.sessionState.catalog.listPartitions(relation.tableMeta.identifier)
+ }
+ prunedPartitions.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
+ }
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ // Using dummyCallSite, as getCallSite can turn out to be expensive with
+ // multiple partitions.
+ val rdd = if (!relation.isPartitioned) {
+ Utils.withDummyCallSite(sparkContext) {
+ hadoopReader.makeRDDForTable(hiveQlTable)
+ }
+ } else {
+ Utils.withDummyCallSite(sparkContext) {
+ hadoopReader.makeRDDForPartitionedTable(prunedPartitions)
+ }
+ }
+ val numOutputRows = longMetric("numOutputRows")
+ // Avoid to serialize MetastoreRelation because schema is lazy. (see SPARK-15649)
+ val outputSchema = schema
+ rdd.mapPartitionsWithIndexInternal {
+ (index, iter) =>
+ val proj = UnsafeProjection.create(outputSchema)
+ proj.initialize(index)
+ iter.map {
+ r =>
+ numOutputRows += 1
+ proj(r)
+ }
+ }
+ }
+
+ // Filters unused DynamicPruningExpression expressions - one which has been replaced
+ // with DynamicPruningExpression(Literal.TrueLiteral) during Physical Planning
+ private def filterUnusedDynamicPruningExpressions(
+ predicates: Seq[Expression]): Seq[Expression] = {
+ predicates.filterNot(_ == DynamicPruningExpression(Literal.TrueLiteral))
+ }
+
+ // Optionally returns a delegate input format based on the provided input format class.
+ // This is currently used to replace SymlinkTextInputFormat with DelegateSymlinkTextInputFormat
+ // in order to fix SPARK-40815.
+ private def getInputFormat(
+ inputFormatClass: Class[_ <: InputFormat[_, _]],
+ conf: SQLConf): Class[_ <: InputFormat[_, _]] = {
+ if (
+ inputFormatClass == classOf[SymlinkTextInputFormat] &&
+ conf != null && conf.getConf(HiveUtils.USE_DELEGATE_FOR_SYMLINK_TEXT_INPUT_FORMAT)
+ ) {
+ classOf[DelegateSymlinkTextInputFormat]
+ } else {
+ inputFormatClass
+ }
+ }
+
+ override def otherCopyArgs: Seq[AnyRef] = Seq(sparkSession)
+
+ def pruneSchema(schema: StructType, requestedFields: Seq[RootField]): StructType = {
+ SchemaPruning.pruneSchema(schema, requestedFields)
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/storage/GlutenPushBasedFetchHelper.scala b/shims/spark42/src/main/scala/org/apache/spark/storage/GlutenPushBasedFetchHelper.scala
new file mode 100644
index 0000000000..c230cd7e8c
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/storage/GlutenPushBasedFetchHelper.scala
@@ -0,0 +1,399 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.storage
+
+import org.apache.spark.MapOutputTracker
+import org.apache.spark.MapOutputTracker.SHUFFLE_PUSH_MAP_ID
+import org.apache.spark.internal.Logging
+import org.apache.spark.network.shuffle.{BlockStoreClient, MergedBlockMeta, MergedBlocksMetaListener}
+import org.apache.spark.shuffle.ShuffleReadMetricsReporter
+import org.apache.spark.storage.BlockManagerId.SHUFFLE_MERGER_IDENTIFIER
+import org.apache.spark.storage.ShuffleBlockFetcherIterator._
+
+import org.roaringbitmap.RoaringBitmap
+
+import java.util.concurrent.TimeUnit
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.util.{Failure, Success}
+
+/**
+ * Helper class for [[ShuffleBlockFetcherIterator]] that encapsulates all the push-based
+ * functionality to fetch push-merged block meta and shuffle chunks. A push-merged block contains
+ * multiple shuffle chunks where each shuffle chunk contains multiple shuffle blocks that belong to
+ * the common reduce partition and were merged by the external shuffle service to that chunk.
+ */
+private class GlutenPushBasedFetchHelper(
+ private val iterator: GlutenShuffleBlockFetcherIterator,
+ private val shuffleClient: BlockStoreClient,
+ private val blockManager: BlockManager,
+ private val mapOutputTracker: MapOutputTracker,
+ private val shuffleMetrics: ShuffleReadMetricsReporter)
+ extends Logging {
+
+ private[this] val startTimeNs = System.nanoTime()
+
+ private[storage] val localShuffleMergerBlockMgrId = BlockManagerId(
+ SHUFFLE_MERGER_IDENTIFIER,
+ blockManager.blockManagerId.host,
+ blockManager.blockManagerId.port,
+ blockManager.blockManagerId.topologyInfo)
+
+ /** A map for storing shuffle chunk bitmap. */
+ private[this] val chunksMetaMap = new mutable.HashMap[ShuffleBlockChunkId, RoaringBitmap]()
+
+ /** Returns true if the address is for a push-merged block. */
+ def isPushMergedShuffleBlockAddress(address: BlockManagerId): Boolean = {
+ SHUFFLE_MERGER_IDENTIFIER == address.executorId
+ }
+
+ /** Returns true if the address is of a remote push-merged block. false otherwise. */
+ def isRemotePushMergedBlockAddress(address: BlockManagerId): Boolean = {
+ isPushMergedShuffleBlockAddress(address) && address.host != blockManager.blockManagerId.host
+ }
+
+ /** Returns true if the address is of a push-merged-local block. false otherwise. */
+ def isLocalPushMergedBlockAddress(address: BlockManagerId): Boolean = {
+ isPushMergedShuffleBlockAddress(address) && address.host == blockManager.blockManagerId.host
+ }
+
+ /**
+ * This is executed by the task thread when the `iterator.next()` is invoked and the iterator
+ * processes a response of type [[ShuffleBlockFetcherIterator.SuccessFetchResult]].
+ *
+ * @param blockId
+ * shuffle chunk id.
+ */
+ def removeChunk(blockId: ShuffleBlockChunkId): Unit = {
+ chunksMetaMap.remove(blockId)
+ }
+
+ /**
+ * This is executed by the task thread when the `iterator.next()` is invoked and the iterator
+ * processes a response of type [[ShuffleBlockFetcherIterator.PushMergedLocalMetaFetchResult]].
+ *
+ * @param blockId
+ * shuffle chunk id.
+ */
+ def addChunk(blockId: ShuffleBlockChunkId, chunkMeta: RoaringBitmap): Unit = {
+ chunksMetaMap(blockId) = chunkMeta
+ }
+
+ /**
+ * Get the RoaringBitMap for a specific ShuffleBlockChunkId
+ *
+ * @param blockId
+ * shuffle chunk id.
+ */
+ def getRoaringBitMap(blockId: ShuffleBlockChunkId): Option[RoaringBitmap] = {
+ chunksMetaMap.get(blockId)
+ }
+
+ /**
+ * Get the number of map blocks in a ShuffleBlockChunk
+ * @param blockId
+ * @return
+ */
+ def getShuffleChunkCardinality(blockId: ShuffleBlockChunkId): Int = {
+ getRoaringBitMap(blockId).map(_.getCardinality).getOrElse(0)
+ }
+
+ /**
+ * This is executed by the task thread when the `iterator.next()` is invoked and the iterator
+ * processes a response of type [[ShuffleBlockFetcherIterator.PushMergedRemoteMetaFetchResult]].
+ *
+ * @param shuffleId
+ * shuffle id.
+ * @param reduceId
+ * reduce id.
+ * @param blockSize
+ * size of the push-merged block.
+ * @param bitmaps
+ * chunk bitmaps, where each bitmap contains all the mapIds that were merged to that chunk.
+ * @return
+ * shuffle chunks to fetch.
+ */
+ def createChunkBlockInfosFromMetaResponse(
+ shuffleId: Int,
+ shuffleMergeId: Int,
+ reduceId: Int,
+ blockSize: Long,
+ bitmaps: Array[RoaringBitmap]): ArrayBuffer[(BlockId, Long, Int)] = {
+ val approxChunkSize = blockSize / bitmaps.length
+ val blocksToFetch = new ArrayBuffer[(BlockId, Long, Int)]()
+ for (i <- bitmaps.indices) {
+ val blockChunkId = ShuffleBlockChunkId(shuffleId, shuffleMergeId, reduceId, i)
+ chunksMetaMap.put(blockChunkId, bitmaps(i))
+ logDebug(s"adding block chunk $blockChunkId of size $approxChunkSize")
+ blocksToFetch += ((blockChunkId, approxChunkSize, SHUFFLE_PUSH_MAP_ID))
+ }
+ blocksToFetch
+ }
+
+ /**
+ * This is executed by the task thread when the iterator is initialized and only if it has
+ * push-merged blocks for which it needs to fetch the metadata.
+ *
+ * @param req
+ * [[ShuffleBlockFetcherIterator.FetchRequest]] that only contains requests to fetch metadata of
+ * push-merged blocks.
+ */
+ def sendFetchMergedStatusRequest(req: FetchRequest): Unit = {
+ val sizeMap = req.blocks.map {
+ case FetchBlockInfo(blockId, size, _) =>
+ val shuffleBlockId = blockId.asInstanceOf[ShuffleMergedBlockId]
+ ((shuffleBlockId.shuffleId, shuffleBlockId.reduceId), size)
+ }.toMap
+ val address = req.address
+ val mergedBlocksMetaListener = new MergedBlocksMetaListener {
+ override def onSuccess(
+ shuffleId: Int,
+ shuffleMergeId: Int,
+ reduceId: Int,
+ meta: MergedBlockMeta): Unit = {
+ logDebug(
+ s"Received the meta of push-merged block for ($shuffleId, $shuffleMergeId," +
+ s" $reduceId) from ${req.address.host}:${req.address.port}")
+ try {
+ iterator.addToResultsQueue(
+ PushMergedRemoteMetaFetchResult(
+ shuffleId,
+ shuffleMergeId,
+ reduceId,
+ sizeMap((shuffleId, reduceId)),
+ meta.readChunkBitmaps(),
+ address))
+ } catch {
+ case exception: Exception =>
+ logError(
+ s"Failed to parse the meta of push-merged block for ($shuffleId, " +
+ s"$shuffleMergeId, $reduceId) from" +
+ s" ${req.address.host}:${req.address.port}",
+ exception
+ )
+ iterator.addToResultsQueue(
+ PushMergedRemoteMetaFailedFetchResult(shuffleId, shuffleMergeId, reduceId, address))
+ }
+ }
+
+ override def onFailure(
+ shuffleId: Int,
+ shuffleMergeId: Int,
+ reduceId: Int,
+ exception: Throwable): Unit = {
+ logError(
+ s"Failed to get the meta of push-merged block for ($shuffleId, $reduceId) " +
+ s"from ${req.address.host}:${req.address.port}",
+ exception)
+ iterator.addToResultsQueue(
+ PushMergedRemoteMetaFailedFetchResult(shuffleId, shuffleMergeId, reduceId, address))
+ }
+ }
+ req.blocks.foreach {
+ block =>
+ val shuffleBlockId = block.blockId.asInstanceOf[ShuffleMergedBlockId]
+ shuffleClient.getMergedBlockMeta(
+ address.host,
+ address.port,
+ shuffleBlockId.shuffleId,
+ shuffleBlockId.shuffleMergeId,
+ shuffleBlockId.reduceId,
+ mergedBlocksMetaListener)
+ }
+ }
+
+ /**
+ * This is executed by the task thread when the iterator is initialized. It fetches all the
+ * outstanding push-merged local blocks.
+ * @param pushMergedLocalBlocks
+ * set of identified merged local blocks and their sizes.
+ */
+ def fetchAllPushMergedLocalBlocks(pushMergedLocalBlocks: mutable.LinkedHashSet[BlockId]): Unit = {
+ if (pushMergedLocalBlocks.nonEmpty) {
+ blockManager.hostLocalDirManager.foreach(fetchPushMergedLocalBlocks(_, pushMergedLocalBlocks))
+ }
+ }
+
+ /**
+ * Fetch the push-merged blocks dirs if they are not in the cache and eventually fetch push-merged
+ * local blocks.
+ */
+ private def fetchPushMergedLocalBlocks(
+ hostLocalDirManager: HostLocalDirManager,
+ pushMergedLocalBlocks: mutable.LinkedHashSet[BlockId]): Unit = {
+ val cachedPushedMergedDirs =
+ hostLocalDirManager.getCachedHostLocalDirsFor(SHUFFLE_MERGER_IDENTIFIER)
+ if (cachedPushedMergedDirs.isDefined) {
+ logDebug(
+ s"Fetch the push-merged-local blocks with cached merged dirs: " +
+ s"${cachedPushedMergedDirs.get.mkString(", ")}")
+ pushMergedLocalBlocks.foreach {
+ blockId =>
+ fetchPushMergedLocalBlock(
+ blockId,
+ cachedPushedMergedDirs.get,
+ localShuffleMergerBlockMgrId)
+ }
+ } else {
+ // Push-based shuffle is only enabled when the external shuffle service is enabled. If the
+ // external shuffle service is not enabled, then there will not be any push-merged blocks
+ // for the iterator to fetch.
+ logDebug(
+ s"Asynchronous fetch the push-merged-local blocks without cached merged " +
+ s"dirs from the external shuffle service")
+ hostLocalDirManager.getHostLocalDirs(
+ blockManager.blockManagerId.host,
+ blockManager.externalShuffleServicePort,
+ Array(SHUFFLE_MERGER_IDENTIFIER)) {
+ case Success(dirs) =>
+ logDebug(
+ s"Fetched merged dirs in " +
+ s"${TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTimeNs)} ms")
+ pushMergedLocalBlocks.foreach {
+ blockId =>
+ logDebug(
+ s"Successfully fetched local dirs: " +
+ s"${dirs.get(SHUFFLE_MERGER_IDENTIFIER).mkString(", ")}")
+ fetchPushMergedLocalBlock(
+ blockId,
+ dirs(SHUFFLE_MERGER_IDENTIFIER),
+ localShuffleMergerBlockMgrId)
+ }
+ case Failure(throwable) =>
+ // If we see an exception with getting the local dirs for push-merged-local blocks,
+ // we fallback to fetch the original blocks. We do not report block fetch failure.
+ logWarning(
+ s"Error while fetching the merged dirs for push-merged-local " +
+ s"blocks: ${pushMergedLocalBlocks.mkString(", ")}. Fetch the original blocks instead",
+ throwable
+ )
+ pushMergedLocalBlocks.foreach {
+ blockId =>
+ iterator.addToResultsQueue(
+ FallbackOnPushMergedFailureResult(
+ blockId,
+ localShuffleMergerBlockMgrId,
+ 0,
+ isNetworkReqDone = false))
+ }
+ }
+ }
+ }
+
+ /**
+ * Fetch a single push-merged-local block generated. This can also be executed by the task thread
+ * as well as the netty thread.
+ * @param blockId
+ * ShuffleBlockId to be fetched
+ * @param localDirs
+ * Local directories where the push-merged shuffle files are stored
+ * @param blockManagerId
+ * BlockManagerId
+ */
+ private[this] def fetchPushMergedLocalBlock(
+ blockId: BlockId,
+ localDirs: Array[String],
+ blockManagerId: BlockManagerId): Unit = {
+ try {
+ val shuffleBlockId = blockId.asInstanceOf[ShuffleMergedBlockId]
+ val chunksMeta = blockManager.getLocalMergedBlockMeta(shuffleBlockId, localDirs)
+ iterator.addToResultsQueue(
+ PushMergedLocalMetaFetchResult(
+ shuffleBlockId.shuffleId,
+ shuffleBlockId.shuffleMergeId,
+ shuffleBlockId.reduceId,
+ chunksMeta.readChunkBitmaps(),
+ localDirs))
+ } catch {
+ case e: Exception =>
+ // If we see an exception with reading a push-merged-local meta, we fallback to
+ // fetch the original blocks. We do not report block fetch failure
+ // and will continue with the remaining local block read.
+ logWarning(
+ s"Error occurred while fetching push-merged-local meta, " +
+ s"prepare to fetch the original blocks",
+ e)
+ iterator.addToResultsQueue(
+ FallbackOnPushMergedFailureResult(blockId, blockManagerId, 0, isNetworkReqDone = false))
+ }
+ }
+
+ /**
+ * This is executed by the task thread when the `iterator.next()` is invoked and the iterator
+ * processes a response of type: 1) [[ShuffleBlockFetcherIterator.SuccessFetchResult]] 2)
+ * [[ShuffleBlockFetcherIterator.FallbackOnPushMergedFailureResult]] 3)
+ * [[ShuffleBlockFetcherIterator.PushMergedRemoteMetaFailedFetchResult]]
+ *
+ * This initiates fetching fallback blocks for a push-merged block or a shuffle chunk that failed
+ * to fetch. It makes a call to the map output tracker to get the list of original blocks for the
+ * given push-merged block/shuffle chunk, split them into remote and local blocks, and process
+ * them accordingly. It also updates the numberOfBlocksToFetch in the iterator as it processes
+ * failed response and finds more push-merged requests to remote and again updates it with
+ * additional requests for original blocks. The fallback happens when:
+ * 1. There is an exception while creating shuffle chunks from push-merged-local shuffle block.
+ * See fetchLocalBlock. 2. There is a failure when fetching remote shuffle chunks. 3. There
+ * is a failure when processing SuccessFetchResult which is for a shuffle chunk (local or
+ * remote). 4. There is a zero-size buffer when processing SuccessFetchResult for a shuffle
+ * chunk (local or remote).
+ */
+ def initiateFallbackFetchForPushMergedBlock(blockId: BlockId, address: BlockManagerId): Unit = {
+ assert(blockId.isInstanceOf[ShuffleMergedBlockId] || blockId.isInstanceOf[ShuffleBlockChunkId])
+ logWarning(s"Falling back to fetch the original blocks for push-merged block $blockId")
+ shuffleMetrics.incMergedFetchFallbackCount(1)
+ // Increase the blocks processed since we will process another block in the next iteration of
+ // the while loop in ShuffleBlockFetcherIterator.next().
+ val fallbackBlocksByAddr: Iterator[(BlockManagerId, collection.Seq[(BlockId, Long, Int)])] =
+ blockId match {
+ case shuffleBlockId: ShuffleMergedBlockId =>
+ iterator.decreaseNumBlocksToFetch(1)
+ mapOutputTracker.getMapSizesForMergeResult(
+ shuffleBlockId.shuffleId,
+ shuffleBlockId.reduceId)
+ case _ =>
+ val shuffleChunkId = blockId.asInstanceOf[ShuffleBlockChunkId]
+ val chunkBitmap: RoaringBitmap = chunksMetaMap.remove(shuffleChunkId).get
+ var blocksProcessed = 1
+ // When there is a failure to fetch a remote shuffle chunk, then we try to
+ // fallback not only for that particular remote shuffle chunk but also for all the
+ // pending chunks that belong to the same host. The reason for doing so is that it
+ // is very likely that the subsequent requests for shuffle chunks from this host will
+ // fail as well. Since, push-based shuffle is best effort and we try not to increase the
+ // delay of the fetches, we immediately fallback for all the pending shuffle chunks in the
+ // fetchRequests queue.
+ if (isRemotePushMergedBlockAddress(address)) {
+ // Fallback for all the pending fetch requests
+ val pendingShuffleChunks = iterator.removePendingChunks(shuffleChunkId, address)
+ pendingShuffleChunks.foreach {
+ pendingBlockId =>
+ logInfo(s"Falling back immediately for shuffle chunk $pendingBlockId")
+ shuffleMetrics.incMergedFetchFallbackCount(1)
+ val bitmapOfPendingChunk: RoaringBitmap = chunksMetaMap.remove(pendingBlockId).get
+ chunkBitmap.or(bitmapOfPendingChunk)
+ }
+ // These blocks were added to numBlocksToFetch so we increment numBlocksProcessed
+ blocksProcessed += pendingShuffleChunks.size
+ }
+ iterator.decreaseNumBlocksToFetch(blocksProcessed)
+ mapOutputTracker.getMapSizesForMergeResult(
+ shuffleChunkId.shuffleId,
+ shuffleChunkId.reduceId,
+ chunkBitmap)
+ }
+ iterator.fallbackFetch(fallbackBlocksByAddr)
+ }
+}
diff --git a/shims/spark42/src/main/scala/org/apache/spark/storage/GlutenShuffleBlockFetcherIterator.scala b/shims/spark42/src/main/scala/org/apache/spark/storage/GlutenShuffleBlockFetcherIterator.scala
new file mode 100644
index 0000000000..3ddc901d49
--- /dev/null
+++ b/shims/spark42/src/main/scala/org/apache/spark/storage/GlutenShuffleBlockFetcherIterator.scala
@@ -0,0 +1,1601 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.storage
+
+import org.apache.spark.{MapOutputTracker, SparkException, TaskContext}
+import org.apache.spark.MapOutputTracker.SHUFFLE_PUSH_MAP_ID
+import org.apache.spark.errors.SparkCoreErrors
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys._
+import org.apache.spark.network.buffer.{FileSegmentManagedBuffer, ManagedBuffer}
+import org.apache.spark.network.shuffle._
+import org.apache.spark.network.shuffle.checksum.{Cause, ShuffleChecksumHelper}
+import org.apache.spark.network.util.TransportConf
+import org.apache.spark.shuffle.ShuffleReadMetricsReporter
+import org.apache.spark.util.{Clock, SystemClock, TaskCompletionListener, Utils}
+
+import io.netty.util.internal.OutOfDirectMemoryError
+
+import javax.annotation.concurrent.GuardedBy
+
+import java.io.{InputStream, IOException}
+import java.nio.channels.ClosedByInterruptException
+import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingDeque, TimeUnit}
+import java.util.zip.CheckedInputStream
+
+import scala.collection.mutable
+import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet, Queue}
+import scala.util.{Failure, Success}
+
+/**
+ * An iterator that fetches multiple blocks. For local blocks, it fetches from the local block
+ * manager. For remote blocks, it fetches them using the provided BlockTransferService.
+ *
+ * This creates an iterator of (BlockID, InputStream) tuples so the caller can handle blocks in a
+ * pipelined fashion as they are received.
+ *
+ * The implementation throttles the remote fetches so they don't exceed maxBytesInFlight to avoid
+ * using too much memory.
+ *
+ * @param context
+ * [[TaskContext]], used for metrics update
+ * @param shuffleClient
+ * [[BlockStoreClient]] for fetching remote blocks
+ * @param blockManager
+ * [[BlockManager]] for reading local blocks
+ * @param blocksByAddress
+ * list of blocks to fetch grouped by the [[BlockManagerId]]. For each block we also require two
+ * info: 1. the size (in bytes as a long field) in order to throttle the memory usage; 2. the
+ * mapIndex for this block, which indicate the index in the map stage. Note that zero-sized blocks
+ * are already excluded, which happened in
+ * [[org.apache.spark.MapOutputTracker.convertMapStatuses]].
+ * @param mapOutputTracker
+ * [[MapOutputTracker]] for falling back to fetching the original blocks if we fail to fetch
+ * shuffle chunks when push based shuffle is enabled.
+ * @param streamWrapper
+ * A function to wrap the returned input stream.
+ * @param maxBytesInFlight
+ * max size (in bytes) of remote blocks to fetch at any given point.
+ * @param maxReqsInFlight
+ * max number of remote requests to fetch blocks at any given point.
+ * @param maxBlocksInFlightPerAddress
+ * max number of shuffle blocks being fetched at any given point for a given remote host:port.
+ * @param maxReqSizeShuffleToMem
+ * max size (in bytes) of a request that can be shuffled to memory.
+ * @param maxAttemptsOnNettyOOM
+ * The max number of a block could retry due to Netty OOM before throwing the fetch failure.
+ * @param detectCorrupt
+ * whether to detect any corruption in fetched blocks.
+ * @param checksumEnabled
+ * whether the shuffle checksum is enabled. When enabled, Spark will try to diagnose the cause of
+ * the block corruption.
+ * @param checksumAlgorithm
+ * the checksum algorithm that is used when calculating the checksum value for the block data.
+ * @param shuffleMetrics
+ * used to report shuffle metrics.
+ * @param doBatchFetch
+ * fetch continuous shuffle blocks from same executor in batch if the server side supports.
+ */
+final class GlutenShuffleBlockFetcherIterator(
+ context: TaskContext,
+ shuffleClient: BlockStoreClient,
+ blockManager: BlockManager,
+ mapOutputTracker: MapOutputTracker,
+ blocksByAddress: Iterator[(BlockManagerId, collection.Seq[(BlockId, Long, Int)])],
+ streamWrapper: (BlockId, InputStream) => InputStream,
+ maxBytesInFlight: Long,
+ maxReqsInFlight: Int,
+ maxBlocksInFlightPerAddress: Int,
+ val maxReqSizeShuffleToMem: Long,
+ maxAttemptsOnNettyOOM: Int,
+ detectCorrupt: Boolean,
+ detectCorruptUseExtraMemory: Boolean,
+ checksumEnabled: Boolean,
+ checksumAlgorithm: String,
+ shuffleMetrics: ShuffleReadMetricsReporter,
+ doBatchFetch: Boolean,
+ clock: Clock = new SystemClock())
+ extends GlutenShuffleBlockFetcherIteratorBase
+ with DownloadFileManager
+ with Logging {
+
+ import ShuffleBlockFetcherIterator._
+
+ // Make remote requests at most maxBytesInFlight / 5 in length; the reason to keep them
+ // smaller than maxBytesInFlight is to allow multiple, parallel fetches from up to 5
+ // nodes, rather than blocking on reading output from one node.
+ private val targetRemoteRequestSize = math.max(maxBytesInFlight / 5, 1L)
+
+ /** Total number of blocks to fetch. */
+ private[this] var numBlocksToFetch = 0
+
+ /**
+ * The number of blocks processed by the caller. The iterator is exhausted when
+ * [[numBlocksProcessed]] == [[numBlocksToFetch]].
+ */
+ private[this] var numBlocksProcessed = 0
+
+ private[this] val startTimeNs = System.nanoTime()
+
+ /** Host local blocks to fetch, excluding zero-sized blocks. */
+ private[this] val hostLocalBlocks = scala.collection.mutable.LinkedHashSet[(BlockId, Int)]()
+
+ /**
+ * A queue to hold our results. This turns the asynchronous model provided by
+ * [[org.apache.spark.network.BlockTransferService]] into a synchronous model (iterator).
+ */
+ private[this] val results = new LinkedBlockingDeque[FetchResult]()
+
+ /**
+ * Current [[FetchResult]] being processed per thread. We track this so we can release the current
+ * buffer in case of a runtime exception when processing the current buffer. Using
+ * ConcurrentHashMap to support concurrent access from multiple threads while allowing cleanup
+ * from any thread.
+ */
+ private[this] val currentResults = ConcurrentHashMap.newKeySet[SuccessFetchResult]()
+
+ /**
+ * Queue of fetch requests to issue; we'll pull requests off this gradually to make sure that the
+ * number of bytes in flight is limited to maxBytesInFlight.
+ */
+ private[this] val fetchRequests = new Queue[FetchRequest]
+
+ /**
+ * Queue of fetch requests which could not be issued the first time they were dequeued. These
+ * requests are tried again when the fetch constraints are satisfied.
+ */
+ private[this] val deferredFetchRequests = new HashMap[BlockManagerId, Queue[FetchRequest]]()
+
+ /** Current bytes in flight from our requests */
+ private[this] var bytesInFlight = 0L
+
+ /** Current number of requests in flight */
+ private[this] var reqsInFlight = 0
+
+ /** Current number of blocks in flight per host:port */
+ private[this] val numBlocksInFlightPerAddress = new HashMap[BlockManagerId, Int]()
+
+ /**
+ * Count the retry times for the blocks due to Netty OOM. The block will stop retry if retry times
+ * has exceeded the [[maxAttemptsOnNettyOOM]].
+ */
+ private[this] val blockOOMRetryCounts = new HashMap[String, Int]
+
+ /**
+ * The blocks that can't be decompressed successfully, it is used to guarantee that we retry at
+ * most once for those corrupted blocks.
+ */
+ private[this] val corruptedBlocks = mutable.HashSet[BlockId]()
+
+ /**
+ * Whether the iterator is still active. If isZombie is true, the callback interface will no
+ * longer place fetched blocks into [[results]].
+ */
+ @GuardedBy("this")
+ private[this] var isZombie = false
+
+ /**
+ * A set to store the files used for shuffling remote huge blocks. Files in this set will be
+ * deleted when cleanup. This is a layer of defensiveness against disk file leaks.
+ */
+ @GuardedBy("this")
+ private[this] val shuffleFilesSet = mutable.HashSet[DownloadFile]()
+
+ private[this] val onCompleteCallback = new GlutenShuffleFetchCompletionListener(this)
+
+ private[this] val pushBasedFetchHelper =
+ new GlutenPushBasedFetchHelper(
+ this,
+ shuffleClient,
+ blockManager,
+ mapOutputTracker,
+ shuffleMetrics)
+
+ initialize()
+
+ private def withFetchWaitTimeTracked[T](f: => T): T = {
+ val startFetchWait = System.nanoTime()
+ val res = f
+ val fetchWaitTime = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startFetchWait)
+ shuffleMetrics.incFetchWaitTime(fetchWaitTime)
+ res
+ }
+
+ override def createTempFile(transportConf: TransportConf): DownloadFile = {
+ // we never need to do any encryption or decryption here, regardless of configs, because that
+ // is handled at another layer in the code. When encryption is enabled, shuffle data is written
+ // to disk encrypted in the first place, and sent over the network still encrypted.
+ new SimpleDownloadFile(blockManager.diskBlockManager.createTempLocalBlock()._2, transportConf)
+ }
+
+ override def registerTempFileToClean(file: DownloadFile): Boolean = synchronized {
+ if (isZombie) {
+ false
+ } else {
+ shuffleFilesSet += file
+ true
+ }
+ }
+
+ /** Mark the iterator as zombie, and release all buffers that haven't been deserialized yet. */
+ private[storage] def cleanup(): Unit = {
+ synchronized {
+ isZombie = true
+ }
+ // Release all current result buffers from all threads
+ while (!currentResults.isEmpty) {
+ currentResults.toArray(new Array[SuccessFetchResult](0)).foreach {
+ result =>
+ if (currentResults.remove(result)) {
+ result.buf.release()
+ }
+ }
+ }
+
+ // Release buffers in the results queue
+ val iter = results.iterator()
+ while (iter.hasNext) {
+ val result = iter.next()
+ result match {
+ case SuccessFetchResult(blockId, mapIndex, address, _, buf, _) =>
+ if (address != blockManager.blockManagerId) {
+ if (
+ pushBasedFetchHelper.isLocalPushMergedBlockAddress(address) ||
+ hostLocalBlocks.contains(blockId -> mapIndex)
+ ) {
+ shuffleMetricsUpdate(blockId, buf, local = true)
+ } else {
+ shuffleMetricsUpdate(blockId, buf, local = false)
+ }
+ }
+ buf.release()
+ case _ =>
+ }
+ }
+ shuffleFilesSet.foreach {
+ file =>
+ if (!file.delete()) {
+ logWarning(log"Failed to cleanup shuffle fetch temp file ${MDC(PATH, file.path())}")
+ }
+ }
+ }
+
+ private[this] def sendRequest(req: FetchRequest): Unit = {
+ logDebug(
+ "Sending request for %d blocks (%s) from %s"
+ .format(req.blocks.size, Utils.bytesToString(req.size), req.address.hostPort))
+ bytesInFlight += req.size
+ reqsInFlight += 1
+
+ // so we can look up the block info of each blockID
+ val infoMap = req.blocks.map {
+ case FetchBlockInfo(blockId, size, mapIndex) => (blockId.toString, (size, mapIndex))
+ }.toMap
+ val remainingBlocks = new HashSet[String]() ++= infoMap.keys
+ val deferredBlocks = new ArrayBuffer[String]()
+ val blockIds = req.blocks.map(_.blockId.toString)
+ val address = req.address
+ val requestStartTime = clock.nanoTime()
+
+ @inline def enqueueDeferredFetchRequestIfNecessary(): Unit = {
+ if (remainingBlocks.isEmpty && deferredBlocks.nonEmpty) {
+ val blocks = deferredBlocks.map {
+ blockId =>
+ val (size, mapIndex) = infoMap(blockId)
+ FetchBlockInfo(BlockId(blockId), size, mapIndex)
+ }
+ results.put(DeferFetchRequestResult(FetchRequest(address, blocks)))
+ deferredBlocks.clear()
+ }
+ }
+
+ @inline def updateMergedReqsDuration(wasReqForMergedChunks: Boolean = false): Unit = {
+ if (remainingBlocks.isEmpty) {
+ val durationMs = TimeUnit.NANOSECONDS.toMillis(clock.nanoTime() - requestStartTime)
+ if (wasReqForMergedChunks) {
+ shuffleMetrics.incRemoteMergedReqsDuration(durationMs)
+ }
+ shuffleMetrics.incRemoteReqsDuration(durationMs)
+ }
+ }
+
+ val blockFetchingListener = new BlockFetchingListener {
+ override def onBlockFetchSuccess(blockId: String, buf: ManagedBuffer): Unit = {
+ // Only add the buffer to results queue if the iterator is not zombie,
+ // i.e. cleanup() has not been called yet.
+ GlutenShuffleBlockFetcherIterator.this.synchronized {
+ if (!isZombie) {
+ // Increment the ref count because we need to pass this to a different thread.
+ // This needs to be released after use.
+ buf.retain()
+ remainingBlocks -= blockId
+ blockOOMRetryCounts.remove(blockId)
+ updateMergedReqsDuration(BlockId(blockId).isShuffleChunk)
+ results.put(
+ SuccessFetchResult(
+ BlockId(blockId),
+ infoMap(blockId)._2,
+ address,
+ infoMap(blockId)._1,
+ buf,
+ remainingBlocks.isEmpty))
+ logDebug("remainingBlocks: " + remainingBlocks)
+ enqueueDeferredFetchRequestIfNecessary()
+ }
+ }
+ logTrace(s"Got remote block $blockId after ${Utils.getUsedTimeNs(startTimeNs)}")
+ }
+
+ override def onBlockFetchFailure(blockId: String, e: Throwable): Unit = {
+ GlutenShuffleBlockFetcherIterator.this.synchronized {
+ logError(
+ log"Failed to get block(s) from " +
+ log"${MDC(HOST, req.address.host)}:${MDC(PORT, req.address.port)}",
+ e)
+ e match {
+ // SPARK-27991: Catch the Netty OOM and set the flag `isNettyOOMOnShuffle` (shared among
+ // tasks) to true as early as possible. The pending fetch requests won't be sent
+ // afterwards until the flag is set to false on:
+ // 1) the Netty free memory >= maxReqSizeShuffleToMem
+ // - we'll check this whenever there's a fetch request succeeds.
+ // 2) the number of in-flight requests becomes 0
+ // - we'll check this in `fetchUpToMaxBytes` whenever it's invoked.
+ // Although Netty memory is shared across multiple modules, e.g., shuffle, rpc, the flag
+ // only takes effect for the shuffle due to the implementation simplicity concern.
+ // And we'll buffer the consecutive block failures caused by the OOM error until there's
+ // no remaining blocks in the current request. Then, we'll package these blocks into
+ // a same fetch request for the retry later. In this way, instead of creating the fetch
+ // request per block, it would help reduce the concurrent connections and data loads
+ // pressure at remote server.
+ // Note that catching OOM and do something based on it is only a workaround for
+ // handling the Netty OOM issue, which is not the best way towards memory management.
+ // We can get rid of it when we find a way to manage Netty's memory precisely.
+ case _: OutOfDirectMemoryError
+ if blockOOMRetryCounts.getOrElseUpdate(blockId, 0) < maxAttemptsOnNettyOOM =>
+ if (!isZombie) {
+ val failureTimes = blockOOMRetryCounts(blockId)
+ blockOOMRetryCounts(blockId) += 1
+ if (isNettyOOMOnShuffle.compareAndSet(false, true)) {
+ // The fetcher can fail remaining blocks in batch for the same error. So we only
+ // log the warning once to avoid flooding the logs.
+ logInfo(log"Block ${MDC(BLOCK_ID, blockId)} has failed " +
+ log"${MDC(FAILURES, failureTimes)} times due to Netty OOM, will retry")
+ }
+ remainingBlocks -= blockId
+ deferredBlocks += blockId
+ enqueueDeferredFetchRequestIfNecessary()
+ }
+
+ case _ =>
+ val block = BlockId(blockId)
+ if (block.isShuffleChunk) {
+ remainingBlocks -= blockId
+ updateMergedReqsDuration(wasReqForMergedChunks = true)
+ results.put(
+ FallbackOnPushMergedFailureResult(
+ block,
+ address,
+ infoMap(blockId)._1,
+ remainingBlocks.isEmpty))
+ } else {
+ results.putFirst(FailureFetchResult(block, infoMap(blockId)._2, address, e))
+ }
+ }
+ }
+ }
+ }
+
+ // Fetch remote shuffle blocks to disk when the request is too large. Since the shuffle data is
+ // already encrypted and compressed over the wire(w.r.t. the related configs), we can just fetch
+ // the data and write it to file directly.
+ if (req.size > maxReqSizeShuffleToMem) {
+ shuffleClient.fetchBlocks(
+ address.host,
+ address.port,
+ address.executorId,
+ blockIds.toArray,
+ blockFetchingListener,
+ this)
+ } else {
+ shuffleClient.fetchBlocks(
+ address.host,
+ address.port,
+ address.executorId,
+ blockIds.toArray,
+ blockFetchingListener,
+ null)
+ }
+ }
+
+ /**
+ * This is called from initialize and also from the fallback which is triggered from
+ * [[PushBasedFetchHelper]].
+ */
+ private[this] def partitionBlocksByFetchMode(
+ blocksByAddress: Iterator[(BlockManagerId, collection.Seq[(BlockId, Long, Int)])],
+ localBlocks: mutable.LinkedHashSet[(BlockId, Int)],
+ hostLocalBlocksByExecutor: mutable.LinkedHashMap[
+ BlockManagerId,
+ collection.Seq[(BlockId, Long, Int)]],
+ pushMergedLocalBlocks: mutable.LinkedHashSet[BlockId]): ArrayBuffer[FetchRequest] = {
+ logDebug(
+ s"maxBytesInFlight: $maxBytesInFlight, targetRemoteRequestSize: "
+ + s"$targetRemoteRequestSize, maxBlocksInFlightPerAddress: $maxBlocksInFlightPerAddress")
+
+ // Partition to local, host-local, push-merged-local, remote (includes push-merged-remote)
+ // blocks.Remote blocks are further split into FetchRequests of size at most maxBytesInFlight
+ // in order to limit the amount of data in flight
+ val collectedRemoteRequests = new ArrayBuffer[FetchRequest]
+ var localBlockBytes = 0L
+ var hostLocalBlockBytes = 0L
+ var numHostLocalBlocks = 0
+ var pushMergedLocalBlockBytes = 0L
+ val prevNumBlocksToFetch = numBlocksToFetch
+
+ val fallback = FallbackStorage.FALLBACK_BLOCK_MANAGER_ID.executorId
+ val localExecIds = Set(blockManager.blockManagerId.executorId, fallback)
+ for ((address, blockInfos) <- blocksByAddress) {
+ checkBlockSizes(blockInfos)
+ if (pushBasedFetchHelper.isPushMergedShuffleBlockAddress(address)) {
+ // These are push-merged blocks or shuffle chunks of these blocks.
+ if (address.host == blockManager.blockManagerId.host) {
+ numBlocksToFetch += blockInfos.size
+ pushMergedLocalBlocks ++= blockInfos.map(_._1)
+ pushMergedLocalBlockBytes += blockInfos.map(_._2).sum
+ } else {
+ collectFetchRequests(address, blockInfos, collectedRemoteRequests)
+ }
+ } else if (localExecIds.contains(address.executorId)) {
+ val mergedBlockInfos = mergeContinuousShuffleBlockIdsIfNeeded(
+ blockInfos.map(info => FetchBlockInfo(info._1, info._2, info._3)),
+ doBatchFetch)
+ numBlocksToFetch += mergedBlockInfos.size
+ localBlocks ++= mergedBlockInfos.map(info => (info.blockId, info.mapIndex))
+ localBlockBytes += mergedBlockInfos.map(_.size).sum
+ } else if (
+ blockManager.hostLocalDirManager.isDefined &&
+ address.host == blockManager.blockManagerId.host
+ ) {
+ val mergedBlockInfos = mergeContinuousShuffleBlockIdsIfNeeded(
+ blockInfos.map(info => FetchBlockInfo(info._1, info._2, info._3)),
+ doBatchFetch)
+ numBlocksToFetch += mergedBlockInfos.size
+ val blocksForAddress =
+ mergedBlockInfos.map(info => (info.blockId, info.size, info.mapIndex))
+ hostLocalBlocksByExecutor += address -> blocksForAddress
+ numHostLocalBlocks += blocksForAddress.size
+ hostLocalBlockBytes += mergedBlockInfos.map(_.size).sum
+ } else {
+ val (_, timeCost) = Utils.timeTakenMs[Unit] {
+ collectFetchRequests(address, blockInfos, collectedRemoteRequests)
+ }
+ logDebug(s"Collected remote fetch requests for $address in $timeCost ms")
+ }
+ }
+ val (remoteBlockBytes, numRemoteBlocks) =
+ collectedRemoteRequests.foldLeft((0L, 0))((x, y) => (x._1 + y.size, x._2 + y.blocks.size))
+ val totalBytes = localBlockBytes + remoteBlockBytes + hostLocalBlockBytes +
+ pushMergedLocalBlockBytes
+ val blocksToFetchCurrentIteration = numBlocksToFetch - prevNumBlocksToFetch
+ assert(
+ blocksToFetchCurrentIteration == localBlocks.size +
+ numHostLocalBlocks + numRemoteBlocks + pushMergedLocalBlocks.size,
+ s"The number of non-empty blocks $blocksToFetchCurrentIteration doesn't equal to the sum " +
+ s"of the number of local blocks ${localBlocks.size} + " +
+ s"the number of host-local blocks $numHostLocalBlocks " +
+ s"the number of push-merged-local blocks ${pushMergedLocalBlocks.size} " +
+ s"+ the number of remote blocks $numRemoteBlocks "
+ )
+ logInfo(
+ log"Getting ${MDC(NUM_BLOCKS, blocksToFetchCurrentIteration)} " +
+ log"(${MDC(TOTAL_SIZE, Utils.bytesToString(totalBytes))}) non-empty blocks including " +
+ log"${MDC(NUM_LOCAL_BLOCKS, localBlocks.size)} " +
+ log"(${MDC(LOCAL_BLOCKS_SIZE, Utils.bytesToString(localBlockBytes))}) local and " +
+ log"${MDC(NUM_HOST_LOCAL_BLOCKS, numHostLocalBlocks)} " +
+ log"(${MDC(HOST_LOCAL_BLOCKS_SIZE, Utils.bytesToString(hostLocalBlockBytes))}) " +
+ log"host-local and ${MDC(NUM_PUSH_MERGED_LOCAL_BLOCKS, pushMergedLocalBlocks.size)} " +
+ log"(${MDC(
+ PUSH_MERGED_LOCAL_BLOCKS_SIZE,
+ Utils.bytesToString(pushMergedLocalBlockBytes))})" +
+ log" push-merged-local and ${MDC(NUM_REMOTE_BLOCKS, numRemoteBlocks)} " +
+ log"(${MDC(REMOTE_BLOCKS_SIZE, Utils.bytesToString(remoteBlockBytes))}) remote blocks")
+ this.hostLocalBlocks ++= hostLocalBlocksByExecutor.values
+ .flatMap(infos => infos.map(info => (info._1, info._3)))
+ collectedRemoteRequests
+ }
+
+ private def createFetchRequest(
+ blocks: collection.Seq[FetchBlockInfo],
+ address: BlockManagerId,
+ forMergedMetas: Boolean): FetchRequest = {
+ logDebug(
+ s"Creating fetch request of ${blocks.map(_.size).sum} at $address "
+ + s"with ${blocks.size} blocks")
+ FetchRequest(address, blocks, forMergedMetas)
+ }
+
+ private def createFetchRequests(
+ curBlocks: collection.Seq[FetchBlockInfo],
+ address: BlockManagerId,
+ isLast: Boolean,
+ collectedRemoteRequests: ArrayBuffer[FetchRequest],
+ enableBatchFetch: Boolean,
+ forMergedMetas: Boolean = false): ArrayBuffer[FetchBlockInfo] = {
+ val mergedBlocks = mergeContinuousShuffleBlockIdsIfNeeded(curBlocks, enableBatchFetch)
+ numBlocksToFetch += mergedBlocks.size
+ val retBlocks = new ArrayBuffer[FetchBlockInfo]
+ if (mergedBlocks.length <= maxBlocksInFlightPerAddress) {
+ collectedRemoteRequests += createFetchRequest(mergedBlocks, address, forMergedMetas)
+ } else {
+ mergedBlocks.grouped(maxBlocksInFlightPerAddress).foreach {
+ blocks =>
+ if (blocks.length == maxBlocksInFlightPerAddress || isLast) {
+ collectedRemoteRequests += createFetchRequest(blocks, address, forMergedMetas)
+ } else {
+ // The last group does not exceed `maxBlocksInFlightPerAddress`. Put it back
+ // to `curBlocks`.
+ retBlocks ++= blocks
+ numBlocksToFetch -= blocks.size
+ }
+ }
+ }
+ retBlocks
+ }
+
+ private def collectFetchRequests(
+ address: BlockManagerId,
+ blockInfos: collection.Seq[(BlockId, Long, Int)],
+ collectedRemoteRequests: ArrayBuffer[FetchRequest]): Unit = {
+ val iterator = blockInfos.iterator
+ var curRequestSize = 0L
+ var curBlocks = new ArrayBuffer[FetchBlockInfo]()
+
+ while (iterator.hasNext) {
+ val (blockId, size, mapIndex) = iterator.next()
+ curBlocks += FetchBlockInfo(blockId, size, mapIndex)
+ curRequestSize += size
+ blockId match {
+ // Either all blocks are push-merged blocks, shuffle chunks, or original blocks.
+ // Based on these types, we decide to do batch fetch and create FetchRequests with
+ // forMergedMetas set.
+ case ShuffleBlockChunkId(_, _, _, _) =>
+ if (
+ curRequestSize >= targetRemoteRequestSize ||
+ curBlocks.size >= maxBlocksInFlightPerAddress
+ ) {
+ curBlocks = createFetchRequests(
+ curBlocks,
+ address,
+ isLast = false,
+ collectedRemoteRequests,
+ enableBatchFetch = false)
+ curRequestSize = curBlocks.map(_.size).sum
+ }
+ case ShuffleMergedBlockId(_, _, _) =>
+ if (curBlocks.size >= maxBlocksInFlightPerAddress) {
+ curBlocks = createFetchRequests(
+ curBlocks,
+ address,
+ isLast = false,
+ collectedRemoteRequests,
+ enableBatchFetch = false,
+ forMergedMetas = true)
+ }
+ case _ =>
+ // For batch fetch, the actual block in flight should count for merged block.
+ val mayExceedsMaxBlocks = !doBatchFetch && curBlocks.size >= maxBlocksInFlightPerAddress
+ if (curRequestSize >= targetRemoteRequestSize || mayExceedsMaxBlocks) {
+ curBlocks = createFetchRequests(
+ curBlocks,
+ address,
+ isLast = false,
+ collectedRemoteRequests,
+ doBatchFetch)
+ curRequestSize = curBlocks.map(_.size).sum
+ }
+ }
+ }
+ // Add in the final request
+ if (curBlocks.nonEmpty) {
+ val (enableBatchFetch, forMergedMetas) = {
+ curBlocks.head.blockId match {
+ case ShuffleBlockChunkId(_, _, _, _) => (false, false)
+ case ShuffleMergedBlockId(_, _, _) => (false, true)
+ case _ => (doBatchFetch, false)
+ }
+ }
+ createFetchRequests(
+ curBlocks,
+ address,
+ isLast = true,
+ collectedRemoteRequests,
+ enableBatchFetch = enableBatchFetch,
+ forMergedMetas = forMergedMetas)
+ }
+ }
+
+ private def assertPositiveBlockSize(blockId: BlockId, blockSize: Long): Unit = {
+ if (blockSize < 0) {
+ throw BlockException(blockId, "Negative block size " + blockSize)
+ } else if (blockSize == 0) {
+ throw BlockException(blockId, "Zero-sized blocks should be excluded.")
+ }
+ }
+
+ private def checkBlockSizes(blockInfos: collection.Seq[(BlockId, Long, Int)]): Unit = {
+ blockInfos.foreach { case (blockId, size, _) => assertPositiveBlockSize(blockId, size) }
+ }
+
+ /**
+ * Fetch the local blocks while we are fetching remote blocks. This is ok because
+ * `ManagedBuffer`'s memory is allocated lazily when we create the input stream, so all we track
+ * in-memory are the ManagedBuffer references themselves.
+ */
+ private[this] def fetchLocalBlocks(localBlocks: mutable.LinkedHashSet[(BlockId, Int)]): Unit = {
+ logDebug(s"Start fetching local blocks: ${localBlocks.mkString(", ")}")
+ val iter = localBlocks.iterator
+ while (iter.hasNext) {
+ val (blockId, mapIndex) = iter.next()
+ try {
+ val buf = blockManager.getLocalBlockData(blockId)
+ shuffleMetrics.incLocalBlocksFetched(1)
+ shuffleMetrics.incLocalBytesRead(buf.size)
+ buf.retain()
+ results.put(
+ SuccessFetchResult(
+ blockId,
+ mapIndex,
+ blockManager.blockManagerId,
+ buf.size(),
+ buf,
+ false))
+ } catch {
+ // If we see an exception, stop immediately.
+ case e: Exception =>
+ e match {
+ // ClosedByInterruptException is an excepted exception when kill task,
+ // don't log the exception stack trace to avoid confusing users.
+ // See: SPARK-28340
+ case ce: ClosedByInterruptException =>
+ logError(
+ log"Error occurred while fetching local blocks, ${MDC(ERROR, ce.getMessage)}")
+ case ex: Exception => logError("Error occurred while fetching local blocks", ex)
+ }
+ results.putFirst(FailureFetchResult(blockId, mapIndex, blockManager.blockManagerId, e))
+ return
+ }
+ }
+ }
+
+ private[this] def fetchHostLocalBlock(
+ blockId: BlockId,
+ mapIndex: Int,
+ localDirs: Array[String],
+ blockManagerId: BlockManagerId): Boolean = {
+ try {
+ val buf = blockManager.getHostLocalShuffleData(blockId, localDirs)
+ buf.retain()
+ results.put(
+ SuccessFetchResult(
+ blockId,
+ mapIndex,
+ blockManagerId,
+ buf.size(),
+ buf,
+ isNetworkReqDone = false))
+ true
+ } catch {
+ case e: Exception =>
+ // If we see an exception, stop immediately.
+ logError(s"Error occurred while fetching local blocks", e)
+ results.putFirst(FailureFetchResult(blockId, mapIndex, blockManagerId, e))
+ false
+ }
+ }
+
+ /**
+ * Fetch the host-local blocks while we are fetching remote blocks. This is ok because
+ * `ManagedBuffer`'s memory is allocated lazily when we create the input stream, so all we track
+ * in-memory are the ManagedBuffer references themselves.
+ */
+ private[this] def fetchHostLocalBlocks(
+ hostLocalDirManager: HostLocalDirManager,
+ hostLocalBlocksByExecutor: mutable.LinkedHashMap[
+ BlockManagerId,
+ collection.Seq[(BlockId, Long, Int)]]): Unit = {
+ val cachedDirsByExec = hostLocalDirManager.getCachedHostLocalDirs
+ val (hostLocalBlocksWithCachedDirs, hostLocalBlocksWithMissingDirs) = {
+ val (hasCache, noCache) = hostLocalBlocksByExecutor.partition {
+ case (hostLocalBmId, _) =>
+ cachedDirsByExec.contains(hostLocalBmId.executorId)
+ }
+ (hasCache.toMap, noCache.toMap)
+ }
+
+ if (hostLocalBlocksWithMissingDirs.nonEmpty) {
+ logDebug(
+ s"Asynchronous fetching host-local blocks without cached executors' dir: " +
+ s"${hostLocalBlocksWithMissingDirs.mkString(", ")}")
+
+ // If the external shuffle service is enabled, we'll fetch the local directories for
+ // multiple executors from the external shuffle service, which located at the same host
+ // with the executors, in once. Otherwise, we'll fetch the local directories from those
+ // executors directly one by one. The fetch requests won't be too much since one host is
+ // almost impossible to have many executors at the same time practically.
+ val dirFetchRequests = if (blockManager.externalShuffleServiceEnabled) {
+ val host = blockManager.blockManagerId.host
+ val port = blockManager.externalShuffleServicePort
+ Seq((host, port, hostLocalBlocksWithMissingDirs.keys.toArray))
+ } else {
+ hostLocalBlocksWithMissingDirs.keys.map(bmId => (bmId.host, bmId.port, Array(bmId))).toSeq
+ }
+
+ dirFetchRequests.foreach {
+ case (host, port, bmIds) =>
+ hostLocalDirManager.getHostLocalDirs(host, port, bmIds.map(_.executorId)) {
+ case Success(dirsByExecId) =>
+ fetchMultipleHostLocalBlocks(
+ hostLocalBlocksWithMissingDirs.filter { case (k, _) => bmIds.contains(k) },
+ dirsByExecId,
+ cached = false)
+
+ case Failure(throwable) =>
+ logError("Error occurred while fetching host local blocks", throwable)
+ val bmId = bmIds.head
+ val blockInfoSeq = hostLocalBlocksWithMissingDirs(bmId)
+ val (blockId, _, mapIndex) = blockInfoSeq.head
+ results.putFirst(FailureFetchResult(blockId, mapIndex, bmId, throwable))
+ }
+ }
+ }
+
+ if (hostLocalBlocksWithCachedDirs.nonEmpty) {
+ logDebug(
+ s"Synchronous fetching host-local blocks with cached executors' dir: " +
+ s"${hostLocalBlocksWithCachedDirs.mkString(", ")}")
+ fetchMultipleHostLocalBlocks(hostLocalBlocksWithCachedDirs, cachedDirsByExec, cached = true)
+ }
+ }
+
+ private def fetchMultipleHostLocalBlocks(
+ bmIdToBlocks: Map[BlockManagerId, collection.Seq[(BlockId, Long, Int)]],
+ localDirsByExecId: Map[String, Array[String]],
+ cached: Boolean): Unit = {
+ // We use `forall` because once there's a failed block fetch, `fetchHostLocalBlock` will put
+ // a `FailureFetchResult` immediately to the `results`. So there's no reason to fetch the
+ // remaining blocks.
+ val allFetchSucceeded = bmIdToBlocks.forall {
+ case (bmId, blockInfos) =>
+ blockInfos.forall {
+ case (blockId, _, mapIndex) =>
+ fetchHostLocalBlock(blockId, mapIndex, localDirsByExecId(bmId.executorId), bmId)
+ }
+ }
+ if (allFetchSucceeded) {
+ logDebug(
+ s"Got host-local blocks from ${bmIdToBlocks.keys.mkString(", ")} " +
+ s"(${if (cached) "with" else "without"} cached executors' dir) " +
+ s"in ${Utils.getUsedTimeNs(startTimeNs)}")
+ }
+ }
+
+ private[this] def initialize(): Unit = {
+ // Add a task completion callback (called in both success case and failure case) to cleanup.
+ context.addTaskCompletionListener(onCompleteCallback)
+ // Local blocks to fetch, excluding zero-sized blocks.
+ val localBlocks = mutable.LinkedHashSet[(BlockId, Int)]()
+ val hostLocalBlocksByExecutor =
+ mutable.LinkedHashMap[BlockManagerId, collection.Seq[(BlockId, Long, Int)]]()
+ val pushMergedLocalBlocks = mutable.LinkedHashSet[BlockId]()
+ // Partition blocks by the different fetch modes: local, host-local, push-merged-local and
+ // remote blocks.
+ val remoteRequests = partitionBlocksByFetchMode(
+ blocksByAddress,
+ localBlocks,
+ hostLocalBlocksByExecutor,
+ pushMergedLocalBlocks)
+ // Add the remote requests into our queue in a random order
+ fetchRequests ++= Utils.randomize(remoteRequests)
+ assert(
+ (0 == reqsInFlight) == (0 == bytesInFlight),
+ "expected reqsInFlight = 0 but found reqsInFlight = " + reqsInFlight +
+ ", expected bytesInFlight = 0 but found bytesInFlight = " + bytesInFlight
+ )
+
+ // Send out initial requests for blocks, up to our maxBytesInFlight
+ withFetchWaitTimeTracked(fetchUpToMaxBytes())
+
+ val numDeferredRequest = deferredFetchRequests.values.map(_.size).sum
+ val numFetches = remoteRequests.size - fetchRequests.size - numDeferredRequest
+ logInfo(log"Started ${MDC(COUNT, numFetches)} remote fetches in " +
+ log"${MDC(DURATION, Utils.getUsedTimeNs(startTimeNs))}" +
+ (if (numDeferredRequest > 0) log", deferred ${MDC(NUM_REQUESTS, numDeferredRequest)} requests"
+ else log""))
+
+ // Get Local Blocks
+ fetchLocalBlocks(localBlocks)
+ logDebug(s"Got local blocks in ${Utils.getUsedTimeNs(startTimeNs)}")
+ // Get host local blocks if any
+ withFetchWaitTimeTracked(fetchAllHostLocalBlocks(hostLocalBlocksByExecutor))
+ pushBasedFetchHelper.fetchAllPushMergedLocalBlocks(pushMergedLocalBlocks)
+ }
+
+ private def fetchAllHostLocalBlocks(
+ hostLocalBlocksByExecutor: mutable.LinkedHashMap[
+ BlockManagerId,
+ collection.Seq[(BlockId, Long, Int)]]): Unit = {
+ if (hostLocalBlocksByExecutor.nonEmpty) {
+ blockManager.hostLocalDirManager.foreach(fetchHostLocalBlocks(_, hostLocalBlocksByExecutor))
+ }
+ }
+
+ private def shuffleMetricsUpdate(blockId: BlockId, buf: ManagedBuffer, local: Boolean): Unit = {
+ if (local) {
+ shuffleLocalMetricsUpdate(blockId, buf)
+ } else {
+ shuffleRemoteMetricsUpdate(blockId, buf)
+ }
+ }
+
+ private def shuffleLocalMetricsUpdate(blockId: BlockId, buf: ManagedBuffer): Unit = {
+ blockId match {
+ case chunkId: ShuffleBlockChunkId =>
+ val chunkCardinality = pushBasedFetchHelper.getShuffleChunkCardinality(chunkId)
+ shuffleMetrics.incLocalMergedChunksFetched(1)
+ shuffleMetrics.incLocalMergedBlocksFetched(chunkCardinality)
+ shuffleMetrics.incLocalMergedBytesRead(buf.size)
+ shuffleMetrics.incLocalBlocksFetched(chunkCardinality)
+ case _ =>
+ shuffleMetrics.incLocalBlocksFetched(1)
+ }
+ shuffleMetrics.incLocalBytesRead(buf.size)
+ }
+
+ private def shuffleRemoteMetricsUpdate(blockId: BlockId, buf: ManagedBuffer): Unit = {
+ blockId match {
+ case chunkId: ShuffleBlockChunkId =>
+ val chunkCardinality = pushBasedFetchHelper.getShuffleChunkCardinality(chunkId)
+ shuffleMetrics.incRemoteMergedChunksFetched(1)
+ shuffleMetrics.incRemoteMergedBlocksFetched(chunkCardinality)
+ shuffleMetrics.incRemoteMergedBytesRead(buf.size)
+ shuffleMetrics.incRemoteBlocksFetched(chunkCardinality)
+ case _ =>
+ shuffleMetrics.incRemoteBlocksFetched(1)
+ }
+ shuffleMetrics.incRemoteBytesRead(buf.size)
+ if (buf.isInstanceOf[FileSegmentManagedBuffer]) {
+ shuffleMetrics.incRemoteBytesReadToDisk(buf.size)
+ }
+ }
+
+ override def hasNext: Boolean = numBlocksProcessed < numBlocksToFetch
+
+ /**
+ * Fetches the next (BlockId, InputStream). If a task fails, the ManagedBuffers underlying each
+ * InputStream will be freed by the cleanup() method registered with the TaskCompletionListener.
+ * However, callers should close() these InputStreams as soon as they are no longer needed, in
+ * order to release memory as early as possible.
+ *
+ * Throws a FetchFailedException if the next block could not be fetched.
+ */
+ override def next(): (BlockId, InputStream) = {
+ if (!hasNext) {
+ throw SparkCoreErrors.noSuchElementError()
+ }
+
+ numBlocksProcessed += 1
+
+ var result: FetchResult = null
+ var input: InputStream = null
+ // This's only initialized when shuffle checksum is enabled.
+ var checkedIn: CheckedInputStream = null
+ var streamCompressedOrEncrypted: Boolean = false
+ // Take the next fetched result and try to decompress it to detect data corruption,
+ // then fetch it one more time if it's corrupt, throw FailureFetchResult if the second fetch
+ // is also corrupt, so the previous stage could be retried.
+ // For local shuffle block, throw FailureFetchResult for the first IOException.
+ while (result == null) {
+ result = withFetchWaitTimeTracked[FetchResult](results.take())
+
+ result match {
+ case SuccessFetchResult(blockId, mapIndex, address, size, buf, isNetworkReqDone) =>
+ if (address != blockManager.blockManagerId) {
+ if (
+ hostLocalBlocks.contains(blockId -> mapIndex) ||
+ pushBasedFetchHelper.isLocalPushMergedBlockAddress(address)
+ ) {
+ // It is a host local block or a local shuffle chunk
+ shuffleMetricsUpdate(blockId, buf, local = true)
+ } else {
+ numBlocksInFlightPerAddress(address) -= 1
+ shuffleMetricsUpdate(blockId, buf, local = false)
+ bytesInFlight -= size
+ }
+ }
+ if (isNetworkReqDone) {
+ reqsInFlight -= 1
+ resetNettyOOMFlagIfPossible(maxReqSizeShuffleToMem)
+ logDebug("Number of requests in flight " + reqsInFlight)
+ }
+
+ val in = if (buf.size == 0) {
+ // We will never legitimately receive a zero-size block. All blocks with zero records
+ // have zero size and all zero-size blocks have no records (and hence should never
+ // have been requested in the first place). This statement relies on behaviors of the
+ // shuffle writers, which are guaranteed by the following test cases:
+ //
+ // - BypassMergeSortShuffleWriterSuite: "write with some empty partitions"
+ // - UnsafeShuffleWriterSuite: "writeEmptyIterator"
+ // - DiskBlockObjectWriterSuite: "commit() and close() without ever opening or writing"
+ //
+ // There is not an explicit test for SortShuffleWriter but the underlying APIs that
+ // uses are shared by the UnsafeShuffleWriter (both writers use DiskBlockObjectWriter
+ // which returns a zero-size from commitAndGet() in case no records were written
+ // since the last call.
+ val msg = log"Received a zero-size buffer for block ${MDC(BLOCK_ID, blockId)} " +
+ log"from ${MDC(URI, address)} " +
+ log"(expectedApproxSize = ${MDC(NUM_BYTES, size)}, " +
+ log"isNetworkReqDone=${MDC(IS_NETWORK_REQUEST_DONE, isNetworkReqDone)})"
+ if (blockId.isShuffleChunk) {
+ // Zero-size block may come from nodes with hardware failures, For shuffle chunks,
+ // the original shuffle blocks that belong to that zero-size shuffle chunk is
+ // available and we can opt to fallback immediately.
+ logWarning(msg)
+ pushBasedFetchHelper.initiateFallbackFetchForPushMergedBlock(blockId, address)
+ shuffleMetrics.incCorruptMergedBlockChunks(1)
+ // Set result to null to trigger another iteration of the while loop to get either.
+ result = null
+ null
+ } else {
+ throwFetchFailedException(blockId, mapIndex, address, new IOException(msg.message))
+ }
+ } else {
+ try {
+ val bufIn = buf.createInputStream()
+ if (checksumEnabled) {
+ val checksum = ShuffleChecksumHelper.getChecksumByAlgorithm(checksumAlgorithm)
+ checkedIn = new CheckedInputStream(bufIn, checksum)
+ checkedIn
+ } else {
+ bufIn
+ }
+ } catch {
+ // The exception could only be throwed by local shuffle block
+ case e: IOException =>
+ assert(buf.isInstanceOf[FileSegmentManagedBuffer])
+ e match {
+ case ce: ClosedByInterruptException =>
+ lazy val error = MDC(ERROR, ce.getMessage)
+ logError(log"Failed to create input stream from local block, $error")
+ case e: IOException =>
+ logError("Failed to create input stream from local block", e)
+ }
+ buf.release()
+ if (blockId.isShuffleChunk) {
+ pushBasedFetchHelper.initiateFallbackFetchForPushMergedBlock(blockId, address)
+ // Set result to null to trigger another iteration of the while loop to get
+ // either.
+ result = null
+ null
+ } else {
+ throwFetchFailedException(blockId, mapIndex, address, e)
+ }
+ }
+ }
+
+ if (in != null) {
+ try {
+ input = streamWrapper(blockId, in)
+ // If the stream is compressed or wrapped, then we optionally decompress/unwrap the
+ // first maxBytesInFlight/3 bytes into memory, to check for corruption in that portion
+ // of the data. But even if 'detectCorruptUseExtraMemory' configuration is off, or if
+ // the corruption is later, we'll still detect the corruption later in the stream.
+ streamCompressedOrEncrypted = !input.eq(in)
+ if (streamCompressedOrEncrypted && detectCorruptUseExtraMemory) {
+ // TODO: manage the memory used here, and spill it into disk in case of OOM.
+ input = Utils.copyStreamUpTo(input, maxBytesInFlight / 3)
+ }
+ } catch {
+ case e: IOException =>
+ // When shuffle checksum is enabled, for a block that is corrupted twice,
+ // we'd calculate the checksum of the block by consuming the remaining data
+ // in the buf. So, we should release the buf later.
+ if (!(checksumEnabled && corruptedBlocks.contains(blockId))) {
+ buf.release()
+ }
+
+ if (blockId.isShuffleChunk) {
+ shuffleMetrics.incCorruptMergedBlockChunks(1)
+ // TODO (SPARK-36284): Add shuffle checksum support for push-based shuffle
+ // Retrying a corrupt block may result again in a corrupt block. For shuffle
+ // chunks, we opt to fallback on the original shuffle blocks that belong to that
+ // corrupt shuffle chunk immediately instead of retrying to fetch the corrupt
+ // chunk. This also makes the code simpler because the chunkMeta corresponding to
+ // a shuffle chunk is always removed from chunksMetaMap whenever a shuffle chunk
+ // gets processed. If we try to re-fetch a corrupt shuffle chunk, then it has to
+ // be added back to the chunksMetaMap.
+ pushBasedFetchHelper.initiateFallbackFetchForPushMergedBlock(blockId, address)
+ // Set result to null to trigger another iteration of the while loop.
+ result = null
+ } else if (buf.isInstanceOf[FileSegmentManagedBuffer]) {
+ throwFetchFailedException(blockId, mapIndex, address, e)
+ } else if (corruptedBlocks.contains(blockId)) {
+ // It's the second time this block is detected corrupted
+ if (checksumEnabled) {
+ // Diagnose the cause of data corruption if shuffle checksum is enabled
+ val diagnosisResponse = diagnoseCorruption(checkedIn, address, blockId)
+ buf.release()
+ logError(diagnosisResponse)
+ throwFetchFailedException(
+ blockId,
+ mapIndex,
+ address,
+ e,
+ Some(diagnosisResponse))
+ } else {
+ throwFetchFailedException(blockId, mapIndex, address, e)
+ }
+ } else {
+ // It's the first time this block is detected corrupted
+ logWarning(
+ log"got an corrupted block ${MDC(BLOCK_ID, blockId)} " +
+ log"from ${MDC(URI, address)}, fetch again",
+ e)
+ corruptedBlocks += blockId
+ fetchRequests += FetchRequest(
+ address,
+ Array(FetchBlockInfo(blockId, size, mapIndex)))
+ result = null
+ }
+ } finally {
+ if (blockId.isShuffleChunk) {
+ pushBasedFetchHelper.removeChunk(blockId.asInstanceOf[ShuffleBlockChunkId])
+ }
+ // TODO: release the buf here to free memory earlier
+ if (input == null) {
+ // Close the underlying stream if there was an issue in wrapping the stream using
+ // streamWrapper
+ in.close()
+ }
+ }
+ }
+
+ case FailureFetchResult(blockId, mapIndex, address, e) =>
+ var errorMsg: String = null
+ if (e.isInstanceOf[OutOfDirectMemoryError]) {
+ val logMessage = log"Block ${MDC(BLOCK_ID, blockId)} fetch failed after " +
+ log"${MDC(MAX_ATTEMPTS, maxAttemptsOnNettyOOM)} retries due to Netty OOM"
+ logError(logMessage)
+ errorMsg = logMessage.message
+ }
+ throwFetchFailedException(blockId, mapIndex, address, e, Some(errorMsg))
+
+ case DeferFetchRequestResult(request) =>
+ val address = request.address
+ numBlocksInFlightPerAddress(address) -= request.blocks.size
+ bytesInFlight -= request.size
+ reqsInFlight -= 1
+ logDebug("Number of requests in flight " + reqsInFlight)
+ val defReqQueue =
+ deferredFetchRequests.getOrElseUpdate(address, new Queue[FetchRequest]())
+ defReqQueue.enqueue(request)
+ result = null
+
+ case FallbackOnPushMergedFailureResult(blockId, address, size, isNetworkReqDone) =>
+ // We get this result in 3 cases:
+ // 1. Failure to fetch the data of a remote shuffle chunk. In this case, the
+ // blockId is a ShuffleBlockChunkId.
+ // 2. Failure to read the push-merged-local meta. In this case, the blockId is
+ // ShuffleBlockId.
+ // 3. Failure to get the push-merged-local directories from the external shuffle service.
+ // In this case, the blockId is ShuffleBlockId.
+ if (pushBasedFetchHelper.isRemotePushMergedBlockAddress(address)) {
+ numBlocksInFlightPerAddress(address) -= 1
+ bytesInFlight -= size
+ }
+ if (isNetworkReqDone) {
+ reqsInFlight -= 1
+ logDebug("Number of requests in flight " + reqsInFlight)
+ }
+ pushBasedFetchHelper.initiateFallbackFetchForPushMergedBlock(blockId, address)
+ // Set result to null to trigger another iteration of the while loop to get either
+ // a SuccessFetchResult or a FailureFetchResult.
+ result = null
+
+ case PushMergedLocalMetaFetchResult(
+ shuffleId,
+ shuffleMergeId,
+ reduceId,
+ bitmaps,
+ localDirs) =>
+ // Fetch push-merged-local shuffle block data as multiple shuffle chunks
+ val shuffleBlockId = ShuffleMergedBlockId(shuffleId, shuffleMergeId, reduceId)
+ try {
+ val bufs: Seq[ManagedBuffer] =
+ blockManager.getLocalMergedBlockData(shuffleBlockId, localDirs)
+ // Since the request for local block meta completed successfully, numBlocksToFetch
+ // is decremented.
+ numBlocksToFetch -= 1
+ // Update total number of blocks to fetch, reflecting the multiple local shuffle
+ // chunks.
+ numBlocksToFetch += bufs.size
+ bufs.zipWithIndex.foreach {
+ case (buf, chunkId) =>
+ buf.retain()
+ val shuffleChunkId =
+ ShuffleBlockChunkId(shuffleId, shuffleMergeId, reduceId, chunkId)
+ pushBasedFetchHelper.addChunk(shuffleChunkId, bitmaps(chunkId))
+ results.put(
+ SuccessFetchResult(
+ shuffleChunkId,
+ SHUFFLE_PUSH_MAP_ID,
+ pushBasedFetchHelper.localShuffleMergerBlockMgrId,
+ buf.size(),
+ buf,
+ isNetworkReqDone = false))
+ }
+ } catch {
+ case e: Exception =>
+ // If we see an exception with reading push-merged-local index file, we fallback
+ // to fetch the original blocks. We do not report block fetch failure
+ // and will continue with the remaining local block read.
+ logWarning(
+ "Error occurred while reading push-merged-local index, " +
+ "prepare to fetch the original blocks",
+ e)
+ pushBasedFetchHelper.initiateFallbackFetchForPushMergedBlock(
+ shuffleBlockId,
+ pushBasedFetchHelper.localShuffleMergerBlockMgrId)
+ }
+ result = null
+
+ case PushMergedRemoteMetaFetchResult(
+ shuffleId,
+ shuffleMergeId,
+ reduceId,
+ blockSize,
+ bitmaps,
+ address) =>
+ // The original meta request is processed so we decrease numBlocksToFetch and
+ // numBlocksInFlightPerAddress by 1. We will collect new shuffle chunks request and the
+ // count of this is added to numBlocksToFetch in collectFetchReqsFromMergedBlocks.
+ numBlocksInFlightPerAddress(address) -= 1
+ numBlocksToFetch -= 1
+ val blocksToFetch = pushBasedFetchHelper.createChunkBlockInfosFromMetaResponse(
+ shuffleId,
+ shuffleMergeId,
+ reduceId,
+ blockSize,
+ bitmaps)
+ val additionalRemoteReqs = new ArrayBuffer[FetchRequest]
+ collectFetchRequests(address, blocksToFetch.toSeq, additionalRemoteReqs)
+ fetchRequests ++= additionalRemoteReqs
+ // Set result to null to force another iteration.
+ result = null
+
+ case PushMergedRemoteMetaFailedFetchResult(shuffleId, shuffleMergeId, reduceId, address) =>
+ // The original meta request failed so we decrease numBlocksInFlightPerAddress by 1.
+ numBlocksInFlightPerAddress(address) -= 1
+ // If we fail to fetch the meta of a push-merged block, we fall back to fetching the
+ // original blocks.
+ pushBasedFetchHelper.initiateFallbackFetchForPushMergedBlock(
+ ShuffleMergedBlockId(shuffleId, shuffleMergeId, reduceId),
+ address)
+ // Set result to null to force another iteration.
+ result = null
+ }
+
+ // Send fetch requests up to maxBytesInFlight
+ withFetchWaitTimeTracked(fetchUpToMaxBytes())
+ }
+
+ val successResult = result.asInstanceOf[SuccessFetchResult]
+ currentResults.add(successResult)
+ (
+ successResult.blockId,
+ new GlutenBufferReleasingInputStream(
+ input,
+ this,
+ () => {
+ val result = successResult
+ if (currentResults.remove(result)) {
+ result.buf.release()
+ }
+ },
+ successResult.blockId,
+ successResult.mapIndex,
+ successResult.address,
+ detectCorrupt && streamCompressedOrEncrypted,
+ successResult.isNetworkReqDone,
+ Option(checkedIn)
+ ))
+ }
+
+ /**
+ * Get the suspect corruption cause for the corrupted block. It should be only invoked when
+ * checksum is enabled and corruption was detected at least once.
+ *
+ * This will firstly consume the rest of stream of the corrupted block to calculate the checksum
+ * of the block. Then, it will raise a synchronized RPC call along with the checksum to ask the
+ * server(where the corrupted block is fetched from) to diagnose the cause of corruption and
+ * return it.
+ *
+ * Any exception raised during the process will result in the [[Cause.UNKNOWN_ISSUE]] of the
+ * corruption cause since corruption diagnosis is only a best effort.
+ *
+ * @param checkedIn
+ * the [[CheckedInputStream]] which is used to calculate the checksum.
+ * @param address
+ * the address where the corrupted block is fetched from.
+ * @param blockId
+ * the blockId of the corrupted block.
+ * @return
+ * The corruption diagnosis response for different causes.
+ */
+ private[storage] def diagnoseCorruption(
+ checkedIn: CheckedInputStream,
+ address: BlockManagerId,
+ blockId: BlockId): String = {
+ logInfo("Start corruption diagnosis.")
+ blockId match {
+ case shuffleBlock: ShuffleBlockId =>
+ val startTimeNs = System.nanoTime()
+ val buffer = new Array[Byte](ShuffleChecksumHelper.CHECKSUM_CALCULATION_BUFFER)
+ // consume the remaining data to calculate the checksum
+ var cause: Cause = null
+ try {
+ while (checkedIn.read(buffer) != -1) {}
+ val checksum = checkedIn.getChecksum.getValue
+ cause = shuffleClient.diagnoseCorruption(
+ address.host,
+ address.port,
+ address.executorId,
+ shuffleBlock.shuffleId,
+ shuffleBlock.mapId,
+ shuffleBlock.reduceId,
+ checksum,
+ checksumAlgorithm)
+ } catch {
+ case e: Exception =>
+ logWarning("Unable to diagnose the corruption cause of the corrupted block", e)
+ cause = Cause.UNKNOWN_ISSUE
+ }
+ val duration = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTimeNs)
+ val diagnosisResponse = cause match {
+ case Cause.UNSUPPORTED_CHECKSUM_ALGORITHM =>
+ s"Block $blockId is corrupted but corruption diagnosis failed due to " +
+ s"unsupported checksum algorithm: $checksumAlgorithm"
+
+ case Cause.CHECKSUM_VERIFY_PASS =>
+ s"Block $blockId is corrupted but checksum verification passed"
+
+ case Cause.UNKNOWN_ISSUE =>
+ s"Block $blockId is corrupted but the cause is unknown"
+
+ case otherCause =>
+ s"Block $blockId is corrupted due to $otherCause"
+ }
+ logInfo(log"Finished corruption diagnosis in ${MDC(DURATION, duration)} ms. " +
+ log"${MDC(STATUS, diagnosisResponse)}")
+ diagnosisResponse
+ case shuffleBlockChunk: ShuffleBlockChunkId =>
+ // TODO SPARK-36284 Add shuffle checksum support for push-based shuffle
+ logWarning(log"BlockChunk ${MDC(SHUFFLE_BLOCK_INFO, shuffleBlockChunk)} " +
+ log"is corrupted but corruption diagnosis is skipped due to lack of shuffle " +
+ log"checksum support for push-based shuffle.")
+ s"BlockChunk $shuffleBlockChunk is corrupted but corruption " +
+ s"diagnosis is skipped due to lack of shuffle checksum support for push-based shuffle."
+ case shuffleBlockBatch: ShuffleBlockBatchId =>
+ logWarning(log"BlockBatch ${MDC(SHUFFLE_BLOCK_INFO, shuffleBlockBatch)} is corrupted " +
+ log"but corruption diagnosis is skipped due to lack of shuffle checksum support for " +
+ log"ShuffleBlockBatchId")
+ s"BlockBatch $shuffleBlockBatch is corrupted but corruption " +
+ s"diagnosis is skipped due to lack of shuffle checksum support for ShuffleBlockBatchId"
+ case unexpected: BlockId =>
+ throw SparkException.internalError(
+ s"Unexpected type of BlockId, $unexpected",
+ category = "STORAGE")
+ }
+ }
+
+ override def onComplete(): Unit = {
+ onCompleteCallback.onComplete(context)
+ }
+
+ private def fetchUpToMaxBytes(): Unit = {
+ if (isNettyOOMOnShuffle.get()) {
+ if (reqsInFlight > 0) {
+ // Return immediately if Netty is still OOMed and there're ongoing fetch requests
+ return
+ } else {
+ resetNettyOOMFlagIfPossible(0)
+ }
+ }
+
+ // Send fetch requests up to maxBytesInFlight. If you cannot fetch from a remote host
+ // immediately, defer the request until the next time it can be processed.
+
+ // Process any outstanding deferred fetch requests if possible.
+ if (deferredFetchRequests.nonEmpty) {
+ for ((remoteAddress, defReqQueue) <- deferredFetchRequests) {
+ while (
+ isRemoteBlockFetchable(defReqQueue) &&
+ !isRemoteAddressMaxedOut(remoteAddress, defReqQueue.front)
+ ) {
+ val request = defReqQueue.dequeue()
+ logDebug(
+ s"Processing deferred fetch request for $remoteAddress with "
+ + s"${request.blocks.length} blocks")
+ send(remoteAddress, request)
+ if (defReqQueue.isEmpty) {
+ deferredFetchRequests -= remoteAddress
+ }
+ }
+ }
+ }
+
+ // Process any regular fetch requests if possible.
+ while (isRemoteBlockFetchable(fetchRequests)) {
+ val request = fetchRequests.dequeue()
+ val remoteAddress = request.address
+ if (isRemoteAddressMaxedOut(remoteAddress, request)) {
+ logDebug(s"Deferring fetch request for $remoteAddress with ${request.blocks.size} blocks")
+ val defReqQueue = deferredFetchRequests.getOrElse(remoteAddress, new Queue[FetchRequest]())
+ defReqQueue.enqueue(request)
+ deferredFetchRequests(remoteAddress) = defReqQueue
+ } else {
+ send(remoteAddress, request)
+ }
+ }
+
+ def send(remoteAddress: BlockManagerId, request: FetchRequest): Unit = {
+ if (request.forMergedMetas) {
+ pushBasedFetchHelper.sendFetchMergedStatusRequest(request)
+ } else {
+ sendRequest(request)
+ }
+ numBlocksInFlightPerAddress(remoteAddress) =
+ numBlocksInFlightPerAddress.getOrElse(remoteAddress, 0) + request.blocks.size
+ }
+
+ def isRemoteBlockFetchable(fetchReqQueue: Queue[FetchRequest]): Boolean = {
+ fetchReqQueue.nonEmpty &&
+ (bytesInFlight == 0 ||
+ (reqsInFlight + 1 <= maxReqsInFlight &&
+ bytesInFlight + fetchReqQueue.front.size <= maxBytesInFlight))
+ }
+
+ // Checks if sending a new fetch request will exceed the max no. of blocks being fetched from a
+ // given remote address.
+ def isRemoteAddressMaxedOut(remoteAddress: BlockManagerId, request: FetchRequest): Boolean = {
+ numBlocksInFlightPerAddress.getOrElse(remoteAddress, 0) + request.blocks.size >
+ maxBlocksInFlightPerAddress
+ }
+ }
+
+ private[storage] def throwFetchFailedException(
+ blockId: BlockId,
+ mapIndex: Int,
+ address: BlockManagerId,
+ e: Throwable,
+ message: Option[String] = None) = {
+ val msg = message.getOrElse(e.getMessage)
+ blockId match {
+ case ShuffleBlockId(shufId, mapId, reduceId) =>
+ throw SparkCoreErrors.fetchFailedError(address, shufId, mapId, mapIndex, reduceId, msg, e)
+ case ShuffleBlockBatchId(shuffleId, mapId, startReduceId, _) =>
+ throw SparkCoreErrors.fetchFailedError(
+ address,
+ shuffleId,
+ mapId,
+ mapIndex,
+ startReduceId,
+ msg,
+ e)
+ case ShuffleBlockChunkId(shuffleId, _, reduceId, _) =>
+ throw SparkCoreErrors.fetchFailedError(
+ address,
+ shuffleId,
+ SHUFFLE_PUSH_MAP_ID.toLong,
+ SHUFFLE_PUSH_MAP_ID,
+ reduceId,
+ msg,
+ e)
+ case _ => throw SparkCoreErrors.failToGetNonShuffleBlockError(blockId, e)
+ }
+ }
+
+ /** All the below methods are used by [[PushBasedFetchHelper]] to communicate with the iterator */
+ private[storage] def addToResultsQueue(result: FetchResult): Unit = {
+ results.put(result)
+ }
+
+ private[storage] def decreaseNumBlocksToFetch(blocksFetched: Int): Unit = {
+ numBlocksToFetch -= blocksFetched
+ }
+
+ /**
+ * Currently used by [[PushBasedFetchHelper]] to fetch fallback blocks when there is a fetch
+ * failure related to a push-merged block or shuffle chunk. This is executed by the task thread
+ * when the `iterator.next()` is invoked and if that initiates fallback.
+ */
+ private[storage] def fallbackFetch(
+ originalBlocksByAddr: Iterator[(BlockManagerId, collection.Seq[(BlockId, Long, Int)])])
+ : Unit = {
+ val originalLocalBlocks = mutable.LinkedHashSet[(BlockId, Int)]()
+ val originalHostLocalBlocksByExecutor =
+ mutable.LinkedHashMap[BlockManagerId, collection.Seq[(BlockId, Long, Int)]]()
+ val originalMergedLocalBlocks = mutable.LinkedHashSet[BlockId]()
+ val originalRemoteReqs = partitionBlocksByFetchMode(
+ originalBlocksByAddr,
+ originalLocalBlocks,
+ originalHostLocalBlocksByExecutor,
+ originalMergedLocalBlocks)
+ // Add the remote requests into our queue in a random order
+ fetchRequests ++= Utils.randomize(originalRemoteReqs)
+ logInfo(log"Created ${MDC(NUM_REQUESTS, originalRemoteReqs.size)} fallback remote requests " +
+ log"for push-merged")
+ // fetch all the fallback blocks that are local.
+ fetchLocalBlocks(originalLocalBlocks)
+ // Merged local blocks should be empty during fallback
+ assert(
+ originalMergedLocalBlocks.isEmpty,
+ "There should be zero push-merged blocks during fallback")
+ // Some of the fallback local blocks could be host local blocks
+ fetchAllHostLocalBlocks(originalHostLocalBlocksByExecutor)
+ }
+
+ /**
+ * Removes all the pending shuffle chunks that are on the same host and have the same reduceId as
+ * the current chunk that had a fetch failure. This is executed by the task thread when the
+ * `iterator.next()` is invoked and if that initiates fallback.
+ *
+ * @return
+ * set of all the removed shuffle chunk Ids.
+ */
+ private[storage] def removePendingChunks(
+ failedBlockId: ShuffleBlockChunkId,
+ address: BlockManagerId): mutable.HashSet[ShuffleBlockChunkId] = {
+ val removedChunkIds = new mutable.HashSet[ShuffleBlockChunkId]()
+
+ def sameShuffleReducePartition(block: BlockId): Boolean = {
+ val chunkId = block.asInstanceOf[ShuffleBlockChunkId]
+ chunkId.shuffleId == failedBlockId.shuffleId && chunkId.reduceId == failedBlockId.reduceId
+ }
+
+ def filterRequests(queue: mutable.Queue[FetchRequest]): Unit = {
+ val fetchRequestsToRemove = new mutable.Queue[FetchRequest]()
+ fetchRequestsToRemove ++= queue.dequeueAll {
+ req =>
+ val firstBlock = req.blocks.head
+ firstBlock.blockId.isShuffleChunk && req.address.equals(address) &&
+ sameShuffleReducePartition(firstBlock.blockId)
+ }
+ fetchRequestsToRemove.foreach {
+ _ =>
+ removedChunkIds ++=
+ fetchRequestsToRemove.flatMap(_.blocks.map(_.blockId.asInstanceOf[ShuffleBlockChunkId]))
+ }
+ }
+
+ filterRequests(fetchRequests)
+ deferredFetchRequests.get(address).foreach {
+ defRequests =>
+ filterRequests(defRequests)
+ if (defRequests.isEmpty) deferredFetchRequests.remove(address)
+ }
+ removedChunkIds
+ }
+}
+
+/**
+ * Helper class that ensures a ManagedBuffer is released upon InputStream.close() and also detects
+ * stream corruption if streamCompressedOrEncrypted is true
+ */
+class GlutenBufferReleasingInputStream(
+ // This is visible for testing
+ val delegate: InputStream,
+ private val iterator: GlutenShuffleBlockFetcherIterator,
+ private val releaseCallback: () => Unit,
+ private val blockId: BlockId,
+ private val mapIndex: Int,
+ private val address: BlockManagerId,
+ private val detectCorruption: Boolean,
+ private val isNetworkReqDone: Boolean,
+ private val checkedInOpt: Option[CheckedInputStream])
+ extends InputStream {
+ private[this] var closed = false
+
+ override def read(): Int =
+ tryOrFetchFailedException(delegate.read())
+
+ override def close(): Unit = {
+ if (!closed) {
+ try {
+ delegate.close()
+ releaseCallback()
+ } finally {
+ // Unset the flag when a remote request finished and free memory is fairly enough.
+ if (isNetworkReqDone) {
+ ShuffleBlockFetcherIterator.resetNettyOOMFlagIfPossible(
+ iterator.maxReqSizeShuffleToMem)
+ }
+ closed = true
+ }
+ }
+ }
+
+ override def available(): Int =
+ tryOrFetchFailedException(delegate.available())
+
+ override def mark(readlimit: Int): Unit = delegate.mark(readlimit)
+
+ override def skip(n: Long): Long =
+ tryOrFetchFailedException(delegate.skip(n))
+
+ override def markSupported(): Boolean = delegate.markSupported()
+
+ override def read(b: Array[Byte]): Int =
+ tryOrFetchFailedException(delegate.read(b))
+
+ override def read(b: Array[Byte], off: Int, len: Int): Int =
+ tryOrFetchFailedException(delegate.read(b, off, len))
+
+ override def reset(): Unit = tryOrFetchFailedException(delegate.reset())
+
+ /**
+ * Execute a block of code that returns a value, close this stream quietly and re-throwing
+ * IOException as FetchFailedException when detectCorruption is true. This method is only used by
+ * the `available`, `read` and `skip` methods inside `BufferReleasingInputStream` currently.
+ */
+ private def tryOrFetchFailedException[T](block: => T): T = {
+ try {
+ block
+ } catch {
+ case e: IOException if detectCorruption =>
+ val diagnosisResponse = checkedInOpt.map {
+ checkedIn => iterator.diagnoseCorruption(checkedIn, address, blockId)
+ }
+ Utils.closeQuietly(this)
+ // We'd never retry the block whatever the cause is since the block has been
+ // partially consumed by downstream RDDs.
+ iterator.throwFetchFailedException(blockId, mapIndex, address, e, diagnosisResponse)
+ }
+ }
+}
+
+/**
+ * A listener to be called at the completion of the ShuffleBlockFetcherIterator
+ * @param data
+ * the ShuffleBlockFetcherIterator to process
+ */
+private class GlutenShuffleFetchCompletionListener(var data: GlutenShuffleBlockFetcherIterator)
+ extends TaskCompletionListener {
+
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ if (data != null) {
+ data.cleanup()
+ // Null out the referent here to make sure we don't keep a reference to this
+ // ShuffleBlockFetcherIterator, after we're done reading from it, to let it be
+ // collected during GC. Otherwise we can hold metadata on block locations(blocksByAddress)
+ data = null
+ }
+ }
+
+ // Just an alias for onTaskCompletion to avoid confusing
+ def onComplete(context: TaskContext): Unit = this.onTaskCompletion(context)
+}
diff --git a/shims/spark42/src/test/scala/org/apache/gluten/sql/shims/spark42/Spark42LocalTableScanStreamSuite.scala b/shims/spark42/src/test/scala/org/apache/gluten/sql/shims/spark42/Spark42LocalTableScanStreamSuite.scala
new file mode 100644
index 0000000000..8edebb618a
--- /dev/null
+++ b/shims/spark42/src/test/scala/org/apache/gluten/sql/shims/spark42/Spark42LocalTableScanStreamSuite.scala
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.sql.shims.spark42
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.connector.read.streaming.{Offset, SparkDataStream}
+import org.apache.spark.sql.execution.LocalTableScanExec
+import org.apache.spark.sql.types.IntegerType
+
+import org.scalatest.funsuite.AnyFunSuite
+
+/**
+ * Spark 4.1+ only: LocalTableScanExec carries an optional streaming source (`stream`). This suite
+ * verifies the Spark42Shims accessor that VeloxSparkPlanExecApi.isSupportLocalTableScanExec uses to
+ * skip offload for streaming sources. The Velox backend test lives here (not in the shared
+ * cross-version suite) because the `stream` constructor parameter does not exist on Spark 3.x.
+ */
+class Spark42LocalTableScanStreamSuite extends AnyFunSuite {
+
+ private val shims = new Spark42Shims
+
+ private def output = Seq(AttributeReference("id", IntegerType)())
+
+ private val stubStream: SparkDataStream = new SparkDataStream {
+ override def initialOffset(): Offset = null
+ override def deserializeOffset(json: String): Offset = null
+ override def commit(end: Offset): Unit = {}
+ override def stop(): Unit = {}
+ }
+
+ test("getLocalTableScanStream returns the stream for a streaming LocalTableScanExec") {
+ val plan = LocalTableScanExec(output, Seq.empty[InternalRow], Some(stubStream))
+ val stream = shims.getLocalTableScanStream(plan)
+ assert(stream.isDefined, "Streaming LocalTableScanExec must be detected as a streaming source")
+ assert(stream.get eq stubStream)
+ }
+
+ test("getLocalTableScanStream returns None for a batch LocalTableScanExec") {
+ val plan = LocalTableScanExec(output, Seq.empty[InternalRow], None)
+ assert(
+ shims.getLocalTableScanStream(plan).isEmpty,
+ "Batch LocalTableScanExec must not be detected as a streaming source")
+ }
+}
diff --git a/tools/gluten-it/pom.xml b/tools/gluten-it/pom.xml
index 637e1fd3f2..825028a8d7 100644
--- a/tools/gluten-it/pom.xml
+++ b/tools/gluten-it/pom.xml
@@ -327,6 +327,18 @@
<delta.version>4.1.0</delta.version>
</properties>
</profile>
+ <profile>
+ <id>spark-4.2</id>
+ <properties>
+ <spark.version>4.2.0</spark.version>
+ <scala.version>2.13.18</scala.version>
+ <scala.binary.version>2.13</scala.binary.version>
+ <!-- No Delta build for Spark 4.2 yet; placeholder inherited from spark-4.1.
+ The delta profile is opt-in and is not enabled on the Spark 4.2 lane. -->
+ <delta.package.name>delta-spark_4.1</delta.package.name>
+ <delta.version>4.1.0</delta.version>
+ </properties>
+ </profile>
<profile>
<id>celeborn-0.5</id>
<properties>
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org