Databricks Community

shubham_007 · ‎02-18-2025

Dear experts,

I am facing difficulty while developing pyspark automation logic on “Developing automation logic to delete/remove display() and cache() method used in scripts in multiple databricks notebooks (tasks)”.

kindly advise on developing automation script.

mark_ott · Friday

To automate the removal of display() and cache() method calls from multiple PySpark scripts in Databricks notebooks, develop a script that programmatically processes exportable notebook source files (usually in .dbc or .ipynb format) using text-based search and replace logic. This approach is effective and well-suited for bulk Notebook modifications, enabling automation across many files without manual intervention.

Automation Approach: Step-by-Step

Export all relevant Databricks notebooks to source file formats (.ipynb or .dbc). Store these in a local or cloud workspace dedicated for batch processing.
Develop a Python script (or use tools like Notepad++ or sed/awk for shell scripting) that reads each notebook as a text file and searches for lines containing display( and .cache(.
Remove, comment out, or replace these lines according to your policy (for instance, remove visualization logic for production, or replace with logging/debug output).
After transformation, re-import cleaned notebooks back into Databricks using the CLI, API, or web interface for downstream use.
(Optional) Version the changes using Git integration or maintain backup copies of the raw notebooks to prevent accidental data loss.

Example Python Script Snippet

python

import os, re

def clean_notebook(notebook_path):
    with open(notebook_path, 'r', encoding='utf-8') as f:
        content = f.read()

    # Remove display() calls
    content = re.sub(r'display\([^)]*\)', '', content)
    # Remove cache() calls
    content = re.sub(r'\.cache\(\)', '', content)

    with open(notebook_path, 'w', encoding='utf-8') as f:
        f.write(content)

# Apply to directory of notebooks
target_dir = "path/to/notebooks"
for fname in os.listdir(target_dir):
    if fname.endswith('.py') or fname.endswith('.ipynb'):
        clean_notebook(os.path.join(target_dir, fname))

This basic logic parses each notebook and removes the designated method calls. For .ipynb files, you may want to process JSON cell structures instead of flat text for higher precision.

Best Practices & Notes

Always backup original notebooks before bulk modification.
Use regular expressions judiciously. Ensure that function calls spanning multiple lines, or embedded in longer statements, are thoroughly detected and removed.
Test the automation logic on a small sample before full application, to avoid accidental removal of valid code.
Consider using the Databricks CLI for notebook export/import operations programmatically for large scale or repeated runs.
For advanced workflows, integrate the logic into your CI/CD pipeline using Python, PowerShell, or Bash scripting for automated enforcement.

This method streamlines the removal of visualization and caching calls from PySpark scripts, making your Databricks pipelines cleaner and more production-ready.