Probably two of most frequently used assets in Declarative Automation Bundles on Databricks are wheel-based jobs and notebooks. But do you know the Lakeflow Jobs component can also run regular Python scripts that don't have to be packaged as a wheel.
What would it take for you to trust your Databricks pipelines in production?
A 3-day bug hunt on a 3-person team costs up to €7,200 in lost engineering time. This workshop teaches you to prevent that — unit tests, data tests, and integration tests for PySpark and Databricks Lakeflow, including Spark Declarative Pipelines.
Konieczny
Use cases
The first question you may ask, why do I need Python scripts while I can deploy wheel-packaged code? True, but what if you need to deploy something really simple, like a generator of dynamic task values for downstream tasks? You can make it a part of your whole package but ultimately it can be just a regular Python script that doesn't rely on any external dependencies.
When not to use Python scripts? I'm always in favor of deploying fully packaged and versioned artifacts, so the wheel-based jobs. Besides providing some clear release logic, they quite naturally impose some engineering rigor when it comes to naming and organizing the code base. Consequently, Python scripts are poor candidates - the same way as notebooks - for production code that must do some business logic, be linted, formatted, and tested properly, as anything reaching the production and real data.
On another side, the simplicity of Python scripts is something we may leverage for simpler things, such as implementing a sensor (cf. Readiness Marker from my Data Engineering Design Patterns book) that verifies if the pipeline can start processing or should wait. It can be also used as a data provider for downstream tasks, e.g. if you want to separate this parameterization part from the actual data processing.
Ultimately, you could also use it as a separation layer between technical and business logic to not mix both aspects in the code base. In a scenario all code touching the data could be implemented as wheel-based PySpark jobs while the parts interacting with Databricks infrastructure could be read and processed as Python scripts. Of course, as long as you are an unlucky user who has to do everything on Databricks and cannot rely on more natural infrastructure abstractions, such as Databricks Terraform provider. That being said, script jobs could somehow be useful to cover missing gaps, such as bootstraping subdirectories in a volume, the same way I described hooks in Running scripts as hooks with Databricks Asset Bundles.
Execution model
There is no difference with the wheel-based jobs or notebooks. A Python script runs on the cluster associated to your job. It automatically has access to the libraries installed on the cluster and the SparkSession instance.
When it comes to the executed code, it can be located in your workspace, on an object store (S3, GCS, Azure Blob), or directly in your Git repository.You don't have to specify the package name as it's the case for the wheel-based task. Instead, this single remote or local path is enough.
A subtle difference with the wheel-based tasks are input parameters. Python script ony supports an array of strings which is less handy than the named_parameters from the wheel task.
Example
Let's see an example where we are going to use a Python script task as an implementation of the Readiness Marker pattern and as a data provider for a downstream foreach task. First, let's take a look at the Declarative Automation Bundle (DAB) definition for our job:
tasks:
- task_key: check_if_file_exists
max_retries: 5
min_retry_interval_millis: 60000
spark_python_task:
python_file: ../src/python_script_task/file_checker.py
parameters: ['--file_to_check', '/Volumes/workspace/default/countries/countries_list.txt']
source: WORKSPACE
environment_key: default
- task_key: generate_countries_list
depends_on:
- task_key: check_if_file_exists
spark_python_task:
python_file: ../src/python_script_task/country_list_generator.py
parameters: ['--file_to_read', '/Volumes/workspace/default/countries/countries_list.txt']
source: WORKSPACE
environment_key: default
- task_key: process_generated_countries
depends_on:
- task_key: generate_countries_list
for_each_task:
inputs: '{{ tasks.generate_countries_list.values.countries_list }}'
concurrency: 2
task:
task_key: process_country
environment_key: default
python_wheel_task:
package_name: python_script_task
entry_point: process_country
named_parameters:
country: "{{input}}"
The full code for this example, alongside the code snippets for past and future blog posts, is in my add repo link databricks-playground repo.
When it comes to the particular tasks, please notice the retry configuration for the first sensor. Whenever the file we are expecting is not present, the task will retry every 5 minutes and eventually fail after the 5th retry. The script logic is simple as is:
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--file_to_check", required=True, help="The file path to look for.")
arg = parser.parse_args()
if not pathlib.Path(arg.file_to_check).is_file():
raise RuntimeError(f'File {arg.file_to_check} not found.')
As you can see, just a few lines without declaring an entrypoint in the pyproject.toml. Once the file is available for processing, another Python script task reads it from the volume and generates a dynamic list of countries our fictional PySpark job must process:
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--file_to_read", required=True, help="The file path to look for.")
arg = parser.parse_args()
countries_list = str(pathlib.Path(arg.file_to_read).read_text(encoding='UTF-8')).split(',')
from databricks.sdk.runtime import dbutils
dbutils.jobs.taskValues.set(key='countries_list', value=countries_list)
Here is how the whole orchestration runs:
As exotic as it may sound, Python scripts are yet another data processing backend you can use in Lakeflow Jobs. For sake of simplicity, reserve them to rather simpler tasks as the ones from my example.
Data Engineering Design Patterns
Looking for a book that defines and solves most common data engineering problems? I wrote
one on that topic! You can read it online
on the O'Reilly platform,
or get a print copy on Amazon.
I also help solve your data engineering problems contact@waitingforcode.com đź“©
Read also about Python script tasks on Databricks here:
Related blog posts:
- SQLFluff, i.e. keeping SQL queries clean
- Ruff and Declarative Automation Bundles
- Managing Unity Catalog resources on Databricks