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

make tracy optional, add readme.md, continue fixing thrashing for sequential reads for workers

Browse Source
This commit is contained in:
Bill
2025-11-22 22:11:08 -07:00
parent 95f88ee6e4
commit c338fcae65
5 changed files with 293 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
Cargo.lock generated
View File

@@ -259,6 +259,18 @@ dependencies = [
"regex-syntax",
]
[[package]]
name = "async-channel"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2"
dependencies = [
"concurrent-queue",
"event-listener-strategy",
"futures-core",
"pin-project-lite",
]
[[package]]
name = "async-trait"
version = "0.1.89"
@@ -374,6 +386,15 @@ dependencies = [
"unicode-width",
]
[[package]]
name = "concurrent-queue"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "const-random"
version = "0.1.18"
@@ -409,6 +430,12 @@ dependencies = [
"libc",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
[[package]]
name = "crunchy"
version = "0.2.4"
@@ -498,6 +525,27 @@ version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
[[package]]
name = "event-listener"
version = "5.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab"
dependencies = [
"concurrent-queue",
"parking",
"pin-project-lite",
]
[[package]]
name = "event-listener-strategy"
version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93"
dependencies = [
"event-listener",
"pin-project-lite",
]
[[package]]
name = "fallible-iterator"
version = "0.2.0"
@@ -1004,6 +1052,12 @@ version = "1.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
[[package]]
name = "parking"
version = "2.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba"
[[package]]
name = "parking_lot"
version = "0.12.5"
@@ -1765,6 +1819,7 @@ version = "0.1.0"
dependencies = [
"anyhow",
"arrow",
"async-channel",
"byteorder",
"bytes",
"chrono",

14
Cargo.toml
View File

@@ -11,8 +11,7 @@ crate-type = ["cdylib", "rlib"]
[dependencies]
# 1. Python Bindings for FFI
pyo3 = { version = "0.27.1", features = ["extension-module", "chrono-tz"] }
pyo3 = { version = "0.27.1", features = ["extension-module"] }
chrono-tz = "0.10"
# 2. Configuration Parsing (YAML)
@@ -26,6 +25,7 @@ arrow = "57.0.0"
# 4. Asynchronous Runtime (Essential for I/O and Postgres)
tokio = { version = "1.37", features = ["full"] }
# 5. Database Connection (Postgres)
tokio-postgres = "0.7"
deadpool-postgres = "0.14"
@@ -54,9 +54,13 @@ pyo3-arrow = "0.15.0"
# 13. System CPU Count (NEW)
num_cpus = "1.16"
tracing-subscriber = { version = "0.3", features = ["registry", "env-filter"] }
tracing-tracy = "=0.11.2"
tracing = "0.1.41"
tracing-subscriber = { version = "0.3", features = ["registry", "env-filter"], optional = true}
tracing-tracy = { version = "=0.11.2", optional = true }
tracing = {version ="0.1.41", optional = true }
async-channel = "2.5.0"
[target.'cfg(not(target_env = "msvc"))'.dependencies]
mimalloc = { version = "0.1.39" }
async-channel = "2.3"

2
benchmark.py
View File

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

89
docs/rEaDMe.md Normal file
View File

@@ -0,0 +1,89 @@
UncheckedIO 🚀
The world's fastest, most dangerous PostgreSQL-to-Arrow loader.
UncheckedIO is an opinionated, high-performance data connector designed for bulk-loading massive datasets (Terabytes+) from PostgreSQL into Python/Apache Arrow. It achieves extreme speed by skipping runtime schema validation and using a zero-copy streaming parser.
⚠️ WARNING: This library assumes you know what you are doing. If your config.yaml schema does not match your database schema, it will produce garbage data or crash. The "Unchecked" in the name is not a suggestion; it is a promise.
Features
🚀 1.5x - 2x Faster than ConnectorX when Benchmarked on 20M+ row dataset in default Colab runtime.
🧵 Parallel Worker Pool: Automatically scales to available CPU cores.
🧠 Auto-Tuning: Dynamically calculates optimal partition sizes (blast_radius).
📉 Low Memory Footprint: Streaming parser with zero-copy string slicing.
🛠️ "Self-Healing" (Nullable): If a partition fails (e.g., network flake), it returns NULLs for that chunk instead of crashing the entire job.
🔬 Optional Profiling: Built-in integration with the Tracy Profiler for deep I/O analysis.
Installation
pip install unchecked-io
Quick Start
Create a config.yaml file:
connection_string: postgresql://user:password@localhost:5432/dbname
# query MUST be a COPY ... TO STDOUT (FORMAT binary) command
query: >
COPY (SELECT id, name, score, created_at FROM users)
TO STDOUT (FORMAT binary)
schema:
- column_name: id
arrow_type: Int64
- column_name: name
arrow_type: Utf8
- column_name: score
arrow_type: Float32
- column_name: created_at
arrow_type: Timestamp(Nanosecond, None)
Run it in Python:
import unchecked_io
# blast_radius=0 enables auto-tuning
arrow_table = unchecked_io.load_data_from_config("config.yaml", blast_radius=0)
# Convert to Pandas (Zero-Copy)
df = unchecked_io.to_pandas_dataframe(arrow_table)
print(df.head())
🔬 Profiling with Tracy
UncheckedIO includes optional instrumentation for the Tracy Profiler to visualize I/O starvation and thread contention.
This feature is DISABLED by default to prevent overhead and crashes in production environments (like Google Colab).
How to Enable Profiling
You must build the library from source with the profiling feature enabled.
Install Rust & Maturin:
curl --proto '=https' --tlsv1.2 -sSf [https://sh.rustup.rs](https://sh.rustup.rs) | sh
pip install maturin
Build with Feature Flag:
# This builds the wheel and installs it in your current venv
maturin develop --release --features profiling
Run Tracy (GUI):
Download Tracy v0.11.1 (Must match protocol v0.11.x).
Open Tracy.exe and click Connect.
Run your Python script. You will see real-time thread timelines and "IO_WAIT_STARVATION" blocks.
Configuration Reference
Supported Types
PostgreSQL Type
Config arrow_type
bigint / int8
Int64
integer / int4
Int32
double precision
Float64
real / float4
Float32
text / varchar / uuid
Utf8
boolean
Boolean
timestamp
Timestamp(Nanosecond, None)
date
Date32
blast_radius Parameter
0: Auto-Tune (Recommended). Calculates partition size based on row count / (cores * 4).
> 0: Manual Override. Sets specific number of rows per partition. Use for fine-tuning on specific hardware.
License
Apache 2.0
Free for non-production use and production use.

190
src/parser.rs
View File

@@ -24,7 +24,10 @@ use chrono::{NaiveDateTime, NaiveDate};
use std::time::Instant;
use tokio::task::JoinSet;
use arrow::compute::concat_batches;
// NEW: Tracing macros for profiling
// NEW: For the Work Stealing Queue
use async_channel;
// NEW: Tracing macros for profiling - Only import if feature is enabled
#[cfg(feature = "profiling")]
use tracing::{span, Level};
// --- Internal Crates ---
@@ -37,7 +40,7 @@ use crate::config::{ConnectorConfig, load_and_validate_config};
const POSTGRES_EPOCH_MICROS_OFFSET: i64 = 946684800000000;
// --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) ---
// --- 1. CORE DATABASE LOGIC (WORKER POOL PATTERN) ---
// 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> {
@@ -46,25 +49,35 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
let start_phase1 = Instant::now();
// NEW: High-level span for the whole operation
#[cfg(feature = "profiling")]
let root_span = span!(Level::INFO, "UncheckedIO_Run");
#[cfg(feature = "profiling")]
let _root_guard = root_span.enter();
println!("UncheckedIO: Starting Query Planner (Blast Radius: {})...", blast_radius);
// --- PHASE 1: SETUP CONNECTION POOL & STATS ---
#[cfg(feature = "profiling")]
let phase1_span = span!(Level::INFO, "Phase1_Setup");
#[cfg(feature = "profiling")]
let _p1_guard = phase1_span.enter();
// 1. Calculate Worker Count (Fixed Parallelism)
let num_workers = num_cpus::get();
println!("UncheckedIO: Detected {} logical cores. Spawning {} worker threads.", num_workers, num_workers);
// 2. Setup Connection Pool
let pg_config: tokio_postgres::Config = PgConfig::from_str(&config.connection_string)
.context("Invalid connection string in config")?;
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)
.max_size(20)
.max_size(num_workers)
.runtime(Runtime::Tokio1)
.build()
.context("Failed to build connection pool")?;
// 3. Query Table Bounds
// 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";
@@ -77,108 +90,168 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
let row = client.query_one(&stats_query, &[]).await?;
let min_id: i64 = row.try_get(0).context("Failed to get MIN(id)")?;
let max_id: i64 = row.try_get(1).context("Failed to get MAX(id)")?;
drop(client); // Return connection to pool immediately
println!("UncheckedIO: ID Range: {} to {}", min_id, max_id);
// --- NEW: DYNAMIC PARTITION SIZING ---
let total_rows = (max_id - min_id + 1).max(1);
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 dynamic_size = total_rows / target_chunks;
// Ensure a sane minimum (e.g., don't make chunks of 1 row)
let size = dynamic_size.max(10_000);
println!("UncheckedIO: Auto-tuned partition size to {} rows (Targeting {} chunks).", size, target_chunks);
size
} else {
println!("UncheckedIO: Using user-defined partition size: {} rows.", blast_radius);
blast_radius
};
// 4. Create Work Queue
// We use a tuple: (index, query, expected_rows) so we can re-sort later
struct PartitionTask {
index: usize,
query: String,
expected_rows: usize
}
let mut partitions: Vec<PartitionTask> = Vec::new();
// Create an unbounded channel.
// tx = transmitter (main thread), rx = receiver (workers)
let (tx, rx) = async_channel::unbounded::<PartitionTask>();
// 5. Populate the Queue (The "Blast Radius" Logic)
let mut current_min = min_id;
let mut idx = 0;
let mut total_partitions = 0;
while current_min <= max_id {
let current_max = (current_min + blast_radius - 1).min(max_id);
let current_max = (current_min + calculated_blast_radius - 1).min(max_id);
let new_query = format!(
"COPY (SELECT * FROM ({}) AS sub WHERE {} BETWEEN {} AND {}) TO STDOUT (FORMAT binary)",
base_query_inner, partition_key, current_min, current_max
);
let estimated_rows = (current_max - current_min + 1) as usize;
partitions.push(PartitionTask { index: idx, query: new_query, expected_rows: estimated_rows });
current_min += blast_radius;
idx += 1;
}
println!("UncheckedIO: Generated {} parallel partitions.", partitions.len());
// 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 };
// Send to queue (non-blocking since it's unbounded)
tx.send(task).await.context("Failed to fill work queue")?;
current_min += calculated_blast_radius;
idx += 1;
total_partitions += 1;
}
// Close the channel so workers know when to stop
tx.close();
println!("UncheckedIO: Queued {} partitions for processing.", total_partitions);
#[cfg(feature = "profiling")]
drop(_p1_guard); // End Phase 1 Span
let duration_phase1 = start_phase1.elapsed();
// --- PHASE 2: PARALLEL EXECUTION (DATA TRANSFER + PARSING) ---
let start_phase2 = Instant::now();
#[cfg(feature = "profiling")]
let phase2_span = span!(Level::INFO, "Phase2_Execution");
#[cfg(feature = "profiling")]
let _p2_guard = phase2_span.enter();
let arrow_schema = Arc::new(build_arrow_schema(&config)?);
let mut join_set = JoinSet::new();
for task in partitions {
// Spawn exactly 'num_workers' long-lived tasks
for worker_id in 0..num_workers {
let worker_rx = rx.clone();
let worker_pool = pool.clone();
let worker_schema = arrow_schema.clone();
join_set.spawn(async move {
let worker_logic = async {
// Instrument each worker thread so they show up as separate tracks in Tracy
let worker_span = span!(Level::INFO, "Worker_Thread", id = task.index);
// VISUALIZATION: Create a "track" for this worker in Tracy
#[cfg(feature = "profiling")]
let worker_span = span!(Level::INFO, "Worker_Thread", id = worker_id);
#[cfg(feature = "profiling")]
let _w_guard = worker_span.enter();
let worker_client = worker_pool.get().await
.context("Worker: Failed to get pool connection")?;
let mut worker_batches: Vec<(usize, RecordBatch)> = Vec::new();
let copy_stream = worker_client.copy_out(task.query.as_str()).await?;
// Worker Loop: Keep grabbing tasks until the queue is empty and closed
while let Ok(task) = worker_rx.recv().await {
// VISUALIZATION: Show exactly which partition is being processed
#[cfg(feature = "profiling")]
let task_span = span!(Level::INFO, "Processing_Task", partition_id = task.index);
#[cfg(feature = "profiling")]
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 client = worker_pool.get().await.context("Pool exhausted")?;
let copy_stream = client.copy_out(task.query.as_str()).await?;
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
parse_data_with_schema(pinned_stream, worker_schema).await
};
// Call Static Dispatch Parser
parse_data_with_schema(pinned_stream, worker_schema.clone()).await
}.await;
let result = worker_logic.await;
(task.index, result, task.expected_rows)
});
}
// Collect results
let mut results: Vec<Option<RecordBatch>> = vec![None; idx];
while let Some(join_result) = join_set.join_next().await {
let (index, parse_result, expected_rows) = join_result.context("Worker thread panic")?;
match parse_result {
Ok(batch) => {
results[index] = Some(batch.1);
match result {
Ok((_rows, batch)) => {
worker_batches.push((task.index, batch));
}
Err(e) => {
// ERROR LOGGING
// We use tracing::error! here so it shows up brightly in the Tracy log view
tracing::error!(
"Partition {} failed! Error: {}. Falling back to NULLs.",
index, e
);
eprintln!("UncheckedIO: Partition {} failed! Error: {}. Falling back to NULLs.", index, e);
#[cfg(feature = "profiling")]
tracing::error!("Worker {}: Partition {} failed! Error: {}", worker_id, task.index, e);
let null_batch = create_null_batch(arrow_schema.clone(), expected_rows)?;
results[index] = Some(null_batch);
// --- SELF-HEALING LOGIC ---
// If a partition fails (e.g. bad data), we log it and return NULLs
eprintln!("UncheckedIO Worker {}: Partition {} failed! Error: {}. Filling NULLs.", worker_id, task.index, e);
let null_batch = create_null_batch(worker_schema.clone(), task.expected_rows)?;
worker_batches.push((task.index, null_batch));
}
}
}
// Return all batches processed by this worker
Ok::<Vec<(usize, RecordBatch)>, anyhow::Error>(worker_batches)
});
}
// --- PHASE 3: AGGREGATION ---
let mut all_results: Vec<(usize, RecordBatch)> = Vec::with_capacity(total_partitions);
while let Some(join_result) = join_set.join_next().await {
match join_result {
Ok(worker_result) => {
match worker_result {
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")]
drop(_p2_guard); // End Phase 2 Span
let duration_phase2 = start_phase2.elapsed();
// --- PHASE 3: CONCATENATION AND FINALIZATION ---
let start_phase3 = Instant::now();
#[cfg(feature = "profiling")]
let phase3_span = span!(Level::INFO, "Phase3_Concat");
#[cfg(feature = "profiling")]
let _p3_guard = phase3_span.enter();
let batches: Vec<RecordBatch> = results.into_iter()
.filter_map(|b| b)
.collect();
if batches.is_empty() {
if all_results.is_empty() {
println!("UncheckedIO: All workers returned empty batches.");
let duration_total = start_total.elapsed();
println!("--- UncheckedIO Internal Timing ---");
@@ -189,9 +262,17 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
return Ok(RecordBatch::new_empty(arrow_schema));
}
let final_batch = concat_batches(&arrow_schema, &batches)
.context("Failed to concatenate parallel batches")?;
// 1. Sort by index to restore original table order
all_results.sort_by_key(|(index, _)| *index);
// 2. Strip indices
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")]
drop(_p3_guard); // End Phase 3 Span
let duration_phase3 = start_phase3.elapsed();
let duration_total = start_total.elapsed();
@@ -218,8 +299,7 @@ fn create_null_batch(schema: Arc<Schema>, num_rows: usize) -> 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| {
// FIX: Force all columns to be nullable.
// This allows our "Self-Healing" logic to insert NULLs if a partition fails.
// FIX: Force all columns to be nullable for safety
let nullable = true;
let arrow_type = match col_cfg.arrow_type.as_str() {
@@ -315,18 +395,22 @@ async fn parse_binary_stream_static(
'stream_loop: loop {
// 1. MEASURE STARVATION (Waiting for Network)
// Level::ERROR makes this show up nicely in filters/colors, visualizing the "gap"
#[cfg(feature = "profiling")]
let wait_span = span!(Level::ERROR, "IO_WAIT_STARVATION");
#[cfg(feature = "profiling")]
let guard = wait_span.enter();
let next_item = stream.next().await;
#[cfg(feature = "profiling")]
drop(guard); // Important: Drop guard immediately when data arrives!
match next_item {
Some(segment_result) => {
// 2. MEASURE WORK (CPU Parsing)
#[cfg(feature = "profiling")]
let work_span = span!(Level::INFO, "CPU_Parse_Chunk");
#[cfg(feature = "profiling")]
let _work_guard = work_span.enter();
let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?;