title: Low parallelism when running Apache Beam wordcount pipeline on Spark with Python SDK
body: <p>I am quite experienced with Spark cluster configuration and running Pyspark pipelines, but I'm just starting with Beam. So, I am trying to do an apple-to-apple comparison between Pyspark and the Beam python SDK on a Spark PortableRunner (running on top of the same small Spark cluster, 4 workers each with 4 cores and 8GB RAM), and I've settled on a wordcount job for a reasonably large dataset, storing the results in a Parquet table.</p>
  <p>I have thus downloaded 50GB of Wikipedia text files, splitted across about 100 uncompressed files, and stored them in the directory <code>/mnt/nfs_drive/wiki_files/</code> (<code>/mnt/nfs_drive</code> is a NFS drive mounted on all workers).</p>
  <p>First, I am running the following Pyspark wordcount script:</p>
  <pre><code>from pyspark.sql import SparkSession, Row
  from operator import add
  wiki_files = '/mnt/nfs_drive/wiki_files/*'
  
  spark = SparkSession.builder.appName(&quot;WordCountSpark&quot;).getOrCreate()
  
  spark_counts = spark.read.text(wiki_files).rdd.map(lambda r: r['value']) \
      .flatMap(lambda x: x.split(' ')) \
      .map(lambda x: (x, 1)) \
      .reduceByKey(add) \
      .map(lambda x: Row(word=x[0], count=x[1]))
  
  spark.createDataFrame(spark_counts).write.parquet(path='/mnt/nfs_drive/spark_output', mode='overwrite')
  </code></pre>
  <p>The script runs perfectly well and outputs the parquet files at the desired location in about 8 minutes. The main stage (reading and splitting tokens) is divided in a reasonable number of tasks, so that the cluster is used efficiently:
  <a href="https://i.stack.imgur.com/SKAiT.png" rel="noreferrer"><img src="https://i.stack.imgur.com/SKAiT.png" alt="enter image description here" /></a></p>
  <p>I am now trying to achieve the same with Beam and the portable runner. First, I have started the Spark job server (on the Spark master node) with the following command:</p>
  <pre><code>docker run --rm --net=host -e SPARK_EXECUTOR_MEMORY=8g apache/beam_spark_job_server:2.25.0 --spark-master-url=spark://localhost:7077
  </code></pre>
  <p>Then, on the master and worker nodes, I am running the SDK Harness as follows:</p>
  <pre><code> docker run --net=host -d --rm -v /mnt/nfs_drive:/mnt/nfs_drive apache/beam_python3.6_sdk:2.25.0 --worker_pool
  </code></pre>
  <p>Now that the Spark cluster is set up to run Beam pipelines, I can submit the following script:</p>
  <pre><code>import apache_beam as beam
  import pyarrow
  from apache_beam.options.pipeline_options import PipelineOptions
  from apache_beam.io import fileio
  
  options = PipelineOptions([
      &quot;--runner=PortableRunner&quot;,
      &quot;--job_endpoint=localhost:8099&quot;,
      &quot;--environment_type=EXTERNAL&quot;,
      &quot;--environment_config=localhost:50000&quot;,
      &quot;--job_name=WordCountBeam&quot;
  ])
  
  
  wiki_files = '/mnt/nfs_drive/wiki_files/*'
  
  p = beam.Pipeline(options=options)
  beam_counts = (
      p
      | fileio.MatchFiles(wiki_files)
      | beam.Map(lambda x: x.path)
      | beam.io.ReadAllFromText()
      | 'ExtractWords' &gt;&gt; beam.FlatMap(lambda x: x.split(' '))
      | beam.combiners.Count.PerElement()
      | beam.Map(lambda x: {'word': x[0], 'count': x[1]})
  )
  
  
  _ = beam_counts | 'Write' &gt;&gt; beam.io.WriteToParquet('/mnt/nfs_drive/beam_output',
        pyarrow.schema(
            [('word', pyarrow.binary()), ('count', pyarrow.int64())]
        )
  )
  
  result = p.run().wait_until_finish()
  </code></pre>
  <p>The code is submitted successfully, I can see the job on the Spark UI and the workers are executing it. However, even if left running for more than 1 hour, it does not produce any output!</p>
  <p>I thus wanted to make sure that there is not a problem with my setup, so I've run the exact same script on a smaller dataset (just 1 Wiki file). This completes successfully in about 3.5 minutes (Spark wordcount on the same dataset takes 16s!).</p>
  <p>I wondered how could Beam be <em>that</em> slower, so I started looking at the DAG submitted to Spark by the Beam pipeline via the job server. I noticed that the Spark job spends most of the time in the following stage:
  <a href="https://i.stack.imgur.com/kSK6X.png" rel="noreferrer"><img src="https://i.stack.imgur.com/kSK6X.png" alt="enter image description here" /></a></p>
  <p>This is just splitted in 2 tasks, as shown here:
  <a href="https://i.stack.imgur.com/OIb6J.png" rel="noreferrer"><img src="https://i.stack.imgur.com/OIb6J.png" alt="enter image description here" /></a></p>
  <p>Printing debugging lines show that this task is where the &quot;heavy-lifting&quot; (i.e. reading lines from the wiki files and splitting tokens) is performed - however, since this happens in 2 tasks only, the work will be distributed on 2 workers at most. What's also interesting is that running on the large 50GB dataset results on <em>exactly</em> the same DAG with <em>exactly</em> the same number of tasks.</p>
  <p>I am quite unsure how to proceed further. It seems like the Beam pipeline has reduced parallelism, but I'm not sure if this is due to sub-optimal translation of the pipeline by the job server, or whether I should specify my PTransforms in some other way to increase the parallelism on Spark.</p>
  <p>Any suggestion appreciated!</p>