dash-synthetic 0.1.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,53 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+ branches: [main]
8
+
9
+ jobs:
10
+ lint:
11
+ runs-on: ubuntu-latest
12
+ steps:
13
+ - uses: actions/checkout@v4
14
+ - uses: actions/setup-python@v5
15
+ with:
16
+ python-version: "3.11"
17
+ - run: pip install ruff
18
+ - run: ruff check dashsynthetic/
19
+
20
+ test:
21
+ runs-on: ubuntu-latest
22
+ needs: lint
23
+ strategy:
24
+ matrix:
25
+ python-version: ["3.9", "3.10", "3.11", "3.12"]
26
+ steps:
27
+ - uses: actions/checkout@v4
28
+ - uses: actions/setup-python@v5
29
+ with:
30
+ python-version: ${{ matrix.python-version }}
31
+ - name: Install
32
+ run: pip install -e ".[dev]" pytest pytest-cov
33
+ - name: Test
34
+ run: pytest tests/ -v --cov=dashsynthetic --cov-report=xml
35
+ - name: Upload coverage
36
+ uses: codecov/codecov-action@v4
37
+ with:
38
+ files: coverage.xml
39
+
40
+ build:
41
+ runs-on: ubuntu-latest
42
+ needs: test
43
+ steps:
44
+ - uses: actions/checkout@v4
45
+ - uses: actions/setup-python@v5
46
+ with:
47
+ python-version: "3.11"
48
+ - run: pip install hatch
49
+ - run: hatch build
50
+ - uses: actions/upload-artifact@v4
51
+ with:
52
+ name: dist
53
+ path: dist/
@@ -0,0 +1,33 @@
1
+ name: Daily Tests
2
+
3
+ on:
4
+ schedule:
5
+ - cron: "0 6 * * *" # Every day 06:00 UTC — tests only, no commit
6
+ workflow_dispatch:
7
+
8
+ jobs:
9
+ test:
10
+ name: Test (Python ${{ matrix.python-version }})
11
+ runs-on: ubuntu-latest
12
+ strategy:
13
+ fail-fast: false
14
+ matrix:
15
+ python-version: ["3.9", "3.10", "3.11", "3.12"]
16
+ steps:
17
+ - uses: actions/checkout@v4
18
+
19
+ - uses: actions/setup-python@v5
20
+ with:
21
+ python-version: ${{ matrix.python-version }}
22
+
23
+ - name: Install
24
+ run: pip install -e ".[dev]" pytest pytest-cov
25
+
26
+ - name: Run tests
27
+ run: pytest tests/ -v --cov=dashsynthetic --cov-report=xml --cov-report=term-missing
28
+
29
+ - name: Upload coverage
30
+ uses: codecov/codecov-action@v4
31
+ with:
32
+ files: coverage.xml
33
+ fail_ci_if_error: false
@@ -0,0 +1,262 @@
1
+ name: Weekly Release
2
+
3
+ on:
4
+ schedule:
5
+ - cron: "0 9 * * 1" # Every Monday 09:00 UTC
6
+ workflow_dispatch:
7
+ inputs:
8
+ release_note:
9
+ description: "Optional release note (shown in GitHub release body)"
10
+ required: false
11
+ default: ""
12
+
13
+ jobs:
14
+ # ── Gate: tests must pass ────────────────────────────────────────────────
15
+ test:
16
+ name: Test (Python ${{ matrix.python-version }})
17
+ runs-on: ubuntu-latest
18
+ strategy:
19
+ fail-fast: true
20
+ matrix:
21
+ python-version: ["3.9", "3.10", "3.11", "3.12"]
22
+ steps:
23
+ - uses: actions/checkout@v4
24
+
25
+ - uses: actions/setup-python@v5
26
+ with:
27
+ python-version: ${{ matrix.python-version }}
28
+
29
+ - name: Install
30
+ run: pip install -e ".[dev]" pytest pytest-cov
31
+
32
+ - name: Run tests
33
+ run: pytest tests/ -v --cov=dashsynthetic --cov-report=xml --cov-report=term-missing
34
+
35
+ - name: Upload coverage
36
+ uses: codecov/codecov-action@v4
37
+ with:
38
+ files: coverage.xml
39
+ fail_ci_if_error: false
40
+
41
+ # ── Generate docs ────────────────────────────────────────────────────────
42
+ docs:
43
+ name: Generate API docs
44
+ runs-on: ubuntu-latest
45
+ needs: test
46
+ steps:
47
+ - uses: actions/checkout@v4
48
+
49
+ - uses: actions/setup-python@v5
50
+ with:
51
+ python-version: "3.11"
52
+
53
+ - name: Install
54
+ run: pip install -e ".[dev]" pdoc
55
+
56
+ - name: Generate docs
57
+ run: |
58
+ pdoc dashsynthetic --output-dir docs/api --docformat google
59
+ echo "Docs generated at $(date -u)" > docs/api/.generated
60
+
61
+ - name: Upload docs artifact
62
+ uses: actions/upload-artifact@v4
63
+ with:
64
+ name: api-docs
65
+ path: docs/api/
66
+
67
+ # ── Release: tag, GitHub release, commit docs ────────────────────────────
68
+ release:
69
+ name: Bump version & release
70
+ runs-on: ubuntu-latest
71
+ needs: [test, docs]
72
+ permissions:
73
+ contents: write
74
+ pull-requests: write
75
+ outputs:
76
+ version: ${{ steps.bump.outputs.version }}
77
+ steps:
78
+ - uses: actions/checkout@v4
79
+ with:
80
+ fetch-depth: 0
81
+
82
+ - uses: actions/setup-python@v5
83
+ with:
84
+ python-version: "3.11"
85
+
86
+ - name: Install build tools
87
+ run: pip install hatch pdoc
88
+
89
+ - name: Bump patch version
90
+ id: bump
91
+ run: |
92
+ current=$(hatch version)
93
+ hatch version patch
94
+ new=$(hatch version)
95
+ echo "version=$new" >> $GITHUB_OUTPUT
96
+ echo "prev_version=$current" >> $GITHUB_OUTPUT
97
+ echo "Bumped $current → $new"
98
+
99
+ - name: Regenerate docs into repo
100
+ run: |
101
+ pip install -e ".[dev]"
102
+ pdoc dashsynthetic --output-dir docs/api --docformat google
103
+
104
+ - name: Build wheel + sdist
105
+ run: hatch build
106
+
107
+ - name: Write release notes
108
+ env:
109
+ VERSION: ${{ steps.bump.outputs.version }}
110
+ PREV_VERSION: ${{ steps.bump.outputs.prev_version }}
111
+ RELEASE_NOTE: ${{ github.event.inputs.release_note }}
112
+ run: |
113
+ cat > RELEASE_NOTES.md << EOF
114
+ ## DashSynthetic — Synthetic Data v${VERSION}
115
+
116
+ **Released:** $(date -u '+%Y-%m-%d')
117
+ **Previous:** v${PREV_VERSION}
118
+
119
+ $( [ -n "${RELEASE_NOTE}" ] && echo "### Notes" && echo "${RELEASE_NOTE}" || true )
120
+
121
+ ### What's included
122
+ - All tests passing across Python 3.9, 3.10, 3.11, 3.12
123
+ - API documentation regenerated (see \`docs/api/\`)
124
+ - Published to PyPI and Databricks Marketplace
125
+
126
+ ### Install
127
+ \`\`\`bash
128
+ pip install dash-synthetic==${VERSION}
129
+ \`\`\`
130
+
131
+ ### Quick Start (Databricks notebook)
132
+ \`\`\`python
133
+ %pip install dash-synthetic==${VERSION}
134
+ import dashsynthetic
135
+ dashsynthetic.launch()
136
+ \`\`\`
137
+ EOF
138
+
139
+ - name: Commit version bump + docs to a release branch
140
+ env:
141
+ VERSION: ${{ steps.bump.outputs.version }}
142
+ run: |
143
+ git config user.name "github-actions[bot]"
144
+ git config user.email "github-actions[bot]@users.noreply.github.com"
145
+ git checkout -b "release/v${VERSION}"
146
+ git add .
147
+ git commit -m "release: v${VERSION} — tests passed, docs updated"
148
+ git tag "v${VERSION}"
149
+ git push origin "release/v${VERSION}" --tags
150
+
151
+ - name: Open and auto-merge release PR
152
+ env:
153
+ GH_TOKEN: ${{ github.token }}
154
+ VERSION: ${{ steps.bump.outputs.version }}
155
+ run: |
156
+ gh pr create --base main --head "release/v${VERSION}" \
157
+ --title "release: v${VERSION}" \
158
+ --body "Automated release PR — tests passed, docs regenerated, version bumped to v${VERSION}."
159
+ gh pr merge "release/v${VERSION}" --merge --auto --delete-branch
160
+
161
+ - name: Create GitHub Release
162
+ uses: softprops/action-gh-release@v2
163
+ with:
164
+ tag_name: "v${{ steps.bump.outputs.version }}"
165
+ body_path: RELEASE_NOTES.md
166
+ files: dist/*
167
+
168
+ - name: Upload dist artifact for PyPI job
169
+ uses: actions/upload-artifact@v4
170
+ with:
171
+ name: dist
172
+ path: dist/
173
+
174
+ # ── Publish to PyPI (Trusted Publisher / OIDC — no token needed) ─────────
175
+ publish-pypi:
176
+ name: Publish to PyPI
177
+ runs-on: ubuntu-latest
178
+ needs: release
179
+ permissions:
180
+ id-token: write # required for OIDC trusted publisher
181
+ environment:
182
+ name: pypi
183
+ url: https://pypi.org/project/dash-synthetic
184
+ steps:
185
+ - name: Download dist
186
+ uses: actions/download-artifact@v4
187
+ with:
188
+ name: dist
189
+ path: dist/
190
+
191
+ - name: Publish to PyPI
192
+ uses: pypa/gh-action-pypi-publish@release/v1
193
+
194
+ # ── Package for Databricks Marketplace ───────────────────────────────────
195
+ publish-databricks:
196
+ name: Package for Databricks Marketplace
197
+ runs-on: ubuntu-latest
198
+ needs: release
199
+ steps:
200
+ - uses: actions/checkout@v4
201
+ with:
202
+ ref: "v${{ needs.release.outputs.version }}"
203
+
204
+ - name: Download dist
205
+ uses: actions/download-artifact@v4
206
+ with:
207
+ name: dist
208
+ path: dist/
209
+
210
+ - name: Build Marketplace bundle
211
+ env:
212
+ VERSION: ${{ needs.release.outputs.version }}
213
+ run: |
214
+ mkdir -p marketplace-bundle/files
215
+
216
+ # Copy wheel
217
+ cp dist/*.whl marketplace-bundle/files/
218
+
219
+ # Generate companion notebook
220
+ cat > marketplace-bundle/files/DashSynthetic — Synthetic Data_Quickstart.py << NBEOF
221
+ # Databricks notebook source
222
+ # MAGIC %md
223
+ # MAGIC # DashSynthetic — Synthetic Data v${VERSION} — Data Quality for Databricks
224
+ # MAGIC Install and launch the interactive UI.
225
+
226
+ # COMMAND ----------
227
+ # MAGIC %pip install dash-synthetic==${VERSION}
228
+
229
+ # COMMAND ----------
230
+ dbutils.library.restartPython()
231
+
232
+ # COMMAND ----------
233
+ import dashsynthetic
234
+ dashsynthetic.launch()
235
+ NBEOF
236
+
237
+ # Generate listing metadata
238
+ cat > marketplace-bundle/listing.json << LEOF
239
+ {
240
+ "listing_name": "DashSynthetic — Synthetic Data — Data Quality",
241
+ "version": "${VERSION}",
242
+ "short_description": "Generate privacy-safe synthetic data from real Databricks tables",
243
+ "long_description": "DashSynthetic — Synthetic Data provides an ipywidgets UI inside Databricks notebooks to run data quality checks on DataFrames, Unity Catalog tables, and SQL queries. No coding required for business users.",
244
+ "categories": ["Synthetic Data,Data Engineering"],
245
+ "tags": ["data-quality", "databricks", "unity-catalog", "pyspark"],
246
+ "provider": "dash-libs",
247
+ "documentation_url": "https://github.com/dash-libs/dash-synthetic",
248
+ "source_url": "https://github.com/dash-libs/dash-synthetic",
249
+ "pypi_package": "dash-synthetic==${VERSION}"
250
+ }
251
+ LEOF
252
+
253
+ # Zip the bundle
254
+ cd marketplace-bundle && zip -r ../dashsynthetic-marketplace-${VERSION}.zip .
255
+ echo "Bundle created: dashsynthetic-marketplace-${VERSION}.zip"
256
+
257
+ - name: Upload Marketplace bundle artifact
258
+ uses: actions/upload-artifact@v4
259
+ with:
260
+ name: marketplace-bundle
261
+ path: dashsynthetic-marketplace-*.zip
262
+ retention-days: 90
@@ -0,0 +1,9 @@
1
+ __pycache__/
2
+ *.pyc
3
+ *.egg-info/
4
+ dist/
5
+ build/
6
+ .coverage
7
+ coverage.xml
8
+ .pytest_cache/
9
+ .ruff_cache/
@@ -0,0 +1,35 @@
1
+ # CLAUDE.md — dash-synthetic
2
+
3
+ Part of the **Dashlibs** suite. See ~/dashlibs for the full context.
4
+
5
+ ## Purpose
6
+ Synthetic data generation from real Databricks tables. generator.py=generation API
7
+ (`SyntheticGenerator` single-table, `MultiTableGenerator` multi-table), profiler.py=source
8
+ profiling, engine.py=numpy-based correlated sampling, relationships.py=`RelationshipGraph`
9
+ (tables, primary/foreign keys, master data columns, dependency order), multi_engine.py=
10
+ FK-aware multi-table orchestration. UI widgets come from the shared `dashui` package
11
+ (PyPI distribution name `dash-uis` — `dash-ui` was already taken).
12
+
13
+ ## Structure
14
+ - `/ui.py` — ipywidgets UI (Single Table + Multi-Table Relationships tabs), `launch()` entrypoint
15
+ - `/relationships.py` — `RelationshipGraph`: define tables, PK/FK, master data columns, topo-sort generation order
16
+ - `/multi_engine.py` — `generate_multi()` / `TableGenSpec`: runs tables in dependency order, samples FKs from parent PKs
17
+ - `/*.py` — core logic
18
+ - `tests/` — pytest, no Spark dependency for unit tests
19
+
20
+ ## Key Design Rules
21
+ - Never import Spark at module level — always inside functions
22
+ - UI calls core classes; never contains business logic
23
+ - `launch()` is always the public entrypoint for business users
24
+ - Reusable widgets (header, source picker, output panel, schema lookups) come from `dashui`
25
+ — don't reimplement them locally; add new generic widgets to the `dash-ui` repo (PyPI: `dash-uis`) instead
26
+ - Foreign-key columns are sampled from the parent table's *already-generated* primary key
27
+ values (`pk_pools` in `multi_engine.py`), so generation order matters — always go through
28
+ `RelationshipGraph.generation_order()` rather than an arbitrary table order
29
+ - "Master data" columns (e.g. currency/country codes) are forced into categorical sampling
30
+ from real distinct source values regardless of dtype, via `force_categorical` in `engine.generate()`
31
+
32
+ ## CI
33
+ - `ci.yml` — PR gate: lint → test → build
34
+ - `daily.yml` — 06:00 UTC: tests + .health/log.txt commit
35
+ - `release.yml`— Monday 09:00 UTC: patch bump + GitHub release
@@ -0,0 +1,78 @@
1
+ Metadata-Version: 2.4
2
+ Name: dash-synthetic
3
+ Version: 0.1.1
4
+ Summary: Synthetic data generation for Databricks — with a built-in notebook UI
5
+ Project-URL: Homepage, https://github.com/darshan-innovation/dash-synthetic
6
+ Author-email: Darshan Shah <darshan.innovation@icloud.com>
7
+ License: Apache-2.0
8
+ Keywords: databricks,delta,pyspark,synthetic-data
9
+ Requires-Python: >=3.9
10
+ Requires-Dist: dash-uis>=0.1.0
11
+ Requires-Dist: ipywidgets>=8.0
12
+ Requires-Dist: numpy>=1.24
13
+ Provides-Extra: dev
14
+ Requires-Dist: hatch; extra == 'dev'
15
+ Requires-Dist: pdoc; extra == 'dev'
16
+ Requires-Dist: pytest; extra == 'dev'
17
+ Requires-Dist: pytest-cov; extra == 'dev'
18
+ Requires-Dist: ruff; extra == 'dev'
19
+ Description-Content-Type: text/markdown
20
+
21
+ # Dashsynthetic — Databricks Library
22
+
23
+ [![CI](https://github.com/darshan-innovation/dash-synthetic/actions/workflows/ci.yml/badge.svg)](https://github.com/darshan-innovation/dash-synthetic/actions)
24
+ [![PyPI](https://img.shields.io/pypi/v/dash-synthetic)](https://pypi.org/project/dash-synthetic/)
25
+ [![License](https://img.shields.io/badge/license-Apache%202.0-blue)](LICENSE)
26
+
27
+ Part of the **[Dashlibs](https://github.com/darshan-innovation)** suite — Databricks libraries built for business users.
28
+
29
+ ## Installation
30
+
31
+ ```bash
32
+ %pip install dash-synthetic
33
+ ```
34
+
35
+ ## Quick Start
36
+
37
+ ```python
38
+ import dashsynthetic
39
+ dashsynthetic.launch() # Opens interactive UI in your Databricks notebook
40
+ ```
41
+
42
+ The UI has two tabs:
43
+ - **Single Table** — profile a source table/DataFrame/SQL query and generate synthetic data from it.
44
+ - **Multi-Table Relationships** — define multiple tables, their primary keys, foreign keys, and
45
+ master data columns (e.g. currency/country codes); the tool figures out the dependency order and
46
+ generates every table with referentially valid foreign keys.
47
+
48
+ ## Python API
49
+
50
+ ```python
51
+ from dashsynthetic import RelationshipGraph, MultiTableGenerator
52
+
53
+ graph = RelationshipGraph()
54
+ graph.add_table("Customer", table="catalog.schema.dim_customer", primary_key="customer_id")
55
+ graph.add_table("Account", table="catalog.schema.fact_account", primary_key="account_id",
56
+ master_data_columns=["currency_code"])
57
+ graph.add_foreign_key("Account", "customer_id", "Customer", "customer_id")
58
+
59
+ gen = MultiTableGenerator(graph)
60
+ gen.configure_table("Customer", n_rows=5000)
61
+ gen.configure_table("Account", n_rows=20000, output_table="catalog.schema.syn_account")
62
+ results = gen.run() # {"Customer": df, "Account": df}, generated in dependency order
63
+ ```
64
+
65
+ ## Part of Dashlibs
66
+
67
+ | Library | Purpose |
68
+ |---|---|
69
+ | dash-dq | Data Quality |
70
+ | dash-synthetic | Synthetic Data Generation |
71
+ | dash-ml | ML Model Monitoring |
72
+ | dash-ingest | Data Ingestion |
73
+ | dash-gov | Data Governance |
74
+ | dash-relate | Ontology & Lineage for AI |
75
+
76
+ ## License
77
+
78
+ Apache 2.0
@@ -0,0 +1,58 @@
1
+ # Dashsynthetic — Databricks Library
2
+
3
+ [![CI](https://github.com/darshan-innovation/dash-synthetic/actions/workflows/ci.yml/badge.svg)](https://github.com/darshan-innovation/dash-synthetic/actions)
4
+ [![PyPI](https://img.shields.io/pypi/v/dash-synthetic)](https://pypi.org/project/dash-synthetic/)
5
+ [![License](https://img.shields.io/badge/license-Apache%202.0-blue)](LICENSE)
6
+
7
+ Part of the **[Dashlibs](https://github.com/darshan-innovation)** suite — Databricks libraries built for business users.
8
+
9
+ ## Installation
10
+
11
+ ```bash
12
+ %pip install dash-synthetic
13
+ ```
14
+
15
+ ## Quick Start
16
+
17
+ ```python
18
+ import dashsynthetic
19
+ dashsynthetic.launch() # Opens interactive UI in your Databricks notebook
20
+ ```
21
+
22
+ The UI has two tabs:
23
+ - **Single Table** — profile a source table/DataFrame/SQL query and generate synthetic data from it.
24
+ - **Multi-Table Relationships** — define multiple tables, their primary keys, foreign keys, and
25
+ master data columns (e.g. currency/country codes); the tool figures out the dependency order and
26
+ generates every table with referentially valid foreign keys.
27
+
28
+ ## Python API
29
+
30
+ ```python
31
+ from dashsynthetic import RelationshipGraph, MultiTableGenerator
32
+
33
+ graph = RelationshipGraph()
34
+ graph.add_table("Customer", table="catalog.schema.dim_customer", primary_key="customer_id")
35
+ graph.add_table("Account", table="catalog.schema.fact_account", primary_key="account_id",
36
+ master_data_columns=["currency_code"])
37
+ graph.add_foreign_key("Account", "customer_id", "Customer", "customer_id")
38
+
39
+ gen = MultiTableGenerator(graph)
40
+ gen.configure_table("Customer", n_rows=5000)
41
+ gen.configure_table("Account", n_rows=20000, output_table="catalog.schema.syn_account")
42
+ results = gen.run() # {"Customer": df, "Account": df}, generated in dependency order
43
+ ```
44
+
45
+ ## Part of Dashlibs
46
+
47
+ | Library | Purpose |
48
+ |---|---|
49
+ | dash-dq | Data Quality |
50
+ | dash-synthetic | Synthetic Data Generation |
51
+ | dash-ml | ML Model Monitoring |
52
+ | dash-ingest | Data Ingestion |
53
+ | dash-gov | Data Governance |
54
+ | dash-relate | Ontology & Lineage for AI |
55
+
56
+ ## License
57
+
58
+ Apache 2.0
@@ -0,0 +1,10 @@
1
+ """
2
+ DashSynthetic — Synthetic data generation for Databricks.
3
+ Launch the UI with dashsynthetic.launch() inside a Databricks notebook.
4
+ """
5
+ from dashsynthetic.generator import MultiTableGenerator, SyntheticGenerator
6
+ from dashsynthetic.relationships import RelationshipGraph
7
+ from dashsynthetic.ui import launch
8
+
9
+ __version__ = "0.1.1"
10
+ __all__ = ["SyntheticGenerator", "MultiTableGenerator", "RelationshipGraph", "launch"]
@@ -0,0 +1,131 @@
1
+ """
2
+ Synthetic data generation engine.
3
+ Uses numpy multivariate normal sampling for correlated columns
4
+ and per-column distribution fitting for marginals.
5
+ """
6
+ from __future__ import annotations
7
+ import numpy as np
8
+ from pyspark.sql import functions as F
9
+
10
+
11
+ def generate(source_df, n_rows: int, preserve_corr: bool,
12
+ preserve_nulls: bool, preserve_distributions: bool,
13
+ force_categorical: list = None, column_value_pools: dict = None):
14
+ """
15
+ force_categorical: column names to sample from real distinct source
16
+ values regardless of dtype — used for "master data" columns (e.g.
17
+ currency/country codes) so synthetic values stay referentially valid.
18
+ column_value_pools: column name -> explicit list of values to sample
19
+ from uniformly — used for foreign-key columns, populated from the
20
+ already-generated parent table's primary key values.
21
+ """
22
+ from pyspark.sql import SparkSession
23
+ spark = SparkSession.getActiveSession()
24
+ force_categorical = set(force_categorical or [])
25
+ column_value_pools = column_value_pools or {}
26
+
27
+ schema = source_df.schema
28
+ numeric_cols = [f.name for f in schema.fields
29
+ if any(t in str(f.dataType)
30
+ for t in ("Int", "Long", "Double", "Float", "Decimal"))
31
+ and f.name not in force_categorical and f.name not in column_value_pools]
32
+ string_cols = [f.name for f in schema.fields
33
+ if "String" in str(f.dataType) and f.name not in column_value_pools]
34
+ bool_cols = [f.name for f in schema.fields if "Boolean" in str(f.dataType)]
35
+ date_cols = [f.name for f in schema.fields
36
+ if any(t in str(f.dataType) for t in ("Date", "Timestamp"))]
37
+ forced_cat_cols = [f.name for f in schema.fields
38
+ if f.name in force_categorical and f.name not in column_value_pools]
39
+
40
+ np.random.seed(42)
41
+ rows = []
42
+
43
+ # Collect stats for numeric columns
44
+ numeric_stats = {}
45
+ if numeric_cols:
46
+ agg_exprs = []
47
+ for c in numeric_cols:
48
+ agg_exprs += [F.mean(c).alias(f"{c}__mean"), F.stddev(c).alias(f"{c}__std")]
49
+ stats_row = source_df.agg(*agg_exprs).collect()[0].asDict()
50
+ for c in numeric_cols:
51
+ mean = stats_row.get(f"{c}__mean") or 0.0
52
+ std = stats_row.get(f"{c}__std") or 1.0
53
+ numeric_stats[c] = (float(mean), float(std))
54
+
55
+ # Collect categorical distributions
56
+ cat_dists = {}
57
+ for c in string_cols + forced_cat_cols:
58
+ total = source_df.count()
59
+ dist = (
60
+ source_df.groupBy(c).count()
61
+ .withColumn("prob", F.col("count") / total)
62
+ .select(c, "prob")
63
+ .collect()
64
+ )
65
+ cat_dists[c] = (
66
+ [r[c] for r in dist],
67
+ [float(r["prob"]) for r in dist],
68
+ )
69
+
70
+ # Null rates
71
+ null_rates = {}
72
+ if preserve_nulls:
73
+ total = source_df.count()
74
+ for f in schema.fields:
75
+ nc = source_df.filter(F.col(f.name).isNull()).count()
76
+ null_rates[f.name] = nc / total if total > 0 else 0.0
77
+
78
+ # Generate numeric data (correlated if requested)
79
+ if numeric_cols and preserve_corr:
80
+ pandas_df = source_df.select(numeric_cols).dropna().toPandas()
81
+ cov = pandas_df.cov().values
82
+ means = pandas_df.mean().values
83
+ try:
84
+ samples = np.random.multivariate_normal(means, cov, n_rows)
85
+ except np.linalg.LinAlgError:
86
+ samples = np.column_stack([
87
+ np.random.normal(numeric_stats[c][0], numeric_stats[c][1], n_rows)
88
+ for c in numeric_cols
89
+ ])
90
+ numeric_samples = {c: samples[:, i] for i, c in enumerate(numeric_cols)}
91
+ elif numeric_cols:
92
+ numeric_samples = {
93
+ c: np.random.normal(numeric_stats[c][0], max(numeric_stats[c][1], 1e-6), n_rows)
94
+ for c in numeric_cols
95
+ }
96
+ else:
97
+ numeric_samples = {}
98
+
99
+ # Build rows
100
+ for i in range(n_rows):
101
+ row = {}
102
+ for f in schema.fields:
103
+ c = f.name
104
+ dtype = str(f.dataType)
105
+ null_roll = np.random.random() < null_rates.get(c, 0.0)
106
+
107
+ if c in column_value_pools:
108
+ pool = column_value_pools[c]
109
+ row[c] = pool[int(np.random.randint(0, len(pool)))] if pool else None
110
+ elif null_roll and preserve_nulls:
111
+ row[c] = None
112
+ elif c in numeric_samples:
113
+ val = float(numeric_samples[c][i])
114
+ if "Int" in dtype or "Long" in dtype:
115
+ row[c] = int(round(val))
116
+ else:
117
+ row[c] = round(val, 4)
118
+ elif c in cat_dists:
119
+ values, probs = cat_dists[c]
120
+ probs_arr = np.array(probs)
121
+ probs_arr = probs_arr / probs_arr.sum()
122
+ row[c] = np.random.choice(values, p=probs_arr)
123
+ elif c in bool_cols:
124
+ row[c] = bool(np.random.random() > 0.5)
125
+ elif c in date_cols:
126
+ row[c] = None # dates handled separately in practice
127
+ else:
128
+ row[c] = None
129
+ rows.append(row)
130
+
131
+ return spark.createDataFrame(rows, schema)