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

Merge pull request #2 from BillTheMaker/gobrrr

Browse Source 
gobrrr
parallelization done.
This commit is contained in:
BillTheMaker
2025-11-18 10:07:16 -07:00
committed by GitHub
parent 84bc53c4b4 b2224a15b8
commit f5624c1ce6
5 changed files with 326 additions and 124 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",

46
Cargo.toml
View File

@@ -1,36 +1,52 @@
[package]
name = "unchecked-io"
version = "0.1.0"
authors = ['billthemaker']
license = "MIT OR Apache-2.0"
edition = "2024"
authors = ["Your Name/Alias"] # Replace with your name or alias
license = "MIT OR Apache-2.0" # Good practice for open-source
edition = "2021"
[lib]
name = "unchecked_io"
crate-type = ["cdylib", "rlib"]
[dependencies]
# 1. Python Bindings for FFI
pyo3 = { version = "0.27.1", features = ["extension-module"] }
# 2. Configuration Parsing
# 2. Configuration Parsing (YAML)
serde = { version = "1.0", features = ["derive"] }
serde_yaml = "0.9"
# 3. Apache Arrow
# 3. Apache Arrow and Data Handling
# FIX: Add the "compute" feature to get concat_batches
arrow = "57.0.0"
# 4. Asynchronous Runtime
tokio = { version = "1.48.0", features = ["full"] }
# 4. Asynchronous Runtime (Essential for I/O and Postgres)
tokio = { version = "1.37", features = ["full"] }
# 5. DB Connection (Postgres)
tokio-postgres = "0.7.15"
anyhow = "1.0.100"
futures-util = "0.3.31"
bytes = "1.11.0"
byteorder = "1.5.0"
# 5. Database Connection (Postgres)
tokio-postgres = "0.7"
# 6. Error Handling Crate (The FIX)
anyhow = "1.0"
# 7. Futures Utilities (The FIX for StreamExt)
futures-util = "0.3"
# 8. Byte Buffer Management
bytes = "1.6"
# 9. Binary Data Reading (NEW)
byteorder = "1.5"
# 10. Timestamp Handling (NEW)
chrono = "0.4"
# 11. UUID Handling (NEW)
uuid = { version = "1.8", features = ["serde", "v4"] }
# 12. Arrow <-> Python Bridge (NEW)
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;

254
src/parser.rs
View File

@@ -1,7 +1,7 @@
// --- External Crates ---
use std::pin::Pin;
use std::sync::Arc;
use tokio_postgres::{NoTls, CopyOutStream};
use tokio_postgres::{NoTls, CopyOutStream, Client, Config as PgConfig};
use anyhow::{Context, Result, anyhow};
use arrow::array::{
ArrayBuilder, ArrayRef,
@@ -20,40 +20,148 @@ use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Cursor, Read};
use std::str;
use chrono::{NaiveDateTime, NaiveDate};
use std::mem; // <-- We need this for mem::take
use std::mem;
use std::str::FromStr;
use tokio::task::JoinSet;
use arrow::compute::concat_batches;
// --- Internal Crates ---
use crate::config::ConnectorConfig;
// --- 1. CORE DATABASE LOGIC (Public API) ---
// This function remains unchanged.
// --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) ---
pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
// 1. Establish the connection
println!("UncheckedIO: Attempting connection...");
let (client, connection) = tokio_postgres::connect(&config.connection_string, NoTls).await
.context("Failed to connect to PostgreSQL database")?;
println!("UncheckedIO: Starting Query Planner...");
// 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!("Postgres connection error: {}", e);
eprintln!("Coordinator connection error: {}", e);
}
});
// 2. Execute the COPY TO STDOUT command
let copy_query = &config.query;
println!("UncheckedIO: Executing user-defined query...");
// 2. Define Partition Strategy
let num_partitions = num_cpus::get().max(2);
let partition_key = "id";
let copy_stream = client.copy_out(copy_query.as_str()).await
.context("Failed to execute COPY TO STDOUT protocol. Check your query syntax and permissions.")?;
// 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")?;
// 3. Handle the Binary Stream
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")");
// Build the Arrow Schema from the config
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 {
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);
// 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);
}
println!("UncheckedIO: Generated {} parallel queries.", partition_queries.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 {
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?;
tokio::spawn(async move {
if let Err(e) = worker_connection.await {
eprintln!("Worker connection error: {}", e);
}
});
let copy_stream = worker_client.copy_out(query.as_str()).await?;
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
let record_batch = parse_binary_stream(pinned_stream, worker_schema).await?;
Ok::<_, anyhow::Error>(record_batch)
});
}
// --- 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 {
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)),
}
}
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)
.context("Failed to concatenate parallel batches")?;
Ok(final_batch)
}
/// 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,
@@ -64,32 +172,17 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
"Boolean" => DataType::Boolean,
"Timestamp(Nanosecond, None)" => DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None),
"Date32" => DataType::Date32,
_ => panic!("Unsupported type in config: {}", col_cfg.arrow_type),
_ => return Err(anyhow!("Unsupported type in config: {}", col_cfg.arrow_type)),
};
Field::new(&col_cfg.column_name, arrow_type, nullable)
}).collect();
let arrow_schema = Arc::new(Schema::new(schema_fields));
Ok(Field::new(&col_cfg.column_name, arrow_type, nullable))
}).collect::<Result<Vec<Field>>>()?;
// Call the parser
// We pass the schema and get back the final RecordBatch
let record_batch = parse_binary_stream(pinned_stream, arrow_schema.clone()).await?;
println!("UncheckedIO: Successfully parsed {} rows via binary stream.", record_batch.num_rows());
// 4. Final Output Confirmation
println!("UncheckedIO: Built RecordBatch with {} rows and {} columns.",
record_batch.num_rows(), record_batch.num_columns());
println!("UncheckedIO: Data transfer complete. We lived.");
Ok(record_batch)
Ok(Schema::new(schema_fields))
}
// --- 2. INTERNAL PARSER IMPLEMENTATION ---
// All the complex logic is now contained in this private section.
// This enum will hold our different builder types
enum DynamicBuilder {
Int64(Box<Int64Builder>),
Int32(Box<Int32Builder>),
@@ -101,19 +194,14 @@ enum DynamicBuilder {
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 refactored streaming state machine.
/// It reads the stream chunk by chunk and parses it.
async fn parse_binary_stream(
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())),
@@ -126,100 +214,74 @@ async fn parse_binary_stream(
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 {
let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?;
// Combine leftover bytes from last chunk with the new chunk
let mut current_chunk: Vec<u8> = mem::take(&mut leftover_buffer);
current_chunk.extend_from_slice(&segment);
// --- Phase 2: Task 5 (Create Cursor) ---
let mut cursor = Cursor::new(&current_chunk[..]);
// --- Phase 2: Task 6 (Handle Header) ---
if !is_header_parsed {
// Check if we have enough bytes for the header (11 + 4 + 4 = 19 bytes)
if current_chunk.len() < 19 {
// Not enough data. Move the chunk back and wait for more.
leftover_buffer = current_chunk;
continue 'stream_loop; // Get next segment
continue 'stream_loop;
}
parse_stream_header(&mut cursor)?;
is_header_parsed = true;
}
// --- Phase 3: Tasks 7-9 (Inner Parsing Loop) ---
'parsing_loop: loop {
// --- Task 8: Implement Safe Read (Row Level) ---
// Save state before attempting to read a row's header
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 => {
// This is NOT an error. It's a partial read.
// We don't have enough data for a full row header.
// Rewind the cursor to the safe position...
cursor.set_position(safe_position);
// ... and save the remaining bytes for the next chunk.
leftover_buffer.extend_from_slice(&current_chunk[safe_position as usize..]);
// Break the *inner* loop to get the next network segment
break 'parsing_loop;
}
Err(e) => return Err(e.into()), // This is a real, unexpected error
Err(e) => return Err(e.into()),
};
// --- Task 10: Verify Stream Trailer ---
if col_count == -1 {
println!("UncheckedIO: Reached end-of-stream trailer.");
leftover_buffer.clear(); // We are done, clear any remaining bytes
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(_) => {
// Row was parsed successfully
rows_processed += 1;
}
// FIX: 'e' is already a std::io::Error. We just check its .kind() directly.
// This resolves the E0599 (method not found) errors.
// 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 => {
// This was a partial row. The error is expected.
// Rewind to the start of the row.
cursor.set_position(safe_position);
leftover_buffer.extend_from_slice(&current_chunk[safe_position as usize..]);
break 'parsing_loop;
}
Err(e) => {
// This was a real, unexpected error.
// FIX: Convert the std::io::Error into an anyhow::Error.
// This resolves the E0308 (mismatched types) error.
// Explicitly convert std::io::Error to anyhow::Error
return Err(e.into());
}
}
} // End inner 'parsing_loop
} // End outer 'stream_loop
// --- Phase 1: Task 3 (Refactor Finalization) ---
if !leftover_buffer.is_empty() {
// We should have broken on the trailer. If we have leftovers, something is wrong.
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,
@@ -233,7 +295,6 @@ async fn parse_binary_stream(
}
}).collect();
// Build the Final RecordBatch
let record_batch = RecordBatch::try_new(
arrow_schema.clone(),
final_columns,
@@ -242,8 +303,6 @@ async fn parse_binary_stream(
Ok(record_batch)
}
/// Helper function to parse the 19-byte Postgres binary header.
/// This advances the cursor.
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")?;
@@ -253,26 +312,19 @@ 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.
/// This function is designed to fail with an `UnexpectedEof` error if the row is partial,
/// allowing the outer loop to handle it.
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
current_chunk: &[u8]
) -> Result<(), std::io::Error> {
for (i, builder) in builders.iter_mut().enumerate() {
// 1. Read field length (4 bytes)
// This will propagate the UnexpectedEof error up if it fails
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(),
@@ -283,21 +335,15 @@ fn parse_row(
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*
// This is a "look-ahead" check *before* we consume bytes.
if (cursor.position() as usize + field_len_usize) > current_chunk.len() {
// Partial read: The field's data is split.
// We return an EOF error to signal the outer loop.
return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial field read"));
}
// 3. We have enough bytes. Parse it.
// These reads will now succeed because we checked the length.
match builder {
DynamicBuilder::Int64(b) => {
let val = cursor.read_i64::<BigEndian>()?;
@@ -318,7 +364,6 @@ fn parse_row(
DynamicBuilder::String(b) => {
let mut str_buf = vec![0; field_len_usize];
cursor.read_exact(&mut str_buf)?;
// This str::from_utf8 is a potential panic! We should handle it.
let val_str = str::from_utf8(&str_buf)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
b.append_value(val_str);
@@ -343,7 +388,6 @@ fn parse_row(
b.append_value(unix_days);
}
}
} // End field loop
Ok(()) // Row was successfully parsed
}
Ok(())
}