Parquetファイルのダミーデータを生成する

Pythonでparquetファイルのダミーデータを生成する方法の備忘録です。

はじめに
Parquetファイルのダミーデータを生成する
補足：S3に配置する
おわりに
参考

はじめに

ちょっとした検証のために、ダミーデータのparquetファイルを用意する機会がありました。

Pythonスクリプトを書いたので、備忘録にメモします。

※ソースコードはこちらに置いています： github.com

Parquetファイルのダミーデータを生成する

Parquetファイルのダミーデータを生成するPythonスクリプト例です。圧縮形式はここではgzipとしています。

from typing import Dict, List
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import boto3


def create_dummy_data(rows: int = 1000) -> pd.DataFrame:
    data: Dict[str, List] = {
        'id': range(1, rows + 1),
        'name': [f'name_{i}' for i in range(1, rows + 1)],
        'value': [i * 100 for i in range(1, rows + 1)]
    }
    return pd.DataFrame(data)


def save_gz_parquet(df: pd.DataFrame, file_path: str) -> None:
    table: pa.Table = pa.Table.from_pandas(df)
    pq.write_table(table, file_path, compression='gzip')


def main() -> None:
    dummy_data: pd.DataFrame = create_dummy_data()

    file_path: str = 'dummy_data.parquet.gz'
    save_gz_parquet(dummy_data, file_path)


if __name__ == '__main__':
    main()

pyarrow.parquetを使うことで、pandas DataFrameから簡単にparquetファイルを書き出すことができます。

上記のコード例では、create_dummy_data関数でダミーのpandas DataFrameを作成し、save_gz_parquet関数でそのデータをparquetファイルで保存しています。

補足：S3に配置する

作成したparquetファイルをそのままS3に配置するには、下記のような形になります。

from typing import Dict, List
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import boto3


def create_dummy_data(rows: int = 1000) -> pd.DataFrame:
    data: Dict[str, List] = {
        'id': range(1, rows + 1),
        'name': [f'name_{i}' for i in range(1, rows + 1)],
        'value': [i * 100 for i in range(1, rows + 1)]
    }
    return pd.DataFrame(data)


def save_gz_parquet(df: pd.DataFrame, file_path: str) -> None:
    table: pa.Table = pa.Table.from_pandas(df)
    pq.write_table(table, file_path, compression='gzip')


def upload_to_s3(bucket: str, s3_key: str, file_path: str) -> None:
    s3 = boto3.client('s3')
    s3.upload_file(file_path, bucket, s3_key)


def main() -> None:
    dummy_data: pd.DataFrame = create_dummy_data()

    file_path: str = 'dummy_data.parquet.gz'
    save_gz_parquet(dummy_data, file_path)

    bucket_name: str = 'your-bucket-name'
    # ".parquet.gz" doesn't work in S3 select.
    s3_key: str = 'path/to/dummy_data.gz.parquet'
    upload_to_s3(bucket_name, s3_key, file_path)


if __name__ == '__main__':
    main()