dbldatagen

The Databricks Labs Data Generator (dbldatagen) is a Spark-based Python library designed to generate realistic synthetic data at scale. It leverages Spark DataFrames and Spark SQL to produce data that can be written to various storage formats, saved as Delta tables, or manipulated using standard Spark APIs.

Key capabilities include:

• Generating billions of rows of data efficiently using Spark clusters.
• Specifying row counts and Spark partition counts for distribution.
• Controlling data via numeric, time, and date ranges.
• Using random or repeatable seed values (including weighting for discrete values).
• Template-based text generation and string formatting.
• Using SQL-based expressions to control or augment column generation.
• Supporting Delta Live Tables (DLT) for both streaming and batch operations.
• Supporting Unity Catalog enabled clusters.
• Providing pluggable standard datasets for quick prototyping.

Use ArrayColumn to produce variable-length arrays of elements. This is ideal for repeated values like tags or scores. Each row's array length is randomly determined within the range [min_length, max_length].

Configuration Fields

• element (required): The generation strategy for each element in the array. This must be a strategy model (not a ColumnSpec).
• min_length: Inclusive minimum length per row (default 1). Set to 0 to allow sometimes-empty arrays.
• max_length: Inclusive maximum length per row (default 5).

Key Constraints

• Maximum length: max_length is capped at 1000. For larger fan-outs, restructure your data as a long/narrow table instead of a wide array.
• Unsupported strategies: FakerColumn and ForeignKeyColumn cannot be used as the element strategy. To use them, wrap them in a StructColumn.
• Null behavior: An array column with null_fraction=1.0 produces Spark NULL values, which are distinct from empty arrays []. There is no per-element null_fraction.
• Fixed length: To produce an array with a fixed length, set min_length == max_length.

from dbldatagen.core import ColumnSpec
from dbldatagen.core.spec.schema import ArrayColumn, ValuesColumn, RangeColumn

# Array of categorical values
ColumnSpec(
    name="tags",
    gen=ArrayColumn(
        element=ValuesColumn(values=["sale", "new", "popular", "clearance"]),
        min_length=1,
        max_length=4,
    ),
)

# Array of numbers from a range
ColumnSpec(
    name="scores",
    gen=ArrayColumn(element=RangeColumn(min=0, max=100), min_length=2, max_length=5),
)

Neither SequenceColumn nor UUIDColumn automatically marks a column as a primary key. You must explicitly declare the key using the PrimaryKey object within the TableSpec.

PrimaryKey Configuration:

• columns: A list[str] containing the names of the columns forming the key, in declaration order.
• Composite Keys: You can provide multiple column names for a composite primary key.
• Constraints:
  • Each name must match a ColumnSpec.name defined in the table's columns.
  • Duplicate column names in the PrimaryKey list are rejected at plan time.
  • Important: While composite PKs are supported, ForeignKeyRef cannot point to a table with a composite PK; foreign keys must target single-column primary keys.
• Requirement: A PrimaryKey is only strictly required if another table in your generation plan will use a ForeignKeyRef to point to this table.

# Example of a single-column PK
TableSpec(
    name="customers",
    primary_key=PrimaryKey(columns=["customer_id"]),
    columns=[
        ColumnSpec(name="customer_id", gen=SequenceColumn()),
    ]
)

# Example of a composite PK
TableSpec(
    name="order_items",
    primary_key=PrimaryKey(columns=["order_id", "item_id"]),
    columns=[
        ColumnSpec(name="order_id", gen=SequenceColumn()),
        ColumnSpec(name="item_id", gen=SequenceColumn()),
    ]
)

To generate complex nested types like struct, array, or map that will be converted to JSON when saved, you can use the expr attribute within withColumn.

Important: When using the expr attribute, it will override other column data generation rules for that specific column.

To simplify the process, you can use dg.INFER_DATATYPE as the type argument, which allows dbldatagen to infer the schema directly from your SQL expression.

import dbldatagen as dg

# Using expr with explicit type
testDataSpec = (
    dg.DataGenerator(spark, name="device_data_set")
    .withColumn("event_info", 
                 dg.StructType([dg.StructField('event_type', dg.StringType()), 
                                dg.StructField('event_ts', dg.TimestampType())]),
                 expr="named_struct('event_type', event_type, 'event_ts', event_ts)",
                 baseColumn=['event_type', 'event_ts'])
)

# Using dg.INFER_DATATYPE to simplify
testDataSpec = (
    dg.DataGenerator(spark, name="device_data_set")
    .withColumn("event_info", 
                 dg.INFER_DATATYPE,
                 expr="named_struct('event_type', event_type, 'event_ts', event_ts)")
)

You can define a DataGenPlan by explicitly constructing ColumnSpec objects and passing specific strategy models (like SequenceColumn or RangeColumn) to each column. This approach is equivalent to the structure used when deserializing a plan from a JSON or YAML specification file. This style provides high granularity and is useful when you want to mirror a static configuration file in your code.

from dbldatagen.core import DataGenPlan, TableSpec, ColumnSpec, PrimaryKey
from dbldatagen.core.spec.schema import SequenceColumn, RangeColumn, DataType

plan = DataGenPlan(
    seed=42,
    tables=[
        TableSpec(
            name="orders",
            rows=1000,
            primary_key=PrimaryKey(columns=["order_id"]),
            columns=[
                ColumnSpec(name="order_id", gen=SequenceColumn(start=1, step=1)),
                ColumnSpec(name="amount", dtype=DataType.DOUBLE,
                           gen=RangeColumn(min=5.0, max=500.0)),
            ],
        ),
    ],
)

To create join-ready data across multiple tables, you can generate predictable foreign keys using the baseColumnType="hash" option. This allows you to derive a unique, repeatable key from an existing column (like a customer_id) without needing to manage complex sequences.

When using hashing to generate keys:

• Use baseColumn to specify the source column.
• Use baseColumnType="hash" to trigger the hashing logic.
• For large datasets, consider using decimal types for IDs to prevent overflow when working with large hashed values.
• Note that while hashing is highly efficient for generating predictable keys, you should still call .dropDuplicates() on the resulting DataFrame to handle potential hash collisions in smaller datasets.

# Generating a device_id that is a hash of the customer_id
customer_dataspec = (
    dg.DataGenerator(spark, rows=50000, partitions=8)
    .withColumn("customer_id", "decimal(10)", minValue=1000, uniqueValues=50000)
    .withColumn("device_id", "decimal(10)", minValue=1000000000, 
                baseColumn="customer_id", baseColumnType="hash")
)

# Ensure uniqueness after generation
df_customers = customer_dataspec.build().dropDuplicates(["device_id"])