diff --git a/benchmark.py b/benchmark.py index 6168dea..cf586e1 100644 --- a/benchmark.py +++ b/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...") diff --git a/config.yaml b/config.yaml index 644767f..4584b4c 100644 --- a/config.yaml +++ b/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 diff --git a/src/config.rs b/src/config.rs index c9c0155..8a6fb4a 100644 --- a/src/config.rs +++ b/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)] diff --git a/src/parser.rs b/src/parser.rs index 0be3c7e..378fe7e 100644 --- a/src/parser.rs +++ b/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::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 = 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 { Ok("Profiler Placeholder".to_string()) } \ No newline at end of file