Databricks Connect로 외부 Spark 실행하기

Databricks pushdown은 Code Repositories에서 PySpark DataFrame 코드를 작성하고 Databricks의 계산 자원에서 실행하는 방식입니다. 입력과 출력은 모두 같은 Databricks 소스에 등록된 가상 테이블이어야 합니다. 일반 Foundry 데이터셋 경로를 TableInput에 넣는 것만으로 가상 테이블로 바뀌지는 않습니다.

환경 준비

가상 테이블 공통 설정을 적용하고 transforms-tables와 databricks-connect를 설치합니다. Databricks Connect는 PyPI 패키지이며 Conda 패키지로 설치하지 않습니다. 저장소의 Libraries에서 추가하거나 기존 conda_recipe/meta.yaml의 pip 요구사항에 넣습니다.

서버리스에서 16.4 계열을 사용할 때는 호환 패치인 16.4.1 이상과 Python 3.12를 사용합니다. 16.4.0의 서버리스 환경은 호환 Python 조합이 없으므로 선택하지 않습니다.

requirements:
  pip:
    - databricks-connect >=16.4.1,<17.0.0

Classic cluster를 사용하면 cluster의 Databricks Runtime이 Connect 패키지 버전 이상이어야 하며, 해당 Runtime에 맞는 Connect 버전을 선택합니다. 16.4 cluster도 Python 3.12를 사용합니다. UDF를 실행할 때는 클라이언트와 외부 계산 환경의 Python minor 버전이 같아야 합니다.

Databricks 계정과 workspace에는 Unity Catalog가 필요합니다. 서버리스는 해당 workspace에서 서버리스를 사용할 수 있어야 하고, cluster는 지원되는 Assigned/Shared 접근 모드를 사용해야 합니다. Databricks Connect는 일반 오픈소스 Spark 서버에 연결하는 클라이언트가 아닙니다.

기존 Spark 논리 옮기기

다음은 id > 1인 행을 남기는 Foundry Spark 변환입니다.

from pyspark.sql.functions import col
from transforms.api import Input, Output, transform


@transform.spark.using(
    source=Input("/Example Project/Inputs/foundry_table"),
    output=Output("/Example Project/Outputs/foundry_filtered"),
)
def compute(source, output):
    output.write_dataframe(source.dataframe().filter(col("id") > 1))

같은 로직을 Databricks에서 실행하려면 가상 테이블 입력과 출력을 연결하고 Databricks 선언을 사용합니다. 아래 경로와 RID 및 외부 이름은 가상의 값입니다.

from pyspark.sql.functions import col
from transforms.api import transform
from transforms.tables import (
    DatabricksInput,
    DatabricksOutput,
    TableInput,
    TableOutput,
)

SOURCE_RID = "ri.magritte..source.00000000-0000-0000-0000-000000000001"


@transform.databricks.using(
    source=TableInput("/Example Project/Inputs/databricks_table"),
    output=TableOutput(
        "/Example Project/Outputs/databricks_filtered",
        SOURCE_RID,
        "CATALOG.SCHEMA.FILTERED_TABLE",
    ),
)
def compute(source: DatabricksInput, output: DatabricksOutput):
    output.write_dataframe(source.dataframe().filter(col("id") > 1))

DatabricksInput.dataframe()은 원격 PySpark DataFrame을, DatabricksOutput.write_dataframe()은 외부 테이블 쓰기를 제공합니다. 위치는 Unity Catalog의 세 단계 식별자인 catalog.schema.table 문자열이나 DatabricksTable(catalog, schema, table)로 지정합니다.

입력과 출력에 TableInput/TableOutput 이외의 자산을 섞을 수 없으며 @incremental 증분 처리도 지원하지 않습니다.

출력 테이블의 세 유형

format과 location은 DatabricksTable에 지정하거나 지원되는 write_dataframe() 인자로 전달합니다. 다음 예제는 서로 다른 Foundry 출력 경로를 사용합니다.

from transforms.api import transform
from transforms.tables import DatabricksTable, TableInput, TableOutput

SOURCE_RID = "ri.magritte..source.00000000-0000-0000-0000-000000000001"


@transform.databricks.using(
    source=TableInput("/Example Project/Inputs/databricks_table"),
    managed_delta=TableOutput(
        "/Example Project/Outputs/managed_delta", SOURCE_RID,
        "CATALOG.SCHEMA.MANAGED_DELTA_TABLE",
    ),
    managed_iceberg=TableOutput(
        "/Example Project/Outputs/managed_iceberg", SOURCE_RID,
        "CATALOG.SCHEMA.MANAGED_ICEBERG_TABLE",
    ),
    external_delta=TableOutput(
        "/Example Project/Outputs/external_delta", SOURCE_RID,
        DatabricksTable(
            catalog="CATALOG", schema="SCHEMA", table="EXTERNAL_DELTA_TABLE",
            location="s3://example-bucket/path/to/table",
        ),
    ),
)
def compute(source, managed_delta, managed_iceberg, external_delta):
    rows = source.dataframe()
    managed_delta.write_dataframe(rows)
    managed_iceberg.write_dataframe(rows, format="iceberg")
    external_delta.write_dataframe(rows)

외부 경로는 해당 소스와 Unity Catalog가 접근할 수 있는 실제 저장 위치여야 합니다. 생성 뒤 소스나 locator를 바꾸는 것은 지원하지 않습니다.

서버리스와 cluster 선택

기본 Databricks Connect 세션은 서버리스 계산을 사용합니다. 기존 databricks(serverless=True) 설정과 같은 선택입니다. 특정 cluster를 쓰려면 선언에 .with_compute(cluster_id="실제-cluster-id")를 연결합니다. cluster ID는 Databricks에서 선택한 대상의 실제 식별자로 지정합니다.

UDF 실행

UDF는 Connect가 직렬화하여 원격 환경으로 보냅니다. 다음 예제는 id의 제곱을 doubled 열에 넣습니다. 반환형은 IntegerType이므로 제곱 결과가 그 범위에 들어가는 입력에 사용하며 NULL은 NULL로 유지합니다.

from pyspark.sql.functions import col, udf
from pyspark.sql.types import IntegerType
from transforms.api import transform
from transforms.tables import TableInput, TableOutput

SOURCE_RID = "ri.magritte..source.00000000-0000-0000-0000-000000000001"


@transform.databricks.using(
    source=TableInput("/Example Project/Inputs/databricks_table"),
    output=TableOutput(
        "/Example Project/Outputs/squared_ids", SOURCE_RID,
        "CATALOG.SCHEMA.SQUARED_IDS",
    ),
)
def compute(source, output):
    @udf(returnType=IntegerType())
    def double(value):
        return None if value is None else value * value

    output.write_dataframe(source.dataframe().withColumn("doubled", double(col("id"))))

함수와 필요한 계산을 transform 안에 정의하면 원격에 없는 myproject 모듈의 함수를 이름으로 참조하는 문제를 줄일 수 있습니다.

UDF의 외부 Python 의존성

.with_dependencies(dependencies=[...])는 UDF가 실행되는 Databricks Python 환경에 추가 패키지를 설치하도록 지정합니다. 이 기능은 Databricks Connect/Runtime 16.4 이상에서 제공되는 Public Preview이며 workspace의 Enhanced Python UDFs in Unity Catalog 기능이 활성화되어 있어야 합니다. 서버리스의 16.3 이하에서는 사용자 패키지를 포함한 UDF를 사용할 수 없습니다.

다음은 원격 환경에 pyjokes<1을 설치하고 1부터 9까지의 ID에 실제 라이브러리가 생성한 문장을 붙입니다. 함수 내부에서 원격 패키지를 import하는 구성이므로 Python 하위 프로젝트의 TLLV를 비활성화합니다.

transformsPython {
    tllv false
}

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
from transforms.api import transform
from transforms.tables import TableOutput

SOURCE_RID = "ri.magritte..source.00000000-0000-0000-0000-000000000001"


@transform.databricks.using(
    output=TableOutput(
        "/Example Project/Outputs/jokes", SOURCE_RID,
        "CATALOG.SCHEMA.JOKES",
    ),
).with_dependencies(dependencies=["pyjokes<1"])
def compute(output):
    @udf(returnType=StringType())
    def make_joke():
        from pyjokes import get_joke
        return get_joke()

    rows = output.spark_session.range(1, 10)
    output.write_dataframe(rows.withColumn("jokes", make_joke()))

저장소에만 있는 함수나 클래스를 UDF에서 참조하면 ModuleNotFoundError: No module named 'myproject'가 발생할 수 있습니다. 필요한 논리를 UDF에 직접 넣거나, UDF와 보조 코드를 transform 안에 함께 정의하거나, UDF를 만들어 반환하는 함수를 transform 안에서 호출하는 방식으로 직렬화 범위를 구성합니다. 원격 패키지 의존성 설치와 로컬 프로젝트 모듈의 전달은 다른 문제입니다.

PySpark 기능의 차이

Databricks Connect는 일반 Spark의 모든 기능을 제공하지 않습니다. SparkContext, RDD 및 Spark JVM에 직접 접근하는 라이브러리는 사용할 수 없습니다. dataframe.display()와 일부 Databricks Utilities의 credentials, library, notebook workflow, widgets도 지원하지 않습니다. 분산 ML 학습과 로컬 환경의 자동 cluster 동기화도 지원하지 않습니다.

CREATE TABLE ... AS SELECT 대신 DataFrame으로 조회한 뒤 write.saveAsTable() 계열을 사용하는 제약도 있습니다. Foundry 가상 출력은 위 예제의 write_dataframe() 계약을 사용합니다. 구버전별로는 13.3 이하의 streaming foreachBatch, 128MB를 초과하는 미해결 논리 계획과 3,600초 초과 쿼리, 15.3 이하 Standard 접근 모드의 ApplyInPandas/Cogroup에 추가 제한이 있습니다. 128MB는 실제 데이터 크기와 구분되는 계획 크기입니다.