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

optimizing by reducing copy

Browse Source
This commit is contained in:
Bill
2025-11-18 11:52:28 -07:00
parent 69c709b77f
commit 6e477933d0
3 changed files with 78 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/lib.rs
View File

@@ -6,33 +6,28 @@ mod parser;
use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
use tokio;
use pyo3::types::{PyModule, PyAny}; // Added PyAny
use pyo3::types::{PyModule, PyAny};
use pyo3::Bound;
// FIX 1: Import PyRecordBatch.
use pyo3_arrow::PyRecordBatch;
// --- Internal Crates ---
use crate::config::{load_and_validate_config, ConnectorConfig};
use crate::parser::run_db_logic;
// --- THE PYTHON-CALLABLE ENTRY POINT ---
// The function signature sets the default for blast_radius and DANGER_MODE
#[pyfunction]
#[pyo3(signature = (config_path, blast_radius=62500, danger_mode=true))]
#[allow(unsafe_code)]
#[allow(unsafe_op_in_unsafe_fn)]
#[allow(rust_2024_compatibility)]
// 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
blast_radius: i64,
danger_mode: bool // Defaults to true if not passed from Python
) -> PyResult<Bound<'py, PyAny>> {
// --- 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))),
@@ -40,36 +35,26 @@ 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::<Vec<_>>());
// --- Phase 2: Run Core Logic ---
let record_batch = py.allow_threads(|| {
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.unwrap()
.block_on(async {
// FIX 3: Pass the 'blast_radius' argument to the parser logic
run_db_logic(config, blast_radius).await
// Passes the defaulted danger_mode value
run_db_logic(config, blast_radius, danger_mode).await
})
}).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?;
// --- Phase 3: Return Data to Python ---
// Create the wrapper. .new() takes one argument in this version.
let py_record_batch = PyRecordBatch::new(record_batch);
// 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 4: Use the two-argument wrap_pyfunction! macro pattern which you confirmed works.
m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
Ok(())
}

136
src/parser.rs
View File

@@ -16,7 +16,7 @@ use futures_util::stream::StreamExt;
use bytes::{Bytes, BytesMut, Buf};
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Cursor, Read};
use std::str; // Ensure this is imported for zero-copy parsing
use std::str;
use chrono::{NaiveDateTime, NaiveDate};
use std::mem;
use std::str::FromStr;
@@ -28,10 +28,9 @@ use crate::config::ConnectorConfig;
// --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) ---
// (This section is unchanged from your previous correct version)
pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<RecordBatch> {
pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64, danger_mode: bool) -> Result<RecordBatch> {
println!("UncheckedIO: Starting Query Planner (Blast Radius: {} rows)...", blast_radius);
println!("UncheckedIO: Starting Query Planner (Blast Radius: {}, Danger Mode: {})...", blast_radius, danger_mode);
// 1. Establish the *coordinator* connection
let pg_config = PgConfig::from_str(&config.connection_string)?;
@@ -73,7 +72,6 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
"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 });
@@ -100,7 +98,8 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
let copy_stream = worker_client.copy_out(task.query.as_str()).await?;
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
parse_binary_stream(pinned_stream, worker_schema).await
// Pass danger_mode to the worker
parse_binary_stream(pinned_stream, worker_schema, danger_mode).await
};
let result = worker_logic.await;
@@ -112,7 +111,6 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
let mut results: Vec<Option<RecordBatch>> = vec![None; idx];
while let Some(join_result) = join_set.join_next().await {
// NOTE: In production, handle JoinError better than .context() panic
let (index, parse_result, expected_rows) = join_result.context("Worker thread panic")?;
match parse_result {
@@ -183,7 +181,8 @@ 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>
arrow_schema: Arc<Schema>,
danger_mode: bool
) -> Result<(usize, RecordBatch)> {
let mut builders: Vec<DynamicBuilder> = arrow_schema.fields().iter().map(|field| {
@@ -221,7 +220,6 @@ async fn parse_binary_stream(
}
'parsing_loop: loop {
// 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>() {
@@ -237,8 +235,7 @@ async fn parse_binary_stream(
break 'stream_loop;
}
// Pass the slice `buffer.as_ref()` explicitly for zero-copy logic
match parse_row(&mut cursor, &mut builders, buffer.as_ref()) {
match parse_row(&mut cursor, &mut builders, buffer.as_ref(), danger_mode) {
Ok(_) => {
rows_processed += 1;
let bytes_consumed = cursor.position();
@@ -290,16 +287,21 @@ 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
#[inline(always)]
fn parse_row(
cursor: &mut Cursor<&[u8]>,
builders: &mut [DynamicBuilder],
current_chunk: &[u8]
current_chunk: &[u8],
danger_mode: bool
) -> Result<(), std::io::Error> {
for builder in builders.iter_mut() {
let field_len_i32 = cursor.read_i32::<BigEndian>()?;
if field_len_i32 == -1 {
// Append NULL
match builder {
DynamicBuilder::Int64(b) => b.append_null(),
DynamicBuilder::Int32(b) => b.append_null(),
@@ -319,61 +321,63 @@ fn parse_row(
return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial field read"));
}
match builder {
DynamicBuilder::Int64(b) => {
let val = cursor.read_i64::<BigEndian>()?;
b.append_value(val);
if danger_mode {
// --- FAST PATH (Unchecked / Panic on Error) ---
match builder {
DynamicBuilder::Int64(b) => b.append_value(cursor.read_i64::<BigEndian>()?),
DynamicBuilder::Int32(b) => b.append_value(cursor.read_i32::<BigEndian>()?),
DynamicBuilder::Float64(b) => b.append_value(cursor.read_f64::<BigEndian>()?),
DynamicBuilder::Float32(b) => b.append_value(cursor.read_f32::<BigEndian>()?),
DynamicBuilder::String(b) => {
let start = cursor.position() as usize;
let end = start + field_len_usize;
let slice = &current_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::<BigEndian>()?;
// 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::<BigEndian>()?;
// Optimization: 10957 days between 1970 and 2000
b.append_value(pg_days + 10957);
}
}
DynamicBuilder::Int32(b) => {
let val = cursor.read_i32::<BigEndian>()?;
b.append_value(val);
}
DynamicBuilder::Float64(b) => {
let val = cursor.read_f64::<BigEndian>()?;
b.append_value(val);
}
DynamicBuilder::Float32(b) => {
let val = cursor.read_f32::<BigEndian>()?;
b.append_value(val);
}
// --- OPTIMIZATION START: ZERO-COPY STRING PARSING ---
DynamicBuilder::String(b) => {
// 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);
}
DynamicBuilder::Timestamp(b) => {
let pg_micros = cursor.read_i64::<BigEndian>()?;
let unix_epoch = NaiveDateTime::from_timestamp_opt(0, 0).unwrap();
let pg_epoch = POSTGRES_EPOCH_NAIVE;
let epoch_delta_micros = (pg_epoch - unix_epoch).num_microseconds().unwrap();
let unix_micros = epoch_delta_micros + pg_micros;
let unix_nanos = unix_micros * 1000;
b.append_value(unix_nanos);
}
DynamicBuilder::Date32(b) => {
let pg_days = cursor.read_i32::<BigEndian>()?;
let epoch_delta_days = (POSTGRES_EPOCH_NAIVE.date() - UNIX_EPOCH_NAIVE_DATE).num_days() as i32;
let unix_days = epoch_delta_days + pg_days;
b.append_value(unix_days);
} 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::<BigEndian>() {
Ok(v) => b.append_value(v),
Err(_) => b.append_null(),
}
},
DynamicBuilder::Int32(b) => {
match cursor.read_i32::<BigEndian>() {
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
}
}
}
}
}