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

debugged parallelization logic

Browse Source
This commit is contained in:
Bill
2025-11-18 09:59:49 -07:00
parent c91c3ad2f7
commit b2224a15b8
5 changed files with 241 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
Cargo.lock generated
View File

@@ -596,6 +596,12 @@ version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
[[package]]
name = "hermit-abi"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
[[package]]
name = "hmac"
version = "0.12.1"
@@ -858,6 +864,16 @@ dependencies = [
"libm",
]
[[package]]
name = "num_cpus"
version = "1.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b"
dependencies = [
"hermit-abi",
"libc",
]
[[package]]
name = "numpy"
version = "0.27.0"
@@ -1529,6 +1545,7 @@ dependencies = [
"bytes",
"chrono",
"futures-util",
"num_cpus",
"pyo3",
"pyo3-arrow",
"serde",

6
Cargo.toml
View File

@@ -11,7 +11,7 @@ crate-type = ["cdylib", "rlib"]
[dependencies]
# 1. Python Bindings for FFI
pyo3 = { version = "0.20", features = ["extension-module"] }
pyo3 = { version = "0.27.1", features = ["extension-module"] }
# 2. Configuration Parsing (YAML)
serde = { version = "1.0", features = ["derive"] }
@@ -19,7 +19,7 @@ serde_yaml = "0.9"
# 3. Apache Arrow and Data Handling
# FIX: Add the "compute" feature to get concat_batches
arrow = { version = "50.0", features = ["compute"] }
arrow = "57.0.0"
# 4. Asynchronous Runtime (Essential for I/O and Postgres)
tokio = { version = "1.37", features = ["full"] }
@@ -46,7 +46,7 @@ chrono = "0.4"
uuid = { version = "1.8", features = ["serde", "v4"] }
# 12. Arrow <-> Python Bridge (NEW)
pyo3-arrow = "0.13"
pyo3-arrow = "0.15.0"
# 13. System CPU Count (NEW)
num_cpus = "1.16"

125
docs/Roadmap11-17 Normal file
View File

@@ -0,0 +1,125 @@
Project `UncheckedIO` - Roadmap & Decision Log
Date: 2025-11-18
Status: MVP (single-threaded) is functional and benchmarks show it is 2x faster than ConnectorX in a cloud (Colab) environment. The parallel-executor branch (gobrrr) is currently failing to build.
## 1. Project Philosophy & Goal
The primary goal of `UncheckedIO` is to create the fastest possible data connector for bulk-loading data from a PostgreSQL database directly into a Python/Arrow environment.
The core philosophy is the **"Configured Opinion"**:
1. **The Problem:** Generic connectors waste significant CPU time at runtime discovering schemas and validating data types.
2. **Our Solution:** We _require_ the user to provide an explicit `config.yaml` file that defines the query and the exact, ordered schema. We **trust this config** to skip expensive checks.
3. **The "Unchecked" Joke:** The name is a warning. The library is "unchecked" (skipping runtime validation) to gain speed. The Rust compiler provides _memory_ safety; the user provides _schema_ safety.
## 2. Core Architectural Decisions
- **Stack:** Rust core (Tokio, Arrow, `tokio-postgres`) with a `PyO3`/`Maturin` bridge to Python.
- **Data Flow:** We _must_ use the PostgreSQL `COPY ... TO STDOUT (FORMAT binary)` protocol. This is the "firehose" that bypasses slow, row-by-row SQL serialization.
- **Parser:** We use a **Streaming Parser** (`parse_binary_stream`). This is the "Sprint 2" logic. It reads the `COPY` stream in small chunks, manages a `leftover_buffer` for partial rows, and parses bytes as they arrive. This has a tiny, constant memory footprint and is the **core engine for all operations**.
- **License:** **Business Source License (BSL) 1.1**.
- The _entire_ codebase (including "Pro" features) is source-available.
- It is **Free** for non-production use and for production use by entities with < $1M annual revenue.
- **Paid** commercial licenses are required for production use by all other entities.
## 3. Current Status & Blockers (The `gobrrr` Branch)
- **`main` Branch (WORKING):** Contains the **Sprint 2** single-threaded streaming parser.
- **Benchmark (Colab):** ~1603 ms (2x faster than ConnectorX).
- **FFI Status:** The complex `pyo3-arrow` FFI issues _on this branch_ are **SOLVED**. The correct, working syntax is:
- `fn open_floodgates<'py>(py: Python<'py>, ...) -> PyResult<Bound<'py, PyAny>>`
- `let py_record_batch = PyRecordBatch::new(record_batch); py_record_batch.into_pyarrow(py)`
- `m.add_function(wrap_pyfunction!(open_floodgates, m)?)?`
- **`gobrrr` Branch (BROKEN):** Contains the **Sprint 3** parallel executor logic.
- **Status:** This branch is **failing to compile**.
- **Unknown Errors:** We have not debugged these build failures yet, as they only appeared after the major refactor to introduce the parallel coordinator. This will be the first task when back at the IDE.
## 🏁 Next Sprints: The Roadmap
### Sprint 3: The Parallel Coordinator (The "Batch King")
This is the "dynamic work queue" architecture. We need to debug the `gobrrr` branch to get this working.
- [ ] **Task 1: Debug Build Failures**
- Get the `gobrrr` branch (which contains the parallel `run_db_logic` coordinator) to compile.
- [ ] **Task 2: Implement the `MIN/MAX` Pre-Query**
- `run_db_logic` runs `SELECT MIN(id), MAX(id) FROM ...` to find table boundaries.
- [ ] **Task 3: Implement the `blast_radius` Parameter**
- Add a `blast_radius: usize` argument to the `open_floodgates` Python function.
- The coordinator uses this to create a "work queue" of queries (e.g., `...WHERE id BETWEEN 1 AND 10000`, etc.).
- [ ] **Task 4: Spawn Parallel Workers**
- Use `tokio::task::JoinSet` to create a pool of workers.
- Each worker opens its _own_ new database connection.
- Each worker pulls a query from the queue, runs its `COPY` command, and parses it using our existing `parse_binary_stream` engine.
- [ ] **Task 5: Implement "Self-Healing" Fault Tolerance (Your New Idea)**
- This is the new, robust error handling model.
- If a worker task panics (due to messy data in `danger_mode`), the coordinator will:
1. Catch the `Err` from the `JoinSet`.
2. Log the error and the failed partition (e.g., "Chunk 4 (rows 40k-50k) failed").
3. **Re-spawn** that _one_ failed query, but this time passing a `danger_mode=False` flag to the parser.
- [ ] **Task 6: Stitch Results**
- Use `arrow::compute::concat_batches` to merge all `RecordBatch`es (from both fast-path and safe-path workers) into one final table.
- Return the final table _and_ a list of the partitions that had to be re-run in safe mode.
### Sprint 4: The "Pro" Features (The Parser Logic)
This involves upgrading the _parser engine_ (`parse_binary_stream` and `parse_row`) to be "mode-aware."
- [ ] **Task 7: Implement `danger_mode` Toggle in Parser**
- Add a `danger_mode: bool` argument to `parse_row`.
- **If `danger_mode=True`:** Keep the current logic. Use `?` for `read_i64`, `read_f32`, etc. This is fast but will panic on a type mismatch.
- **If `danger_mode=False` (Safe Mode):** Refactor `parse_row`. Instead of `?`, use `match` for _every_ field read. If the bytes don't match the expected type (e.g., string in an int column), append `NULL` to the builder and continue.
- [ ] **Task 8: Implement the `profile_data` Utility**
- Create a _new_ `PyFunction` called `profile_data`.
- This function runs a `SELECT * ... LIMIT 10000`.
- It _only_ runs the "Safe Mode" parser (`danger_mode=False`).
- It returns a Python dictionary (a report) of how many `NULL`s were appended per column.

2
setup_db.sql
View File

@@ -36,7 +36,7 @@ SELECT
(random() * 10 + 100)::INT AS course_id,
(NOW() - (random() * '1000 days'::INTERVAL))::DATE AS start_date,
(random() * 5)::FLOAT8 AS rating
FROM generate_series(1, 1000000) s(i);
FROM generate_series(1, 5000000) s(i);
-- 4. Analyze the table for better query planning (good practice)
ANALYZE benchmark_table;

226
src/parser.rs
View File

@@ -1,98 +1,94 @@
// --- External Crates ---
use std::pin::Pin;
use std::pin::Pin;
use std::sync::Arc;
use tokio_postgres::{NoTls, CopyOutStream, Client, Config as PgConfig}; // FIX: Import Client and PgConfig
use anyhow::{Context, Result, anyhow};
use tokio_postgres::{NoTls, CopyOutStream, Client, Config as PgConfig};
use anyhow::{Context, Result, anyhow};
use arrow::array::{
ArrayBuilder, ArrayRef,
ArrayBuilder, ArrayRef,
Int64Builder, Float64Builder, Float32Builder, StringBuilder, BooleanBuilder,
TimestampNanosecondBuilder, Date32Builder, Int32Builder
};
};
use arrow::datatypes::{
DataType, Field, Schema,
DataType, Field, Schema,
Float64Type, Float32Type, Int64Type, Int32Type, Utf8Type, BooleanType, TimestampNanosecondType,
Date32Type
};
};
use arrow::record_batch::RecordBatch;
use futures_util::stream::StreamExt;
use futures_util::stream::StreamExt;
use bytes::Bytes;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Cursor, Read};
use std::io::{Cursor, Read};
use std::str;
use chrono::{NaiveDateTime, NaiveDate};
use std::mem;
use std::str::FromStr; // FIX: Import FromStr for PgConfig
use tokio::task::JoinSet; // FIX: Import JoinSet for parallelism
use arrow::compute::concat_batches; // FIX: Import the final stitching function
use std::mem;
use std::str::FromStr;
use tokio::task::JoinSet;
use arrow::compute::concat_batches;
// --- Internal Crates ---
use crate::config::ConnectorConfig;
use crate::config::ConnectorConfig;
// --- 1. CORE DATABASE LOGIC (REFACTORED TO A *PARALLEL COORDINATOR*) ---
// --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) ---
pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
// --- Phase 1: Task 1-3 (Query Planner) ---
println!("UncheckedIO: Starting Query Planner...");
// 1. Establish the *coordinator* connection
// FIX: Parse the connection string into a config we can clone for workers
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);
}
});
// 2. Define Partition Strategy (Task 3)
let num_partitions = num_cpus::get().max(2); // Use all cores, but at least 2
let partition_key = "id"; // Task 1: Hardcode partition key
// 2. Define Partition Strategy
let num_partitions = num_cpus::get().max(2);
let partition_key = "id";
// 3. Query for Table Bounds (Task 2)
// 3. Query for Table Bounds
let (base_query, _) = config.query
.trim()
.split_once("TO STDOUT (FORMAT binary)")
.context("Failed to parse base query from config")?;
// We need just the table/query part, remove the COPY (...)
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
);
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 {
// Return an empty, valid RecordBatch
println!("UncheckedIO: Table has no rows (COUNT=0). Returning empty batch.");
let arrow_schema = build_arrow_schema(&config)?;
return Ok(RecordBatch::new_empty(arrow_schema));
// 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);
// --- Phase 1: Task 4 (Generate Partitioned Queries) ---
// 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;
break;
}
// Construct the new, partitioned query
let new_query = format!(
"COPY (SELECT * FROM ({}) AS sub WHERE {} BETWEEN {} AND {}) TO STDOUT (FORMAT binary)",
base_query_inner,
@@ -105,42 +101,34 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
println!("UncheckedIO: Generated {} parallel queries.", partition_queries.len());
// --- Phase 2: Tasks 5-9 (Parallel Execution) ---
// Build the Arrow Schema *once* and wrap it in an Arc so it can be cloned cheaply
// --- Phase 2: Parallel Execution ---
let arrow_schema = Arc::new(build_arrow_schema(&config)?);
// Task 5: Spawn Tokio Tasks
let mut join_set = JoinSet::new();
for query in partition_queries {
let worker_pg_config = pg_config.clone(); // Clone config for the new task
let worker_schema = arrow_schema.clone(); // Clone Arc (cheap) for the new task
let worker_pg_config = pg_config.clone();
let worker_schema = arrow_schema.clone();
join_set.spawn(async move {
// Task 6: Isolate Worker Connections
let (worker_client, worker_connection) = worker_pg_config.connect(NoTls).await?;
// Spawn the connection task for this worker
tokio::spawn(async move {
if let Err(e) = worker_connection.await {
eprintln!("Worker connection error: {}", e);
}
});
// Task 7: Execute Partition in Parallel
let copy_stream = worker_client.copy_out(query.as_str()).await?;
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
// Call our existing streaming parser!
let record_batch = parse_binary_stream(pinned_stream, worker_schema).await?;
// Return the finished batch
Ok::<_, anyhow::Error>(record_batch)
});
}
// --- Phase 3: Task 8 (Collect and Concatenate Results) ---
// --- Phase 3: Collect and Concatenate ---
let mut batches: Vec<RecordBatch> = Vec::new();
while let Some(join_result) = join_set.join_next().await {
match join_result {
@@ -157,14 +145,14 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
Err(e) => return Err(anyhow!("A tokio task failed to join: {}", e)),
}
}
if batches.is_empty() {
println!("UncheckedIO: All partitions returned empty. Returning empty batch.");
return Ok(RecordBatch::new_empty(arrow_schema));
}
// Task 8: Stitch all the RecordBatches together into one
let final_batch = concat_batches(&batches[0].schema(), &batches)?
// FIX 2: Removed the '?' after concat_batches so .context() applies to the Result
let final_batch = concat_batches(&batches[0].schema(), &batches)
.context("Failed to concatenate parallel batches")?;
Ok(final_batch)
@@ -173,56 +161,47 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
/// Helper function to build the Arrow Schema from the config
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";
let nullable = col_cfg.column_name == "notes"; // Hack for MVP
let arrow_type = match col_cfg.arrow_type.as_str() {
"Int64" => DataType::Int64,
"Int32" => DataType::Int32,
"Float64" => DataType::Float64,
"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,
"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>>>()?; // Collect the results, propagating errors
}).collect::<Result<Vec<Field>>>()?;
Ok(Schema::new(schema_fields))
}
// --- 2. INTERNAL PARSER IMPLEMENTATION ---
// This section (parse_binary_stream, parse_stream_header, parse_row)
// remains 100% UNCHANGED from our Sprint 2.
// It is now the "worker" logic that will be called by our parallel tasks.
// This enum will hold our different builder types
enum DynamicBuilder {
Int64(Box<Int64Builder>),
Int32(Box<Int32Builder>),
Float64(Box<Float64Builder>),
Int32(Box<Int32Builder>),
Float64(Box<Float64Builder>),
Float32(Box<Float32Builder>),
String(Box<StringBuilder>),
Boolean(Box<BooleanBuilder>),
Timestamp(Box<TimestampNanosecondBuilder>),
Date32(Box<Date32Builder>),
Boolean(Box<BooleanBuilder>),
Timestamp(Box<TimestampNanosecondBuilder>),
Date32(Box<Date32Builder>),
}
// Postgres Epoch for timestamps
const POSTGRES_EPOCH_NAIVE: NaiveDateTime = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap().and_hms_opt(0, 0, 0).unwrap();
// Unix Epoch for dates
const UNIX_EPOCH_NAIVE_DATE: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
/// This is the streaming "worker" function.
/// It reads a stream chunk by chunk and parses it.
async fn parse_binary_stream(
mut stream: Pin<Box<CopyOutStream>>,
mut stream: Pin<Box<CopyOutStream>>,
arrow_schema: Arc<Schema>
) -> Result<RecordBatch> {
// --- Phase 1: Task 1 (Initialize Builders) ---
let mut builders: Vec<DynamicBuilder> = arrow_schema.fields().iter().map(|field| {
match field.data_type() {
DataType::Int64 => DynamicBuilder::Int64(Box::new(Int64Builder::new())),
@@ -230,45 +209,40 @@ async fn parse_binary_stream(
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::Boolean => DynamicBuilder::Boolean(Box::new(BooleanBuilder::new())),
DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) => {
DynamicBuilder::Timestamp(Box::new(TimestampNanosecondBuilder::new()))
DynamicBuilder::Timestamp(Box::new(TimestampNanosecondBuilder::new()))
},
DataType::Date32 => DynamicBuilder::Date32(Box::new(Date32Builder::new())),
_ => panic!("Unsupported type in builder creation!"), // Will improve later
_ => panic!("Unsupported type in builder creation!"),
}
}).collect();
// --- Phase 1: Task 2 (Initialize State) ---
let mut leftover_buffer: Vec<u8> = Vec::new();
let mut is_header_parsed: bool = false;
let mut rows_processed: usize = 0;
// --- Phase 2: Task 4 (Streaming Loop) ---
'stream_loop: while let Some(segment_result) = stream.next().await {
'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);
let mut current_chunk: Vec<u8> = mem::take(&mut leftover_buffer);
current_chunk.extend_from_slice(&segment);
let mut cursor = Cursor::new(&current_chunk[..]);
// --- Phase 2: Task 6 (Handle Header) ---
if !is_header_parsed {
if current_chunk.len() < 19 {
leftover_buffer = current_chunk;
continue 'stream_loop;
continue 'stream_loop;
}
parse_stream_header(&mut cursor)?;
is_header_parsed = true;
}
// --- Phase 3: Tasks 7-9 (Inner Parsing Loop) ---
'parsing_loop: loop {
let safe_position = cursor.position();
// 1. Try to read the 2-byte row header (column count)
let col_count = match cursor.read_i16::<BigEndian>() {
Ok(count) => count,
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
@@ -276,38 +250,38 @@ async fn parse_binary_stream(
leftover_buffer.extend_from_slice(&current_chunk[safe_position as usize..]);
break 'parsing_loop;
}
Err(e) => return Err(e.into()),
Err(e) => return Err(e.into()),
};
// 2. Check for end-of-stream trailer
if col_count == -1 {
println!("UncheckedIO: Worker reached end-of-stream trailer.");
leftover_buffer.clear();
break 'stream_loop; // Break the *outer* loop
// println!("UncheckedIO: Worker reached end-of-stream trailer.");
leftover_buffer.clear();
break 'stream_loop;
}
// 3. Try to parse all fields for this row
match parse_row(&mut cursor, &mut builders, &current_chunk) {
Ok(_) => {
rows_processed += 1;
}
Err(e) if e.is::<std::io::Error>() && e.downcast_ref::<std::io::Error>().unwrap().kind() == std::io::ErrorKind::UnexpectedEof => {
// 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) => {
return Err(e);
// Explicitly convert std::io::Error to anyhow::Error
return Err(e.into());
}
}
} // End inner 'parsing_loop
} // End outer 'stream_loop
if !leftover_buffer.is_empty() {
return Err(anyhow!("Stream ended with leftover bytes ({}) but no trailer. Data is corrupt.", leftover_buffer.len()));
}
}
if !leftover_buffer.is_empty() {
return Err(anyhow!("Stream ended with leftover bytes ({}) but no trailer.", leftover_buffer.len()));
}
// Finalize all the builders
let final_columns: Vec<ArrayRef> = builders.into_iter().map(|builder| {
match builder {
DynamicBuilder::Int64(mut b) => Arc::new(b.finish()) as ArrayRef,
@@ -315,22 +289,20 @@ async fn parse_binary_stream(
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::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();
// Build the Final RecordBatch
let record_batch = RecordBatch::try_new(
arrow_schema.clone(),
final_columns,
).context("Failed to create final Arrow RecordBatch")?;
Ok(record_batch)
}
/// Helper function to parse the 19-byte Postgres binary header.
fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> {
let mut magic_signature = [0u8; 11];
cursor.read_exact(&mut magic_signature).context("Failed to read magic signature")?;
@@ -339,46 +311,39 @@ 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")?;
println!("UncheckedIO: Postgres binary header validated.");
Ok(())
}
/// Helper function to parse one full row of data from the cursor.
fn parse_row(
cursor: &mut Cursor<&[u8]>,
builders: &mut [DynamicBuilder],
current_chunk: &[u8] // Needed for partial string reads
) -> Result<(), std::io::Error> { // Returns a specific IO Error
cursor: &mut Cursor<&[u8]>,
builders: &mut [DynamicBuilder],
current_chunk: &[u8]
) -> Result<(), std::io::Error> {
for (i, builder) in builders.iter_mut().enumerate() {
// 1. Read field length (4 bytes)
let field_len_i32 = cursor.read_i32::<BigEndian>()?;
let field_len_i32 = cursor.read_i32::<BigEndian>()?;
if field_len_i32 == -1 {
// Handle NULLs
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::Boolean(b) => b.append_null(),
DynamicBuilder::Timestamp(b) => b.append_null(),
DynamicBuilder::Date32(b) => b.append_null(),
}
continue; // Go to the next field in this row
continue;
}
let field_len_usize = field_len_i32 as usize;
// 2. Check if we have enough bytes in *this chunk* for the *entire field*
if (cursor.position() as usize + field_len_usize) > current_chunk.len() {
// Partial read: The field's data is split.
return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial field read"));
}
// 3. We have enough bytes. Parse it.
match builder {
DynamicBuilder::Int64(b) => {
let val = cursor.read_i64::<BigEndian>()?;
@@ -405,7 +370,7 @@ fn parse_row(
}
DynamicBuilder::Boolean(b) => {
let val_bool = cursor.read_u8()?;
b.append_value(val_bool != 0);
b.append_value(val_bool != 0);
}
DynamicBuilder::Timestamp(b) => {
let pg_micros = cursor.read_i64::<BigEndian>()?;
@@ -423,7 +388,6 @@ fn parse_row(
b.append_value(unix_days);
}
}
} // End field loop
Ok(()) // Row was successfully parsed
}
Ok(())
}