DuckDB에서 Ibis와 SQLFrame 사용하기

Ibis는 여러 데이터베이스 백엔드에 DataFrame 표현식을 적용하는 라이브러리이고, SQLFrame은 PySpark와 유사한 DataFrame 코드를 SQL 엔진에서 실행하도록 변환하는 라이브러리입니다. 두 라이브러리는 Palantir와 별개의 제삼자 의존성이며 Foundry의 제품 지원 보장이 적용되지 않습니다. 사용하려는 연산과 설치한 버전의 지원 범위를 확인합니다.

Ibis 연결

Code Repositories의 패키지 환경에 ibis-framework와 ibis-duckdb를 추가합니다. Foundry가 준비한 DuckDB 연결을 ibis.duckdb.from_connection()에 넘기면 이미 등록된 입력을 사용할 수 있습니다.

다음 예제는 가상의 입력을 Ibis 테이블 표현식으로 읽고 DuckDB SQL로 변환하여 출력에 복사합니다.

import ibis
from transforms.api import Input, Output, transform


@transform.using(
    source=Input("/Example Project/Inputs/source"),
    output=Output("/Example Project/Outputs/ibis_copy"),
)
def compute(ctx, source, output):
    connection = ctx.duckdb().conn
    ibis_connection = ibis.duckdb.from_connection(connection)
    table = ibis_connection.table("source")
    sql = ibis.to_sql(table, dialect="duckdb")
    output.write_table(connection.sql(sql))

Ibis는 약 20개 백엔드를 다루는 DataFrame API를 제공하지만, 특정 연산을 다른 백엔드에서도 동일하게 실행할 수 있는지는 해당 백엔드의 구현을 확인해야 합니다.

SQLFrame으로 Spark 계산 옮기기

환경에 sqlframe을 추가합니다. 입력은 Foundry 연결에 decorator의 입력 이름으로 등록되어 있으므로 session.table("입력_이름")으로 읽습니다. 계산을 마친 DataFrame의 sql(optimize=True, dialect="duckdb") 결과를 DuckDB 연결에서 실행하고 write_table()로 씁니다.

다음 Spark 변환은 운석의 질량이 같은 분류의 평균보다 큰 행을 출력합니다. meteorite_landings에는 name, class, mass가 있고 meteorite_stats에는 class, avg_mass_per_class가 있어야 합니다. mass와 평균 열은 비교 가능한 숫자형이어야 합니다.

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


@transform.spark.using(
    output=Output("/Example Project/Outputs/meteorite_enriched_spark"),
    meteorite_landings=Input("/Example Project/Inputs/meteorite_landings"),
    meteorite_stats=Input("/Example Project/Inputs/meteorite_stats"),
)
def compute(output, meteorite_landings, meteorite_stats):
    landings = meteorite_landings.dataframe()
    stats = meteorite_stats.dataframe()
    result = (
        landings.join(stats, "class")
        .withColumn(
            "greater_mass", F.col("mass") > F.col("avg_mass_per_class")
        )
        .filter(F.col("greater_mass"))
        .select("name", "class", "mass", "greater_mass")
    )
    output.write_dataframe(result)

이를 단일 노드 DuckDB로 옮길 때 다음 부분이 달라집니다.

SQLFrame을 직접 가져오는 방식은 PySpark 모듈을 전역 교체하지 않고 사용합니다.

from sqlframe.duckdb import DuckDBSession
from sqlframe.duckdb import functions as F
from transforms.api import Input, Output, transform


@transform.using(
    output=Output("/Example Project/Outputs/meteorite_enriched_duckdb"),
    meteorite_landings=Input("/Example Project/Inputs/meteorite_landings"),
    meteorite_stats=Input("/Example Project/Inputs/meteorite_stats"),
)
def compute(ctx, output, meteorite_landings, meteorite_stats):
    connection = ctx.duckdb().conn
    session = DuckDBSession(conn=connection)
    landings = session.table("meteorite_landings")
    stats = session.table("meteorite_stats")
    result = (
        landings.join(stats, "class")
        .withColumn(
            "greater_mass", F.col("mass") > F.col("avg_mass_per_class")
        )
        .filter(F.col("greater_mass"))
        .select("name", "class", "mass", "greater_mass")
    )
    sql = result.sql(optimize=True, dialect="duckdb")
    output.write_table(connection.sql(sql))

조인과 필터의 논리를 유지해도 실행 엔진은 Spark에서 단일 노드 DuckDB로 바뀝니다. Spark 전용 기능이나 임의의 Python UDF까지 자동 이식된다고 가정하지 않습니다. 이 예제처럼 SQLFrame이 지원하는 DataFrame 연산을 사용합니다.

PySpark import를 유지하는 activate 방식

SQLFrame의 activate()로 PySpark import를 전환하는 방식도 사용할 수 있습니다. 이 경우 activate("duckdb", conn=connection)를 먼저 호출하고 그 뒤에 pyspark를 가져와야 합니다.

함수 내부 import와 다른 모듈의 교체를 사용하는 구성이므로 Python 하위 프로젝트 build.gradle에서 TLLV를 비활성화합니다.

transformsPython {
    tllv false
}

다음은 등록된 입력을 SQLFrame DataFrame으로 읽어 복사하는 전체 변환입니다.

from sqlframe import activate
from transforms.api import Input, Output, transform


@transform.using(
    source=Input("/Example Project/Inputs/source"),
    output=Output("/Example Project/Outputs/sqlframe_copy"),
)
def compute(ctx, source, output):
    connection = ctx.duckdb().conn
    activate("duckdb", conn=connection)

    from pyspark.sql import SparkSession

    session = SparkSession.builder.getOrCreate()
    result = session.table("source")
    sql = result.sql(optimize=True, dialect="duckdb")
    output.write_table(connection.sql(sql))

운석 계산에도 이 방식을 적용할 수 있습니다. 활성화 뒤에 SparkSession과 functions as F를 함께 가져오고, 위의 조인부터 select까지 같은 연산을 수행합니다. 활성화 전에 다른 위치에서 PySpark를 가져오지 않도록 import 순서를 관리합니다. 직접 SQLFrame을 가져오는 구성에서는 모듈 수준 import를 유지할 수 있습니다.