hangshuo652
34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
import struct, json
|
|
from pathlib import Path
|
|
from data.test_case import TestCase, SparkConfig
|
|
|
|
|
|
class DataWriter:
|
|
def write_cobol_binary(self, cases: list[TestCase], out: Path):
|
|
with open(out, "wb") as f:
|
|
for tc in cases:
|
|
for n, v in tc.fields.items():
|
|
if isinstance(v, int):
|
|
f.write(struct.pack(">q", v))
|
|
elif isinstance(v, float):
|
|
f.write(struct.pack(">d", v))
|
|
elif isinstance(v, str):
|
|
f.write(v.encode("ascii", errors="replace").ljust(10, b" ")[:10])
|
|
|
|
def write_spark_json(self, cases: list[TestCase], cfg: SparkConfig, d: Path):
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
base = cases[0].fields if cases else {}
|
|
recs = []
|
|
for i in range(cfg.num_records):
|
|
r = dict(base)
|
|
if cfg.key_field in r:
|
|
r[cfg.key_field] = f"{r[cfg.key_field]}-{i:04d}"
|
|
recs.append(r)
|
|
(d / "part-00000.json").write_text("\n".join(json.dumps(r) for r in recs))
|
|
|
|
def write_native_json(self, cases: list[TestCase], out: Path):
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(out, "w") as f:
|
|
for tc in cases:
|
|
f.write(json.dumps(tc.fields) + "\n")
|