62 lines
2.3 KiB
Python
Executable File
62 lines
2.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Generates the sharded order CSVs this module partitions over.
|
|
|
|
Usage:
|
|
generate-shards.py <output-dir> <total-rows> <num-shards> [--corrupt-shard N] [--seed S]
|
|
|
|
Each shard is orders-shard-NN.csv with header "orderId,customerId,amountCents,region".
|
|
Row counts are split as evenly as possible across shards (the last shard absorbs the remainder).
|
|
--corrupt-shard N replaces one row near the middle of shard N (0-indexed) with a non-numeric
|
|
amountCents field, so FlatFileItemReader's FieldSetMapper throws NumberFormatException wrapped in
|
|
FlatFileParseException when that shard is read -- a real, reproducible parse failure, not a
|
|
simulated one.
|
|
"""
|
|
import argparse
|
|
import os
|
|
import random
|
|
import sys
|
|
|
|
REGIONS = ["NORTH", "SOUTH", "EAST", "WEST", "CENTRAL"]
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("output_dir")
|
|
ap.add_argument("total_rows", type=int)
|
|
ap.add_argument("num_shards", type=int)
|
|
ap.add_argument("--corrupt-shard", type=int, default=-1)
|
|
ap.add_argument("--seed", type=int, default=42)
|
|
args = ap.parse_args()
|
|
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
rnd = random.Random(args.seed)
|
|
|
|
base = args.total_rows // args.num_shards
|
|
counts = [base] * args.num_shards
|
|
counts[-1] += args.total_rows - base * args.num_shards
|
|
|
|
order_id = 1
|
|
for shard_idx, count in enumerate(counts):
|
|
path = os.path.join(args.output_dir, f"orders-shard-{shard_idx:02d}.csv")
|
|
corrupt_at = -1
|
|
if shard_idx == args.corrupt_shard:
|
|
corrupt_at = count // 2
|
|
with open(path, "w", newline="") as f:
|
|
f.write("orderId,customerId,amountCents,region\n")
|
|
for i in range(count):
|
|
customer_id = rnd.randint(1, 2_000_000)
|
|
amount_cents = rnd.randint(500, 9_999_999)
|
|
region = REGIONS[rnd.randint(0, len(REGIONS) - 1)]
|
|
if i == corrupt_at:
|
|
f.write(f"{order_id},{customer_id},NOT_A_NUMBER,{region}\n")
|
|
else:
|
|
f.write(f"{order_id},{customer_id},{amount_cents},{region}\n")
|
|
order_id += 1
|
|
print(f"wrote {path}: {count} rows" + (" (1 corrupt row)" if corrupt_at >= 0 else ""))
|
|
|
|
print(f"total: {order_id - 1} rows across {args.num_shards} shards in {args.output_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|