PySpark项目从Databricks UI迁移至VSCode（Databricks Connect）的实践难题

问题1：.py文件中%run魔法命令无法识别的规避方案 #

无需修改现有代码，通过预处理脚本+VSCode运行配置实现兼容：

• 编写预处理脚本（比如run_processor.py），读取目标.py文件内容并自动转换%run语句为Python可执行逻辑：

  import sys
  import os
  
  def process_run_statements(file_path):
      with open(file_path, 'r') as f:
          content = f.read()
      # 将%run ./xxx转换为exec读取对应文件的代码
      processed_content = content.replace('%run ', 'exec(open("').replace('\n', '").read())\n')
      exec(processed_content, globals())
  
  if __name__ == "__main__":
      if len(sys.argv) != 2:
          print("Usage: python run_processor.py <path-to-your-script.py>")
          sys.exit(1)
      target_file = sys.argv[1]
      os.chdir(os.path.dirname(os.path.abspath(target_file)))
      process_run_statements(target_file)
  

• 在VSCode中配置运行任务（.vscode/launch.json），把默认Python运行命令替换为调用该预处理脚本：

  {
      "version": "0.2.0",
      "configurations": [
          {
              "name": "Python: Databricks Connect Run",
              "type": "python",
              "request": "launch",
              "program": "${workspaceFolder}/run_processor.py",
              "args": ["${file}"],
              "console": "integratedTerminal"
          }
      ]
  }
  

后续在VSCode中运行.py文件时，选择这个配置即可自动处理%run语句，原代码完全无需改动。

问题2：自动初始化SparkSession与dbutils的方案 #

通过Python启动钩子实现全局自动初始化，不用在每个文件中重复添加代码：

方案1：使用PYTHONSTARTUP环境变量 #

• 在项目根目录创建init_spark.py，写入适配Databricks Connect的初始化代码（替换为你的集群配置）：

  from pyspark.sql import SparkSession
  from pyspark.dbutils import DBUtils
  
  spark = SparkSession.builder \
      .remote("sc://<your-workspace-url>:443") \
      .config("spark.databricks.service.token", "<your-access-token>") \
      .config("spark.databricks.service.cluster.id", "<your-cluster-id>") \
      .getOrCreate()
  
  dbutils = DBUtils(spark)
  

• 在VSCode的settings.json中配置终端默认环境变量，让每次启动终端时自动加载初始化脚本：

  {
      "terminal.integrated.env.linux": {
          "PYTHONSTARTUP": "${workspaceFolder}/init_spark.py"
      },
      "terminal.integrated.env.windows": {
          "PYTHONSTARTUP": "${workspaceFolder}/init_spark.py"
      },
      "terminal.integrated.env.osx": {
          "PYTHONSTARTUP": "${workspaceFolder}/init_spark.py"
      }
  }
  

之后在VSCode终端中运行或调试代码时，spark和dbutils会自动全局生成，和Databricks UI的使用体验完全一致。

方案2：使用sitecustomize.py #

如果希望在整个Python环境中生效，可在环境的site-packages目录下创建sitecustomize.py，内容和上述init_spark.py一致。这样无论在哪个终端或编辑器中运行该环境的Python，都会自动完成初始化。

内容的提问来源于stack exchange，提问作者tendaitakas