Azure Databricks中如何在执行前分析Spark物理计划并拦截特定路径读取？

要实现查询执行前检查物理计划并终止违规查询，可以借助Spark的QueryStagePrepRule——这是Spark在物理计划生成完成后、实际执行前会运行的规则，正好能满足你的需求。

核心思路 #

• 自定义一个继承QueryStagePrepRule的规则类，遍历物理计划的所有执行节点
• 识别出FileSourceScanExec类型的节点（这类节点对应文件读取操作，比如spark.read.parquet），提取其关联的文件路径
• 检查路径是否包含你要阻止的目标路径，若存在则抛出异常终止查询
• 将规则注册到SparkSession，仅在需要限制的集群上执行注册逻辑

代码实现（Scala） #

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.execution.QueryStagePrepRule
import org.apache.spark.sql.execution.datasources.FileSourceScanExec

// 自定义物理计划检查规则
class BlockPathRule(targetPath: String) extends QueryStagePrepRule {
  override def apply(plan: org.apache.spark.sql.execution.SparkPlan): org.apache.spark.sql.execution.SparkPlan = {
    // 递归遍历所有物理计划节点
    plan.transform {
      case scan: FileSourceScanExec =>
        // 获取扫描的所有路径
        val scannedPaths = scan.metadata("Path").split(",").map(_.trim)
        if (scannedPaths.exists(_.contains(targetPath))) {
          throw new IllegalArgumentException(s"查询访问了受限路径: $targetPath，已终止执行")
        }
        scan
    }
  }
}

// 注册规则到SparkSession
val spark = SparkSession.builder().getOrCreate()
spark.extensions.registerQueryStagePrepRule(new BlockPathRule("/path/to/block"))

关键说明 #

• 为什么用QueryStagePrepRule：这个规则是在物理计划生成后、查询执行前触发的，完美契合你"执行前检查物理计划"的需求，不像QueryExecutionListener是事后回调，也不像普通Rule只能处理逻辑计划。
• 路径识别：FileSourceScanExec节点的metadata里存储了实际扫描的文件路径，正好能拿到spark.read.parquet指定的路径（这部分在逻辑计划里是看不到的）。
• 集群范围控制：你只需要在需要限制的集群上运行这段注册代码——比如把代码放到集群的初始化脚本里，或者在该集群的Notebook开头执行，其他集群不执行注册，就不会受影响，避免了全局访问限制的问题。

注意事项 #

• 如果用Python开发，需要通过Py4J调用Scala的API，或者直接用Scala编写规则打包成Jar，在Python环境中引入后注册。
• 抛出的异常会直接终止查询，你可以根据需要调整异常类型或提示信息，让用户更清晰地知道原因。

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