debugged parallelization logic
This commit is contained in:
Bill
17
Cargo.lock
generated
17
Cargo.lock
generated
@@ -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",
|
||||
|
||||
@@ -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
125
docs/Roadmap11-17
Normal 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.
|
||||
@@ -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;
|
||||
|
||||
108
src/parser.rs
108
src/parser.rs
@@ -1,7 +1,7 @@
|
||||
// --- External Crates ---
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use tokio_postgres::{NoTls, CopyOutStream, Client, Config as PgConfig}; // FIX: Import Client and PgConfig
|
||||
use tokio_postgres::{NoTls, CopyOutStream, Client, Config as PgConfig};
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use arrow::array::{
|
||||
ArrayBuilder, ArrayRef,
|
||||
@@ -21,22 +21,20 @@ 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::str::FromStr;
|
||||
use tokio::task::JoinSet;
|
||||
use arrow::compute::concat_batches;
|
||||
|
||||
// --- Internal Crates ---
|
||||
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
|
||||
@@ -48,17 +46,16 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
||||
}
|
||||
});
|
||||
|
||||
// 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!(
|
||||
@@ -73,16 +70,16 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
||||
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);
|
||||
@@ -92,7 +89,6 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
||||
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) ---
|
||||
// --- Phase 2: Parallel Execution ---
|
||||
|
||||
// Build the Arrow Schema *once* and wrap it in an Arc so it can be cloned cheaply
|
||||
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);
|
||||
|
||||
// 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 {
|
||||
@@ -163,8 +151,8 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
||||
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,7 +161,7 @@ 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,
|
||||
@@ -187,18 +175,14 @@ fn build_arrow_schema(config: &ConnectorConfig) -> Result<Schema> {
|
||||
_ => 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>),
|
||||
@@ -210,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 streaming "worker" function.
|
||||
/// It reads a 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())),
|
||||
@@ -235,16 +214,14 @@ 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")?;
|
||||
|
||||
@@ -253,7 +230,6 @@ async fn parse_binary_stream(
|
||||
|
||||
let mut cursor = Cursor::new(¤t_chunk[..]);
|
||||
|
||||
// --- Phase 2: Task 6 (Handle Header) ---
|
||||
if !is_header_parsed {
|
||||
if current_chunk.len() < 19 {
|
||||
leftover_buffer = current_chunk;
|
||||
@@ -264,11 +240,9 @@ async fn parse_binary_stream(
|
||||
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 => {
|
||||
@@ -279,35 +253,35 @@ async fn parse_binary_stream(
|
||||
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.");
|
||||
// println!("UncheckedIO: Worker reached end-of-stream trailer.");
|
||||
leftover_buffer.clear();
|
||||
break 'stream_loop; // Break the *outer* loop
|
||||
break 'stream_loop;
|
||||
}
|
||||
|
||||
// 3. Try to parse all fields for this row
|
||||
match parse_row(&mut cursor, &mut builders, ¤t_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(¤t_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()));
|
||||
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,
|
||||
@@ -321,7 +295,6 @@ async fn parse_binary_stream(
|
||||
}
|
||||
}).collect();
|
||||
|
||||
// Build the Final RecordBatch
|
||||
let record_batch = RecordBatch::try_new(
|
||||
arrow_schema.clone(),
|
||||
final_columns,
|
||||
@@ -330,7 +303,6 @@ async fn parse_binary_stream(
|
||||
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")?;
|
||||
@@ -340,23 +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.
|
||||
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)
|
||||
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(),
|
||||
@@ -367,18 +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*
|
||||
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>()?;
|
||||
@@ -423,7 +388,6 @@ fn parse_row(
|
||||
b.append_value(unix_days);
|
||||
}
|
||||
}
|
||||
} // End field loop
|
||||
|
||||
Ok(()) // Row was successfully parsed
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user