From 27100825c65ce6e53f2b15e2411229ef6689d727 Mon Sep 17 00:00:00 2001 From: hall-alex Date: Mon, 6 Apr 2026 16:33:01 +0200 Subject: [PATCH 1/2] Make sure we hand up ordered batches per file --- iceberg_rust_ffi/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/iceberg_rust_ffi/src/lib.rs b/iceberg_rust_ffi/src/lib.rs index dee2aad..7eede7c 100644 --- a/iceberg_rust_ffi/src/lib.rs +++ b/iceberg_rust_ffi/src/lib.rs @@ -174,7 +174,7 @@ pub(crate) fn transform_stream_with_parallel_serialization( )), } }) - .buffer_unordered(concurrency) + .buffered(concurrency) .boxed() } From e49a4da9c1b31163a012848f7c721a9aee630646 Mon Sep 17 00:00:00 2001 From: hall-alex Date: Mon, 6 Apr 2026 18:29:38 +0200 Subject: [PATCH 2/2] File-parallel pipeline with strict ordering and backpressure Replace the to_arrow() code path with a custom file-parallel pipeline that processes N parquet files concurrently while preserving strict file-then-row ordering for Julia consumption. Architecture: - plan_files() provides an ordered list of FileScanTasks - Each file task runs as a background tokio task with its own mpsc channel and per-file Semaphore(100MB) for backpressure - FuturesOrdered yields per-file receivers in file order; the consumer drains file 0 batch-by-batch, then file 1, etc. - Each file uses the same iceberg-rs ArrowReader code path as before Changes: - New: ordered_file_pipeline.rs (~420 lines incl. profiling stats) - full.rs: IcebergScan gains file_io, batch_size, file_concurrency fields; iceberg_arrow_stream uses plan_files() + pipeline - scan_common.rs / incremental.rs: pass through new fields in macros - Temporary PipelineStats with FFI print_summary for benchmarking Co-Authored-By: Claude Opus 4.6 (1M context) --- iceberg_rust_ffi/src/full.rs | 116 +++- iceberg_rust_ffi/src/incremental.rs | 7 + iceberg_rust_ffi/src/lib.rs | 3 + iceberg_rust_ffi/src/ordered_file_pipeline.rs | 499 ++++++++++++++++++ iceberg_rust_ffi/src/scan_common.rs | 12 + 5 files changed, 608 insertions(+), 29 deletions(-) create mode 100644 iceberg_rust_ffi/src/ordered_file_pipeline.rs diff --git a/iceberg_rust_ffi/src/full.rs b/iceberg_rust_ffi/src/full.rs index 1cb7285..0324610 100644 --- a/iceberg_rust_ffi/src/full.rs +++ b/iceberg_rust_ffi/src/full.rs @@ -1,50 +1,91 @@ use std::ffi::{c_char, c_void, CStr}; use std::ptr; +use iceberg::io::FileIO; use iceberg::scan::{TableScan, TableScanBuilder}; use object_store_ffi::{ export_runtime_op, with_cancellation, CResult, NotifyGuard, ResponseGuard, RT, }; -use tokio::sync::Mutex as AsyncMutex; - use crate::scan_common::*; use crate::{IcebergArrowStream, IcebergArrowStreamResponse, IcebergTable}; -/// Struct for regular (full) scan builder and scan +/// Holds state for a full table scan across its lifecycle: +/// 1. Construction: `builder` is set, everything else is None/0. +/// 2. Configuration: Julia calls with_* methods that transform `builder`. +/// 3. Build: `builder` is consumed → `scan` is populated. +/// 4. Stream: `scan` + `file_io` + `batch_size` are consumed by +/// `iceberg_arrow_stream` to create the file-parallel pipeline. #[repr(C)] pub struct IcebergScan { pub builder: Option>, pub scan: Option, - /// 0 = auto-detect (num_cpus) + /// Serialization thread count (0 = auto-detect). Currently unused by the + /// pipeline (kept for future use / backward compat with incremental scan). pub serialization_concurrency: usize, + /// Cloned from the Table at construction time. Passed to the pipeline so + /// each per-file ArrowReader can open its own parquet file. + pub file_io: Option, + /// Captured when Julia calls with_batch_size. Forwarded to each per-file + /// ArrowReaderBuilder inside the pipeline. + pub batch_size: Option, + /// How many parquet files to process concurrently in the pipeline. + /// Set by with_data_file_concurrency_limit (0 = auto-detect). + pub file_concurrency: usize, } unsafe impl Send for IcebergScan {} -/// Create a new scan builder +/// Create a new scan builder from an opened table. +/// Captures `file_io` from the table for later use by the pipeline. #[no_mangle] pub extern "C" fn iceberg_new_scan(table: *mut IcebergTable) -> *mut IcebergScan { if table.is_null() { return ptr::null_mut(); } let table_ref = unsafe { &*table }; + let file_io = table_ref.table.file_io().clone(); let scan_builder = table_ref.table.scan(); Box::into_raw(Box::new(IcebergScan { builder: Some(scan_builder), scan: None, serialization_concurrency: 0, + file_io: Some(file_io), + batch_size: None, + file_concurrency: 0, })) } -// Use macros from scan_common for shared functionality +// ── Scan builder configuration (via macros from scan_common.rs) ───────── + impl_select_columns!(iceberg_select_columns, IcebergScan); -impl_scan_builder_method!( - iceberg_scan_with_data_file_concurrency_limit, - IcebergScan, - with_data_file_concurrency_limit, - n: usize -); +/// Set file concurrency. Hand-written (not macro-generated) because we need +/// to both (a) forward the value to the iceberg-rs scan builder and (b) +/// capture it in `file_concurrency` for our pipeline. +#[no_mangle] +pub extern "C" fn iceberg_scan_with_data_file_concurrency_limit( + scan: &mut *mut IcebergScan, + n: usize, +) -> CResult { + if scan.is_null() || (*scan).is_null() { + return CResult::Error; + } + let scan_ref = unsafe { Box::from_raw(*scan) }; + if scan_ref.builder.is_none() { + return CResult::Error; + } + *scan = Box::into_raw(Box::new(IcebergScan { + builder: scan_ref + .builder + .map(|b| b.with_data_file_concurrency_limit(n)), + scan: scan_ref.scan, + serialization_concurrency: scan_ref.serialization_concurrency, + file_io: scan_ref.file_io, + batch_size: scan_ref.batch_size, + file_concurrency: n, + })); + CResult::Ok +} impl_scan_builder_method!( iceberg_scan_with_manifest_file_concurrency_limit, @@ -80,7 +121,16 @@ impl_scan_builder_method!( snapshot_id: i64 ); -// Async function to initialize stream from a table scan +// ── Stream creation ───────────────────────────────────────────────────── +// +// Instead of calling iceberg-rs's `scan.to_arrow()` (which uses +// `try_for_each_concurrent` internally and interleaves batches across +// files in arbitrary order), we: +// 1. Call `scan.plan_files()` to get an ordered list of FileScanTasks. +// 2. Feed them into our own file-parallel pipeline +// (ordered_file_pipeline.rs) which processes N files concurrently +// but yields batches in strict file-then-row order. + export_runtime_op!( iceberg_arrow_stream, IcebergArrowStreamResponse, @@ -88,39 +138,47 @@ export_runtime_op!( if scan.is_null() { return Err(anyhow::anyhow!("Null scan pointer provided")); } - let scan_ptr = unsafe { &*scan }; + let scan_ptr = unsafe { &mut *scan }; let scan_ref = &scan_ptr.scan; if scan_ref.is_none() { return Err(anyhow::anyhow!("Scan not initialized")); } - // Determine concurrency (0 = auto-detect) - let serialization_concurrency = scan_ptr.serialization_concurrency; - let serialization_concurrency = if serialization_concurrency == 0 { + // File pipeline concurrency (0 = auto-detect from available CPUs) + let concurrency = scan_ptr.file_concurrency; + let concurrency = if concurrency == 0 { std::thread::available_parallelism() .map(|n| n.get()) .unwrap_or(1) } else { - serialization_concurrency + concurrency }; - Ok((scan_ref.as_ref().unwrap(), serialization_concurrency)) + // Take file_io — it's moved into the pipeline (one clone per file task). + let file_io = scan_ptr + .file_io + .take() + .ok_or_else(|| anyhow::anyhow!("file_io not available"))?; + let batch_size = scan_ptr.batch_size; + + Ok((scan_ref.as_ref().unwrap(), concurrency, file_io, batch_size)) }, result_tuple, async { - let (scan_ref, serialization_concurrency) = result_tuple; + let (scan_ref, concurrency, file_io, batch_size) = result_tuple; - let stream = scan_ref.to_arrow().await?; + // Collect the ordered file task list from iceberg-rs. + use futures::TryStreamExt; + let tasks: Vec = + scan_ref.plan_files().await?.try_collect().await?; - // Transform stream: RecordBatch -> ArrowBatch with parallel serialization - let serialized_stream = crate::transform_stream_with_parallel_serialization( - stream, - serialization_concurrency - ); + // Hand off to the file-parallel pipeline. + let stream = crate::ordered_file_pipeline::create_pipeline( + tasks, file_io, batch_size, concurrency, + ) + .await?; - Ok::(IcebergArrowStream { - stream: AsyncMutex::new(serialized_stream), - }) + Ok::(stream) }, scan: *mut IcebergScan ); diff --git a/iceberg_rust_ffi/src/incremental.rs b/iceberg_rust_ffi/src/incremental.rs index 9e1d1fe..8a281c2 100644 --- a/iceberg_rust_ffi/src/incremental.rs +++ b/iceberg_rust_ffi/src/incremental.rs @@ -22,6 +22,10 @@ pub struct IcebergIncrementalScan { pub scan: Option, /// 0 = auto-detect (num_cpus) pub serialization_concurrency: usize, + // Present for macro compatibility with IcebergScan; unused for incremental. + pub file_io: Option, + pub batch_size: Option, + pub file_concurrency: usize, } unsafe impl Send for IcebergIncrementalScan {} @@ -102,6 +106,9 @@ pub extern "C" fn iceberg_new_incremental_scan( builder: Some(scan_builder), scan: None, serialization_concurrency: 0, + file_io: None, + batch_size: None, + file_concurrency: 0, })) } diff --git a/iceberg_rust_ffi/src/lib.rs b/iceberg_rust_ffi/src/lib.rs index 7eede7c..be2dd10 100644 --- a/iceberg_rust_ffi/src/lib.rs +++ b/iceberg_rust_ffi/src/lib.rs @@ -31,6 +31,9 @@ mod writer; // Column-based writer module (zero-copy from Julia) mod writer_columns; +// Ordered file-parallel pipeline +mod ordered_file_pipeline; + // Response types module mod response; diff --git a/iceberg_rust_ffi/src/ordered_file_pipeline.rs b/iceberg_rust_ffi/src/ordered_file_pipeline.rs new file mode 100644 index 0000000..d091b63 --- /dev/null +++ b/iceberg_rust_ffi/src/ordered_file_pipeline.rs @@ -0,0 +1,499 @@ +//! File-parallel pipeline for reading Iceberg tables with strict ordering. +//! +//! # Problem +//! iceberg-rs's `to_arrow()` uses `try_for_each_concurrent` which interleaves +//! batches from different files in arbitrary order. We need strict +//! file-then-row ordering. +//! +//! # Approach +//! We call `plan_files()` to get an ordered list of FileScanTasks, then +//! process N files concurrently while yielding batches in strict file order. +//! +//! Each file task is a background tokio task that: +//! 1. Builds a per-file ArrowReader (same iceberg-rs code path as to_arrow) +//! 2. Reads row groups sequentially → RecordBatch stream +//! 3. Serializes each batch to Arrow IPC (via spawn_blocking) +//! 4. Sends serialized batches into a per-file mpsc channel +//! +//! A consumer (`run`) uses FuturesOrdered to drain files in order: +//! - FuturesOrdered yields per-file receivers in the order files were pushed +//! - The consumer drains file 0's channel batch-by-batch, then file 1's, etc. +//! - Each drained batch is forwarded to the outer channel (read by Julia) +//! +//! # Memory bounding +//! Each file task has its own Semaphore(MAX_BUFFERED_BYTES_PER_TASK). After +//! serializing a batch, the task acquires byte_len permits. If the budget is +//! exhausted, the task yields (async, not blocking) until the consumer drains +//! batches and releases permits. This caps each file's buffered output to +//! ~100MB independently. + +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Instant; + +use futures::StreamExt; +use iceberg::arrow::ArrowReaderBuilder; +use iceberg::io::FileIO; +use iceberg::scan::FileScanTask; +use tokio::sync::{mpsc, Mutex as AsyncMutex, Semaphore}; + +use crate::table::{ArrowBatch, IcebergArrowStream}; + +/// Per-file cap on serialized bytes buffered ahead of the consumer. +const MAX_BUFFERED_BYTES_PER_TASK: usize = 100 * 1024 * 1024; + +/// Hard cap on file-level concurrency to keep total memory bounded. +const MAX_FILE_CONCURRENCY: usize = 16; + +// =========================================================================== +// Temporary profiling — will be removed before merging to production. +// +// Global atomic counters accumulate timing/size data across all file tasks. +// Called from Julia via `@ccall iceberg_print_pipeline_stats()`. +// =========================================================================== + +/// Accumulates profiling data across all file tasks in a pipeline run. +/// All fields are atomics so concurrent file tasks can update without locks. +struct PipelineStats { + // ── Overall ── + pipeline_wall_ns: AtomicU64, + files_completed: AtomicUsize, + batches_produced: AtomicUsize, + bytes_produced: AtomicU64, + + // ── Concurrency ── + peak_concurrency: AtomicUsize, + active_tasks: AtomicUsize, + + // ── Per-phase cumulative time (ns), summed across all file tasks ── + /// Time to build ArrowReader and call reader.read() (opens parquet, + /// loads metadata, resolves schema, loads delete files, builds filters). + reader_setup_ns: AtomicU64, + /// Time inside batch_stream.next() — fetches compressed pages from + /// storage, decompresses (ZSTD), decodes pages, assembles columns. + fetch_decode_ns: AtomicU64, + /// Time in spawn_blocking(serialize_record_batch) — writes RecordBatch + /// to Arrow IPC wire format for transfer to Julia. + serialize_ns: AtomicU64, + /// Time blocked on per-file semaphore (backpressure from consumer). + semaphore_wait_ns: AtomicU64, + + // ── Consumer ── + /// Time the consumer spends waiting on FuturesOrdered for the next + /// file's receiver (ordering stall — head file not ready yet). + consumer_wait_ns: AtomicU64, + + // ── Memory ── + /// Live counter of serialized bytes buffered across all file tasks. + buffered_bytes: AtomicU64, + /// High-water mark of buffered_bytes. + peak_buffered_bytes: AtomicU64, +} + +impl PipelineStats { + const fn new() -> Self { + Self { + pipeline_wall_ns: AtomicU64::new(0), + files_completed: AtomicUsize::new(0), + batches_produced: AtomicUsize::new(0), + bytes_produced: AtomicU64::new(0), + peak_concurrency: AtomicUsize::new(0), + active_tasks: AtomicUsize::new(0), + reader_setup_ns: AtomicU64::new(0), + fetch_decode_ns: AtomicU64::new(0), + serialize_ns: AtomicU64::new(0), + semaphore_wait_ns: AtomicU64::new(0), + consumer_wait_ns: AtomicU64::new(0), + buffered_bytes: AtomicU64::new(0), + peak_buffered_bytes: AtomicU64::new(0), + } + } + + fn reset(&self) { + self.pipeline_wall_ns.store(0, Ordering::Relaxed); + self.files_completed.store(0, Ordering::Relaxed); + self.batches_produced.store(0, Ordering::Relaxed); + self.bytes_produced.store(0, Ordering::Relaxed); + self.peak_concurrency.store(0, Ordering::Relaxed); + self.active_tasks.store(0, Ordering::Relaxed); + self.reader_setup_ns.store(0, Ordering::Relaxed); + self.fetch_decode_ns.store(0, Ordering::Relaxed); + self.serialize_ns.store(0, Ordering::Relaxed); + self.semaphore_wait_ns.store(0, Ordering::Relaxed); + self.consumer_wait_ns.store(0, Ordering::Relaxed); + self.buffered_bytes.store(0, Ordering::Relaxed); + self.peak_buffered_bytes.store(0, Ordering::Relaxed); + } + + fn track_task_start(&self) { + let prev = self.active_tasks.fetch_add(1, Ordering::Relaxed); + self.peak_concurrency.fetch_max(prev + 1, Ordering::Relaxed); + } + + fn track_task_end(&self) { + self.active_tasks.fetch_sub(1, Ordering::Relaxed); + } + + fn track_buffer_add(&self, bytes: u64) { + let prev = self.buffered_bytes.fetch_add(bytes, Ordering::Relaxed); + self.peak_buffered_bytes + .fetch_max(prev + bytes, Ordering::Relaxed); + } + + fn track_buffer_release(&self, bytes: u64) { + self.buffered_bytes.fetch_sub(bytes, Ordering::Relaxed); + } + + fn print_summary(&self) { + let wall_ms = self.pipeline_wall_ns.load(Ordering::Relaxed) as f64 / 1e6; + let files = self.files_completed.load(Ordering::Relaxed); + let batches = self.batches_produced.load(Ordering::Relaxed); + let bytes = self.bytes_produced.load(Ordering::Relaxed); + let peak = self.peak_concurrency.load(Ordering::Relaxed); + let setup_ms = self.reader_setup_ns.load(Ordering::Relaxed) as f64 / 1e6; + let fd_ms = self.fetch_decode_ns.load(Ordering::Relaxed) as f64 / 1e6; + let ser_ms = self.serialize_ns.load(Ordering::Relaxed) as f64 / 1e6; + let sem_ms = self.semaphore_wait_ns.load(Ordering::Relaxed) as f64 / 1e6; + let con_ms = self.consumer_wait_ns.load(Ordering::Relaxed) as f64 / 1e6; + let peak_buf = self.peak_buffered_bytes.load(Ordering::Relaxed) as f64 / (1024.0 * 1024.0); + + let bytes_mb = bytes as f64 / (1024.0 * 1024.0); + let throughput = if wall_ms > 0.0 { + bytes_mb / (wall_ms / 1000.0) + } else { + 0.0 + }; + let file_wall_ms = setup_ms + fd_ms + ser_ms + sem_ms; + let parallelism = if wall_ms > 0.0 { + file_wall_ms / wall_ms + } else { + 0.0 + }; + let limit_mb = MAX_BUFFERED_BYTES_PER_TASK / (1024 * 1024); + + // Box layout: "│ " + content + padding + " │", total = BOX chars. + // All content uses ASCII only so byte len = display width. + const BOX: usize = 68; // total width including borders + + let row = |content: &str| { + // "│ " = 3, " │" = 3 => content area = BOX - 6 + let pad = (BOX - 6).saturating_sub(content.len()); + println!("│ {}{:pad$} │", content, "", pad = pad); + }; + let dashes = |n: usize| -> String { "─".repeat(n) }; + let sep = |label: &str| { + // "│ ── label ──── │" — label is ASCII, dashes are multi-byte. + // Display width: 3 + 3 + label.len() + 1 + fill + 1 = BOX + // => fill = BOX - 3 - 3 - label.len() - 3 - 1 = BOX - 10 - label.len() + let fill = (BOX - 10).saturating_sub(label.len()); + println!("│ {} {} {} │", dashes(2), label, dashes(fill)); + }; + let border = |left: char, right: char| { + println!("{}{}{}", left, dashes(BOX - 2), right); + }; + + border('┌', '┐'); + row(&format!("Pipeline Stats")); + row(""); + row(&format!("wall time: {:>9.1} ms", wall_ms)); + row(&format!("files: {:>9} peak concurrency: {}", files, peak)); + row(&format!("batches: {:>9} serialized: {:.1} MB", batches, bytes_mb)); + row(&format!("throughput: {:>9.1} MB/s parallelism: {:.1}x", throughput, parallelism)); + sep("time across all file tasks (sum)"); + row(&format!("reader setup: {:>9.1} ms (open, metadata, deletes)", setup_ms)); + row(&format!("fetch+decode: {:>9.1} ms (I/O + ZSTD + decode)", fd_ms)); + row(&format!("serialize IPC: {:>9.1} ms (RecordBatch -> Arrow IPC)", ser_ms)); + row(&format!("semaphore wait: {:>9.1} ms (backpressure)", sem_ms)); + sep("consumer"); + row(&format!("ordering stall: {:>9.1} ms (waiting for head file)", con_ms)); + sep("memory"); + row(&format!("peak buffered: {:>9.1} MB (limit: {} MB/task)", peak_buf, limit_mb)); + border('└', '┘'); + } +} + +static STATS: PipelineStats = PipelineStats::new(); + +// ── FFI exports for profiling (called from Julia benchmark teardown) ───── + +#[no_mangle] +pub extern "C" fn iceberg_print_pipeline_stats() { + STATS.print_summary(); +} + +#[no_mangle] +pub extern "C" fn iceberg_reset_pipeline_stats() { + STATS.reset(); +} + +// =========================================================================== +// Pipeline implementation +// =========================================================================== + +/// A serialized Arrow IPC batch bundled with its byte size and a reference +/// to the originating file task's semaphore. The consumer releases permits +/// after forwarding the batch to Julia, unblocking the producer. +struct BufferedBatch { + batch: ArrowBatch, + byte_len: usize, + semaphore: Arc, +} + +/// Create the file-parallel pipeline and return it as an IcebergArrowStream. +/// +/// Spawns a background `run` task that processes files concurrently and +/// feeds serialized batches into an mpsc channel. The returned stream +/// wraps the receiving end of that channel. +pub async fn create_pipeline( + tasks: Vec, + file_io: FileIO, + batch_size: Option, + concurrency: usize, +) -> anyhow::Result { + assert!( + concurrency <= MAX_FILE_CONCURRENCY, + "file concurrency {concurrency} exceeds hard cap {MAX_FILE_CONCURRENCY}" + ); + + STATS.reset(); + + // Outer channel: run() → Julia (via IcebergArrowStream). + let (tx, rx) = mpsc::channel(concurrency * 2); + + tokio::spawn(run(tasks, file_io, batch_size, concurrency, tx)); + + // Wrap the mpsc receiver as a BoxStream for IcebergArrowStream. + let stream = futures::stream::unfold(rx, |mut rx| async move { + rx.recv().await.map(|item| (item, rx)) + }) + .boxed(); + + Ok(IcebergArrowStream { + stream: AsyncMutex::new(stream), + }) +} + +/// Spawn a single file task. Returns a future that resolves immediately +/// to the receiving end of the file's batch channel. +/// +/// The actual work (parquet I/O, decode, serialize) happens in the +/// background tokio task. The future resolves to the Receiver so that +/// FuturesOrdered can yield receivers in file order. +fn spawn_file_task( + task: FileScanTask, + file_io: FileIO, + batch_size: Option, +) -> impl std::future::Future< + Output = Result>, iceberg::Error>, +> { + // Each file gets its own semaphore for independent backpressure. + let sem = Arc::new(Semaphore::new(MAX_BUFFERED_BYTES_PER_TASK)); + let (file_tx, file_rx) = mpsc::channel(8); + + tokio::spawn(process_file(task, file_io, batch_size, sem, file_tx)); + + async move { Ok(file_rx) } +} + +/// Main consumer loop. Orchestrates file-level parallelism while +/// maintaining strict file ordering. +/// +/// Uses FuturesOrdered to poll N file tasks concurrently but yield their +/// receivers in push order. For each file, drains its channel batch-by-batch, +/// forwarding to the outer channel and releasing semaphore permits. +async fn run( + tasks: Vec, + file_io: FileIO, + batch_size: Option, + concurrency: usize, + tx: mpsc::Sender>, +) { + use futures::stream::FuturesOrdered; + + let pipeline_start = Instant::now(); + let mut in_flight = FuturesOrdered::new(); + let mut task_iter = tasks.into_iter(); + + // Seed the first N file tasks. + for _ in 0..concurrency { + if let Some(task) = task_iter.next() { + in_flight.push_back(spawn_file_task(task, file_io.clone(), batch_size)); + } + } + + // FuturesOrdered::next() yields results in push order (file 0, 1, 2, ...). + while let Some(file_result) = in_flight.next().await { + // Eagerly start the next file to keep N tasks in flight. + if let Some(task) = task_iter.next() { + in_flight.push_back(spawn_file_task(task, file_io.clone(), batch_size)); + } + + let mut file_rx = match file_result { + Ok(rx) => rx, + Err(e) => { + let _ = tx.send(Err(e)).await; + return; + } + }; + + // Drain this file's batches in row order, forwarding to Julia. + while let Some(batch_result) = file_rx.recv().await { + match batch_result { + Ok(buf) => { + let sem = buf.semaphore.clone(); + let byte_len = buf.byte_len; + if tx.send(Ok(buf.batch)).await.is_err() { + return; // Julia side dropped the stream + } + // Release permits so the file task can produce more. + sem.add_permits(byte_len); + STATS.track_buffer_release(byte_len as u64); + } + Err(e) => { + let _ = tx.send(Err(e)).await; + return; + } + } + } + // file_rx exhausted → file task finished → move to next file + } + + STATS + .pipeline_wall_ns + .store(pipeline_start.elapsed().as_nanos() as u64, Ordering::Relaxed); +} + +/// Wrapper around process_file_inner that ensures errors are sent to the +/// channel and stats are updated even on failure. When this function +/// returns, `tx` is dropped, causing the consumer's `file_rx.recv()` to +/// return None — signaling that this file is done. +async fn process_file( + task: FileScanTask, + file_io: FileIO, + batch_size: Option, + semaphore: Arc, + tx: mpsc::Sender>, +) { + let result = process_file_inner(task, file_io, batch_size, semaphore, &tx).await; + if let Err(e) = result { + let _ = tx.send(Err(e)).await; + } + STATS.track_task_end(); +} + +/// Process a single parquet file through four timed phases: +/// 1. Reader setup — build ArrowReader, open file, resolve schema +/// 2. Fetch + decode — I/O, ZSTD decompression, column assembly +/// 3. Serialize — RecordBatch → Arrow IPC for transfer to Julia +/// 4. Backpressure — wait on semaphore if buffered too far ahead +/// +/// Each phase's cumulative time is recorded in STATS for the summary. +async fn process_file_inner( + task: FileScanTask, + file_io: FileIO, + batch_size: Option, + semaphore: Arc, + tx: &mpsc::Sender>, +) -> Result<(), iceberg::Error> { + STATS.track_task_start(); + + // ── Phase 1: Reader setup ─────────────────────────────────────────── + // Builds a per-file ArrowReader using the same iceberg-rs code path as + // to_arrow(): opens parquet metadata, resolves schema, loads delete + // files, builds row filters. with_data_file_concurrency_limit(1) means + // this reader processes one file (the one we give it). + let setup_start = Instant::now(); + let mut builder = ArrowReaderBuilder::new(file_io).with_data_file_concurrency_limit(1); + if let Some(bs) = batch_size { + builder = builder.with_batch_size(bs); + } + let reader = builder.build(); + let task_stream = Box::pin(futures::stream::once(async { Ok(task) })); + let batch_stream = reader + .read(task_stream) + .map_err(|e| iceberg::Error::new(iceberg::ErrorKind::Unexpected, e.to_string()))?; + STATS + .reader_setup_ns + .fetch_add(setup_start.elapsed().as_nanos() as u64, Ordering::Relaxed); + + tokio::pin!(batch_stream); + + loop { + // ── Phase 2: Fetch + decode ───────────────────────────────────── + // Each .next() call fetches compressed parquet pages from storage, + // decompresses (ZSTD), decodes column encodings, and assembles a + // RecordBatch. These are inseparable without forking parquet-rs. + let fd_start = Instant::now(); + let batch_opt = batch_stream.next().await; + STATS + .fetch_decode_ns + .fetch_add(fd_start.elapsed().as_nanos() as u64, Ordering::Relaxed); + + let batch = match batch_opt { + Some(Ok(b)) => b, + Some(Err(e)) => return Err(e), + None => break, // end of file + }; + + // ── Phase 3: Serialize to Arrow IPC ───────────────────────────── + // CPU-bound work, offloaded to the blocking thread pool to avoid + // starving the tokio executor. + let ser_start = Instant::now(); + let serialized = tokio::task::spawn_blocking(move || crate::serialize_record_batch(batch)) + .await + .map_err(|e| { + iceberg::Error::new( + iceberg::ErrorKind::Unexpected, + format!("serialize panicked: {e}"), + ) + })? + .map_err(|e| iceberg::Error::new(iceberg::ErrorKind::Unexpected, e.to_string()))?; + STATS + .serialize_ns + .fetch_add(ser_start.elapsed().as_nanos() as u64, Ordering::Relaxed); + + let byte_len = serialized.length; + STATS.batches_produced.fetch_add(1, Ordering::Relaxed); + STATS + .bytes_produced + .fetch_add(byte_len as u64, Ordering::Relaxed); + + // ── Phase 4: Backpressure ─────────────────────────────────────── + // Acquire permits equal to the serialized size. If this file task + // has produced > MAX_BUFFERED_BYTES_PER_TASK ahead of the consumer, + // this yields (async, not thread-blocking) until permits are freed. + let sem_start = Instant::now(); + let _permit = semaphore + .acquire_many(byte_len as u32) + .await + .map_err(|e| { + iceberg::Error::new( + iceberg::ErrorKind::Unexpected, + format!("semaphore: {e}"), + ) + })?; + STATS + .semaphore_wait_ns + .fetch_add(sem_start.elapsed().as_nanos() as u64, Ordering::Relaxed); + // Detach the permit — the consumer releases it via add_permits(). + std::mem::forget(_permit); + + STATS.track_buffer_add(byte_len as u64); + + // Send the batch to this file's channel. + if tx + .send(Ok(BufferedBatch { + batch: serialized, + byte_len, + semaphore: semaphore.clone(), + })) + .await + .is_err() + { + return Ok(()); // consumer dropped the receiver + } + } + + STATS.files_completed.fetch_add(1, Ordering::Relaxed); + Ok(()) +} diff --git a/iceberg_rust_ffi/src/scan_common.rs b/iceberg_rust_ffi/src/scan_common.rs index daac382..f8c6737 100644 --- a/iceberg_rust_ffi/src/scan_common.rs +++ b/iceberg_rust_ffi/src/scan_common.rs @@ -38,6 +38,9 @@ macro_rules! impl_select_columns { builder: scan_ref.builder.map(|b| b.select(columns)), scan: scan_ref.scan, serialization_concurrency: scan_ref.serialization_concurrency, + file_io: scan_ref.file_io, + batch_size: scan_ref.batch_size, + file_concurrency: scan_ref.file_concurrency, })); CResult::Ok @@ -84,6 +87,9 @@ macro_rules! impl_scan_builder_method { builder: scan_ref.builder.map(|b| b.$builder_method($($param),*)), scan: scan_ref.scan, serialization_concurrency: scan_ref.serialization_concurrency, + file_io: scan_ref.file_io, + batch_size: scan_ref.batch_size, + file_concurrency: scan_ref.file_concurrency, })); CResult::Ok @@ -111,6 +117,9 @@ macro_rules! impl_with_batch_size { builder: scan_ref.builder.map(|b| b.with_batch_size(Some(n))), scan: None, serialization_concurrency: scan_ref.serialization_concurrency, + file_io: scan_ref.file_io, + batch_size: Some(n), + file_concurrency: scan_ref.file_concurrency, })); CResult::Ok @@ -138,6 +147,9 @@ macro_rules! impl_scan_build { builder: None, scan: Some(built_scan), serialization_concurrency: scan_ref.serialization_concurrency, + file_io: scan_ref.file_io, + batch_size: scan_ref.batch_size, + file_concurrency: scan_ref.file_concurrency, })); CResult::Ok }