bill/unchecked-io
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

implemented zero-copy via cursor movement and reduced overhead.

Browse Source
This commit is contained in:
Bill
2025-11-18 11:27:10 -07:00
parent b2224a15b8
commit 69c709b77f
4 changed files with 221 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
benchmark.py
View File

@@ -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")

87
docs/sprint4.md Normal file
View File

@@ -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(&current_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.

27
src/lib.rs
View File

@@ -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<Bound<'py, PyAny>> {
// 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<Bound<'py, PyAny>> {
// --- 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<Bound<'py, PyAny>>
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(())

230
src/parser.rs
View File

@@ -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<RecordBatch> {
// (This section is unchanged from your previous correct version)
pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<RecordBatch> {
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<String> = 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<PartitionTask> = 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<CopyOutStream>> = Box::pin(copy_stream);
let copy_stream = worker_client.copy_out(query.as_str()).await?;
let pinned_stream: Pin<Box<CopyOutStream>> = 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<RecordBatch> = Vec::new();
// --- Phase 3: Collect, Order, and Stitch ---
let mut results: Vec<Option<RecordBatch>> = 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<RecordBatch> = 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<Schema>, num_rows: usize) -> Result<RecordBatch> {
let columns: Vec<ArrayRef> = 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<Schema> {
let schema_fields: Vec<Field> = 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::<Result<Vec<Field>>>()?;
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<Box<CopyOutStream>>,
arrow_schema: Arc<Schema>
) -> Result<RecordBatch> {
) -> Result<(usize, RecordBatch)> {
let mut builders: Vec<DynamicBuilder> = 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<u8> = 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<u8> = mem::take(&mut leftover_buffer);
current_chunk.extend_from_slice(&segment);
let mut cursor = Cursor::new(&current_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::<BigEndian>() {
Ok(count) => count,
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
cursor.set_position(safe_position);
leftover_buffer.extend_from_slice(&current_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, &current_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(&current_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<ArrayRef> = 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::<BigEndian>().context("Failed to read flags")?;
let _header_ext_len = cursor.read_u32::<BigEndian>().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::<BigEndian>()?;
if field_len_i32 == -1 {
@@ -361,13 +336,26 @@ fn parse_row(
let val = cursor.read_f32::<BigEndian>()?;
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 = &current_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);