From 0be5094cec4444fe0106aca93dc66fb7511ee567 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 29 Jul 2026 13:35:04 +0100 Subject: [PATCH] Even more extreme morsels Signed-off-by: Adam Gutglick --- Cargo.lock | 1 + vortex-datafusion/Cargo.toml | 1 + vortex-datafusion/src/persistent/morsel.rs | 236 ++++++++++++++++----- vortex-layout/src/scan/repeated_scan.rs | 13 +- vortex-utils/src/parallelism.rs | 1 + 5 files changed, 196 insertions(+), 56 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 51188781c9a..c3161a59832 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9657,6 +9657,7 @@ dependencies = [ "insta", "itertools 0.14.0", "object_store", + "parking_lot", "rstest", "tempfile", "tokio", diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 49fb22d4f59..4fc0bbcc702 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -33,6 +33,7 @@ datafusion-pruning = { workspace = true } futures = { workspace = true } itertools = { workspace = true } object_store = { workspace = true } +parking_lot = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "fs"] } tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } diff --git a/vortex-datafusion/src/persistent/morsel.rs b/vortex-datafusion/src/persistent/morsel.rs index 4ddbc1c1549..bdf2dfe93a8 100644 --- a/vortex-datafusion/src/persistent/morsel.rs +++ b/vortex-datafusion/src/persistent/morsel.rs @@ -6,8 +6,11 @@ //! Morsel-driven I/O support for Vortex files. use std::ops::Range; +use std::pin::Pin; use std::sync::Arc; use std::sync::Weak; +use std::task::Context; +use std::task::Poll; use arrow_array::RecordBatch; use arrow_array::RecordBatchOptions; @@ -41,11 +44,14 @@ use datafusion_physical_expr_adapter::replace_columns_with_literals; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::metrics::MetricBuilder; use datafusion_pruning::FilePruner; +use futures::Stream; use futures::StreamExt; use futures::TryStreamExt; +use futures::stream; use futures::stream::BoxStream; use itertools::Itertools; use object_store::path::Path; +use parking_lot::Mutex; use tracing::Instrument; use vortex::array::VortexSessionExecute; use vortex::dtype::FieldMask; @@ -75,7 +81,6 @@ use crate::convert::schema::calculate_physical_schema; use crate::metrics::PARTITION_LABEL; use crate::metrics::PATH_LABEL; use crate::persistent::cache::CachedVortexMetadata; -use crate::persistent::stream::PrunableStream; use crate::reader::VortexReaderFactory; /// Creates morsel planners for Vortex files. @@ -319,8 +324,9 @@ enum State { vxf: VortexFile, }, PreparedScan { - scan: RepeatedScan, - file_pruner: Option, + scan: Arc, + file_pruner: Option>>, + morsel_ranges: Vec>>, output_schema: SchemaRef, session: VortexSession, stream_target_field: Field, @@ -363,6 +369,7 @@ impl State { State::PreparedScan { scan, file_pruner, + morsel_ranges, output_schema, session, stream_target_field, @@ -371,6 +378,7 @@ impl State { } => Ok(State::PreparedScan { scan, file_pruner, + morsel_ranges, output_schema, session, stream_target_field, @@ -401,6 +409,7 @@ impl std::fmt::Debug for State { .finish(), Self::PreparedScan { file_pruner, + morsel_ranges, output_schema, stream_target_field, file_location, @@ -413,6 +422,7 @@ impl std::fmt::Debug for State { "file_pruner", &file_pruner.as_ref().map(|_| ""), ) + .field("morsel_ranges", morsel_ranges) .field("output_schema", output_schema) .field("stream_target_field", stream_target_field) .field("file_location", file_location) @@ -451,18 +461,140 @@ impl std::fmt::Debug for FileOpenState { } struct VortexStreamMorsel { - inner: BoxStream<'static, DFResult>, + scan: Arc, + row_range: Option>, + file_pruner: Option>>, + output_schema: SchemaRef, + session: VortexSession, + stream_target_field: Field, + file_location: Path, + projector: Projector, } impl std::fmt::Debug for VortexStreamMorsel { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("VortexStreamMorsel").finish_non_exhaustive() + f.debug_struct("VortexStreamMorsel") + .field("row_range", &self.row_range) + .field( + "file_pruner", + &self.file_pruner.as_ref().map(|_| ""), + ) + .field("output_schema", &self.output_schema) + .field("stream_target_field", &self.stream_target_field) + .field("file_location", &self.file_location) + .field("projector", &self.projector) + .finish_non_exhaustive() } } impl Morsel for VortexStreamMorsel { fn into_stream(self: Box) -> BoxStream<'static, DFResult> { - self.inner + let Self { + scan, + row_range, + file_pruner, + output_schema, + session, + stream_target_field, + file_location, + projector, + } = *self; + + let stream = match scan.execute_array_stream(row_range) { + Ok(stream) => stream, + Err(error) => { + return stream::once(async move { + Err(exec_datafusion_err!( + "Failed to create Vortex stream: {error}" + )) + }) + .boxed(); + } + }; + + let stream = stream + // Convert to Arrow inline on the polling thread: DataFusion sources are expected + // to do their CPU work inside `poll_next`, and spawning this onto the blocking + // pool oversubscribes the CPU. + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + chunk.and_then(|chunk| { + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) + }) + .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e)) + .map(move |batch| { + let batch = if projector.projection().as_ref().is_empty() { + batch + } else { + batch.and_then(|b| projector.project_batch(&b)) + }?; + + let (_, columns, row_count) = batch.into_parts(); + RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + ) + .map_err(Into::into) + }) + .boxed(); + + if let Some(file_pruner) = file_pruner { + SharedPrunableStream::new(file_pruner, stream).boxed() + } else { + stream + } + } +} + +struct SharedPrunableStream { + file_pruner: Arc>, + stream: Option>>, +} + +impl SharedPrunableStream { + fn new( + file_pruner: Arc>, + stream: BoxStream<'static, DFResult>, + ) -> Self { + Self { + file_pruner, + stream: Some(stream), + } + } +} + +impl Stream for SharedPrunableStream { + type Item = DFResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let should_prune = { + let file_pruner = Arc::clone(&self.file_pruner); + let mut file_pruner = file_pruner.lock(); + file_pruner.should_prune() + }; + + match should_prune { + Ok(true) => { + self.stream.take(); + Poll::Ready(None) + } + Ok(false) => match self.stream.as_mut() { + Some(stream) => stream.poll_next_unpin(cx), + None => Poll::Ready(None), + }, + Err(error) => { + self.stream.take(); + Poll::Ready(Some(Err(error))) + } + } } } @@ -721,9 +853,8 @@ impl MorselPlanner for VortexMorselPlanner { scan_builder = scan_builder.with_limit(limit); } - if let Some(concurrency) = scan_concurrency { - scan_builder = scan_builder.with_concurrency(concurrency); - } + let splits_per_morsel = scan_concurrency.unwrap_or(1).max(1); + scan_builder = scan_builder.with_concurrency(1); let scan = scan_builder .with_metrics_registry(metrics_registry) @@ -732,16 +863,30 @@ impl MorselPlanner for VortexMorselPlanner { .with_ordered(has_output_ordering) .prepare() .map_err(|e| exec_datafusion_err!("Failed to prepare Vortex scan: {e}"))?; + let morsel_ranges = if scan.has_limit() { + vec![None] + } else { + split_ranges_into_morsel_ranges(scan.split_ranges(None), splits_per_morsel) + .into_iter() + .map(Some) + .collect() + }; + if morsel_ranges.is_empty() { + return Ok(None); + } let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); let file_location = file.object_meta.location; + let scan = Arc::new(scan); + let file_pruner = file_pruner.map(|file_pruner| Arc::new(Mutex::new(file_pruner))); Ok(Some(MorselPlan::new().with_planners(vec![Box::new( Self { state: State::PreparedScan { scan, file_pruner, + morsel_ranges, output_schema, session, stream_target_field, @@ -754,57 +899,30 @@ impl MorselPlanner for VortexMorselPlanner { State::PreparedScan { scan, file_pruner, + morsel_ranges, output_schema, session, stream_target_field, file_location, projector, } => { - let stream = scan - .execute_array_stream(None) - .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - // Convert to Arrow inline on the polling thread: DataFusion sources are expected - // to do their CPU work inside `poll_next`, and spawning this onto the blocking - // pool oversubscribes the CPU. - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - chunk.and_then(|chunk| { - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) + let morsels = morsel_ranges + .into_iter() + .map(|row_range| { + Box::new(VortexStreamMorsel { + scan: Arc::clone(&scan), + row_range, + file_pruner: file_pruner.as_ref().map(Arc::clone), + output_schema: Arc::clone(&output_schema), + session: session.clone(), + stream_target_field: stream_target_field.clone(), + file_location: file_location.clone(), + projector: projector.clone(), + }) as Box }) - .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e)) - .map(move |batch| { - let batch = if projector.projection().as_ref().is_empty() { - batch - } else { - batch.and_then(|b| projector.project_batch(&b)) - }?; - - let (_, columns, row_count) = batch.into_parts(); - RecordBatch::try_new_with_options( - Arc::clone(&output_schema), - columns, - &RecordBatchOptions::new().with_row_count(Some(row_count)), - ) - .map_err(Into::into) - }) - .boxed(); - - let stream = if let Some(file_pruner) = file_pruner { - PrunableStream::new(file_pruner, stream).boxed() - } else { - stream - }; + .collect(); - Ok(Some(MorselPlan::new().with_morsels(vec![ - Box::new(VortexStreamMorsel { inner: stream }) as Box, - ]))) + Ok(Some(MorselPlan::new().with_morsels(morsels))) } State::Done => Ok(None), new_state => Ok(Some( @@ -814,6 +932,20 @@ impl MorselPlanner for VortexMorselPlanner { } } +fn split_ranges_into_morsel_ranges( + split_ranges: Vec>, + splits_per_morsel: usize, +) -> Vec> { + split_ranges + .chunks(splits_per_morsel.max(1)) + .filter_map(|chunk| { + let first = chunk.first()?; + let last = chunk.last()?; + (first.start < last.end).then_some(first.start..last.end) + }) + .collect() +} + fn natural_split_ranges_for_file( natural_split_ranges: &DashMap]>>, path: &Path, diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 7c3ea7973b4..b907e5b0b73 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -31,7 +31,6 @@ use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; use vortex_scan::selection::Selection; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; @@ -287,7 +286,8 @@ impl RepeatedScan { } } - fn split_ranges(&self, row_range: Option>) -> Vec> { + /// Returns the row ranges this prepared scan will execute for the requested row range. + pub fn split_ranges(&self, row_range: Option>) -> Vec> { let selection_range: Option> = match &self.selection { Selection::IncludeByIndex(buf) if !buf.is_empty() => { Some(buf[0]..buf[buf.len() - 1] + 1) @@ -343,6 +343,11 @@ impl RepeatedScan { } } + /// Returns whether this prepared scan has a row limit. + pub fn has_limit(&self) -> bool { + self.limit.is_some() || self.row_limit.is_some() + } + fn task_context(&self) -> Arc { Arc::new(TaskContext { filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), @@ -383,12 +388,12 @@ impl RepeatedScan { &self, row_range: Option>, ) -> VortexResult { - let num_workers = get_available_parallelism().unwrap_or(1); + // let num_workers = get_available_parallelism().unwrap_or(1); let row_limit = self .row_limit .clone() .or_else(|| self.limit.map(RowLimit::new)); - let concurrency = self.concurrency * num_workers; + let concurrency = self.concurrency; let handle = self.session.handle(); // With both a filter and a limit we cannot know each split's output row count ahead of diff --git a/vortex-utils/src/parallelism.rs b/vortex-utils/src/parallelism.rs index 14b251e8bfb..75f85b0e3a5 100644 --- a/vortex-utils/src/parallelism.rs +++ b/vortex-utils/src/parallelism.rs @@ -10,6 +10,7 @@ use std::sync::LazyLock; /// This is currently implemented using [`std::thread::available_parallelism`], but might change in the future. /// /// Returns `None` if the underlying functions fails. +#[allow(dead_code)] pub fn get_available_parallelism() -> Option { #[allow(clippy::disallowed_methods)] static PARALLELISM: LazyLock> =