File size: 1,605 Bytes
d561d1e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """
Add a cropped-spectrum column (channels 700:1018, 318 channels) to the
unknown-data parquet, processed row-group by row-group to bound memory.
Input : Data Parquet/unknown_data_after_noise_filter.parquet
Output: Data Parquet/unknown_data_after_noise_filter_cropped.parquet
"""
import os
import pyarrow.compute as pc
import pyarrow.parquet as pq
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SRC = os.path.join(ROOT, 'Data Parquet', 'unknown_data_after_noise_filter.parquet')
DST = os.path.join(ROOT, 'Data Parquet', 'unknown_data_after_noise_filter_cropped.parquet')
START, END = 700, 1018 # channels 700..1017 (inclusive) -> 318 channels
def main():
pf = pq.ParquetFile(SRC)
n_rgs = pf.metadata.num_row_groups
print(f'{SRC}\n rows: {pf.metadata.num_rows:,} row groups: {n_rgs}')
writer = None
for i in range(n_rgs):
t = pf.read_row_group(i)
cropped = pc.list_slice(t.column('spectrum'), START, END)
t = t.append_column('spectrum_cropped', cropped)
if writer is None:
writer = pq.ParquetWriter(DST, t.schema)
writer.write_table(t)
if (i + 1) % 4 == 0 or i + 1 == n_rgs:
print(f' row group {i + 1}/{n_rgs}')
writer.close()
print(f'Wrote {DST}')
# Verify
out = pq.ParquetFile(DST)
print(f'Output rows: {out.metadata.num_rows:,}')
sample = out.read_row_group(0, columns=['spectrum_cropped']).slice(0, 3).to_pydict()['spectrum_cropped']
print('sample cropped lengths:', [len(s) for s in sample])
if __name__ == '__main__':
main() |