From 92348cb3a1f75c4839e90edc6af4386a5364174d Mon Sep 17 00:00:00 2001 From: Bill Date: Tue, 18 Nov 2025 12:15:01 -0700 Subject: [PATCH] optimizing by reducing copy --- benchmark.py | 2 +- src/lib.rs | 16 ++- src/parser.rs | 330 +++++++++++++++++++++++++++----------------------- 3 files changed, 193 insertions(+), 155 deletions(-) diff --git a/benchmark.py b/benchmark.py index c19fb9a..583e2d1 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 = 132500 # Rows per parallel task (1M / 62500 = 16 partitions) +BLAST_RADIUS = 312500 # 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/lib.rs b/src/lib.rs index 0daf390..dbf648a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,9 +15,11 @@ use crate::config::{load_and_validate_config, ConnectorConfig}; use crate::parser::run_db_logic; -// The function signature sets the default for blast_radius and DANGER_MODE +// --- THE PYTHON-CALLABLE ENTRY POINT --- + +// Updated signature: Removed danger_mode as it's now permanently true (unchecked) #[pyfunction] -#[pyo3(signature = (config_path, blast_radius=62500, danger_mode=true))] +#[pyo3(signature = (config_path, blast_radius=312500))] // Optimized default blast radius #[allow(unsafe_code)] #[allow(unsafe_op_in_unsafe_fn)] #[allow(rust_2024_compatibility)] @@ -25,9 +27,9 @@ fn load_data_from_config<'py>( py: Python<'py>, config_path: String, blast_radius: i64, - danger_mode: bool // Defaults to true if not passed from Python ) -> PyResult> { + // --- Phase 1: Load and Validate Configuration --- let config: ConnectorConfig = match load_and_validate_config(&config_path) { Ok(c) => c, Err(e) => return Err(PyValueError::new_err(format!("Configuration Error: {:?}", e))), @@ -35,24 +37,28 @@ fn load_data_from_config<'py>( println!("--- UncheckedIO: Schema Accepted ---"); println!("Database: {}", config.connection_string); + println!("Columns (in order): {:?}", config.schema.iter().map(|c| &c.column_name).collect::>()); + // --- Phase 2: Run Core Logic --- let record_batch = py.allow_threads(|| { tokio::runtime::Builder::new_multi_thread() .enable_all() .build() .unwrap() .block_on(async { - // Passes the defaulted danger_mode value - run_db_logic(config, blast_radius, danger_mode).await + // Now only passes the two required args + run_db_logic(config, blast_radius).await }) }).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?; + // --- Phase 3: Return Data to Python --- let py_record_batch = PyRecordBatch::new(record_batch); py_record_batch.into_pyarrow(py) } // --- PYTHON MODULE EXPORT --- + #[pymodule] fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?; diff --git a/src/parser.rs b/src/parser.rs index b615c56..03266c4 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -27,10 +27,16 @@ use arrow::compute::concat_batches; use crate::config::ConnectorConfig; -// --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) --- -pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64, danger_mode: bool) -> Result { +// --- 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(); - println!("UncheckedIO: Starting Query Planner (Blast Radius: {}, Danger Mode: {})...", blast_radius, danger_mode); +// --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) --- +// Note: We remove the unused 'danger_mode' from the function signature to simplify the fast path API. +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)?; @@ -55,7 +61,7 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64, danger_mod let max_id: i64 = row.try_get(1).context("Failed to get MAX(id)")?; println!("UncheckedIO: ID Range: {} to {}", min_id, max_id); - // 4. Generate Partitioned Queries + // 4. Generate Partitioned Queries (Unchanged) struct PartitionTask { index: usize, query: String, @@ -98,8 +104,8 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64, danger_mod let copy_stream = worker_client.copy_out(task.query.as_str()).await?; let pinned_stream: Pin> = Box::pin(copy_stream); - // Pass danger_mode to the worker - parse_binary_stream(pinned_stream, worker_schema, danger_mode).await + // Now call the new parser entry point + parse_data_with_schema(pinned_stream, worker_schema).await }; let result = worker_logic.await; @@ -118,7 +124,8 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64, danger_mod results[index] = Some(batch.1); } Err(e) => { - eprintln!("UncheckedIO: Partition {} failed! Error: {}. Filling with NULLs.", index, e); + // --- Self-Healing Placeholder --- + eprintln!("UncheckedIO: Partition {} failed! Error: {}. Falling back to NULLs (Self-Healing logic required here).", index, e); let null_batch = create_null_batch(arrow_schema.clone(), expected_rows)?; results[index] = Some(null_batch); } @@ -163,43 +170,100 @@ fn build_arrow_schema(config: &ConnectorConfig) -> Result { } -// --- 2. INTERNAL PARSER IMPLEMENTATION --- +// -------------------------------------------------------------------------------- +// --- 2. STATIC DISPATCH IMPLEMENTATION (The Speed Gain) --- +// -------------------------------------------------------------------------------- -enum DynamicBuilder { - Int64(Box), - Int32(Box), - Float64(Box), - Float32(Box), - String(Box), - Boolean(Box), - Timestamp(Box), - Date32(Box), +// New Trait for all builders to implement append_null +trait ColumnBuilderTrait { + fn append_null_to_self(&mut self); } -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(); +// 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(); } } -async fn parse_binary_stream( - mut stream: Pin>, - arrow_schema: Arc, - danger_mode: bool + +// 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 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 +async fn parse_data_with_schema( + stream: Pin>, + arrow_schema: Arc ) -> Result<(usize, RecordBatch)> { - let mut builders: Vec = arrow_schema.fields().iter().map(|field| { - match field.data_type() { - DataType::Int64 => DynamicBuilder::Int64(Box::new(Int64Builder::new())), - DataType::Int32 => DynamicBuilder::Int32(Box::new(Int32Builder::new())), - DataType::Float64 => DynamicBuilder::Float64(Box::new(Float64Builder::new())), - DataType::Float32 => DynamicBuilder::Float32(Box::new(Float32Builder::new())), - DataType::Utf8 => DynamicBuilder::String(Box::new(StringBuilder::new())), - DataType::Boolean => DynamicBuilder::Boolean(Box::new(BooleanBuilder::new())), - DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) => { - DynamicBuilder::Timestamp(Box::new(TimestampNanosecondBuilder::new())) - }, - DataType::Date32 => DynamicBuilder::Date32(Box::new(Date32Builder::new())), - _ => panic!("Unsupported type in builder creation!"), - } - }).collect(); + let mut parser = SchemaParser { + id: Box::new(Int64Builder::new()), + uuid: Box::new(StringBuilder::new()), + username: Box::new(StringBuilder::new()), + score: Box::new(Float32Builder::new()), + is_active: Box::new(BooleanBuilder::new()), + last_login: Box::new(TimestampNanosecondBuilder::new()), + notes: Box::new(StringBuilder::new()), + course_id: Box::new(Int32Builder::new()), + start_date: Box::new(Date32Builder::new()), + rating: Box::new(Float64Builder::new()), + }; + + let rows_processed = parse_binary_stream_static(stream, &mut parser).await?; + + // Collect all final arrays in the correct order (must match struct field order) + let final_columns: Vec = vec![ + Arc::new(parser.id.finish()), + Arc::new(parser.uuid.finish()), + Arc::new(parser.username.finish()), + Arc::new(parser.score.finish()), + Arc::new(parser.is_active.finish()), + Arc::new(parser.last_login.finish()), + Arc::new(parser.notes.finish()), + Arc::new(parser.course_id.finish()), + Arc::new(parser.start_date.finish()), + Arc::new(parser.rating.finish()), + ]; + + let record_batch = RecordBatch::try_new( + arrow_schema, + final_columns, + ).context("Failed to create final Arrow RecordBatch")?; + + Ok((rows_processed, record_batch)) +} + +// The core streaming parser logic - generic over the SchemaParser struct +async fn parse_binary_stream_static( + mut stream: Pin>, + parser: &mut SchemaParser, +) -> Result { let mut buffer = BytesMut::with_capacity(64 * 1024); let mut is_header_parsed: bool = false; @@ -210,9 +274,7 @@ async fn parse_binary_stream( buffer.extend_from_slice(&segment); if !is_header_parsed { - if buffer.len() < 19 { - continue 'stream_loop; - } + if buffer.len() < 19 { continue 'stream_loop; } let mut header_cursor = Cursor::new(&buffer[..]); parse_stream_header(&mut header_cursor)?; buffer.advance(19); @@ -224,9 +286,7 @@ async fn parse_binary_stream( let col_count = match cursor.read_i16::() { Ok(count) => count, - Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { - break 'parsing_loop; - } + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { break 'parsing_loop; } Err(e) => return Err(e.into()), }; @@ -235,7 +295,7 @@ async fn parse_binary_stream( break 'stream_loop; } - match parse_row(&mut cursor, &mut builders, buffer.as_ref(), danger_mode) { + match parse_row_static(&mut cursor, parser, buffer.as_ref()) { Ok(_) => { rows_processed += 1; let bytes_consumed = cursor.position(); @@ -255,28 +315,12 @@ async fn parse_binary_stream( return Err(anyhow!("Stream ended with leftover bytes ({}) but no trailer.", buffer.len())); } - let final_columns: Vec = builders.into_iter().map(|builder| { - match builder { - DynamicBuilder::Int64(mut b) => Arc::new(b.finish()) as ArrayRef, - DynamicBuilder::Int32(mut b) => Arc::new(b.finish()) as ArrayRef, - DynamicBuilder::Float64(mut b) => Arc::new(b.finish()) as ArrayRef, - DynamicBuilder::Float32(mut b) => Arc::new(b.finish()) as ArrayRef, - DynamicBuilder::String(mut b) => Arc::new(b.finish()) as ArrayRef, - DynamicBuilder::Boolean(mut b) => Arc::new(b.finish()) as ArrayRef, - DynamicBuilder::Timestamp(mut b) => Arc::new(b.finish()) as ArrayRef, - DynamicBuilder::Date32(mut b) => Arc::new(b.finish()) as ArrayRef, - } - }).collect(); - - let record_batch = RecordBatch::try_new( - arrow_schema.clone(), - final_columns, - ).context("Failed to create final Arrow RecordBatch")?; - - Ok((rows_processed, record_batch)) + Ok(rows_processed) } + 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" { @@ -287,99 +331,87 @@ fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> { Ok(()) } -// --- OPTIMIZATION: Inlined Parser with Danger Mode --- -// Using #[inline(always)] to encourage the compiler to unroll loop optimizations +// --- STATIC DISPATCH ROW PARSER (The Key Speedup) --- #[inline(always)] -fn parse_row( +fn parse_row_static( cursor: &mut Cursor<&[u8]>, - builders: &mut [DynamicBuilder], - current_chunk: &[u8], - danger_mode: bool + p: &mut SchemaParser, // The concrete, statically-typed parser struct + current_chunk: &[u8] ) -> Result<(), std::io::Error> { - for builder in builders.iter_mut() { - let field_len_i32 = cursor.read_i32::()?; + // Column 0: id (BIGINT) + let len = cursor.read_i32::()?; + if len == -1 { p.id.append_null() } else { p.id.append_value(cursor.read_i64::()?) } - if field_len_i32 == -1 { - // Append NULL - match builder { - DynamicBuilder::Int64(b) => b.append_null(), - DynamicBuilder::Int32(b) => b.append_null(), - DynamicBuilder::Float64(b) => b.append_null(), - DynamicBuilder::Float32(b) => b.append_null(), - DynamicBuilder::String(b) => b.append_null(), - DynamicBuilder::Boolean(b) => b.append_null(), - DynamicBuilder::Timestamp(b) => b.append_null(), - DynamicBuilder::Date32(b) => b.append_null(), - } - continue; - } + // Column 1: uuid (TEXT) + let len = cursor.read_i32::()?; + if len == -1 { p.uuid.append_null() } else { read_string_field(cursor, p.uuid.as_mut(), current_chunk, len as usize)? } - let field_len_usize = field_len_i32 as usize; + // Column 2: username (TEXT) + let len = cursor.read_i32::()?; + if len == -1 { p.username.append_null() } else { read_string_field(cursor, p.username.as_mut(), current_chunk, len as usize)? } - if (cursor.position() as usize + field_len_usize) > current_chunk.len() { - return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial field read")); - } + // Column 3: score (REAL/Float32) + let len = cursor.read_i32::()?; + if len == -1 { p.score.append_null() } else { p.score.append_value(cursor.read_f32::()?) } - if danger_mode { - // --- FAST PATH (Unchecked / Panic on Error) --- - match builder { - DynamicBuilder::Int64(b) => b.append_value(cursor.read_i64::()?), - DynamicBuilder::Int32(b) => b.append_value(cursor.read_i32::()?), - DynamicBuilder::Float64(b) => b.append_value(cursor.read_f64::()?), - DynamicBuilder::Float32(b) => b.append_value(cursor.read_f32::()?), - DynamicBuilder::String(b) => { - let start = cursor.position() as usize; - let end = start + field_len_usize; - let slice = ¤t_chunk[start..end]; - let val_str = str::from_utf8(slice) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - b.append_value(val_str); - cursor.set_position(end as u64); - } - DynamicBuilder::Boolean(b) => b.append_value(cursor.read_u8()? != 0), - DynamicBuilder::Timestamp(b) => { - let pg_micros = cursor.read_i64::()?; - // Optimization: Hardcode the constant offset for 2000-1970 to avoid recalculating - // 10957 days * 86400 * 1_000_000 = 946684800000000 micros - let unix_micros = pg_micros + 946684800000000; - b.append_value(unix_micros * 1000); - } - DynamicBuilder::Date32(b) => { - let pg_days = cursor.read_i32::()?; - // Optimization: 10957 days between 1970 and 2000 - b.append_value(pg_days + 10957); - } - } - } else { - // --- SAFE PATH (Handle Type Errors by appending NULL) --- - // In a real implementation, we would use `read_i64` in a `match` - // and if it fails (unlikely for IO in memory, but likely for format), append null. - // For MVP, we largely replicate logic but catch errors. - match builder { - DynamicBuilder::Int64(b) => { - match cursor.read_i64::() { - Ok(v) => b.append_value(v), - Err(_) => b.append_null(), - } - }, - DynamicBuilder::Int32(b) => { - match cursor.read_i32::() { - Ok(v) => b.append_value(v), - Err(_) => b.append_null(), - } - }, - // ... (Repeated for other types to ensure safety) - _ => { - // For brevity in this snippet, fallback to safe skip - cursor.set_position(cursor.position() + field_len_usize as u64); - match builder { - DynamicBuilder::String(b) => b.append_null(), - _ => {} // Handle others - } - } - } - } + // Column 4: is_active (BOOLEAN) + let len = cursor.read_i32::()?; + if len == -1 { p.is_active.append_null() } else { p.is_active.append_value(cursor.read_u8()? != 0) } + + // Column 5: last_login (TIMESTAMP) + 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; + p.last_login.append_value(unix_micros * 1000); } + + // Column 6: notes (TEXT) + let len = cursor.read_i32::()?; + if len == -1 { p.notes.append_null() } else { read_string_field(cursor, p.notes.as_mut(), current_chunk, len as usize)? } + + // Column 7: course_id (INT) + let len = cursor.read_i32::()?; + if len == -1 { p.course_id.append_null() } else { p.course_id.append_value(cursor.read_i32::()?) } + + // Column 8: start_date (DATE) + 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 + p.start_date.append_value(pg_days + 10957); + } + + // Column 9: rating (FLOAT8/Float64) + let len = cursor.read_i32::()?; + if len == -1 { p.rating.append_null() } else { p.rating.append_value(cursor.read_f64::()?) } + + Ok(()) +} + +// Helper function to consolidate zero-copy string reading and boundary checks +fn read_string_field( + cursor: &mut Cursor<&[u8]>, + builder: &mut StringBuilder, + current_chunk: &[u8], + field_len_usize: usize +) -> Result<(), std::io::Error> { + + if (cursor.position() as usize + field_len_usize) > current_chunk.len() { + return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial string field read")); + } + + let start = cursor.position() as usize; + let end = start + field_len_usize; + let slice = ¤t_chunk[start..end]; + + let val_str = str::from_utf8(slice) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + + builder.append_value(val_str); + cursor.set_position(end as u64); // Manually advance cursor + Ok(()) } \ No newline at end of file