Spark 3.4: Take shuffle partitions into account for parallelism (#8327)

apache · Aug 15, 2023 · bfa0529 · bfa0529
1 parent 86737e6
commit bfa0529
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java
@@ -275,4 +275,10 @@ public boolean adaptiveSplitSizeEnabled() {
         .defaultValue(TableProperties.ADAPTIVE_SPLIT_SIZE_ENABLED_DEFAULT)
         .parse();
   }
+
+  public int parallelism() {
+    int defaultParallelism = spark.sparkContext().defaultParallelism();
+    int numShufflePartitions = spark.sessionState().conf().numShufflePartitions();
+    return Math.max(defaultParallelism, numShufflePartitions);
+  }
 }
diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkScan.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkScan.java
@@ -227,7 +227,7 @@ public CustomMetric[] supportedCustomMetrics() {
   protected long adjustSplitSize(List<? extends ScanTask> tasks, long splitSize) {
     if (readConf.splitSizeOption() == null && readConf.adaptiveSplitSizeEnabled()) {
       long scanSize = tasks.stream().mapToLong(ScanTask::sizeBytes).sum();
-      int parallelism = sparkContext.defaultParallelism();
+      int parallelism = readConf.parallelism();
       return TableScanUtil.adjustSplitSize(scanSize, parallelism, splitSize);
     } else {
       return splitSize;