Initial implentation of process_month.py

This commit is contained in:
Watchful1 2024-05-17 21:35:00 -07:00
parent b54a2483dc
commit fa5f6316fb
10 changed files with 703 additions and 369 deletions

View file

@ -8,39 +8,39 @@ from datetime import datetime
import json
import argparse
log = discord_logging.init_logging()
log = discord_logging.get_logger(init=True)
import utils
NEWLINE_ENCODED = "\n".encode('utf-8')
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Take a zst_blocks file and split it by minute chunks")
parser.add_argument('--input', help='Input file', required=True)
parser.add_argument('--output', help='Output folder', required=True)
args = parser.parse_args()
def split_by_minutes(input_file, output_file):
file_type = "comments" if "RC" in input_file else "submissions"
# input_file = r"\\MYCLOUDPR4100\Public\reddit\blocks\RS_2023-10.zst_blocks"
# output_folder = r"\\MYCLOUDPR4100\Public\ingest\download"
file_type = "comments" if "RC" in args.input else "submissions"
log.info(f"Input file: {args.input}")
log.info(f"Output folder: {args.output}")
log.info(f"{file_type}: Input file: {input_file}")
log.info(f"{file_type}: Output folder: {output_file}")
previous_minute, output_handle, created_utc = None, None, None
count_objects, count_minute = 0, 0
for obj in utils.read_obj_zst_blocks(args.input):
if input_file.endswith(".zst"):
reader = utils.read_obj_zst(input_file)
elif input_file.endswith(".zst_blocks"):
reader = utils.read_obj_zst_blocks(input_file)
else:
log.error(f"{file_type}: Unsupported file type: {input_file}")
return
for obj in reader:
created_utc = datetime.utcfromtimestamp(obj["created_utc"])
current_minute = created_utc.replace(second=0)
if previous_minute is None or current_minute > previous_minute:
log.info(f"{created_utc.strftime('%y-%m-%d_%H-%M')}: {count_objects:,} : {count_minute: ,}")
log.info(f"{file_type}: {created_utc.strftime('%y-%m-%d_%H-%M')}: {count_objects:,} : {count_minute: ,}")
previous_minute = current_minute
count_minute = 0
if output_handle is not None:
output_handle.close()
output_path = os.path.join(args.output, file_type, created_utc.strftime('%y-%m-%d'))
output_path = os.path.join(output_file, file_type, created_utc.strftime('%y-%m-%d'))
if not os.path.exists(output_path):
os.makedirs(output_path)
output_path = os.path.join(output_path, f"{('RC' if file_type == 'comments' else 'RS')}_{created_utc.strftime('%y-%m-%d_%H-%M')}.zst")
@ -51,6 +51,18 @@ if __name__ == "__main__":
output_handle.write(json.dumps(obj, sort_keys=True).encode('utf-8'))
output_handle.write(NEWLINE_ENCODED)
log.info(f"{created_utc.strftime('%y-%m-%d_%H-%M')}: {count_objects:,} : {count_minute: ,}")
if created_utc is None:
log.error(f"{file_type}: {input_file} appears to be empty")
sys.exit(1)
log.info(f"{file_type}: {created_utc.strftime('%y-%m-%d_%H-%M')}: {count_objects:,} : {count_minute: ,}")
if output_handle is not None:
output_handle.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Take a zst_blocks file and split it by minute chunks")
parser.add_argument('--input', help='Input file', required=True)
parser.add_argument('--output', help='Output folder', required=True)
args = parser.parse_args()
split_by_minutes(args.input, args.output)