Obtenir le contexte de la tâche dans une UDF

Utilisez l'API PySpark TaskContext pour obtenir des informations de contexte lors de l'exécution d'une UDF Python de Unity Catalog par lot ou d'une UDF PySpark.

Par exemple, des informations contextuelles telles que l'identité de l'utilisateur et les Cluster Tags peuvent vérifier l'identité d'un utilisateur pour accéder aux services externes.

Exigences

• TaskContext est pris en charge sur les versions de Databricks Runtime 16.3 et ultérieures.

• TaskContext est pris en charge sur les types d'UDF suivants :

  • UDF Python de batch Unity Catalog
  • UDFs scalaires Python

Utiliser TaskContext pour obtenir des informations contextuelles

Sélectionnez un tab pour voir des exemples de TaskContext pour les UDF PySpark ou les UDF Python Unity Catalog par batch.

PySpark UDF

L'exemple de PySpark UDF suivant affiche le contexte de l'utilisateur :

Python

@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
  }

Batch Unity Catalog Python UDF

L'exemple suivant de Batch Unity Catalog Python UDF obtient l'identité de l'utilisateur pour appeler une fonction AWS Lambda à l'aide d'informations d'identification de service :

Python

%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"])
$$;

Appelez l'UDF après son enregistrement :

SQL

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

Propriétés TaskContext

La méthode TaskContext.getLocalProperty() possède les clés de propriété suivantes :

Clé de propriété	Description	Exemple d'utilisation
user	L'utilisateur qui exécute actuellement l'UDF	tc.getLocalProperty("user") -> "alice"
spark.jobGroup.id	L'ID de groupe de job Spark associé à l'UDF actuelle	tc.getLocalProperty("spark.jobGroup.id") -> "jobGroup-92318"
spark.databricks.clusterUsageTags.clusterAllTags	Balises de métadonnées de cluster sous forme de paires clé-valeur formatées comme une représentation sous forme de chaîne d'un dictionnaire JSON	tc.getLocalProperty("spark.databricks.clusterUsageTags.clusterAllTags") -> [{"Department": "Finance"}]
spark.databricks.clusterUsageTags.region	La région où réside le Workspace	tc.getLocalProperty("spark.databricks.clusterUsageTags.region") -> "us-west-2"
accountId	ID de compte Databricks pour le contexte d'exécution	tc.getLocalProperty("accountId") -> "1234567890123456"
orgId	ID du Workspace (non disponible sur DBSQL)	tc.getLocalProperty("orgId") -> "987654321"
spark.databricks.clusterUsageTags.sparkVersion	Version de Databricks Runtime pour le cluster (dans les environnements non-DBSQL)	tc.getLocalProperty("spark.databricks.clusterUsageTags.sparkVersion") -> "16.3"
spark.databricks.clusterUsageTags.dbsqlVersion	Version DBSQL (sur les environnements DBSQL)	tc.getLocalProperty("spark.databricks.clusterUsageTags.dbsqlVersion") -> "2024.35"

Clé de propriété	Description	Exemple d'utilisation
user	L'utilisateur qui exécute actuellement l'UDF	tc.getLocalProperty("user") -> "alice"
spark.jobGroup.id	L'ID de groupe de job Spark associé à l'UDF actuelle	tc.getLocalProperty("spark.jobGroup.id") -> "jobGroup-92318"
spark.databricks.clusterUsageTags.clusterAllTags	Balises de métadonnées de cluster sous forme de paires clé-valeur formatées comme une représentation sous forme de chaîne d'un dictionnaire JSON	tc.getLocalProperty("spark.databricks.clusterUsageTags.clusterAllTags") -> [{"Department": "Finance"}]
spark.databricks.clusterUsageTags.region	La région où réside le Workspace	tc.getLocalProperty("spark.databricks.clusterUsageTags.region") -> "us-west-2"
accountId	ID de compte Databricks pour le contexte d'exécution	tc.getLocalProperty("accountId") -> "1234567890123456"
orgId	ID du Workspace (non disponible sur DBSQL)	tc.getLocalProperty("orgId") -> "987654321"
spark.databricks.clusterUsageTags.sparkVersion	Version de Databricks Runtime pour le cluster (dans les environnements non-DBSQL)	tc.getLocalProperty("spark.databricks.clusterUsageTags.sparkVersion") -> "16.3"
spark.databricks.clusterUsageTags.dbsqlVersion	Version DBSQL (sur les environnements DBSQL)	tc.getLocalProperty("spark.databricks.clusterUsageTags.dbsqlVersion") -> "2024.35"