setup_sqlite.py

import pandas as pd
import sqlite3
import glob
import time
import os

# Config
DATA_DIR = "data"
SQLITE_DB = f"{DATA_DIR}/taxi.db"
PARQUET_GLOB = f"{DATA_DIR}/yellow_tripdata_2024-*.parquet"

def load_sqlite():
    print("🛠️  Setting up SQLite for 'Fair Comparison'...")
    
    # Check if DB exists and delete to ensure clean timing
    if os.path.exists(SQLITE_DB):
        os.remove(SQLITE_DB)

    con = sqlite3.connect(SQLITE_DB)
    files = sorted(glob.glob(PARQUET_GLOB))
    
    total_start = time.time()
    total_rows = 0

    for f in files:
        print(f"   📖 Reading {f}...")
        # Read parquet into Pandas first (SQLite can't read parquet natively)
        df = pd.read_parquet(f)
        total_rows += len(df)
        
        print(f"   💾 Writing {len(df):,} rows to SQLite...")
        # Write to SQLite
        df.to_sql("trips", con, if_exists='append', index=False)
        
        # Free memory immediately
        del df

    # Indexing (Crucial for fair SQL comparison)
    print("   📇 Creating Indices (Fairness)...")
    con.execute("CREATE INDEX idx_pu ON trips(PULocationID)")
    con.execute("CREATE INDEX idx_do ON trips(DOLocationID)")
    con.execute("CREATE INDEX idx_dist_fare ON trips(trip_distance, fare_amount)")
    
    total_time = time.time() - total_start
    print(f"\n✅ SQLite Setup Complete.")
    print(f"   Total Rows: {total_rows:,}")
    print(f"   Total Setup Time: {total_time:.2f}s")
    
    # Save this number! You need it for your report.
    with open("sqlite_timing.txt", "w") as f:
        f.write(str(total_time))

if __name__ == "__main__":
    load_sqlite()