diff --git a/manifest.json b/manifest.json index 07ee4daa..bdcfe63f 100644 --- a/manifest.json +++ b/manifest.json @@ -435,7 +435,7 @@ "version": "0.1.0" }, "databricks-synthetic-data-gen": { - "description": "Generate realistic synthetic data using Spark + Faker (strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions...", + "description": "Generate realistic synthetic data using Spark + dbldatagen (Databricks Labs Data Generator, strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), a...", "files": [ "SKILL.md", "agents/openai.yaml", @@ -446,7 +446,7 @@ "scripts/generate_synthetic_data.py" ], "repo_dir": "skills", - "version": "0.1.0" + "version": "0.2.0" }, "databricks-unity-catalog": { "description": "Unity Catalog system tables and volumes.", diff --git a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/SKILL.md b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/SKILL.md index f534f60f..fe2b70c6 100644 --- a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/SKILL.md +++ b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/SKILL.md @@ -1,6 +1,6 @@ --- name: databricks-synthetic-data-gen -description: "Generate realistic synthetic data using Spark + Faker (strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'Faker', or 'sample data'." +description: "Generate realistic synthetic data using Spark + dbldatagen (Databricks Labs Data Generator, strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'dbldatagen', or 'sample data'." compatibility: Requires databricks CLI (>= v1.0.0) metadata: version: "0.1.0" @@ -11,7 +11,9 @@ parent: databricks-core # Databricks Synthetic Data Generation -Generate realistic, story-driven synthetic data for Databricks using **Spark + Faker + Pandas UDFs** (strongly recommended). +Generate realistic, story-driven synthetic data for Databricks using **Spark + [dbldatagen](https://databrickslabs.github.io/dbldatagen/public_docs/index.html)** (the Databricks Labs Data Generator — strongly recommended). + +> **Use the public `dbldatagen` API.** Import the top-level package (`import dbldatagen as dg`) and its public submodules (`dbldatagen.distributions`, `dbldatagen.constraints`). Everything below is built from the public API. ## Data Must Tell a Business Story @@ -44,11 +46,11 @@ Synthetic data should demonstrate how Databricks helps solve real business probl 5. **Enough data for trends** — ~100K+ rows for main tables so patterns survive aggregation 6. **Ask for catalog/schema** — Never default, always confirm before generating 7. **Present plan for approval** — Show tables, distributions, assumptions before writing code -8. **Master tables first** — Generate parent tables, write to Delta, then create children with valid FKs -9. **Use Spark + Faker + Pandas UDFs** — Scalable, parallel. Polars only if user explicitly wants local + <30K rows +8. **Parent tables first** — Generate parent tables, write to Delta, then create children with valid FKs +9. **Use Spark + dbldatagen** — Scalable, parallel, declarative. Build a `dg.DataGenerator` spec and `.build()` it into a Spark DataFrame. Use a `pandas_udf` only for logic dbldatagen can't express 10. **Use Databricks Connect Serverless by default to generate data** — Update databricks-connect on python 3.12 if required (avoid using execute_code unless instructed to not use Databricks Connect) 11. **No `.cache()` or `.persist()`** — Not supported on serverless. Write to Delta, read back for joins -12. **No Python loops or `.collect()`** — Use Spark parallelism. No driver-side iteration, avoid Pandas↔Spark conversions +12. **No Python loops or `.collect()`** — Use Spark parallelism and dbldatagen specs. No driver-side iteration, avoid Pandas↔Spark conversions ## Generation Planning Workflow @@ -142,29 +144,57 @@ SELECT COUNT(*) FROM parquet.\`/Volumes/CATALOG/SCHEMA/raw_data/customers\` See [references/2-troubleshooting.md](references/2-troubleshooting.md) for full validation examples. -## Use Databricks Connect Spark + Faker Pattern +## Use Databricks Connect + dbldatagen Pattern + +A `DataGenerator` spec declares each column; `.build()` returns a Spark DataFrame. No UDFs, no driver loops — generation is fully distributed across `partitions`. Install `dbldatagen` and `faker` locally first (see Setup). ```python from databricks.connect import DatabricksSession -from pyspark.sql import functions as F -from pyspark.sql.types import StringType -import pandas as pd +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import company, person # providers we draw from -# Setup serverless Spark session +# Setup serverless Spark session (deps installed locally) spark = DatabricksSession.builder.serverless(True).getOrCreate() -# Pandas UDF pattern - import lib INSIDE the function (libs must be installed locally) -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - from faker import Faker # Import inside UDF - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) - -# Generate with spark.range, apply UDFs -customers_df = spark.range(0, 10000, numPartitions=16).select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), +CATALOG, SCHEMA = "", "" # always user-supplied +N_CUSTOMERS = 10_000 + +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company]) + +customers = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, + randomSeed=42, randomSeedMethod="hash_fieldname") + # surrogate key 0..N-1 (the implicit contiguous seed column) — drives FK joins + .withColumn("customer_sk", "long", expr="id") + # business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # realistic text via the shared Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + # email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(replace(name, ' ', '.')), '@example.com')") + # skewed categories — NEVER uniform (weights are relative) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=[40, 25, 20, 15], random=True) + # right-skewed ARR correlated to tier: log-normal = exp() of a standard normal. + # The hidden helper column (_z, omit=True) is computed and reusable as a base column. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") + .withColumn("created_at", "date", + data_range=dg.DateRange("2023-01-01 00:00:00", "2024-12-31 00:00:00", "days=1"), + random=True) ) +customers_df = customers.build() # Write to Volume as Parquet (default for raw data) # Path is a folder with table name: /Volumes/catalog/schema/raw_data/customers/ @@ -173,7 +203,7 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") customers_df.write.mode("overwrite").parquet(f"/Volumes/{CATALOG}/{SCHEMA}/raw_data/customers") ``` -**Partitions by scale:** `spark.range(N, numPartitions=P)` +**Partitions by scale:** `dg.DataGenerator(..., rows=N, partitions=P)` - <100K rows: 8 partitions - 100K-500K: 16 partitions - 500K-1M: 32 partitions @@ -190,32 +220,79 @@ Generated scripts must be highly performant. **Never** do these: | Anti-Pattern | Why It's Slow | Do This Instead | |--------------|---------------|-----------------| -| Python loops on driver | Single-threaded, no parallelism | Use `spark.range()` + Spark operations | +| Python loops on driver | Single-threaded, no parallelism | Declare columns in a `dg.DataGenerator` spec and `.build()` | | `.collect()` then iterate | Brings all data to driver memory | Keep data in Spark, use DataFrame ops | -| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark, use `pandas_udf` only for UDFs | +| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark; let dbldatagen generate columns | | Read/write temp files | Unnecessary I/O | Chain DataFrame transformations | -| Scalar UDFs | Row-by-row processing | Use `pandas_udf` for batch processing | +| Scalar UDFs | Row-by-row processing | Use dbldatagen `expr`/`template`/`text`; `pandas_udf` only when unavoidable | -**Good pattern:** `spark.range()` → Spark transforms → `pandas_udf` for Faker → write directly +**Good pattern:** `dg.DataGenerator(rows, partitions)` → declarative `.withColumn(...)` specs → `.build()` → write directly ## Common Patterns +All snippets use the public `dbldatagen` API (`import dbldatagen as dg`, `import dbldatagen.distributions as dist`). + ### Weighted Categories (never uniform) +`weights` are relative frequencies — they don't need to sum to 100: ```python -F.when(F.rand() < 0.6, "Free").when(F.rand() < 0.9, "Pro").otherwise("Enterprise") +.withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) ``` -### Log-Normal Amounts (in a pandas UDF) -Use `np.random.lognormal(mean, sigma)` — always positive, long tail: -- Enterprise: `lognormal(7.5, 0.8)` → ~$1800 median -- Pro: `lognormal(5.5, 0.7)` → ~$245 median -- Free: `lognormal(4.0, 0.6)` → ~$55 median +### Skewed / Long-Tailed Amounts +Apply a continuous distribution to a numeric range — `Gamma`/`Exponential` give the always-positive long tail you'd want from log-normal: +```python +.withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) +``` +For a true log-normal, generate over a normal and exponentiate via `expr`: +```python +.withColumn("amount", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # standard normal, hidden +.withColumn("arr", "double", baseColumn="amount", + expr="round(exp(5.5 + 0.8 * amount), 2)") # median ~$245 +``` -### Date Range (Last 6 Months) +### Coherent Rows (correlated attributes via `expr` + `baseColumn`) +Derive dependent columns from earlier ones so each row makes business sense — no UDF needed: ```python -END_DATE = datetime.now() +.withColumn("priority", "string", values=["Critical", "High", "Medium", "Low"], + weights=[5, 15, 50, 30], random=True) +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*8 " + "WHEN 'High' THEN rand()*24 WHEN 'Medium' THEN rand()*72 " + "ELSE rand()*120 END, 1)") +.withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 4 " + "WHEN resolution_hours<72 THEN 3 ELSE 2 END") +``` + +### Date / Timestamp Range (Last 6 Months) +Use `dg.DateRange(begin, end, interval)` with the `data_range` option: +```python +from datetime import datetime, timedelta +END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +FMT = "%Y-%m-%d %H:%M:%S" + +.withColumn("order_ts", "timestamp", + data_range=dg.DateRange(START_DATE.strftime(FMT), END_DATE.strftime(FMT), "minutes=30"), + random=True) +``` + +### Realistic Text (Faker provider plugin) +Build one `FakerTextFactory` with a default locale and the faker providers you need, then reuse it: +```python +from faker.providers import person, company, internet +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, internet]) + +.withColumn("name", "string", text=FakerText("name")) # uses the 'person' provider +.withColumn("company", "string", text=FakerText("company")) # uses the 'company' provider +.withColumn("ip", "string", text=FakerText("ipv4_private")) # uses the 'internet' provider ``` +dbldatagen alternatives that need no library: +* `template=r"\w.\w@\w.com"` (templated text) +* `text=dg.ILText(paragraphs=(1, 3), sentences=(2, 5))` (lorem-ipsum free text). ### Infrastructure (always create in script) ```python @@ -224,27 +301,33 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") ``` ### Referential Integrity (FK pattern) -Write master table to Delta first, then read back for FK joins (no `.cache()` on serverless): +Generate the FK as an integer over the parent's surrogate-key range so **every value is valid by construction**, then write the parent to Delta and read it back to attach coherent parent attributes (no `.cache()` on serverless): ```python -# 1. Write master table +# 1. Parent table: surrogate key 0..N-1, then write to Delta customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") -# 2. Read back for FK lookup -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_idx", "customer_id") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +# 2. Child table: FK drawn from the parent key range (valid by construction). +# A distribution skews it 80/20 (a few customers place most orders). +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("order_id", "string", + expr="concat('ORD-', lpad(cast(id as string), 8, '0'))", baseColumn="id") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") + +# 3. Join to the parent (read back from Delta) to pull coherent parent attributes +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "customer_id", "tier") +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") ``` ## Setup -Requires Python 3.12 and databricks-connect>=16.4. Use `uv`: +Requires Python 3.12 and databricks-connect>=16.4. Install dependencies locally with `uv`: ```bash -uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays +uv pip install "databricks-connect>=16.4,<17.4" dbldatagen faker ``` ## Related Skills @@ -256,10 +339,11 @@ uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays | Issue | Solution | |-------|----------| -| `ModuleNotFoundError: faker` | Install locally: `uv pip install faker`, import inside UDF | -| Faker UDF is slow | Use `pandas_udf` for batch processing | -| Out of memory | Increase `numPartitions` in `spark.range()` | -| Referential integrity errors | Write master table to Delta first, read back for FK joins | +| `ModuleNotFoundError: dbldatagen` (or `faker`) | Install locally: `uv pip install dbldatagen faker` | +| `FakerText`/provider not found | Pass the provider to `dg.FakerTextFactory(providers=[...])` and import it from `faker.providers` | +| All rows identical / not random | Set `random=True` on the column (default is deterministic), or set `randomSeed`/`randomSeedMethod` on the generator | +| Out of memory | Increase `partitions` in `dg.DataGenerator(..., partitions=P)` | +| Referential integrity errors | Draw the FK from the parent key range (`minValue=0, maxValue=N-1`); write parent to Delta first, read back to join attributes | | `PERSIST TABLE is not supported on serverless` | **NEVER use `.cache()` or `.persist()` with serverless** - write to Delta table first, then read back | | `F.window` vs `Window` confusion | Use `from pyspark.sql.window import Window` for `row_number()`, `rank()`, etc. `F.window` is for streaming only. | | Broadcast variables not supported | **NEVER use `spark.sparkContext.broadcast()` with serverless** | diff --git a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/1-data-patterns.md b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/1-data-patterns.md index eba64916..87774896 100644 --- a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/1-data-patterns.md +++ b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/1-data-patterns.md @@ -28,15 +28,31 @@ Real data is never uniformly distributed. Use appropriate distributions: | **Exponential** | Time between events | Support resolution time, session duration | | **Weighted categorical** | Skewed categories | Status (70% complete, 5% failed), tiers | -```python -# Log-normal for amounts (long tail, always positive) -amount = np.random.lognormal(mean=5.5, sigma=0.8) # ~$245 median +dbldatagen applies a distribution to a column's random draw via the `distribution=` option +(`import dbldatagen.distributions as dist`). Supported: `Normal`, `Gamma`, `Beta`, `Exponential`. -# Pareto for power-law (few large, many small) -value = (np.random.pareto(a=1.5) + 1) * base_value +```python +import dbldatagen as dg +import dbldatagen.distributions as dist + +spec = ( + dg.DataGenerator(spark, name="amounts", rows=100_000, partitions=8, randomSeed=42) + # Use Gamma for right-tailed distributions like amounts + .withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) + # Use Exponential for waiting times + .withColumn("resolution_hours", "double", minValue=0, maxValue=240, + random=True, distribution=dist.Exponential(rate=1.0 / 24)) +) -# Exponential for time-to-event -hours = np.random.exponential(scale=24) # avg 24h, skewed right +# Exponentiate a standard normal via expr for true log-normal distributions +spec = ( + dg.DataGenerator(spark, name="ln", rows=100_000, partitions=8, randomSeed=42) + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # hidden helper column + .withColumn("amount", "double", baseColumn="_z", + expr="round(exp(5.5 + 0.8 * _z), 2)") # ~$245 median +) ``` ### 3. Row Coherence @@ -51,20 +67,24 @@ Attributes within a row must make business sense together. Generate correlated a | Large transaction + unusual hour | Higher fraud probability | | Fast resolution | Higher CSAT score | +In dbldatagen, chain `baseColumn` + `expr` when columns derive from other columns +in the same dataset. The generator resolves the dependency order automatically: + ```python -@F.pandas_udf("struct") -def generate_coherent_ticket(tiers: pd.Series) -> pd.DataFrame: - """All attributes correlate logically within each row.""" - results = [] - for tier in tiers: - # Priority depends on tier - priority = "Critical" if tier == "Enterprise" and random() < 0.3 else "Medium" - # Resolution depends on priority - resolution = np.random.exponential(4 if priority == "Critical" else 36) - # CSAT depends on resolution - csat = 5 if resolution < 4 else (3 if resolution < 24 else 2) - results.append({"priority": priority, "resolution_hours": resolution, "csat": csat}) - return pd.DataFrame(results) +spec = ( + dg.DataGenerator(spark, name="tickets", rows=80_000, partitions=8, randomSeed=42) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + # Priority depends on tier + .withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' AND rand()<0.3 THEN 'Critical' ELSE 'Medium' END") + # Resolution depends on priority + .withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 ELSE rand()*36 END, 1)") + # CSAT depends on resolution + .withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 3 ELSE 2 END") +) ``` ### 4. The 80/20 Rule @@ -75,7 +95,13 @@ Apply power-law distributions where appropriate: - **20% of products** account for 80% of sales - **20% of support agents** handle 80% of tickets -Implementation: Use weighted sampling when assigning FKs, not uniform random. +Implementation: Skew FKs by drawing from a non-uniform distribution. + +```python +# Use a Gamma distribution to skew FK values +.withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) +``` ### 5. Time-Based Patterns @@ -86,14 +112,22 @@ Most data has temporal patterns: - **Seasonality** — Q4 retail spike, summer travel peak - **Trends** — Growth over time, degradation curves +Bias *when* events happen by weighting time buckets, then derive the timestamp from the bucket. +The example below clusters tickets into business hours and skews volume toward weekdays: + ```python -def get_volume_multiplier(date): - multiplier = 1.0 - if date.weekday() >= 5: multiplier *= 0.6 # Weekend drop - if date.month in [11, 12]: multiplier *= 1.5 # Holiday spike - return multiplier +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], # 9am–5pm peak + random=True, omit=True) +.withColumn("day_offset", "int", minValue=0, maxValue=179, random=True, + distribution=dist.Normal(60, 30), omit=True) # volume ramps over the window +.withColumn("event_ts", "timestamp", baseColumn=["day_offset", "hour"], + expr="date_add(timestamp'2025-01-01 00:00:00', day_offset) + make_interval(0,0,0,0,hour,0,0)") ``` +For holiday/seasonal spikes, weight a month or week-of-year bucket the same way, or post-filter +the built DataFrame with a Spark `expr` multiplier. + ### 6. ML-Ready Data If data will train ML models, ensure: @@ -105,20 +139,32 @@ If data will train ML models, ensure: ## Referential Integrity -Generate master tables first, write to Delta, then join for FKs: +Give the parent a surrogate key over `0..N-1`, draw the child's FK from that same range (valid by +construction), then write the parent to Delta and read it back to attach parent attributes: ```python -# 1. Generate and write master table +# 1. # 1. Generate and write parent table with surrogate keys from 0..N-1 +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_id", "string", baseColumn="id", # References the built-in sequential 'id' column + expr="concat('CUST-', lpad(cast(id as string), 5, '0'))") + # ... other attributes ... + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK joins (NOT cache - unsupported on serverless) -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") + +# 3. Generate child table with FKs drawn from the parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True, omit=True) + .withColumn("customer_id", "string", baseColumn="customer_sk", # References the 'customer_sk' column + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") +orders_with_fk = orders_df.join(customer_lookup, on="customer_id", how="inner") ``` ## Data Volume diff --git a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md index 793b64f7..8a19ebc8 100644 --- a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md +++ b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md @@ -4,23 +4,24 @@ Common issues and solutions for synthetic data generation. ## Environment Issues -### ModuleNotFoundError: faker (or other library) +### ModuleNotFoundError: dbldatagen (or faker) -**Problem:** Dependencies not available in execution environment. +**Problem:** Dependencies not available in execution environment. `dbldatagen` is required, and +`faker` too if you use `FakerTextFactory`/`fakerText`. **Solutions by execution mode:** | Mode | Solution | |------|----------| -| **DB Connect with Serverless** | Install libs locally (`uv pip install faker`), use `DatabricksSession.builder.serverless(True)` | -| **Databricks Runtime** | Use Databricks CLI to install `faker holidays` | -| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "faker"}}, {"pypi": {"package": "holidays"}}]}'` | +| **DB Connect with Serverless** | Install libs locally (`uv pip install dbldatagen faker`), use `DatabricksSession.builder.serverless(True)` | +| **Databricks Runtime** | `%pip install dbldatagen faker` at the top of the notebook | +| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "dbldatagen"}}, {"pypi": {"package": "faker"}}]}'` | ```python # For DB Connect with serverless from databricks.connect import DatabricksSession -# Install dependencies locally first: uv pip install faker pandas numpy holidays +# Install dependencies locally first: uv pip install dbldatagen faker spark = DatabricksSession.builder.serverless(True).getOrCreate() ``` @@ -50,21 +51,21 @@ AnalysisException: [NOT_SUPPORTED_WITH_SERVERLESS] PERSIST TABLE is not supporte **Why this happens:** Serverless compute does not support caching DataFrames in memory. This is a fundamental limitation of the serverless architecture. -**Solution:** Write master tables to Delta first, then read them back for FK joins: +**Solution:** Write parent tables to Delta first, then read them back for FK joins: ```python # BAD - will fail on serverless -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.cache() # ❌ FAILS: "PERSIST TABLE is not supported on serverless compute" # GOOD - write to Delta, then read back -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") # ✓ Read from Delta ``` **Best practice for referential integrity:** -1. Generate master table (e.g., customers) +1. Generate parent table (e.g., customers) 2. Write to Delta table 3. Read back for FK lookup joins 4. Generate child tables (e.g., orders, tickets) with valid FKs @@ -114,7 +115,7 @@ spark = DatabricksSession.builder.serverless(True).getOrCreate() "environment_key": "datagen_env", "spec": { "client": "4", # Required! - "dependencies": ["faker", "numpy", "pandas"] + "dependencies": ["dbldatagen", "faker"] } }] } @@ -153,34 +154,37 @@ contacts_df = contacts_df.withColumn( --- -### Faker UDF is slow +### Generation is slow -**Problem:** Single-row UDFs don't parallelize well. +**Problem:** Row-by-row Python UDFs (or driver loops) don't parallelize well. -**Solution:** Use `pandas_udf` for batch processing: +**Solution:** Let dbldatagen generate columns declaratively — it builds the whole DataFrame in +parallel across `partitions`. Prefer `expr`, `template`, `values`/`weights`, and `distribution` +over UDFs. For names/addresses use the Faker plugin (`FakerTextFactory`) instead of hand-rolled UDFs: ```python -# SLOW - scalar UDF -@F.udf(returnType=StringType()) -def slow_fake_name(): - return Faker().name() - -# FAST - pandas UDF (batch processing) -@F.pandas_udf(StringType()) -def fast_fake_name(ids: pd.Series) -> pd.Series: - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) +import dbldatagen as dg +from faker.providers import person + +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person]) + +spec = ( + dg.DataGenerator(spark, name="people", rows=1_000_000, partitions=64, randomSeed=42) + .withColumn("name", "string", text=FakerText("name")) # batched by dbldatagen + .withColumn("status", "string", values=["active", "churned"], weights=[85, 15], random=True) +) +df = spec.build() ``` ### Out of memory with large data **Problem:** Not enough partitions for data size. -**Solution:** Increase partitions: +**Solution:** Increase `partitions` on the generator: ```python # For large datasets (1M+ rows) -customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from default +spec = dg.DataGenerator(spark, name="big", rows=N_CUSTOMERS, partitions=64, randomSeed=42) ``` | Data Size | Recommended Partitions | @@ -205,21 +209,28 @@ customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from de **Problem:** Foreign keys reference non-existent parent records. -**Solution:** Write master table to Delta first, then read back for FK joins: +**Solution:** Sample the FK from the parent's surrogate-key range, write the parent +to Delta, then read it back for joins: ```python -# 1. Generate and WRITE master table (do NOT use cache with serverless!) -customers_df = spark.range(0, N_CUSTOMERS)... +# 1. Generate and WRITE parent table (do NOT use cache with serverless!) +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_sk", "long", expr="id") # References the built-in sequential 'id' column + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], weights=[60, 30, 10], random=True) + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK lookups -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") - -# 3. Generate child table with valid FKs -orders_df = spark.range(0, N_ORDERS).join( - customer_lookup, - on=, - how="left" +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "tier") + +# 3. Generate child table with FK in the valid parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True) + .build() + .join(customer_lookup, on="customer_sk", how="inner") ) ``` @@ -233,52 +244,52 @@ orders_df = spark.range(0, N_ORDERS).join( **Problem:** All customers have similar order counts, amounts are evenly distributed. -**Solution:** Use non-linear distributions: +**Solution:** Apply a `distribution=` to the column instead of leaving it uniform: ```python -# BAD - uniform -amounts = np.random.uniform(10, 1000, N) +import dbldatagen.distributions as dist -# GOOD - log-normal (realistic) -amounts = np.random.lognormal(mean=5, sigma=0.8, N) +# BAD - uniform (no distribution) +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True) + +# GOOD - skewed, realistic +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True, + distribution=dist.Gamma(1.0, 2.0)) ``` ### Missing time-based patterns **Problem:** Data doesn't reflect weekday/weekend or seasonal patterns. -**Solution:** Add multipliers: +**Solution:** Weight time buckets so volume clusters realistically, then derive the timestamp: ```python -import holidays - -US_HOLIDAYS = holidays.US(years=[2024, 2025]) - -def get_multiplier(date): - mult = 1.0 - if date.weekday() >= 5: # Weekend - mult *= 0.6 - if date in US_HOLIDAYS: - mult *= 0.3 - return mult +# Cluster events into business hours +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], + random=True, omit=True) +.withColumn("event_ts", "timestamp", baseColumn="hour", + expr="date'2025-06-01' + make_interval(0,0,0,0,hour,0,0)") ``` +For weekend/holiday dips, weight a day-of-week or day bucket the same way, or post-filter the built +DataFrame with a Spark `expr` (e.g. `dayofweek(event_ts) IN (1,7)`). + ### Incoherent row attributes **Problem:** Enterprise customer has low-value orders, critical ticket has slow resolution. -**Solution:** Correlate attributes: +**Solution:** Correlate attributes with `baseColumn` + `expr` so each derives from the previous: ```python # Priority based on tier -if tier == 'Enterprise': - priority = np.random.choice(['Critical', 'High'], p=[0.4, 0.6]) -else: - priority = np.random.choice(['Medium', 'Low'], p=[0.6, 0.4]) - -# Resolution based on priority -resolution_scale = {'Critical': 4, 'High': 12, 'Medium': 36, 'Low': 72} -resolution_hours = np.random.exponential(scale=resolution_scale[priority]) +.withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' THEN (CASE WHEN rand()<0.4 THEN 'Critical' ELSE 'High' END) " + "ELSE (CASE WHEN rand()<0.6 THEN 'Medium' ELSE 'Low' END) END") +# Resolution scaled by priority +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 WHEN 'High' THEN rand()*12 " + "WHEN 'Medium' THEN rand()*36 ELSE rand()*72 END, 1)") ``` --- diff --git a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py index b36edb8e..185e4a18 100644 --- a/plugins/databricks/claude/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py +++ b/plugins/databricks/claude/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py @@ -1,22 +1,24 @@ -"""Generate synthetic data using Spark + Faker + Pandas UDFs. +"""Generate synthetic data using Spark + dbldatagen (Databricks Labs Data Generator). This is the recommended approach for ALL data generation tasks: - Scales from thousands to millions of rows -- Parallel execution via Spark +- Declarative, parallel generation (no driver loops, no row-by-row UDFs) - Direct write to Unity Catalog - Works with serverless and classic compute +Uses ONLY the public dbldatagen API: +- `dbldatagen.DataGenerator` / `.withColumn(...)` / `.build()` +- `dbldatagen.distributions` for skew +- `dbldatagen.FakerTextFactory` for realistic names/companies/addresses + Prerequisites: -- Install dependencies locally: uv pip install faker pandas numpy holidays databricks-connect +- Install dependencies locally: uv pip install dbldatagen faker databricks-connect - Configure ~/.databrickscfg with serverless_compute_id = auto """ import sys import os from pyspark.sql import functions as F from pyspark.sql.window import Window -from pyspark.sql.types import StringType, DoubleType, StructType, StructField, IntegerType -import numpy as np -import pandas as pd from datetime import datetime, timedelta # ============================================================================= @@ -39,6 +41,7 @@ # Date range - last 6 months from today END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +DATE_FMT = "%Y-%m-%d %H:%M:%S" # Write mode - "overwrite" for one-time, "append" for incremental WRITE_MODE = "overwrite" @@ -51,14 +54,18 @@ "orphan_fk_rate": 0.01, # 1% orphan foreign keys } -# Reproducibility +# Reproducibility - same seed => same data SEED = 42 -# Tier distribution: Free 60%, Pro 30%, Enterprise 10% -TIER_PROBS = [0.6, 0.3, 0.1] +# Tier distribution: Free 60%, Pro 30%, Enterprise 10% (relative weights) +TIER_WEIGHTS = [60, 30, 10] + +# Region distribution (relative weights) +REGION_WEIGHTS = [40, 25, 20, 15] -# Region distribution -REGION_PROBS = [0.4, 0.25, 0.2, 0.15] +# Order status distribution (relative weights) +STATUS_VALUES = ["delivered", "shipped", "processing", "pending", "cancelled"] +STATUS_WEIGHTS = [65, 15, 10, 5, 5] # ============================================================================= # SESSION CREATION @@ -79,59 +86,14 @@ spark = DatabricksSession.builder.clusterId(CLUSTER_ID).getOrCreate() print(f"Connected to cluster {CLUSTER_ID}") -# Import Faker for UDF definitions -from faker import Faker - -# ============================================================================= -# DEFINE PANDAS UDFs FOR FAKER DATA -# ============================================================================= +# Import the public dbldatagen API (installed locally) +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import person, company, address as faker_address -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - """Generate realistic person names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.name() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_company(ids: pd.Series) -> pd.Series: - """Generate realistic company names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.company() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_address(ids: pd.Series) -> pd.Series: - """Generate realistic addresses.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.address().replace('\n', ', ') for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_email(names: pd.Series) -> pd.Series: - """Generate email based on name.""" - emails = [] - for name in names: - if name: - domain = name.lower().replace(" ", ".").replace(",", "")[:20] - emails.append(f"{domain}@example.com") - else: - emails.append("unknown@example.com") - return pd.Series(emails) - -@F.pandas_udf(DoubleType()) -def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: - """Generate amount based on tier using log-normal distribution.""" - np.random.seed(SEED) - amounts = [] - for tier in tiers: - if tier == "Enterprise": - amounts.append(float(np.random.lognormal(mean=7.5, sigma=0.8))) # ~$1800 avg - elif tier == "Pro": - amounts.append(float(np.random.lognormal(mean=5.5, sigma=0.7))) # ~$245 avg - else: - amounts.append(float(np.random.lognormal(mean=4.0, sigma=0.6))) # ~$55 avg - return pd.Series(amounts) +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, faker_address]) # ============================================================================= # CREATE INFRASTRUCTURE @@ -142,39 +104,50 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f"Infrastructure ready: {VOLUME_PATH}") # ============================================================================= -# GENERATE CUSTOMERS (Master Table) +# GENERATE CUSTOMERS (Parent Table) # ============================================================================= print(f"\nGenerating {N_CUSTOMERS:,} customers...") -customers_df = ( - spark.range(0, N_CUSTOMERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), - fake_company(F.col("id")).alias("company"), - fake_address(F.col("id")).alias("address"), - # Tier distribution: Free 60%, Pro 30%, Enterprise 10% - F.when(F.rand(SEED) < TIER_PROBS[0], "Free") - .when(F.rand(SEED) < TIER_PROBS[0] + TIER_PROBS[1], "Pro") - .otherwise("Enterprise").alias("tier"), - # Region distribution - F.when(F.rand(SEED) < REGION_PROBS[0], "North") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1], "South") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1] + REGION_PROBS[2], "East") - .otherwise("West").alias("region"), - # Created date (within last 2 years before start date) - F.date_sub(F.lit(START_DATE.date()), (F.rand(SEED) * 730).cast("int")).alias("created_at"), - ) +customers_spec = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + # Surrogate key 0..N-1 (the implicit contiguous seed column) — stable join key + .withColumn("customer_sk", "long", expr="id") + # Business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # Realistic text via the Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + .withColumn("address", "string", text=FakerText("address")) + # Email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(regexp_replace(name, '[^A-Za-z]+', '.')), '@example.com')") + # Skewed categories (never uniform) — weights are relative frequencies + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=TIER_WEIGHTS, random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=REGION_WEIGHTS, random=True) + # Account created within 2 years before the analysis window + .withColumn("created_at", "date", + data_range=dg.DateRange( + (START_DATE - timedelta(days=730)).strftime(DATE_FMT), + START_DATE.strftime(DATE_FMT), + "days=1"), + random=True) + # Tier-based ARR via a log-normal (exp of a standard normal) — long tail, always positive. + # The omitted helper column '_z' keeps ARR reproducible under randomSeed. + # Enterprise ~ $1800, Pro ~ $245, Free ~ $55 median. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") ) -# Add tier-based ARR and email -customers_df = ( - customers_df - .withColumn("arr", F.round(generate_lognormal_amount(F.col("tier")), 2)) - .withColumn("email", fake_email(F.col("name"))) -) +customers_df = customers_spec.build() -# Save customers +# Save customers as raw Parquet customers_df.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/customers") print(f" Saved customers to {VOLUME_PATH}/customers") @@ -183,46 +156,45 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: # ============================================================================= print(f"\nGenerating {N_ORDERS:,} orders with referential integrity...") -# Write customer lookup to temp Delta table (no .cache() on serverless!) +# Write a customer lookup to a temp Delta table (no .cache() on serverless!) customers_tmp_table = f"{CATALOG}.{SCHEMA}._tmp_customers_lookup" -customers_df.select("customer_id", "tier").write.mode("overwrite").saveAsTable(customers_tmp_table) +(customers_df.select("customer_sk", "customer_id", "tier") + .write.mode("overwrite").saveAsTable(customers_tmp_table)) customer_lookup = spark.table(customers_tmp_table) -# Generate orders base -orders_df = ( - spark.range(0, N_ORDERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("ORD-"), F.lpad(F.col("id").cast("string"), 6, "0")).alias("order_id"), - # Generate customer_idx for FK join (hash-based distribution) - (F.abs(F.hash(F.col("id"), F.lit(SEED))) % N_CUSTOMERS).alias("customer_idx"), - # Order status - F.when(F.rand(SEED) < 0.65, "delivered") - .when(F.rand(SEED) < 0.80, "shipped") - .when(F.rand(SEED) < 0.90, "processing") - .when(F.rand(SEED) < 0.95, "pending") - .otherwise("cancelled").alias("status"), - # Order date within date range - F.date_add(F.lit(START_DATE.date()), (F.rand(SEED) * 180).cast("int")).alias("order_date"), - ) +# Generate orders. The FK (customer_sk) is drawn from the parent key range, so every +# value is valid by construction. A Gamma distribution skews it 80/20 (a few customers +# place most orders). +orders_spec = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + .withColumn("order_id", "string", baseColumn="id", + expr="concat('ORD-', lpad(cast(id as string), 6, '0'))") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .withColumn("status", "string", values=STATUS_VALUES, weights=STATUS_WEIGHTS, random=True) + .withColumn("order_date", "date", + data_range=dg.DateRange(START_DATE.strftime(DATE_FMT), + END_DATE.strftime(DATE_FMT), "days=1"), + random=True) ) +orders_df = orders_spec.build() -# Add customer_idx to lookup for join -customer_lookup_with_idx = customer_lookup.withColumn( - "customer_idx", - (F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) - 1).cast("int") -) - -# Join to get customer_id and tier as foreign key -orders_with_fk = ( - orders_df - .join(customer_lookup_with_idx, on="customer_idx", how="left") - .drop("customer_idx") -) +# Join to the parent (read back from Delta) to attach valid customer_id + tier +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") -# Add tier-based amount +# Tier-based amount (log-normal), coherent with the customer's tier orders_with_fk = orders_with_fk.withColumn( "amount", - F.round(generate_lognormal_amount(F.col("tier")), 2) + F.round( + F.expr( + "CASE tier " + "WHEN 'Enterprise' THEN exp(7.5 + 0.8 * randn()) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * randn()) " + "ELSE exp(4.0 + 0.6 * randn()) END" + ), + 2, + ), ) # ============================================================================= @@ -236,19 +208,19 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: outlier_count = int(N_ORDERS * BAD_DATA_CONFIG["outlier_rate"]) orphan_count = int(N_ORDERS * BAD_DATA_CONFIG["orphan_fk_rate"]) - # Add bad data flags + # Add a row number to target specific rows (no cache/persist on serverless) orders_with_fk = orders_with_fk.withColumn( "row_num", F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) ) - # Inject nulls in customer_id for first null_count rows + # Inject nulls in customer_id for the first null_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when(F.col("row_num") <= null_count, None).otherwise(F.col("customer_id")) ) - # Inject negative amounts for next outlier_count rows + # Inject negative amounts for the next outlier_count rows orders_with_fk = orders_with_fk.withColumn( "amount", F.when( @@ -257,7 +229,7 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: ).otherwise(F.col("amount")) ) - # Inject orphan FKs for next orphan_count rows + # Inject orphan FKs for the next orphan_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when( @@ -273,8 +245,8 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f" Injected {outlier_count} negative amounts") print(f" Injected {orphan_count} orphan foreign keys") -# Drop tier column (not needed in final output) -orders_final = orders_with_fk.drop("tier") +# Drop join-only columns (not needed in final output) +orders_final = orders_with_fk.drop("tier", "customer_sk") # Save orders orders_final.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/orders") diff --git a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/SKILL.md b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/SKILL.md index f534f60f..fe2b70c6 100644 --- a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/SKILL.md +++ b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/SKILL.md @@ -1,6 +1,6 @@ --- name: databricks-synthetic-data-gen -description: "Generate realistic synthetic data using Spark + Faker (strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'Faker', or 'sample data'." +description: "Generate realistic synthetic data using Spark + dbldatagen (Databricks Labs Data Generator, strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'dbldatagen', or 'sample data'." compatibility: Requires databricks CLI (>= v1.0.0) metadata: version: "0.1.0" @@ -11,7 +11,9 @@ parent: databricks-core # Databricks Synthetic Data Generation -Generate realistic, story-driven synthetic data for Databricks using **Spark + Faker + Pandas UDFs** (strongly recommended). +Generate realistic, story-driven synthetic data for Databricks using **Spark + [dbldatagen](https://databrickslabs.github.io/dbldatagen/public_docs/index.html)** (the Databricks Labs Data Generator — strongly recommended). + +> **Use the public `dbldatagen` API.** Import the top-level package (`import dbldatagen as dg`) and its public submodules (`dbldatagen.distributions`, `dbldatagen.constraints`). Everything below is built from the public API. ## Data Must Tell a Business Story @@ -44,11 +46,11 @@ Synthetic data should demonstrate how Databricks helps solve real business probl 5. **Enough data for trends** — ~100K+ rows for main tables so patterns survive aggregation 6. **Ask for catalog/schema** — Never default, always confirm before generating 7. **Present plan for approval** — Show tables, distributions, assumptions before writing code -8. **Master tables first** — Generate parent tables, write to Delta, then create children with valid FKs -9. **Use Spark + Faker + Pandas UDFs** — Scalable, parallel. Polars only if user explicitly wants local + <30K rows +8. **Parent tables first** — Generate parent tables, write to Delta, then create children with valid FKs +9. **Use Spark + dbldatagen** — Scalable, parallel, declarative. Build a `dg.DataGenerator` spec and `.build()` it into a Spark DataFrame. Use a `pandas_udf` only for logic dbldatagen can't express 10. **Use Databricks Connect Serverless by default to generate data** — Update databricks-connect on python 3.12 if required (avoid using execute_code unless instructed to not use Databricks Connect) 11. **No `.cache()` or `.persist()`** — Not supported on serverless. Write to Delta, read back for joins -12. **No Python loops or `.collect()`** — Use Spark parallelism. No driver-side iteration, avoid Pandas↔Spark conversions +12. **No Python loops or `.collect()`** — Use Spark parallelism and dbldatagen specs. No driver-side iteration, avoid Pandas↔Spark conversions ## Generation Planning Workflow @@ -142,29 +144,57 @@ SELECT COUNT(*) FROM parquet.\`/Volumes/CATALOG/SCHEMA/raw_data/customers\` See [references/2-troubleshooting.md](references/2-troubleshooting.md) for full validation examples. -## Use Databricks Connect Spark + Faker Pattern +## Use Databricks Connect + dbldatagen Pattern + +A `DataGenerator` spec declares each column; `.build()` returns a Spark DataFrame. No UDFs, no driver loops — generation is fully distributed across `partitions`. Install `dbldatagen` and `faker` locally first (see Setup). ```python from databricks.connect import DatabricksSession -from pyspark.sql import functions as F -from pyspark.sql.types import StringType -import pandas as pd +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import company, person # providers we draw from -# Setup serverless Spark session +# Setup serverless Spark session (deps installed locally) spark = DatabricksSession.builder.serverless(True).getOrCreate() -# Pandas UDF pattern - import lib INSIDE the function (libs must be installed locally) -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - from faker import Faker # Import inside UDF - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) - -# Generate with spark.range, apply UDFs -customers_df = spark.range(0, 10000, numPartitions=16).select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), +CATALOG, SCHEMA = "", "" # always user-supplied +N_CUSTOMERS = 10_000 + +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company]) + +customers = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, + randomSeed=42, randomSeedMethod="hash_fieldname") + # surrogate key 0..N-1 (the implicit contiguous seed column) — drives FK joins + .withColumn("customer_sk", "long", expr="id") + # business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # realistic text via the shared Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + # email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(replace(name, ' ', '.')), '@example.com')") + # skewed categories — NEVER uniform (weights are relative) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=[40, 25, 20, 15], random=True) + # right-skewed ARR correlated to tier: log-normal = exp() of a standard normal. + # The hidden helper column (_z, omit=True) is computed and reusable as a base column. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") + .withColumn("created_at", "date", + data_range=dg.DateRange("2023-01-01 00:00:00", "2024-12-31 00:00:00", "days=1"), + random=True) ) +customers_df = customers.build() # Write to Volume as Parquet (default for raw data) # Path is a folder with table name: /Volumes/catalog/schema/raw_data/customers/ @@ -173,7 +203,7 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") customers_df.write.mode("overwrite").parquet(f"/Volumes/{CATALOG}/{SCHEMA}/raw_data/customers") ``` -**Partitions by scale:** `spark.range(N, numPartitions=P)` +**Partitions by scale:** `dg.DataGenerator(..., rows=N, partitions=P)` - <100K rows: 8 partitions - 100K-500K: 16 partitions - 500K-1M: 32 partitions @@ -190,32 +220,79 @@ Generated scripts must be highly performant. **Never** do these: | Anti-Pattern | Why It's Slow | Do This Instead | |--------------|---------------|-----------------| -| Python loops on driver | Single-threaded, no parallelism | Use `spark.range()` + Spark operations | +| Python loops on driver | Single-threaded, no parallelism | Declare columns in a `dg.DataGenerator` spec and `.build()` | | `.collect()` then iterate | Brings all data to driver memory | Keep data in Spark, use DataFrame ops | -| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark, use `pandas_udf` only for UDFs | +| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark; let dbldatagen generate columns | | Read/write temp files | Unnecessary I/O | Chain DataFrame transformations | -| Scalar UDFs | Row-by-row processing | Use `pandas_udf` for batch processing | +| Scalar UDFs | Row-by-row processing | Use dbldatagen `expr`/`template`/`text`; `pandas_udf` only when unavoidable | -**Good pattern:** `spark.range()` → Spark transforms → `pandas_udf` for Faker → write directly +**Good pattern:** `dg.DataGenerator(rows, partitions)` → declarative `.withColumn(...)` specs → `.build()` → write directly ## Common Patterns +All snippets use the public `dbldatagen` API (`import dbldatagen as dg`, `import dbldatagen.distributions as dist`). + ### Weighted Categories (never uniform) +`weights` are relative frequencies — they don't need to sum to 100: ```python -F.when(F.rand() < 0.6, "Free").when(F.rand() < 0.9, "Pro").otherwise("Enterprise") +.withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) ``` -### Log-Normal Amounts (in a pandas UDF) -Use `np.random.lognormal(mean, sigma)` — always positive, long tail: -- Enterprise: `lognormal(7.5, 0.8)` → ~$1800 median -- Pro: `lognormal(5.5, 0.7)` → ~$245 median -- Free: `lognormal(4.0, 0.6)` → ~$55 median +### Skewed / Long-Tailed Amounts +Apply a continuous distribution to a numeric range — `Gamma`/`Exponential` give the always-positive long tail you'd want from log-normal: +```python +.withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) +``` +For a true log-normal, generate over a normal and exponentiate via `expr`: +```python +.withColumn("amount", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # standard normal, hidden +.withColumn("arr", "double", baseColumn="amount", + expr="round(exp(5.5 + 0.8 * amount), 2)") # median ~$245 +``` -### Date Range (Last 6 Months) +### Coherent Rows (correlated attributes via `expr` + `baseColumn`) +Derive dependent columns from earlier ones so each row makes business sense — no UDF needed: ```python -END_DATE = datetime.now() +.withColumn("priority", "string", values=["Critical", "High", "Medium", "Low"], + weights=[5, 15, 50, 30], random=True) +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*8 " + "WHEN 'High' THEN rand()*24 WHEN 'Medium' THEN rand()*72 " + "ELSE rand()*120 END, 1)") +.withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 4 " + "WHEN resolution_hours<72 THEN 3 ELSE 2 END") +``` + +### Date / Timestamp Range (Last 6 Months) +Use `dg.DateRange(begin, end, interval)` with the `data_range` option: +```python +from datetime import datetime, timedelta +END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +FMT = "%Y-%m-%d %H:%M:%S" + +.withColumn("order_ts", "timestamp", + data_range=dg.DateRange(START_DATE.strftime(FMT), END_DATE.strftime(FMT), "minutes=30"), + random=True) +``` + +### Realistic Text (Faker provider plugin) +Build one `FakerTextFactory` with a default locale and the faker providers you need, then reuse it: +```python +from faker.providers import person, company, internet +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, internet]) + +.withColumn("name", "string", text=FakerText("name")) # uses the 'person' provider +.withColumn("company", "string", text=FakerText("company")) # uses the 'company' provider +.withColumn("ip", "string", text=FakerText("ipv4_private")) # uses the 'internet' provider ``` +dbldatagen alternatives that need no library: +* `template=r"\w.\w@\w.com"` (templated text) +* `text=dg.ILText(paragraphs=(1, 3), sentences=(2, 5))` (lorem-ipsum free text). ### Infrastructure (always create in script) ```python @@ -224,27 +301,33 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") ``` ### Referential Integrity (FK pattern) -Write master table to Delta first, then read back for FK joins (no `.cache()` on serverless): +Generate the FK as an integer over the parent's surrogate-key range so **every value is valid by construction**, then write the parent to Delta and read it back to attach coherent parent attributes (no `.cache()` on serverless): ```python -# 1. Write master table +# 1. Parent table: surrogate key 0..N-1, then write to Delta customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") -# 2. Read back for FK lookup -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_idx", "customer_id") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +# 2. Child table: FK drawn from the parent key range (valid by construction). +# A distribution skews it 80/20 (a few customers place most orders). +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("order_id", "string", + expr="concat('ORD-', lpad(cast(id as string), 8, '0'))", baseColumn="id") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") + +# 3. Join to the parent (read back from Delta) to pull coherent parent attributes +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "customer_id", "tier") +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") ``` ## Setup -Requires Python 3.12 and databricks-connect>=16.4. Use `uv`: +Requires Python 3.12 and databricks-connect>=16.4. Install dependencies locally with `uv`: ```bash -uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays +uv pip install "databricks-connect>=16.4,<17.4" dbldatagen faker ``` ## Related Skills @@ -256,10 +339,11 @@ uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays | Issue | Solution | |-------|----------| -| `ModuleNotFoundError: faker` | Install locally: `uv pip install faker`, import inside UDF | -| Faker UDF is slow | Use `pandas_udf` for batch processing | -| Out of memory | Increase `numPartitions` in `spark.range()` | -| Referential integrity errors | Write master table to Delta first, read back for FK joins | +| `ModuleNotFoundError: dbldatagen` (or `faker`) | Install locally: `uv pip install dbldatagen faker` | +| `FakerText`/provider not found | Pass the provider to `dg.FakerTextFactory(providers=[...])` and import it from `faker.providers` | +| All rows identical / not random | Set `random=True` on the column (default is deterministic), or set `randomSeed`/`randomSeedMethod` on the generator | +| Out of memory | Increase `partitions` in `dg.DataGenerator(..., partitions=P)` | +| Referential integrity errors | Draw the FK from the parent key range (`minValue=0, maxValue=N-1`); write parent to Delta first, read back to join attributes | | `PERSIST TABLE is not supported on serverless` | **NEVER use `.cache()` or `.persist()` with serverless** - write to Delta table first, then read back | | `F.window` vs `Window` confusion | Use `from pyspark.sql.window import Window` for `row_number()`, `rank()`, etc. `F.window` is for streaming only. | | Broadcast variables not supported | **NEVER use `spark.sparkContext.broadcast()` with serverless** | diff --git a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/1-data-patterns.md b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/1-data-patterns.md index eba64916..87774896 100644 --- a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/1-data-patterns.md +++ b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/1-data-patterns.md @@ -28,15 +28,31 @@ Real data is never uniformly distributed. Use appropriate distributions: | **Exponential** | Time between events | Support resolution time, session duration | | **Weighted categorical** | Skewed categories | Status (70% complete, 5% failed), tiers | -```python -# Log-normal for amounts (long tail, always positive) -amount = np.random.lognormal(mean=5.5, sigma=0.8) # ~$245 median +dbldatagen applies a distribution to a column's random draw via the `distribution=` option +(`import dbldatagen.distributions as dist`). Supported: `Normal`, `Gamma`, `Beta`, `Exponential`. -# Pareto for power-law (few large, many small) -value = (np.random.pareto(a=1.5) + 1) * base_value +```python +import dbldatagen as dg +import dbldatagen.distributions as dist + +spec = ( + dg.DataGenerator(spark, name="amounts", rows=100_000, partitions=8, randomSeed=42) + # Use Gamma for right-tailed distributions like amounts + .withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) + # Use Exponential for waiting times + .withColumn("resolution_hours", "double", minValue=0, maxValue=240, + random=True, distribution=dist.Exponential(rate=1.0 / 24)) +) -# Exponential for time-to-event -hours = np.random.exponential(scale=24) # avg 24h, skewed right +# Exponentiate a standard normal via expr for true log-normal distributions +spec = ( + dg.DataGenerator(spark, name="ln", rows=100_000, partitions=8, randomSeed=42) + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # hidden helper column + .withColumn("amount", "double", baseColumn="_z", + expr="round(exp(5.5 + 0.8 * _z), 2)") # ~$245 median +) ``` ### 3. Row Coherence @@ -51,20 +67,24 @@ Attributes within a row must make business sense together. Generate correlated a | Large transaction + unusual hour | Higher fraud probability | | Fast resolution | Higher CSAT score | +In dbldatagen, chain `baseColumn` + `expr` when columns derive from other columns +in the same dataset. The generator resolves the dependency order automatically: + ```python -@F.pandas_udf("struct") -def generate_coherent_ticket(tiers: pd.Series) -> pd.DataFrame: - """All attributes correlate logically within each row.""" - results = [] - for tier in tiers: - # Priority depends on tier - priority = "Critical" if tier == "Enterprise" and random() < 0.3 else "Medium" - # Resolution depends on priority - resolution = np.random.exponential(4 if priority == "Critical" else 36) - # CSAT depends on resolution - csat = 5 if resolution < 4 else (3 if resolution < 24 else 2) - results.append({"priority": priority, "resolution_hours": resolution, "csat": csat}) - return pd.DataFrame(results) +spec = ( + dg.DataGenerator(spark, name="tickets", rows=80_000, partitions=8, randomSeed=42) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + # Priority depends on tier + .withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' AND rand()<0.3 THEN 'Critical' ELSE 'Medium' END") + # Resolution depends on priority + .withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 ELSE rand()*36 END, 1)") + # CSAT depends on resolution + .withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 3 ELSE 2 END") +) ``` ### 4. The 80/20 Rule @@ -75,7 +95,13 @@ Apply power-law distributions where appropriate: - **20% of products** account for 80% of sales - **20% of support agents** handle 80% of tickets -Implementation: Use weighted sampling when assigning FKs, not uniform random. +Implementation: Skew FKs by drawing from a non-uniform distribution. + +```python +# Use a Gamma distribution to skew FK values +.withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) +``` ### 5. Time-Based Patterns @@ -86,14 +112,22 @@ Most data has temporal patterns: - **Seasonality** — Q4 retail spike, summer travel peak - **Trends** — Growth over time, degradation curves +Bias *when* events happen by weighting time buckets, then derive the timestamp from the bucket. +The example below clusters tickets into business hours and skews volume toward weekdays: + ```python -def get_volume_multiplier(date): - multiplier = 1.0 - if date.weekday() >= 5: multiplier *= 0.6 # Weekend drop - if date.month in [11, 12]: multiplier *= 1.5 # Holiday spike - return multiplier +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], # 9am–5pm peak + random=True, omit=True) +.withColumn("day_offset", "int", minValue=0, maxValue=179, random=True, + distribution=dist.Normal(60, 30), omit=True) # volume ramps over the window +.withColumn("event_ts", "timestamp", baseColumn=["day_offset", "hour"], + expr="date_add(timestamp'2025-01-01 00:00:00', day_offset) + make_interval(0,0,0,0,hour,0,0)") ``` +For holiday/seasonal spikes, weight a month or week-of-year bucket the same way, or post-filter +the built DataFrame with a Spark `expr` multiplier. + ### 6. ML-Ready Data If data will train ML models, ensure: @@ -105,20 +139,32 @@ If data will train ML models, ensure: ## Referential Integrity -Generate master tables first, write to Delta, then join for FKs: +Give the parent a surrogate key over `0..N-1`, draw the child's FK from that same range (valid by +construction), then write the parent to Delta and read it back to attach parent attributes: ```python -# 1. Generate and write master table +# 1. # 1. Generate and write parent table with surrogate keys from 0..N-1 +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_id", "string", baseColumn="id", # References the built-in sequential 'id' column + expr="concat('CUST-', lpad(cast(id as string), 5, '0'))") + # ... other attributes ... + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK joins (NOT cache - unsupported on serverless) -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") + +# 3. Generate child table with FKs drawn from the parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True, omit=True) + .withColumn("customer_id", "string", baseColumn="customer_sk", # References the 'customer_sk' column + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") +orders_with_fk = orders_df.join(customer_lookup, on="customer_id", how="inner") ``` ## Data Volume diff --git a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md index 793b64f7..8a19ebc8 100644 --- a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md +++ b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md @@ -4,23 +4,24 @@ Common issues and solutions for synthetic data generation. ## Environment Issues -### ModuleNotFoundError: faker (or other library) +### ModuleNotFoundError: dbldatagen (or faker) -**Problem:** Dependencies not available in execution environment. +**Problem:** Dependencies not available in execution environment. `dbldatagen` is required, and +`faker` too if you use `FakerTextFactory`/`fakerText`. **Solutions by execution mode:** | Mode | Solution | |------|----------| -| **DB Connect with Serverless** | Install libs locally (`uv pip install faker`), use `DatabricksSession.builder.serverless(True)` | -| **Databricks Runtime** | Use Databricks CLI to install `faker holidays` | -| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "faker"}}, {"pypi": {"package": "holidays"}}]}'` | +| **DB Connect with Serverless** | Install libs locally (`uv pip install dbldatagen faker`), use `DatabricksSession.builder.serverless(True)` | +| **Databricks Runtime** | `%pip install dbldatagen faker` at the top of the notebook | +| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "dbldatagen"}}, {"pypi": {"package": "faker"}}]}'` | ```python # For DB Connect with serverless from databricks.connect import DatabricksSession -# Install dependencies locally first: uv pip install faker pandas numpy holidays +# Install dependencies locally first: uv pip install dbldatagen faker spark = DatabricksSession.builder.serverless(True).getOrCreate() ``` @@ -50,21 +51,21 @@ AnalysisException: [NOT_SUPPORTED_WITH_SERVERLESS] PERSIST TABLE is not supporte **Why this happens:** Serverless compute does not support caching DataFrames in memory. This is a fundamental limitation of the serverless architecture. -**Solution:** Write master tables to Delta first, then read them back for FK joins: +**Solution:** Write parent tables to Delta first, then read them back for FK joins: ```python # BAD - will fail on serverless -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.cache() # ❌ FAILS: "PERSIST TABLE is not supported on serverless compute" # GOOD - write to Delta, then read back -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") # ✓ Read from Delta ``` **Best practice for referential integrity:** -1. Generate master table (e.g., customers) +1. Generate parent table (e.g., customers) 2. Write to Delta table 3. Read back for FK lookup joins 4. Generate child tables (e.g., orders, tickets) with valid FKs @@ -114,7 +115,7 @@ spark = DatabricksSession.builder.serverless(True).getOrCreate() "environment_key": "datagen_env", "spec": { "client": "4", # Required! - "dependencies": ["faker", "numpy", "pandas"] + "dependencies": ["dbldatagen", "faker"] } }] } @@ -153,34 +154,37 @@ contacts_df = contacts_df.withColumn( --- -### Faker UDF is slow +### Generation is slow -**Problem:** Single-row UDFs don't parallelize well. +**Problem:** Row-by-row Python UDFs (or driver loops) don't parallelize well. -**Solution:** Use `pandas_udf` for batch processing: +**Solution:** Let dbldatagen generate columns declaratively — it builds the whole DataFrame in +parallel across `partitions`. Prefer `expr`, `template`, `values`/`weights`, and `distribution` +over UDFs. For names/addresses use the Faker plugin (`FakerTextFactory`) instead of hand-rolled UDFs: ```python -# SLOW - scalar UDF -@F.udf(returnType=StringType()) -def slow_fake_name(): - return Faker().name() - -# FAST - pandas UDF (batch processing) -@F.pandas_udf(StringType()) -def fast_fake_name(ids: pd.Series) -> pd.Series: - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) +import dbldatagen as dg +from faker.providers import person + +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person]) + +spec = ( + dg.DataGenerator(spark, name="people", rows=1_000_000, partitions=64, randomSeed=42) + .withColumn("name", "string", text=FakerText("name")) # batched by dbldatagen + .withColumn("status", "string", values=["active", "churned"], weights=[85, 15], random=True) +) +df = spec.build() ``` ### Out of memory with large data **Problem:** Not enough partitions for data size. -**Solution:** Increase partitions: +**Solution:** Increase `partitions` on the generator: ```python # For large datasets (1M+ rows) -customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from default +spec = dg.DataGenerator(spark, name="big", rows=N_CUSTOMERS, partitions=64, randomSeed=42) ``` | Data Size | Recommended Partitions | @@ -205,21 +209,28 @@ customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from de **Problem:** Foreign keys reference non-existent parent records. -**Solution:** Write master table to Delta first, then read back for FK joins: +**Solution:** Sample the FK from the parent's surrogate-key range, write the parent +to Delta, then read it back for joins: ```python -# 1. Generate and WRITE master table (do NOT use cache with serverless!) -customers_df = spark.range(0, N_CUSTOMERS)... +# 1. Generate and WRITE parent table (do NOT use cache with serverless!) +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_sk", "long", expr="id") # References the built-in sequential 'id' column + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], weights=[60, 30, 10], random=True) + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK lookups -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") - -# 3. Generate child table with valid FKs -orders_df = spark.range(0, N_ORDERS).join( - customer_lookup, - on=, - how="left" +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "tier") + +# 3. Generate child table with FK in the valid parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True) + .build() + .join(customer_lookup, on="customer_sk", how="inner") ) ``` @@ -233,52 +244,52 @@ orders_df = spark.range(0, N_ORDERS).join( **Problem:** All customers have similar order counts, amounts are evenly distributed. -**Solution:** Use non-linear distributions: +**Solution:** Apply a `distribution=` to the column instead of leaving it uniform: ```python -# BAD - uniform -amounts = np.random.uniform(10, 1000, N) +import dbldatagen.distributions as dist -# GOOD - log-normal (realistic) -amounts = np.random.lognormal(mean=5, sigma=0.8, N) +# BAD - uniform (no distribution) +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True) + +# GOOD - skewed, realistic +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True, + distribution=dist.Gamma(1.0, 2.0)) ``` ### Missing time-based patterns **Problem:** Data doesn't reflect weekday/weekend or seasonal patterns. -**Solution:** Add multipliers: +**Solution:** Weight time buckets so volume clusters realistically, then derive the timestamp: ```python -import holidays - -US_HOLIDAYS = holidays.US(years=[2024, 2025]) - -def get_multiplier(date): - mult = 1.0 - if date.weekday() >= 5: # Weekend - mult *= 0.6 - if date in US_HOLIDAYS: - mult *= 0.3 - return mult +# Cluster events into business hours +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], + random=True, omit=True) +.withColumn("event_ts", "timestamp", baseColumn="hour", + expr="date'2025-06-01' + make_interval(0,0,0,0,hour,0,0)") ``` +For weekend/holiday dips, weight a day-of-week or day bucket the same way, or post-filter the built +DataFrame with a Spark `expr` (e.g. `dayofweek(event_ts) IN (1,7)`). + ### Incoherent row attributes **Problem:** Enterprise customer has low-value orders, critical ticket has slow resolution. -**Solution:** Correlate attributes: +**Solution:** Correlate attributes with `baseColumn` + `expr` so each derives from the previous: ```python # Priority based on tier -if tier == 'Enterprise': - priority = np.random.choice(['Critical', 'High'], p=[0.4, 0.6]) -else: - priority = np.random.choice(['Medium', 'Low'], p=[0.6, 0.4]) - -# Resolution based on priority -resolution_scale = {'Critical': 4, 'High': 12, 'Medium': 36, 'Low': 72} -resolution_hours = np.random.exponential(scale=resolution_scale[priority]) +.withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' THEN (CASE WHEN rand()<0.4 THEN 'Critical' ELSE 'High' END) " + "ELSE (CASE WHEN rand()<0.6 THEN 'Medium' ELSE 'Low' END) END") +# Resolution scaled by priority +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 WHEN 'High' THEN rand()*12 " + "WHEN 'Medium' THEN rand()*36 ELSE rand()*72 END, 1)") ``` --- diff --git a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py index b36edb8e..185e4a18 100644 --- a/plugins/databricks/codex/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py +++ b/plugins/databricks/codex/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py @@ -1,22 +1,24 @@ -"""Generate synthetic data using Spark + Faker + Pandas UDFs. +"""Generate synthetic data using Spark + dbldatagen (Databricks Labs Data Generator). This is the recommended approach for ALL data generation tasks: - Scales from thousands to millions of rows -- Parallel execution via Spark +- Declarative, parallel generation (no driver loops, no row-by-row UDFs) - Direct write to Unity Catalog - Works with serverless and classic compute +Uses ONLY the public dbldatagen API: +- `dbldatagen.DataGenerator` / `.withColumn(...)` / `.build()` +- `dbldatagen.distributions` for skew +- `dbldatagen.FakerTextFactory` for realistic names/companies/addresses + Prerequisites: -- Install dependencies locally: uv pip install faker pandas numpy holidays databricks-connect +- Install dependencies locally: uv pip install dbldatagen faker databricks-connect - Configure ~/.databrickscfg with serverless_compute_id = auto """ import sys import os from pyspark.sql import functions as F from pyspark.sql.window import Window -from pyspark.sql.types import StringType, DoubleType, StructType, StructField, IntegerType -import numpy as np -import pandas as pd from datetime import datetime, timedelta # ============================================================================= @@ -39,6 +41,7 @@ # Date range - last 6 months from today END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +DATE_FMT = "%Y-%m-%d %H:%M:%S" # Write mode - "overwrite" for one-time, "append" for incremental WRITE_MODE = "overwrite" @@ -51,14 +54,18 @@ "orphan_fk_rate": 0.01, # 1% orphan foreign keys } -# Reproducibility +# Reproducibility - same seed => same data SEED = 42 -# Tier distribution: Free 60%, Pro 30%, Enterprise 10% -TIER_PROBS = [0.6, 0.3, 0.1] +# Tier distribution: Free 60%, Pro 30%, Enterprise 10% (relative weights) +TIER_WEIGHTS = [60, 30, 10] + +# Region distribution (relative weights) +REGION_WEIGHTS = [40, 25, 20, 15] -# Region distribution -REGION_PROBS = [0.4, 0.25, 0.2, 0.15] +# Order status distribution (relative weights) +STATUS_VALUES = ["delivered", "shipped", "processing", "pending", "cancelled"] +STATUS_WEIGHTS = [65, 15, 10, 5, 5] # ============================================================================= # SESSION CREATION @@ -79,59 +86,14 @@ spark = DatabricksSession.builder.clusterId(CLUSTER_ID).getOrCreate() print(f"Connected to cluster {CLUSTER_ID}") -# Import Faker for UDF definitions -from faker import Faker - -# ============================================================================= -# DEFINE PANDAS UDFs FOR FAKER DATA -# ============================================================================= +# Import the public dbldatagen API (installed locally) +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import person, company, address as faker_address -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - """Generate realistic person names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.name() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_company(ids: pd.Series) -> pd.Series: - """Generate realistic company names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.company() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_address(ids: pd.Series) -> pd.Series: - """Generate realistic addresses.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.address().replace('\n', ', ') for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_email(names: pd.Series) -> pd.Series: - """Generate email based on name.""" - emails = [] - for name in names: - if name: - domain = name.lower().replace(" ", ".").replace(",", "")[:20] - emails.append(f"{domain}@example.com") - else: - emails.append("unknown@example.com") - return pd.Series(emails) - -@F.pandas_udf(DoubleType()) -def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: - """Generate amount based on tier using log-normal distribution.""" - np.random.seed(SEED) - amounts = [] - for tier in tiers: - if tier == "Enterprise": - amounts.append(float(np.random.lognormal(mean=7.5, sigma=0.8))) # ~$1800 avg - elif tier == "Pro": - amounts.append(float(np.random.lognormal(mean=5.5, sigma=0.7))) # ~$245 avg - else: - amounts.append(float(np.random.lognormal(mean=4.0, sigma=0.6))) # ~$55 avg - return pd.Series(amounts) +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, faker_address]) # ============================================================================= # CREATE INFRASTRUCTURE @@ -142,39 +104,50 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f"Infrastructure ready: {VOLUME_PATH}") # ============================================================================= -# GENERATE CUSTOMERS (Master Table) +# GENERATE CUSTOMERS (Parent Table) # ============================================================================= print(f"\nGenerating {N_CUSTOMERS:,} customers...") -customers_df = ( - spark.range(0, N_CUSTOMERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), - fake_company(F.col("id")).alias("company"), - fake_address(F.col("id")).alias("address"), - # Tier distribution: Free 60%, Pro 30%, Enterprise 10% - F.when(F.rand(SEED) < TIER_PROBS[0], "Free") - .when(F.rand(SEED) < TIER_PROBS[0] + TIER_PROBS[1], "Pro") - .otherwise("Enterprise").alias("tier"), - # Region distribution - F.when(F.rand(SEED) < REGION_PROBS[0], "North") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1], "South") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1] + REGION_PROBS[2], "East") - .otherwise("West").alias("region"), - # Created date (within last 2 years before start date) - F.date_sub(F.lit(START_DATE.date()), (F.rand(SEED) * 730).cast("int")).alias("created_at"), - ) +customers_spec = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + # Surrogate key 0..N-1 (the implicit contiguous seed column) — stable join key + .withColumn("customer_sk", "long", expr="id") + # Business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # Realistic text via the Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + .withColumn("address", "string", text=FakerText("address")) + # Email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(regexp_replace(name, '[^A-Za-z]+', '.')), '@example.com')") + # Skewed categories (never uniform) — weights are relative frequencies + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=TIER_WEIGHTS, random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=REGION_WEIGHTS, random=True) + # Account created within 2 years before the analysis window + .withColumn("created_at", "date", + data_range=dg.DateRange( + (START_DATE - timedelta(days=730)).strftime(DATE_FMT), + START_DATE.strftime(DATE_FMT), + "days=1"), + random=True) + # Tier-based ARR via a log-normal (exp of a standard normal) — long tail, always positive. + # The omitted helper column '_z' keeps ARR reproducible under randomSeed. + # Enterprise ~ $1800, Pro ~ $245, Free ~ $55 median. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") ) -# Add tier-based ARR and email -customers_df = ( - customers_df - .withColumn("arr", F.round(generate_lognormal_amount(F.col("tier")), 2)) - .withColumn("email", fake_email(F.col("name"))) -) +customers_df = customers_spec.build() -# Save customers +# Save customers as raw Parquet customers_df.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/customers") print(f" Saved customers to {VOLUME_PATH}/customers") @@ -183,46 +156,45 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: # ============================================================================= print(f"\nGenerating {N_ORDERS:,} orders with referential integrity...") -# Write customer lookup to temp Delta table (no .cache() on serverless!) +# Write a customer lookup to a temp Delta table (no .cache() on serverless!) customers_tmp_table = f"{CATALOG}.{SCHEMA}._tmp_customers_lookup" -customers_df.select("customer_id", "tier").write.mode("overwrite").saveAsTable(customers_tmp_table) +(customers_df.select("customer_sk", "customer_id", "tier") + .write.mode("overwrite").saveAsTable(customers_tmp_table)) customer_lookup = spark.table(customers_tmp_table) -# Generate orders base -orders_df = ( - spark.range(0, N_ORDERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("ORD-"), F.lpad(F.col("id").cast("string"), 6, "0")).alias("order_id"), - # Generate customer_idx for FK join (hash-based distribution) - (F.abs(F.hash(F.col("id"), F.lit(SEED))) % N_CUSTOMERS).alias("customer_idx"), - # Order status - F.when(F.rand(SEED) < 0.65, "delivered") - .when(F.rand(SEED) < 0.80, "shipped") - .when(F.rand(SEED) < 0.90, "processing") - .when(F.rand(SEED) < 0.95, "pending") - .otherwise("cancelled").alias("status"), - # Order date within date range - F.date_add(F.lit(START_DATE.date()), (F.rand(SEED) * 180).cast("int")).alias("order_date"), - ) +# Generate orders. The FK (customer_sk) is drawn from the parent key range, so every +# value is valid by construction. A Gamma distribution skews it 80/20 (a few customers +# place most orders). +orders_spec = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + .withColumn("order_id", "string", baseColumn="id", + expr="concat('ORD-', lpad(cast(id as string), 6, '0'))") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .withColumn("status", "string", values=STATUS_VALUES, weights=STATUS_WEIGHTS, random=True) + .withColumn("order_date", "date", + data_range=dg.DateRange(START_DATE.strftime(DATE_FMT), + END_DATE.strftime(DATE_FMT), "days=1"), + random=True) ) +orders_df = orders_spec.build() -# Add customer_idx to lookup for join -customer_lookup_with_idx = customer_lookup.withColumn( - "customer_idx", - (F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) - 1).cast("int") -) - -# Join to get customer_id and tier as foreign key -orders_with_fk = ( - orders_df - .join(customer_lookup_with_idx, on="customer_idx", how="left") - .drop("customer_idx") -) +# Join to the parent (read back from Delta) to attach valid customer_id + tier +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") -# Add tier-based amount +# Tier-based amount (log-normal), coherent with the customer's tier orders_with_fk = orders_with_fk.withColumn( "amount", - F.round(generate_lognormal_amount(F.col("tier")), 2) + F.round( + F.expr( + "CASE tier " + "WHEN 'Enterprise' THEN exp(7.5 + 0.8 * randn()) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * randn()) " + "ELSE exp(4.0 + 0.6 * randn()) END" + ), + 2, + ), ) # ============================================================================= @@ -236,19 +208,19 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: outlier_count = int(N_ORDERS * BAD_DATA_CONFIG["outlier_rate"]) orphan_count = int(N_ORDERS * BAD_DATA_CONFIG["orphan_fk_rate"]) - # Add bad data flags + # Add a row number to target specific rows (no cache/persist on serverless) orders_with_fk = orders_with_fk.withColumn( "row_num", F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) ) - # Inject nulls in customer_id for first null_count rows + # Inject nulls in customer_id for the first null_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when(F.col("row_num") <= null_count, None).otherwise(F.col("customer_id")) ) - # Inject negative amounts for next outlier_count rows + # Inject negative amounts for the next outlier_count rows orders_with_fk = orders_with_fk.withColumn( "amount", F.when( @@ -257,7 +229,7 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: ).otherwise(F.col("amount")) ) - # Inject orphan FKs for next orphan_count rows + # Inject orphan FKs for the next orphan_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when( @@ -273,8 +245,8 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f" Injected {outlier_count} negative amounts") print(f" Injected {orphan_count} orphan foreign keys") -# Drop tier column (not needed in final output) -orders_final = orders_with_fk.drop("tier") +# Drop join-only columns (not needed in final output) +orders_final = orders_with_fk.drop("tier", "customer_sk") # Save orders orders_final.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/orders") diff --git a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/SKILL.md b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/SKILL.md index f534f60f..fe2b70c6 100644 --- a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/SKILL.md +++ b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/SKILL.md @@ -1,6 +1,6 @@ --- name: databricks-synthetic-data-gen -description: "Generate realistic synthetic data using Spark + Faker (strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'Faker', or 'sample data'." +description: "Generate realistic synthetic data using Spark + dbldatagen (Databricks Labs Data Generator, strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'dbldatagen', or 'sample data'." compatibility: Requires databricks CLI (>= v1.0.0) metadata: version: "0.1.0" @@ -11,7 +11,9 @@ parent: databricks-core # Databricks Synthetic Data Generation -Generate realistic, story-driven synthetic data for Databricks using **Spark + Faker + Pandas UDFs** (strongly recommended). +Generate realistic, story-driven synthetic data for Databricks using **Spark + [dbldatagen](https://databrickslabs.github.io/dbldatagen/public_docs/index.html)** (the Databricks Labs Data Generator — strongly recommended). + +> **Use the public `dbldatagen` API.** Import the top-level package (`import dbldatagen as dg`) and its public submodules (`dbldatagen.distributions`, `dbldatagen.constraints`). Everything below is built from the public API. ## Data Must Tell a Business Story @@ -44,11 +46,11 @@ Synthetic data should demonstrate how Databricks helps solve real business probl 5. **Enough data for trends** — ~100K+ rows for main tables so patterns survive aggregation 6. **Ask for catalog/schema** — Never default, always confirm before generating 7. **Present plan for approval** — Show tables, distributions, assumptions before writing code -8. **Master tables first** — Generate parent tables, write to Delta, then create children with valid FKs -9. **Use Spark + Faker + Pandas UDFs** — Scalable, parallel. Polars only if user explicitly wants local + <30K rows +8. **Parent tables first** — Generate parent tables, write to Delta, then create children with valid FKs +9. **Use Spark + dbldatagen** — Scalable, parallel, declarative. Build a `dg.DataGenerator` spec and `.build()` it into a Spark DataFrame. Use a `pandas_udf` only for logic dbldatagen can't express 10. **Use Databricks Connect Serverless by default to generate data** — Update databricks-connect on python 3.12 if required (avoid using execute_code unless instructed to not use Databricks Connect) 11. **No `.cache()` or `.persist()`** — Not supported on serverless. Write to Delta, read back for joins -12. **No Python loops or `.collect()`** — Use Spark parallelism. No driver-side iteration, avoid Pandas↔Spark conversions +12. **No Python loops or `.collect()`** — Use Spark parallelism and dbldatagen specs. No driver-side iteration, avoid Pandas↔Spark conversions ## Generation Planning Workflow @@ -142,29 +144,57 @@ SELECT COUNT(*) FROM parquet.\`/Volumes/CATALOG/SCHEMA/raw_data/customers\` See [references/2-troubleshooting.md](references/2-troubleshooting.md) for full validation examples. -## Use Databricks Connect Spark + Faker Pattern +## Use Databricks Connect + dbldatagen Pattern + +A `DataGenerator` spec declares each column; `.build()` returns a Spark DataFrame. No UDFs, no driver loops — generation is fully distributed across `partitions`. Install `dbldatagen` and `faker` locally first (see Setup). ```python from databricks.connect import DatabricksSession -from pyspark.sql import functions as F -from pyspark.sql.types import StringType -import pandas as pd +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import company, person # providers we draw from -# Setup serverless Spark session +# Setup serverless Spark session (deps installed locally) spark = DatabricksSession.builder.serverless(True).getOrCreate() -# Pandas UDF pattern - import lib INSIDE the function (libs must be installed locally) -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - from faker import Faker # Import inside UDF - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) - -# Generate with spark.range, apply UDFs -customers_df = spark.range(0, 10000, numPartitions=16).select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), +CATALOG, SCHEMA = "", "" # always user-supplied +N_CUSTOMERS = 10_000 + +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company]) + +customers = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, + randomSeed=42, randomSeedMethod="hash_fieldname") + # surrogate key 0..N-1 (the implicit contiguous seed column) — drives FK joins + .withColumn("customer_sk", "long", expr="id") + # business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # realistic text via the shared Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + # email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(replace(name, ' ', '.')), '@example.com')") + # skewed categories — NEVER uniform (weights are relative) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=[40, 25, 20, 15], random=True) + # right-skewed ARR correlated to tier: log-normal = exp() of a standard normal. + # The hidden helper column (_z, omit=True) is computed and reusable as a base column. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") + .withColumn("created_at", "date", + data_range=dg.DateRange("2023-01-01 00:00:00", "2024-12-31 00:00:00", "days=1"), + random=True) ) +customers_df = customers.build() # Write to Volume as Parquet (default for raw data) # Path is a folder with table name: /Volumes/catalog/schema/raw_data/customers/ @@ -173,7 +203,7 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") customers_df.write.mode("overwrite").parquet(f"/Volumes/{CATALOG}/{SCHEMA}/raw_data/customers") ``` -**Partitions by scale:** `spark.range(N, numPartitions=P)` +**Partitions by scale:** `dg.DataGenerator(..., rows=N, partitions=P)` - <100K rows: 8 partitions - 100K-500K: 16 partitions - 500K-1M: 32 partitions @@ -190,32 +220,79 @@ Generated scripts must be highly performant. **Never** do these: | Anti-Pattern | Why It's Slow | Do This Instead | |--------------|---------------|-----------------| -| Python loops on driver | Single-threaded, no parallelism | Use `spark.range()` + Spark operations | +| Python loops on driver | Single-threaded, no parallelism | Declare columns in a `dg.DataGenerator` spec and `.build()` | | `.collect()` then iterate | Brings all data to driver memory | Keep data in Spark, use DataFrame ops | -| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark, use `pandas_udf` only for UDFs | +| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark; let dbldatagen generate columns | | Read/write temp files | Unnecessary I/O | Chain DataFrame transformations | -| Scalar UDFs | Row-by-row processing | Use `pandas_udf` for batch processing | +| Scalar UDFs | Row-by-row processing | Use dbldatagen `expr`/`template`/`text`; `pandas_udf` only when unavoidable | -**Good pattern:** `spark.range()` → Spark transforms → `pandas_udf` for Faker → write directly +**Good pattern:** `dg.DataGenerator(rows, partitions)` → declarative `.withColumn(...)` specs → `.build()` → write directly ## Common Patterns +All snippets use the public `dbldatagen` API (`import dbldatagen as dg`, `import dbldatagen.distributions as dist`). + ### Weighted Categories (never uniform) +`weights` are relative frequencies — they don't need to sum to 100: ```python -F.when(F.rand() < 0.6, "Free").when(F.rand() < 0.9, "Pro").otherwise("Enterprise") +.withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) ``` -### Log-Normal Amounts (in a pandas UDF) -Use `np.random.lognormal(mean, sigma)` — always positive, long tail: -- Enterprise: `lognormal(7.5, 0.8)` → ~$1800 median -- Pro: `lognormal(5.5, 0.7)` → ~$245 median -- Free: `lognormal(4.0, 0.6)` → ~$55 median +### Skewed / Long-Tailed Amounts +Apply a continuous distribution to a numeric range — `Gamma`/`Exponential` give the always-positive long tail you'd want from log-normal: +```python +.withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) +``` +For a true log-normal, generate over a normal and exponentiate via `expr`: +```python +.withColumn("amount", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # standard normal, hidden +.withColumn("arr", "double", baseColumn="amount", + expr="round(exp(5.5 + 0.8 * amount), 2)") # median ~$245 +``` -### Date Range (Last 6 Months) +### Coherent Rows (correlated attributes via `expr` + `baseColumn`) +Derive dependent columns from earlier ones so each row makes business sense — no UDF needed: ```python -END_DATE = datetime.now() +.withColumn("priority", "string", values=["Critical", "High", "Medium", "Low"], + weights=[5, 15, 50, 30], random=True) +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*8 " + "WHEN 'High' THEN rand()*24 WHEN 'Medium' THEN rand()*72 " + "ELSE rand()*120 END, 1)") +.withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 4 " + "WHEN resolution_hours<72 THEN 3 ELSE 2 END") +``` + +### Date / Timestamp Range (Last 6 Months) +Use `dg.DateRange(begin, end, interval)` with the `data_range` option: +```python +from datetime import datetime, timedelta +END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +FMT = "%Y-%m-%d %H:%M:%S" + +.withColumn("order_ts", "timestamp", + data_range=dg.DateRange(START_DATE.strftime(FMT), END_DATE.strftime(FMT), "minutes=30"), + random=True) +``` + +### Realistic Text (Faker provider plugin) +Build one `FakerTextFactory` with a default locale and the faker providers you need, then reuse it: +```python +from faker.providers import person, company, internet +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, internet]) + +.withColumn("name", "string", text=FakerText("name")) # uses the 'person' provider +.withColumn("company", "string", text=FakerText("company")) # uses the 'company' provider +.withColumn("ip", "string", text=FakerText("ipv4_private")) # uses the 'internet' provider ``` +dbldatagen alternatives that need no library: +* `template=r"\w.\w@\w.com"` (templated text) +* `text=dg.ILText(paragraphs=(1, 3), sentences=(2, 5))` (lorem-ipsum free text). ### Infrastructure (always create in script) ```python @@ -224,27 +301,33 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") ``` ### Referential Integrity (FK pattern) -Write master table to Delta first, then read back for FK joins (no `.cache()` on serverless): +Generate the FK as an integer over the parent's surrogate-key range so **every value is valid by construction**, then write the parent to Delta and read it back to attach coherent parent attributes (no `.cache()` on serverless): ```python -# 1. Write master table +# 1. Parent table: surrogate key 0..N-1, then write to Delta customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") -# 2. Read back for FK lookup -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_idx", "customer_id") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +# 2. Child table: FK drawn from the parent key range (valid by construction). +# A distribution skews it 80/20 (a few customers place most orders). +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("order_id", "string", + expr="concat('ORD-', lpad(cast(id as string), 8, '0'))", baseColumn="id") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") + +# 3. Join to the parent (read back from Delta) to pull coherent parent attributes +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "customer_id", "tier") +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") ``` ## Setup -Requires Python 3.12 and databricks-connect>=16.4. Use `uv`: +Requires Python 3.12 and databricks-connect>=16.4. Install dependencies locally with `uv`: ```bash -uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays +uv pip install "databricks-connect>=16.4,<17.4" dbldatagen faker ``` ## Related Skills @@ -256,10 +339,11 @@ uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays | Issue | Solution | |-------|----------| -| `ModuleNotFoundError: faker` | Install locally: `uv pip install faker`, import inside UDF | -| Faker UDF is slow | Use `pandas_udf` for batch processing | -| Out of memory | Increase `numPartitions` in `spark.range()` | -| Referential integrity errors | Write master table to Delta first, read back for FK joins | +| `ModuleNotFoundError: dbldatagen` (or `faker`) | Install locally: `uv pip install dbldatagen faker` | +| `FakerText`/provider not found | Pass the provider to `dg.FakerTextFactory(providers=[...])` and import it from `faker.providers` | +| All rows identical / not random | Set `random=True` on the column (default is deterministic), or set `randomSeed`/`randomSeedMethod` on the generator | +| Out of memory | Increase `partitions` in `dg.DataGenerator(..., partitions=P)` | +| Referential integrity errors | Draw the FK from the parent key range (`minValue=0, maxValue=N-1`); write parent to Delta first, read back to join attributes | | `PERSIST TABLE is not supported on serverless` | **NEVER use `.cache()` or `.persist()` with serverless** - write to Delta table first, then read back | | `F.window` vs `Window` confusion | Use `from pyspark.sql.window import Window` for `row_number()`, `rank()`, etc. `F.window` is for streaming only. | | Broadcast variables not supported | **NEVER use `spark.sparkContext.broadcast()` with serverless** | diff --git a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/1-data-patterns.md b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/1-data-patterns.md index eba64916..87774896 100644 --- a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/1-data-patterns.md +++ b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/1-data-patterns.md @@ -28,15 +28,31 @@ Real data is never uniformly distributed. Use appropriate distributions: | **Exponential** | Time between events | Support resolution time, session duration | | **Weighted categorical** | Skewed categories | Status (70% complete, 5% failed), tiers | -```python -# Log-normal for amounts (long tail, always positive) -amount = np.random.lognormal(mean=5.5, sigma=0.8) # ~$245 median +dbldatagen applies a distribution to a column's random draw via the `distribution=` option +(`import dbldatagen.distributions as dist`). Supported: `Normal`, `Gamma`, `Beta`, `Exponential`. -# Pareto for power-law (few large, many small) -value = (np.random.pareto(a=1.5) + 1) * base_value +```python +import dbldatagen as dg +import dbldatagen.distributions as dist + +spec = ( + dg.DataGenerator(spark, name="amounts", rows=100_000, partitions=8, randomSeed=42) + # Use Gamma for right-tailed distributions like amounts + .withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) + # Use Exponential for waiting times + .withColumn("resolution_hours", "double", minValue=0, maxValue=240, + random=True, distribution=dist.Exponential(rate=1.0 / 24)) +) -# Exponential for time-to-event -hours = np.random.exponential(scale=24) # avg 24h, skewed right +# Exponentiate a standard normal via expr for true log-normal distributions +spec = ( + dg.DataGenerator(spark, name="ln", rows=100_000, partitions=8, randomSeed=42) + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # hidden helper column + .withColumn("amount", "double", baseColumn="_z", + expr="round(exp(5.5 + 0.8 * _z), 2)") # ~$245 median +) ``` ### 3. Row Coherence @@ -51,20 +67,24 @@ Attributes within a row must make business sense together. Generate correlated a | Large transaction + unusual hour | Higher fraud probability | | Fast resolution | Higher CSAT score | +In dbldatagen, chain `baseColumn` + `expr` when columns derive from other columns +in the same dataset. The generator resolves the dependency order automatically: + ```python -@F.pandas_udf("struct") -def generate_coherent_ticket(tiers: pd.Series) -> pd.DataFrame: - """All attributes correlate logically within each row.""" - results = [] - for tier in tiers: - # Priority depends on tier - priority = "Critical" if tier == "Enterprise" and random() < 0.3 else "Medium" - # Resolution depends on priority - resolution = np.random.exponential(4 if priority == "Critical" else 36) - # CSAT depends on resolution - csat = 5 if resolution < 4 else (3 if resolution < 24 else 2) - results.append({"priority": priority, "resolution_hours": resolution, "csat": csat}) - return pd.DataFrame(results) +spec = ( + dg.DataGenerator(spark, name="tickets", rows=80_000, partitions=8, randomSeed=42) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + # Priority depends on tier + .withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' AND rand()<0.3 THEN 'Critical' ELSE 'Medium' END") + # Resolution depends on priority + .withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 ELSE rand()*36 END, 1)") + # CSAT depends on resolution + .withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 3 ELSE 2 END") +) ``` ### 4. The 80/20 Rule @@ -75,7 +95,13 @@ Apply power-law distributions where appropriate: - **20% of products** account for 80% of sales - **20% of support agents** handle 80% of tickets -Implementation: Use weighted sampling when assigning FKs, not uniform random. +Implementation: Skew FKs by drawing from a non-uniform distribution. + +```python +# Use a Gamma distribution to skew FK values +.withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) +``` ### 5. Time-Based Patterns @@ -86,14 +112,22 @@ Most data has temporal patterns: - **Seasonality** — Q4 retail spike, summer travel peak - **Trends** — Growth over time, degradation curves +Bias *when* events happen by weighting time buckets, then derive the timestamp from the bucket. +The example below clusters tickets into business hours and skews volume toward weekdays: + ```python -def get_volume_multiplier(date): - multiplier = 1.0 - if date.weekday() >= 5: multiplier *= 0.6 # Weekend drop - if date.month in [11, 12]: multiplier *= 1.5 # Holiday spike - return multiplier +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], # 9am–5pm peak + random=True, omit=True) +.withColumn("day_offset", "int", minValue=0, maxValue=179, random=True, + distribution=dist.Normal(60, 30), omit=True) # volume ramps over the window +.withColumn("event_ts", "timestamp", baseColumn=["day_offset", "hour"], + expr="date_add(timestamp'2025-01-01 00:00:00', day_offset) + make_interval(0,0,0,0,hour,0,0)") ``` +For holiday/seasonal spikes, weight a month or week-of-year bucket the same way, or post-filter +the built DataFrame with a Spark `expr` multiplier. + ### 6. ML-Ready Data If data will train ML models, ensure: @@ -105,20 +139,32 @@ If data will train ML models, ensure: ## Referential Integrity -Generate master tables first, write to Delta, then join for FKs: +Give the parent a surrogate key over `0..N-1`, draw the child's FK from that same range (valid by +construction), then write the parent to Delta and read it back to attach parent attributes: ```python -# 1. Generate and write master table +# 1. # 1. Generate and write parent table with surrogate keys from 0..N-1 +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_id", "string", baseColumn="id", # References the built-in sequential 'id' column + expr="concat('CUST-', lpad(cast(id as string), 5, '0'))") + # ... other attributes ... + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK joins (NOT cache - unsupported on serverless) -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") + +# 3. Generate child table with FKs drawn from the parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True, omit=True) + .withColumn("customer_id", "string", baseColumn="customer_sk", # References the 'customer_sk' column + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") +orders_with_fk = orders_df.join(customer_lookup, on="customer_id", how="inner") ``` ## Data Volume diff --git a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md index 793b64f7..8a19ebc8 100644 --- a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md +++ b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md @@ -4,23 +4,24 @@ Common issues and solutions for synthetic data generation. ## Environment Issues -### ModuleNotFoundError: faker (or other library) +### ModuleNotFoundError: dbldatagen (or faker) -**Problem:** Dependencies not available in execution environment. +**Problem:** Dependencies not available in execution environment. `dbldatagen` is required, and +`faker` too if you use `FakerTextFactory`/`fakerText`. **Solutions by execution mode:** | Mode | Solution | |------|----------| -| **DB Connect with Serverless** | Install libs locally (`uv pip install faker`), use `DatabricksSession.builder.serverless(True)` | -| **Databricks Runtime** | Use Databricks CLI to install `faker holidays` | -| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "faker"}}, {"pypi": {"package": "holidays"}}]}'` | +| **DB Connect with Serverless** | Install libs locally (`uv pip install dbldatagen faker`), use `DatabricksSession.builder.serverless(True)` | +| **Databricks Runtime** | `%pip install dbldatagen faker` at the top of the notebook | +| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "dbldatagen"}}, {"pypi": {"package": "faker"}}]}'` | ```python # For DB Connect with serverless from databricks.connect import DatabricksSession -# Install dependencies locally first: uv pip install faker pandas numpy holidays +# Install dependencies locally first: uv pip install dbldatagen faker spark = DatabricksSession.builder.serverless(True).getOrCreate() ``` @@ -50,21 +51,21 @@ AnalysisException: [NOT_SUPPORTED_WITH_SERVERLESS] PERSIST TABLE is not supporte **Why this happens:** Serverless compute does not support caching DataFrames in memory. This is a fundamental limitation of the serverless architecture. -**Solution:** Write master tables to Delta first, then read them back for FK joins: +**Solution:** Write parent tables to Delta first, then read them back for FK joins: ```python # BAD - will fail on serverless -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.cache() # ❌ FAILS: "PERSIST TABLE is not supported on serverless compute" # GOOD - write to Delta, then read back -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") # ✓ Read from Delta ``` **Best practice for referential integrity:** -1. Generate master table (e.g., customers) +1. Generate parent table (e.g., customers) 2. Write to Delta table 3. Read back for FK lookup joins 4. Generate child tables (e.g., orders, tickets) with valid FKs @@ -114,7 +115,7 @@ spark = DatabricksSession.builder.serverless(True).getOrCreate() "environment_key": "datagen_env", "spec": { "client": "4", # Required! - "dependencies": ["faker", "numpy", "pandas"] + "dependencies": ["dbldatagen", "faker"] } }] } @@ -153,34 +154,37 @@ contacts_df = contacts_df.withColumn( --- -### Faker UDF is slow +### Generation is slow -**Problem:** Single-row UDFs don't parallelize well. +**Problem:** Row-by-row Python UDFs (or driver loops) don't parallelize well. -**Solution:** Use `pandas_udf` for batch processing: +**Solution:** Let dbldatagen generate columns declaratively — it builds the whole DataFrame in +parallel across `partitions`. Prefer `expr`, `template`, `values`/`weights`, and `distribution` +over UDFs. For names/addresses use the Faker plugin (`FakerTextFactory`) instead of hand-rolled UDFs: ```python -# SLOW - scalar UDF -@F.udf(returnType=StringType()) -def slow_fake_name(): - return Faker().name() - -# FAST - pandas UDF (batch processing) -@F.pandas_udf(StringType()) -def fast_fake_name(ids: pd.Series) -> pd.Series: - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) +import dbldatagen as dg +from faker.providers import person + +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person]) + +spec = ( + dg.DataGenerator(spark, name="people", rows=1_000_000, partitions=64, randomSeed=42) + .withColumn("name", "string", text=FakerText("name")) # batched by dbldatagen + .withColumn("status", "string", values=["active", "churned"], weights=[85, 15], random=True) +) +df = spec.build() ``` ### Out of memory with large data **Problem:** Not enough partitions for data size. -**Solution:** Increase partitions: +**Solution:** Increase `partitions` on the generator: ```python # For large datasets (1M+ rows) -customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from default +spec = dg.DataGenerator(spark, name="big", rows=N_CUSTOMERS, partitions=64, randomSeed=42) ``` | Data Size | Recommended Partitions | @@ -205,21 +209,28 @@ customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from de **Problem:** Foreign keys reference non-existent parent records. -**Solution:** Write master table to Delta first, then read back for FK joins: +**Solution:** Sample the FK from the parent's surrogate-key range, write the parent +to Delta, then read it back for joins: ```python -# 1. Generate and WRITE master table (do NOT use cache with serverless!) -customers_df = spark.range(0, N_CUSTOMERS)... +# 1. Generate and WRITE parent table (do NOT use cache with serverless!) +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_sk", "long", expr="id") # References the built-in sequential 'id' column + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], weights=[60, 30, 10], random=True) + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK lookups -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") - -# 3. Generate child table with valid FKs -orders_df = spark.range(0, N_ORDERS).join( - customer_lookup, - on=, - how="left" +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "tier") + +# 3. Generate child table with FK in the valid parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True) + .build() + .join(customer_lookup, on="customer_sk", how="inner") ) ``` @@ -233,52 +244,52 @@ orders_df = spark.range(0, N_ORDERS).join( **Problem:** All customers have similar order counts, amounts are evenly distributed. -**Solution:** Use non-linear distributions: +**Solution:** Apply a `distribution=` to the column instead of leaving it uniform: ```python -# BAD - uniform -amounts = np.random.uniform(10, 1000, N) +import dbldatagen.distributions as dist -# GOOD - log-normal (realistic) -amounts = np.random.lognormal(mean=5, sigma=0.8, N) +# BAD - uniform (no distribution) +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True) + +# GOOD - skewed, realistic +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True, + distribution=dist.Gamma(1.0, 2.0)) ``` ### Missing time-based patterns **Problem:** Data doesn't reflect weekday/weekend or seasonal patterns. -**Solution:** Add multipliers: +**Solution:** Weight time buckets so volume clusters realistically, then derive the timestamp: ```python -import holidays - -US_HOLIDAYS = holidays.US(years=[2024, 2025]) - -def get_multiplier(date): - mult = 1.0 - if date.weekday() >= 5: # Weekend - mult *= 0.6 - if date in US_HOLIDAYS: - mult *= 0.3 - return mult +# Cluster events into business hours +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], + random=True, omit=True) +.withColumn("event_ts", "timestamp", baseColumn="hour", + expr="date'2025-06-01' + make_interval(0,0,0,0,hour,0,0)") ``` +For weekend/holiday dips, weight a day-of-week or day bucket the same way, or post-filter the built +DataFrame with a Spark `expr` (e.g. `dayofweek(event_ts) IN (1,7)`). + ### Incoherent row attributes **Problem:** Enterprise customer has low-value orders, critical ticket has slow resolution. -**Solution:** Correlate attributes: +**Solution:** Correlate attributes with `baseColumn` + `expr` so each derives from the previous: ```python # Priority based on tier -if tier == 'Enterprise': - priority = np.random.choice(['Critical', 'High'], p=[0.4, 0.6]) -else: - priority = np.random.choice(['Medium', 'Low'], p=[0.6, 0.4]) - -# Resolution based on priority -resolution_scale = {'Critical': 4, 'High': 12, 'Medium': 36, 'Low': 72} -resolution_hours = np.random.exponential(scale=resolution_scale[priority]) +.withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' THEN (CASE WHEN rand()<0.4 THEN 'Critical' ELSE 'High' END) " + "ELSE (CASE WHEN rand()<0.6 THEN 'Medium' ELSE 'Low' END) END") +# Resolution scaled by priority +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 WHEN 'High' THEN rand()*12 " + "WHEN 'Medium' THEN rand()*36 ELSE rand()*72 END, 1)") ``` --- diff --git a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py index b36edb8e..185e4a18 100644 --- a/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py +++ b/plugins/databricks/copilot/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py @@ -1,22 +1,24 @@ -"""Generate synthetic data using Spark + Faker + Pandas UDFs. +"""Generate synthetic data using Spark + dbldatagen (Databricks Labs Data Generator). This is the recommended approach for ALL data generation tasks: - Scales from thousands to millions of rows -- Parallel execution via Spark +- Declarative, parallel generation (no driver loops, no row-by-row UDFs) - Direct write to Unity Catalog - Works with serverless and classic compute +Uses ONLY the public dbldatagen API: +- `dbldatagen.DataGenerator` / `.withColumn(...)` / `.build()` +- `dbldatagen.distributions` for skew +- `dbldatagen.FakerTextFactory` for realistic names/companies/addresses + Prerequisites: -- Install dependencies locally: uv pip install faker pandas numpy holidays databricks-connect +- Install dependencies locally: uv pip install dbldatagen faker databricks-connect - Configure ~/.databrickscfg with serverless_compute_id = auto """ import sys import os from pyspark.sql import functions as F from pyspark.sql.window import Window -from pyspark.sql.types import StringType, DoubleType, StructType, StructField, IntegerType -import numpy as np -import pandas as pd from datetime import datetime, timedelta # ============================================================================= @@ -39,6 +41,7 @@ # Date range - last 6 months from today END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +DATE_FMT = "%Y-%m-%d %H:%M:%S" # Write mode - "overwrite" for one-time, "append" for incremental WRITE_MODE = "overwrite" @@ -51,14 +54,18 @@ "orphan_fk_rate": 0.01, # 1% orphan foreign keys } -# Reproducibility +# Reproducibility - same seed => same data SEED = 42 -# Tier distribution: Free 60%, Pro 30%, Enterprise 10% -TIER_PROBS = [0.6, 0.3, 0.1] +# Tier distribution: Free 60%, Pro 30%, Enterprise 10% (relative weights) +TIER_WEIGHTS = [60, 30, 10] + +# Region distribution (relative weights) +REGION_WEIGHTS = [40, 25, 20, 15] -# Region distribution -REGION_PROBS = [0.4, 0.25, 0.2, 0.15] +# Order status distribution (relative weights) +STATUS_VALUES = ["delivered", "shipped", "processing", "pending", "cancelled"] +STATUS_WEIGHTS = [65, 15, 10, 5, 5] # ============================================================================= # SESSION CREATION @@ -79,59 +86,14 @@ spark = DatabricksSession.builder.clusterId(CLUSTER_ID).getOrCreate() print(f"Connected to cluster {CLUSTER_ID}") -# Import Faker for UDF definitions -from faker import Faker - -# ============================================================================= -# DEFINE PANDAS UDFs FOR FAKER DATA -# ============================================================================= +# Import the public dbldatagen API (installed locally) +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import person, company, address as faker_address -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - """Generate realistic person names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.name() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_company(ids: pd.Series) -> pd.Series: - """Generate realistic company names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.company() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_address(ids: pd.Series) -> pd.Series: - """Generate realistic addresses.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.address().replace('\n', ', ') for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_email(names: pd.Series) -> pd.Series: - """Generate email based on name.""" - emails = [] - for name in names: - if name: - domain = name.lower().replace(" ", ".").replace(",", "")[:20] - emails.append(f"{domain}@example.com") - else: - emails.append("unknown@example.com") - return pd.Series(emails) - -@F.pandas_udf(DoubleType()) -def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: - """Generate amount based on tier using log-normal distribution.""" - np.random.seed(SEED) - amounts = [] - for tier in tiers: - if tier == "Enterprise": - amounts.append(float(np.random.lognormal(mean=7.5, sigma=0.8))) # ~$1800 avg - elif tier == "Pro": - amounts.append(float(np.random.lognormal(mean=5.5, sigma=0.7))) # ~$245 avg - else: - amounts.append(float(np.random.lognormal(mean=4.0, sigma=0.6))) # ~$55 avg - return pd.Series(amounts) +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, faker_address]) # ============================================================================= # CREATE INFRASTRUCTURE @@ -142,39 +104,50 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f"Infrastructure ready: {VOLUME_PATH}") # ============================================================================= -# GENERATE CUSTOMERS (Master Table) +# GENERATE CUSTOMERS (Parent Table) # ============================================================================= print(f"\nGenerating {N_CUSTOMERS:,} customers...") -customers_df = ( - spark.range(0, N_CUSTOMERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), - fake_company(F.col("id")).alias("company"), - fake_address(F.col("id")).alias("address"), - # Tier distribution: Free 60%, Pro 30%, Enterprise 10% - F.when(F.rand(SEED) < TIER_PROBS[0], "Free") - .when(F.rand(SEED) < TIER_PROBS[0] + TIER_PROBS[1], "Pro") - .otherwise("Enterprise").alias("tier"), - # Region distribution - F.when(F.rand(SEED) < REGION_PROBS[0], "North") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1], "South") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1] + REGION_PROBS[2], "East") - .otherwise("West").alias("region"), - # Created date (within last 2 years before start date) - F.date_sub(F.lit(START_DATE.date()), (F.rand(SEED) * 730).cast("int")).alias("created_at"), - ) +customers_spec = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + # Surrogate key 0..N-1 (the implicit contiguous seed column) — stable join key + .withColumn("customer_sk", "long", expr="id") + # Business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # Realistic text via the Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + .withColumn("address", "string", text=FakerText("address")) + # Email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(regexp_replace(name, '[^A-Za-z]+', '.')), '@example.com')") + # Skewed categories (never uniform) — weights are relative frequencies + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=TIER_WEIGHTS, random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=REGION_WEIGHTS, random=True) + # Account created within 2 years before the analysis window + .withColumn("created_at", "date", + data_range=dg.DateRange( + (START_DATE - timedelta(days=730)).strftime(DATE_FMT), + START_DATE.strftime(DATE_FMT), + "days=1"), + random=True) + # Tier-based ARR via a log-normal (exp of a standard normal) — long tail, always positive. + # The omitted helper column '_z' keeps ARR reproducible under randomSeed. + # Enterprise ~ $1800, Pro ~ $245, Free ~ $55 median. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") ) -# Add tier-based ARR and email -customers_df = ( - customers_df - .withColumn("arr", F.round(generate_lognormal_amount(F.col("tier")), 2)) - .withColumn("email", fake_email(F.col("name"))) -) +customers_df = customers_spec.build() -# Save customers +# Save customers as raw Parquet customers_df.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/customers") print(f" Saved customers to {VOLUME_PATH}/customers") @@ -183,46 +156,45 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: # ============================================================================= print(f"\nGenerating {N_ORDERS:,} orders with referential integrity...") -# Write customer lookup to temp Delta table (no .cache() on serverless!) +# Write a customer lookup to a temp Delta table (no .cache() on serverless!) customers_tmp_table = f"{CATALOG}.{SCHEMA}._tmp_customers_lookup" -customers_df.select("customer_id", "tier").write.mode("overwrite").saveAsTable(customers_tmp_table) +(customers_df.select("customer_sk", "customer_id", "tier") + .write.mode("overwrite").saveAsTable(customers_tmp_table)) customer_lookup = spark.table(customers_tmp_table) -# Generate orders base -orders_df = ( - spark.range(0, N_ORDERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("ORD-"), F.lpad(F.col("id").cast("string"), 6, "0")).alias("order_id"), - # Generate customer_idx for FK join (hash-based distribution) - (F.abs(F.hash(F.col("id"), F.lit(SEED))) % N_CUSTOMERS).alias("customer_idx"), - # Order status - F.when(F.rand(SEED) < 0.65, "delivered") - .when(F.rand(SEED) < 0.80, "shipped") - .when(F.rand(SEED) < 0.90, "processing") - .when(F.rand(SEED) < 0.95, "pending") - .otherwise("cancelled").alias("status"), - # Order date within date range - F.date_add(F.lit(START_DATE.date()), (F.rand(SEED) * 180).cast("int")).alias("order_date"), - ) +# Generate orders. The FK (customer_sk) is drawn from the parent key range, so every +# value is valid by construction. A Gamma distribution skews it 80/20 (a few customers +# place most orders). +orders_spec = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + .withColumn("order_id", "string", baseColumn="id", + expr="concat('ORD-', lpad(cast(id as string), 6, '0'))") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .withColumn("status", "string", values=STATUS_VALUES, weights=STATUS_WEIGHTS, random=True) + .withColumn("order_date", "date", + data_range=dg.DateRange(START_DATE.strftime(DATE_FMT), + END_DATE.strftime(DATE_FMT), "days=1"), + random=True) ) +orders_df = orders_spec.build() -# Add customer_idx to lookup for join -customer_lookup_with_idx = customer_lookup.withColumn( - "customer_idx", - (F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) - 1).cast("int") -) - -# Join to get customer_id and tier as foreign key -orders_with_fk = ( - orders_df - .join(customer_lookup_with_idx, on="customer_idx", how="left") - .drop("customer_idx") -) +# Join to the parent (read back from Delta) to attach valid customer_id + tier +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") -# Add tier-based amount +# Tier-based amount (log-normal), coherent with the customer's tier orders_with_fk = orders_with_fk.withColumn( "amount", - F.round(generate_lognormal_amount(F.col("tier")), 2) + F.round( + F.expr( + "CASE tier " + "WHEN 'Enterprise' THEN exp(7.5 + 0.8 * randn()) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * randn()) " + "ELSE exp(4.0 + 0.6 * randn()) END" + ), + 2, + ), ) # ============================================================================= @@ -236,19 +208,19 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: outlier_count = int(N_ORDERS * BAD_DATA_CONFIG["outlier_rate"]) orphan_count = int(N_ORDERS * BAD_DATA_CONFIG["orphan_fk_rate"]) - # Add bad data flags + # Add a row number to target specific rows (no cache/persist on serverless) orders_with_fk = orders_with_fk.withColumn( "row_num", F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) ) - # Inject nulls in customer_id for first null_count rows + # Inject nulls in customer_id for the first null_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when(F.col("row_num") <= null_count, None).otherwise(F.col("customer_id")) ) - # Inject negative amounts for next outlier_count rows + # Inject negative amounts for the next outlier_count rows orders_with_fk = orders_with_fk.withColumn( "amount", F.when( @@ -257,7 +229,7 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: ).otherwise(F.col("amount")) ) - # Inject orphan FKs for next orphan_count rows + # Inject orphan FKs for the next orphan_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when( @@ -273,8 +245,8 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f" Injected {outlier_count} negative amounts") print(f" Injected {orphan_count} orphan foreign keys") -# Drop tier column (not needed in final output) -orders_final = orders_with_fk.drop("tier") +# Drop join-only columns (not needed in final output) +orders_final = orders_with_fk.drop("tier", "customer_sk") # Save orders orders_final.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/orders") diff --git a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/SKILL.md b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/SKILL.md index f534f60f..fe2b70c6 100644 --- a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/SKILL.md +++ b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/SKILL.md @@ -1,6 +1,6 @@ --- name: databricks-synthetic-data-gen -description: "Generate realistic synthetic data using Spark + Faker (strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'Faker', or 'sample data'." +description: "Generate realistic synthetic data using Spark + dbldatagen (Databricks Labs Data Generator, strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'dbldatagen', or 'sample data'." compatibility: Requires databricks CLI (>= v1.0.0) metadata: version: "0.1.0" @@ -11,7 +11,9 @@ parent: databricks-core # Databricks Synthetic Data Generation -Generate realistic, story-driven synthetic data for Databricks using **Spark + Faker + Pandas UDFs** (strongly recommended). +Generate realistic, story-driven synthetic data for Databricks using **Spark + [dbldatagen](https://databrickslabs.github.io/dbldatagen/public_docs/index.html)** (the Databricks Labs Data Generator — strongly recommended). + +> **Use the public `dbldatagen` API.** Import the top-level package (`import dbldatagen as dg`) and its public submodules (`dbldatagen.distributions`, `dbldatagen.constraints`). Everything below is built from the public API. ## Data Must Tell a Business Story @@ -44,11 +46,11 @@ Synthetic data should demonstrate how Databricks helps solve real business probl 5. **Enough data for trends** — ~100K+ rows for main tables so patterns survive aggregation 6. **Ask for catalog/schema** — Never default, always confirm before generating 7. **Present plan for approval** — Show tables, distributions, assumptions before writing code -8. **Master tables first** — Generate parent tables, write to Delta, then create children with valid FKs -9. **Use Spark + Faker + Pandas UDFs** — Scalable, parallel. Polars only if user explicitly wants local + <30K rows +8. **Parent tables first** — Generate parent tables, write to Delta, then create children with valid FKs +9. **Use Spark + dbldatagen** — Scalable, parallel, declarative. Build a `dg.DataGenerator` spec and `.build()` it into a Spark DataFrame. Use a `pandas_udf` only for logic dbldatagen can't express 10. **Use Databricks Connect Serverless by default to generate data** — Update databricks-connect on python 3.12 if required (avoid using execute_code unless instructed to not use Databricks Connect) 11. **No `.cache()` or `.persist()`** — Not supported on serverless. Write to Delta, read back for joins -12. **No Python loops or `.collect()`** — Use Spark parallelism. No driver-side iteration, avoid Pandas↔Spark conversions +12. **No Python loops or `.collect()`** — Use Spark parallelism and dbldatagen specs. No driver-side iteration, avoid Pandas↔Spark conversions ## Generation Planning Workflow @@ -142,29 +144,57 @@ SELECT COUNT(*) FROM parquet.\`/Volumes/CATALOG/SCHEMA/raw_data/customers\` See [references/2-troubleshooting.md](references/2-troubleshooting.md) for full validation examples. -## Use Databricks Connect Spark + Faker Pattern +## Use Databricks Connect + dbldatagen Pattern + +A `DataGenerator` spec declares each column; `.build()` returns a Spark DataFrame. No UDFs, no driver loops — generation is fully distributed across `partitions`. Install `dbldatagen` and `faker` locally first (see Setup). ```python from databricks.connect import DatabricksSession -from pyspark.sql import functions as F -from pyspark.sql.types import StringType -import pandas as pd +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import company, person # providers we draw from -# Setup serverless Spark session +# Setup serverless Spark session (deps installed locally) spark = DatabricksSession.builder.serverless(True).getOrCreate() -# Pandas UDF pattern - import lib INSIDE the function (libs must be installed locally) -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - from faker import Faker # Import inside UDF - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) - -# Generate with spark.range, apply UDFs -customers_df = spark.range(0, 10000, numPartitions=16).select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), +CATALOG, SCHEMA = "", "" # always user-supplied +N_CUSTOMERS = 10_000 + +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company]) + +customers = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, + randomSeed=42, randomSeedMethod="hash_fieldname") + # surrogate key 0..N-1 (the implicit contiguous seed column) — drives FK joins + .withColumn("customer_sk", "long", expr="id") + # business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # realistic text via the shared Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + # email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(replace(name, ' ', '.')), '@example.com')") + # skewed categories — NEVER uniform (weights are relative) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=[40, 25, 20, 15], random=True) + # right-skewed ARR correlated to tier: log-normal = exp() of a standard normal. + # The hidden helper column (_z, omit=True) is computed and reusable as a base column. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") + .withColumn("created_at", "date", + data_range=dg.DateRange("2023-01-01 00:00:00", "2024-12-31 00:00:00", "days=1"), + random=True) ) +customers_df = customers.build() # Write to Volume as Parquet (default for raw data) # Path is a folder with table name: /Volumes/catalog/schema/raw_data/customers/ @@ -173,7 +203,7 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") customers_df.write.mode("overwrite").parquet(f"/Volumes/{CATALOG}/{SCHEMA}/raw_data/customers") ``` -**Partitions by scale:** `spark.range(N, numPartitions=P)` +**Partitions by scale:** `dg.DataGenerator(..., rows=N, partitions=P)` - <100K rows: 8 partitions - 100K-500K: 16 partitions - 500K-1M: 32 partitions @@ -190,32 +220,79 @@ Generated scripts must be highly performant. **Never** do these: | Anti-Pattern | Why It's Slow | Do This Instead | |--------------|---------------|-----------------| -| Python loops on driver | Single-threaded, no parallelism | Use `spark.range()` + Spark operations | +| Python loops on driver | Single-threaded, no parallelism | Declare columns in a `dg.DataGenerator` spec and `.build()` | | `.collect()` then iterate | Brings all data to driver memory | Keep data in Spark, use DataFrame ops | -| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark, use `pandas_udf` only for UDFs | +| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark; let dbldatagen generate columns | | Read/write temp files | Unnecessary I/O | Chain DataFrame transformations | -| Scalar UDFs | Row-by-row processing | Use `pandas_udf` for batch processing | +| Scalar UDFs | Row-by-row processing | Use dbldatagen `expr`/`template`/`text`; `pandas_udf` only when unavoidable | -**Good pattern:** `spark.range()` → Spark transforms → `pandas_udf` for Faker → write directly +**Good pattern:** `dg.DataGenerator(rows, partitions)` → declarative `.withColumn(...)` specs → `.build()` → write directly ## Common Patterns +All snippets use the public `dbldatagen` API (`import dbldatagen as dg`, `import dbldatagen.distributions as dist`). + ### Weighted Categories (never uniform) +`weights` are relative frequencies — they don't need to sum to 100: ```python -F.when(F.rand() < 0.6, "Free").when(F.rand() < 0.9, "Pro").otherwise("Enterprise") +.withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) ``` -### Log-Normal Amounts (in a pandas UDF) -Use `np.random.lognormal(mean, sigma)` — always positive, long tail: -- Enterprise: `lognormal(7.5, 0.8)` → ~$1800 median -- Pro: `lognormal(5.5, 0.7)` → ~$245 median -- Free: `lognormal(4.0, 0.6)` → ~$55 median +### Skewed / Long-Tailed Amounts +Apply a continuous distribution to a numeric range — `Gamma`/`Exponential` give the always-positive long tail you'd want from log-normal: +```python +.withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) +``` +For a true log-normal, generate over a normal and exponentiate via `expr`: +```python +.withColumn("amount", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # standard normal, hidden +.withColumn("arr", "double", baseColumn="amount", + expr="round(exp(5.5 + 0.8 * amount), 2)") # median ~$245 +``` -### Date Range (Last 6 Months) +### Coherent Rows (correlated attributes via `expr` + `baseColumn`) +Derive dependent columns from earlier ones so each row makes business sense — no UDF needed: ```python -END_DATE = datetime.now() +.withColumn("priority", "string", values=["Critical", "High", "Medium", "Low"], + weights=[5, 15, 50, 30], random=True) +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*8 " + "WHEN 'High' THEN rand()*24 WHEN 'Medium' THEN rand()*72 " + "ELSE rand()*120 END, 1)") +.withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 4 " + "WHEN resolution_hours<72 THEN 3 ELSE 2 END") +``` + +### Date / Timestamp Range (Last 6 Months) +Use `dg.DateRange(begin, end, interval)` with the `data_range` option: +```python +from datetime import datetime, timedelta +END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +FMT = "%Y-%m-%d %H:%M:%S" + +.withColumn("order_ts", "timestamp", + data_range=dg.DateRange(START_DATE.strftime(FMT), END_DATE.strftime(FMT), "minutes=30"), + random=True) +``` + +### Realistic Text (Faker provider plugin) +Build one `FakerTextFactory` with a default locale and the faker providers you need, then reuse it: +```python +from faker.providers import person, company, internet +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, internet]) + +.withColumn("name", "string", text=FakerText("name")) # uses the 'person' provider +.withColumn("company", "string", text=FakerText("company")) # uses the 'company' provider +.withColumn("ip", "string", text=FakerText("ipv4_private")) # uses the 'internet' provider ``` +dbldatagen alternatives that need no library: +* `template=r"\w.\w@\w.com"` (templated text) +* `text=dg.ILText(paragraphs=(1, 3), sentences=(2, 5))` (lorem-ipsum free text). ### Infrastructure (always create in script) ```python @@ -224,27 +301,33 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") ``` ### Referential Integrity (FK pattern) -Write master table to Delta first, then read back for FK joins (no `.cache()` on serverless): +Generate the FK as an integer over the parent's surrogate-key range so **every value is valid by construction**, then write the parent to Delta and read it back to attach coherent parent attributes (no `.cache()` on serverless): ```python -# 1. Write master table +# 1. Parent table: surrogate key 0..N-1, then write to Delta customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") -# 2. Read back for FK lookup -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_idx", "customer_id") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +# 2. Child table: FK drawn from the parent key range (valid by construction). +# A distribution skews it 80/20 (a few customers place most orders). +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("order_id", "string", + expr="concat('ORD-', lpad(cast(id as string), 8, '0'))", baseColumn="id") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") + +# 3. Join to the parent (read back from Delta) to pull coherent parent attributes +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "customer_id", "tier") +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") ``` ## Setup -Requires Python 3.12 and databricks-connect>=16.4. Use `uv`: +Requires Python 3.12 and databricks-connect>=16.4. Install dependencies locally with `uv`: ```bash -uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays +uv pip install "databricks-connect>=16.4,<17.4" dbldatagen faker ``` ## Related Skills @@ -256,10 +339,11 @@ uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays | Issue | Solution | |-------|----------| -| `ModuleNotFoundError: faker` | Install locally: `uv pip install faker`, import inside UDF | -| Faker UDF is slow | Use `pandas_udf` for batch processing | -| Out of memory | Increase `numPartitions` in `spark.range()` | -| Referential integrity errors | Write master table to Delta first, read back for FK joins | +| `ModuleNotFoundError: dbldatagen` (or `faker`) | Install locally: `uv pip install dbldatagen faker` | +| `FakerText`/provider not found | Pass the provider to `dg.FakerTextFactory(providers=[...])` and import it from `faker.providers` | +| All rows identical / not random | Set `random=True` on the column (default is deterministic), or set `randomSeed`/`randomSeedMethod` on the generator | +| Out of memory | Increase `partitions` in `dg.DataGenerator(..., partitions=P)` | +| Referential integrity errors | Draw the FK from the parent key range (`minValue=0, maxValue=N-1`); write parent to Delta first, read back to join attributes | | `PERSIST TABLE is not supported on serverless` | **NEVER use `.cache()` or `.persist()` with serverless** - write to Delta table first, then read back | | `F.window` vs `Window` confusion | Use `from pyspark.sql.window import Window` for `row_number()`, `rank()`, etc. `F.window` is for streaming only. | | Broadcast variables not supported | **NEVER use `spark.sparkContext.broadcast()` with serverless** | diff --git a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/1-data-patterns.md b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/1-data-patterns.md index eba64916..87774896 100644 --- a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/1-data-patterns.md +++ b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/1-data-patterns.md @@ -28,15 +28,31 @@ Real data is never uniformly distributed. Use appropriate distributions: | **Exponential** | Time between events | Support resolution time, session duration | | **Weighted categorical** | Skewed categories | Status (70% complete, 5% failed), tiers | -```python -# Log-normal for amounts (long tail, always positive) -amount = np.random.lognormal(mean=5.5, sigma=0.8) # ~$245 median +dbldatagen applies a distribution to a column's random draw via the `distribution=` option +(`import dbldatagen.distributions as dist`). Supported: `Normal`, `Gamma`, `Beta`, `Exponential`. -# Pareto for power-law (few large, many small) -value = (np.random.pareto(a=1.5) + 1) * base_value +```python +import dbldatagen as dg +import dbldatagen.distributions as dist + +spec = ( + dg.DataGenerator(spark, name="amounts", rows=100_000, partitions=8, randomSeed=42) + # Use Gamma for right-tailed distributions like amounts + .withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) + # Use Exponential for waiting times + .withColumn("resolution_hours", "double", minValue=0, maxValue=240, + random=True, distribution=dist.Exponential(rate=1.0 / 24)) +) -# Exponential for time-to-event -hours = np.random.exponential(scale=24) # avg 24h, skewed right +# Exponentiate a standard normal via expr for true log-normal distributions +spec = ( + dg.DataGenerator(spark, name="ln", rows=100_000, partitions=8, randomSeed=42) + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # hidden helper column + .withColumn("amount", "double", baseColumn="_z", + expr="round(exp(5.5 + 0.8 * _z), 2)") # ~$245 median +) ``` ### 3. Row Coherence @@ -51,20 +67,24 @@ Attributes within a row must make business sense together. Generate correlated a | Large transaction + unusual hour | Higher fraud probability | | Fast resolution | Higher CSAT score | +In dbldatagen, chain `baseColumn` + `expr` when columns derive from other columns +in the same dataset. The generator resolves the dependency order automatically: + ```python -@F.pandas_udf("struct") -def generate_coherent_ticket(tiers: pd.Series) -> pd.DataFrame: - """All attributes correlate logically within each row.""" - results = [] - for tier in tiers: - # Priority depends on tier - priority = "Critical" if tier == "Enterprise" and random() < 0.3 else "Medium" - # Resolution depends on priority - resolution = np.random.exponential(4 if priority == "Critical" else 36) - # CSAT depends on resolution - csat = 5 if resolution < 4 else (3 if resolution < 24 else 2) - results.append({"priority": priority, "resolution_hours": resolution, "csat": csat}) - return pd.DataFrame(results) +spec = ( + dg.DataGenerator(spark, name="tickets", rows=80_000, partitions=8, randomSeed=42) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + # Priority depends on tier + .withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' AND rand()<0.3 THEN 'Critical' ELSE 'Medium' END") + # Resolution depends on priority + .withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 ELSE rand()*36 END, 1)") + # CSAT depends on resolution + .withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 3 ELSE 2 END") +) ``` ### 4. The 80/20 Rule @@ -75,7 +95,13 @@ Apply power-law distributions where appropriate: - **20% of products** account for 80% of sales - **20% of support agents** handle 80% of tickets -Implementation: Use weighted sampling when assigning FKs, not uniform random. +Implementation: Skew FKs by drawing from a non-uniform distribution. + +```python +# Use a Gamma distribution to skew FK values +.withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) +``` ### 5. Time-Based Patterns @@ -86,14 +112,22 @@ Most data has temporal patterns: - **Seasonality** — Q4 retail spike, summer travel peak - **Trends** — Growth over time, degradation curves +Bias *when* events happen by weighting time buckets, then derive the timestamp from the bucket. +The example below clusters tickets into business hours and skews volume toward weekdays: + ```python -def get_volume_multiplier(date): - multiplier = 1.0 - if date.weekday() >= 5: multiplier *= 0.6 # Weekend drop - if date.month in [11, 12]: multiplier *= 1.5 # Holiday spike - return multiplier +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], # 9am–5pm peak + random=True, omit=True) +.withColumn("day_offset", "int", minValue=0, maxValue=179, random=True, + distribution=dist.Normal(60, 30), omit=True) # volume ramps over the window +.withColumn("event_ts", "timestamp", baseColumn=["day_offset", "hour"], + expr="date_add(timestamp'2025-01-01 00:00:00', day_offset) + make_interval(0,0,0,0,hour,0,0)") ``` +For holiday/seasonal spikes, weight a month or week-of-year bucket the same way, or post-filter +the built DataFrame with a Spark `expr` multiplier. + ### 6. ML-Ready Data If data will train ML models, ensure: @@ -105,20 +139,32 @@ If data will train ML models, ensure: ## Referential Integrity -Generate master tables first, write to Delta, then join for FKs: +Give the parent a surrogate key over `0..N-1`, draw the child's FK from that same range (valid by +construction), then write the parent to Delta and read it back to attach parent attributes: ```python -# 1. Generate and write master table +# 1. # 1. Generate and write parent table with surrogate keys from 0..N-1 +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_id", "string", baseColumn="id", # References the built-in sequential 'id' column + expr="concat('CUST-', lpad(cast(id as string), 5, '0'))") + # ... other attributes ... + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK joins (NOT cache - unsupported on serverless) -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") + +# 3. Generate child table with FKs drawn from the parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True, omit=True) + .withColumn("customer_id", "string", baseColumn="customer_sk", # References the 'customer_sk' column + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") +orders_with_fk = orders_df.join(customer_lookup, on="customer_id", how="inner") ``` ## Data Volume diff --git a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md index 793b64f7..8a19ebc8 100644 --- a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md +++ b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md @@ -4,23 +4,24 @@ Common issues and solutions for synthetic data generation. ## Environment Issues -### ModuleNotFoundError: faker (or other library) +### ModuleNotFoundError: dbldatagen (or faker) -**Problem:** Dependencies not available in execution environment. +**Problem:** Dependencies not available in execution environment. `dbldatagen` is required, and +`faker` too if you use `FakerTextFactory`/`fakerText`. **Solutions by execution mode:** | Mode | Solution | |------|----------| -| **DB Connect with Serverless** | Install libs locally (`uv pip install faker`), use `DatabricksSession.builder.serverless(True)` | -| **Databricks Runtime** | Use Databricks CLI to install `faker holidays` | -| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "faker"}}, {"pypi": {"package": "holidays"}}]}'` | +| **DB Connect with Serverless** | Install libs locally (`uv pip install dbldatagen faker`), use `DatabricksSession.builder.serverless(True)` | +| **Databricks Runtime** | `%pip install dbldatagen faker` at the top of the notebook | +| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "dbldatagen"}}, {"pypi": {"package": "faker"}}]}'` | ```python # For DB Connect with serverless from databricks.connect import DatabricksSession -# Install dependencies locally first: uv pip install faker pandas numpy holidays +# Install dependencies locally first: uv pip install dbldatagen faker spark = DatabricksSession.builder.serverless(True).getOrCreate() ``` @@ -50,21 +51,21 @@ AnalysisException: [NOT_SUPPORTED_WITH_SERVERLESS] PERSIST TABLE is not supporte **Why this happens:** Serverless compute does not support caching DataFrames in memory. This is a fundamental limitation of the serverless architecture. -**Solution:** Write master tables to Delta first, then read them back for FK joins: +**Solution:** Write parent tables to Delta first, then read them back for FK joins: ```python # BAD - will fail on serverless -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.cache() # ❌ FAILS: "PERSIST TABLE is not supported on serverless compute" # GOOD - write to Delta, then read back -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") # ✓ Read from Delta ``` **Best practice for referential integrity:** -1. Generate master table (e.g., customers) +1. Generate parent table (e.g., customers) 2. Write to Delta table 3. Read back for FK lookup joins 4. Generate child tables (e.g., orders, tickets) with valid FKs @@ -114,7 +115,7 @@ spark = DatabricksSession.builder.serverless(True).getOrCreate() "environment_key": "datagen_env", "spec": { "client": "4", # Required! - "dependencies": ["faker", "numpy", "pandas"] + "dependencies": ["dbldatagen", "faker"] } }] } @@ -153,34 +154,37 @@ contacts_df = contacts_df.withColumn( --- -### Faker UDF is slow +### Generation is slow -**Problem:** Single-row UDFs don't parallelize well. +**Problem:** Row-by-row Python UDFs (or driver loops) don't parallelize well. -**Solution:** Use `pandas_udf` for batch processing: +**Solution:** Let dbldatagen generate columns declaratively — it builds the whole DataFrame in +parallel across `partitions`. Prefer `expr`, `template`, `values`/`weights`, and `distribution` +over UDFs. For names/addresses use the Faker plugin (`FakerTextFactory`) instead of hand-rolled UDFs: ```python -# SLOW - scalar UDF -@F.udf(returnType=StringType()) -def slow_fake_name(): - return Faker().name() - -# FAST - pandas UDF (batch processing) -@F.pandas_udf(StringType()) -def fast_fake_name(ids: pd.Series) -> pd.Series: - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) +import dbldatagen as dg +from faker.providers import person + +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person]) + +spec = ( + dg.DataGenerator(spark, name="people", rows=1_000_000, partitions=64, randomSeed=42) + .withColumn("name", "string", text=FakerText("name")) # batched by dbldatagen + .withColumn("status", "string", values=["active", "churned"], weights=[85, 15], random=True) +) +df = spec.build() ``` ### Out of memory with large data **Problem:** Not enough partitions for data size. -**Solution:** Increase partitions: +**Solution:** Increase `partitions` on the generator: ```python # For large datasets (1M+ rows) -customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from default +spec = dg.DataGenerator(spark, name="big", rows=N_CUSTOMERS, partitions=64, randomSeed=42) ``` | Data Size | Recommended Partitions | @@ -205,21 +209,28 @@ customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from de **Problem:** Foreign keys reference non-existent parent records. -**Solution:** Write master table to Delta first, then read back for FK joins: +**Solution:** Sample the FK from the parent's surrogate-key range, write the parent +to Delta, then read it back for joins: ```python -# 1. Generate and WRITE master table (do NOT use cache with serverless!) -customers_df = spark.range(0, N_CUSTOMERS)... +# 1. Generate and WRITE parent table (do NOT use cache with serverless!) +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_sk", "long", expr="id") # References the built-in sequential 'id' column + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], weights=[60, 30, 10], random=True) + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK lookups -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") - -# 3. Generate child table with valid FKs -orders_df = spark.range(0, N_ORDERS).join( - customer_lookup, - on=, - how="left" +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "tier") + +# 3. Generate child table with FK in the valid parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True) + .build() + .join(customer_lookup, on="customer_sk", how="inner") ) ``` @@ -233,52 +244,52 @@ orders_df = spark.range(0, N_ORDERS).join( **Problem:** All customers have similar order counts, amounts are evenly distributed. -**Solution:** Use non-linear distributions: +**Solution:** Apply a `distribution=` to the column instead of leaving it uniform: ```python -# BAD - uniform -amounts = np.random.uniform(10, 1000, N) +import dbldatagen.distributions as dist -# GOOD - log-normal (realistic) -amounts = np.random.lognormal(mean=5, sigma=0.8, N) +# BAD - uniform (no distribution) +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True) + +# GOOD - skewed, realistic +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True, + distribution=dist.Gamma(1.0, 2.0)) ``` ### Missing time-based patterns **Problem:** Data doesn't reflect weekday/weekend or seasonal patterns. -**Solution:** Add multipliers: +**Solution:** Weight time buckets so volume clusters realistically, then derive the timestamp: ```python -import holidays - -US_HOLIDAYS = holidays.US(years=[2024, 2025]) - -def get_multiplier(date): - mult = 1.0 - if date.weekday() >= 5: # Weekend - mult *= 0.6 - if date in US_HOLIDAYS: - mult *= 0.3 - return mult +# Cluster events into business hours +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], + random=True, omit=True) +.withColumn("event_ts", "timestamp", baseColumn="hour", + expr="date'2025-06-01' + make_interval(0,0,0,0,hour,0,0)") ``` +For weekend/holiday dips, weight a day-of-week or day bucket the same way, or post-filter the built +DataFrame with a Spark `expr` (e.g. `dayofweek(event_ts) IN (1,7)`). + ### Incoherent row attributes **Problem:** Enterprise customer has low-value orders, critical ticket has slow resolution. -**Solution:** Correlate attributes: +**Solution:** Correlate attributes with `baseColumn` + `expr` so each derives from the previous: ```python # Priority based on tier -if tier == 'Enterprise': - priority = np.random.choice(['Critical', 'High'], p=[0.4, 0.6]) -else: - priority = np.random.choice(['Medium', 'Low'], p=[0.6, 0.4]) - -# Resolution based on priority -resolution_scale = {'Critical': 4, 'High': 12, 'Medium': 36, 'Low': 72} -resolution_hours = np.random.exponential(scale=resolution_scale[priority]) +.withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' THEN (CASE WHEN rand()<0.4 THEN 'Critical' ELSE 'High' END) " + "ELSE (CASE WHEN rand()<0.6 THEN 'Medium' ELSE 'Low' END) END") +# Resolution scaled by priority +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 WHEN 'High' THEN rand()*12 " + "WHEN 'Medium' THEN rand()*36 ELSE rand()*72 END, 1)") ``` --- diff --git a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py index b36edb8e..185e4a18 100644 --- a/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py +++ b/plugins/databricks/cursor/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py @@ -1,22 +1,24 @@ -"""Generate synthetic data using Spark + Faker + Pandas UDFs. +"""Generate synthetic data using Spark + dbldatagen (Databricks Labs Data Generator). This is the recommended approach for ALL data generation tasks: - Scales from thousands to millions of rows -- Parallel execution via Spark +- Declarative, parallel generation (no driver loops, no row-by-row UDFs) - Direct write to Unity Catalog - Works with serverless and classic compute +Uses ONLY the public dbldatagen API: +- `dbldatagen.DataGenerator` / `.withColumn(...)` / `.build()` +- `dbldatagen.distributions` for skew +- `dbldatagen.FakerTextFactory` for realistic names/companies/addresses + Prerequisites: -- Install dependencies locally: uv pip install faker pandas numpy holidays databricks-connect +- Install dependencies locally: uv pip install dbldatagen faker databricks-connect - Configure ~/.databrickscfg with serverless_compute_id = auto """ import sys import os from pyspark.sql import functions as F from pyspark.sql.window import Window -from pyspark.sql.types import StringType, DoubleType, StructType, StructField, IntegerType -import numpy as np -import pandas as pd from datetime import datetime, timedelta # ============================================================================= @@ -39,6 +41,7 @@ # Date range - last 6 months from today END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +DATE_FMT = "%Y-%m-%d %H:%M:%S" # Write mode - "overwrite" for one-time, "append" for incremental WRITE_MODE = "overwrite" @@ -51,14 +54,18 @@ "orphan_fk_rate": 0.01, # 1% orphan foreign keys } -# Reproducibility +# Reproducibility - same seed => same data SEED = 42 -# Tier distribution: Free 60%, Pro 30%, Enterprise 10% -TIER_PROBS = [0.6, 0.3, 0.1] +# Tier distribution: Free 60%, Pro 30%, Enterprise 10% (relative weights) +TIER_WEIGHTS = [60, 30, 10] + +# Region distribution (relative weights) +REGION_WEIGHTS = [40, 25, 20, 15] -# Region distribution -REGION_PROBS = [0.4, 0.25, 0.2, 0.15] +# Order status distribution (relative weights) +STATUS_VALUES = ["delivered", "shipped", "processing", "pending", "cancelled"] +STATUS_WEIGHTS = [65, 15, 10, 5, 5] # ============================================================================= # SESSION CREATION @@ -79,59 +86,14 @@ spark = DatabricksSession.builder.clusterId(CLUSTER_ID).getOrCreate() print(f"Connected to cluster {CLUSTER_ID}") -# Import Faker for UDF definitions -from faker import Faker - -# ============================================================================= -# DEFINE PANDAS UDFs FOR FAKER DATA -# ============================================================================= +# Import the public dbldatagen API (installed locally) +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import person, company, address as faker_address -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - """Generate realistic person names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.name() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_company(ids: pd.Series) -> pd.Series: - """Generate realistic company names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.company() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_address(ids: pd.Series) -> pd.Series: - """Generate realistic addresses.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.address().replace('\n', ', ') for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_email(names: pd.Series) -> pd.Series: - """Generate email based on name.""" - emails = [] - for name in names: - if name: - domain = name.lower().replace(" ", ".").replace(",", "")[:20] - emails.append(f"{domain}@example.com") - else: - emails.append("unknown@example.com") - return pd.Series(emails) - -@F.pandas_udf(DoubleType()) -def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: - """Generate amount based on tier using log-normal distribution.""" - np.random.seed(SEED) - amounts = [] - for tier in tiers: - if tier == "Enterprise": - amounts.append(float(np.random.lognormal(mean=7.5, sigma=0.8))) # ~$1800 avg - elif tier == "Pro": - amounts.append(float(np.random.lognormal(mean=5.5, sigma=0.7))) # ~$245 avg - else: - amounts.append(float(np.random.lognormal(mean=4.0, sigma=0.6))) # ~$55 avg - return pd.Series(amounts) +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, faker_address]) # ============================================================================= # CREATE INFRASTRUCTURE @@ -142,39 +104,50 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f"Infrastructure ready: {VOLUME_PATH}") # ============================================================================= -# GENERATE CUSTOMERS (Master Table) +# GENERATE CUSTOMERS (Parent Table) # ============================================================================= print(f"\nGenerating {N_CUSTOMERS:,} customers...") -customers_df = ( - spark.range(0, N_CUSTOMERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), - fake_company(F.col("id")).alias("company"), - fake_address(F.col("id")).alias("address"), - # Tier distribution: Free 60%, Pro 30%, Enterprise 10% - F.when(F.rand(SEED) < TIER_PROBS[0], "Free") - .when(F.rand(SEED) < TIER_PROBS[0] + TIER_PROBS[1], "Pro") - .otherwise("Enterprise").alias("tier"), - # Region distribution - F.when(F.rand(SEED) < REGION_PROBS[0], "North") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1], "South") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1] + REGION_PROBS[2], "East") - .otherwise("West").alias("region"), - # Created date (within last 2 years before start date) - F.date_sub(F.lit(START_DATE.date()), (F.rand(SEED) * 730).cast("int")).alias("created_at"), - ) +customers_spec = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + # Surrogate key 0..N-1 (the implicit contiguous seed column) — stable join key + .withColumn("customer_sk", "long", expr="id") + # Business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # Realistic text via the Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + .withColumn("address", "string", text=FakerText("address")) + # Email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(regexp_replace(name, '[^A-Za-z]+', '.')), '@example.com')") + # Skewed categories (never uniform) — weights are relative frequencies + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=TIER_WEIGHTS, random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=REGION_WEIGHTS, random=True) + # Account created within 2 years before the analysis window + .withColumn("created_at", "date", + data_range=dg.DateRange( + (START_DATE - timedelta(days=730)).strftime(DATE_FMT), + START_DATE.strftime(DATE_FMT), + "days=1"), + random=True) + # Tier-based ARR via a log-normal (exp of a standard normal) — long tail, always positive. + # The omitted helper column '_z' keeps ARR reproducible under randomSeed. + # Enterprise ~ $1800, Pro ~ $245, Free ~ $55 median. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") ) -# Add tier-based ARR and email -customers_df = ( - customers_df - .withColumn("arr", F.round(generate_lognormal_amount(F.col("tier")), 2)) - .withColumn("email", fake_email(F.col("name"))) -) +customers_df = customers_spec.build() -# Save customers +# Save customers as raw Parquet customers_df.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/customers") print(f" Saved customers to {VOLUME_PATH}/customers") @@ -183,46 +156,45 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: # ============================================================================= print(f"\nGenerating {N_ORDERS:,} orders with referential integrity...") -# Write customer lookup to temp Delta table (no .cache() on serverless!) +# Write a customer lookup to a temp Delta table (no .cache() on serverless!) customers_tmp_table = f"{CATALOG}.{SCHEMA}._tmp_customers_lookup" -customers_df.select("customer_id", "tier").write.mode("overwrite").saveAsTable(customers_tmp_table) +(customers_df.select("customer_sk", "customer_id", "tier") + .write.mode("overwrite").saveAsTable(customers_tmp_table)) customer_lookup = spark.table(customers_tmp_table) -# Generate orders base -orders_df = ( - spark.range(0, N_ORDERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("ORD-"), F.lpad(F.col("id").cast("string"), 6, "0")).alias("order_id"), - # Generate customer_idx for FK join (hash-based distribution) - (F.abs(F.hash(F.col("id"), F.lit(SEED))) % N_CUSTOMERS).alias("customer_idx"), - # Order status - F.when(F.rand(SEED) < 0.65, "delivered") - .when(F.rand(SEED) < 0.80, "shipped") - .when(F.rand(SEED) < 0.90, "processing") - .when(F.rand(SEED) < 0.95, "pending") - .otherwise("cancelled").alias("status"), - # Order date within date range - F.date_add(F.lit(START_DATE.date()), (F.rand(SEED) * 180).cast("int")).alias("order_date"), - ) +# Generate orders. The FK (customer_sk) is drawn from the parent key range, so every +# value is valid by construction. A Gamma distribution skews it 80/20 (a few customers +# place most orders). +orders_spec = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + .withColumn("order_id", "string", baseColumn="id", + expr="concat('ORD-', lpad(cast(id as string), 6, '0'))") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .withColumn("status", "string", values=STATUS_VALUES, weights=STATUS_WEIGHTS, random=True) + .withColumn("order_date", "date", + data_range=dg.DateRange(START_DATE.strftime(DATE_FMT), + END_DATE.strftime(DATE_FMT), "days=1"), + random=True) ) +orders_df = orders_spec.build() -# Add customer_idx to lookup for join -customer_lookup_with_idx = customer_lookup.withColumn( - "customer_idx", - (F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) - 1).cast("int") -) - -# Join to get customer_id and tier as foreign key -orders_with_fk = ( - orders_df - .join(customer_lookup_with_idx, on="customer_idx", how="left") - .drop("customer_idx") -) +# Join to the parent (read back from Delta) to attach valid customer_id + tier +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") -# Add tier-based amount +# Tier-based amount (log-normal), coherent with the customer's tier orders_with_fk = orders_with_fk.withColumn( "amount", - F.round(generate_lognormal_amount(F.col("tier")), 2) + F.round( + F.expr( + "CASE tier " + "WHEN 'Enterprise' THEN exp(7.5 + 0.8 * randn()) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * randn()) " + "ELSE exp(4.0 + 0.6 * randn()) END" + ), + 2, + ), ) # ============================================================================= @@ -236,19 +208,19 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: outlier_count = int(N_ORDERS * BAD_DATA_CONFIG["outlier_rate"]) orphan_count = int(N_ORDERS * BAD_DATA_CONFIG["orphan_fk_rate"]) - # Add bad data flags + # Add a row number to target specific rows (no cache/persist on serverless) orders_with_fk = orders_with_fk.withColumn( "row_num", F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) ) - # Inject nulls in customer_id for first null_count rows + # Inject nulls in customer_id for the first null_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when(F.col("row_num") <= null_count, None).otherwise(F.col("customer_id")) ) - # Inject negative amounts for next outlier_count rows + # Inject negative amounts for the next outlier_count rows orders_with_fk = orders_with_fk.withColumn( "amount", F.when( @@ -257,7 +229,7 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: ).otherwise(F.col("amount")) ) - # Inject orphan FKs for next orphan_count rows + # Inject orphan FKs for the next orphan_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when( @@ -273,8 +245,8 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f" Injected {outlier_count} negative amounts") print(f" Injected {orphan_count} orphan foreign keys") -# Drop tier column (not needed in final output) -orders_final = orders_with_fk.drop("tier") +# Drop join-only columns (not needed in final output) +orders_final = orders_with_fk.drop("tier", "customer_sk") # Save orders orders_final.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/orders") diff --git a/skills/databricks-synthetic-data-gen/SKILL.md b/skills/databricks-synthetic-data-gen/SKILL.md index f534f60f..fe2b70c6 100755 --- a/skills/databricks-synthetic-data-gen/SKILL.md +++ b/skills/databricks-synthetic-data-gen/SKILL.md @@ -1,6 +1,6 @@ --- name: databricks-synthetic-data-gen -description: "Generate realistic synthetic data using Spark + Faker (strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'Faker', or 'sample data'." +description: "Generate realistic synthetic data using Spark + dbldatagen (Databricks Labs Data Generator, strongly recommended). Supports serverless execution, multiple output formats (Parquet/JSON/CSV/Delta), and scales from thousands to millions of rows. For small datasets (<10K rows), can optionally generate locally and upload to volumes. Use when user mentions 'synthetic data', 'test data', 'generate data', 'demo dataset', 'dbldatagen', or 'sample data'." compatibility: Requires databricks CLI (>= v1.0.0) metadata: version: "0.1.0" @@ -11,7 +11,9 @@ parent: databricks-core # Databricks Synthetic Data Generation -Generate realistic, story-driven synthetic data for Databricks using **Spark + Faker + Pandas UDFs** (strongly recommended). +Generate realistic, story-driven synthetic data for Databricks using **Spark + [dbldatagen](https://databrickslabs.github.io/dbldatagen/public_docs/index.html)** (the Databricks Labs Data Generator — strongly recommended). + +> **Use the public `dbldatagen` API.** Import the top-level package (`import dbldatagen as dg`) and its public submodules (`dbldatagen.distributions`, `dbldatagen.constraints`). Everything below is built from the public API. ## Data Must Tell a Business Story @@ -44,11 +46,11 @@ Synthetic data should demonstrate how Databricks helps solve real business probl 5. **Enough data for trends** — ~100K+ rows for main tables so patterns survive aggregation 6. **Ask for catalog/schema** — Never default, always confirm before generating 7. **Present plan for approval** — Show tables, distributions, assumptions before writing code -8. **Master tables first** — Generate parent tables, write to Delta, then create children with valid FKs -9. **Use Spark + Faker + Pandas UDFs** — Scalable, parallel. Polars only if user explicitly wants local + <30K rows +8. **Parent tables first** — Generate parent tables, write to Delta, then create children with valid FKs +9. **Use Spark + dbldatagen** — Scalable, parallel, declarative. Build a `dg.DataGenerator` spec and `.build()` it into a Spark DataFrame. Use a `pandas_udf` only for logic dbldatagen can't express 10. **Use Databricks Connect Serverless by default to generate data** — Update databricks-connect on python 3.12 if required (avoid using execute_code unless instructed to not use Databricks Connect) 11. **No `.cache()` or `.persist()`** — Not supported on serverless. Write to Delta, read back for joins -12. **No Python loops or `.collect()`** — Use Spark parallelism. No driver-side iteration, avoid Pandas↔Spark conversions +12. **No Python loops or `.collect()`** — Use Spark parallelism and dbldatagen specs. No driver-side iteration, avoid Pandas↔Spark conversions ## Generation Planning Workflow @@ -142,29 +144,57 @@ SELECT COUNT(*) FROM parquet.\`/Volumes/CATALOG/SCHEMA/raw_data/customers\` See [references/2-troubleshooting.md](references/2-troubleshooting.md) for full validation examples. -## Use Databricks Connect Spark + Faker Pattern +## Use Databricks Connect + dbldatagen Pattern + +A `DataGenerator` spec declares each column; `.build()` returns a Spark DataFrame. No UDFs, no driver loops — generation is fully distributed across `partitions`. Install `dbldatagen` and `faker` locally first (see Setup). ```python from databricks.connect import DatabricksSession -from pyspark.sql import functions as F -from pyspark.sql.types import StringType -import pandas as pd +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import company, person # providers we draw from -# Setup serverless Spark session +# Setup serverless Spark session (deps installed locally) spark = DatabricksSession.builder.serverless(True).getOrCreate() -# Pandas UDF pattern - import lib INSIDE the function (libs must be installed locally) -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - from faker import Faker # Import inside UDF - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) - -# Generate with spark.range, apply UDFs -customers_df = spark.range(0, 10000, numPartitions=16).select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), +CATALOG, SCHEMA = "", "" # always user-supplied +N_CUSTOMERS = 10_000 + +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company]) + +customers = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, + randomSeed=42, randomSeedMethod="hash_fieldname") + # surrogate key 0..N-1 (the implicit contiguous seed column) — drives FK joins + .withColumn("customer_sk", "long", expr="id") + # business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # realistic text via the shared Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + # email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(replace(name, ' ', '.')), '@example.com')") + # skewed categories — NEVER uniform (weights are relative) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=[40, 25, 20, 15], random=True) + # right-skewed ARR correlated to tier: log-normal = exp() of a standard normal. + # The hidden helper column (_z, omit=True) is computed and reusable as a base column. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") + .withColumn("created_at", "date", + data_range=dg.DateRange("2023-01-01 00:00:00", "2024-12-31 00:00:00", "days=1"), + random=True) ) +customers_df = customers.build() # Write to Volume as Parquet (default for raw data) # Path is a folder with table name: /Volumes/catalog/schema/raw_data/customers/ @@ -173,7 +203,7 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") customers_df.write.mode("overwrite").parquet(f"/Volumes/{CATALOG}/{SCHEMA}/raw_data/customers") ``` -**Partitions by scale:** `spark.range(N, numPartitions=P)` +**Partitions by scale:** `dg.DataGenerator(..., rows=N, partitions=P)` - <100K rows: 8 partitions - 100K-500K: 16 partitions - 500K-1M: 32 partitions @@ -190,32 +220,79 @@ Generated scripts must be highly performant. **Never** do these: | Anti-Pattern | Why It's Slow | Do This Instead | |--------------|---------------|-----------------| -| Python loops on driver | Single-threaded, no parallelism | Use `spark.range()` + Spark operations | +| Python loops on driver | Single-threaded, no parallelism | Declare columns in a `dg.DataGenerator` spec and `.build()` | | `.collect()` then iterate | Brings all data to driver memory | Keep data in Spark, use DataFrame ops | -| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark, use `pandas_udf` only for UDFs | +| Pandas → Spark → Pandas | Serialization overhead, defeats distribution | Stay in Spark; let dbldatagen generate columns | | Read/write temp files | Unnecessary I/O | Chain DataFrame transformations | -| Scalar UDFs | Row-by-row processing | Use `pandas_udf` for batch processing | +| Scalar UDFs | Row-by-row processing | Use dbldatagen `expr`/`template`/`text`; `pandas_udf` only when unavoidable | -**Good pattern:** `spark.range()` → Spark transforms → `pandas_udf` for Faker → write directly +**Good pattern:** `dg.DataGenerator(rows, partitions)` → declarative `.withColumn(...)` specs → `.build()` → write directly ## Common Patterns +All snippets use the public `dbldatagen` API (`import dbldatagen as dg`, `import dbldatagen.distributions as dist`). + ### Weighted Categories (never uniform) +`weights` are relative frequencies — they don't need to sum to 100: ```python -F.when(F.rand() < 0.6, "Free").when(F.rand() < 0.9, "Pro").otherwise("Enterprise") +.withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) ``` -### Log-Normal Amounts (in a pandas UDF) -Use `np.random.lognormal(mean, sigma)` — always positive, long tail: -- Enterprise: `lognormal(7.5, 0.8)` → ~$1800 median -- Pro: `lognormal(5.5, 0.7)` → ~$245 median -- Free: `lognormal(4.0, 0.6)` → ~$55 median +### Skewed / Long-Tailed Amounts +Apply a continuous distribution to a numeric range — `Gamma`/`Exponential` give the always-positive long tail you'd want from log-normal: +```python +.withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) +``` +For a true log-normal, generate over a normal and exponentiate via `expr`: +```python +.withColumn("amount", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # standard normal, hidden +.withColumn("arr", "double", baseColumn="amount", + expr="round(exp(5.5 + 0.8 * amount), 2)") # median ~$245 +``` -### Date Range (Last 6 Months) +### Coherent Rows (correlated attributes via `expr` + `baseColumn`) +Derive dependent columns from earlier ones so each row makes business sense — no UDF needed: ```python -END_DATE = datetime.now() +.withColumn("priority", "string", values=["Critical", "High", "Medium", "Low"], + weights=[5, 15, 50, 30], random=True) +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*8 " + "WHEN 'High' THEN rand()*24 WHEN 'Medium' THEN rand()*72 " + "ELSE rand()*120 END, 1)") +.withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 4 " + "WHEN resolution_hours<72 THEN 3 ELSE 2 END") +``` + +### Date / Timestamp Range (Last 6 Months) +Use `dg.DateRange(begin, end, interval)` with the `data_range` option: +```python +from datetime import datetime, timedelta +END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +FMT = "%Y-%m-%d %H:%M:%S" + +.withColumn("order_ts", "timestamp", + data_range=dg.DateRange(START_DATE.strftime(FMT), END_DATE.strftime(FMT), "minutes=30"), + random=True) +``` + +### Realistic Text (Faker provider plugin) +Build one `FakerTextFactory` with a default locale and the faker providers you need, then reuse it: +```python +from faker.providers import person, company, internet +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, internet]) + +.withColumn("name", "string", text=FakerText("name")) # uses the 'person' provider +.withColumn("company", "string", text=FakerText("company")) # uses the 'company' provider +.withColumn("ip", "string", text=FakerText("ipv4_private")) # uses the 'internet' provider ``` +dbldatagen alternatives that need no library: +* `template=r"\w.\w@\w.com"` (templated text) +* `text=dg.ILText(paragraphs=(1, 3), sentences=(2, 5))` (lorem-ipsum free text). ### Infrastructure (always create in script) ```python @@ -224,27 +301,33 @@ spark.sql(f"CREATE VOLUME IF NOT EXISTS {CATALOG}.{SCHEMA}.raw_data") ``` ### Referential Integrity (FK pattern) -Write master table to Delta first, then read back for FK joins (no `.cache()` on serverless): +Generate the FK as an integer over the parent's surrogate-key range so **every value is valid by construction**, then write the parent to Delta and read it back to attach coherent parent attributes (no `.cache()` on serverless): ```python -# 1. Write master table +# 1. Parent table: surrogate key 0..N-1, then write to Delta customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") -# 2. Read back for FK lookup -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_idx", "customer_id") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +# 2. Child table: FK drawn from the parent key range (valid by construction). +# A distribution skews it 80/20 (a few customers place most orders). +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("order_id", "string", + expr="concat('ORD-', lpad(cast(id as string), 8, '0'))", baseColumn="id") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") + +# 3. Join to the parent (read back from Delta) to pull coherent parent attributes +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "customer_id", "tier") +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") ``` ## Setup -Requires Python 3.12 and databricks-connect>=16.4. Use `uv`: +Requires Python 3.12 and databricks-connect>=16.4. Install dependencies locally with `uv`: ```bash -uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays +uv pip install "databricks-connect>=16.4,<17.4" dbldatagen faker ``` ## Related Skills @@ -256,10 +339,11 @@ uv pip install "databricks-connect>=16.4,<17.4" faker numpy pandas holidays | Issue | Solution | |-------|----------| -| `ModuleNotFoundError: faker` | Install locally: `uv pip install faker`, import inside UDF | -| Faker UDF is slow | Use `pandas_udf` for batch processing | -| Out of memory | Increase `numPartitions` in `spark.range()` | -| Referential integrity errors | Write master table to Delta first, read back for FK joins | +| `ModuleNotFoundError: dbldatagen` (or `faker`) | Install locally: `uv pip install dbldatagen faker` | +| `FakerText`/provider not found | Pass the provider to `dg.FakerTextFactory(providers=[...])` and import it from `faker.providers` | +| All rows identical / not random | Set `random=True` on the column (default is deterministic), or set `randomSeed`/`randomSeedMethod` on the generator | +| Out of memory | Increase `partitions` in `dg.DataGenerator(..., partitions=P)` | +| Referential integrity errors | Draw the FK from the parent key range (`minValue=0, maxValue=N-1`); write parent to Delta first, read back to join attributes | | `PERSIST TABLE is not supported on serverless` | **NEVER use `.cache()` or `.persist()` with serverless** - write to Delta table first, then read back | | `F.window` vs `Window` confusion | Use `from pyspark.sql.window import Window` for `row_number()`, `rank()`, etc. `F.window` is for streaming only. | | Broadcast variables not supported | **NEVER use `spark.sparkContext.broadcast()` with serverless** | diff --git a/skills/databricks-synthetic-data-gen/references/1-data-patterns.md b/skills/databricks-synthetic-data-gen/references/1-data-patterns.md index eba64916..87774896 100755 --- a/skills/databricks-synthetic-data-gen/references/1-data-patterns.md +++ b/skills/databricks-synthetic-data-gen/references/1-data-patterns.md @@ -28,15 +28,31 @@ Real data is never uniformly distributed. Use appropriate distributions: | **Exponential** | Time between events | Support resolution time, session duration | | **Weighted categorical** | Skewed categories | Status (70% complete, 5% failed), tiers | -```python -# Log-normal for amounts (long tail, always positive) -amount = np.random.lognormal(mean=5.5, sigma=0.8) # ~$245 median +dbldatagen applies a distribution to a column's random draw via the `distribution=` option +(`import dbldatagen.distributions as dist`). Supported: `Normal`, `Gamma`, `Beta`, `Exponential`. -# Pareto for power-law (few large, many small) -value = (np.random.pareto(a=1.5) + 1) * base_value +```python +import dbldatagen as dg +import dbldatagen.distributions as dist + +spec = ( + dg.DataGenerator(spark, name="amounts", rows=100_000, partitions=8, randomSeed=42) + # Use Gamma for right-tailed distributions like amounts + .withColumn("order_amount", "decimal(10,2)", minValue=5, maxValue=25_000, + random=True, distribution=dist.Gamma(1.0, 2.0)) + # Use Exponential for waiting times + .withColumn("resolution_hours", "double", minValue=0, maxValue=240, + random=True, distribution=dist.Exponential(rate=1.0 / 24)) +) -# Exponential for time-to-event -hours = np.random.exponential(scale=24) # avg 24h, skewed right +# Exponentiate a standard normal via expr for true log-normal distributions +spec = ( + dg.DataGenerator(spark, name="ln", rows=100_000, partitions=8, randomSeed=42) + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) # hidden helper column + .withColumn("amount", "double", baseColumn="_z", + expr="round(exp(5.5 + 0.8 * _z), 2)") # ~$245 median +) ``` ### 3. Row Coherence @@ -51,20 +67,24 @@ Attributes within a row must make business sense together. Generate correlated a | Large transaction + unusual hour | Higher fraud probability | | Fast resolution | Higher CSAT score | +In dbldatagen, chain `baseColumn` + `expr` when columns derive from other columns +in the same dataset. The generator resolves the dependency order automatically: + ```python -@F.pandas_udf("struct") -def generate_coherent_ticket(tiers: pd.Series) -> pd.DataFrame: - """All attributes correlate logically within each row.""" - results = [] - for tier in tiers: - # Priority depends on tier - priority = "Critical" if tier == "Enterprise" and random() < 0.3 else "Medium" - # Resolution depends on priority - resolution = np.random.exponential(4 if priority == "Critical" else 36) - # CSAT depends on resolution - csat = 5 if resolution < 4 else (3 if resolution < 24 else 2) - results.append({"priority": priority, "resolution_hours": resolution, "csat": csat}) - return pd.DataFrame(results) +spec = ( + dg.DataGenerator(spark, name="tickets", rows=80_000, partitions=8, randomSeed=42) + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=[60, 30, 10], random=True) + # Priority depends on tier + .withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' AND rand()<0.3 THEN 'Critical' ELSE 'Medium' END") + # Resolution depends on priority + .withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 ELSE rand()*36 END, 1)") + # CSAT depends on resolution + .withColumn("csat", "int", baseColumn="resolution_hours", + expr="CASE WHEN resolution_hours<4 THEN 5 WHEN resolution_hours<24 THEN 3 ELSE 2 END") +) ``` ### 4. The 80/20 Rule @@ -75,7 +95,13 @@ Apply power-law distributions where appropriate: - **20% of products** account for 80% of sales - **20% of support agents** handle 80% of tickets -Implementation: Use weighted sampling when assigning FKs, not uniform random. +Implementation: Skew FKs by drawing from a non-uniform distribution. + +```python +# Use a Gamma distribution to skew FK values +.withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) +``` ### 5. Time-Based Patterns @@ -86,14 +112,22 @@ Most data has temporal patterns: - **Seasonality** — Q4 retail spike, summer travel peak - **Trends** — Growth over time, degradation curves +Bias *when* events happen by weighting time buckets, then derive the timestamp from the bucket. +The example below clusters tickets into business hours and skews volume toward weekdays: + ```python -def get_volume_multiplier(date): - multiplier = 1.0 - if date.weekday() >= 5: multiplier *= 0.6 # Weekend drop - if date.month in [11, 12]: multiplier *= 1.5 # Holiday spike - return multiplier +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], # 9am–5pm peak + random=True, omit=True) +.withColumn("day_offset", "int", minValue=0, maxValue=179, random=True, + distribution=dist.Normal(60, 30), omit=True) # volume ramps over the window +.withColumn("event_ts", "timestamp", baseColumn=["day_offset", "hour"], + expr="date_add(timestamp'2025-01-01 00:00:00', day_offset) + make_interval(0,0,0,0,hour,0,0)") ``` +For holiday/seasonal spikes, weight a month or week-of-year bucket the same way, or post-filter +the built DataFrame with a Spark `expr` multiplier. + ### 6. ML-Ready Data If data will train ML models, ensure: @@ -105,20 +139,32 @@ If data will train ML models, ensure: ## Referential Integrity -Generate master tables first, write to Delta, then join for FKs: +Give the parent a surrogate key over `0..N-1`, draw the child's FK from that same range (valid by +construction), then write the parent to Delta and read it back to attach parent attributes: ```python -# 1. Generate and write master table +# 1. # 1. Generate and write parent table with surrogate keys from 0..N-1 +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_id", "string", baseColumn="id", # References the built-in sequential 'id' column + expr="concat('CUST-', lpad(cast(id as string), 5, '0'))") + # ... other attributes ... + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK joins (NOT cache - unsupported on serverless) -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") - -# 3. Generate child table with valid FKs via join -orders_df = spark.range(N_ORDERS).select( - (F.abs(F.hash(F.col("id"))) % N_CUSTOMERS).alias("customer_idx") +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") + +# 3. Generate child table with FKs drawn from the parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True, omit=True) + .withColumn("customer_id", "string", baseColumn="customer_sk", # References the 'customer_sk' column + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + .build() ) -orders_with_fk = orders_df.join(customer_lookup, on="customer_idx") +orders_with_fk = orders_df.join(customer_lookup, on="customer_id", how="inner") ``` ## Data Volume diff --git a/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md b/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md index 793b64f7..8a19ebc8 100755 --- a/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md +++ b/skills/databricks-synthetic-data-gen/references/2-troubleshooting.md @@ -4,23 +4,24 @@ Common issues and solutions for synthetic data generation. ## Environment Issues -### ModuleNotFoundError: faker (or other library) +### ModuleNotFoundError: dbldatagen (or faker) -**Problem:** Dependencies not available in execution environment. +**Problem:** Dependencies not available in execution environment. `dbldatagen` is required, and +`faker` too if you use `FakerTextFactory`/`fakerText`. **Solutions by execution mode:** | Mode | Solution | |------|----------| -| **DB Connect with Serverless** | Install libs locally (`uv pip install faker`), use `DatabricksSession.builder.serverless(True)` | -| **Databricks Runtime** | Use Databricks CLI to install `faker holidays` | -| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "faker"}}, {"pypi": {"package": "holidays"}}]}'` | +| **DB Connect with Serverless** | Install libs locally (`uv pip install dbldatagen faker`), use `DatabricksSession.builder.serverless(True)` | +| **Databricks Runtime** | `%pip install dbldatagen faker` at the top of the notebook | +| **Classic cluster** | Use Databricks CLI to install libraries. `databricks libraries install --json '{"cluster_id": "", "libraries": [{"pypi": {"package": "dbldatagen"}}, {"pypi": {"package": "faker"}}]}'` | ```python # For DB Connect with serverless from databricks.connect import DatabricksSession -# Install dependencies locally first: uv pip install faker pandas numpy holidays +# Install dependencies locally first: uv pip install dbldatagen faker spark = DatabricksSession.builder.serverless(True).getOrCreate() ``` @@ -50,21 +51,21 @@ AnalysisException: [NOT_SUPPORTED_WITH_SERVERLESS] PERSIST TABLE is not supporte **Why this happens:** Serverless compute does not support caching DataFrames in memory. This is a fundamental limitation of the serverless architecture. -**Solution:** Write master tables to Delta first, then read them back for FK joins: +**Solution:** Write parent tables to Delta first, then read them back for FK joins: ```python # BAD - will fail on serverless -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.cache() # ❌ FAILS: "PERSIST TABLE is not supported on serverless compute" # GOOD - write to Delta, then read back -customers_df = spark.range(0, N_CUSTOMERS)... +customers_df = dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8).build() customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers") # ✓ Read from Delta ``` **Best practice for referential integrity:** -1. Generate master table (e.g., customers) +1. Generate parent table (e.g., customers) 2. Write to Delta table 3. Read back for FK lookup joins 4. Generate child tables (e.g., orders, tickets) with valid FKs @@ -114,7 +115,7 @@ spark = DatabricksSession.builder.serverless(True).getOrCreate() "environment_key": "datagen_env", "spec": { "client": "4", # Required! - "dependencies": ["faker", "numpy", "pandas"] + "dependencies": ["dbldatagen", "faker"] } }] } @@ -153,34 +154,37 @@ contacts_df = contacts_df.withColumn( --- -### Faker UDF is slow +### Generation is slow -**Problem:** Single-row UDFs don't parallelize well. +**Problem:** Row-by-row Python UDFs (or driver loops) don't parallelize well. -**Solution:** Use `pandas_udf` for batch processing: +**Solution:** Let dbldatagen generate columns declaratively — it builds the whole DataFrame in +parallel across `partitions`. Prefer `expr`, `template`, `values`/`weights`, and `distribution` +over UDFs. For names/addresses use the Faker plugin (`FakerTextFactory`) instead of hand-rolled UDFs: ```python -# SLOW - scalar UDF -@F.udf(returnType=StringType()) -def slow_fake_name(): - return Faker().name() - -# FAST - pandas UDF (batch processing) -@F.pandas_udf(StringType()) -def fast_fake_name(ids: pd.Series) -> pd.Series: - fake = Faker() - return pd.Series([fake.name() for _ in range(len(ids))]) +import dbldatagen as dg +from faker.providers import person + +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person]) + +spec = ( + dg.DataGenerator(spark, name="people", rows=1_000_000, partitions=64, randomSeed=42) + .withColumn("name", "string", text=FakerText("name")) # batched by dbldatagen + .withColumn("status", "string", values=["active", "churned"], weights=[85, 15], random=True) +) +df = spec.build() ``` ### Out of memory with large data **Problem:** Not enough partitions for data size. -**Solution:** Increase partitions: +**Solution:** Increase `partitions` on the generator: ```python # For large datasets (1M+ rows) -customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from default +spec = dg.DataGenerator(spark, name="big", rows=N_CUSTOMERS, partitions=64, randomSeed=42) ``` | Data Size | Recommended Partitions | @@ -205,21 +209,28 @@ customers_df = spark.range(0, N_CUSTOMERS, numPartitions=64) # Increase from de **Problem:** Foreign keys reference non-existent parent records. -**Solution:** Write master table to Delta first, then read back for FK joins: +**Solution:** Sample the FK from the parent's surrogate-key range, write the parent +to Delta, then read it back for joins: ```python -# 1. Generate and WRITE master table (do NOT use cache with serverless!) -customers_df = spark.range(0, N_CUSTOMERS)... +# 1. Generate and WRITE parent table (do NOT use cache with serverless!) +customers_df = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=8, randomSeed=42) + .withColumn("customer_sk", "long", expr="id") # References the built-in sequential 'id' column + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], weights=[60, 30, 10], random=True) + .build() +) customers_df.write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.customers") # 2. Read back for FK lookups -customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_id", "tier") - -# 3. Generate child table with valid FKs -orders_df = spark.range(0, N_ORDERS).join( - customer_lookup, - on=, - how="left" +customer_lookup = spark.table(f"{CATALOG}.{SCHEMA}.customers").select("customer_sk", "tier") + +# 3. Generate child table with FK in the valid parent key range +orders_df = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=16, randomSeed=42) + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, random=True) + .build() + .join(customer_lookup, on="customer_sk", how="inner") ) ``` @@ -233,52 +244,52 @@ orders_df = spark.range(0, N_ORDERS).join( **Problem:** All customers have similar order counts, amounts are evenly distributed. -**Solution:** Use non-linear distributions: +**Solution:** Apply a `distribution=` to the column instead of leaving it uniform: ```python -# BAD - uniform -amounts = np.random.uniform(10, 1000, N) +import dbldatagen.distributions as dist -# GOOD - log-normal (realistic) -amounts = np.random.lognormal(mean=5, sigma=0.8, N) +# BAD - uniform (no distribution) +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True) + +# GOOD - skewed, realistic +.withColumn("amount", "double", minValue=10, maxValue=1000, random=True, + distribution=dist.Gamma(1.0, 2.0)) ``` ### Missing time-based patterns **Problem:** Data doesn't reflect weekday/weekend or seasonal patterns. -**Solution:** Add multipliers: +**Solution:** Weight time buckets so volume clusters realistically, then derive the timestamp: ```python -import holidays - -US_HOLIDAYS = holidays.US(years=[2024, 2025]) - -def get_multiplier(date): - mult = 1.0 - if date.weekday() >= 5: # Weekend - mult *= 0.6 - if date in US_HOLIDAYS: - mult *= 0.3 - return mult +# Cluster events into business hours +.withColumn("hour", "int", values=list(range(24)), + weights=[1,1,1,1,1,1,2,4,8,10,10,9,8,9,10,9,7,5,3,2,2,1,1,1], + random=True, omit=True) +.withColumn("event_ts", "timestamp", baseColumn="hour", + expr="date'2025-06-01' + make_interval(0,0,0,0,hour,0,0)") ``` +For weekend/holiday dips, weight a day-of-week or day bucket the same way, or post-filter the built +DataFrame with a Spark `expr` (e.g. `dayofweek(event_ts) IN (1,7)`). + ### Incoherent row attributes **Problem:** Enterprise customer has low-value orders, critical ticket has slow resolution. -**Solution:** Correlate attributes: +**Solution:** Correlate attributes with `baseColumn` + `expr` so each derives from the previous: ```python # Priority based on tier -if tier == 'Enterprise': - priority = np.random.choice(['Critical', 'High'], p=[0.4, 0.6]) -else: - priority = np.random.choice(['Medium', 'Low'], p=[0.6, 0.4]) - -# Resolution based on priority -resolution_scale = {'Critical': 4, 'High': 12, 'Medium': 36, 'Low': 72} -resolution_hours = np.random.exponential(scale=resolution_scale[priority]) +.withColumn("priority", "string", baseColumn="tier", + expr="CASE WHEN tier='Enterprise' THEN (CASE WHEN rand()<0.4 THEN 'Critical' ELSE 'High' END) " + "ELSE (CASE WHEN rand()<0.6 THEN 'Medium' ELSE 'Low' END) END") +# Resolution scaled by priority +.withColumn("resolution_hours", "double", baseColumn="priority", + expr="round(CASE priority WHEN 'Critical' THEN rand()*4 WHEN 'High' THEN rand()*12 " + "WHEN 'Medium' THEN rand()*36 ELSE rand()*72 END, 1)") ``` --- diff --git a/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py b/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py index b36edb8e..185e4a18 100755 --- a/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py +++ b/skills/databricks-synthetic-data-gen/scripts/generate_synthetic_data.py @@ -1,22 +1,24 @@ -"""Generate synthetic data using Spark + Faker + Pandas UDFs. +"""Generate synthetic data using Spark + dbldatagen (Databricks Labs Data Generator). This is the recommended approach for ALL data generation tasks: - Scales from thousands to millions of rows -- Parallel execution via Spark +- Declarative, parallel generation (no driver loops, no row-by-row UDFs) - Direct write to Unity Catalog - Works with serverless and classic compute +Uses ONLY the public dbldatagen API: +- `dbldatagen.DataGenerator` / `.withColumn(...)` / `.build()` +- `dbldatagen.distributions` for skew +- `dbldatagen.FakerTextFactory` for realistic names/companies/addresses + Prerequisites: -- Install dependencies locally: uv pip install faker pandas numpy holidays databricks-connect +- Install dependencies locally: uv pip install dbldatagen faker databricks-connect - Configure ~/.databrickscfg with serverless_compute_id = auto """ import sys import os from pyspark.sql import functions as F from pyspark.sql.window import Window -from pyspark.sql.types import StringType, DoubleType, StructType, StructField, IntegerType -import numpy as np -import pandas as pd from datetime import datetime, timedelta # ============================================================================= @@ -39,6 +41,7 @@ # Date range - last 6 months from today END_DATE = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) START_DATE = END_DATE - timedelta(days=180) +DATE_FMT = "%Y-%m-%d %H:%M:%S" # Write mode - "overwrite" for one-time, "append" for incremental WRITE_MODE = "overwrite" @@ -51,14 +54,18 @@ "orphan_fk_rate": 0.01, # 1% orphan foreign keys } -# Reproducibility +# Reproducibility - same seed => same data SEED = 42 -# Tier distribution: Free 60%, Pro 30%, Enterprise 10% -TIER_PROBS = [0.6, 0.3, 0.1] +# Tier distribution: Free 60%, Pro 30%, Enterprise 10% (relative weights) +TIER_WEIGHTS = [60, 30, 10] + +# Region distribution (relative weights) +REGION_WEIGHTS = [40, 25, 20, 15] -# Region distribution -REGION_PROBS = [0.4, 0.25, 0.2, 0.15] +# Order status distribution (relative weights) +STATUS_VALUES = ["delivered", "shipped", "processing", "pending", "cancelled"] +STATUS_WEIGHTS = [65, 15, 10, 5, 5] # ============================================================================= # SESSION CREATION @@ -79,59 +86,14 @@ spark = DatabricksSession.builder.clusterId(CLUSTER_ID).getOrCreate() print(f"Connected to cluster {CLUSTER_ID}") -# Import Faker for UDF definitions -from faker import Faker - -# ============================================================================= -# DEFINE PANDAS UDFs FOR FAKER DATA -# ============================================================================= +# Import the public dbldatagen API (installed locally) +import dbldatagen as dg +import dbldatagen.distributions as dist +from faker.providers import person, company, address as faker_address -@F.pandas_udf(StringType()) -def fake_name(ids: pd.Series) -> pd.Series: - """Generate realistic person names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.name() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_company(ids: pd.Series) -> pd.Series: - """Generate realistic company names.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.company() for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_address(ids: pd.Series) -> pd.Series: - """Generate realistic addresses.""" - fake = Faker() - Faker.seed(SEED) - return pd.Series([fake.address().replace('\n', ', ') for _ in range(len(ids))]) - -@F.pandas_udf(StringType()) -def fake_email(names: pd.Series) -> pd.Series: - """Generate email based on name.""" - emails = [] - for name in names: - if name: - domain = name.lower().replace(" ", ".").replace(",", "")[:20] - emails.append(f"{domain}@example.com") - else: - emails.append("unknown@example.com") - return pd.Series(emails) - -@F.pandas_udf(DoubleType()) -def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: - """Generate amount based on tier using log-normal distribution.""" - np.random.seed(SEED) - amounts = [] - for tier in tiers: - if tier == "Enterprise": - amounts.append(float(np.random.lognormal(mean=7.5, sigma=0.8))) # ~$1800 avg - elif tier == "Pro": - amounts.append(float(np.random.lognormal(mean=5.5, sigma=0.7))) # ~$245 avg - else: - amounts.append(float(np.random.lognormal(mean=4.0, sigma=0.6))) # ~$55 avg - return pd.Series(amounts) +# One shared Faker text factory: default locale + the providers we use. +# The factory builds the Faker instance internally (no `from faker import Faker` needed). +FakerText = dg.FakerTextFactory(locale=["en_US"], providers=[person, company, faker_address]) # ============================================================================= # CREATE INFRASTRUCTURE @@ -142,39 +104,50 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f"Infrastructure ready: {VOLUME_PATH}") # ============================================================================= -# GENERATE CUSTOMERS (Master Table) +# GENERATE CUSTOMERS (Parent Table) # ============================================================================= print(f"\nGenerating {N_CUSTOMERS:,} customers...") -customers_df = ( - spark.range(0, N_CUSTOMERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("CUST-"), F.lpad(F.col("id").cast("string"), 5, "0")).alias("customer_id"), - fake_name(F.col("id")).alias("name"), - fake_company(F.col("id")).alias("company"), - fake_address(F.col("id")).alias("address"), - # Tier distribution: Free 60%, Pro 30%, Enterprise 10% - F.when(F.rand(SEED) < TIER_PROBS[0], "Free") - .when(F.rand(SEED) < TIER_PROBS[0] + TIER_PROBS[1], "Pro") - .otherwise("Enterprise").alias("tier"), - # Region distribution - F.when(F.rand(SEED) < REGION_PROBS[0], "North") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1], "South") - .when(F.rand(SEED) < REGION_PROBS[0] + REGION_PROBS[1] + REGION_PROBS[2], "East") - .otherwise("West").alias("region"), - # Created date (within last 2 years before start date) - F.date_sub(F.lit(START_DATE.date()), (F.rand(SEED) * 730).cast("int")).alias("created_at"), - ) +customers_spec = ( + dg.DataGenerator(spark, name="customers", rows=N_CUSTOMERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + # Surrogate key 0..N-1 (the implicit contiguous seed column) — stable join key + .withColumn("customer_sk", "long", expr="id") + # Business key derived from the surrogate + .withColumn("customer_id", "string", baseColumn="customer_sk", + expr="concat('CUST-', lpad(cast(customer_sk as string), 5, '0'))") + # Realistic text via the Faker provider factory + .withColumn("name", "string", text=FakerText("name")) + .withColumn("company", "string", text=FakerText("company")) + .withColumn("address", "string", text=FakerText("address")) + # Email derived from the generated name (name as base column) + .withColumn("email", "string", baseColumn="name", + expr="concat(lower(regexp_replace(name, '[^A-Za-z]+', '.')), '@example.com')") + # Skewed categories (never uniform) — weights are relative frequencies + .withColumn("tier", "string", values=["Free", "Pro", "Enterprise"], + weights=TIER_WEIGHTS, random=True) + .withColumn("region", "string", values=["North", "South", "East", "West"], + weights=REGION_WEIGHTS, random=True) + # Account created within 2 years before the analysis window + .withColumn("created_at", "date", + data_range=dg.DateRange( + (START_DATE - timedelta(days=730)).strftime(DATE_FMT), + START_DATE.strftime(DATE_FMT), + "days=1"), + random=True) + # Tier-based ARR via a log-normal (exp of a standard normal) — long tail, always positive. + # The omitted helper column '_z' keeps ARR reproducible under randomSeed. + # Enterprise ~ $1800, Pro ~ $245, Free ~ $55 median. + .withColumn("_z", "double", minValue=-1, maxValue=1, random=True, + distribution=dist.Normal(0.0, 1.0), omit=True) + .withColumn("arr", "double", baseColumn=["tier", "_z"], + expr="round(CASE tier WHEN 'Enterprise' THEN exp(7.5 + 0.8 * _z) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * _z) ELSE exp(4.0 + 0.6 * _z) END, 2)") ) -# Add tier-based ARR and email -customers_df = ( - customers_df - .withColumn("arr", F.round(generate_lognormal_amount(F.col("tier")), 2)) - .withColumn("email", fake_email(F.col("name"))) -) +customers_df = customers_spec.build() -# Save customers +# Save customers as raw Parquet customers_df.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/customers") print(f" Saved customers to {VOLUME_PATH}/customers") @@ -183,46 +156,45 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: # ============================================================================= print(f"\nGenerating {N_ORDERS:,} orders with referential integrity...") -# Write customer lookup to temp Delta table (no .cache() on serverless!) +# Write a customer lookup to a temp Delta table (no .cache() on serverless!) customers_tmp_table = f"{CATALOG}.{SCHEMA}._tmp_customers_lookup" -customers_df.select("customer_id", "tier").write.mode("overwrite").saveAsTable(customers_tmp_table) +(customers_df.select("customer_sk", "customer_id", "tier") + .write.mode("overwrite").saveAsTable(customers_tmp_table)) customer_lookup = spark.table(customers_tmp_table) -# Generate orders base -orders_df = ( - spark.range(0, N_ORDERS, numPartitions=PARTITIONS) - .select( - F.concat(F.lit("ORD-"), F.lpad(F.col("id").cast("string"), 6, "0")).alias("order_id"), - # Generate customer_idx for FK join (hash-based distribution) - (F.abs(F.hash(F.col("id"), F.lit(SEED))) % N_CUSTOMERS).alias("customer_idx"), - # Order status - F.when(F.rand(SEED) < 0.65, "delivered") - .when(F.rand(SEED) < 0.80, "shipped") - .when(F.rand(SEED) < 0.90, "processing") - .when(F.rand(SEED) < 0.95, "pending") - .otherwise("cancelled").alias("status"), - # Order date within date range - F.date_add(F.lit(START_DATE.date()), (F.rand(SEED) * 180).cast("int")).alias("order_date"), - ) +# Generate orders. The FK (customer_sk) is drawn from the parent key range, so every +# value is valid by construction. A Gamma distribution skews it 80/20 (a few customers +# place most orders). +orders_spec = ( + dg.DataGenerator(spark, name="orders", rows=N_ORDERS, partitions=PARTITIONS, + randomSeed=SEED, randomSeedMethod="hash_fieldname") + .withColumn("order_id", "string", baseColumn="id", + expr="concat('ORD-', lpad(cast(id as string), 6, '0'))") + .withColumn("customer_sk", "long", minValue=0, maxValue=N_CUSTOMERS - 1, + random=True, distribution=dist.Gamma(0.4, 1.0)) + .withColumn("status", "string", values=STATUS_VALUES, weights=STATUS_WEIGHTS, random=True) + .withColumn("order_date", "date", + data_range=dg.DateRange(START_DATE.strftime(DATE_FMT), + END_DATE.strftime(DATE_FMT), "days=1"), + random=True) ) +orders_df = orders_spec.build() -# Add customer_idx to lookup for join -customer_lookup_with_idx = customer_lookup.withColumn( - "customer_idx", - (F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) - 1).cast("int") -) - -# Join to get customer_id and tier as foreign key -orders_with_fk = ( - orders_df - .join(customer_lookup_with_idx, on="customer_idx", how="left") - .drop("customer_idx") -) +# Join to the parent (read back from Delta) to attach valid customer_id + tier +orders_with_fk = orders_df.join(customer_lookup, on="customer_sk", how="inner") -# Add tier-based amount +# Tier-based amount (log-normal), coherent with the customer's tier orders_with_fk = orders_with_fk.withColumn( "amount", - F.round(generate_lognormal_amount(F.col("tier")), 2) + F.round( + F.expr( + "CASE tier " + "WHEN 'Enterprise' THEN exp(7.5 + 0.8 * randn()) " + "WHEN 'Pro' THEN exp(5.5 + 0.7 * randn()) " + "ELSE exp(4.0 + 0.6 * randn()) END" + ), + 2, + ), ) # ============================================================================= @@ -236,19 +208,19 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: outlier_count = int(N_ORDERS * BAD_DATA_CONFIG["outlier_rate"]) orphan_count = int(N_ORDERS * BAD_DATA_CONFIG["orphan_fk_rate"]) - # Add bad data flags + # Add a row number to target specific rows (no cache/persist on serverless) orders_with_fk = orders_with_fk.withColumn( "row_num", F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) ) - # Inject nulls in customer_id for first null_count rows + # Inject nulls in customer_id for the first null_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when(F.col("row_num") <= null_count, None).otherwise(F.col("customer_id")) ) - # Inject negative amounts for next outlier_count rows + # Inject negative amounts for the next outlier_count rows orders_with_fk = orders_with_fk.withColumn( "amount", F.when( @@ -257,7 +229,7 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: ).otherwise(F.col("amount")) ) - # Inject orphan FKs for next orphan_count rows + # Inject orphan FKs for the next orphan_count rows orders_with_fk = orders_with_fk.withColumn( "customer_id", F.when( @@ -273,8 +245,8 @@ def generate_lognormal_amount(tiers: pd.Series) -> pd.Series: print(f" Injected {outlier_count} negative amounts") print(f" Injected {orphan_count} orphan foreign keys") -# Drop tier column (not needed in final output) -orders_final = orders_with_fk.drop("tier") +# Drop join-only columns (not needed in final output) +orders_final = orders_with_fk.drop("tier", "customer_sk") # Save orders orders_final.write.mode(WRITE_MODE).parquet(f"{VOLUME_PATH}/orders")