[beam] branch spark-runner_structured-streaming updated: Use "sparkMaster" in local mode to obtain number of shuffle partitions + spotless apply

[ar...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

aromanenko pushed a commit to branch spark-runner_structured-streaming
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/spark-runner_structured-streaming by this push:
     new b63892e  Use "sparkMaster" in local mode to obtain number of shuffle partitions + spotless apply
b63892e is described below

commit b63892ee1f5c76594716d276845ed89816c783bc
Author: Alexey Romanenko <ar...@gmail.com>
AuthorDate: Fri Jul 19 15:48:32 2019 +0200

    Use "sparkMaster" in local mode to obtain number of shuffle partitions
    + spotless apply
---
 .../structuredstreaming/translation/TranslationContext.java   | 11 +++++++++++
 .../translation/batch/ParDoTranslatorBatch.java               |  8 +++++---
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TranslationContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TranslationContext.java
index 9f2d53c..fcea789 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TranslationContext.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TranslationContext.java
@@ -78,6 +78,17 @@ public class TranslationContext {
       sparkConf.setJars(options.getFilesToStage().toArray(new String[0]));
     }
 
+    // By default, Spark defines 200 as a number of sql partitions. This seems too much for local
+    // mode, so try to align with value of "sparkMaster" option in this case.
+    String sparkMaster = options.getSparkMaster();
+    if (sparkMaster != null && sparkMaster.startsWith("local[")) {
+      int numPartitions =
+          Integer.valueOf(sparkMaster.substring("local[".length(), sparkMaster.length() - 1));
+      if (numPartitions > 0) {
+        sparkConf.set("spark.sql.shuffle.partitions", "4");
+      }
+    }
+
     this.sparkSession = SparkSession.builder().config(sparkConf).getOrCreate();
     this.serializablePipelineOptions = new SerializablePipelineOptions(options);
     this.datasets = new HashMap<>();
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java
index 5da0acd..827d77a 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java
@@ -137,9 +137,11 @@ class ParDoTranslatorBatch<InputT, OutputT>
         pruneOutputFilteredByTag(context, allOutputs, output);
       }
     } else {
-      Dataset<WindowedValue<?>> outputDataset = allOutputs.map(
-          (MapFunction<Tuple2<TupleTag<?>, WindowedValue<?>>, WindowedValue<?>>) value -> value._2,
-          EncoderHelpers.windowedValueEncoder());
+      Dataset<WindowedValue<?>> outputDataset =
+          allOutputs.map(
+              (MapFunction<Tuple2<TupleTag<?>, WindowedValue<?>>, WindowedValue<?>>)
+                  value -> value._2,
+              EncoderHelpers.windowedValueEncoder());
       context.putDatasetWildcard(outputs.entrySet().iterator().next().getValue(), outputDataset);
     }
   }