⚡ Optimize pandas numeric column conversion

benchmark.py

@@ -0,0 +1,46 @@
+import pandas as pd
+import numpy as np
+import timeit
+
+# Generate sample data
+np.random.seed(0)
+n_rows = 100000
+
+data = {
+    'B19013_001E': np.random.randint(1, 100000, n_rows).astype(str),
+    'B19301_001E': np.random.randint(1, 100000, n_rows).astype(str),
+    'B23025_002E': np.random.randint(1, 100000, n_rows).astype(str),
+    'B23025_003E': np.random.randint(1, 100000, n_rows).astype(str),
+    'B23025_004E': np.random.randint(1, 100000, n_rows).astype(str),
+    'B23025_005E': np.random.randint(1, 100000, n_rows).astype(str),
+    'state': np.random.randint(1, 50, n_rows).astype(str),
+}
+
+df_original = pd.DataFrame(data)
+
+def original_method(df):
+    numeric_columns = ['B19013_001E', 'B19301_001E', 'B23025_002E',
+                      'B23025_003E', 'B23025_004E', 'B23025_005E']
+    for col in numeric_columns:
+        df[col] = pd.to_numeric(df[col], errors='coerce')
+    df['state'] = pd.to_numeric(df['state'], errors='coerce')
+    return df
+
+def optimized_method(df):
+    numeric_columns = ['B19013_001E', 'B19301_001E', 'B23025_002E',
+                      'B23025_003E', 'B23025_004E', 'B23025_005E']
+    df[numeric_columns] = df[numeric_columns].apply(pd.to_numeric, errors='coerce')
+    df['state'] = pd.to_numeric(df['state'], errors='coerce')
+    return df
+
+# Test for equivalence
+df1 = original_method(df_original.copy())
+df2 = optimized_method(df_original.copy())
+pd.testing.assert_frame_equal(df1, df2)
+
+time_original = timeit.timeit("original_method(df_original.copy())", globals=globals(), number=10)
+time_optimized = timeit.timeit("optimized_method(df_original.copy())", globals=globals(), number=10)
+
+print(f"Original: {time_original:.4f}s")
+print(f"Optimized: {time_optimized:.4f}s")
+print(f"Improvement: {(time_original - time_optimized) / time_original * 100:.2f}%")

data_engineering/data_sources/census_acs/census_extractor.py

@@ -257,8 +257,7 @@ def _clean_economic_data(self, df: pd.DataFrame) -> pd.DataFrame:
         numeric_columns = ['B19013_001E', 'B19301_001E', 'B23025_002E',
                           'B23025_003E', 'B23025_004E', 'B23025_005E']
 
-        for col in numeric_columns:
-            df[col] = pd.to_numeric(df[col], errors='coerce')
+        df[numeric_columns] = df[numeric_columns].apply(pd.to_numeric, errors='coerce')
 
         # Convert state code to numeric
         df['state'] = pd.to_numeric(df['state'], errors='coerce')

test_census_extractor.py

@@ -0,0 +1,31 @@
+import pandas as pd
+from data_engineering.data_sources.census_acs.census_extractor import CensusExtractor
+
+def test_clean_economic_data():
+    extractor = CensusExtractor()
+    data = {
+        'B19013_001E': ['1000', '2000', 'invalid'],
+        'B19301_001E': ['3000', '4000', '5000'],
+        'B23025_002E': ['100', '200', '300'],
+        'B23025_003E': ['90', '180', '270'],
+        'B23025_004E': ['80', '160', '240'],
+        'B23025_005E': ['10', '20', '30'],
+        'state': ['01', '02', '03'],
+        'NAME': ['State 1', 'State 2', 'State 3'],
+        'year': [2021, 2021, 2021],
+        'extracted_at': ['2023-01-01', '2023-01-01', '2023-01-01']
+    }
+    df = pd.DataFrame(data)
+
+    # Run the cleaning function
+    cleaned_df = extractor._clean_economic_data(df)
+
+    # Assertions
+    assert len(cleaned_df) == 2  # The row with 'invalid' should be dropped due to dropna()
+    assert list(cleaned_df['state_code']) == [1, 2]
+    assert list(cleaned_df['median_household_income']) == [1000.0, 2000.0]
+
+    print("Test passed successfully!")
+
+if __name__ == "__main__":
+    test_clean_economic_data()