def get_spark_context(spark_conf):
"""
outputs: SparkContext
Returns appropriate SparkContext based on whether invoked via a Notebook or a Flyte workflow.
"""
# We run in cluster-mode in Flyte.
# Ref https://github.com/lyft/flyteplugins/blob/master/go/tasks/v1/flytek8s/k8s_resource_adds.go#L46
if "FLYTE_INTERNAL_EXECUTION_ID" in _os.environ:
return _pyspark.SparkContext()
# Add system spark-conf for local/notebook based execution.
spark_conf.add(("spark.master", "local"))
conf = _pyspark.SparkConf().setAll(spark_conf)
return _pyspark.SparkContext(conf=conf)
def __enter__(self):
GlobalSparkContext._SPARK_CONTEXT = _pyspark.SparkContext()
GlobalSparkContext._SPARK_SESSION = _pyspark.sql.SparkSession.builder.appName(
"Flyte Spark SQL Context").getOrCreate()
return self
def __enter__(self):
GlobalSparkContext._SPARK_CONTEXT = _pyspark.SparkContext()
return self