Create a shim for Spark 3.2.0 development

api_validation/pom.xml

@@ -46,6 +46,12 @@
                 <spark.version>${spark311.version}</spark.version>
             </properties>
         </profile>
+        <profile>
+            <id>spark320</id>
+            <properties>
+                <spark.version>${spark320.version}</spark.version>
+            </properties>
+        </profile>
     </profiles>
 
     <dependencies>

docs/configs.md

@@ -36,7 +36,7 @@ Name | Description | Default Value
 <a name="memory.gpu.direct.storage.spill.enabled"></a>spark.rapids.memory.gpu.direct.storage.spill.enabled|Should GPUDirect Storage (GDS) be used to spill GPU memory buffers directly to disk. GDS must be enabled and the directory `spark.local.dir` must support GDS. This is an experimental feature. For more information on GDS, see https://docs.nvidia.com/gpudirect-storage/.|false
 <a name="memory.gpu.maxAllocFraction"></a>spark.rapids.memory.gpu.maxAllocFraction|The fraction of total GPU memory that limits the maximum size of the RMM pool. The value must be greater than or equal to the setting for spark.rapids.memory.gpu.allocFraction. Note that this limit will be reduced by the reserve memory configured in spark.rapids.memory.gpu.reserve.|1.0
 <a name="memory.gpu.oomDumpDir"></a>spark.rapids.memory.gpu.oomDumpDir|The path to a local directory where a heap dump will be created if the GPU encounters an unrecoverable out-of-memory (OOM) error. The filename will be of the form: "gpu-oom-<pid>.hprof" where <pid> is the process ID.|None
-<a name="memory.gpu.pool"></a>spark.rapids.memory.gpu.pool|Select the RMM pooling allocator to use. Valid values are "DEFAULT", "ARENA", and "NONE". With "DEFAULT", `rmm::mr::pool_memory_resource` is used; with "ARENA", `rmm::mr::arena_memory_resource` is used. If set to "NONE", pooling is disabled and RMM just passes through to CUDA memory allocation directly. Note: "ARENA" is the recommended pool allocator if CUDF is built with Per-thread Default Stream (PTDS), as "DEFAULT" is known to be unstable (https://github.com/NVIDIA/spark-rapids/issues/1141)|ARENA
+<a name="memory.gpu.pool"></a>spark.rapids.memory.gpu.pool|Select the RMM pooling allocator to use. Valid values are "DEFAULT", "ARENA", and "NONE". With "DEFAULT", `rmm::mr::pool_memory_resource` is used; with "ARENA", `rmm::mr::arena_memory_resource` is used. If set to "NONE", pooling is disabled and RMM just passes through to CUDA memory allocation directly. Note: "ARENA" is the recommended pool allocator if CUDF is built with Per-Thread Default Stream (PTDS), as "DEFAULT" is known to be unstable (https://github.com/NVIDIA/spark-rapids/issues/1141)|ARENA
 <a name="memory.gpu.pooling.enabled"></a>spark.rapids.memory.gpu.pooling.enabled|Should RMM act as a pooling allocator for GPU memory, or should it just pass through to CUDA memory allocation directly. DEPRECATED: please use spark.rapids.memory.gpu.pool instead.|true
 <a name="memory.gpu.reserve"></a>spark.rapids.memory.gpu.reserve|The amount of GPU memory that should remain unallocated by RMM and left for system use such as memory needed for kernels, kernel launches or JIT compilation.|1073741824
 <a name="memory.host.spillStorageSize"></a>spark.rapids.memory.host.spillStorageSize|Amount of off-heap host memory to use for buffering spilled GPU data before spilling to local disk|1073741824

docs/get-started/getting-started-on-prem.md

@@ -435,6 +435,7 @@ that matches your Spark version. Currently we support
    - Spark 3.0.1 (com.nvidia.spark.rapids.spark301.RapidsShuffleManager) 
    - Spark 3.0.2 (com.nvidia.spark.rapids.spark302.RapidsShuffleManager) 
    - Spark 3.1.1 (com.nvidia.spark.rapids.spark311.RapidsShuffleManager)
+   - Spark 3.2.0 (com.nvidia.spark.rapids.spark320.RapidsShuffleManager)
 
 ```shell
 ...

integration_tests/pom.xml

@@ -29,31 +29,37 @@
     <version>0.4.0-SNAPSHOT</version>
 
     <properties>
-        <spark.test.version>3.0.0</spark.test.version>
+        <spark.test.version>${spark300.version}</spark.test.version>
     </properties>
     <profiles>
         <profile>
             <id>spark301dbtests</id>
             <properties>
-                <spark.test.version>3.0.1-databricks</spark.test.version>
+                <spark.test.version>${spark301db.version}</spark.test.version>
             </properties>
         </profile>
         <profile>
             <id>spark301tests</id>
             <properties>
-                <spark.test.version>3.0.1</spark.test.version>
+                <spark.test.version>${spark301.version}</spark.test.version>
             </properties>
         </profile>
         <profile>
             <id>spark302tests</id>
             <properties>
-                <spark.test.version>3.0.2-SNAPSHOT</spark.test.version>
+                <spark.test.version>${spark302.version}</spark.test.version>
             </properties>
         </profile>
         <profile>
             <id>spark311tests</id>
             <properties>
-                <spark.test.version>3.1.1-SNAPSHOT</spark.test.version>
+                <spark.test.version>${spark311.version}</spark.test.version>
+            </properties>
+        </profile>
+        <profile>
+            <id>spark320tests</id>
+            <properties>
+                <spark.test.version>${spark320.version}</spark.test.version>
             </properties>
         </profile>
     </profiles>

pom.xml

@@ -140,12 +140,18 @@
                 <module>tests-spark310+</module>
             </modules>
         </profile>
+        <profile>
+            <id>spark320tests</id>
+            <modules>
+                <module>tests-spark310+</module>
+            </modules>
+        </profile>
     </profiles>
 
     <properties>
         <maven.compiler.source>1.8</maven.compiler.source>
         <maven.compiler.target>1.8</maven.compiler.target>
-        <spark.version>3.0.0</spark.version>
+        <spark.version>${spark300.version}</spark.version>
         <cuda.version>cuda10-1</cuda.version>
         <cudf.version>0.18-SNAPSHOT</cudf.version>
         <scala.binary.version>2.12</scala.binary.version>
@@ -174,8 +180,10 @@
          If you update a dependendy version so it is no longer a SNAPSHOT
          please update the snapshot-shims profile as well so it is accurate -->
         <spark301.version>3.0.1</spark301.version>
+        <spark301db.version>3.0.1-databricks</spark301db.version>
         <spark302.version>3.0.2-SNAPSHOT</spark302.version>
         <spark311.version>3.1.1-SNAPSHOT</spark311.version>
+        <spark320.version>3.2.0-SNAPSHOT</spark320.version>
         <mockito.version>3.6.0</mockito.version>
         <scala.plugin.version>4.3.0</scala.plugin.version>
         <maven.jar.plugin.version>3.2.0</maven.jar.plugin.version>

shims/aggregator/pom.xml

@@ -62,6 +62,12 @@
                 <activeByDefault>true</activeByDefault>
             </activation>
             <dependencies>
+                <dependency>
+                    <groupId>com.nvidia</groupId>
+                    <artifactId>rapids-4-spark-shims-spark320_${scala.binary.version}</artifactId>
+                    <version>${project.version}</version>
+                    <scope>compile</scope>
+                </dependency>
                 <dependency>
                     <groupId>com.nvidia</groupId>
                     <artifactId>rapids-4-spark-shims-spark311_${scala.binary.version}</artifactId>

shims/pom.xml

@@ -47,6 +47,7 @@
             <modules>
                 <module>spark302</module>
                 <module>spark311</module>
+                <module>spark320</module>
             </modules>
         </profile>
     </profiles>

shims/spark300/src/main/scala/com/nvidia/spark/rapids/shims/spark300/Spark300Shims.scala

@@ -44,7 +44,7 @@ import org.apache.spark.sql.execution.datasources.{FileIndex, FilePartition, Fil
 import org.apache.spark.sql.execution.datasources.rapids.GpuPartitioningUtils
 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.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, SortMergeJoinExec}
 import org.apache.spark.sql.execution.joins.ShuffledHashJoinExec
 import org.apache.spark.sql.execution.python.WindowInPandasExec
@@ -580,4 +580,9 @@ class Spark300Shims extends SparkShims {
     }
     recurse(plan, predicate, new ListBuffer[SparkPlan]())
   }
+
+  override def reusedExchangeExecPfn: PartialFunction[SparkPlan, ReusedExchangeExec] = {
+    case ShuffleQueryStageExec(_, e: ReusedExchangeExec) => e
+    case BroadcastQueryStageExec(_, e: ReusedExchangeExec) => e
+  }
 }

shims/spark301db/pom.xml

@@ -73,7 +73,6 @@
 
     <properties>
         <parquet.version>1.10.1</parquet.version>
-        <spark301db.version>3.0.1-databricks</spark301db.version>
         <arrow.version>0.15.1</arrow.version>
     </properties>
 

shims/spark311/src/main/scala/com/nvidia/spark/rapids/shims/spark311/SparkShimServiceProvider.scala

@@ -22,7 +22,8 @@ object SparkShimServiceProvider {
   // keep 3.1.0 snapshot version for now until 3.1.1 is released
   val VERSION310 = SparkShimVersion(3, 1, 0)
   val VERSION = SparkShimVersion(3, 1, 1)
-  val VERSIONNAMES = Seq(s"$VERSION310-SNAPSHOT", s"$VERSION", s"$VERSION-SNAPSHOT")
+  val VERSIONNAMES: Seq[String] = Seq(VERSION310, VERSION)
+    .flatMap(v => Seq(s"$v", s"$v-SNAPSHOT"))
 }
 
 class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {

shims/spark320/pom.xml

@@ -0,0 +1,91 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  Copyright (c) 2021, NVIDIA CORPORATION.
+
+  Licensed 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>com.nvidia</groupId>
+        <artifactId>rapids-4-spark-shims_2.12</artifactId>
+        <version>0.4.0-SNAPSHOT</version>
+	<relativePath>../pom.xml</relativePath>
+    </parent>
+    <groupId>com.nvidia</groupId>
+    <artifactId>rapids-4-spark-shims-spark320_2.12</artifactId>
+    <name>RAPIDS Accelerator for Apache Spark SQL Plugin Spark 3.2.0 Shim</name>
+    <description>The RAPIDS SQL plugin for Apache Spark 3.2.0 Shim</description>
+    <version>0.4.0-SNAPSHOT</version>
+
+    <!-- Set 'spark.version' for the shims layer -->
+    <!-- Create a separate file 'SPARK_VER.properties' in the jar to save cudf & spark version info -->
+    <build>
+        <plugins>
+            <plugin>
+                <artifactId>maven-antrun-plugin</artifactId>
+                <executions>
+                    <execution>
+                        <id>dependency</id>
+                        <phase>generate-resources</phase>
+                        <configuration>
+                            <target>
+                                <mkdir dir="${project.build.directory}/extra-resources"/>
+                                <exec executable="bash" failonerror="true" output="${project.build.directory}/extra-resources/spark-${spark320.version}-info.properties">
+                                    <arg value="${user.dir}/build/dependency-info.sh"/>
+                                    <arg value="${cudf.version}"/>
+                                    <arg value="${cuda.version}"/>
+                                    <arg value="${spark320.version}"/>
+                                </exec>
+                            </target>
+                        </configuration>
+                        <goals>
+                            <goal>run</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
+            <plugin>
+                <groupId>org.scalastyle</groupId>
+                <artifactId>scalastyle-maven-plugin</artifactId>
+            </plugin>
+        </plugins>
+
+        <resources>
+          <resource>
+            <!-- Include the properties file to provide the build information. -->
+            <directory>${project.build.directory}/extra-resources</directory>
+          </resource>
+          <resource>
+            <directory>src/main/resources</directory>
+          </resource>
+        </resources>
+    </build>
+
+    <dependencies>
+	<dependency>
+            <groupId>com.nvidia</groupId>
+            <artifactId>rapids-4-spark-shims-spark311_${scala.binary.version}</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.spark</groupId>
+            <artifactId>spark-sql_${scala.binary.version}</artifactId>
+            <version>${spark320.version}</version>
+	    <scope>provided</scope>
+        </dependency>
+    </dependencies>
+</project>

shims/spark320/src/main/resources/META-INF/services/com.nvidia.spark.rapids.SparkShimServiceProvider

@@ -0,0 +1 @@
+com.nvidia.spark.rapids.shims.spark320.SparkShimServiceProvider

shims/spark320/src/main/scala/com/nvidia/spark/rapids/shims/spark320/Spark320Shims.scala

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION.
+ *
+ * Licensed 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 com.nvidia.spark.rapids.shims.spark320
+
+import com.nvidia.spark.rapids.shims.spark311.Spark311Shims
+
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+
+class Spark320Shims extends Spark311Shims {
+  override def reusedExchangeExecPfn: PartialFunction[SparkPlan, ReusedExchangeExec] = {
+    case ShuffleQueryStageExec(_, e: ReusedExchangeExec, _) => e
+    case BroadcastQueryStageExec(_, e: ReusedExchangeExec, _) => e
+  }
+}

shims/spark320/src/main/scala/com/nvidia/spark/rapids/shims/spark320/SparkShimServiceProvider.scala

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION.
+ *
+ * Licensed 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 com.nvidia.spark.rapids.shims.spark320
+
+import com.nvidia.spark.rapids.{SparkShims, SparkShimVersion}
+
+object SparkShimServiceProvider {
+  val VERSION320 = SparkShimVersion(3, 2, 0)
+  val VERSIONNAMES: Seq[String] = Seq(VERSION320)
+    .flatMap(v => Seq(s"$v", s"$v-SNAPSHOT"))
+}
+
+class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
+
+  def matchesVersion(version: String): Boolean = {
+    SparkShimServiceProvider.VERSIONNAMES.contains(version)
+  }
+
+  def buildShim: SparkShims = {
+    new Spark320Shims()
+  }
+}

shims/spark320/src/main/scala/com/nvidia/spark/rapids/spark320/RapidsShuffleManager.scala

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ *
+ * Licensed 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 com.nvidia.spark.rapids.spark320
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.rapids.shims.spark311.RapidsShuffleInternalManager
+
+/** A shuffle manager optimized for the RAPIDS Plugin for Apache Spark. */
+sealed class RapidsShuffleManager(
+    conf: SparkConf,
+    isDriver: Boolean) extends RapidsShuffleInternalManager(conf, isDriver) {
+}

sql-plugin/src/main/scala/com/nvidia/spark/rapids/SparkShims.scala

@@ -34,7 +34,7 @@ import org.apache.spark.sql.connector.read.Scan
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec
 import org.apache.spark.sql.execution.datasources.{FileIndex, FilePartition, HadoopFsRelation, PartitionDirectory, PartitionedFile}
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins._
 import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, ShuffleManagerShimBase}
 import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExecBase, GpuBroadcastNestedLoopJoinExecBase, GpuShuffleExchangeExecBase}
@@ -188,4 +188,6 @@ trait SparkShims {
   def shouldFailDivByZero(): Boolean
 
   def findOperators(plan: SparkPlan, predicate: SparkPlan => Boolean): Seq[SparkPlan]
+
+  def reusedExchangeExecPfn: PartialFunction[SparkPlan, ReusedExchangeExec]
 }

tests-spark310+/pom.xml

@@ -29,7 +29,7 @@
     <version>0.4.0-SNAPSHOT</version>
 
     <properties>
-        <spark.test.version>3.1.0-SNAPSHOT</spark.test.version>
+        <spark.test.version>${spark311.version}</spark.test.version>
     </properties>
 
     <dependencies>