Add leafNodeDefaultParallelism support

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

@@ -252,6 +252,8 @@ trait SparkShims {
   def isCustomReaderExec(x: SparkPlan): Boolean
 
   def aqeShuffleReaderExec: ExecRule[_ <: SparkPlan]
+
+  def leafNodeDefaultParallelism(ss: SparkSession): Int
 }
 
 abstract class SparkCommonShims extends SparkShims {

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

@@ -378,7 +378,8 @@ case class GpuRangeExec(range: org.apache.spark.sql.catalyst.plans.logical.Range
   val start: Long = range.start
   val end: Long = range.end
   val step: Long = range.step
-  val numSlices: Int = range.numSlices.getOrElse(sparkContext.defaultParallelism)
+  val numSlices: Int = range.numSlices.getOrElse(ShimLoader.getSparkShims
+    .leafNodeDefaultParallelism(sparkSession))
   val numElements: BigInt = range.numElements
   val isEmptyRange: Boolean = start == end || (start < end ^ 0 < step)
 

sql-plugin/src/main/spark30+all/scala/com/nvidia/spark/rapids/shims/v2/Spark30XShims.scala

@@ -96,4 +96,9 @@ trait Spark30XShims extends SparkShims {
     ExecChecks((TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL_64 + TypeSig.ARRAY +
         TypeSig.STRUCT + TypeSig.MAP).nested(), TypeSig.all),
     (exec, conf, p, r) => new GpuCustomShuffleReaderMeta(exec, conf, p, r))
+
+  override def leafNodeDefaultParallelism(ss: SparkSession): Int = {
+    ss.sparkContext.defaultParallelism
+  }
+
 }

sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/v2/Spark32XShims.scala

@@ -19,10 +19,10 @@ package com.nvidia.spark.rapids.shims.v2
 import com.nvidia.spark.rapids._
 import com.nvidia.spark.rapids.GpuOverrides.exec
 import com.nvidia.spark.rapids.shims._
-
 import org.apache.hadoop.fs.FileStatus
 import org.apache.parquet.schema.MessageType
 
+import org.apache.spark.sql.Spark32XShimsUtils
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier}
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
@@ -131,6 +131,11 @@ trait Spark32XShims extends SparkShims {
       // we will need to change the API to pass these values in.
       enableAddPartitions = true,
       enableDropPartitions = false)
+
+  override def leafNodeDefaultParallelism(ss: SparkSession): Int = {
+    Spark32XShimsUtils.leafNodeDefaultParallelism(ss)
+  }
+
 }
 
 // TODO dedupe utils inside shims

sql-plugin/src/main/spark320/scala/org/apache/spark/sql/Spark32XShimsUtils.scala

@@ -0,0 +1,26 @@
+/*
+ * 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 org.apache.spark.sql
+
+object Spark32XShimsUtils {
+
+  def leafNodeDefaultParallelism(ss: SparkSession): Int = {
+    ss.leafNodeDefaultParallelism
+  }
+
+}
+

tests/src/test/scala/org/apache/spark/sql/GpuSparkPlanSuite.scala

@@ -0,0 +1,45 @@
+/*
+ * 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 org.apache.spark.sql
+
+import java.util.{Locale, TimeZone}
+
+import com.nvidia.spark.rapids.{ShimLoader, SparkSessionHolder}
+import org.scalatest.FunSuite
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.catalyst.plans.logical.Range
+
+class GpuSparkPlanSuite extends FunSuite {
+
+  test("leafNodeDefaultParallelism for GpuRangeExec") {
+
+    val conf = new SparkConf()
+      .set("spark.sql.leafNodeDefaultParallelism", "7")
+      .set("spark.rapids.sql.enabled", "true")
+
+    SparkSessionHolder.withSparkSession(conf, spark => {
+      val defaultSlice = ShimLoader.getSparkShims.leafNodeDefaultParallelism(spark)
+      val ds = new Dataset(spark, Range(0, 20, 1, None), Encoders.LONG)
+      val partitions = ds.rdd.getNumPartitions
+      assert(partitions == defaultSlice)
+    })
+
+  }
+
+}
+