Try building monthly count files

personal/diagnostic/count_subreddits_multiprocess.py

@@ -29,9 +29,10 @@ log.addHandler(log_file_handler)
 
 # convenience object used to pass status information between processes
 class FileConfig:
-	def __init__(self, input_path, output_path=None, complete=False, lines_processed=0, error_lines=0):
+	def __init__(self, input_path, output_path=None, complete=False, lines_processed=0, error_lines=0, count_file_path=None):
 		self.input_path = input_path
 		self.output_path = output_path
+		self.count_file_path = count_file_path
 		self.file_size = os.stat(input_path).st_size
 		self.complete = complete
 		self.bytes_processed = self.file_size if complete else 0
@@ -60,16 +61,17 @@ class Queue:
 
 # save file information and progress to a json file
 # we don't want to save the whole FileConfig object, since some info resets if we restart
-def save_file_list(input_files, working_folder, status_json, script_type):
+def save_file_list(input_files, working_folder, status_json, script_type, stage):
 	if not os.path.exists(working_folder):
 		os.makedirs(working_folder)
 	simple_file_list = []
 	for file in input_files:
-		simple_file_list.append([file.input_path, file.output_path, file.complete, file.lines_processed, file.error_lines])
+		simple_file_list.append([file.input_path, file.output_path, file.complete, file.lines_processed, file.error_lines, file.monthly_count_file])
 	with open(status_json, 'w') as status_json_file:
 		output_dict = {
 			"files": simple_file_list,
 			"type": script_type,
+			"stage": stage,
 		}
 		status_json_file.write(json.dumps(output_dict, indent=4))
 
@@ -82,11 +84,11 @@ def load_file_list(status_json):
 			input_files = []
 			for simple_file in output_dict["files"]:
 				input_files.append(
-					FileConfig(simple_file[0], simple_file[1], simple_file[2], simple_file[3], simple_file[4])
+					FileConfig(simple_file[0], simple_file[1], simple_file[2], simple_file[3], simple_file[4], simple_file[5] if len(simple_file) > 5 else None)
 				)
-			return input_files, output_dict["type"]
+			return input_files, output_dict["type"], output_dict["stage"]
 	else:
-		return None, None
+		return None, None, "count"
 
 
 # recursively decompress and decode a chunk of bytes. If there's a decode error then read another chunk and try with that, up to a limit of max_window_size bytes
@@ -158,6 +160,7 @@ if __name__ == '__main__':
 	parser.add_argument("input", help="The input folder to recursively read files from")
 	parser.add_argument("--output", help="Name of the output file", default="field_counts")
 	parser.add_argument("--working", help="The folder to store temporary files in", default="pushshift_working")
+	parser.add_argument("--monthly_count_folder", help="The folder to store monthly count files in", default="pushshift_counts")
 	parser.add_argument("--field", help="Which field to count", default="subreddit")
 	parser.add_argument("--min_count", help="Dont write any counts below this number", default=1000, type=int)
 	parser.add_argument("--processes", help="Number of processes to use", default=10, type=int)
@@ -183,12 +186,13 @@ if __name__ == '__main__':
 	multiprocessing.set_start_method('spawn')
 	queue = multiprocessing.Manager().Queue()
 	status_json = os.path.join(args.working, "status.json")
-	input_files, saved_type = load_file_list(status_json)
+	input_files, saved_type, stage = load_file_list(status_json)
 
 	if saved_type and saved_type != script_type:
 		log.warning(f"Script type doesn't match type from json file. Delete working folder")
 		sys.exit(0)
 
+	if stage == "count":
 		# if the file list wasn't loaded from the json, this is the first run, find what files we need to process
 		if input_files is None:
 			input_files = []
@@ -200,7 +204,7 @@ if __name__ == '__main__':
 						output_path = os.path.join(args.working, file_name[:-4])
 						input_files.append(FileConfig(input_path, output_path=output_path))
 
-		save_file_list(input_files, args.working, status_json, script_type)
+			save_file_list(input_files, args.working, status_json, script_type, "count")
 		else:
 			log.info(f"Existing input file was read, if this is not correct you should delete the {args.working} folder and run this script again")
 
@@ -277,11 +281,15 @@ if __name__ == '__main__':
 						f"{(str(days_left) + 'd ' if days_left > 0 else '')}{hours_left - (days_left * 24)}:{minutes_left - (hours_left * 60):02}:{seconds_left - (minutes_left * 60):02} remaining")
 
 		log.info(f"{total_lines_processed:,}, {total_lines_errored} errored : {(total_bytes_processed / (2**30)):.2f} gb, {(total_bytes_processed / total_bytes) * 100:.0f}% : {files_processed}/{len(input_files)}")
+		stage = "sum"
+		save_file_list(input_files, args.working, status_json, script_type, stage)
 
-	working_file_paths = []
+	if stage == "sum":
+		#working_file_paths = []
 		count_incomplete = 0
 		# build a list of output files to combine
-	for file in sorted(input_files, key=lambda item: os.path.split(item.output_path)[1]):
+		input_files = sorted(input_files, key=lambda item: os.path.split(item.output_path)[1])
+		for file in input_files:
 			if not file.complete:
 				if file.error_message is not None:
 					log.info(f"File {file.input_path} errored {file.error_message}")
@@ -294,26 +302,44 @@ if __name__ == '__main__':
 						f"File {file.input_path} has {file.error_lines:,} errored lines out of {file.lines_processed:,}, "
 						f"{(file.error_lines / file.lines_processed) * (args.error_rate * 0.01):.2f}% which is above the limit of {args.error_rate}%")
 					count_incomplete += 1
-			elif file.output_path is not None:
-				if os.path.exists(file.output_path):
-					working_file_paths.append(file.output_path)
 
 		if count_incomplete > 0:
 			log.info(f"{count_incomplete} files were not completed, errored or don't exist, something went wrong. Aborting")
 			sys.exit()
 
-	log.info(f"Processing complete, combining {len(working_file_paths)} result files")
+		log.info(f"Processing complete, combining {len(input_files)} result files")
 
 		input_lines = 0
-	files_combined = 0
-	field_counts = defaultdict(int)
-	for working_file_path in working_file_paths:
-		files_combined += 1
-		log.info(f"Reading {files_combined}/{len(working_file_paths)} : {input_lines:,} : {len(field_counts):,} : {os.path.split(working_file_path)[1]}")
-		with open(working_file_path, 'r') as input_file:
+		files_counted = 0
+		monthly_count_folder_paths = []
+		for file in input_files:
+			files_counted += 1
+			if not os.path.exists(file.output_path):
+				log.info(f"Output file {file.output_path} does not exist, skipping")
+				continue
+			monthly_counts = defaultdict(int)
+			log.info(f"Reading {files_counted}/{len(input_files)} : {input_lines:,} : {os.path.split(file.output_path)[1]}")
+			with open(file.output_path, 'r') as input_file:
 				for line in input_file:
 					input_lines += 1
-				field_counts[line.strip()] += 1
+					monthly_counts[line.strip()] += 1
+
+			file.monthly_count_file = os.path.join(args.monthly_count_folder, os.path.basename(file.output_path))
+			with open(file.monthly_count_file, 'w') as output_handle:
+				for field, count in sorted(monthly_counts.items(), key=lambda item: item[1], reverse=True):
+					output_handle.write(f"{field}	{count}\n")
+
+		log.info(f"Finished combining files into monthlies, {input_lines:,} lines read. Combining into result output")
+		stage = "agg"
+		save_file_list(input_files, args.working, status_json, script_type, stage)
+
+	if stage == "agg":
+		field_counts = defaultdict(int)
+		for file in input_files:
+			with open(file.monthly_count_file, 'r') as input_handle:
+				for line in input_handle:
+					field, count = line.strip().split("\t")
+					field_counts[field] = count
 
 		sorted_counts = sorted(field_counts.items(), key=lambda item: item[1], reverse=True)
 
@@ -324,4 +350,4 @@ if __name__ == '__main__':
 					output_counts += 1
 					output_handle.write(f"{field}	{count}\n")
 
-	log.info(f"Finished combining files, {input_lines:,} lines read, {output_counts:,} field counts written")
+		log.info(f"Finished combining files, {output_counts:,} field counts written")