具有 spark/dbutils 依赖项的 azure databricks 中的自定义 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

documentation for Databricks Connect shows the example如何实现。该示例将 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)

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

Custom python module in azure databricks with spark/dbutils dependencies

python

python-3.x

databricks

azure-databricks