MockTableGenerators.jl

The MockTableGenerators.jl package provides an interface for composing the the generation of multiple dependent Tables.jl to produce realistic mock datasets.

Users should define subtypes of TableGenerator and extend the table_key, num_rows, and emit! functions. Special row generators may also need to make use of visit! for introducing state or dependency_key for multiple TableGenerator types which creating rows for the same table. Instances of TableGenerators can be constructed into a DAG which defines dependences between generators.

Methods for functions that may introduce randomness (i.e., num_rows, emit!, and visit!) must accept a random number generator as the first argument in order to support reproducible generation. In cases where visit! introduces randomness in the generated state and emit! and num_rows only consume this state, they still have to accept it but may ignore it.

An example showing row generation including the use of variable number of rows, state, and conditional dependencies:

using MockTableGenerators, Dates, StableRNGs, UUIDs

const FIRST_NAMES = ["Alice", "Bob", "Carol", "David"]
const LAST_NAMES = ["Smith", "Johnson", "Williams", "Brown"]

struct PersonGenerator <: TableGenerator
    num::AbstractRange{Int}
end

PersonGenerator(num::Integer) = PersonGenerator(range(num))

MockTableGenerators.table_key(g::PersonGenerator) = :person
MockTableGenerators.num_rows(rng, g::PersonGenerator) = rand(rng, g.num)

function MockTableGenerators.emit!(rng, g::PersonGenerator, deps)
    return (; id=uuid4(rng), 
            first_name=rand(rng, FIRST_NAMES), 
            last_name=rand(rng, LAST_NAMES))
end

struct VisitGenerator <: TableGenerator
    num::AbstractRange{Int}
end

function MockTableGenerators.visit!(rng, g::VisitGenerator, deps)
    n = rand(rng, g.num)
    visits = sort!(rand(rng, Date(1970):Day(1):Date(2000), n))
    return Dict(:i => 1, :visits => visits, :n => n)
end

MockTableGenerators.table_key(g::VisitGenerator) = :visit
MockTableGenerators.num_rows(rng, g::VisitGenerator, state) = state[:n]

function MockTableGenerators.emit!(rng, g::VisitGenerator, deps, state)
    visit = popfirst!(state[:visits])

    row = (; id=uuid4(rng), person_id=deps[:person].id, index=state[:i], date=visit)

    state[:i] += 1
    return row
end

const LIGHT_SYMPTOMS = ["Fever", "Chills", "Fatigue", "Runny nose", "Cough"]
const SEVERE_SYMPTOMS = ["Weakness", "Muscle Loss", "Fainting"]

struct SymptomGenerator <: TableGenerator
    num::AbstractRange{Int}
end

function MockTableGenerators.visit!(rng, g::SymptomGenerator, deps)
    # Number of symptoms increase, on average, with number of visits
    n = rand(rng, min(deps[:visit].index, last(g.num)):last(g.num))
    return (; n)
end

MockTableGenerators.table_key(g::SymptomGenerator) = :symptom
MockTableGenerators.num_rows(rng, g::SymptomGenerator, state) = state.n

function MockTableGenerators.emit!(rng, g::SymptomGenerator, deps, state)
    # Conditional generation based upon number of visits
    symptoms = deps[:visit].index > 2 ? SEVERE_SYMPTOMS : LIGHT_SYMPTOMS
    return (; visit_id=deps[:visit].id, symptom=rand(rng, symptoms))
end

const DAG = [PersonGenerator(3:5) => [VisitGenerator(1:4) => [SymptomGenerator(1:2)]]]
# pass RNG for reproducible generation:
results = collect(MockTableGenerators.generate(StableRNG(11), DAG))

# Alternatively, since v0.2.1, linear DAGs can be also constructed in a flat representation:
const FLAT_DAG = PersonGenerator(3:5) => VisitGenerator(1:4) => SymptomGenerator(1:2)
flat_results = collect(MockTableGenerators.generate(StableRNG(11), FLAT_DAG))

@assert results == flat_results