DuckDB API

Foundry は、Polars lazy API と同様に、DuckDB を使用したストリーミングによる遅延実行をサポートしています。

@transform.using(...) で定義したトランスフォームは、デフォルトで Lightweight トランスフォームになるため、@lightweight デコレーターは不要です。注入される ctx は LightweightContext であり、出力は path_for_write_table や path_for_object_store_write_table などのプロパティを持つ LightweightOutput オブジェクトです。

トランスフォーム内で事前設定済みの DuckDB 接続にアクセスするには、以下に示すようにコンテキストオブジェクトを使用します。

Copied!
1
2
3
4
5
6
7
8
9
10
from transforms.api import transform, Input, Output

@transform.using(
    output=Output("/Users/jsmith/output"),
    my_input=Input("/Users/jsmith/input"),
)
def compute(ctx, output, my_input):
    conn = ctx.duckdb().conn
    query = conn.sql("""SELECT * FROM my_input""")
    output.write_table(query)

DuckDB 接続では、すべての入力データセットが、入力名をテーブル名として、マテリアライズされていない DuckDB テーブルとして登録済みです。これにより、手動で設定することなく入力データに簡単にアクセスできます。

DuckDB を使用する場面

DuckDB は高性能で効率的なクエリエンジンであり、多くのワークロードで Pandas や Polars などのほかのシングルノードコンピューティングエンジンを上回るパフォーマンスを発揮することがよくあります。次のような場合は、DuckDB の使用を検討してください。

• 非常に大規模なデータを処理するシングルノードのトランスフォーム。
• メモリー使用量を厳格に制御する必要があるトランスフォーム。

ただし、DuckDB は Polars や Pandas のような Python DataFrame API をまだサポートしておらず、データ操作には生の SQL クエリが必要です。Ibis などのサードパーティーライブラリを使用すると、DuckDB 上でデータフレーム API を利用できます。以下の DuckDB で Ibis を使用するを参照してください。

インクリメンタル（差分処理）ワークフロー

ほかのコンピューティングエンジンと同様に、DuckDB API では入力のインクリメンタル読み取りモードを設定できます。

ほかのコンピューティングエンジンと同様に、インクリメンタルに実行する場合のデフォルトは added、 インクリメンタルではない実行の場合のデフォルトは current です。

Copied!
1
2
3
4
5
6
7
8
9
10
11
from transforms.api import transform, Input, Output, incremental

@incremental()
@transform.using(
    output=Output("/Users/jsmith/output"),
    my_input=Input("/Users/jsmith/input"),
)
def compute(ctx, output, my_input):
    ddb = ctx.duckdb(read_modes={my_input: "current"})
    result = ddb.conn.sql("""SELECT * FROM my_input""")
    output.write_table(result)

.sql ファイルから読み取る

DuckDB のワークフローによっては、別の .sql ファイルで管理する方が適した複雑な SQL クエリを使用する場合があります。DuckDB 接続の sql_from_file メソッドを使用して、このようなファイルを参照できます。パスは、sql_from_file の呼び出しを含むファイルからの相対パスです。

Copied!
1
2
# -- query.sql
SELECT * FROM my_input WHERE column_a > 100;

Copied!
1
2
3
4
5
6
7
8
9
10
11
from transforms.api import transform, Input, Output, incremental

@incremental()
@transform.using(
    output=Output("/Users/jsmith/output"),
    my_input=Input("/Users/jsmith/input"),
)
def compute(ctx, output, my_input):
    conn = ctx.duckdb().conn
    query = conn.sql_from_file("query.sql")
    output.write_table(query)

DuckDB で Ibis を使用する

Ibis ↗ は、20近くの異なるバックエンドで使用できるポータブルなデータフレーム API を提供する Python ライブラリです。DuckDB もサポートされており、本来 SQL を使用するバックエンド上でデータフレーム API を利用できます。

Ibis を使用するには、まず環境に ibis-framework および ibis-duckdb パッケージをインストールします。リポジトリへのパッケージの追加については、パッケージタブのドキュメントを参照してください。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import ibis
from transforms.api import Input, Output, transform, LightweightInput, LightweightOutput, LightweightContext

@transform.using(
    output=Output("/Users/jsmith/output"),
    my_input=Input("/Users/jsmith/input"),
)
def compute(ctx, output, my_input) -> None:
    conn = ctx.duckdb().conn
    ibis_conn = ibis.duckdb.from_connection(conn)
    table = ibis_conn.table("my_input")
    # Ibis データフレームの変換
    sql = ibis.to_sql(table, dialect="duckdb")
    output.write_table(conn.sql(sql))

Ibis のデータフレーム変換の記述について詳しくは、Ibis の公式ドキュメント ↗ を参照してください。

DuckDB で SQLFrame を使用する

SQLFrame ↗ は、DuckDB を含む SQL ベースのクエリエンジンで PySpark DataFrame API のコードを実行できる Python ライブラリです。

SQLFrame を使用するには、まず環境に sqlframe パッケージをインストールします。リポジトリへのパッケージの追加については、パッケージタブのドキュメントを参照してください。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from sqlframe import activate
from transforms.api import Input, Output, transform, LightweightInput, LightweightOutput, LightweightContext

@transform.using(
    output=Output("/Users/jsmith/output"),
    my_input=Input("/Users/jsmith/input"),
)
def compute(ctx, output, my_input) -> None:
    conn = ctx.duckdb().conn

    activate("duckdb", conn=conn)

    from pyspark.sql import SparkSession
    # その他の pyspark のインポートは、activate の後に行う必要があります

    session = SparkSession.builder.getOrCreate()
    input_df = session.table("my_input")  # 入力はエイリアスで参照します

    # Pyspark データフレームの変換

    sql = result_df.sql(optimize=True, dialect="duckdb")
    output.write_table(conn.sql(sql))

SQLFrame のデータフレーム変換の記述について詳しくは、SQLFrame の公式ドキュメント ↗ を参照してください。

SQLFrame を使用して Spark トランスフォームを DuckDB に移行する

既存の Spark トランスフォームをリファクタリングして、SQLFrame を使用して DuckDB 上で実行できます。このセクションでは、以下の Spark の例をリファクタリングします。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from transforms.api import transform, Input, Output
from pyspark.sql import functions as F

@transform.spark.using(
    output=Output("/Users/jsmith/meteorite_enriched"),
    meteorite_landings=Input("/Users/jsmith/meteorite_landings"),
    meteorite_stats=Input("/Users/jsmith/meteorite_stats"),
)
def enriched(output, meteorite_landings, meteorite_stats):
    df_landings = meteorite_landings.dataframe()
    df_stats = meteorite_stats.dataframe()

    enriched_together = df_landings.join(df_stats, "class")
    greater_mass = enriched_together.withColumn(
        'greater_mass', F.col("mass") > F.col("avg_mass_per_class")
    )
    result = greater_mass.filter(F.col("greater_mass")).select(
        F.col("name"),
        F.col("class"),
        F.col("mass"),
        F.col("greater_mass")
    )
    output.write_dataframe(result)

まず、デフォルトでシングルノードのトランスフォームになるようにデコレーターを変更し、compute 関数に LightweightContext を追加します。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
- from transforms.api import transform, Input, Output
+ from transforms.api import transform, Input, Output, LightweightContext
+ from sqlframe import activate
  from pyspark.sql import functions as F

- @transform.spark.using(
+ @transform.using(
      output=Output("/Users/jsmith/meteorite_enriched"),
      meteorite_landings=Input("/Users/jsmith/meteorite_landings"),
      meteorite_stats=Input("/Users/jsmith/meteorite_stats"),
  )
- def enriched(output, meteorite_landings, meteorite_stats):
+ def enriched(ctx: LightweightContext, output, meteorite_landings, meteorite_stats):

次に、DuckDB 接続を使用して SQLFrame を有効にし、PySpark のインポートを compute 関数内に移動します。

Copied!
1
2
3
4
5
6
7
8
9
10
  from transforms.api import transform, Input, Output, LightweightContext
  from sqlframe import activate
- from pyspark.sql import functions as F

  def enriched(ctx: LightweightContext, output, meteorite_landings, meteorite_stats):
+     conn = ctx.duckdb().conn
+     activate("duckdb", conn=conn)
+
+     # 有効化した後に PySpark をインポートします
+     from pyspark.sql import SparkSession, functions as F

これで、SQLFrame の SparkSession から入力 DataFrame を生成できます。

Copied!
1
2
3
4
-     df_landings = meteorite_landings.dataframe()
-     df_stats = meteorite_stats.dataframe()
+     df_landings = session.table("meteorite_landings")  # デコレーターのエイリアスを使用します
+     df_stats = session.table("meteorite_stats")

最後に、結果を SQL に変換して出力として書き込みます。

Copied!
1
2
3
4
      result = greater_mass.filter("greater_mass")
-     output.write_dataframe(result)
+     sql = result.sql(optimize=True, dialect="duckdb")
+     output.write_table(conn.sql(sql))

リファクタリングしたトランスフォームの全体を以下に示します。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from transforms.api import transform, Input, Output, LightweightContext
from sqlframe import activate

@transform.using(
    output=Output("/Users/jsmith/meteorite_enriched"),
    meteorite_landings=Input("/Users/jsmith/meteorite_landings"),
    meteorite_stats=Input("/Users/jsmith/meteorite_stats"),
)
def enriched(ctx: LightweightContext, output, meteorite_landings, meteorite_stats):
    conn = ctx.duckdb().conn
    activate("duckdb", conn=conn)

    from pyspark.sql import SparkSession, functions as F

    session = SparkSession.builder.getOrCreate()

    df_landings = session.table("meteorite_landings")
    df_stats = session.table("meteorite_stats")

    # PySpark DataFrame の変換はそのまま動作します
    enriched_together = df_landings.join(df_stats, "class")
    greater_mass = enriched_together.withColumn(
        'greater_mass', F.col("mass") > F.col("avg_mass_per_class")
    )
    result = greater_mass.filter(F.col("greater_mass")).select(
        F.col("name"),
        F.col("class"),
        F.col("mass"),
        F.col("greater_mass")
    )

    # SQL に変換して書き込みます
    sql = result.sql(optimize=True, dialect="duckdb")
    output.write_table(conn.sql(sql))

DuckDB の設定

ほかの多くのシングルノードコンピューティングエンジンとは異なり、DuckDB はメモリー使用量と並列度を制御するためのリソース設定をサポートしており、さまざまなワークロードに合わせて細かく最適化できます。これは、メモリーが制限された状況で特に重要です。DuckDB は、パフォーマンスを犠牲にする代わりに、メモリー消費量を自ら制限してメモリー不足エラーを回避できます。

これらのオプションは、コンテキストオブジェクトの duckdb メソッドを使用して DuckDB 接続を初期化するときに設定できます。設定オプションの完全な一覧は、公式ドキュメント ↗ を参照してください。

Copied!
1
2
3
conn = ctx.duckdb().conn
conn.execute("SET memory_limit='2GB';")
conn.execute("SET threads=4;")

高度なユースケース

DuckDB のワークフローによっては、入力データセットを単に参照するだけでなく、低レベル API へのアクセスが必要になります。例として、次のものがあります。

• DuckDB の read_parquet ↗ または read_csv ↗ メソッドの設定を必要とするカスタム書式
• パーティション分割や圧縮オプションなどのカスタム書き込み設定 ↗

これらのユースケースでは、<input_name>_files() という設定済みの DuckDB UDF を使用して、データセットの生の Parquet または CSV フィールドを読み取れます。 この関数は、入力データセットの基になるデータファイルのファイルパスのリストを返します。このリストをカスタムパラメーターとともに DuckDB の読み取り関数に渡せます。

その後、出力をディスク上の中間的な場所にコピーするか、DuckDB の COPY TO コマンドを使用して出力を Foundry に直接ストリーミングできます。

このセクションの例では、LightweightOutput の output.path_for_write_table と output.path_for_object_store_write_table を使用します。

write_table を呼び出さずに COPY TO で手動アップロードしたファイルからスキーマを推論するには、Output オブジェクトの put_metadata() メソッドを使用できます。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import uuid

from transforms.api import transform, Input, Output

@transform.using(
    output=Output("/Users/jsmith/output"),
    my_input=Input("/Users/jsmith/input"),
)
def compute(ctx, output, my_input):
    conn = ctx.duckdb().conn

    # ディスクに書き込み、その後ディスクからデータセットをアップロードします
    query = conn.sql(f"""
    COPY (
        SELECT *
        FROM read_parquet(my_input_files(), file_row_number=True)
        WHERE my_column = 'abc'
        LIMIT 10
    ) TO '{output.path_for_write_table}'
    (FORMAT 'parquet', PER_THREAD_OUTPUT TRUE)
    """)
    output.write_table(output.path_for_write_table)

    # 入力と出力を直接ストリーミングします
    path_uuid = str(uuid.uuid4())
    conn.execute(f"""
        COPY (
            SELECT *
            FROM read_parquet(my_input_files(), file_row_number=True)
            WHERE my_column = 'abc'
            LIMIT 10
        ) TO '{output.path_for_object_store_write_table}'
        (FORMAT PARQUET, PARTITION_BY (date), WRITE_PARTITION_COLUMNS true, FILENAME_PATTERN 'file_{path_uuid}')
        """
    )
    output.put_metadata()

インクリメンタル（差分処理）出力

ユーザーは COPY TO 構文を使用して、出力ファイルの名前も含め、データセットへの出力を直接制御できます。インクリメンタルワークフローでは、ファイル間の競合や上書きを避けるため、連続する書き込みで過去のトランザクションと同じファイル名を使用しないようにしてください。DuckDB では出力にファイル名パターン ↗を指定できます。一意性を確保するため、上の例のようにファイル名に UUID を追加することを推奨します。

パーティション分割された出力

DuckDB のパーティション分割の動作は、Polars や Spark などの一般的なライブラリのデフォルトとは異なります。パーティション分割されたデータセットを書き込む場合は、ほかのトランスフォームとの互換性を確保するため、上の例のように WRITE_PARTITION_COLUMNS を true に設定することを強く推奨します。また、DuckDB は Hive ファイルパス内の欠損値に対して、Hive 標準の __HIVE_DEFAULT_PARTITION__ ではなく文字列 null を使用します。そのため、下流のトランスフォームがこれらのパーティションを正しくパースするよう、特に注意してください。通常は、DuckDB で出力を書き込む前にパーティション列の null 値を補完します。

一括ダウンロード

デフォルトでは、Foundry の DuckDB 統合は、パフォーマンスとリソース使用量を最適化するために、入力データセットの遅延ストリーミングダウンロードを使用します。ただし、シナリオによっては、トランスフォームの開始時にすべての入力データをディスクに一括ダウンロードしたい場合があります。入力で path() を呼び出すと、ディスクにダウンロードして入力をクエリできます。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import duckdb
from transforms.api import transform, Output, Input


@transform.using(my_input=Input('my-input'), my_output=Output('my-output'))
def my_duckdb_transform(my_input, my_output):
    duckdb.connect(database=':memory:').execute(f"""
        COPY
        (
            SELECT *
            FROM parquet_scan('{my_input.path()}/**/*.parquet')
            WHERE Name LIKE 'John%'
        )
        TO '{my_output.path_for_write_table}'
        (FORMAT 'parquet', PER_THREAD_OUTPUT TRUE)
    """)  # スレッドごとに個別の Parquet ファイルを並列で書き込み、パフォーマンスを最適化します

    my_output.write_table(my_output.path_for_write_table)