diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index 5582e4397f..9ee981f595 100644 --- a/crates/iceberg/src/arrow/caching_delete_file_loader.rs +++ b/crates/iceberg/src/arrow/caching_delete_file_loader.rs @@ -101,7 +101,11 @@ impl CachingDeleteFileLoader { ) -> Self { let scan_metrics = ScanMetrics::new(); CachingDeleteFileLoader { - basic_delete_file_loader: BasicDeleteFileLoader::new(file_io, scan_metrics), + basic_delete_file_loader: BasicDeleteFileLoader::new( + file_io, + scan_metrics, + runtime.clone(), + ), concurrency_limit_data_files, delete_filter: DeleteFilter::new(runtime.clone()), runtime, @@ -112,6 +116,7 @@ impl CachingDeleteFileLoader { self.basic_delete_file_loader = BasicDeleteFileLoader::new( self.basic_delete_file_loader.file_io().clone(), scan_metrics, + self.runtime.clone(), ); self } @@ -204,7 +209,10 @@ impl CachingDeleteFileLoader { let del_filter = self.delete_filter.clone(); let concurrency_limit_data_files = self.concurrency_limit_data_files; let basic_delete_file_loader = self.basic_delete_file_loader.clone(); - self.runtime.io().spawn(async move { + let io = self.runtime.io().clone(); + // Decode Parquet batches and construct delete filters on CPU workers. + // Storage operations are dispatched separately to the I/O runtime. + self.runtime.cpu().spawn(async move { let result = async move { let mut del_filter = del_filter; let basic_delete_file_loader = basic_delete_file_loader.clone(); @@ -212,12 +220,14 @@ impl CachingDeleteFileLoader { let mut results_stream = task_stream .map(move |(task, file_io, del_filter, schema)| { let basic_delete_file_loader = basic_delete_file_loader.clone(); + let io = io.clone(); async move { Self::load_file_for_task( &task, basic_delete_file_loader.clone(), del_filter, schema, + io, ) .await } @@ -262,13 +272,20 @@ impl CachingDeleteFileLoader { basic_delete_file_loader: BasicDeleteFileLoader, del_filter: DeleteFilter, schema: SchemaRef, + io: crate::runtime::RuntimeHandle, ) -> Result { match task.file_type() { DataContentType::PositionDeletes => { // A V3 deletion vector arrives as a PositionDeletes entry whose deletes live in // a Puffin blob, not in a positional-delete parquet file. if task.file_format() == DataFileFormat::Puffin { - return Self::load_deletion_vector(task, basic_delete_file_loader).await; + let task = task.clone(); + return io + .spawn(async move { + Self::load_deletion_vector(&task, basic_delete_file_loader).await + }) + .abort_on_drop() + .await?; } match del_filter.try_start_pos_del_load(task.file_path()) { @@ -840,6 +857,41 @@ mod tests { use crate::spec::{DataContentType, DataFileFormat, Schema}; use crate::test_utils::encode_dv_blob; + #[test] + fn delete_loading_requires_cpu_runtime() { + let io = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap(); + // Deliberately not driven until after the I/O runtime has had a chance + // to run. A loader accidentally spawned on I/O would finish early. + let cpu = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let runtime = Runtime::new_with_split(&io, &cpu); + let dir = TempDir::new().unwrap(); + let tasks = setup(dir.path()); + let loader = CachingDeleteFileLoader::new(FileIO::new_with_fs(), 2, runtime); + let mut loaded = loader.load_deletes(tasks[0].deletes(), tasks[0].schema_ref()); + io.block_on(async { + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut loaded) + .await + .is_err() + ); + }); + let filter = cpu.block_on(async { + tokio::time::timeout(std::time::Duration::from_secs(5), loaded) + .await + .unwrap() + .unwrap() + .unwrap() + }); + assert!(filter.get_delete_vector(&tasks[0]).is_some()); + } + #[test] fn test_deserialized_deletion_vector_rejects_missing_fields() { // A deletion vector that omits its required coordinates fails validation @@ -869,7 +921,7 @@ mod tests { let eq_delete_file_path = setup_write_equality_delete_file_1(table_location); let basic_delete_file_loader = - BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new()); + BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new(), Runtime::current()); let record_batch_stream = basic_delete_file_loader .parquet_to_batch_stream( &eq_delete_file_path, @@ -1307,7 +1359,7 @@ mod tests { let file_io = FileIO::new_with_fs(); let basic_delete_file_loader = - BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new()); + BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new(), Runtime::current()); let batch_stream = basic_delete_file_loader .parquet_to_batch_stream( @@ -1488,7 +1540,7 @@ mod tests { writer.close().unwrap(); let basic_delete_file_loader = - BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new()); + BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new(), Runtime::current()); let record_batch_stream = basic_delete_file_loader .parquet_to_batch_stream(&path, std::fs::metadata(&path).unwrap().len(), None) .await diff --git a/crates/iceberg/src/arrow/delete_file_loader.rs b/crates/iceberg/src/arrow/delete_file_loader.rs index ec61710a42..1f3f32c75e 100644 --- a/crates/iceberg/src/arrow/delete_file_loader.rs +++ b/crates/iceberg/src/arrow/delete_file_loader.rs @@ -47,14 +47,16 @@ pub trait DeleteFileLoader { pub(crate) struct BasicDeleteFileLoader { file_io: FileIO, scan_metrics: ScanMetrics, + runtime: crate::Runtime, } #[allow(unused_variables)] impl BasicDeleteFileLoader { - pub fn new(file_io: FileIO, scan_metrics: ScanMetrics) -> Self { + pub fn new(file_io: FileIO, scan_metrics: ScanMetrics, runtime: crate::Runtime) -> Self { BasicDeleteFileLoader { file_io, scan_metrics, + runtime, } } @@ -78,6 +80,7 @@ impl BasicDeleteFileLoader { let (parquet_file_reader, arrow_metadata) = ArrowReader::open_parquet_file( data_file_path, &self.file_io, + &self.runtime, file_size_in_bytes, parquet_read_options, self.scan_metrics.bytes_read_counter(), @@ -156,7 +159,8 @@ mod tests { let file_io = FileIO::new_with_fs(); let scan_metrics = ScanMetrics::new(); - let delete_file_loader = BasicDeleteFileLoader::new(file_io.clone(), scan_metrics); + let delete_file_loader = + BasicDeleteFileLoader::new(file_io.clone(), scan_metrics, crate::Runtime::current()); let file_scan_tasks = setup(table_location); @@ -241,7 +245,8 @@ mod tests { .unwrap(); let scan_metrics = ScanMetrics::new(); - let delete_file_loader = BasicDeleteFileLoader::new(file_io, scan_metrics); + let delete_file_loader = + BasicDeleteFileLoader::new(file_io, scan_metrics, crate::Runtime::current()); let result = delete_file_loader .read_delete_file(&task, schema) @@ -317,7 +322,8 @@ mod tests { .unwrap(); let scan_metrics = ScanMetrics::new(); - let delete_file_loader = BasicDeleteFileLoader::new(file_io, scan_metrics); + let delete_file_loader = + BasicDeleteFileLoader::new(file_io, scan_metrics, crate::Runtime::current()); let result = delete_file_loader .read_delete_file(&task, schema) diff --git a/crates/iceberg/src/arrow/mod.rs b/crates/iceberg/src/arrow/mod.rs index 28cf6e94a4..cea7c849d7 100644 --- a/crates/iceberg/src/arrow/mod.rs +++ b/crates/iceberg/src/arrow/mod.rs @@ -33,6 +33,7 @@ mod reader; pub mod record_batch_projector; pub(crate) mod record_batch_transformer; pub(crate) use record_batch_transformer::build_partition_constant; +mod runtime_file_read; mod scan_metrics; mod value; diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 1ddee493c4..9e54fa297c 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -42,6 +42,7 @@ use crate::arrow::build_partition_constant; use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader; use crate::arrow::int96::coerce_int96_timestamps; use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder; +use crate::arrow::runtime_file_read::RuntimeFileRead; use crate::arrow::scan_metrics::{CountingFileRead, ScanMetrics, ScanResult}; use crate::encryption::StandardKeyMetadata; use crate::error::{Result, invalid_data}; @@ -169,6 +170,7 @@ impl FileScanTaskReader { let file_size_in_bytes = task.file_size_in_bytes(); let bytes_read = Arc::clone(self.scan_metrics.bytes_read_counter()); let key_metadata = task.key_metadata().map(|km| km.to_vec()); + let runtime = self.runtime.clone(); let opening = self .runtime .io() @@ -176,6 +178,7 @@ impl FileScanTaskReader { ArrowReader::open_parquet_file( &path, &file_io, + &runtime, file_size_in_bytes, parquet_read_options, &bytes_read, @@ -815,14 +818,14 @@ impl ArrowReader { pub(crate) async fn open_parquet_file( data_file_path: &str, file_io: &FileIO, + runtime: &Runtime, file_size_in_bytes: u64, parquet_read_options: ParquetReadOptions, bytes_read: &Arc, key_metadata: Option<&[u8]>, ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> { - let parquet_file = file_io.new_input(data_file_path)?; - let counting_reader = - CountingFileRead::new(parquet_file.reader().await?, Arc::clone(bytes_read)); + let reader = RuntimeFileRead::open(file_io, data_file_path, runtime.io()).await?; + let counting_reader = CountingFileRead::new(reader, Arc::clone(bytes_read)); Self::build_parquet_reader( Box::new(counting_reader), file_size_in_bytes, diff --git a/crates/iceberg/src/arrow/runtime_file_read.rs b/crates/iceberg/src/arrow/runtime_file_read.rs new file mode 100644 index 0000000000..d205ba3d43 --- /dev/null +++ b/crates/iceberg/src/arrow/runtime_file_read.rs @@ -0,0 +1,152 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Dispatch storage operations without moving Parquet decoding off its caller's runtime. + +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; + +use crate::Result; +use crate::io::{FileIO, FileRead}; +use crate::runtime::RuntimeHandle; + +pub(super) struct RuntimeFileRead { + inner: Arc, + io: RuntimeHandle, +} + +impl RuntimeFileRead { + pub(super) async fn open(file_io: &FileIO, path: &str, io: &RuntimeHandle) -> Result { + let file_io = file_io.clone(); + let path = path.to_owned(); + let inner = io + .spawn(async move { file_io.new_input(&path)?.reader().await }) + .abort_on_drop() + .await??; + Ok(Self { + inner: Arc::from(inner), + io: io.clone(), + }) + } +} + +#[async_trait::async_trait] +impl FileRead for RuntimeFileRead { + async fn read(&self, range: Range) -> Result { + let inner = self.inner.clone(); + self.io + .spawn(async move { inner.read(range).await }) + .abort_on_drop() + .await? + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{Error, ErrorKind, Runtime}; + + struct CheckedRead; + + #[async_trait::async_trait] + impl FileRead for CheckedRead { + async fn read(&self, range: Range) -> Result { + assert_eq!(std::thread::current().name(), Some("parquet-io")); + tokio::task::yield_now().await; + assert_eq!(std::thread::current().name(), Some("parquet-io")); + if range.is_empty() { + return Err(Error::new(ErrorKind::DataInvalid, "empty test range")); + } + Ok(Bytes::from_static(b"data")) + } + } + + #[test] + fn reads_run_on_io_and_errors_reach_cpu() { + let io = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .thread_name("parquet-io") + .enable_all() + .build() + .unwrap(); + let cpu = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .thread_name("parquet-cpu") + .enable_all() + .build() + .unwrap(); + let runtime = Runtime::new_with_split(&io, &cpu); + let reader = RuntimeFileRead { + inner: Arc::new(CheckedRead), + io: runtime.io().clone(), + }; + cpu.block_on(runtime.cpu().spawn(async move { + assert_eq!( + reader.read(0..4).await.unwrap(), + Bytes::from_static(b"data") + ); + assert_eq!(std::thread::current().name(), Some("parquet-cpu")); + assert_eq!( + reader.read(0..0).await.unwrap_err().kind(), + ErrorKind::DataInvalid + ); + })) + .unwrap(); + } + + struct PendingRead { + started: tokio::sync::Notify, + dropped: Arc, + } + + #[async_trait::async_trait] + impl FileRead for PendingRead { + async fn read(&self, _range: Range) -> Result { + struct OnDrop(Arc); + impl Drop for OnDrop { + fn drop(&mut self) { + self.0.notify_one(); + } + } + let _guard = OnDrop(self.dropped.clone()); + self.started.notify_one(); + std::future::pending().await + } + } + + #[tokio::test] + async fn dropping_read_aborts_storage_request() { + let dropped = Arc::new(tokio::sync::Notify::new()); + let inner = Arc::new(PendingRead { + started: tokio::sync::Notify::new(), + dropped: dropped.clone(), + }); + let reader = RuntimeFileRead { + inner: inner.clone(), + io: Runtime::current().io().clone(), + }; + let mut read = Box::pin(reader.read(0..4)); + assert!(futures::poll!(&mut read).is_pending()); + inner.started.notified().await; + drop(read); + tokio::time::timeout(std::time::Duration::from_secs(5), dropped.notified()) + .await + .unwrap(); + } +}