Snowpark for Python UDF 性能検証レポート Part1（vs Vectorized UDF）

はじめに

今回は、下記の興味深いクイックスタートに取り組んでみたので紹介です。

上記のクイックスタートでは、まず Snowpark UDF、UDTF、Vectorized UDF/UDTF の数値処理・非数値処理における性能を確認しています。さらに、キャッシュ、ローカルテストについての確認も行っていますが、本記事では割愛します。

本記事では、第一弾として、Snowpark UDF vs Vectorized UDF について紹介します。
第二弾はこちら(to be updated...)

また、UDF を使った機械学習モデルによる推論について 1年半ほど前に LT 登壇しているため、そちらのスライドも展開しておきます。

結論

数値処理では Vectorized UDF、非数値処理ではスカラー UDF が良い可能性が高いよ。

検証

今回の検証は、数値計算、文字列操作、正規表現マスキング、タイムスタンプ操作というユースケースに対して行います。（クイックスタート内ではさらに、データ件数 1500万件、1億5000万件、ウェアハウスサイズ Medium、Large での検証も行っていましたが、本記事では詳細は割愛します。）

Snowpark UDF / Vectorized UDF

UDF と Vectorized UDF の概要

UDF とは、SQL から呼び出すことのできる"スカラー関数"です。ストアドプロシージャのように複数の処理を実行するものではなく、単一の値あるいは表形式の値を処理することができます。例えば、add_one(x) という UDF があったとしたら、SELECT add_one(col1) FROM TABLE; のように呼び出すことができ、各行の値に 1を加えます。

次に、Vectorized UDF は、通常の UDF が 1行に対して 1行を返していたのと挙動は同じですが、内部的には複数の行をより効率的に一括で処理しています。

この一括処理は、Pandas DataFrame が担っています。これにより高度に高速化された数値計算処理が行えるため、1行ずつ Python 関数を呼び出しながら処理するスカラー UDF よりも効率的な処理が可能になるわけです。

ここで"数値計算"と明に記載したのは、Vectorized UDF の場合、非数値計算で性能劣化の可能性が生じるためです。

今回の記事では、その点について検証を行っています。

UDF と Vectorized UDF の性能比較（数値計算）

数値計算の検証では、簡単な四則演算の計算を行います。入力行ごとに、あらかじめ計算しておいた customer テーブルの口座残高 C_ACCTBAL の平均値と標準偏差を足し引きし、最後に 10000 をかけます。

まず、下記のコードで Snowpark DataFrame を作成します。

df_customer_100 = session.table("SNOWFLAKE_SAMPLE_DATA.TPCH_SF100.customer")

また、事前計算として、平均値と標準偏差を求めます。

df_customer_100_mean = float(df_customer_100.agg(avg("C_ACCTBAL")).to_pandas().values[0])
df_customer_100_stddev = float(df_customer_100.agg(stddev("C_ACCTBAL")).to_pandas().values[0])

続いて、スカラー UDF と Vectorized UDF の定義を行います。

## df_customer_100
def basic_compute_100(inp):
    return (inp - df_customer_100_mean + df_customer_100_stddev) * 10000.0

# Let's create a UDF on top of that
udf_bc_100 = udf(basic_compute_100, return_type=FloatType(), input_types=[FloatType()])

# Let's vectorise the UDF
@udf()
def vect_udf_bc_100(inp: PandasSeries[float]) -> PandasSeries[float]:
    return (inp - df_customer_100_mean + df_customer_100_stddev) * 10000.0

ここで、UDF と Vectorized UDF の定義方法の違いを確認しておきましょう。基本的に、異なる部分は一箇所だけで、それは引数と返り値の型ヒントに PandasSeries を指定している点です。こうすることで、Vectorized UDF として実行されます。

それでは、準備が整ったので、実際に計算・実行時間の確認を行います。

st=datetime.now()
df_customer_100.select(udf_bc_100("C_ACCTBAL").alias("bal_from_mean")).agg(avg("bal_from_mean")).show()
et=datetime.now()
print(f"Scalar UDF:={(et-st).total_seconds()}")

st=datetime.now()
df_customer_100.select(vect_udf_bc_100("C_ACCTBAL").alias("bal_from_mean")).agg(avg("bal_from_mean")).show()
et=datetime.now()
print(f"Vectorized UDF:={(et-st).total_seconds()}")

上記の実行結果は次のようになりました。上半分がスカラー UDF の実行結果で、下半分が Vectorized UDF の実行結果です。概ね 2倍ほど高速に処理されていますね🎉

UDF と Vectorized UDF の性能比較（文字列操作）

続いては文字列操作です。

同様に UDF を定義します。今回は入力されたテキストを # で区切る処理です。

def str_manipulate(inp):
    return inp.split('#')[1]

# Let's create a UDF on top of that
udf_sm = udf(str_manipulate, return_type=StringType(), input_types=[StringType()])

# Let's vectorise the same UDF
@udf()
def vect_udf_sm(inp: PandasSeries[str]) -> PandasSeries[str]:
    return inp.str.split('#', expand=True)[1]

それでは実行してみましょう。

st=datetime.now() 
df_customer_100.select(udf_sm("C_NAME").alias("CustIDs")).write.mode("overwrite").save_as_table("SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs")
et=datetime.now()
print(f"Scalar UDF:={(et-st).total_seconds()}")

st=datetime.now() 
df_customer_100.select(vect_udf_sm("C_NAME").alias("CustIDs")).write.mode("overwrite").save_as_table("SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs")
et=datetime.now()
print(f"Vectorized UDF:={(et-st).total_seconds()}")

ここで、先ほどのコードとは違い、show() ではなく、save_as_table() を使用しています。これは、show() や collect() メソッドだと、結果をローカルに持ってきてしまうため非効率な IO が発生するためだと考えられます。

上記の実行結果を下図に示します。

今回は、スカラー UDF のほうが遅くなってしまいましたね。こちらはきちんと再現性のある実行結果です。ではなぜこのような結果になったのでしょうか？

それは、Python や Pandas がテキストの処理を高速に行うことができないため、バッチ化した分遅くなってしまったからです。テキストの処理であれば、並列分散で処理されるスカラー UDF のほうが高速に行えるということです。

これは、きちんと覚えておきたいポイントですね！私も何となくそうだろうとは予測していましたが、しっかり計測したことはなかったので非常に参考になりました。

ちなみに、数値計算であれば高速に行える理由は、Pandas や Numpy がよりマシンレベルにベクトル処理してくれているからです。この速度が、Python Sandbox を都度呼び出すよりは高速だということですね。

UDF と Vectorized UDF の性能比較（正規表現）

それでは、正規表現による処理も確認しましょう。ここまでの内容が理解できれば、結果も予測できそうです。

正規表現の UDF の定義です。

# Let's create a slightly complex function using native python regex for string replacement/masking
def mask_data(inp):
    return re.sub('\d{4}$', '****', inp)

# Let's create a UDF on top of that
udf_md = udf(mask_data, return_type=StringType(), input_types=[StringType()])

# Let's vectorise the same UDF
@udf()
def vect_udf_md(inp: PandasSeries[str]) -> PandasSeries[str]:
    return inp.apply(lambda x: mask_data(x))

それでは、実行してみましょう。

st=datetime.now()
df_customer_100.select(udf_md("C_PHONE").alias("masked_phone_nums")).write.mode("overwrite").save_as_table("SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data")
et=datetime.now()
print(f"Scalar UDF:={(et-st).total_seconds()}")

st=datetime.now()
df_customer_100.select(vect_udf_md("C_PHONE").alias("masked_phone_nums")).write.mode("overwrite").save_as_table("SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data")
et=datetime.now()
print(f"Vectorized UDF:={(et-st).total_seconds()}")

下図に実行結果を示します。やはり、スカラー UDF のほうが高速でした。特に、Pandas DataFrame の apply() メソッドは内部的にはただの for 文と変わらないため、注意が必要ですね。

UDF と Vectorized UDF の性能比較（タイムスタンプ）

最後に、タイムスタンプ操作についても確認します。

タイムスタンプ列の値に 10日を加える UDF の定義です。

# Let's try some timestamp manipulation
def change_format(inp):
    instantiate = datetime.strptime(inp, '%Y-%m-%d')
    change_format = datetime.strptime(instantiate.strftime('%m/%d/%Y'), '%m/%d/%Y')
    dt_add10_days = change_format + timedelta(days=10)
    return dt_add10_days

# Let's create a UDF on top of that
udf_cf = udf(change_format, return_type=StringType(), input_types=[StringType()])

# Let's vectorise the UDF - still using native python
@udf()
def vect_udf_cf(inp: PandasSeries[str]) -> PandasSeries[str]:
    return inp.apply(lambda x: change_format(x))

同様に、実行してみます。

st=datetime.now() 
df_order_100.select(udf_cf("O_ORDERDATE").alias("NY_OrderDate")).write.mode("overwrite").save_as_table("SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates")
et=datetime.now()
print(f"Scalar UDF:={(et-st).total_seconds()}")

st=datetime.now()  
df_order_100.select(vect_udf_cf("O_ORDERDATE").alias("NY_OrderDate")).write.mode("overwrite").save_as_table("SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates")
et=datetime.now()
print(f"Vectorized UDF:={(et-st).total_seconds()}")

下図に実行結果を示します。

こちらも同様に、スカラー UDF のほうが高速で、Vectorized UDF のほうが遅くなってしまいました。

検証のまとめ

今回の検証から、スカラー UDF と Vectorized UDF は、適材適所で活用するべきということが理解いただけたかと思います。最も簡単には、よりネイティブにベクトル処理が行える数値計算については Vectorized UDF が適しているが、内部的には for 文となってしまう処理についてはスカラー UDF のほうが適している可能性が高いということです。

ただし、ユースケースによっては変化しうるものだと思いますので、非常に重たい処理などについては慎重に処理方式を検討するようにしてください。

おわりに

今回は、ユースケースによって UDF と Vectorized UDF で性能に差が生じるという、意外だけど非常に重要な結果の得られた検証でした。

こういった差をしっかりと抑えて、楽しい Snowpark ライフを送りましょう！🏃

第二弾の検証もお楽しみに！第二弾では、UDTF や Vectorized UDTF、さらには UDTF 内でのマルチスレッド利用について検討しています。