keep postgress connections open and switch tasks between them. Setup builds 20 million row data tables now

benchmark.py

@@ -16,7 +16,7 @@ DB_PORT = "5433" # <-- Your local Docker port
 DB_NAME = "postgres"
 
 # Global Configuration
-BLAST_RADIUS = 125000 # Rows per parallel task (1M / 62500 = 16 partitions)
+BLAST_RADIUS = 1250000 # Rows per parallel task (1M / 62500 = 16 partitions)
 
 # SQLAlchemy connection string (for Pandas)
 sqlalchemy_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
@@ -57,8 +57,9 @@ print(f"UncheckedIO Config: {config_file} (dynamically created for local test)")
 # --- 3. Define Benchmark Functions ---
 
 def test_pandas():
-    df = pd.read_sql(sql_query, engine)
+#     df = pd.read_sql(sql_query, engine)
-    return df
+#     return df
+    pass
 
 def test_connectorx():
     df = cx.read_sql(connectorx_conn_str, sql_query, return_type="arrow")
@@ -70,16 +71,16 @@ def test_unchecked_io():
     return arrow_table
 
 # --- 4. Run Benchmarks ---
-run_count = 3
+run_count = 1
-print(f"Running benchmarks for 5,000,000 rows (average of {run_count} runs)...")
+print(f"Running benchmarks for 20,000,000 rows (average of {run_count} runs)...")
 print(f"Blast Radius: {BLAST_RADIUS} rows per task")
 
 # --- Pandas ---
-print("\nRunning Pandas warmup...")
+# print("\nRunning Pandas warmup...")
-_ = test_pandas()
+# _ = test_pandas()
-print("Timing pandas.read_sql...")
+# print("Timing pandas.read_sql...")
-pandas_time = timeit.timeit(test_pandas, number=run_count) / run_count
+# pandas_time = timeit.timeit(test_pandas, number=run_count) / run_count
-print(f"Pandas Average Time: {pandas_time * 1000:.2f} ms")
+# print(f"Pandas Average Time: {pandas_time * 1000:.2f} ms")
 
 # --- ConnectorX ---
 print("\nRunning ConnectorX warmup...")
@@ -97,15 +98,15 @@ print(f"UncheckedIO Average Time: {unchecked_io_time * 1000:.2f} ms")
 
 # --- 5. Print Results ---
 print("\n" + "---" * 10)
-print("--- Benchmark Results (5,000,000 Rows) ---")
+print("--- Benchmark Results (20,000,000 Rows) ---")
-print(f"Pandas:       {pandas_time * 1000:>10.2f} ms")
+# print(f"Pandas:       {pandas_time * 1000:>10.2f} ms")
 print(f"ConnectorX:   {connectorx_time * 1000:>10.2f} ms")
 print(f"UncheckedIO:  {unchecked_io_time * 1000:>10.2f} ms")
 print("---" * 10)
 
 print("\n--- Ratios ---")
 if unchecked_io_time > 0:
-    print(f"UncheckedIO is {pandas_time / unchecked_io_time:.2f}x faster than Pandas")
+#     print(f"UncheckedIO is {pandas_time / unchecked_io_time:.2f}x faster than Pandas")
     print(f"UncheckedIO is {connectorx_time / unchecked_io_time:.2f}x faster than ConnectorX")
 else:
     print("UncheckedIO was too fast to measure accurately!")

run_setup.py

@@ -0,0 +1,54 @@
+import sqlalchemy
+import os
+import time
+
+# --- Configuration (Must match benchmark.py) ---
+DB_USER = "postgres"
+DB_PASS = "mysecretpassword"
+DB_HOST = "localhost"
+DB_PORT = "5433"
+DB_NAME = "postgres"
+
+# Build the SQLAlchemy connection string
+sqlalchemy_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
+
+# Define the path to your setup SQL file
+# Adjust this path if you moved the setup_db.sql file
+# NOTE: This path should be correct based on the file structure you uploaded:
+sql_file_path = "billthemaker/unchecked-io/unchecked-io-f5624c1ce64a916629b9d01def2cfe6de0d08c63/setup_db.sql"
+
+
+def run_sql_setup(engine, path):
+    """Executes the SQL file content against the database."""
+    print(f"Connecting to database at {DB_HOST}:{DB_PORT}...")
+    try:
+        # 1. Read the raw SQL content
+        with open(path, 'r') as f:
+            sql_content = f.read()
+
+        # 2. Establish connection and execute
+        with engine.connect() as connection:
+            print(f"Executing SQL file: {path}")
+            # Use begin/commit block for safety
+            with connection.begin():
+                connection.exec_driver_sql(sql_content)
+
+            print("Successfully executed setup script!")
+            print("Starting ANALYZE (may take a moment for 20M rows)...")
+
+            # Execute ANALYZE separately for proper commit timing
+            with connection.begin():
+                connection.exec_driver_sql("ANALYZE benchmark_table")
+
+            print("Database setup complete.")
+
+    except Exception as e:
+        print(f"FATAL ERROR during database setup: {e}")
+        print("Please ensure your PostgreSQL server is running and accessible.")
+
+if __name__ == "__main__":
+    engine = sqlalchemy.create_engine(sqlalchemy_conn_str)
+    start_time = time.time()
+    run_sql_setup(engine, sql_file_path)
+    end_time = time.time()
+    print(f"Total time taken for setup: {end_time - start_time:.2f} seconds.")