azure databricks 中具有 Spark/dbutils 依赖项的自定义 python 模块

Question

我最近在 Azure Databricks 上启用了预览功能“存储库中的文件”，这样我就可以将许多常规功能从笔记本转移到模块，并消除为单个作业运行大量笔记本带来的开销。

但是，我的一些函数直接依赖于 dbutils 或 Spark/pyspark 函数（例如dbutils.secrets.get()和spark.conf.set()）。由于这些模块是在笔记本的后台导入的，并且直接与底层会话相关联，因此我完全不知道如何在自定义模块中引用这些模块。

对于我的小示例模块，我通过将 dbutils 设置为参数来修复它，如下例所示：

class Connection:
    def __init__(self, dbutils):
        token = dbutils.secrets.get(scope="my-scope", key="mykey")
        ...

然而，对所有现有函数执行此操作将需要大量重写函数和调用它们的行。我怎样才能避免这个过程并以更干净的方式进行？

Answer 1

Databricks Connect 的文档展示了如何实现它的示例。该示例具有SparkSession显式参数，但可以对其进行修改以完全避免这种情况，如下所示：

def get_dbutils():
  from pyspark.sql import SparkSession
  spark = SparkSession.getActiveSession()
  if spark.conf.get("spark.databricks.service.client.enabled") == "true":
    from pyspark.dbutils import DBUtils
    return DBUtils(spark)
  else:
    import IPython
    return IPython.get_ipython().user_ns["dbutils"]

然后在你的函数中你可以做这样的事情：

def get_my_secret(scope, key):
  return get_dbutils().secrets.get(scope, key)