60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
import os
|
|
import yaml
|
|
|
|
import awkward as ak
|
|
import pandas as pd
|
|
import numpy as np
|
|
import tqdm
|
|
|
|
scorecard_dir = "data/scorecard"
|
|
scorecard_dir = os.path.join(scorecard_dir, os.listdir(scorecard_dir)[0])
|
|
|
|
print("Loading data.yaml")
|
|
with open(os.path.join(scorecard_dir, 'data.yaml'), 'r') as file:
|
|
data = yaml.safe_load(file)
|
|
|
|
print("Loading CSVs to dataframes")
|
|
files = [f'MERGED{i}_{(i + 1) % 100:02}_PP.csv' for i in tqdm.trange(1996, 2024)]
|
|
dataframes = [pd.read_csv(os.path.join(scorecard_dir, file)) for file in tqdm.tqdm(files)]
|
|
|
|
print("Appending extra rows where needed")
|
|
unit_ids = np.unique(np.hstack([frame.UNITID.to_numpy() for frame in tqdm.tqdm(dataframes)]))
|
|
for i, frame in tqdm.tqdm(enumerate(dataframes)):
|
|
new_rows = pd.DataFrame({"UNITID": unit_ids[~np.isin(unit_ids, frame.UNITID)]})
|
|
dataframes[i] = pd.concat([frame, new_rows]).sort_values(by=["UNITID", "OPEID"])
|
|
|
|
print("Converting to Results Array")
|
|
result = {}
|
|
for key, sec in tqdm.tqdm(data['dictionary'].items()):
|
|
if 'calculate' in sec:
|
|
continue
|
|
|
|
data_key = sec['source']
|
|
if data_key not in dataframes[0]:
|
|
continue
|
|
|
|
parts = key.split('.')
|
|
section = result
|
|
for i in range(len(parts) - 1):
|
|
part = parts[i]
|
|
if part not in section:
|
|
section[part] = {}
|
|
section = section[part]
|
|
|
|
obj = np.vstack([frame[data_key] for frame in dataframes]).T
|
|
for frame in dataframes:
|
|
del frame[data_key] # Memory cleanup
|
|
|
|
if obj.dtype == object:
|
|
obj = obj.astype(str)
|
|
section[parts[-1]] = obj
|
|
|
|
print("Cleanup: Deleting Dataframes from Memory")
|
|
del dataframes # Memory cleanup
|
|
|
|
print("Converting to Awkward Array")
|
|
a = ak.Array(result)
|
|
del result # Memory cleanup
|
|
|
|
print("Writing to Disk")
|
|
ak.to_parquet(a, os.path.join(scorecard_dir, "merged.parquet")) |