Databricks series : Setting up a Databricks project in Visual Studio Code (part III)

Introduction

In this blogpost I'll describe the Visual Studio Project in more detail. 

Previous blogposts : 

• Databricks series : Starting a new python project in Visual Studio Code (Part I) 
• Databricks series : setting up a connection to Databricks from local IDE (Part II)

Project structure

Lets create a project structure  where we can store our code and also the data files

I've created the following structure in Visual studio Code

databricks-transfer/
│
├── databricks.yml
├── resources/
│   └── transfer_job.yml
│
├── src/
│   ├── __init__.py
│   └── main.py
│
├── data/
│   ├── customers.csv
│   └── transactions.csv
│
├── requirements.txt
└── README.md

What is the function of the different folders?

Some test data

For this blogseries I've created a couple of testfiles with help of chatgpt :

customers.csv

transactions.csv

Python code

The application performs three tasks:

1. Receives parameters from the Databricks Job.

2. Reads the specified CSV file using Spark.

3. Writes the data to a Delta table.

The advantage is that the Python code is independent of the specific dataset being processed.

import argparse
from pyspark.sql import SparkSession


def main():

    parser = argparse.ArgumentParser()

    parser.add_argument("--entity", required=True)
    parser.add_argument("--source-path", required=True)
    parser.add_argument("--target-table", required=True)

    args = parser.parse_args()

    spark = SparkSession.builder.getOrCreate()

    print(f"Processing entity: {args.entity}")

    df = (
        spark.read
        .option("header", "true")
        .option("inferSchema", "true")
        .csv(args.source_path)
    )

    (
        df.write
        .format("delta")
        .mode("overwrite")
        .saveAsTable(args.target_table)
    )

    print(f"Successfully processed {df.count()} records")


if __name__ == "__main__":
    main()

DAB

databricks.yml

bundle:
  name: DatabricksProject

include:
  - resources/*.yml

targets:
  dev:
    mode: development
    default: true
    workspace:
      host: https://dbc-xxxxxxx.cloud.databricks.com

This is an example of a serverless Python Job. The environment is linked to the task via environment_key. The Job has three parameters:

Create a transfer_job.yml in the resources folder

resources:
  jobs:
    transfer_job:

      name: python-data-transfer

      parameters:
        - name: entity
          default: customers

        - name: source_path
          default: /Volumes/main/demo/raw/customers.csv

        - name: target_table
          default: main.demo.customers

      tasks:
        - task_key: transfer_data

          spark_python_task:
            python_file: ../src/main.py

            parameters:
              - --entity
              - "{{job.parameters.entity}}"
              - --source-path
              - "{{job.parameters.source_path}}"
              - --target-table
              - "{{job.parameters.target_table}}"

          environment_key: default

      environments:
        - environment_key: default
          spec:
            environment_version: "2"

The Job passes these parameters to main.py. Pay special attention to:"{{job.parameters.entity}}"

This is a dynamic reference to a Job parameter. Databricks populates this value when the Job is executed. This means we do not need to modify our Python code when we want to process a different dataset. It is also possible to perform an explicit sync action.

bundle:
  name: DatabricksProject

include:
  - resources/*.yml

sync:
  include:
    - src/**
    - requirements.txt
    
targets:
  dev:
    mode: development
    default: true
    workspace:
      host: https://dbc-6270be54-d21b.cloud.databricks.com

  ## Optionally, there could be 'staging' or 'prod' targets here.
  #
  # prod:
  #   workspace:
  #     host: https://dbc-6270be54-d21b.cloud.databricks.com

This allows you to specify which project files you want to synchronize. We deliberately exclude the `data/` directory: our CSV files are ultimately stored in a Unity Catalog Volume.

Hennie