go brrrrr
This commit is contained in:
106
benchmark.py
Normal file
106
benchmark.py
Normal file
@@ -0,0 +1,106 @@
|
|||||||
|
import timeit
|
||||||
|
import pandas as pd
|
||||||
|
import sqlalchemy
|
||||||
|
import connectorx as cx
|
||||||
|
import unchecked_io
|
||||||
|
import os
|
||||||
|
import yaml # We need pyyaml for this
|
||||||
|
import time
|
||||||
|
|
||||||
|
# --- 1. Define Connection Strings and Query ---
|
||||||
|
# These must match your local Docker setup
|
||||||
|
DB_USER = "postgres"
|
||||||
|
DB_PASS = "mysecretpassword"
|
||||||
|
DB_HOST = "localhost"
|
||||||
|
DB_PORT = "5433" # <-- Your local Docker port
|
||||||
|
DB_NAME = "postgres"
|
||||||
|
|
||||||
|
# SQLAlchemy connection string (for Pandas)
|
||||||
|
sqlalchemy_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
|
||||||
|
engine = sqlalchemy.create_engine(sqlalchemy_conn_str)
|
||||||
|
|
||||||
|
# ConnectorX connection string
|
||||||
|
connectorx_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
|
||||||
|
|
||||||
|
# The query for Pandas and ConnectorX (non-COPY)
|
||||||
|
sql_query = "SELECT * FROM benchmark_table"
|
||||||
|
|
||||||
|
# --- 2. Dynamically create the config.yaml for UncheckedIO ---
|
||||||
|
config_file = "config.yaml"
|
||||||
|
unchecked_io_query = "COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score::REAL, is_active::BOOLEAN, last_login::TIMESTAMP, notes::TEXT, course_id::INT, start_date::DATE, rating::FLOAT8 FROM benchmark_table) TO STDOUT (FORMAT binary)"
|
||||||
|
|
||||||
|
config_data = {
|
||||||
|
'connection_string': connectorx_conn_str,
|
||||||
|
'query': unchecked_io_query,
|
||||||
|
'schema': [
|
||||||
|
{'column_name': 'id', 'arrow_type': 'Int64'},
|
||||||
|
{'column_name': 'uuid', 'arrow_type': 'Utf8'},
|
||||||
|
{'column_name': 'username', 'arrow_type': 'Utf8'},
|
||||||
|
{'column_name': 'score', 'arrow_type': 'Float32'},
|
||||||
|
{'column_name': 'is_active', 'arrow_type': 'Boolean'},
|
||||||
|
{'column_name': 'last_login', 'arrow_type': 'Timestamp(Nanosecond, None)'},
|
||||||
|
{'column_name': 'notes', 'arrow_type': 'Utf8'},
|
||||||
|
{'column_name': 'course_id', 'arrow_type': 'Int32'},
|
||||||
|
{'column_name': 'start_date', 'arrow_type': 'Date32'},
|
||||||
|
{'column_name': 'rating', 'arrow_type': 'Float64'}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
with open(config_file, 'w') as f:
|
||||||
|
yaml.dump(config_data, f)
|
||||||
|
|
||||||
|
print(f"UncheckedIO Config: {config_file} (dynamically created for local test)")
|
||||||
|
|
||||||
|
# --- 3. Define Benchmark Functions ---
|
||||||
|
|
||||||
|
def test_pandas():
|
||||||
|
df = pd.read_sql(sql_query, engine)
|
||||||
|
return df
|
||||||
|
|
||||||
|
def test_connectorx():
|
||||||
|
df = cx.read_sql(connectorx_conn_str, sql_query, return_type="arrow")
|
||||||
|
return df
|
||||||
|
|
||||||
|
def test_unchecked_io():
|
||||||
|
arrow_table = unchecked_io.load_data_from_config(config_file)
|
||||||
|
return arrow_table
|
||||||
|
|
||||||
|
# --- 4. Run Benchmarks ---
|
||||||
|
run_count = 3
|
||||||
|
print(f"Running benchmarks for 1,000,000 rows (average of {run_count} runs)...")
|
||||||
|
|
||||||
|
# --- Pandas ---
|
||||||
|
print("\nRunning Pandas warmup...")
|
||||||
|
_ = test_pandas()
|
||||||
|
print("Timing pandas.read_sql...")
|
||||||
|
pandas_time = timeit.timeit(test_pandas, number=run_count) / run_count
|
||||||
|
print(f"Pandas Average Time: {pandas_time * 1000:.2f} ms")
|
||||||
|
|
||||||
|
# --- ConnectorX ---
|
||||||
|
print("\nRunning ConnectorX warmup...")
|
||||||
|
_ = test_connectorx()
|
||||||
|
print("Timing connectorx.read_sql...")
|
||||||
|
connectorx_time = timeit.timeit(test_connectorx, number=run_count) / run_count
|
||||||
|
print(f"ConnectorX Average Time: {connectorx_time * 1000:.2f} ms")
|
||||||
|
|
||||||
|
# --- UncheckedIO ---
|
||||||
|
print("\nRunning UncheckedIO warmup...")
|
||||||
|
_ = test_unchecked_io()
|
||||||
|
print("Timing unchecked_io.load_data_from_config...")
|
||||||
|
unchecked_io_time = timeit.timeit(test_unchecked_io, number=run_count) / run_count
|
||||||
|
print(f"UncheckedIO Average Time: {unchecked_io_time * 1000:.2f} ms")
|
||||||
|
|
||||||
|
# --- 5. Print Results ---
|
||||||
|
print("\n" + "---" * 10)
|
||||||
|
print("--- Benchmark Results (1,000,000 Rows) ---")
|
||||||
|
print(f"Pandas: {pandas_time * 1000:>10.2f} ms")
|
||||||
|
print(f"ConnectorX: {connectorx_time * 1000:>10.2f} ms")
|
||||||
|
print(f"UncheckedIO: {unchecked_io_time * 1000:>10.2f} ms")
|
||||||
|
print("---" * 10)
|
||||||
|
|
||||||
|
print("\n--- Ratios ---")
|
||||||
|
if unchecked_io_time > 0:
|
||||||
|
print(f"UncheckedIO is {pandas_time / unchecked_io_time:.2f}x faster than Pandas")
|
||||||
|
print(f"UncheckedIO is {connectorx_time / unchecked_io_time:.2f}x faster than ConnectorX")
|
||||||
|
else:
|
||||||
|
print("UncheckedIO was too fast to measure accurately!")
|
||||||
49
config.yaml
49
config.yaml
@@ -1,28 +1,25 @@
|
|||||||
connection_string: postgresql://postgres:mysecretpassword@localhost:5433/postgres
|
connection_string: postgresql://postgres:mysecretpassword@localhost:5433/postgres
|
||||||
|
query: COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score::REAL, is_active::BOOLEAN,
|
||||||
# The query to execute. This MUST be a COPY...TO STDOUT (FORMAT binary) command.
|
last_login::TIMESTAMP, notes::TEXT, course_id::INT, start_date::DATE, rating::FLOAT8
|
||||||
# We cast all columns to match the binary types our parser expects.
|
FROM benchmark_table) TO STDOUT (FORMAT binary)
|
||||||
query: "COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score::REAL, is_active::BOOLEAN, last_login::TIMESTAMP, notes::TEXT, course_id::INT, start_date::DATE, rating::FLOAT8 FROM benchmark_table) TO STDOUT (FORMAT binary)"
|
|
||||||
|
|
||||||
# 'schema' is the ordered list of columns *exactly* as they appear in the query's SELECT statement.
|
|
||||||
schema:
|
schema:
|
||||||
- column_name: id
|
- arrow_type: Int64
|
||||||
arrow_type: Int64
|
column_name: id
|
||||||
- column_name: uuid
|
- arrow_type: Utf8
|
||||||
arrow_type: Utf8
|
column_name: uuid
|
||||||
- column_name: username
|
- arrow_type: Utf8
|
||||||
arrow_type: Utf8
|
column_name: username
|
||||||
- column_name: score
|
- arrow_type: Float32
|
||||||
arrow_type: Float32
|
column_name: score
|
||||||
- column_name: is_active
|
- arrow_type: Boolean
|
||||||
arrow_type: Boolean
|
column_name: is_active
|
||||||
- column_name: last_login
|
- arrow_type: Timestamp(Nanosecond, None)
|
||||||
arrow_type: Timestamp(Nanosecond, None)
|
column_name: last_login
|
||||||
- column_name: notes
|
- arrow_type: Utf8
|
||||||
arrow_type: Utf8
|
column_name: notes
|
||||||
- column_name: course_id
|
- arrow_type: Int32
|
||||||
arrow_type: Int32
|
column_name: course_id
|
||||||
- column_name: start_date
|
- arrow_type: Date32
|
||||||
arrow_type: Date32
|
column_name: start_date
|
||||||
- column_name: rating
|
- arrow_type: Float64
|
||||||
arrow_type: Float64
|
column_name: rating
|
||||||
|
|||||||
105
docs/sprint2.md
Normal file
105
docs/sprint2.md
Normal file
@@ -0,0 +1,105 @@
|
|||||||
|
Sprint 2: Streaming Parser
|
||||||
|
|
||||||
|
# 🏁 MVP Sprint 2: Streaming Parser Implementation
|
||||||
|
|
||||||
|
**Goal:** Refactor `src/parser.rs` to parse the Postgres binary stream _as it arrives_, rather than buffering the entire 130MB+ result set into a `Vec<u8>` first. This will significantly reduce memory overhead and improve processing latency.
|
||||||
|
|
||||||
|
### Phase 1: Refactor `handle_binary_copy` State
|
||||||
|
|
||||||
|
This phase involves changing the _structure_ of the function to support streaming.
|
||||||
|
|
||||||
|
- [ ] **Task 1: Initialize Builders First**
|
||||||
|
|
||||||
|
- In `handle_binary_copy`, move the `builders: Vec<DynamicBuilder>` creation to the _top_ of the function, before the stream reading loop.
|
||||||
|
|
||||||
|
- [ ] **Task 2: Initialize State Variables**
|
||||||
|
|
||||||
|
- Create a `leftover_buffer: Vec<u8> = Vec::new();` to hold partial row data between chunks.
|
||||||
|
|
||||||
|
- Create a `is_header_parsed: bool = false;` flag.
|
||||||
|
|
||||||
|
- Create a `rows_processed: usize = 0;` counter.
|
||||||
|
|
||||||
|
- [ ] **Task 3: Refactor Finalization**
|
||||||
|
|
||||||
|
- Move the "Finalize the Arrays" and "Build the Final RecordBatch" logic (Steps 4 and 5 in the old parser) to the _very end_ of the function, after the main stream loop.
|
||||||
|
|
||||||
|
|
||||||
|
### Phase 2: Implement the Streaming Loop
|
||||||
|
|
||||||
|
This phase rewrites the core `while let Some...` loop to parse _in-place_.
|
||||||
|
|
||||||
|
- [ ] **Task 4: Prepare the Chunk Buffer**
|
||||||
|
|
||||||
|
- Inside the `while let Some(segment_result) = stream.next().await` loop:
|
||||||
|
|
||||||
|
- Create a _new_ `current_chunk: Vec<u8>` by combining `leftover_buffer` + the new `segment` (from `segment.extend_from_slice(...)`).
|
||||||
|
|
||||||
|
- Clear the `leftover_buffer` (`leftover_buffer.clear()`).
|
||||||
|
|
||||||
|
- [ ] **Task 5: Create a Cursor for the Chunk**
|
||||||
|
|
||||||
|
- Create a `cursor = Cursor::new(¤t_chunk)`.
|
||||||
|
|
||||||
|
- [ ] **Task 6: Handle the Header (Once)**
|
||||||
|
|
||||||
|
- Add an `if !is_header_parsed` block.
|
||||||
|
|
||||||
|
- Inside this block, parse the 19-byte Postgres header (magic signature, flags) using the `cursor`.
|
||||||
|
|
||||||
|
- If the header is incomplete (hits `UnexpectedEof`), copy the _entire_ `current_chunk` back into `leftover_buffer` and `continue` to the next segment.
|
||||||
|
|
||||||
|
- If successful, set `is_header_parsed = true`.
|
||||||
|
|
||||||
|
|
||||||
|
### Phase 3: Implement the Inner Parsing Loop (The Core Logic)
|
||||||
|
|
||||||
|
This is the most complex part: parsing rows until the chunk runs out of data.
|
||||||
|
|
||||||
|
- [ ] **Task 7: Create the Inner `loop`**
|
||||||
|
|
||||||
|
- Inside the `while` loop (after the header check), create a new, inner `loop { ... }`. This loop will repeatedly try to parse rows from the `cursor`.
|
||||||
|
|
||||||
|
- [ ] **Task 8: Implement Safe Read (Handling Partial Data)**
|
||||||
|
|
||||||
|
- Before attempting to read a row, save the current position: `let safe_position = cursor.position();`.
|
||||||
|
|
||||||
|
- Attempt to read the row's 2-byte column count (`read_i16`).
|
||||||
|
|
||||||
|
- **If it fails with `UnexpectedEof`:**
|
||||||
|
|
||||||
|
- This chunk is finished, but the data is partial.
|
||||||
|
|
||||||
|
- Copy the remaining bytes from `safe_position` to the end of `current_chunk` into `leftover_buffer`.
|
||||||
|
|
||||||
|
- `break;` the _inner_ loop (to go get the next network segment).
|
||||||
|
|
||||||
|
- **If it succeeds:**
|
||||||
|
|
||||||
|
- Check for the `-1` trailer (end of stream). If found, `break` _both_ loops.
|
||||||
|
|
||||||
|
- Increment `rows_processed`.
|
||||||
|
|
||||||
|
- [ ] **Task 9: Implement Streaming Field Parsing**
|
||||||
|
|
||||||
|
- Inside the `for (i, builder) in builders.iter_mut().enumerate()` loop:
|
||||||
|
|
||||||
|
- Save the position _again_ (`let field_safe_pos = cursor.position();`).
|
||||||
|
|
||||||
|
- Attempt to read the 4-byte field length.
|
||||||
|
|
||||||
|
- Attempt to read the field data (e.g., `read_i64`).
|
||||||
|
|
||||||
|
- **If** _**any**_ **read fails with `UnexpectedEof`:**
|
||||||
|
|
||||||
|
- The row is partial and split across chunks.
|
||||||
|
|
||||||
|
- Restore the cursor: `cursor.set_position(safe_position);` (rewind to the start of the row).
|
||||||
|
|
||||||
|
- Copy _all_ remaining bytes from `safe_position` into `leftover_buffer`.
|
||||||
|
|
||||||
|
- `break;` the _inner_ loop.
|
||||||
|
|
||||||
|
- **If it succeeds:**
|
||||||
|
|
||||||
|
- Append the value to the correct Arrow builder.
|
||||||
315
src/parser.rs
315
src/parser.rs
@@ -20,13 +20,14 @@ use byteorder::{BigEndian, ReadBytesExt};
|
|||||||
use std::io::{Cursor, Read};
|
use std::io::{Cursor, Read};
|
||||||
use std::str;
|
use std::str;
|
||||||
use chrono::{NaiveDateTime, NaiveDate};
|
use chrono::{NaiveDateTime, NaiveDate};
|
||||||
|
use std::mem; // <-- We need this for mem::take
|
||||||
|
|
||||||
// --- Internal Crates ---
|
// --- Internal Crates ---
|
||||||
use crate::config::ConnectorConfig;
|
use crate::config::ConnectorConfig;
|
||||||
|
|
||||||
|
|
||||||
// --- 1. CORE DATABASE LOGIC ---
|
// --- 1. CORE DATABASE LOGIC (Public API) ---
|
||||||
// FIX: Change return type from Result<()> to Result<RecordBatch>
|
// This function remains unchanged.
|
||||||
pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
||||||
// 1. Establish the connection
|
// 1. Establish the connection
|
||||||
println!("UncheckedIO: Attempting connection...");
|
println!("UncheckedIO: Attempting connection...");
|
||||||
@@ -70,9 +71,10 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
|||||||
let arrow_schema = Arc::new(Schema::new(schema_fields));
|
let arrow_schema = Arc::new(Schema::new(schema_fields));
|
||||||
|
|
||||||
// Call the parser
|
// Call the parser
|
||||||
let (rows_processed, record_batch) = handle_binary_copy(pinned_stream, arrow_schema.clone()).await?;
|
// 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.", rows_processed);
|
println!("UncheckedIO: Successfully parsed {} rows via binary stream.", record_batch.num_rows());
|
||||||
|
|
||||||
// 4. Final Output Confirmation
|
// 4. Final Output Confirmation
|
||||||
println!("UncheckedIO: Built RecordBatch with {} rows and {} columns.",
|
println!("UncheckedIO: Built RecordBatch with {} rows and {} columns.",
|
||||||
@@ -80,12 +82,12 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
|
|||||||
|
|
||||||
println!("UncheckedIO: Data transfer complete. We lived.");
|
println!("UncheckedIO: Data transfer complete. We lived.");
|
||||||
|
|
||||||
// FIX: Return the finished RecordBatch
|
|
||||||
Ok(record_batch)
|
Ok(record_batch)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
// --- 2. BINARY STREAM HANDLER (THE DYNAMIC PARSER) ---
|
// --- 2. INTERNAL PARSER IMPLEMENTATION ---
|
||||||
|
// All the complex logic is now contained in this private section.
|
||||||
|
|
||||||
// This enum will hold our different builder types
|
// This enum will hold our different builder types
|
||||||
enum DynamicBuilder {
|
enum DynamicBuilder {
|
||||||
@@ -104,41 +106,14 @@ const POSTGRES_EPOCH_NAIVE: NaiveDateTime = NaiveDate::from_ymd_opt(2000, 1, 1).
|
|||||||
// Unix Epoch for dates
|
// Unix Epoch for dates
|
||||||
const UNIX_EPOCH_NAIVE_DATE: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
|
const UNIX_EPOCH_NAIVE_DATE: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
|
||||||
|
|
||||||
// This function is now private to this module (it's not 'pub')
|
/// This is the refactored streaming state machine.
|
||||||
async fn handle_binary_copy(
|
/// It reads the 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>
|
arrow_schema: Arc<Schema>
|
||||||
) -> Result<(usize, RecordBatch)> {
|
) -> Result<RecordBatch> {
|
||||||
|
|
||||||
let mut binary_buffer: Vec<u8> = Vec::new();
|
// --- Phase 1: Task 1 (Initialize Builders) ---
|
||||||
|
|
||||||
// 1. Aggregate all bytes from the stream into our buffer
|
|
||||||
while let Some(segment_result) = stream.next().await {
|
|
||||||
let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?;
|
|
||||||
binary_buffer.extend_from_slice(&segment);
|
|
||||||
}
|
|
||||||
|
|
||||||
let bytes_received = binary_buffer.len();
|
|
||||||
if bytes_received == 0 {
|
|
||||||
return Err(anyhow!("Received zero bytes from COPY stream. Check query and permissions."));
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("UncheckedIO: Received {} total bytes from stream.", bytes_received);
|
|
||||||
|
|
||||||
// --- PARSE POSTGRES BINARY HEADER ---
|
|
||||||
let mut cursor = Cursor::new(&binary_buffer[..]);
|
|
||||||
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" {
|
|
||||||
return Err(anyhow!("Invalid Postgres COPY binary signature."));
|
|
||||||
}
|
|
||||||
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.");
|
|
||||||
|
|
||||||
// --- DYNAMIC FIELD DESERIALIZATION ---
|
|
||||||
|
|
||||||
// 1. Create a dynamic list of builders based on the schema
|
|
||||||
let mut builders: Vec<DynamicBuilder> = arrow_schema.fields().iter().map(|field| {
|
let mut builders: Vec<DynamicBuilder> = arrow_schema.fields().iter().map(|field| {
|
||||||
match field.data_type() {
|
match field.data_type() {
|
||||||
DataType::Int64 => DynamicBuilder::Int64(Box::new(Int64Builder::new())),
|
DataType::Int64 => DynamicBuilder::Int64(Box::new(Int64Builder::new())),
|
||||||
@@ -151,102 +126,100 @@ async fn handle_binary_copy(
|
|||||||
DynamicBuilder::Timestamp(Box::new(TimestampNanosecondBuilder::new()))
|
DynamicBuilder::Timestamp(Box::new(TimestampNanosecondBuilder::new()))
|
||||||
},
|
},
|
||||||
DataType::Date32 => DynamicBuilder::Date32(Box::new(Date32Builder::new())),
|
DataType::Date32 => DynamicBuilder::Date32(Box::new(Date32Builder::new())),
|
||||||
_ => panic!("Unsupported type in builder creation!"),
|
_ => panic!("Unsupported type in builder creation!"), // Will improve later
|
||||||
}
|
}
|
||||||
}).collect();
|
}).collect();
|
||||||
|
|
||||||
let mut rows_processed = 0;
|
// --- 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;
|
||||||
|
|
||||||
// 2. Loop through the rest of the buffer until we hit the 2-byte trailer
|
// --- Phase 2: Task 4 (Streaming Loop) ---
|
||||||
loop {
|
'stream_loop: while let Some(segment_result) = stream.next().await {
|
||||||
let col_count = match cursor.read_i16::<BigEndian>() {
|
let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?;
|
||||||
Ok(count) => count,
|
|
||||||
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
|
// Combine leftover bytes from last chunk with the new chunk
|
||||||
break; // Reached end of buffer *before* trailer, assume it's the end
|
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(¤t_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
|
||||||
}
|
}
|
||||||
Err(e) => return Err(e.into()),
|
|
||||||
};
|
|
||||||
|
|
||||||
// Check for the 2-byte trailer (-1) which signals the end of the data
|
parse_stream_header(&mut cursor)?;
|
||||||
if col_count == -1 {
|
is_header_parsed = true;
|
||||||
println!("UncheckedIO: Reached end-of-stream trailer.");
|
|
||||||
break; // End of stream
|
|
||||||
}
|
}
|
||||||
|
|
||||||
rows_processed += 1;
|
// --- 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();
|
||||||
|
|
||||||
// 3. Dynamic Row-Parsing Loop
|
// 1. Try to read the 2-byte row header (column count)
|
||||||
for (i, builder) in builders.iter_mut().enumerate() {
|
let col_count = match cursor.read_i16::<BigEndian>() {
|
||||||
let field_len_i32 = cursor.read_i32::<BigEndian>()
|
Ok(count) => count,
|
||||||
.context(format!("Failed to read field length for col {}", i))?;
|
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
|
||||||
|
// This is NOT an error. It's a partial read.
|
||||||
if field_len_i32 == -1 {
|
// We don't have enough data for a full row header.
|
||||||
// Handle NULLs
|
// Rewind the cursor to the safe position...
|
||||||
match builder {
|
cursor.set_position(safe_position);
|
||||||
DynamicBuilder::Int64(b) => b.append_null(),
|
// ... and save the remaining bytes for the next chunk.
|
||||||
DynamicBuilder::Float64(b) => b.append_null(),
|
leftover_buffer.extend_from_slice(¤t_chunk[safe_position as usize..]);
|
||||||
DynamicBuilder::Int32(b) => b.append_null(),
|
// Break the *inner* loop to get the next network segment
|
||||||
DynamicBuilder::Float32(b) => b.append_null(),
|
break 'parsing_loop;
|
||||||
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; // Go to the next field
|
Err(e) => return Err(e.into()), // This is a real, unexpected error
|
||||||
|
};
|
||||||
|
|
||||||
|
// --- 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
|
||||||
}
|
}
|
||||||
|
|
||||||
let field_len_usize = field_len_i32 as usize;
|
// 3. Try to parse all fields for this row
|
||||||
|
match parse_row(&mut cursor, &mut builders, ¤t_chunk) {
|
||||||
match builder {
|
Ok(_) => {
|
||||||
DynamicBuilder::Int64(b) => {
|
// Row was parsed successfully
|
||||||
let val = cursor.read_i64::<BigEndian>()?;
|
rows_processed += 1;
|
||||||
b.append_value(val);
|
|
||||||
}
|
}
|
||||||
DynamicBuilder::Int32(b) => {
|
// FIX: 'e' is already a std::io::Error. We just check its .kind() directly.
|
||||||
let val = cursor.read_i32::<BigEndian>()?;
|
// This resolves the E0599 (method not found) errors.
|
||||||
b.append_value(val);
|
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(¤t_chunk[safe_position as usize..]);
|
||||||
|
break 'parsing_loop;
|
||||||
}
|
}
|
||||||
DynamicBuilder::Float64(b) => {
|
Err(e) => {
|
||||||
let val = cursor.read_f64::<BigEndian>()?;
|
// This was a real, unexpected error.
|
||||||
b.append_value(val);
|
// FIX: Convert the std::io::Error into an anyhow::Error.
|
||||||
}
|
// This resolves the E0308 (mismatched types) error.
|
||||||
DynamicBuilder::Float32(b) => {
|
return Err(e.into());
|
||||||
let val = cursor.read_f32::<BigEndian>()?;
|
|
||||||
b.append_value(val);
|
|
||||||
}
|
|
||||||
DynamicBuilder::String(b) => {
|
|
||||||
let mut str_buf = vec![0; field_len_usize];
|
|
||||||
cursor.read_exact(&mut str_buf)?;
|
|
||||||
let val_str = str::from_utf8(&str_buf)
|
|
||||||
.context(format!("Failed to parse UTF-8 string for col {}", i))?;
|
|
||||||
b.append_value(val_str);
|
|
||||||
}
|
|
||||||
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) => {
|
|
||||||
// Postgres binary DATE is i32 days since 2000-01-01
|
|
||||||
let pg_days = cursor.read_i32::<BigEndian>()?;
|
|
||||||
// Arrow Date32 is i32 days since 1970-01-01 (Unix Epoch)
|
|
||||||
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);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
} // 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()));
|
||||||
}
|
}
|
||||||
|
|
||||||
// 4. Finalize the Arrays
|
// Finalize all the builders
|
||||||
let final_columns: Vec<ArrayRef> = builders.into_iter().map(|builder| {
|
let final_columns: Vec<ArrayRef> = builders.into_iter().map(|builder| {
|
||||||
match builder {
|
match builder {
|
||||||
DynamicBuilder::Int64(mut b) => Arc::new(b.finish()) as ArrayRef,
|
DynamicBuilder::Int64(mut b) => Arc::new(b.finish()) as ArrayRef,
|
||||||
@@ -260,11 +233,117 @@ async fn handle_binary_copy(
|
|||||||
}
|
}
|
||||||
}).collect();
|
}).collect();
|
||||||
|
|
||||||
// 5. Build the Final RecordBatch
|
// Build the Final RecordBatch
|
||||||
let record_batch = RecordBatch::try_new(
|
let record_batch = RecordBatch::try_new(
|
||||||
arrow_schema.clone(),
|
arrow_schema.clone(),
|
||||||
final_columns,
|
final_columns,
|
||||||
).context("Failed to create final Arrow RecordBatch")?;
|
).context("Failed to create final Arrow RecordBatch")?;
|
||||||
|
|
||||||
Ok((rows_processed, record_batch))
|
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")?;
|
||||||
|
if &magic_signature != b"PGCOPY\n\xff\r\n\0" {
|
||||||
|
return Err(anyhow!("Invalid Postgres COPY binary signature."));
|
||||||
|
}
|
||||||
|
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
|
||||||
|
|
||||||
|
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(),
|
||||||
|
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; // Go to the next field in this row
|
||||||
|
}
|
||||||
|
|
||||||
|
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>()?;
|
||||||
|
b.append_value(val);
|
||||||
|
}
|
||||||
|
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);
|
||||||
|
}
|
||||||
|
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);
|
||||||
|
}
|
||||||
|
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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} // End field loop
|
||||||
|
|
||||||
|
Ok(()) // Row was successfully parsed
|
||||||
}
|
}
|
||||||
Reference in New Issue
Block a user