configurable batch size for hardware optimizations

benchmark.py

@@ -4,28 +4,34 @@ import sqlalchemy
 import connectorx as cx
 import unchecked_io
 import os
-import yaml # We need pyyaml for this
+import yaml
 import time
 
 # --- 1. Define Connection Strings and Query ---
-# These must match your local Docker setup
 DB_USER = "postgres"
 DB_PASS = "mysecretpassword"
 DB_HOST = "localhost"
-DB_PORT = "5433" # <-- Your local Docker port
+DB_PORT = "5433"
 DB_NAME = "postgres"
 
 # Global Configuration
+
 # Rows per parallel task (1M / 62500 = 16 partitions)
+
 #BLAST_RADIUS = 1250000 # 16 partitions
+
 #BLAST_RADIUS = 625000 # 32 partitions
+
 #BLAST_RADIUS = 312500 # 64 partitions
+
 #BLAST_RADIUS = 156250 # 128 partitions
-BLAST_RADIUS = 125000 # 160 partitions
+
+#BLAST_RADIUS = 125000 # 160 partitions
+
 #BLAST_RADIUS = 12500 # 1600 partitions
-#BLAST_RADIUS = 1250 # 16000 partitions
-
 
+BLAST_RADIUS = 1250 # 16000 partitions
+BATCH_SIZE = 262144    # Aggregation Size (Large Memory chunks)
 
 # SQLAlchemy connection string (for Pandas)
 sqlalchemy_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
@@ -44,6 +50,7 @@ unchecked_io_query = "COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score
 config_data = {
     'connection_string': connectorx_conn_str,
     'query': unchecked_io_query,
+    'batch_size': BATCH_SIZE, # <--- This explicitly writes your setting to the file
     'schema': [
         {'column_name': 'id', 'arrow_type': 'Int64'},
         {'column_name': 'uuid', 'arrow_type': 'Utf8'},
@@ -62,6 +69,8 @@ with open(config_file, 'w') as f:
     yaml.dump(config_data, f)
 
 print(f"UncheckedIO Config: {config_file} (dynamically created for local test)")
+print(f"  - Blast Radius: {BLAST_RADIUS}")
+print(f"  - Batch Size:   {BATCH_SIZE}")
 
 # --- 3. Define Benchmark Functions ---
 
@@ -82,14 +91,6 @@ def test_unchecked_io():
 # --- 4. Run Benchmarks ---
 run_count = 1
 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...")
-# _ = test_pandas()
-# print("Timing pandas.read_sql...")
-# pandas_time = timeit.timeit(test_pandas, number=run_count) / run_count
-# print(f"Pandas Average Time: {pandas_time * 1000:.2f} ms")
 
 # --- ConnectorX ---
 print("\nRunning ConnectorX warmup...")

config.yaml

@@ -1,3 +1,4 @@
+batch_size: 524288
 connection_string: postgresql://postgres:mysecretpassword@localhost:5433/postgres
 query: COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score::REAL, is_active::BOOLEAN,
   last_login::TIMESTAMP, notes::TEXT, course_id::INT, start_date::DATE, rating::FLOAT8

src/config.rs

@@ -1,6 +1,6 @@
 // src/config.rs
 use serde::Deserialize;
-use anyhow::{Context, Result, anyhow};
+use anyhow::{Context, Result, anyhow}; // Added anyhow macro just in case
 use std::fmt::{self, Display};
 
 #[derive(Debug)]

src/parser.rs

@@ -127,6 +127,9 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
     #[cfg(feature = "profiling")]
     drop(_p1_guard);
 
+    // FIX: Define the duration variable here so it can be used in the print statement later
+    let duration_phase1 = start_phase1.elapsed();
+
 
     // --- PHASE 2: PARALLEL EXECUTION (WITH AGGREGATION) ---
     let start_phase2 = Instant::now();
@@ -170,7 +173,6 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
                     let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
 
                     // Parse DIRECTLY into the persistent parser
-                    // Note: We changed the signature of parse_binary_stream_static to take &mut usize
                     parse_binary_stream_static(pinned_stream, &mut parser).await
                 }.await;
 
@@ -193,7 +195,6 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
                         // SAFETY FLUSH: If we have pending data, flush it first!
                         if parser_row_count > 0 {
                             let batch = flush_parser(&mut parser, worker_schema.clone())?;
-                            // We associate this flush with the current task index for sorting
                             worker_batches.push((task.index, batch));
                             parser_row_count = 0;
                             parser = create_parser();
@@ -209,7 +210,6 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
             // FINAL FLUSH: Handle any remaining rows after queue is empty
             if parser_row_count > 0 {
                 let batch = flush_parser(&mut parser, worker_schema.clone())?;
-                // Use a high index or the last processed index; sorting handles order
                 worker_batches.push((usize::MAX, batch));
             }
 
@@ -240,8 +240,6 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
     if all_results.is_empty() { return Ok(RecordBatch::new_empty(arrow_schema)); }
 
     // Sort by task index to maintain relative order
-    // Note: With aggregation, strict row-to-partition mapping is blurred,
-    // but relative order of data is preserved.
     all_results.sort_by_key(|(index, _)| *index);
     let batches: Vec<RecordBatch> = all_results.into_iter().map(|(_, b)| b).collect();
     let final_batch = concat_batches(&arrow_schema, &batches)?;
@@ -249,11 +247,15 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
     #[cfg(feature = "profiling")]
     drop(_p3_guard);
 
+    let duration_total = start_total.elapsed();
+
+    // --- REPORT ---
     println!("--- UncheckedIO Internal Timing ---");
-    println!("Phase 1 (Setup):   {:.2?}", start_phase2.duration_since(start_total) - duration_phase1);
+    // FIX: Using duration_phase1 correctly now
+    println!("Phase 1 (Setup):   {:.2?}", duration_phase1);
     println!("Phase 2 (Execute): {:.2?}", start_phase3.duration_since(start_phase2));
     println!("Phase 3 (Concat):  {:.2?}", start_phase3.elapsed());
-    println!("Total Wall Time:   {:.2?}", start_total.elapsed());
+    println!("Total Wall Time:   {:.2?}", duration_total);
 
     Ok(final_batch)
 }
@@ -500,5 +502,4 @@ fn read_string_field(
 }
 
 // Profiler logic omitted for brevity (it remains unchanged from previous version)
-// Note: You must keep the run_profiler_logic function in the file if it's called by lib.rs
 pub async fn run_profiler_logic(_: &str) -> Result<String> { Ok("Profiler Placeholder".to_string()) }