apache / apache/airflow

Support `@task.spark_submit` or `@task.pyspark_submit`

Open
#40,566 1 comment 2 reactions 0 assignees View on GitHub
area:providers kind:feature provider:apache-spark
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 10h
Merged PRs (30d)
483

Description

### Description

Hello,
I'm an user of Airflow for managing Spark jobs.
I was wondering if there's any plan to provide `@task.spark_submit` decorator in the future.

Actually, I’ve been using a custom provider that I made myself.
In my situation, I write and execute pyspark code locally, and thenI transform the code to a form that can be submitted to YARN finally.
Transforming the code into a form suitable for YARN (just creating separate .py file) is not particularly difficult.
However, when the task code and the YARN submission code are separated, the folder structure within the DAG can become complex, and the task code and PySpark job files cannot reference each other in the IDE.

by `@task.spark_submit` (or `@task.pyspark_submit`?), it would improve development productivity significantly and very convenient.
This is because we could replace the existing @task.pyspark decorators by simply changing the decorator name and injecting the YARN options.

> all I need to do is just changing the decorator name and inject the YARN options.

### Use case/motivation

The idea is to borrow the approach from the existing `@task.pyspark` and `@task.docker` decorators.
It works by writing a temporary py file to /tmp and then submitting that file.

```python
class _SparkSubmitDecoratedOperator(DecoratedOperator, SparkSubmitOperator):
custom_operator_name = "@task.spark_submit"

# In the constructor, Some codes like removing the sc and spark arguments are located .

...
def execute(self, context: Context):
with TemporaryDirectory() as tmp_dir:
input_filename = os.path.join(tmp_dir, "script.in")
script_filename = os.path.join(tmp_dir, "script.py")
output_filename = os.path.join(tmp_dir, "script.out")
error_filename = os.path.join(tmp_dir, "script.err")

with open(input_filename, "w", encoding="utf-8") as file:
if self.op_args or self.op_kwargs:
self.pickling_library.dump(
{"args": self.op_args, "kwargs": self.op_kwargs}, file
)

py_source = self.get_python_source()
write_python_script(
jinja_context={
"op_args": self.op_args,
"op_kwargs": self.op_kwargs,
"pickling_library": self.pickling_library.__name__,
"python_callable": self.python_callable.__name__,
"python_callable_source": py_source,
"expect_airflow": True,
"string_args_global": False,
},
filename=script_filename,
)
self.application = script_filename
self.application_args = [
input_filename,
output_filename,
"--no-string-args",
error_filename,
]
return super().execute(context)
```

### Related issues

_No response_

### Are you willing to submit a PR?

- [X] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)

Contributor guide

Open the contributing guide

Research direction

Start by reading the existing @task.pyspark and @task.docker decorators and SparkSubmitOperator, then compare their handling of callable source, arguments, and temporary scripts. Define the supported decorator name and YARN options, implement the integration, and add coverage showing that a decorated PySpark callable is rendered and submitted correctly.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, spark
Domain
data-engineering
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.