From 2888c0b5e76df5d1ed40cce1a6f140d949f954f6 Mon Sep 17 00:00:00 2001 From: Bill Date: Tue, 18 Nov 2025 13:08:32 -0700 Subject: [PATCH] keep postgress connections open and switch tasks between them --- benchmark.py | 2 +- src/parser.rs | 115 ++++++++++++++++++++++++-------------------------- 2 files changed, 56 insertions(+), 61 deletions(-) diff --git a/benchmark.py b/benchmark.py index 583e2d1..8ad756e 100644 --- a/benchmark.py +++ b/benchmark.py @@ -16,7 +16,7 @@ DB_PORT = "5433" # <-- Your local Docker port DB_NAME = "postgres" # Global Configuration -BLAST_RADIUS = 312500 # Rows per parallel task (1M / 62500 = 16 partitions) +BLAST_RADIUS = 125000 # 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}" diff --git a/src/parser.rs b/src/parser.rs index 03266c4..f64cc92 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -1,7 +1,9 @@ // --- External Crates --- use std::pin::Pin; use std::sync::Arc; +// FIX: Replace direct tokio_postgres connections with deadpool use tokio_postgres::{NoTls, CopyOutStream, Config as PgConfig}; +use deadpool_postgres::{Pool, Manager, Runtime}; use anyhow::{Context, Result, anyhow}; use arrow::array::{ ArrayBuilder, ArrayRef, @@ -28,40 +30,48 @@ use crate::config::ConnectorConfig; // --- CONSTANTS --- -const POSTGRES_EPOCH_NAIVE: NaiveDateTime = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap().and_hms_opt(0, 0, 0).unwrap(); -const UNIX_EPOCH_NAIVE_DATE: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); +// Optimized calculation of epoch delta (2000-01-01 00:00:00 to 1970-01-01 00:00:00) +// 10957 days * 86400 seconds/day * 1,000,000 micros/second = 946684800000000 micros +const POSTGRES_EPOCH_MICROS_OFFSET: i64 = 946684800000000; + // --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) --- -// Note: We remove the unused 'danger_mode' from the function signature to simplify the fast path API. +// Note: We use blast_radius from Python config as the partitioning strategy pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result { - // Removed danger_mode print, now defaults to fast path println!("UncheckedIO: Starting Query Planner (Blast Radius: {})...", blast_radius); - // 1. Establish the *coordinator* connection - let pg_config = PgConfig::from_str(&config.connection_string)?; - let (client, connection) = pg_config.connect(NoTls).await - .context("Coordinator: Failed to connect to PostgreSQL")?; - tokio::spawn(async move { - if let Err(e) = connection.await { eprintln!("Coordinator connection error: {}", e); } - }); + // --- PHASE 1: SETUP CONNECTION POOL --- + let pg_config: tokio_postgres::Config = PgConfig::from_str(&config.connection_string) + .context("Invalid connection string in config")?; - // 2. Define Partition Strategy + // Initialize the Manager and Pool + let manager = Manager::new(pg_config.clone(), NoTls); + // Set pool size higher than the expected partition count to ensure connections are always available. + let pool = Pool::builder(manager) + .max_size(20) + .runtime(Runtime::Tokio1) + .build() + .context("Failed to build connection pool")?; + + // 2. Query for Table Bounds (using pool connection) + let client = pool.get().await.context("Failed to get pool connection for stats query")?; let partition_key = "id"; - // 3. Query for Table Bounds let (base_query, _) = config.query.trim().split_once("TO STDOUT (FORMAT binary)") .context("Failed to parse base query from config")?; let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")"); + // NOTE: We rely on MIN/MAX here, assuming dense key for benchmark data. let stats_query = format!("SELECT MIN({}), MAX({}) FROM ({}) AS subquery", partition_key, partition_key, base_query_inner); let row = client.query_one(&stats_query, &[]).await?; let min_id: i64 = row.try_get(0).context("Failed to get MIN(id)")?; let max_id: i64 = row.try_get(1).context("Failed to get MAX(id)")?; + // Connection is returned to the pool when 'client' is dropped here. println!("UncheckedIO: ID Range: {} to {}", min_id, max_id); - // 4. Generate Partitioned Queries (Unchanged) + // 4. Generate Partitioned Queries (Based on blast_radius from Python) struct PartitionTask { index: usize, query: String, @@ -91,20 +101,19 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result< let mut join_set = JoinSet::new(); for task in partitions { - let worker_pg_config = pg_config.clone(); + let worker_pool = pool.clone(); // Pass the pool handle let worker_schema = arrow_schema.clone(); join_set.spawn(async move { let worker_logic = async { - let (worker_client, worker_connection) = worker_pg_config.connect(NoTls).await?; - tokio::spawn(async move { - if let Err(e) = worker_connection.await { eprintln!("Worker connection error: {}", e); } - }); + // Get connection from pool (This is the speedup) + let worker_client = worker_pool.get().await + .context("Worker: Failed to get pool connection")?; let copy_stream = worker_client.copy_out(task.query.as_str()).await?; let pinned_stream: Pin> = Box::pin(copy_stream); - // Now call the new parser entry point + // Call the static dispatch parser parse_data_with_schema(pinned_stream, worker_schema).await }; @@ -153,67 +162,45 @@ fn create_null_batch(schema: Arc, num_rows: usize) -> Result Result { let schema_fields: Vec = config.schema.iter().map(|col_cfg| { - let nullable = col_cfg.column_name == "notes"; + let nullable = col_cfg.column_name == "notes"; // Hack for MVP + let arrow_type = match col_cfg.arrow_type.as_str() { - "Int64" => DataType::Int64, "Int32" => DataType::Int32, - "Float64" => DataType::Float64, "Float32" => DataType::Float32, - "Utf8" | "String" => DataType::Utf8, "Boolean" => DataType::Boolean, + "Int64" => DataType::Int64, + "Int32" => DataType::Int32, + "Float64" => DataType::Float64, + "Float32" => DataType::Float32, + "Utf8" | "String" => DataType::Utf8, + "Boolean" => DataType::Boolean, "Timestamp(Nanosecond, None)" => DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), "Date32" => DataType::Date32, _ => return Err(anyhow!("Unsupported type in config: {}", col_cfg.arrow_type)), }; Ok(Field::new(&col_cfg.column_name, arrow_type, nullable)) }).collect::>>()?; + Ok(Schema::new(schema_fields)) } // -------------------------------------------------------------------------------- -// --- 2. STATIC DISPATCH IMPLEMENTATION (The Speed Gain) --- +// --- 2. STATIC DISPATCH IMPLEMENTATION (The Fast Parser) --- // -------------------------------------------------------------------------------- -// New Trait for all builders to implement append_null -trait ColumnBuilderTrait { - fn append_null_to_self(&mut self); -} - -// Implement the trait for the Boxed builders (which were in the DynamicBuilder enum) -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } -impl ColumnBuilderTrait for Box { fn append_null_to_self(&mut self) { self.append_null(); } } - - -// New struct to hold the builders in a statically-known, fixed order -// Note: We use the exact types from the benchmark schema to simplify the MVP +// Struct to hold the builders in a statically-known, fixed order (eliminates DynamicBuilder enum) struct SchemaParser { - // Column 0: id id: Box, - // Column 1: uuid uuid: Box, - // Column 2: username username: Box, - // Column 3: score score: Box, - // Column 4: is_active is_active: Box, - // Column 5: last_login last_login: Box, - // Column 6: notes notes: Box, - // Column 7: course_id course_id: Box, - // Column 8: start_date start_date: Box, - // Column 9: rating rating: Box, - // Note: This struct MUST match the order of the query result. } // Helper to construct and parse data using the static SchemaParser @@ -259,7 +246,7 @@ async fn parse_data_with_schema( Ok((rows_processed, record_batch)) } -// The core streaming parser logic - generic over the SchemaParser struct +// The core streaming parser logic async fn parse_binary_stream_static( mut stream: Pin>, parser: &mut SchemaParser, @@ -276,13 +263,14 @@ async fn parse_binary_stream_static( if !is_header_parsed { if buffer.len() < 19 { continue 'stream_loop; } let mut header_cursor = Cursor::new(&buffer[..]); - parse_stream_header(&mut header_cursor)?; + parse_stream_header(&mut header_cursor).context("Failed to parse stream header")?; buffer.advance(19); is_header_parsed = true; } 'parsing_loop: loop { let mut cursor = Cursor::new(&buffer[..]); + let safe_position = cursor.position(); let col_count = match cursor.read_i16::() { Ok(count) => count, @@ -302,6 +290,14 @@ async fn parse_binary_stream_static( buffer.advance(bytes_consumed as usize); } Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { + cursor.set_position(safe_position); + // Copy remaining bytes back to the buffer for the next chunk + let remaining_slice = &buffer.as_ref()[safe_position as usize..]; + let mut leftover_buffer_vec = Vec::new(); + leftover_buffer_vec.extend_from_slice(remaining_slice); + buffer.clear(); + buffer.extend_from_slice(&leftover_buffer_vec); + break 'parsing_loop; } Err(e) => { @@ -320,7 +316,6 @@ async fn parse_binary_stream_static( fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> { - // (Unchanged) let mut magic_signature = [0u8; 11]; cursor.read_exact(&mut magic_signature).context("Failed to read magic signature")?; if &magic_signature != b"PGCOPY\n\xff\r\n\0" { @@ -363,8 +358,8 @@ fn parse_row_static( let len = cursor.read_i32::()?; if len == -1 { p.last_login.append_null() } else { let pg_micros = cursor.read_i64::()?; - // 10957 days between 1970 and 2000 => 946684800000000 micros - let unix_micros = pg_micros + 946684800000000; + // Optimization: Constant offset applied + let unix_micros = pg_micros + POSTGRES_EPOCH_MICROS_OFFSET; p.last_login.append_value(unix_micros * 1000); } @@ -380,7 +375,7 @@ fn parse_row_static( let len = cursor.read_i32::()?; if len == -1 { p.start_date.append_null() } else { let pg_days = cursor.read_i32::()?; - // 10957 days between 1970 and 2000 + // Optimization: 10957 days between 1970 and 2000 p.start_date.append_value(pg_days + 10957); }