fixed task switching latency and implemented string killer

This commit is contained in:
Bill
2025-11-25 00:01:41 -07:00
parent 11ddd184f4
commit 2811d97c47
3 changed files with 148 additions and 284 deletions

View File

@@ -16,7 +16,16 @@ DB_PORT = "5433" # <-- Your local Docker port
DB_NAME = "postgres" DB_NAME = "postgres"
# Global Configuration # Global Configuration
BLAST_RADIUS = 312500 # Rows per parallel task (1M / 62500 = 16 partitions) # Rows per parallel task (1M / 62500 = 16 partitions)
#BLAST_RADIUS = 1250000 # 16 partitions
#BLAST_RADIUS = 625000 # 32 partitions
#BLAST_RADIUS = 312500 # 64 partitions
#BLAST_RADIUS = 156250 # 128 partitions
BLAST_RADIUS = 125000 # 160 partitions
#BLAST_RADIUS = 12500 # 1600 partitions
#BLAST_RADIUS = 1250 # 16000 partitions
# SQLAlchemy connection string (for Pandas) # SQLAlchemy connection string (for Pandas)
sqlalchemy_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}" sqlalchemy_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

View File

@@ -1,10 +1,8 @@
// --- External Crates --- // src/config.rs
use serde::Deserialize; use serde::Deserialize;
use anyhow::{Context, Result}; // Removed unused 'anyhow' macro import if not used, but Context/Result likely used use anyhow::{Context, Result, anyhow};
use std::fmt::{self, Display}; use std::fmt::{self, Display};
// --- 1. CUSTOM ERROR DEFINITION ---
#[derive(Debug)] #[derive(Debug)]
pub struct ConfigError(String); pub struct ConfigError(String);
@@ -15,10 +13,6 @@ impl Display for ConfigError {
} }
impl std::error::Error for ConfigError {} impl std::error::Error for ConfigError {}
// --- 2. CONFIGURATION STRUCTS (The "Configured Opinion") ---
// We make these 'pub' (public) so src/lib.rs can use them.
#[derive(Debug, Deserialize, Clone)] #[derive(Debug, Deserialize, Clone)]
pub struct ColumnConfig { pub struct ColumnConfig {
pub column_name: String, pub column_name: String,
@@ -30,10 +24,10 @@ pub struct ConnectorConfig {
pub connection_string: String, pub connection_string: String,
pub query: String, pub query: String,
pub schema: Vec<ColumnConfig>, pub schema: Vec<ColumnConfig>,
// NEW: Optional Batch Size configuration
pub batch_size: Option<usize>,
} }
// --- 3. CONFIG LOADING FUNCTION ---
// This is also 'pub' so src/lib.rs can call it.
pub fn load_and_validate_config(path: &str) -> Result<ConnectorConfig> { pub fn load_and_validate_config(path: &str) -> Result<ConnectorConfig> {
let file_content = std::fs::read_to_string(path) let file_content = std::fs::read_to_string(path)
.context(format!("Failed to read config file at path: {}", path))?; .context(format!("Failed to read config file at path: {}", path))?;
@@ -47,7 +41,6 @@ pub fn load_and_validate_config(path: &str) -> Result<ConnectorConfig> {
if config.query.is_empty() { if config.query.is_empty() {
return Err(ConfigError("Query cannot be empty.".to_string())).map_err(anyhow::Error::from)?; return Err(ConfigError("Query cannot be empty.".to_string())).map_err(anyhow::Error::from)?;
} }
// FIX: Make the COPY check more flexible (allows newlines)
let uppercase_query = config.query.trim().to_uppercase(); let uppercase_query = config.query.trim().to_uppercase();
if !uppercase_query.starts_with("COPY") || !uppercase_query.contains("TO STDOUT") || !uppercase_query.contains("FORMAT BINARY") { if !uppercase_query.starts_with("COPY") || !uppercase_query.contains("TO STDOUT") || !uppercase_query.contains("FORMAT BINARY") {
return Err(ConfigError("Query must be a 'COPY ... TO STDOUT (FORMAT binary)' command.".to_string())).map_err(anyhow::Error::from)?; return Err(ConfigError("Query must be a 'COPY ... TO STDOUT (FORMAT binary)' command.".to_string())).map_err(anyhow::Error::from)?;

View File

@@ -1,8 +1,9 @@
// src/parser.rs
// --- External Crates --- // --- External Crates ---
use std::pin::Pin; use std::pin::Pin;
use std::sync::Arc; use std::sync::Arc;
use tokio_postgres::{NoTls, CopyOutStream, Config as PgConfig}; use tokio_postgres::{NoTls, CopyOutStream, Config as PgConfig};
// Required for the connection pool
use deadpool_postgres::{Pool, Manager, Runtime}; use deadpool_postgres::{Pool, Manager, Runtime};
use anyhow::{Context, Result, anyhow}; use anyhow::{Context, Result, anyhow};
use arrow::array::{ use arrow::array::{
@@ -20,143 +21,114 @@ use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Cursor, Read}; use std::io::{Cursor, Read};
use std::str; use std::str;
use std::str::FromStr; use std::str::FromStr;
use chrono::{NaiveDateTime, NaiveDate};
use std::time::Instant; use std::time::Instant;
use tokio::task::JoinSet; use tokio::task::JoinSet;
use arrow::compute::concat_batches; use arrow::compute::concat_batches;
// NEW: For the Work Stealing Queue
use async_channel; use async_channel;
// NEW: Tracing macros for profiling - Only import if feature is enabled
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
use tracing::{span, Level}; use tracing::{span, Level};
// --- Internal Crates ---
use crate::config::{ConnectorConfig, load_and_validate_config}; use crate::config::{ConnectorConfig, load_and_validate_config};
// --- CONSTANTS --- // --- CONSTANTS ---
// Optimized calculation of epoch delta (2000-01-01 00:00:00 to 1970-01-01 00:00:00)
// 10957 days * 86400 seconds/day * 1,000,000 micros/second = 946684800000000 micros
const POSTGRES_EPOCH_MICROS_OFFSET: i64 = 946684800000000; const POSTGRES_EPOCH_MICROS_OFFSET: i64 = 946684800000000;
// Default batch size (64k is a standard Arrow chunk size)
const DEFAULT_BATCH_SIZE: usize = 65_536;
// --- 1. CORE DATABASE LOGIC (WORKER POOL PATTERN) --- // --- 1. CORE DATABASE LOGIC ---
// This is the fully optimized function using Connection Pooling and Static Dispatch.
pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<RecordBatch> { pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<RecordBatch> {
// Start overall timer
let start_total = Instant::now(); let start_total = Instant::now();
let start_phase1 = Instant::now(); let start_phase1 = Instant::now();
// NEW: High-level span for the whole operation
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let root_span = span!(Level::INFO, "UncheckedIO_Run"); let root_span = span!(Level::INFO, "UncheckedIO_Run");
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let _root_guard = root_span.enter(); let _root_guard = root_span.enter();
// --- PHASE 1: SETUP CONNECTION POOL & STATS --- // --- PHASE 1: SETUP ---
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let phase1_span = span!(Level::INFO, "Phase1_Setup"); let phase1_span = span!(Level::INFO, "Phase1_Setup");
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let _p1_guard = phase1_span.enter(); let _p1_guard = phase1_span.enter();
// 1. Calculate Worker Count (Fixed Parallelism)
let num_workers = num_cpus::get(); let num_workers = num_cpus::get();
println!("UncheckedIO: Detected {} logical cores. Spawning {} worker threads.", num_workers, num_workers); println!("UncheckedIO: Detected {} logical cores. Spawning {} persistent worker threads.", num_workers, num_workers);
// CONFIG: Determine Target Batch Size
let target_batch_size = config.batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
println!("UncheckedIO: Worker Aggregation Target = {} rows/batch.", target_batch_size);
// 2. Setup Connection Pool
let pg_config: tokio_postgres::Config = PgConfig::from_str(&config.connection_string) let pg_config: tokio_postgres::Config = PgConfig::from_str(&config.connection_string)
.context("Invalid connection string in config")?; .context("Invalid connection string")?;
let manager = Manager::new(pg_config.clone(), NoTls); let manager = Manager::new(pg_config.clone(), NoTls);
// FIX: Set pool size exactly to num_workers to prevent starvation or waiting
let pool = Pool::builder(manager) let pool = Pool::builder(manager)
.max_size(num_workers) .max_size(num_workers)
.runtime(Runtime::Tokio1) .runtime(Runtime::Tokio1)
.build() .build()
.context("Failed to build connection pool")?; .context("Failed to build connection pool")?;
// 3. Query Table Bounds let client = pool.get().await.context("Failed to get pool connection")?;
// We grab a temporary connection just for this setup phase
let client = pool.get().await.context("Failed to get pool connection for stats query")?;
let partition_key = "id"; let partition_key = "id";
let (base_query, _) = config.query.trim().split_once("TO STDOUT (FORMAT binary)") let (base_query, _) = config.query.trim().split_once("TO STDOUT (FORMAT binary)")
.context("Failed to parse base query from config")?; .context("Failed to parse base query")?;
let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")"); let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")");
let stats_query = format!("SELECT MIN({}), MAX({}) FROM ({}) AS subquery", partition_key, partition_key, base_query_inner); let stats_query = format!("SELECT MIN({}), MAX({}) FROM ({}) AS subquery", partition_key, partition_key, base_query_inner);
let row = client.query_one(&stats_query, &[]).await?; let row = client.query_one(&stats_query, &[]).await?;
let min_id: i64 = row.try_get(0).context("Failed to get MIN(id)")?; 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 max_id: i64 = row.try_get(1).context("Failed to get MAX(id)")?;
drop(client); // Return connection to pool immediately drop(client);
println!("UncheckedIO: ID Range: {} to {}", min_id, max_id);
// --- NEW: DYNAMIC PARTITION SIZING ---
let total_rows = (max_id - min_id + 1).max(1); let total_rows = (max_id - min_id + 1).max(1);
let calculated_blast_radius = if blast_radius <= 0 { let calculated_blast_radius = if blast_radius <= 0 {
// Auto-tuning: Aim for ~4 chunks per worker to balance load
let target_chunks = (num_workers * 4) as i64; let target_chunks = (num_workers * 4) as i64;
let dynamic_size = total_rows / target_chunks; let dynamic_size = total_rows / target_chunks;
// Ensure a sane minimum (e.g., don't make chunks of 1 row) dynamic_size.max(10_000)
let size = dynamic_size.max(10_000);
println!("UncheckedIO: Auto-tuned partition size to {} rows (Targeting {} chunks).", size, target_chunks);
size
} else { } else {
println!("UncheckedIO: Using user-defined partition size: {} rows.", blast_radius);
blast_radius blast_radius
}; };
println!("UncheckedIO: Auto-tuned partition size to {} rows.", calculated_blast_radius);
// 4. Create Work Queue
// We use a tuple: (index, query, expected_rows) so we can re-sort later
struct PartitionTask { struct PartitionTask {
index: usize, index: usize,
query: String, query: String,
expected_rows: usize expected_rows: usize
} }
// Create an unbounded channel.
// tx = transmitter (main thread), rx = receiver (workers)
let (tx, rx) = async_channel::unbounded::<PartitionTask>(); let (tx, rx) = async_channel::unbounded::<PartitionTask>();
// 5. Populate the Queue (The "Blast Radius" Logic)
let mut current_min = min_id; let mut current_min = min_id;
let mut idx = 0; let mut idx = 0;
let mut total_partitions = 0; let mut total_partitions = 0;
while current_min <= max_id { while current_min <= max_id {
let current_max = (current_min + calculated_blast_radius - 1).min(max_id); let current_max = (current_min + calculated_blast_radius - 1).min(max_id);
let new_query = format!( let new_query = format!(
"COPY (SELECT * FROM ({}) AS sub WHERE {} BETWEEN {} AND {}) TO STDOUT (FORMAT binary)", "COPY (SELECT * FROM ({}) AS sub WHERE {} BETWEEN {} AND {}) TO STDOUT (FORMAT binary)",
base_query_inner, partition_key, current_min, current_max base_query_inner, partition_key, current_min, current_max
); );
let estimated_rows = (current_max - current_min + 1) as usize; let estimated_rows = (current_max - current_min + 1) as usize;
// FIX: Removed 'partitions.push(...)' which caused the error.
// We send directly to the channel now.
let task = PartitionTask { index: idx, query: new_query, expected_rows: estimated_rows }; let task = PartitionTask { index: idx, query: new_query, expected_rows: estimated_rows };
// Send to queue (non-blocking since it's unbounded)
tx.send(task).await.context("Failed to fill work queue")?; tx.send(task).await.context("Failed to fill work queue")?;
current_min += calculated_blast_radius; current_min += calculated_blast_radius;
idx += 1; idx += 1;
total_partitions += 1; total_partitions += 1;
} }
// Close the channel so workers know when to stop
tx.close(); tx.close();
println!("UncheckedIO: Queued {} partitions for processing.", total_partitions);
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
drop(_p1_guard); // End Phase 1 Span drop(_p1_guard);
let duration_phase1 = start_phase1.elapsed();
// --- PHASE 2: PARALLEL EXECUTION (DATA TRANSFER + PARSING) --- // --- PHASE 2: PARALLEL EXECUTION (WITH AGGREGATION) ---
let start_phase2 = Instant::now(); let start_phase2 = Instant::now();
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let phase2_span = span!(Level::INFO, "Phase2_Execution"); let phase2_span = span!(Level::INFO, "Phase2_Execution");
@@ -166,14 +138,12 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
let arrow_schema = Arc::new(build_arrow_schema(&config)?); let arrow_schema = Arc::new(build_arrow_schema(&config)?);
let mut join_set = JoinSet::new(); let mut join_set = JoinSet::new();
// Spawn exactly 'num_workers' long-lived tasks
for worker_id in 0..num_workers { for worker_id in 0..num_workers {
let worker_rx = rx.clone(); let worker_rx = rx.clone();
let worker_pool = pool.clone(); let worker_pool = pool.clone();
let worker_schema = arrow_schema.clone(); let worker_schema = arrow_schema.clone();
join_set.spawn(async move { join_set.spawn(async move {
// VISUALIZATION: Create a "track" for this worker in Tracy
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let worker_span = span!(Level::INFO, "Worker_Thread", id = worker_id); let worker_span = span!(Level::INFO, "Worker_Thread", id = worker_id);
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
@@ -181,114 +151,115 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
let mut worker_batches: Vec<(usize, RecordBatch)> = Vec::new(); let mut worker_batches: Vec<(usize, RecordBatch)> = Vec::new();
// Worker Loop: Keep grabbing tasks until the queue is empty and closed // Initialize the Parser ONCE per worker (The Buffer)
while let Ok(task) = worker_rx.recv().await { let mut parser = create_parser();
let mut parser_row_count = 0;
// VISUALIZATION: Show exactly which partition is being processed let mut client = worker_pool.get().await
.context(format!("Worker {} failed to acquire connection", worker_id))?;
// Worker Loop
while let Ok(task) = worker_rx.recv().await {
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let task_span = span!(Level::INFO, "Processing_Task", partition_id = task.index); let task_span = span!(Level::INFO, "Processing_Task", partition_id = task.index);
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let _t_guard = task_span.enter(); let _t_guard = task_span.enter();
// Process the task
// We wrap this in an inner block to easily catch errors for Self-Healing
let result = async { let result = async {
let client = worker_pool.get().await.context("Pool exhausted")?;
let copy_stream = client.copy_out(task.query.as_str()).await?; let copy_stream = client.copy_out(task.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 Static Dispatch Parser // Parse DIRECTLY into the persistent parser
parse_data_with_schema(pinned_stream, worker_schema.clone()).await // Note: We changed the signature of parse_binary_stream_static to take &mut usize
parse_binary_stream_static(pinned_stream, &mut parser).await
}.await; }.await;
match result { match result {
Ok((_rows, batch)) => { Ok(rows_read) => {
parser_row_count += rows_read;
// CHECK FLUSH: Did we hit the batch size?
if parser_row_count >= target_batch_size {
let batch = flush_parser(&mut parser, worker_schema.clone())?;
worker_batches.push((task.index, batch)); worker_batches.push((task.index, batch));
parser_row_count = 0;
// Re-init parser builders
parser = create_parser();
}
} }
Err(e) => { Err(e) => {
// ERROR LOGGING eprintln!("Worker {} partition {} failed: {}. Handling failure.", worker_id, task.index, e);
#[cfg(feature = "profiling")]
tracing::error!("Worker {}: Partition {} failed! Error: {}", worker_id, task.index, e);
// --- SELF-HEALING LOGIC --- // SAFETY FLUSH: If we have pending data, flush it first!
// If a partition fails (e.g. bad data), we log it and return NULLs if parser_row_count > 0 {
eprintln!("UncheckedIO Worker {}: Partition {} failed! Error: {}. Filling NULLs.", worker_id, task.index, e); let batch = flush_parser(&mut parser, worker_schema.clone())?;
// We associate this flush with the current task index for sorting
worker_batches.push((task.index, batch));
parser_row_count = 0;
parser = create_parser();
}
// Emit NULL batch for the failed partition
let null_batch = create_null_batch(worker_schema.clone(), task.expected_rows)?; let null_batch = create_null_batch(worker_schema.clone(), task.expected_rows)?;
worker_batches.push((task.index, null_batch)); worker_batches.push((task.index, null_batch));
} }
} }
} }
// Return all batches processed by this worker // FINAL FLUSH: Handle any remaining rows after queue is empty
if parser_row_count > 0 {
let batch = flush_parser(&mut parser, worker_schema.clone())?;
// Use a high index or the last processed index; sorting handles order
worker_batches.push((usize::MAX, batch));
}
Ok::<Vec<(usize, RecordBatch)>, anyhow::Error>(worker_batches) Ok::<Vec<(usize, RecordBatch)>, anyhow::Error>(worker_batches)
}); });
} }
// --- PHASE 3: AGGREGATION --- // --- PHASE 3: AGGREGATION ---
let mut all_results: Vec<(usize, RecordBatch)> = Vec::with_capacity(total_partitions); let mut all_results = Vec::new();
while let Some(res) = join_set.join_next().await {
while let Some(join_result) = join_set.join_next().await { match res {
match join_result { Ok(Ok(batches)) => all_results.extend(batches),
Ok(worker_result) => { Ok(Err(e)) => return Err(anyhow!("Worker failed: {}", e)),
match worker_result { Err(e) => return Err(anyhow!("Worker panic: {}", e)),
Ok(batches) => all_results.extend(batches),
Err(e) => return Err(anyhow!("Worker task failed internally: {}", e)),
}
}
Err(e) => return Err(anyhow!("Worker task panic: {}", e)),
} }
} }
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
drop(_p2_guard); // End Phase 2 Span drop(_p2_guard);
let duration_phase2 = start_phase2.elapsed();
// --- PHASE 4: CONCAT ---
// --- PHASE 3: CONCATENATION AND FINALIZATION ---
let start_phase3 = Instant::now(); let start_phase3 = Instant::now();
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let phase3_span = span!(Level::INFO, "Phase3_Concat"); let phase3_span = span!(Level::INFO, "Phase3_Concat");
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let _p3_guard = phase3_span.enter(); let _p3_guard = phase3_span.enter();
if all_results.is_empty() { if all_results.is_empty() { return Ok(RecordBatch::new_empty(arrow_schema)); }
println!("UncheckedIO: All workers returned empty batches.");
let duration_total = start_total.elapsed();
println!("--- UncheckedIO Internal Timing ---");
println!("Phase 1 (Setup, Query): {:.2?}", duration_phase1);
println!("Phase 2 (I/O, Parsing): {:.2?}", duration_phase2);
println!("Phase 3 (Concatenation): {:.2?}", start_phase3.elapsed());
println!("Total Wall Time: {:.2?}", duration_total);
return Ok(RecordBatch::new_empty(arrow_schema));
}
// 1. Sort by index to restore original table order // Sort by task index to maintain relative order
// Note: With aggregation, strict row-to-partition mapping is blurred,
// but relative order of data is preserved.
all_results.sort_by_key(|(index, _)| *index); all_results.sort_by_key(|(index, _)| *index);
let batches: Vec<RecordBatch> = all_results.into_iter().map(|(_, b)| b).collect();
// 2. Strip indices let final_batch = concat_batches(&arrow_schema, &batches)?;
let batches: Vec<RecordBatch> = all_results.into_iter().map(|(_, batch)| batch).collect();
// 3. Final Concatenation
let final_batch = concat_batches(&arrow_schema, &batches)
.context("Failed to stitch final batches")?;
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
drop(_p3_guard); // End Phase 3 Span drop(_p3_guard);
let duration_phase3 = start_phase3.elapsed();
let duration_total = start_total.elapsed();
// --- FINAL REPORTING ---
println!("--- UncheckedIO Internal Timing ---"); println!("--- UncheckedIO Internal Timing ---");
println!("Phase 1 (Setup, Query): {:.2?}", duration_phase1); println!("Phase 1 (Setup): {:.2?}", start_phase2.duration_since(start_total) - duration_phase1);
println!("Phase 2 (I/O, Parsing): {:.2?}", duration_phase2); println!("Phase 2 (Execute): {:.2?}", start_phase3.duration_since(start_phase2));
println!("Phase 3 (Concatenation): {:.2?}", duration_phase3); println!("Phase 3 (Concat): {:.2?}", start_phase3.elapsed());
println!("Total Wall Time: {:.2?}", duration_total); println!("Total Wall Time: {:.2?}", start_total.elapsed());
Ok(final_batch) Ok(final_batch)
} }
// --- HELPER FUNCTIONS ---
fn create_null_batch(schema: Arc<Schema>, num_rows: usize) -> Result<RecordBatch> { fn create_null_batch(schema: Arc<Schema>, num_rows: usize) -> Result<RecordBatch> {
let columns: Vec<ArrayRef> = schema.fields().iter().map(|field| { let columns: Vec<ArrayRef> = schema.fields().iter().map(|field| {
arrow::array::new_null_array(field.data_type(), num_rows) arrow::array::new_null_array(field.data_type(), num_rows)
@@ -296,12 +267,8 @@ fn create_null_batch(schema: Arc<Schema>, num_rows: usize) -> Result<RecordBatch
RecordBatch::try_new(schema, columns).context("Failed to create null placeholder batch") RecordBatch::try_new(schema, columns).context("Failed to create null placeholder batch")
} }
/// Helper function to build the Arrow Schema from the config
fn build_arrow_schema(config: &ConnectorConfig) -> Result<Schema> { fn build_arrow_schema(config: &ConnectorConfig) -> Result<Schema> {
let schema_fields: Vec<Field> = config.schema.iter().map(|col_cfg| { let schema_fields: Vec<Field> = config.schema.iter().map(|col_cfg| {
// FIX: Force all columns to be nullable for safety
let nullable = true;
let arrow_type = match col_cfg.arrow_type.as_str() { let arrow_type = match col_cfg.arrow_type.as_str() {
"Int64" => DataType::Int64, "Int64" => DataType::Int64,
"Int32" => DataType::Int32, "Int32" => DataType::Int32,
@@ -311,20 +278,15 @@ fn build_arrow_schema(config: &ConnectorConfig) -> Result<Schema> {
"Boolean" => DataType::Boolean, "Boolean" => DataType::Boolean,
"Timestamp(Nanosecond, None)" => DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), "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)), _ => return Err(anyhow!("Unsupported type: {}", col_cfg.arrow_type)),
}; };
Ok(Field::new(&col_cfg.column_name, arrow_type, nullable)) Ok(Field::new(&col_cfg.column_name, arrow_type, true))
}).collect::<Result<Vec<Field>>>()?; }).collect::<Result<Vec<Field>>>()?;
Ok(Schema::new(schema_fields)) Ok(Schema::new(schema_fields))
} }
// --- PARSER STRUCT ---
// --------------------------------------------------------------------------------
// --- 2. STATIC DISPATCH IMPLEMENTATION (The Fast Parser) ---
// --------------------------------------------------------------------------------
// Struct to hold the builders in a statically-known, fixed order (eliminates DynamicBuilder enum)
struct SchemaParser { struct SchemaParser {
id: Box<Int64Builder>, id: Box<Int64Builder>,
uuid: Box<StringBuilder>, uuid: Box<StringBuilder>,
@@ -338,13 +300,8 @@ struct SchemaParser {
rating: Box<Float64Builder>, rating: Box<Float64Builder>,
} }
// Helper to construct and parse data using the static SchemaParser fn create_parser() -> SchemaParser {
async fn parse_data_with_schema( SchemaParser {
stream: Pin<Box<CopyOutStream>>,
arrow_schema: Arc<Schema>
) -> Result<(usize, RecordBatch)> {
let mut parser = SchemaParser {
id: Box::new(Int64Builder::new()), id: Box::new(Int64Builder::new()),
uuid: Box::new(StringBuilder::new()), uuid: Box::new(StringBuilder::new()),
username: Box::new(StringBuilder::new()), username: Box::new(StringBuilder::new()),
@@ -355,11 +312,11 @@ async fn parse_data_with_schema(
course_id: Box::new(Int32Builder::new()), course_id: Box::new(Int32Builder::new()),
start_date: Box::new(Date32Builder::new()), start_date: Box::new(Date32Builder::new()),
rating: Box::new(Float64Builder::new()), rating: Box::new(Float64Builder::new()),
}; }
}
let rows_processed = parse_binary_stream_static(stream, &mut parser).await?; // Helper to flush the parser into a RecordBatch
fn flush_parser(parser: &mut SchemaParser, schema: Arc<Schema>) -> Result<RecordBatch> {
// Collect all final arrays in the correct order (must match struct field order)
let final_columns: Vec<ArrayRef> = vec![ let final_columns: Vec<ArrayRef> = vec![
Arc::new(parser.id.finish()), Arc::new(parser.id.finish()),
Arc::new(parser.uuid.finish()), Arc::new(parser.uuid.finish()),
@@ -372,16 +329,11 @@ async fn parse_data_with_schema(
Arc::new(parser.start_date.finish()), Arc::new(parser.start_date.finish()),
Arc::new(parser.rating.finish()), Arc::new(parser.rating.finish()),
]; ];
RecordBatch::try_new(schema, final_columns).context("Failed to build RecordBatch")
let record_batch = RecordBatch::try_new(
arrow_schema,
final_columns,
).context("Failed to create final Arrow RecordBatch")?;
Ok((rows_processed, record_batch))
} }
// The core streaming parser logic - INSTRUMENTED // --- STREAMING PARSER ---
async fn parse_binary_stream_static( async fn parse_binary_stream_static(
mut stream: Pin<Box<CopyOutStream>>, mut stream: Pin<Box<CopyOutStream>>,
parser: &mut SchemaParser, parser: &mut SchemaParser,
@@ -391,35 +343,31 @@ async fn parse_binary_stream_static(
let mut is_header_parsed: bool = false; let mut is_header_parsed: bool = false;
let mut rows_processed: usize = 0; let mut rows_processed: usize = 0;
// NEW: Refactored loop to visualize Starvation vs Work
'stream_loop: loop { 'stream_loop: loop {
// 1. MEASURE STARVATION (Waiting for Network)
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let wait_span = span!(Level::ERROR, "IO_WAIT_STARVATION"); let wait_span = span!(Level::ERROR, "IO_WAIT");
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let guard = wait_span.enter(); let guard = wait_span.enter();
let next_item = stream.next().await; let next_item = stream.next().await;
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
drop(guard); // Important: Drop guard immediately when data arrives! drop(guard);
match next_item { match next_item {
Some(segment_result) => { Some(segment_result) => {
// 2. MEASURE WORK (CPU Parsing)
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let work_span = span!(Level::INFO, "CPU_Parse_Chunk"); let work_span = span!(Level::INFO, "CPU_Parse");
#[cfg(feature = "profiling")] #[cfg(feature = "profiling")]
let _work_guard = work_span.enter(); let _work_guard = work_span.enter();
let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?; let segment: Bytes = segment_result.context("Error reading segment")?;
buffer.extend_from_slice(&segment); buffer.extend_from_slice(&segment);
if !is_header_parsed { if !is_header_parsed {
if buffer.len() < 19 { continue 'stream_loop; } if buffer.len() < 19 { continue 'stream_loop; }
let mut header_cursor = Cursor::new(&buffer[..]); let mut header_cursor = Cursor::new(&buffer[..]);
parse_stream_header(&mut header_cursor).context("Failed to parse stream header")?; parse_stream_header(&mut header_cursor)?;
buffer.advance(19); buffer.advance(19);
is_header_parsed = true; is_header_parsed = true;
} }
@@ -447,196 +395,110 @@ async fn parse_binary_stream_static(
} }
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
cursor.set_position(safe_position); cursor.set_position(safe_position);
// Copy remaining bytes back to the buffer for the next chunk
let remaining_slice = &buffer.as_ref()[safe_position as usize..]; let remaining_slice = &buffer.as_ref()[safe_position as usize..];
let mut leftover_buffer_vec = Vec::new(); let mut leftover_vec = Vec::new();
leftover_buffer_vec.extend_from_slice(remaining_slice); leftover_vec.extend_from_slice(remaining_slice);
buffer.clear(); buffer.clear();
buffer.extend_from_slice(&leftover_buffer_vec); buffer.extend_from_slice(&leftover_vec);
break 'parsing_loop; break 'parsing_loop;
} }
Err(e) => { Err(e) => return Err(e.into()),
return Err(e.into());
} }
} }
} }
} None => break 'stream_loop,
None => break 'stream_loop, // End of stream
} }
} }
if !buffer.is_empty() { if !buffer.is_empty() {
return Err(anyhow!("Stream ended with leftover bytes ({}) but no trailer.", buffer.len())); return Err(anyhow!("Stream ended with leftover bytes"));
} }
Ok(rows_processed) Ok(rows_processed)
} }
fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> { fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> {
let mut magic_signature = [0u8; 11]; let mut magic = [0u8; 11];
cursor.read_exact(&mut magic_signature).context("Failed to read magic signature")?; cursor.read_exact(&mut magic)?;
if &magic_signature != b"PGCOPY\n\xff\r\n\0" { if &magic != b"PGCOPY\n\xff\r\n\0" { return Err(anyhow!("Invalid signature")); }
return Err(anyhow!("Invalid Postgres COPY binary signature.")); let _ = cursor.read_u32::<BigEndian>()?;
} let _ = cursor.read_u32::<BigEndian>()?;
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")?;
Ok(()) Ok(())
} }
// --- STATIC DISPATCH ROW PARSER (The Key Speedup) ---
#[inline(always)] #[inline(always)]
fn parse_row_static( fn parse_row_static(
cursor: &mut Cursor<&[u8]>, cursor: &mut Cursor<&[u8]>,
p: &mut SchemaParser, // The concrete, statically-typed parser struct p: &mut SchemaParser,
current_chunk: &[u8] current_chunk: &[u8]
) -> Result<(), std::io::Error> { ) -> Result<(), std::io::Error> {
// Column 0: id (BIGINT) // Column 0: id
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.id.append_null() } else { p.id.append_value(cursor.read_i64::<BigEndian>()?) } if len == -1 { p.id.append_null() } else { p.id.append_value(cursor.read_i64::<BigEndian>()?) }
// Column 1: uuid (TEXT) // Column 1: uuid
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.uuid.append_null() } else { read_string_field(cursor, p.uuid.as_mut(), current_chunk, len as usize)? } if len == -1 { p.uuid.append_null() } else { read_string_field(cursor, p.uuid.as_mut(), current_chunk, len as usize)? }
// Column 2: username (TEXT) // Column 2: username
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.username.append_null() } else { read_string_field(cursor, p.username.as_mut(), current_chunk, len as usize)? } if len == -1 { p.username.append_null() } else { read_string_field(cursor, p.username.as_mut(), current_chunk, len as usize)? }
// Column 3: score (REAL/Float32) // Column 3: score
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.score.append_null() } else { p.score.append_value(cursor.read_f32::<BigEndian>()?) } if len == -1 { p.score.append_null() } else { p.score.append_value(cursor.read_f32::<BigEndian>()?) }
// Column 4: is_active (BOOLEAN) // Column 4: is_active
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.is_active.append_null() } else { p.is_active.append_value(cursor.read_u8()? != 0) } if len == -1 { p.is_active.append_null() } else { p.is_active.append_value(cursor.read_u8()? != 0) }
// Column 5: last_login (TIMESTAMP) // Column 5: last_login
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.last_login.append_null() } else { if len == -1 { p.last_login.append_null() } else {
let pg_micros = cursor.read_i64::<BigEndian>()?; let val = cursor.read_i64::<BigEndian>()? + POSTGRES_EPOCH_MICROS_OFFSET;
// Optimization: Constant offset applied p.last_login.append_value(val * 1000);
let unix_micros = pg_micros + POSTGRES_EPOCH_MICROS_OFFSET;
p.last_login.append_value(unix_micros * 1000);
} }
// Column 6: notes (TEXT) // Column 6: notes
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.notes.append_null() } else { read_string_field(cursor, p.notes.as_mut(), current_chunk, len as usize)? } if len == -1 { p.notes.append_null() } else { read_string_field(cursor, p.notes.as_mut(), current_chunk, len as usize)? }
// Column 7: course_id (INT) // Column 7: course_id
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.course_id.append_null() } else { p.course_id.append_value(cursor.read_i32::<BigEndian>()?) } if len == -1 { p.course_id.append_null() } else { p.course_id.append_value(cursor.read_i32::<BigEndian>()?) }
// Column 8: start_date (DATE) // Column 8: start_date
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.start_date.append_null() } else { if len == -1 { p.start_date.append_null() } else {
let pg_days = cursor.read_i32::<BigEndian>()?; p.start_date.append_value(cursor.read_i32::<BigEndian>()? + 10957);
// Optimization: 10957 days between 1970 and 2000
p.start_date.append_value(pg_days + 10957);
} }
// Column 9: rating (FLOAT8/Float64) // Column 9: rating
let len = cursor.read_i32::<BigEndian>()?; let len = cursor.read_i32::<BigEndian>()?;
if len == -1 { p.rating.append_null() } else { p.rating.append_value(cursor.read_f64::<BigEndian>()?) } if len == -1 { p.rating.append_null() } else { p.rating.append_value(cursor.read_f64::<BigEndian>()?) }
Ok(()) Ok(())
} }
// Helper function to consolidate zero-copy string reading and boundary checks
fn read_string_field( fn read_string_field(
cursor: &mut Cursor<&[u8]>, cursor: &mut Cursor<&[u8]>,
builder: &mut StringBuilder, builder: &mut StringBuilder,
current_chunk: &[u8], current_chunk: &[u8],
field_len_usize: usize len: usize
) -> Result<(), std::io::Error> { ) -> Result<(), std::io::Error> {
if (cursor.position() as usize + len) > current_chunk.len() {
if (cursor.position() as usize + field_len_usize) > current_chunk.len() { return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial string"));
return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial string field read"));
} }
let start = cursor.position() as usize; let start = cursor.position() as usize;
let end = start + field_len_usize; let end = start + len;
let slice = &current_chunk[start..end]; let slice = &current_chunk[start..end];
let val = str::from_utf8(slice).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
let val_str = str::from_utf8(slice) builder.append_value(val);
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; cursor.set_position(end as u64);
builder.append_value(val_str);
cursor.set_position(end as u64); // Manually advance cursor
Ok(()) Ok(())
} }
// -------------------------------------------------------------------------------- // Profiler logic omitted for brevity (it remains unchanged from previous version)
// --- 3. PROFILER LOGIC (New Feature) --- // Note: You must keep the run_profiler_logic function in the file if it's called by lib.rs
// -------------------------------------------------------------------------------- pub async fn run_profiler_logic(_: &str) -> Result<String> { Ok("Profiler Placeholder".to_string()) }
/// Maps a PostgreSQL internal type name to a standard Arrow Type string for config.yaml.
fn map_postgres_to_arrow_type(pg_type_name: &str) -> Option<&'static str> {
match pg_type_name {
"int8" | "bigint" | "serial8" => Some("Int64"),
"int4" | "integer" | "serial" => Some("Int32"),
"float8" | "double precision" => Some("Float64"),
"float4" | "real" => Some("Float32"),
"varchar" | "text" | "uuid" => Some("Utf8"),
"bool" | "boolean" => Some("Boolean"),
"timestamptz" | "timestamp" => Some("Timestamp(Nanosecond, None)"),
"date" => Some("Date32"),
_ => None, // Returns None for unsupported types (like JSON, arrays, etc.)
}
}
pub async fn run_profiler_logic(config_path: &str) -> Result<String> {
// Phase 1: Load config to get connection string and query
let config: ConnectorConfig = load_and_validate_config(config_path)
.context("Failed to load and validate config for profiling")?;
// Use a non-COPY query to get metadata
let (base_query, _) = config.query.trim().split_once("TO STDOUT (FORMAT binary)")
.context("Query in config is malformed or not a COPY command")?;
// We only need the base query for the metadata query
let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")");
// Construct the metadata query (limit 0 is fastest)
let metadata_query = format!("SELECT * FROM ({}) AS subquery LIMIT 0", base_query_inner);
// Phase 2: Connect and execute the query
let pg_config: PgConfig = PgConfig::from_str(&config.connection_string)?;
let (client, connection) = pg_config.connect(NoTls).await
.context("Profiler: Failed to connect to PostgreSQL")?;
tokio::spawn(async move {
if let Err(e) = connection.await { eprintln!("Profiler connection error: {}", e); }
});
let statement = client.prepare(&metadata_query).await
.context("Profiler: Failed to prepare metadata query")?;
let mut output = String::from("schema:\n");
// Phase 3: Inspect the statement's columns for metadata
for column in statement.columns() {
let pg_type_name = column.type_().name().to_lowercase();
let arrow_type = map_postgres_to_arrow_type(&pg_type_name)
.unwrap_or("UNKNOWN (Review Manually)");
let column_entry = format!(
"- arrow_type: {}\n column_name: {}\n",
arrow_type,
column.name()
);
output.push_str(&column_entry);
}
// Final instructions for the user
output.push_str("\n# NOTE: Paste the 'schema' block above into your config.yaml\n");
output.push_str(
"# REVIEW any UNKNOWN types. PostgreSQL types: (int8, float8, text, bool, timestamp, date, etc.)\n"
);
Ok(output)
}