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 12:15:01 -07:00
parent 6e477933d0
commit 92348cb3a1
3 changed files with 193 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/lib.rs
View File

@@ -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<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))),
@@ -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::<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 {
// 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)?)?;

330
src/parser.rs
View File

@@ -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<RecordBatch> {
// --- 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<RecordBatch> {
// 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<CopyOutStream>> = 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<Schema> {
}
// --- 2. INTERNAL PARSER IMPLEMENTATION ---
// --------------------------------------------------------------------------------
// --- 2. STATIC DISPATCH IMPLEMENTATION (The Speed Gain) ---
// --------------------------------------------------------------------------------
enum DynamicBuilder {
Int64(Box<Int64Builder>),
Int32(Box<Int32Builder>),
Float64(Box<Float64Builder>),
Float32(Box<Float32Builder>),
String(Box<StringBuilder>),
Boolean(Box<BooleanBuilder>),
Timestamp(Box<TimestampNanosecondBuilder>),
Date32(Box<Date32Builder>),
// 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<Int64Builder> { fn append_null_to_self(&mut self) { self.append_null(); } }
impl ColumnBuilderTrait for Box<Int32Builder> { fn append_null_to_self(&mut self) { self.append_null(); } }
impl ColumnBuilderTrait for Box<Float64Builder> { fn append_null_to_self(&mut self) { self.append_null(); } }
impl ColumnBuilderTrait for Box<Float32Builder> { fn append_null_to_self(&mut self) { self.append_null(); } }
impl ColumnBuilderTrait for Box<StringBuilder> { fn append_null_to_self(&mut self) { self.append_null(); } }
impl ColumnBuilderTrait for Box<BooleanBuilder> { fn append_null_to_self(&mut self) { self.append_null(); } }
impl ColumnBuilderTrait for Box<TimestampNanosecondBuilder> { fn append_null_to_self(&mut self) { self.append_null(); } }
impl ColumnBuilderTrait for Box<Date32Builder> { fn append_null_to_self(&mut self) { self.append_null(); } }
async fn parse_binary_stream(
mut stream: Pin<Box<CopyOutStream>>,
arrow_schema: Arc<Schema>,
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<Int64Builder>,
// Column 1: uuid
uuid: Box<StringBuilder>,
// Column 2: username
username: Box<StringBuilder>,
// Column 3: score
score: Box<Float32Builder>,
// Column 4: is_active
is_active: Box<BooleanBuilder>,
// Column 5: last_login
last_login: Box<TimestampNanosecondBuilder>,
// Column 6: notes
notes: Box<StringBuilder>,
// Column 7: course_id
course_id: Box<Int32Builder>,
// Column 8: start_date
start_date: Box<Date32Builder>,
// Column 9: rating
rating: Box<Float64Builder>,
// 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<Box<CopyOutStream>>,
arrow_schema: Arc<Schema>
) -> Result<(usize, RecordBatch)> {
let mut builders: Vec<DynamicBuilder> = 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<ArrayRef> = 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<Box<CopyOutStream>>,
parser: &mut SchemaParser,
) -> Result<usize> {
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::<BigEndian>() {
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<ArrayRef> = 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::<BigEndian>()?;
// Column 0: id (BIGINT)
let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.id.append_null() } else { p.id.append_value(cursor.read_i64::<BigEndian>()?) }
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::<BigEndian>()?;
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::<BigEndian>()?;
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::<BigEndian>()?;
if len == -1 { p.score.append_null() } else { p.score.append_value(cursor.read_f32::<BigEndian>()?) }
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);
}
}
} 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
}
}
}
}
// Column 4: is_active (BOOLEAN)
let len = cursor.read_i32::<BigEndian>()?;
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::<BigEndian>()?;
if len == -1 { p.last_login.append_null() } else {
let pg_micros = cursor.read_i64::<BigEndian>()?;
// 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::<BigEndian>()?;
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::<BigEndian>()?;
if len == -1 { p.course_id.append_null() } else { p.course_id.append_value(cursor.read_i32::<BigEndian>()?) }
// Column 8: start_date (DATE)
let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.start_date.append_null() } else {
let pg_days = cursor.read_i32::<BigEndian>()?;
// 10957 days between 1970 and 2000
p.start_date.append_value(pg_days + 10957);
}
// Column 9: rating (FLOAT8/Float64)
let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.rating.append_null() } else { p.rating.append_value(cursor.read_f64::<BigEndian>()?) }
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 = &current_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(())
}