PythonSubmit

id: spark_python_submit
namespace: company.team

tasks:
  - id: python_submit
    type: io.kestra.plugin.spark.PythonSubmit
    taskRunner:
      type: io.kestra.plugin.scripts.runner.docker.Docker
      networkMode: host
      user: root
    master: spark://localhost:7077
    args:
      - "10"
    mainScript: |
      import sys
      from random import random
      from operator import add
      from pyspark.sql import SparkSession


      if __name__ == "__main__":
          spark = SparkSession.builder.appName("PythonPi").getOrCreate()

          partitions = int(sys.argv[1]) if len(sys.argv) > 1 else 2
          n = 100000 * partitions

          def f(_: int) -> float:
              x = random() * 2 - 1
              y = random() * 2 - 1
              return 1 if x ** 2 + y ** 2 <= 1 else 0

          count = spark.sparkContext.parallelize(range(1, n + 1), partitions).map(f).reduce(add)
          print("Pi is roughly %f" % (4.0 * count / n))

          spark.stop()