Получение контекста задачи в пользовательской функции

Используйте API TaskContext PySpark для получения сведений о контексте во время запуска пакетного каталога Unity, Python UDF или PySpark UDF.

Например, контекстная информация, такая как удостоверение пользователя и теги кластера, может проверить личность пользователя для доступа к внешним службам.

Требования

• TaskContext поддерживается в Databricks Runtime версии 16.3 и более поздних.

• TaskContext поддерживается в следующих типах пользовательских функций:

  • Пакетная обработка пользовательских функций Python Catalog Unity
  • Скалярные пользовательские функции Python

Использование TaskContext для получения информации о контексте

Выберите вкладку, чтобы просмотреть примеры TaskContext для пользовательских функций PySpark или пользовательских функций Python из пакетного каталога Unity.

Пользовательская функция PySpark

В следующем примере пользовательской функции PySpark выводится контекст пользователя:

@udf
def log_context():
  import json
  from pyspark.taskcontext import TaskContext
  tc = TaskContext.get()

  # Returns current user executing the UDF
  session_user = tc.getLocalProperty("user")

  # Returns cluster tags
  tags = dict(item.values() for item in json.loads(tc.getLocalProperty("spark.databricks.clusterUsageTags.clusterAllTags  ") or "[]"))

  # Returns current version details
  current_version = {
    "dbr_version": tc.getLocalProperty("spark.databricks.clusterUsageTags.sparkVersion"),
    "dbsql_version": tc.getLocalProperty("spark.databricks.clusterUsageTags.dbsqlVersion")
  }

  return {
    "user": session_user,
    "job_group_id": job_group_id,
    "tags": tags,
    "current_version": current_version
  }

UDF пакетного каталога Unity

В следующем примере пользовательской функции Python в каталоге Unity пакетной службы идентификатор пользователя получает удостоверение пользователя для вызова функции AWS Lambda с использованием учетных данных службы:

%sql
CREATE OR REPLACE FUNCTION main.test.call_lambda_func(data STRING, debug BOOLEAN) RETURNS STRING LANGUAGE PYTHON
PARAMETER STYLE PANDAS
HANDLER 'batchhandler'
CREDENTIALS (
  `batch-udf-service-creds-example-cred` DEFAULT
)
AS $$
import boto3
import json
import pandas as pd
import base64
from pyspark.taskcontext import TaskContext


def batchhandler(it):
  # Automatically picks up DEFAULT credential:
  session = boto3.Session()

  client = session.client("lambda", region_name="us-west-2")

  # Can propagate TaskContext information to lambda context:
  user_ctx = {"custom": {"user": TaskContext.get().getLocalProperty("user")}}

  for vals, is_debug in it:
    payload = json.dumps({"values": vals.to_list(), "is_debug": bool(is_debug[0])})

    res = client.invoke(
      FunctionName="HashValuesFunction",
      InvocationType="RequestResponse",
      ClientContext=base64.b64encode(json.dumps(user_ctx).encode("utf-8")).decode(
        "utf-8"
      ),
      Payload=payload,
    )

    response_payload = json.loads(res["Payload"].read().decode("utf-8"))
    if "errorMessage" in response_payload:
      raise Exception(str(response_payload))

    yield pd.Series(response_payload["values"])
$$;

Вызовите UDF после его регистрации:

SELECT main.test.call_lambda_func(data, false)
FROM VALUES
('abc'),
('def')
AS t(data)

Свойства TaskContext

Метод TaskContext.getLocalProperty() имеет следующие ключи свойств:

Ключ свойства	Описание	Пример использования
user	Пользователь, выполняющий в данный момент пользовательскую функцию	tc.getLocalProperty("user") ->"alice"
spark.jobGroup.id	Идентификатор группы заданий Spark, связанный с текущей пользовательской функцией	tc.getLocalProperty("spark.jobGroup.id") ->"jobGroup-92318"
spark.databricks.clusterUsageTags.clusterAllTags	Кластеризуйте теги метаданных в виде пар ключ-значение, отформатированных как строковое представление словаря JSON	tc.getLocalProperty("spark.databricks.clusterUsageTags.clusterAllTags") ->[{"Department": "Finance"}]
spark.databricks.clusterUsageTags.region	Регион, в котором находится рабочая область	tc.getLocalProperty("spark.databricks.clusterUsageTags.region") ->"us-west-2"
accountId	Идентификатор учетной записи Databricks для контекста выполнения	tc.getLocalProperty("accountId") ->"1234567890123456"
orgId	Идентификатор рабочей области (недоступен в DBSQL)	tc.getLocalProperty("orgId") ->"987654321"
spark.databricks.clusterUsageTags.sparkVersion	Версия среды выполнения Databricks для кластера (в средах, отличных от DBSQL)	tc.getLocalProperty("spark.databricks.clusterUsageTags.sparkVersion") ->"16.3"
spark.databricks.clusterUsageTags.dbsqlVersion	Версия DBSQL (в средах DBSQL)	tc.getLocalProperty("spark.databricks.clusterUsageTags.dbsqlVersion") ->"2024.35"