Databricks Labs による Data Generator Tool (dbldatagen)の紹介

はじめに

コードの検証をしたいけど、データが手に入らない なんて事はないでしょうか？　

Databricks Labsでは様々な便利ツールを開発しておりますが、今回は dbldatagenというデータ生成ツールを見つけて使ってみましたので、メモしておきたいと思います。

Databricks Labs Data Generator

詳細はこちらをご覧ください。
Getting started with the Databricks Labs Data Generator

こちらの特徴は、Spark APIを使ってデータ生成するため、複数サーバーのリソースを使って分散処理させることができます。上記ドキュメントによると１０億行のデータ生成に、8台のクラスターを使って２分程度で作成できるようです。

また数値データや文字データ、Datetimeなど様々なデータ生成に対応しております。

ライブラリーインストール

Wheel化されているため、ノートブック上から以下のコード実行でインストールできます。

%pip install git+https://github.com/databrickslabs/dbldatagen.git

サンプルコード

あとはコードを実行するだけです。シンプルなので、以下のサンプルコードを見ればなんとなくイメージ湧くのではないでしょうか？

from pyspark.sql.types import *
import dbldatagen as dg

row_count = 1000 * 100
column_count = 3
testDataSpec = (dg.DataGenerator(spark, name="test_data_set1", rows=row_count,
                                  partitions=8, randomSeedMethod='hash_fieldname', 
                                  verbose=True)
                .withIdOutput()
                .withColumn("r", FloatType(), expr="floor(rand() * 350) * (86400 + 3600)",
                                    numColumns=column_count)
                .withColumn("code1", IntegerType(), minValue=100, maxValue=200)
                .withColumn("code2", IntegerType(), minValue=0, maxValue=10, random=True)
                .withColumn("code3", StringType(), values=['online', 'offline', 'unknown'])
                .withColumn("code4", StringType(), values=['a', 'b', 'c'], random=True,percentNulls=0.05)
                .withColumn("code5", StringType(), values=['a', 'b', 'c'], random=True,weights=[9, 1, 1])
                .withColumn("purchase_date1", "date", uniqueValues=300,random=True)
                .withColumn("purchase_date2", "date", data_range=dg.DateRange("2021-08-01 00:00:00","2021-10-06 11:55:00","hours=1"), 
                               random=True)
                .withColumn("email", template=r'\\w.\\w@\\w.com')
               )
dfTestData = testDataSpec.build()
display(dfTestData)

コード解説

それでは簡単にコード解説します。

データ数

row_count = 1000 * 100
testDataSpec = (dg.DataGenerator(spark, name="test_data_set1", rows=row_count,
                                  partitions=8, randomSeedMethod='hash_fieldname', 
                                  verbose=True)

↑ 最初の引数の中で、データ数を指定できます。今回は１０万行でしたが、たとえ数億行データでも生成可能です。

id

.withIdOutput()

↑ 0から順番にindexを割り振ってくれます。

複数のカラム生成

column_count = 3

.withColumn("r", FloatType(), expr="floor(rand() * 350) * (86400 + 3600)",numColumns=column_count)

↑ r_xx のように、複数のカラムを一気に生成できます。ただしカラム名はルールによって自動生成されます。生成する絡む数を column_countで指定します。

数値カラム生成

.withColumn("code1", IntegerType(), minValue=100, maxValue=200)
.withColumn("code2", IntegerType(), minValue=0, maxValue=10, random=True)

↑ 数値の最小値と最大値を指定してデータ生成します。random=True をつけないと１づつ増加し、maxValueに到達すると、またMinValueから増えていきます。

カテゴリデータ生成

.withColumn("code3", StringType(), values=['online', 'offline', 'unknown'])
.withColumn("code4", StringType(), values=['a', 'b', 'c'], random=True,percentNulls=0.05)
.withColumn("code5", StringType(), values=['a', 'b', 'c'], random=True,weights=[9, 1, 1])

↑ カテゴリーデータも生成可能です。percentNullsを指定するとnullの割合も指定できます。またそれぞれの重みづけもweightsにより可能です。

日付・時間

.withColumn("purchase_date1", "date", uniqueValues=300,random=True)
.withColumn("purchase_date2", "datetime", data_range=dg.DateRange("2021-08-01 00:00:00","2021-10-06 11:55:00","hours=1"),random=True)

↑ 日付や時間なども生成できます。data_rangeを指定すると、範囲指定ができるのと、hours=1というのは１時間単位でランダムに生成するという意味です。他にdays,や minsなども可能です。

テキストデータ

.withColumn("email", template=r'\\w.\\w@\\w.com')

↑ 名前やメールアドレスなどテキストデータも生成できます。正規表現の \wオプションを使う事で文字列を生成できます。

他にもデータ生成オプションはあります。詳細はドキュメントをご覧ください。

今回のサンプルノートブックは[こちら](https://sajpstorage.blob.core.windows.net/maruyama/public_share/data generator (aka. dbldatagen).html)です。