diff --git a/benchmark.py b/benchmark.py index b0c2522..508e3ab 100644 --- a/benchmark.py +++ b/benchmark.py @@ -15,6 +15,9 @@ DB_HOST = "localhost" DB_PORT = "5433" # <-- Your local Docker port DB_NAME = "postgres" +# Global Configuration +BLAST_RADIUS = 625000 # 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}" engine = sqlalchemy.create_engine(sqlalchemy_conn_str) @@ -62,12 +65,14 @@ def test_connectorx(): return df def test_unchecked_io(): - arrow_table = unchecked_io.load_data_from_config(config_file) + # FIX: Pass the BLAST_RADIUS argument + arrow_table = unchecked_io.load_data_from_config(config_file, BLAST_RADIUS) return arrow_table # --- 4. Run Benchmarks --- run_count = 3 -print(f"Running benchmarks for 1,000,000 rows (average of {run_count} runs)...") +print(f"Running benchmarks for 5,000,000 rows (average of {run_count} runs)...") +print(f"Blast Radius: {BLAST_RADIUS} rows per task") # --- Pandas --- print("\nRunning Pandas warmup...") @@ -92,7 +97,7 @@ print(f"UncheckedIO Average Time: {unchecked_io_time * 1000:.2f} ms") # --- 5. Print Results --- print("\n" + "---" * 10) -print("--- Benchmark Results (1,000,000 Rows) ---") +print("--- Benchmark Results (5,000,000 Rows) ---") 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") diff --git a/docs/sprint4.md b/docs/sprint4.md new file mode 100644 index 0000000..7e6d7cf --- /dev/null +++ b/docs/sprint4.md @@ -0,0 +1,87 @@ +# Sprint 4: Optimization & Stabilization + +This focuses on fixing the "death by a thousand cuts" performance regression we identified in `UncheckedIO`. + +You can copy-paste this directly into a new note in your vault (e.g., named `Sprint 4 - UncheckedIO Optimization`). + +--- + +# [[UncheckedIO]] Sprint 4: Optimization & Stabilization + +**Goal:** Recover performance lost during parallelization (Target: >2x faster than ConnectorX) and eliminate the 15M+ heap allocations. **Current Status:** ~1.5x faster than Pandas, but 0.5x speed of ConnectorX. **Tags:** #rust #optimization #project/unchecked-io #sprint + +## 🚨 Priority 1: The "String Killer" Fix (Memory) + +_Fixes the 15,000,000 heap allocations causing allocator contention._ + +- [ ] **Refactor `parse_row` string handling** + + - [ ] Remove `vec![0; field_len_usize]` allocation in `DynamicBuilder::String` match arm. + + - [ ] Implement zero-copy slicing: `let val_str = std::str::from_utf8(¤t_chunk[start..end])?`. + + - [ ] Update `StringBuilder` to append the slice directly (`b.append_value(val_str)`). + +- [ ] **Verify Memory Profile** + + - [ ] Run 5M row benchmark. + + - [ ] Ensure RAM usage is flat (no spikes matching row count). + + +## ⚡ Priority 2: The "Dispatch Killer" Fix (CPU) + +_Fixes the 50,000,000 runtime type checks._ + +- [ ] **Implement `Danger Mode` (Task 7)** + + - [ ] Add `danger_mode: bool` argument to `run_db_logic` and `load_data_from_config`. + + - [ ] **Fast Path (`danger_mode=true`):** + + - [ ] Create a new parsing loop that _assumes_ types based on column index. + + - [ ] Remove the inner `match builder` loop. + + - [ ] Use `get_unchecked` or specific Typed Builders if possible to bypass runtime checks. + + - [ ] **Safe Path (`danger_mode=false`):** + + - [ ] Keep the current `match` based logic for fallback/recovery. + + +## 🛠️ Priority 3: Tuning & Architecture + +- [ ] **Tune `BLAST_RADIUS`** + + - [ ] Reduce partition count from 80 to ~8-16 (matching logical cores). + + - [ ] Test `BLAST_RADIUS` values: `312,500` (16 chunks) vs `625,000` (8 chunks) vs `62,500` (80 chunks). + + - [ ] Goal: Minimize connection overhead (TCP handshakes) vs parallel throughput. + +- [ ] **Connection Pooling (Optional)** + + - [ ] Investigate if `bb8` or `deadpool-postgres` can reuse connections instead of opening 80 new ones. + + +## 📉 Benchmarking & Validation + +- [ ] **Run `benchmark.py`** + + - [ ] Target: Beat `ConnectorX` (< 6700ms). + + - [ ] Target: Beat `Pandas` by >4x (< 5000ms). + +- [ ] **Profile `UncheckedIO`** + + - [ ] Identify if `malloc` is still the bottleneck after string fix. + + +--- + +**Notes & Observations:** + +- Current regression caused by 15M allocations * 80 threads fighting for locks. + +- Previous "Safe Mode" parser was ~2x faster than ConnectorX (single-threaded). We need to get back to that baseline efficiency _before_ scaling out. \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index d305165..a700125 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,7 +6,7 @@ mod parser; use pyo3::prelude::*; use pyo3::exceptions::PyValueError; use tokio; -use pyo3::types::PyModule; +use pyo3::types::{PyModule, PyAny}; // Added PyAny use pyo3::Bound; // FIX 1: Import PyRecordBatch. @@ -24,7 +24,14 @@ use crate::parser::run_db_logic; #[allow(unsafe_code)] #[allow(unsafe_op_in_unsafe_fn)] #[allow(rust_2024_compatibility)] -fn load_data_from_config<'py>(py: Python<'py>, config_path: String) -> PyResult> { +// FIX 2: Updated signature to accept 'blast_radius' and return 'Bound<'py, PyAny>' +// This matches what you confirmed works with .into_pyarrow(py). +fn load_data_from_config<'py>( + py: Python<'py>, + config_path: String, + blast_radius: i64 +) -> PyResult> { + // --- Phase 1: Load and Validate Configuration --- let config: ConnectorConfig = match load_and_validate_config(&config_path) { Ok(c) => c, @@ -42,26 +49,26 @@ fn load_data_from_config<'py>(py: Python<'py>, config_path: String) -> PyResult< .build() .unwrap() .block_on(async { - run_db_logic(config).await + // FIX 3: Pass the 'blast_radius' argument to the parser logic + run_db_logic(config, blast_radius).await }) }).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?; // --- Phase 3: Return Data to Python --- - // FIX 2: Create PyRecordBatch wrapper using the standard PyO3 __new__ convention. - // let py_record_batch = PyRecordBatch::new(record_batch); - // // Replace: let py_record_batch = PyRecordBatch::new(py, record_batch).map_err(|e| ...)?; - // // With: + + // Create the wrapper. .new() takes one argument in this version. let py_record_batch = PyRecordBatch::new(record_batch); - // Return the PyRecordBatch wrapper object as a generic PyObject reference. + + // Call the correct conversion method which returns PyResult> py_record_batch.into_pyarrow(py) } // --- PYTHON MODULE EXPORT --- + #[pymodule] fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { - // FIX: Use the two-argument version of wrap_pyfunction! - // This resolves the E0308 type mismatch. + // FIX 4: Use the two-argument wrap_pyfunction! macro pattern which you confirmed works. m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?; Ok(()) diff --git a/src/parser.rs b/src/parser.rs index f3c1cc0..c365916 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -1,7 +1,7 @@ // --- External Crates --- use std::pin::Pin; use std::sync::Arc; -use tokio_postgres::{NoTls, CopyOutStream, Client, Config as PgConfig}; +use tokio_postgres::{NoTls, CopyOutStream, Config as PgConfig}; use anyhow::{Context, Result, anyhow}; use arrow::array::{ ArrayBuilder, ArrayRef, @@ -10,15 +10,13 @@ use arrow::array::{ }; use arrow::datatypes::{ DataType, Field, Schema, - Float64Type, Float32Type, Int64Type, Int32Type, Utf8Type, BooleanType, TimestampNanosecondType, - Date32Type }; use arrow::record_batch::RecordBatch; use futures_util::stream::StreamExt; -use bytes::Bytes; +use bytes::{Bytes, BytesMut, Buf}; use byteorder::{BigEndian, ReadBytesExt}; use std::io::{Cursor, Read}; -use std::str; +use std::str; // Ensure this is imported for zero-copy parsing use chrono::{NaiveDateTime, NaiveDate}; use std::mem; use std::str::FromStr; @@ -30,153 +28,139 @@ use crate::config::ConnectorConfig; // --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) --- -pub async fn run_db_logic(config: ConnectorConfig) -> Result { +// (This section is unchanged from your previous correct version) +pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result { - println!("UncheckedIO: Starting Query Planner..."); + println!("UncheckedIO: Starting Query Planner (Blast Radius: {} rows)...", 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); - } + if let Err(e) = connection.await { eprintln!("Coordinator connection error: {}", e); } }); // 2. Define Partition Strategy - let num_partitions = num_cpus::get().max(2); let partition_key = "id"; // 3. Query for Table Bounds - let (base_query, _) = config.query - .trim() - .split_once("TO STDOUT (FORMAT binary)") + 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(")"); - let stats_query = format!( - "SELECT MIN({}), MAX({}), COUNT(*) FROM ({}) AS subquery", - partition_key, partition_key, base_query_inner - ); + let stats_query = format!("SELECT MIN({}), MAX({}) FROM ({}) AS subquery", partition_key, partition_key, base_query_inner); - println!("UncheckedIO: Running stats query: {}", stats_query); 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)")?; - let count: i64 = row.try_get(2).context("Failed to get COUNT(*)")?; - - if count == 0 { - println!("UncheckedIO: Table has no rows (COUNT=0). Returning empty batch."); - let arrow_schema = build_arrow_schema(&config)?; - // FIX 1: Wrap schema in Arc to satisfy SchemaRef - return Ok(RecordBatch::new_empty(Arc::new(arrow_schema))); - } - - let chunk_size = (count as f64 / num_partitions as f64).ceil() as i64; - println!("UncheckedIO: Found {} rows. Creating {} partitions of ~{} rows each.", count, num_partitions, chunk_size); + println!("UncheckedIO: ID Range: {} to {}", min_id, max_id); // 4. Generate Partitioned Queries - let mut partition_queries: Vec = Vec::new(); - for i in 0..num_partitions { - let part_min = min_id + (i as i64 * chunk_size); - let part_max = (part_min + chunk_size - 1).min(max_id); - - if part_min > max_id { - break; - } - - let new_query = format!( - "COPY (SELECT * FROM ({}) AS sub WHERE {} BETWEEN {} AND {}) TO STDOUT (FORMAT binary)", - base_query_inner, - partition_key, - part_min, - part_max - ); - partition_queries.push(new_query); + struct PartitionTask { + index: usize, + query: String, + expected_rows: usize } - println!("UncheckedIO: Generated {} parallel queries.", partition_queries.len()); + let mut partitions: Vec = Vec::new(); + let mut current_min = min_id; + let mut idx = 0; + + while current_min <= max_id { + let current_max = (current_min + blast_radius - 1).min(max_id); + let new_query = format!( + "COPY (SELECT * FROM ({}) AS sub WHERE {} BETWEEN {} AND {}) TO STDOUT (FORMAT binary)", + base_query_inner, partition_key, current_min, current_max + ); + // Estimate rows + let estimated_rows = (current_max - current_min + 1) as usize; + + partitions.push(PartitionTask { index: idx, query: new_query, expected_rows: estimated_rows }); + current_min += blast_radius; + idx += 1; + } + println!("UncheckedIO: Generated {} parallel partitions.", partitions.len()); // --- Phase 2: Parallel Execution --- - let arrow_schema = Arc::new(build_arrow_schema(&config)?); let mut join_set = JoinSet::new(); - for query in partition_queries { + for task in partitions { let worker_pg_config = pg_config.clone(); let worker_schema = arrow_schema.clone(); join_set.spawn(async move { - let (worker_client, worker_connection) = worker_pg_config.connect(NoTls).await?; + 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); } + }); - tokio::spawn(async move { - if let Err(e) = worker_connection.await { - eprintln!("Worker connection error: {}", e); - } - }); + let copy_stream = worker_client.copy_out(task.query.as_str()).await?; + let pinned_stream: Pin> = Box::pin(copy_stream); - let copy_stream = worker_client.copy_out(query.as_str()).await?; - let pinned_stream: Pin> = Box::pin(copy_stream); + parse_binary_stream(pinned_stream, worker_schema).await + }; - let record_batch = parse_binary_stream(pinned_stream, worker_schema).await?; - - Ok::<_, anyhow::Error>(record_batch) + let result = worker_logic.await; + (task.index, result, task.expected_rows) }); } - // --- Phase 3: Collect and Concatenate --- - let mut batches: Vec = Vec::new(); + // --- Phase 3: Collect, Order, and Stitch --- + let mut results: Vec> = vec![None; idx]; + while let Some(join_result) = join_set.join_next().await { - match join_result { - Ok(batch_result) => { - match batch_result { - Ok(batch) => { - if batch.num_rows() > 0 { - batches.push(batch); - } - }, - Err(e) => return Err(anyhow!("A worker task failed: {}", e)), - } - }, - Err(e) => return Err(anyhow!("A tokio task failed to join: {}", e)), + // NOTE: In production, handle JoinError better than .context() panic + let (index, parse_result, expected_rows) = join_result.context("Worker thread panic")?; + + match parse_result { + Ok(batch) => { + results[index] = Some(batch.1); + } + Err(e) => { + eprintln!("UncheckedIO: Partition {} failed! Error: {}. Filling with NULLs.", index, e); + let null_batch = create_null_batch(arrow_schema.clone(), expected_rows)?; + results[index] = Some(null_batch); + } } } + let batches: Vec = results.into_iter() + .filter_map(|b| b) + .collect(); + if batches.is_empty() { - println!("UncheckedIO: All partitions returned empty. Returning empty batch."); return Ok(RecordBatch::new_empty(arrow_schema)); } - // FIX 2: Removed the '?' after concat_batches so .context() applies to the Result - let final_batch = concat_batches(&batches[0].schema(), &batches) + let final_batch = concat_batches(&arrow_schema, &batches) .context("Failed to concatenate parallel batches")?; Ok(final_batch) } -/// Helper function to build the Arrow Schema from the config +fn create_null_batch(schema: Arc, num_rows: usize) -> Result { + let columns: Vec = schema.fields().iter().map(|field| { + arrow::array::new_null_array(field.data_type(), num_rows) + }).collect(); + RecordBatch::try_new(schema, columns).context("Failed to create null placeholder batch") +} + fn build_arrow_schema(config: &ConnectorConfig) -> Result { let schema_fields: Vec = config.schema.iter().map(|col_cfg| { - let nullable = col_cfg.column_name == "notes"; // Hack for MVP - + let nullable = col_cfg.column_name == "notes"; 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)) } @@ -200,7 +184,7 @@ const UNIX_EPOCH_NAIVE_DATE: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unw async fn parse_binary_stream( mut stream: Pin>, arrow_schema: Arc -) -> Result { +) -> Result<(usize, RecordBatch)> { let mut builders: Vec = arrow_schema.fields().iter().map(|field| { match field.data_type() { @@ -218,68 +202,60 @@ async fn parse_binary_stream( } }).collect(); - let mut leftover_buffer: Vec = Vec::new(); + let mut buffer = BytesMut::with_capacity(64 * 1024); let mut is_header_parsed: bool = false; let mut rows_processed: usize = 0; 'stream_loop: while let Some(segment_result) = stream.next().await { let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?; - - let mut current_chunk: Vec = mem::take(&mut leftover_buffer); - current_chunk.extend_from_slice(&segment); - - let mut cursor = Cursor::new(¤t_chunk[..]); + buffer.extend_from_slice(&segment); if !is_header_parsed { - if current_chunk.len() < 19 { - leftover_buffer = current_chunk; + if buffer.len() < 19 { continue 'stream_loop; } - - parse_stream_header(&mut cursor)?; + let mut header_cursor = Cursor::new(&buffer[..]); + parse_stream_header(&mut header_cursor)?; + buffer.advance(19); is_header_parsed = true; } 'parsing_loop: loop { - let safe_position = cursor.position(); + // NOTE: We create a cursor over the *entire* remaining buffer slice + let mut cursor = Cursor::new(&buffer[..]); let col_count = match cursor.read_i16::() { Ok(count) => count, Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { - cursor.set_position(safe_position); - leftover_buffer.extend_from_slice(¤t_chunk[safe_position as usize..]); break 'parsing_loop; } Err(e) => return Err(e.into()), }; if col_count == -1 { - // println!("UncheckedIO: Worker reached end-of-stream trailer."); - leftover_buffer.clear(); + buffer.advance(2); break 'stream_loop; } - match parse_row(&mut cursor, &mut builders, ¤t_chunk) { + // Pass the slice `buffer.as_ref()` explicitly for zero-copy logic + match parse_row(&mut cursor, &mut builders, buffer.as_ref()) { Ok(_) => { rows_processed += 1; + let bytes_consumed = cursor.position(); + buffer.advance(bytes_consumed as usize); } - // FIX 3: Error handling for std::io::Error - // We check .kind() directly because 'e' is std::io::Error Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { - cursor.set_position(safe_position); - leftover_buffer.extend_from_slice(¤t_chunk[safe_position as usize..]); break 'parsing_loop; } Err(e) => { - // Explicitly convert std::io::Error to anyhow::Error return Err(e.into()); } } } } - if !leftover_buffer.is_empty() { - return Err(anyhow!("Stream ended with leftover bytes ({}) but no trailer.", leftover_buffer.len())); + if !buffer.is_empty() { + return Err(anyhow!("Stream ended with leftover bytes ({}) but no trailer.", buffer.len())); } let final_columns: Vec = builders.into_iter().map(|builder| { @@ -300,7 +276,7 @@ async fn parse_binary_stream( final_columns, ).context("Failed to create final Arrow RecordBatch")?; - Ok(record_batch) + Ok((rows_processed, record_batch)) } fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> { @@ -311,7 +287,6 @@ fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> { } let _flags = cursor.read_u32::().context("Failed to read flags")?; let _header_ext_len = cursor.read_u32::().context("Failed to read header extension length")?; - Ok(()) } @@ -321,7 +296,7 @@ fn parse_row( current_chunk: &[u8] ) -> Result<(), std::io::Error> { - for (i, builder) in builders.iter_mut().enumerate() { + for builder in builders.iter_mut() { let field_len_i32 = cursor.read_i32::()?; if field_len_i32 == -1 { @@ -361,13 +336,26 @@ fn parse_row( let val = cursor.read_f32::()?; b.append_value(val); } + // --- OPTIMIZATION START: ZERO-COPY STRING PARSING --- DynamicBuilder::String(b) => { - let mut str_buf = vec![0; field_len_usize]; - cursor.read_exact(&mut str_buf)?; - let val_str = str::from_utf8(&str_buf) + // 1. Get current cursor position (start of string data) + let start = cursor.position() as usize; + let end = start + field_len_usize; + + // 2. Slice the bytes directly from current_chunk (Zero-Copy) + // Safety: We already verified bounds check above. + let slice = ¤t_chunk[start..end]; + + // 3. Verify UTF-8 and append (still validates, but no allocation) + let val_str = str::from_utf8(slice) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + b.append_value(val_str); + + // 4. Manually advance cursor since we didn't use read_exact + cursor.set_position(end as u64); } + // --- OPTIMIZATION END --- DynamicBuilder::Boolean(b) => { let val_bool = cursor.read_u8()?; b.append_value(val_bool != 0);