optimize reading records

dionhaefner · May 14, 2020 · 0abda56 · 0abda56
1 parent 223bf34
commit 0abda56
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/fowd/processing.py b/fowd/processing.py
@@ -63,16 +63,13 @@ def read_pickled_record_chunks(input_file):
 
 def read_pickle_outfile(input_file):
     reader = read_pickled_record_chunks(input_file)
-    try:
-        out = next(reader)
-    except StopIteration:
-        return {}
 
+    records = collections.defaultdict(list)
     for row in reader:
-        assert set(out.keys()) == set(row.keys())
         for key, val in row.items():
-            out[key] = np.concatenate((out[key], val))
+            records[key].append(val)
 
+    out = {key: np.concatenate(val) for key, val in records.items()}
     return out