🐍

pysparkのDataFrameのサイズが0かを判定する方法の比較検証

2025/01/25に公開

結論

len(DataFrame.take(1)) がいいぞ。

背景

記事の概要

pysparkのDataFrameのサイズが0かを判定する方法をいくつか試して実行速度を確認する

扱う内容

  • 後述する判定方法の実行速度比較

扱わない内容

  • メモリ使用量など、実行時間以外の性能評価

pyspark で DataFrame のサイズが 0 かを判定する主な方法

ググったりChatGPTに聞いた結果、以下が出てきた。

性能比較・評価

以下の3パターンで確認。3回実行して平均をとる

  • 行が多い
  • 行・列が多い

検証用コード

import time
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType


def setup():
    # SparkSessionの作成
    spark = SparkSession.builder \
        .appName("CompareDataFrameCountMethods") \
        .getOrCreate()

    # ------------------------------------------------
    # 1. スキーマ定義
    # ------------------------------------------------
    row_count = 100
    schema = StructType([
        StructField(f"col{x}", StringType(), True)
        for x in range(1, row_count + 1)
    ])

    # ------------------------------------------------
    # 2-A. 非空のDataFrameを作成
    # ------------------------------------------------
    # num_rows = 10000
    # data = []
    # for i in range(num_rows):
    #     # col1(ID), col2~col10(文字列)
    #     row = (f"name_{i}", ) * row_count
    #     data.append(row)
    # df = spark.createDataFrame(data, schema)

    # ------------------------------------------------
    # 2-B. 空のDataFrameを作成(同じスキーマ)
    # ------------------------------------------------
    df = spark.createDataFrame([], schema)

    return df, spark


def tearDown(spark):
    spark.stop()

# 同じ状態で検証したいので、都度セッションを作り直す
# ------------------------------------------------
# 1. DataFrame.count() の実行時間計測
# ------------------------------------------------
df, spark = setup()
start_time = time.time()
count_result = df.count()
count_time = time.time() - start_time
tearDown(spark)

# ------------------------------------------------
# 2. DataFrame.limit(1).count() の実行時間計測
# ------------------------------------------------
df, spark = setup()
start_time = time.time()
limit_count_result = df.limit(1).count()
limit_count_time = time.time() - start_time
tearDown(spark)

# ------------------------------------------------
# 3. len(DataFrame.take(1)) の実行時間計測
# ------------------------------------------------
df, spark = setup()
start_time = time.time()
take_result = len(df.take(1))
take_time = time.time() - start_time
tearDown(spark)

# ------------------------------------------------
# 4. 結果を表示
# ------------------------------------------------
print("============================================")
print(f"DataFrame.count() -> 実行結果: {count_result}, 実行時間: {count_time:.6f} 秒")
print(
    f"DataFrame.limit(1).count() -> 実行結果: {limit_count_result}, 実行時間: {limit_count_time:.6f} 秒"
)
print(
    f"len(DataFrame.take(1)) -> 取得件数: {take_result}, 実行時間: {take_time:.6f} 秒")
print("============================================")

検証結果

行が多い

カラム数10, 行数10,000

Method Iteration 1 (s) Iteration 2 (s) Iteration 3 (s) Average (s)
DataFrame.count() 1.751275 1.745480 1.704119 1.733625
DataFrame.limit(1).count() 0.793862 0.796269 0.804845 0.798325
len(DataFrame.take(1)) 0.540404 0.549774 0.521649 0.537276

行・列が多い

カラム数1,000, 行数1,000

Method Iteration 1 (s) Iteration 2 (s) Iteration 3 (s) Average (s)
DataFrame.count() 2.687237 2.584143 2.711557 2.660979
DataFrame.limit(1).count() 1.249905 1.138883 1.290905 1.226564
len(DataFrame.take(1)) 0.626042 0.616781 0.663114 0.635312

カラム数1,000, 行数0

Method Iteration 1 (s) Iteration 2 (s) Iteration 3 (s) Average (s)
DataFrame.count() 2.567037 2.572245 2.524549 2.554610
DataFrame.limit(1).count() 1.224313 1.151787 1.210922 1.195674
len(DataFrame.take(1)) 1.061445 1.035899 1.020977 1.039440

まとめ

  • len(DataFrame.take(1)) が早い
  • カラム数が多いほど効果があるみたい

Discussion