ingest: drop all unmapped columns at once, which is much faster

opencybersecurityalliance · Oct 30, 2023 · fc05669 · fc05669
1 parent 797f1cd
commit fc05669
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/firepit/aio/ingest.py b/firepit/aio/ingest.py
@@ -247,6 +247,9 @@ def translate(
     # columns we need to "unwrap"
     unwrap = set()
 
+    # "Unmapped" columns that we need to drop
+    unmapped = []
+
     logger.debug('columns: %s', cols)
     for col in cols:
         logger.debug('column: %s', col)
@@ -298,7 +301,10 @@ def translate(
         else:
             # Drop unmapped columns
             logger.debug('DROP unmapped column "%s"', col)
-            df = df.drop(col, axis=1)
+            unmapped.append(col)
+
+    # Drop any columns that weren't mapped
+    df = df.drop(unmapped, axis=1)
 
     # Run transformers
     for txf_col, txf_name in txf_cols.items():