Something went wrong. Try again.
A training to familiarize TIDE analysts with parquet datasets and python
Something went wrong. Try again.
13 kB · 349 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350# pattern: Mixed (unavoidable) -- generates synthetic data using I/O (file writes) and random module"""Generate synthetic SCDM v8.2.2 Parquet files for the Parquet + Polars workshop.
Produces three Parquet files in data/: - demographic.parquet (10,000 patients, one row per PatID) - enrollment.parquet (~15,000 enrollment spans) - dispensing.parquet (~200,000 pharmacy dispensings)
All data is deterministic: random.seed(42) ensures identical output across runs.Column names, dtypes, and value sets match the SCDM v8.2.2 specification."""
import randomimport datetimefrom pathlib import Path
import polars as pl
SEED = 42random.seed(SEED)
NUM_PATIENTS = 10_000NO_DRUG_COVERAGE_START = 9_751 # patients 9751-10000: all spans DrugCov='N'
DATA_DIR = Path(__file__).parent
# ── SCDM value-set constants ──────────────────────────────────────────────
SEX_VALUES = "AFMU"SEX_WEIGHTS = [52, 47, 0.5, 0.5]
HISPANIC_VALUES = "NYU"HISPANIC_WEIGHTS = [80, 15, 5]
RACE_VALUES = "012345M"RACE_WEIGHTS = [5, 3, 10, 18, 2, 60, 2]
MEDCOV_VALUES = "YNUA"MEDCOV_WEIGHTS = [85, 8, 4, 3]
DRUGCOV_VALUES = "YNU"DRUGCOV_WEIGHTS = [90, 8, 2]
CHART_VALUES = "YN"CHART_WEIGHTS = [75, 25]
PLANTYPE_VALUES = "FHEPSMOU"PLANTYPE_WEIGHTS = [20, 25, 5, 20, 10, 8, 7, 5]
PAYERTYPE_VALUES = ["CO", "MC", "MD", "OG", "OT", "UN"]PAYERTYPE_WEIGHTS = [50, 25, 15, 3, 4, 3]
RX_CODETYPE_VALUE = "ND" # all generated Rx values are NDCs
# ── Polars schema definitions (SCDM type mapping) ──────────────────────────
DEMOGRAPHIC_SCHEMA: dict[str, pl.DataType] = { "PatID": pl.Int64, "Birth_Date": pl.Date, "Sex": pl.Utf8, "Hispanic": pl.Utf8, "Race": pl.Utf8, "PostalCode": pl.Utf8, "PostalCode_Date": pl.Date, "ImputedRace": pl.Utf8, "ImputedHispanic": pl.Utf8,}
ENROLLMENT_SCHEMA: dict[str, pl.DataType] = { "PatID": pl.Int64, "Enr_Start": pl.Date, "Enr_End": pl.Date, "MedCov": pl.Utf8, "DrugCov": pl.Utf8, "Chart": pl.Utf8, "PlanType": pl.Utf8, "PayerType": pl.Utf8,}
DISPENSING_SCHEMA: dict[str, pl.DataType] = { "PatID": pl.Int64, "ProviderID": pl.Int64, "RxDate": pl.Date, "Rx": pl.Utf8, "Rx_CodeType": pl.Utf8, "RxSup": pl.Int32, "RxAmt": pl.Int32,}
# ── Date helper ────────────────────────────────────────────────────────────
DATE_MIN = datetime.date(2018, 1, 1)DATE_MAX = datetime.date(2023, 12, 31)BIRTH_MIN = datetime.date(1940, 1, 1)BIRTH_MAX = datetime.date(2010, 12, 31)ENR_START_MIN = datetime.date(2018, 1, 1)ENR_START_MAX = datetime.date(2023, 6, 1)ENR_END_MAX = datetime.date(2023, 12, 31)POSTAL_DATE_MIN = datetime.date(2018, 1, 1)POSTAL_DATE_MAX = datetime.date(2023, 6, 30)
def random_date(start: datetime.date, end: datetime.date) -> datetime.date: """Return a random date between start and end (inclusive).""" delta = (end - start).days return start + datetime.timedelta(days=random.randint(0, delta))
def weighted_choice(values: str | list[str], weights: list[float]) -> str: """Pick one value from values using proportional weights.""" return random.choices(values, weights=weights, k=1)[0]
# ── Demographic generation ────────────────────────────────────────────────
def generate_demographic() -> list[dict]: """Generate demographic rows: one per patient, PatID 1..10000.""" rows: list[dict] = [] for pat_id in range(1, NUM_PATIENTS + 1): row = { "PatID": pat_id, "Birth_Date": random_date(BIRTH_MIN, BIRTH_MAX), "Sex": weighted_choice(SEX_VALUES, SEX_WEIGHTS), "Hispanic": weighted_choice(HISPANIC_VALUES, HISPANIC_WEIGHTS), "Race": weighted_choice(RACE_VALUES, RACE_WEIGHTS), "PostalCode": f"{random.randint(1, 99950):05d}" if random.random() > 0.02 else None, "PostalCode_Date": ( random_date(POSTAL_DATE_MIN, POSTAL_DATE_MAX) if random.random() > 0.02 else None ), "ImputedRace": ( weighted_choice(RACE_VALUES, RACE_WEIGHTS) if random.random() > 0.30 else None ), "ImputedHispanic": ( weighted_choice(HISPANIC_VALUES, HISPANIC_WEIGHTS) if random.random() > 0.30 else None ), } rows.append(row) return rows
# ── Enrollment generation ─────────────────────────────────────────────────
def generate_enrollment() -> list[dict]: """Generate enrollment spans: 1-3 per patient, ~15,000 total.""" rows: list[dict] = [] for pat_id in range(1, NUM_PATIENTS + 1): no_drug_coverage = pat_id >= NO_DRUG_COVERAGE_START n_spans = random.choices([1, 2, 3], weights=[60, 30, 10])[0]
prev_end: datetime.date | None = None for span_idx in range(n_spans): if prev_end is not None: gap = random.randint(1, 60) start = prev_end + datetime.timedelta(days=gap) if start > ENR_START_MAX: break else: start = random_date(ENR_START_MIN, ENR_START_MAX)
# span length: 3-12 months months = random.randint(3, 12) end = start + datetime.timedelta(days=months * 30 + random.randint(0, 10)) if end > ENR_END_MAX: end = ENR_END_MAX if end <= start: end = start + datetime.timedelta(days=90)
if no_drug_coverage: drug_cov = "N" else: # For drug-covered patients, ensure at least the first span # has DrugCov='Y' so they have drug coverage for their # dispensings to match against. drug_cov = ( "Y" if span_idx == 0 else weighted_choice(DRUGCOV_VALUES, DRUGCOV_WEIGHTS) )
row = { "PatID": pat_id, "Enr_Start": start, "Enr_End": end, "MedCov": weighted_choice(MEDCOV_VALUES, MEDCOV_WEIGHTS), "DrugCov": drug_cov, "Chart": weighted_choice(CHART_VALUES, CHART_WEIGHTS), "PlanType": weighted_choice(PLANTYPE_VALUES, PLANTYPE_WEIGHTS), "PayerType": weighted_choice(PAYERTYPE_VALUES, PAYERTYPE_WEIGHTS), } rows.append(row) prev_end = end return rows
# ── Dispensing generation ─────────────────────────────────────────────────
def _fabricated_ndc(n: int) -> str: """Produce a fabricated 11-character NDC in 5-4-2 format.""" labeler = n % 100000 product = (n // 100000) % 10000 package = (n // 1_000_000_000) % 100 return f"{labeler:05d}{product:04d}{package:02d}"
def _date_in_spans( d: datetime.date, spans: list[tuple[datetime.date, datetime.date]]) -> bool: """Return True if d falls within any of the given (start, end) spans.""" return any(s <= d <= e for s, e in spans)
def generate_dispensing(enrollment: list[dict]) -> list[dict]: """Generate ~20 dispensings per patient, ~200,000 total.
Coverage design (~5% uncovered): - Patients 9751-10000 have all spans DrugCov='N': all their ~5,000 dispensings are uncovered (no drug-covered span to match). - ~2.5% of dispensings for drug-covered patients have RxDate set outside all enrollment spans (date mismatch): ~5,000 uncovered. - Remaining ~190,000 dispensings have RxDate within a drug-covered enrollment span and are covered. """ # Build per-patient enrollment span lookup pat_spans: dict[int, list[tuple[datetime.date, datetime.date, str]]] = {} for enr in enrollment: pat_spans.setdefault(enr["PatID"], []).append( (enr["Enr_Start"], enr["Enr_End"], enr["DrugCov"]) ) for spans in pat_spans.values(): spans.sort(key=lambda s: s[0])
rows: list[dict] = [] ndc_counter = 1
for pat_id in range(1, NUM_PATIENTS + 1): n_disp = random.randint(15, 25) # ~20 per patient average spans = pat_spans.get(pat_id, [])
# Separate drug-covered spans from all spans drug_cov_spans = [(s, e) for s, e, dc in spans if dc == "Y"] all_spans = [(s, e) for s, e, dc in spans] has_drug_cov = len(drug_cov_spans) > 0
for _ in range(n_disp): if has_drug_cov and random.random() < 0.025: # ~2.5% of drug-covered patient dispensings: date outside # all enrollment spans but still within 2018-2023. rx_date = random_date(DATE_MIN, DATE_MAX) for _ in range(10): if not _date_in_spans(rx_date, all_spans): break rx_date = random_date(DATE_MIN, DATE_MAX) else: # Fallback: nudge just outside first or last span if random.random() < 0.5 and spans[0][0] > DATE_MIN: rx_date = spans[0][0] - datetime.timedelta(days=1) elif spans[-1][1] < DATE_MAX: rx_date = spans[-1][1] + datetime.timedelta(days=1) else: rx_date = random_date(DATE_MIN, DATE_MAX) elif has_drug_cov: # Pick a date within a random drug-covered span span_start, span_end = random.choice(drug_cov_spans) rx_date = random_date(span_start, span_end) else: # No drug coverage: any date in 2018-2023 rx_date = random_date(DATE_MIN, DATE_MAX)
# ProviderID: ~10% null (representing SAS .U special missing) # Note: Parquet/Polars have no SAS special missing values. # We represent .U as null (None) here. provider_id = ( random.randint(10000, 99999) if random.random() > 0.10 else None )
row = { "PatID": pat_id, "ProviderID": provider_id, "RxDate": rx_date, "Rx": _fabricated_ndc(ndc_counter), "Rx_CodeType": RX_CODETYPE_VALUE, "RxSup": ( random.choices( [7, 14, 30, 60, 90], weights=[10, 15, 50, 15, 10], )[0] if random.random() > 0.02 else None ), "RxAmt": ( random.choices( [30, 60, 90, 120, 180, 15, 45], weights=[30, 25, 15, 10, 5, 10, 5], )[0] if random.random() > 0.02 else None ), } rows.append(row) ndc_counter += 1
return rows
# ── Main ───────────────────────────────────────────────────────────────────
def build_and_write() -> None: """Generate all three tables, sort per SCDM spec, write Parquet files.""" # Generate demo_data = generate_demographic() enr_data = generate_enrollment() disp_data = generate_dispensing(enr_data)
# Build DataFrames with explicit schemas df_demo = pl.DataFrame(demo_data, schema=DEMOGRAPHIC_SCHEMA).sort("PatID") df_enr = pl.DataFrame(enr_data, schema=ENROLLMENT_SCHEMA).sort( "PatID", "Enr_Start", "Enr_End", "MedCov", "DrugCov", "Chart" ) df_disp = pl.DataFrame(disp_data, schema=DISPENSING_SCHEMA).sort("PatID", "RxDate")
# Write Parquet df_demo.write_parquet(DATA_DIR / "demographic.parquet") df_enr.write_parquet(DATA_DIR / "enrollment.parquet") df_disp.write_parquet(DATA_DIR / "dispensing.parquet")
# Print summary print("Data generation complete.") print(f" demographic.parquet : {df_demo.shape[0]:>8,} rows") print(f" enrollment.parquet : {df_enr.shape[0]:>8,} rows") print(f" dispensing.parquet : {df_disp.shape[0]:>8,} rows") print() print("Schemas:") print(f" Demographic: {df_demo.schema}") print(f" Enrollment: {df_enr.schema}") print(f" Dispensing: {df_disp.schema}")
if __name__ == "__main__": build_and_write()