Python スカラーユーザー定義関数 (UDF)

Python スカラー UDF を使用すると、Databricks 上の SQL クエリー内でカスタムの Python ロジックを実行できます。このページでは、それらの登録と呼び出し方法、サービス認証情報とシークレットの使用方法、および Spark SQL での部分式の評価順序に関する注意事項の処理方法について説明します。

要件

Databricks Runtime12.2LTS 以前では、標準アクセスモードを使用するUnity Catalogコンピュートでは、Python UDF と Pandas UDF はサポートされていません。
Scalar Python UDF と Pandas UDF は、Databricks Runtime 13.3 LTS 以降ですべてのアクセスモードでサポートされています。
Unity Catalog 対応クラスター上の Python UDF に対する Graviton インスタンスのサポートには、Databricks Runtime 15.2 以上が必要です。

Databricks Runtime14.0 以前では、標準アクセスPython PandasUnity Catalogモードを使用するクラスターでは、 UDF と UDF はサポートされていません。スカラー Python UDF と Pandas UDF は、Databricks Runtime 14.1 以降のすべてのアクセスモードでサポートされています。

Databricks Runtime 14.1 以降では、SQL構文を使用してスカラーPython UDFをUnity Catalogに登録できます。Unity Catalog の SQL および Python ユーザー定義関数 (UDF)を参照してください。

関数を UDF として登録する

Python
def squared(s):
  return s * s
spark.udf.register("squaredWithPython", squared)

オプションで、UDF の戻り値の型を設定できます。デフォルトの戻り値の型は StringTypeです。

Python
from pyspark.sql.types import LongType
def squared_typed(s):
  return s * s
spark.udf.register("squaredWithPython", squared_typed, LongType())

Spark SQL で UDF を呼び出す

Python
spark.range(1, 20).createOrReplaceTempView("test")

SQL
%sql select id, squaredWithPython(id) as id_squared from test

UDF と DataFrames の併用

Python
from pyspark.sql.functions import udf
from pyspark.sql.types import LongType
squared_udf = udf(squared, LongType())
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

または、アノテーション構文を使用して同じ UDF を宣言することもできます。

Python
from pyspark.sql.functions import udf

@udf("long")
def squared_udf(s):
  return s * s
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

UDF を使用したバリアント

バリアントの PySpark 型はVariantTypeであり、値の型はVariantValです。バリアントに関する情報については、「バリアントデータのクエリ」を参照してください。

Python
from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import VariantType, VariantVal

# Return Variant
@udf(returnType = VariantType())
def toVariant(jsonString):
  return VariantVal.parseJson(jsonString)

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toVariant(col("json"))).display()

+---------------+
|toVariant(json)|
+---------------+
|        {"a":1}|
+---------------+

Python
from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import StructField, StructType, VariantType, VariantVal

# Return Struct<Variant>
@udf(returnType = StructType([StructField("v", VariantType(), True)]))
def toStructVariant(jsonString):
  return {"v": VariantVal.parseJson(jsonString)}

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toStructVariant(col("json"))).display()

+---------------------+
|toStructVariant(json)|
+---------------------+
|        {"v":{"a":1}}|
+---------------------+

Python
from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import ArrayType, VariantType, VariantVal

# Return Array<Variant>
@udf(returnType = ArrayType(VariantType()))
def toArrayVariant(jsonString):
  return [VariantVal.parseJson(jsonString)]

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toArrayVariant(col("json"))).display()

+--------------------+
|toArrayVariant(json)|
+--------------------+
|           [{"a":1}]|
+--------------------+

Python
from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import MapType, StringType, VariantType, VariantVal

# Return Map<String, Variant>
@udf(returnType = MapType(StringType(), VariantType(), True))
def toMapVariant(jsonString):
  return {"v1": VariantVal.parseJson(jsonString), "v2": VariantVal.parseJson("[" + jsonString + "]")}

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toMapVariant(col("json"))).display()

+-----------------------------+
|           toMapVariant(json)|
+-----------------------------+
|{"v2":[{"a":1}],"v1":{"a":1}}|
+-----------------------------+

UDF を含むファイル

備考

ベータ版

この機能はベータ版です。

ファイルの PySpark タイプは FileType です。UDF のパラメーターまたは戻り値の型として、トップレベルの型またはネストされた型として使用します。型、そのネスト規則、および FileRef API については、ファイルタイプを参照してください。

UDF でファイルの内容を読み取るには、file.as_local_file() を呼び出して開くことができるローカルパスを取得するか、file.open() を呼び出してそのバイトをストリームとして読み取ります。画像処理、ファイルタイプ検出、ビデオフレーム抽出など、Python、Scala、SQL の例については、UDF を使用したファイルの処理を参照してください。タイプのリファレンスについては、FILE タイプを参照してください。

Unity Catalogに登録された UDF (CREATE FUNCTION) は FILE のメタデータを読み取ることができますが、コンテンツは読み取れず、ファイルを作成することもできません。セッションスコープの UDF を使用して、ファイルの内容 (open, as_local_file) を読み取るか、ファイルを作成します (from_bytes, from_local_file)。

UDF 本文では、各 FILE 値は FileRef オブジェクトです。

Python
from pyspark.sql.functions import col, udf
from pyspark.sql.types import FileRef, StringType

@udf(returnType=StringType())
def file_content_type(file: FileRef) -> str:
  return file.content_type

df = spark.table("documents")
display(df.select(col("file").uri, file_content_type(col("file"))))

評価順序と null チェック

Spark SQL ( SQL 、データフレーム、データセット APIを含む)は、部分式の評価。特に、演算子または関数の入力はそうではありません必ず左から右に評価するか、その他の固定された順序で評価されます。たとえば、論理 AND また、 OR 式には、左から右への "短絡" セマンティクスはありません。

そのため、Booleanの副作用や評価の順番に頼るのは危険です式、および WHERE 節と HAVING 節の順序は、クエリの最適化と計画中に並べ替えられます。具体的には、UDF が null チェックのために SQL の短絡セマンティクスに依存している場合、 UDF を呼び出す前にヌルチェックが行われることを保証します。例えば

Python
spark.udf.register("strlen", lambda s: len(s), "int")
spark.sql("select s from test1 where s is not null and strlen(s) > 1") # no guarantee

この WHERE 句は、null をフィルターで除外した後に strlen UDF が呼び出されることを保証するものではありません。

適切な null チェックを実行するには、次のいずれかを実行することをお勧めします。

UDF 自体を null 対応にし、UDF 自体の内部で null チェックを行います
IF式またはCASE WHEN式を使用してヌルチェックを行い、条件分岐でUDFを呼び出します

Python
spark.udf.register("strlen_nullsafe", lambda s: len(s) if not s is None else -1, "int")
spark.sql("select s from test1 where s is not null and strlen_nullsafe(s) > 1") # ok
spark.sql("select s from test1 where if(s is not null, strlen(s), null) > 1")   # ok

Unity Catalog のシークレットへのアクセス

セッションスコープの Python UDF から Unity Catalog のシークレットにアクセスするには、セッションスコープの Python UDF でシークレットを使用するを参照してください。スカラーまたはバッチ Unity Catalog Python UDF から宣言されたシークレットにアクセスするには、Python UDF でシークレットを使用するを参照してください。

Python UDF におけるサービス資格情報

セッションスコープの Python スカラー UDF およびスカラー Unity Catalog Python UDF では、Unity Catalog サービス認証情報を使用して外部クラウドサービスに安全にアクセスできます。これは、クラウドベースのトークン化、暗号化、シークレット管理などの操作をデータ変換に直接統合する場合に便利です。

要件は、UDF の種類とコンピュートによって異なります。Python UDF でのサービス資格情報の使用を参照してください。

サービス資格情報を作成するには、「サービス資格情報の作成」を参照してください。

セッションスコープのスカラー Python UDF でのサービス認証情報の使用

サービス認証情報にアクセスするには、UDF ロジックの databricks.service_credentials.getServiceCredentialsProvider() ユーティリティを使用して、適切な認証情報でクラウド SDK を初期化します。すべてのコードは UDF 本体にカプセル化する必要があります。

Python
@udf
def use_service_credential():
    from databricks.service_credentials import getServiceCredentialsProvider
    import boto3

    # Assuming there is a service credential named 'testcred' set up in Unity Catalog
    boto3_session = boto3.Session(botocore_session=getServiceCredentialsProvider('testcred'))
    # Use the S3 session to perform operations

サービス資格情報のアクセス許可

セッションスコープの UDF は、呼び出し元の権限を使用します。必要な権限については、Python UDF でのサービス認証情報の使用を参照してください。

セッションスコープの UDF のコンピュートレベルのdefault認証情報

スカラー Python UDF で使用される場合、Databricks はコンピュート環境変数からdefaultのサービス認証情報を自動的に使用します。この動作により、UDF コード内で認証情報のエイリアスを明示的に管理することなく、外部サービスを安全に参照できます。コンピュートリソースのdefaultサービス認証情報の指定を参照してください

Default資格情報のサポートは、StandardおよびDedicatedアクセスモードのクラスターでのみ利用できます。SQLウェアハウスでは使用できません。

Python
@udf
def use_service_credential():
    from databricks.service_credentials import getServiceCredentialsProvider
    import boto3

    # The default service credential for the compute is automatically used
    boto3_session = boto3.Session()
    # Use the S3 client to perform operations

セッションスコープのPython UDF AWS Lambda関数の例

次の例では、サービス資格情報を使用して、Python UDFからAWS Lambda関数を呼び出します。実行される処理：

Databricksサービス資格情報プロバイダーを使用して、デフォルトの資格情報を取得します。
boto3セッションを設定します。
Lambda 関数を呼び出して、入力文字列を処理します。

Python
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

@udf(StringType())
def call_lambda_udf(input_str):
    import boto3
    import json
    import base64
    from databricks.service_credentials import getServiceCredentialsProvider
    from pyspark.taskcontext import TaskContext

    # Create a session using the default Unity Catalog service credential
    session = boto3.Session()
    client = session.client("lambda", region_name="us-west-2")

    # Optionally attach Spark TaskContext metadata to the Lambda request
    user_ctx = {"custom": {"user": TaskContext.get().getLocalProperty("user")}}

    # Build the Lambda payload
    payload = json.dumps({
        "values": [input_str],
        "is_debug": False
    })

    # Encode context for Lambda's client context
    encoded_ctx = base64.b64encode(json.dumps(user_ctx).encode("utf-8")).decode("utf-8")

    # Call the Lambda function
    response = client.invoke(
        FunctionName="HashValuesFunction",
        InvocationType="RequestResponse",
        ClientContext=encoded_ctx,
        Payload=payload,
    )

    response_payload = json.loads(response["Payload"].read().decode("utf-8"))

    if "errorMessage" in response_payload:
        raise Exception(response_payload["errorMessage"])

    return response_payload["values"][0]

スカラー Unity Catalog Python UDF でのサービス資格情報の使用

UDF 定義の CREDENTIALS 句でサービス資格情報を指定します。1 つの資格情報を DEFAULT としてマークし、パッチ適用済みクラウド SDK がそれを自動的に使用するようにすることができます。クラシックコンピュートでは、この機能には Databricks Runtime 18.1 以降が必要です。サーバーレスコンピュート、ならびに Pro および Serverless SQL Warehouse では、UDF の environment_version を 6 以上に明示的に設定します。コンピュート、ネットワーク、および権限の要件の詳細については、Python UDF でのサービス資格情報の使用を参照してください。

サービス資格情報の例：AWS Lambda 関数

次の例では、サービス資格情報を使用して、スカラー Unity Catalog Python UDFからAWS Lambda関数を呼び出します。この例では、Serverless環境バージョン6を使用します。

SQL
CREATE OR REPLACE FUNCTION main.test.call_lambda_func(data STRING, debug BOOLEAN)
RETURNS STRING
LANGUAGE PYTHON
CREDENTIALS (
  `scalar-uc-udf-service-creds-example-cred` DEFAULT
)
ENVIRONMENT (
  dependencies = '["boto3"]',
  environment_version = '6'
)
AS $$
import base64
import boto3
import json
from pyspark.taskcontext import TaskContext

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

user_context = {"custom": {"user": TaskContext.get().getLocalProperty("user")}}
payload = json.dumps({"values": [data], "is_debug": debug})

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

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

return response_payload["values"][0]
$$;

登録後にUDFを呼び出します。

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

タスク実行コンテキストを取得する

TaskContext PySpark API を使用して、ユーザーの ID、クラスタータグ、spark ジョブ ID などのコンテキスト情報を取得します。 UDF でタスクコンテキストを取得するを参照してください。

制限

PySpark UDF には、次の制限が適用されます。

ファイルアクセス制限: Databricks Runtime 14.2 以前では、共有クラスタリング上の PySpark UDF は、Git フォルダ、ワークスペースファイル、または Unity Catalog ボリュームにアクセスできません。
ブロードキャスト変数: 標準アクセスモードクラスターとサーバレスコンピュートの PySpark UDF は、ブロードキャスト変数をサポートしていません。
インスタンスプロファイル: 標準アクセスモードの PySpark UDF クラスターとサーバレスコンピュートはインスタンスプロファイルをサポートしていません。
サーバレスのメモリ制限 : サーバレスコンピュートの PySpark UDF には、 PySpark UDFあたり 1GB のメモリ制限があります。この制限を超えると、タイプ UDF_PYSPARK_USER_CODE_ERROR のエラーが発生します。MEMORY_LIMIT_SERVERLESS。
標準アクセスモードのメモリ制限 : 標準アクセスモードの PySpark UDF には、選択したインスタンスタイプの使用可能なメモリに基づいてメモリ制限があります。使用可能なメモリを超えると、タイプ UDF_PYSPARK_USER_CODE_ERROR のエラーが発生します。MEMORY_LIMIT。
サーバレスSQLウェアハウスのネットワークアクセス : 一応、サーバレスSQLウェアハウスのPython UDFは送信ネットワークリクエストを行うことができず、ネットワーク呼び出しを試みるクエリは無期限にハングします。アウトバウンドネットワークアクセスを有効にするには、ワークスペースの [プレビュー]ページ で、パブリックプレビュー機能[サーバーレスSQL ウェアハウスの分離されたワークロードのネットワークを有効にする] を有効にします。それ以外の場合は、ネットワークアクセスを必要とする UDF にはサーバレスコンピュートまたはクラシックコンピュートを使用します。

要件​

関数を UDF として登録する​

Spark SQL で UDF を呼び出す​

UDF と DataFrames の併用​

UDF を使用したバリアント​

UDF を含むファイル​

評価順序と null チェック​

Unity Catalog のシークレットへのアクセス​

Python UDF におけるサービス資格情報​

セッションスコープのスカラー Python UDF でのサービス認証情報の使用​

サービス資格情報のアクセス許可​

セッションスコープの UDF のコンピュートレベルのdefault認証情報​

セッションスコープのPython UDF AWS Lambda関数の例​

スカラー Unity Catalog Python UDF でのサービス資格情報の使用​

サービス資格情報の例：AWS Lambda 関数​

タスク実行コンテキストを取得する​

制限​

要件