chunk data gen

generate_data.py

@@ -105,23 +105,25 @@ def calculate_distance(pair):
     distance = get_distance(city1["name"], city2["name"])
     return city1["name"], city2["name"], distance
 
-
 def main():
     cities = list(us_cities.values())
     print(f"Num cities: {len(cities)}")
     city_combinations = list(itertools.combinations(cities, 2))
+    chunk_size = 800  # adjust this as needed
+    num_chunks = len(city_combinations) // chunk_size + 1
+    output_file = args.output_file
 
-    with open(args.output_file, "w", newline="") as csvfile:
+    with open(output_file, "w", newline="") as csvfile:
         fieldnames = ["city_from", "city_to", "distance"]
         writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
         writer.writeheader()
 
         try:
             executor = concurrent.futures.ProcessPoolExecutor(max_workers=args.workers)
-            # results = executor.map(calculate_distance, city_combinations)
+            for i in range(num_chunks):
+                chunk = city_combinations[i * chunk_size : (i + 1) * chunk_size]
                 futures = {
-                executor.submit(calculate_distance, pair): pair
+                    executor.submit(calculate_distance, pair): pair for pair in chunk
-                for pair in city_combinations
                 }
                 for future in as_completed(futures):
                     city_from, city_to, distance = future.result()
@@ -133,6 +135,7 @@ def main():
                                 "distance": distance,
                             }
                         )
+                        csvfile.flush()  # write to disk immediately
         except KeyboardInterrupt:
             print("Interrupted. Terminating processes...")
             executor.shutdown(wait=False)