1
0

ingest_csv.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. #!/usr/bin/env python3
  2. import pandas as pd
  3. import myloginpath
  4. import urllib.parse
  5. from sqlalchemy import create_engine, text
  6. from sqlalchemy.types import VARCHAR, TEXT, DOUBLE
  7. import os
  8. import sys
  9. from snmp_notifier import notifier
  10. def get_loader_engine():
  11. try:
  12. import os
  13. db_host = os.environ.get('DB_HOST')
  14. db_user = os.environ.get('DB_USER')
  15. db_pass = os.environ.get('DB_PASS')
  16. if db_host and db_user and db_pass:
  17. password = urllib.parse.quote_plus(db_pass)
  18. conn_str = f"mysql+pymysql://{db_user}:{password}@{db_host}/food_db?charset=utf8mb4"
  19. return create_engine(conn_str)
  20. conf = myloginpath.parse('app_loader')
  21. user = conf.get('user')
  22. password = urllib.parse.quote_plus(conf.get('password'))
  23. host = conf.get('host', '127.0.0.1')
  24. database = 'food_db'
  25. conn_str = f"mysql+pymysql://{user}:{password}@{host}/{database}?charset=utf8mb4"
  26. return create_engine(conn_str)
  27. except Exception as e:
  28. print(f"❌ Failed to parse myloginpath or create engine: {e}")
  29. sys.exit(1)
  30. def ingest_file(filename, engine):
  31. if not os.path.exists(filename):
  32. print(f"File {filename} not found locally.")
  33. return False
  34. print(f"\n🚀 Found {filename}! Starting grouped vertical partition ingestion...")
  35. chunk_size = 10000
  36. total_processed = 0
  37. # Define the groupings
  38. groups = {
  39. 'products_core': ['code', 'product_name', 'generic_name', 'brands', 'ingredients_text'],
  40. 'products_allergens': ['code', 'allergens'],
  41. 'products_macros': ['code', 'energy-kcal_100g', 'proteins_100g', 'fat_100g', 'carbohydrates_100g', 'sugars_100g', 'fiber_100g', 'sodium_100g', 'salt_100g', 'cholesterol_100g'],
  42. 'products_vitamins': ['code', 'vitamin-a_100g', 'vitamin-b1_100g', 'vitamin-b2_100g', 'vitamin-pp_100g', 'vitamin-b6_100g', 'vitamin-b9_100g', 'vitamin-b12_100g', 'vitamin-c_100g', 'vitamin-d_100g', 'vitamin-e_100g', 'vitamin-k_100g'],
  43. 'products_minerals': ['code', 'calcium_100g', 'iron_100g', 'magnesium_100g', 'potassium_100g', 'zinc_100g']
  44. }
  45. # Pre-calculate what to read
  46. all_required_cols = list(set([col for cols in groups.values() for col in cols]))
  47. for chunk in pd.read_csv(filename, sep='\t', dtype=str, chunksize=chunk_size, on_bad_lines='skip', low_memory=False, encoding='utf-8'):
  48. try:
  49. # Drop rows with missing codes
  50. if 'code' not in chunk.columns:
  51. continue
  52. df = chunk.dropna(subset=['code']).drop_duplicates(subset=['code']).copy()
  53. # Ensure all required columns exist in the chunk (fill with None if missing)
  54. for col in all_required_cols:
  55. if col not in df.columns:
  56. df[col] = None
  57. for table_name, columns in groups.items():
  58. slice_df = df[columns].copy()
  59. # Cast datatypes: core and allergens are TEXT, others are DOUBLE
  60. if table_name in ['products_core', 'products_allergens']:
  61. sql_dtypes = {col: TEXT() for col in columns if col != 'code'}
  62. sql_dtypes['code'] = VARCHAR(50)
  63. else:
  64. # Convert to numeric (double) safely
  65. for col in columns:
  66. if col != 'code':
  67. slice_df[col] = pd.to_numeric(slice_df[col], errors='coerce')
  68. sql_dtypes = {col: DOUBLE() for col in columns if col != 'code'}
  69. sql_dtypes['code'] = VARCHAR(50)
  70. # Write to temp table
  71. temp_name = f"temp_{table_name}"
  72. slice_df.to_sql(temp_name, con=engine, if_exists='replace', index=False, dtype=sql_dtypes)
  73. # INSERT IGNORE into final table
  74. with engine.begin() as conn:
  75. conn.execute(text(f"CREATE TABLE IF NOT EXISTS {table_name} LIKE {temp_name}"))
  76. cols_str = ", ".join([f"`{c}`" for c in columns])
  77. conn.execute(text(f"INSERT IGNORE INTO {table_name} ({cols_str}) SELECT {cols_str} FROM {temp_name}"))
  78. conn.execute(text(f"DROP TABLE IF EXISTS {temp_name}"))
  79. total_processed += len(df)
  80. print(f" Successfully appended {total_processed} rows into grouped tables...", end="\r")
  81. if total_processed % 50000 == 0:
  82. notifier.send_alert(f"Ingestion Milestone: {total_processed} rows processed")
  83. except BaseException as e:
  84. notifier.send_alert(f"Ingestion Exception: {str(e)}")
  85. print(f"\n [Warning] Chunk skipped due to error: {e}")
  86. notifier.send_alert(f"Ingestion Finished: {filename}")
  87. print(f"\n✅ Finished importing {filename}.")
  88. return True
  89. if __name__ == "__main__":
  90. print("Initiating OpenFoodFacts Grouped Vertical Ingestion Process...")
  91. engine = get_loader_engine()
  92. processed_en = ingest_file('data/en.openfoodfacts.org.products.csv', engine)
  93. processed_fr = ingest_file('data/fr.openfoodfacts.org.products.csv', engine)
  94. if not processed_en and not processed_fr:
  95. print("\n❌ Could not find CSVs.")
  96. else:
  97. print("\n🎉 Full database reload complete! Ready for AI RAG.")