diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 1a83d374fcb..3d3983f43ee 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -120,7 +120,7 @@ struct Args { #[arg(long, value_delimiter = ',', value_parser = value_parser!(Format))] formats: Vec, - /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H). + /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H, TPC-DS and ClickBench). #[arg(long, action = clap::ArgAction::Set)] scan_v2: Option, @@ -137,8 +137,10 @@ async fn main() -> anyhow::Result<()> { let args = Args::parse(); let opts = Opts::from(args.options); let scan_v2 = args.scan_v2.unwrap_or_else(|| { - matches!(args.benchmark, BenchmarkArg::TpcH) - && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0") + matches!( + args.benchmark, + BenchmarkArg::TpcH | BenchmarkArg::TpcDS | BenchmarkArg::ClickBench + ) && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0") }); set_join_set_tracer(get_static_tracer())?; diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 1444b1879ae..36d2c32ab07 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -4,6 +4,7 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::sync::Arc; +use std::sync::Weak; use arrow_schema::DataType; use arrow_schema::Schema; @@ -35,6 +36,7 @@ use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr::LexRequirement; use datafusion_physical_plan::ExecutionPlan; @@ -67,9 +69,11 @@ use vortex::scalar::Scalar; use vortex::scalar::ScalarValue as VortexScalarValue; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; +use vortex_utils::aliases::dash_map::DashMap; use super::cache::CachedVortexMetadata; use super::sink::VortexSink; +use super::source::SharedScanState; use super::source::VortexSource; use crate::PrecisionExt as _; use crate::convert::ExpressionConverter; @@ -127,6 +131,7 @@ const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usi /// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html /// [`ListingOptions`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingOptions.html pub struct VortexFormat { + scan_states: DashMap<(ObjectStoreUrl, usize), Weak>, session: VortexSession, opts: VortexTableOptions, expression_converter: Option>, @@ -459,12 +464,27 @@ impl VortexFormat { /// Creates a format with explicit [`VortexTableOptions`]. pub fn new_with_options(session: VortexSession, opts: VortexTableOptions) -> Self { Self { + scan_states: DashMap::default(), session, opts, expression_converter: None, } } + fn shared_scan_state( + &self, + store_url: ObjectStoreUrl, + store: Arc, + ) -> Arc { + let store_id = Arc::as_ptr(&store) as *const () as usize; + let mut entry = self.scan_states.entry((store_url, store_id)).or_default(); + entry.upgrade().unwrap_or_else(|| { + let shared = Arc::new(SharedScanState::for_store(store)); + *entry = Arc::downgrade(&shared); + shared + }) + } + /// Returns the format-specific configuration that will be copied into the /// [`VortexSource`] created for a scan. pub fn options(&self) -> &VortexTableOptions { @@ -748,6 +768,11 @@ impl FileFormat for VortexFormat { source = source .with_file_metadata_cache(state.runtime_env().cache_manager.get_file_metadata_cache()); + if source.options().scan_v2 && source.vortex_reader_factory.is_none() { + let store_url = file_scan_config.object_store_url.clone(); + let store = state.runtime_env().object_store(&store_url)?; + source = source.with_scan_state(self.shared_scan_state(store_url, store)); + } let conf = FileScanConfigBuilder::from(file_scan_config) .with_source(Arc::new(source)) @@ -821,6 +846,7 @@ mod tests { use datafusion_physical_expr::expressions as df_expr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_plan::filter_pushdown::PushedDown; + use object_store::memory::InMemory; use rstest::rstest; use vortex::array::stats::StatsSet; use vortex::dtype::Nullability; @@ -834,6 +860,28 @@ mod tests { use crate::convert::ProcessedProjection; use crate::convert::scalar_from_df; + #[test] + fn scan_state_is_scoped_to_store_instance_and_released() { + let format = VortexFormat::new(VortexSession::default()); + let url = ObjectStoreUrl::local_filesystem(); + let store: Arc = Arc::new(InMemory::new()); + let weak_store = Arc::downgrade(&store); + let first = format.shared_scan_state(url.clone(), Arc::clone(&store)); + let second = format.shared_scan_state(url.clone(), Arc::clone(&store)); + assert!(Arc::ptr_eq(&first, &second)); + + let other = format.shared_scan_state(url, Arc::new(InMemory::new())); + assert!(!Arc::ptr_eq(&first, &other)); + + let weak_state = Arc::downgrade(&first); + drop(store); + assert!(weak_store.upgrade().is_some()); + drop(first); + drop(second); + assert!(weak_state.upgrade().is_none()); + assert!(weak_store.upgrade().is_none()); + } + #[test] fn lossy_bound_cast_does_not_prove_constant() -> VortexResult<()> { let dtype = DType::Extension( diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 208691f0b0e..f83e3170a4c 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -45,6 +45,7 @@ use futures::FutureExt; use futures::StreamExt; use futures::TryStreamExt; use futures::stream; +use object_store::ObjectMeta; use object_store::path::Path; use tracing::Instrument; use vortex::array::VortexSessionExecute; @@ -83,6 +84,8 @@ use crate::metrics::PATH_LABEL; use crate::persistent::cache::CachedVortexMetadata; use crate::persistent::reader::VortexReaderFactory; use crate::persistent::scan_v2::ScanDriver; +use crate::persistent::scan_v2::rebalance_batches; +use crate::persistent::source::FileVersion; use crate::persistent::stream::PrunableStream; #[derive(Clone)] @@ -120,10 +123,10 @@ pub(crate) struct VortexOpener { /// Partitions of one file read through one source, so a chunk straddling a partition /// boundary and the file's metadata segments are read once, as they are for V1 through the /// shared layout reader. - pub segment_sources: Arc>>, + pub segment_sources: Arc>>, /// The V2 scan cache each open file's partitions share: its lowered layout, filter and /// projection plans, pruned zones and dictionary values, built once per file. - pub scan_caches: Arc>, + pub scan_caches: Arc>, /// Shared V1 full-file natural splits keyed by file path. pub natural_splits: Arc>>, /// Shared physical Arrow schemas of each file, keyed by file path. @@ -140,6 +143,7 @@ pub(crate) struct VortexOpener { /// How many partitions the engine runs the scan on, so a V2 scan parallelises inside a /// partition only as far as the engine has not. pub partitions: usize, + pub batch_size: usize, } impl FileOpener for VortexOpener { @@ -179,6 +183,7 @@ impl FileOpener for VortexOpener { let has_output_ordering = self.has_output_ordering; let scan_concurrency = self.scan_concurrency; let partitions = self.partitions.max(1); + let batch_size = self.batch_size; let scan_v2 = self.scan_v2; let expr_converter = Arc::clone(&self.expression_converter); @@ -277,6 +282,34 @@ impl FileOpener for VortexOpener { return Ok(stream::empty().boxed()); } + let byte_range = file + .range + .as_ref() + .map(|range| -> DFResult> { + Ok(Range { + start: u64::try_from(range.start).map_err(|_| { + exec_datafusion_err!("Vortex file range start is negative") + })?, + end: u64::try_from(range.end).map_err(|_| { + exec_datafusion_err!("Vortex file range end is negative") + })?, + }) + }) + .transpose()?; + + // Coarse row boundaries are independent of the projection and predicates. Empty + // byte partitions can finish before adapting or converting those expressions. + let v2_rows = if scan_v2 { + byte_range + .clone() + .map(|bytes| partition_row_range(bytes, file.object_meta.size, vxf.row_count())) + } else { + None + }; + if v2_rows.as_ref().is_some_and(Range::is_empty) { + return Ok(stream::empty().boxed()); + } + // This is the expected arrow types of the actual columns in the file, which might have different types // from the unified logical schema or miss. Every partition opening this file computes // the same schema, so it is shared across them. @@ -403,21 +436,6 @@ impl FileOpener for VortexOpener { .transpose()?; let projector = leftover_projection.make_projector(&stream_schema)?; - let byte_range = file - .range - .as_ref() - .map(|range| -> DFResult> { - Ok(Range { - start: u64::try_from(range.start).map_err(|_| { - exec_datafusion_err!("Vortex file range start is negative") - })?, - end: u64::try_from(range.end).map_err(|_| { - exec_datafusion_err!("Vortex file range end is negative") - })?, - }) - }) - .transpose()?; - let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); let residual_path = file.path().clone(); let native_limit = limit @@ -433,23 +451,16 @@ impl FileOpener for VortexOpener { )?; Ok(RecordBatch::from(arrow.as_struct().clone())) }; + let rebalance = limit.is_none() && native_filter.is_none() && residual_filter.is_none(); let stream = if scan_v2 { let mut scan_file = vxf.scan_file(); - scan_file.segments = shared_segment_source( - &segment_sources, - &file.object_meta.location, - scan_file.segments, - ); - scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta.location)); + scan_file.segments = + shared_segment_source(&segment_sources, &file.object_meta, scan_file.segments); + scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta)); let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file) .with_projection(scan_projection) .with_some_filter(native_filter); - if let Some(bytes) = byte_range { - // Every partition maps the same byte boundary to the same row boundary. - let rows = partition_row_range(bytes, file.object_meta.size, vxf.row_count()); - if rows.is_empty() { - return Ok(stream::empty().boxed()); - } + if let Some(rows) = v2_rows { builder = builder.with_row_range(rows); } if let Some(selection) = file @@ -467,6 +478,7 @@ impl FileOpener for VortexOpener { .div_ceil(partitions) .max(1) }); + builder = builder.with_prefetch_batches(if concurrency <= 1 { 4 } else { 1 }); let handle = vxf.session().handle(); // Lowering and planning a wide file is real CPU: it runs where compute runs. let scan = handle @@ -495,7 +507,12 @@ impl FileOpener for VortexOpener { }; let driver = ScanDriver::new(handle, concurrency, has_output_ordering); - Ok(driver.run(scan, root, map_batch)) + let batches = driver.run(scan, root, map_batch); + Ok(if rebalance { + rebalance_batches(batches, batch_size) + } else { + batches + }) } else { // Share V1 layout readers across partitions of the file. let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { @@ -693,10 +710,15 @@ fn contains_lambda(projection: &ProjectionExprs) -> DFResult { Ok(false) } -/// Returns the scan cache other partitions of `path` are planning with, or registers a new one. -/// A scan holds its cache until its stream ends, so the cache lives as long as one scan. -fn shared_scan_cache(scan_caches: &DashMap, path: &Path) -> ScanCache { - match scan_caches.entry(path.clone()) { +/// Returns the scan cache for this object version, or registers a new one. +/// Live streams for this file version own the cache; the registry holds only a weak reference. +/// The last stream releases the cache, even if the multi-file scan is still running. +fn shared_scan_cache( + scan_caches: &DashMap, + meta: &ObjectMeta, +) -> ScanCache { + let path = &meta.location; + match scan_caches.entry(FileVersion(meta.clone())) { Entry::Occupied(mut entry) => { if let Some(shared) = entry.get().upgrade() { tracing::trace!("reusing scan cache for {path}"); @@ -715,14 +737,15 @@ fn shared_scan_cache(scan_caches: &DashMap, path: &Path) -> } } -/// Returns the segment source other partitions of `path` are reading through, or registers +/// Returns the segment source other scans of this object version read through, or registers /// `fresh` as that source. Like the layout readers, a source lives only as long as one scan. fn shared_segment_source( - segment_sources: &DashMap>, - path: &Path, + segment_sources: &DashMap>, + meta: &ObjectMeta, fresh: Arc, ) -> Arc { - match segment_sources.entry(path.clone()) { + let path = &meta.location; + match segment_sources.entry(FileVersion(meta.clone())) { Entry::Occupied(mut entry) => { if let Some(shared) = entry.get().upgrade() { tracing::trace!("reusing segment source for {path}"); @@ -1148,6 +1171,7 @@ mod tests { segment_sources: Default::default(), scan_caches: Default::default(), partitions: 1, + batch_size: 8192, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1558,6 +1582,7 @@ mod tests { segment_sources: Default::default(), scan_caches: Default::default(), partitions: 1, + batch_size: 8192, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1650,6 +1675,7 @@ mod tests { segment_sources: Default::default(), scan_caches: Default::default(), partitions: 1, + batch_size: 8192, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1833,6 +1859,7 @@ mod tests { segment_sources: Default::default(), scan_caches: Default::default(), partitions: 1, + batch_size: 8192, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1898,6 +1925,7 @@ mod tests { segment_sources: Default::default(), scan_caches: Default::default(), partitions: 1, + batch_size: 8192, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -2114,6 +2142,7 @@ mod tests { segment_sources: Default::default(), scan_caches: Default::default(), partitions: 1, + batch_size: 8192, natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, diff --git a/vortex-datafusion/src/persistent/reader.rs b/vortex-datafusion/src/persistent/reader.rs index 66afb43d75a..b0f2aa8b620 100644 --- a/vortex-datafusion/src/persistent/reader.rs +++ b/vortex-datafusion/src/persistent/reader.rs @@ -69,11 +69,14 @@ impl VortexReaderFactory for DefaultVortexReaderFactory { file: &PartitionedFile, session: &VortexSession, ) -> DFResult> { - Ok(Arc::new(ObjectStoreReadAt::new_with_allocator( - Arc::clone(&self.object_store), - file.path().clone(), - session.handle(), - session.allocator(), - )) as _) + Ok(Arc::new( + ObjectStoreReadAt::new_with_allocator( + Arc::clone(&self.object_store), + file.path().clone(), + session.handle(), + session.allocator(), + ) + .with_cached_file_handle(), + ) as _) } } diff --git a/vortex-datafusion/src/persistent/scan_v2.rs b/vortex-datafusion/src/persistent/scan_v2.rs index 10ea7b969ea..87bef2e84e5 100644 --- a/vortex-datafusion/src/persistent/scan_v2.rs +++ b/vortex-datafusion/src/persistent/scan_v2.rs @@ -11,11 +11,13 @@ //! of the file's segment source. The DataFusion thread polling the stream does no scan work //! itself. +use std::collections::VecDeque; use std::future; use std::sync::Arc; use std::task::Poll; use arrow_array::RecordBatch; +use datafusion_common::arrow::compute::concat_batches; use futures::SinkExt; use futures::StreamExt; use futures::TryStreamExt; @@ -241,3 +243,112 @@ fn chunk_stream( .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); stream::select(recv, errors).boxed() } + +/// Keep fewer than three target batches so EOF does not leave a small fragment for downstream joins. +pub(crate) fn rebalance_batches( + input: BoxStream<'static, VortexResult>, + batch_size: usize, +) -> BoxStream<'static, VortexResult> { + if batch_size == 0 || batch_size > usize::MAX / 3 { + return input; + } + let state = BatchRebalance { + input, + batch_size, + pending: VecDeque::new(), + pending_rows: 0, + ready: VecDeque::new(), + done: false, + }; + stream::try_unfold(state, |mut state| async move { + loop { + if let Some(batch) = state.ready.pop_front() { + return Ok(Some((batch, state))); + } + if state.done { + return Ok(None); + } + match state.input.next().await { + Some(batch) => state.push(batch?)?, + None => { + state.done = true; + let count = state.pending_rows.div_ceil(state.batch_size); + if let Some(rows) = state.pending_rows.checked_div(count) { + let extra = state.pending_rows % count; + for index in 0..count { + let batch = state.take(rows + usize::from(index < extra))?; + state.ready.push_back(batch); + } + } + } + } + } + }) + .boxed() +} + +struct BatchRebalance { + input: BoxStream<'static, VortexResult>, + batch_size: usize, + pending: VecDeque, + pending_rows: usize, + ready: VecDeque, + done: bool, +} + +impl BatchRebalance { + fn take(&mut self, rows: usize) -> VortexResult { + let mut remaining = rows; + let mut pieces = Vec::new(); + while remaining > 0 { + let batch = self + .pending + .pop_front() + .ok_or_else(|| vortex_err!("Missing pending scan rows"))?; + let take = remaining.min(batch.num_rows()); + if take == batch.num_rows() { + pieces.push(batch); + } else { + pieces.push(batch.slice(0, take)); + self.pending + .push_front(batch.slice(take, batch.num_rows() - take)); + } + remaining -= take; + } + self.pending_rows -= rows; + if pieces.len() == 1 { + return Ok(pieces.remove(0)); + } + concat_batches(&pieces[0].schema(), &pieces) + .map_err(|error| vortex_err!("Failed to rebalance scan batch tails: {error}")) + } + + fn push(&mut self, batch: RecordBatch) -> VortexResult<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + self.pending_rows += batch.num_rows(); + self.pending.push_back(batch); + if self.pending_rows < self.batch_size * 3 { + return Ok(()); + } + let mut emit = + (self.pending_rows - self.batch_size * 2) / self.batch_size * self.batch_size; + while emit > 0 { + let front = self.pending[0].num_rows(); + // Preserve full spans as slices; only a boundary crossing needs concatenation. + let rows = if front >= self.batch_size { + (front / self.batch_size * self.batch_size).min(emit) + } else { + self.batch_size.min(emit) + }; + let batch = self.take(rows)?; + self.ready.push_back(batch); + emit -= rows; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-datafusion/src/persistent/scan_v2/tests.rs b/vortex-datafusion/src/persistent/scan_v2/tests.rs new file mode 100644 index 00000000000..000f082f4a7 --- /dev/null +++ b/vortex-datafusion/src/persistent/scan_v2/tests.rs @@ -0,0 +1,157 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::Int32Array; +use arrow_array::RecordBatch; +use arrow_array::RecordBatchOptions; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::Schema; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use rstest::rstest; +use vortex::error::VortexResult; +use vortex::error::vortex_err; + +use super::rebalance_batches; + +#[rstest] +#[case(0, vec![3, 0, 10, 2])] +#[case(1, vec![1, 2, 8, 2, 1, 1])] +#[case(4, vec![4, 4, 4, 3])] +#[case(32, vec![15])] +#[tokio::test] +async fn rebalancing_preserves_rows_and_nulls( + #[case] batch_size: usize, + #[case] expected_lengths: Vec, +) -> VortexResult<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + true, + )])); + let values = (0..15) + .map(|i| (i % 3 != 0).then_some(i)) + .collect::>(); + let expected = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))])?; + let input = vec![ + Ok(expected.slice(0, 3)), + Ok(expected.slice(3, 0)), + Ok(expected.slice(3, 10)), + Ok(expected.slice(13, 2)), + ]; + let batches: Vec = rebalance_batches(stream::iter(input).boxed(), batch_size) + .try_collect() + .await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + expected_lengths + ); + let mut offset = 0; + for batch in batches { + assert_eq!(batch, expected.slice(offset, batch.num_rows())); + offset += batch.num_rows(); + } + assert_eq!(offset, expected.num_rows()); + Ok(()) +} + +#[tokio::test] +async fn rebalancing_preserves_zero_column_row_counts() -> VortexResult<()> { + let batch = RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(3)), + )?; + let input = stream::iter(vec![Ok(batch.clone()), Ok(batch)]).boxed(); + let batches: Vec = rebalance_batches(input, 4).try_collect().await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![3, 3] + ); + assert!(batches.iter().all(|batch| batch.num_columns() == 0)); + Ok(()) +} + +#[tokio::test] +async fn rebalancing_propagates_errors_with_a_pending_tail() -> VortexResult<()> { + let batch = RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(3)), + )?; + let input = stream::iter(vec![Ok(batch), Err(vortex_err!("input failed"))]).boxed(); + let mut aligned = rebalance_batches(input, 4); + let error = aligned + .next() + .await + .expect("expected input error") + .expect_err("input failed"); + assert!(error.to_string().contains("input failed")); + assert!(aligned.next().await.is_none()); + Ok(()) +} + +#[tokio::test] +async fn rebalancing_avoids_a_one_row_final_batch() -> VortexResult<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let expected = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from_iter_values(0..17))])?; + let batches: Vec = + rebalance_batches(stream::iter(vec![Ok(expected.clone())]).boxed(), 8) + .try_collect() + .await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![6, 6, 5] + ); + assert_eq!(batches[0], expected.slice(0, 6)); + assert_eq!(batches[1], expected.slice(6, 6)); + assert_eq!(batches[2], expected.slice(12, 5)); + Ok(()) +} + +#[test] +fn pending_rows_stay_below_three_target_batches() -> VortexResult<()> { + let batch = RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(1)), + )?; + let mut state = super::BatchRebalance { + input: stream::empty().boxed(), + batch_size: 8, + pending: std::collections::VecDeque::new(), + pending_rows: 0, + ready: std::collections::VecDeque::new(), + done: false, + }; + let mut emitted = 0; + for count in 1..=1000 { + state.push(batch.clone())?; + assert!(state.pending_rows < 24); + assert!(state.pending.len() < 24); + for batch in state.ready.drain(..) { + assert_eq!(batch.num_rows() % 8, 0); + emitted += batch.num_rows(); + } + assert_eq!(emitted + state.pending_rows, count); + } + Ok(()) +} diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 644fea4e71c..9d7d918fbc5 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use std::sync::Weak; @@ -30,6 +32,7 @@ use datafusion_physical_plan::filter_pushdown::FilterPushdownPropagation; use datafusion_physical_plan::filter_pushdown::PushedDown; use datafusion_physical_plan::filter_pushdown::PushedDownPredicate; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use object_store::ObjectMeta; use object_store::ObjectStore; use object_store::path::Path; use vortex::file::VORTEX_FILE_EXTENSION; @@ -49,6 +52,38 @@ use crate::convert::exprs::ExpressionConverter; use crate::persistent::reader::DefaultVortexReaderFactory; use crate::persistent::reader::VortexReaderFactory; +#[derive(Clone, PartialEq, Eq)] +pub(crate) struct FileVersion(pub ObjectMeta); + +impl Hash for FileVersion { + fn hash(&self, state: &mut H) { + self.0.location.hash(state); + self.0.last_modified.hash(state); + self.0.size.hash(state); + self.0.e_tag.hash(state); + self.0.version.hash(state); + } +} + +#[derive(Default)] +/// File state shared by concurrent scans of one object store. Weak entries release the +/// plans and retained bytes when the last scan finishes. +pub(crate) struct SharedScanState { + // Keep the store alive while its pointer identifies this state in the format's cache. + _object_store: Option>, + segment_sources: Arc>>, + scan_caches: Arc>, +} + +impl SharedScanState { + pub(super) fn for_store(object_store: Arc) -> Self { + Self { + _object_store: Some(object_store), + ..Self::default() + } + } +} + /// File scan implementation for reading one or more `.vortex` files. /// /// `VortexSource` is the lower-level read component underneath @@ -202,9 +237,7 @@ pub struct VortexSource { /// /// Sharing the readers allows us to only read every layout once from the file, even across partitions. layout_readers: Arc>>, - /// Shared V2 segment sources, keyed by path and alive only while a scan reads the file. - segment_sources: Arc>>, - scan_caches: Arc>, + scan_state: Arc, /// Shared full-file natural splits keyed by path. natural_splits: Arc>>, /// Shared physical Arrow schemas of each file, keyed by path. @@ -216,9 +249,15 @@ pub struct VortexSource { file_metadata_cache: Option>, /// Options controlling scan planning and execution behavior. options: VortexTableOptions, + batch_size: usize, } impl VortexSource { + pub(super) fn with_scan_state(mut self, state: Arc) -> Self { + self.scan_state = state; + self + } + /// Creates a new `VortexSource` for a table schema and [`VortexSession`]. /// /// The new source starts with: @@ -241,8 +280,7 @@ impl VortexSource { vortex_predicate: None, df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), - segment_sources: Arc::new(DashMap::default()), - scan_caches: Arc::new(DashMap::default()), + scan_state: Arc::new(SharedScanState::default()), natural_splits: Arc::new(DashMap::default()), file_schemas: Arc::new(DashMap::default()), expression_converter, @@ -251,6 +289,7 @@ impl VortexSource { file_metadata_cache: None, ordered: false, options: VortexTableOptions::default(), + batch_size: 8192, } } @@ -368,8 +407,8 @@ impl VortexSource { metrics_registry: Arc::clone(&self.vx_metrics_registry), df_metrics: self.df_metrics.clone(), layout_readers: Arc::clone(&self.layout_readers), - segment_sources: Arc::clone(&self.segment_sources), - scan_caches: Arc::clone(&self.scan_caches), + segment_sources: Arc::clone(&self.scan_state.segment_sources), + scan_caches: Arc::clone(&self.scan_state.scan_caches), natural_splits: Arc::clone(&self.natural_splits), file_schemas: Arc::clone(&self.file_schemas), has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered, @@ -379,6 +418,7 @@ impl VortexSource { scan_concurrency: self.options.scan_concurrency, scan_v2: self.options.scan_v2, partitions: base_config.file_groups.len(), + batch_size: self.batch_size, }; Ok(opener) @@ -399,9 +439,10 @@ impl FileSource for VortexSource { )?)) } - fn with_batch_size(&self, _batch_size: usize) -> Arc { - // DataSourceExec applies BatchSplitStream after the FileSource stream. - Arc::new(self.clone()) + fn with_batch_size(&self, batch_size: usize) -> Arc { + let mut source = self.clone(); + source.batch_size = batch_size; + Arc::new(source) } fn filter(&self) -> Option> { diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index b354ea705e8..a50af288ea3 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -1,8 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(not(target_arch = "wasm32"))] +use std::fs::File; use std::io; use std::sync::Arc; +use std::sync::OnceLock; use futures::FutureExt; use futures::SinkExt; @@ -35,6 +38,12 @@ use crate::std_file::read_exact_at; /// Default number of concurrent requests to allow. pub const DEFAULT_CONCURRENCY: usize = 192; +struct LocalFile { + #[cfg(not(target_arch = "wasm32"))] + file: File, + size: u64, +} + /// An object store backed I/O source. pub struct ObjectStoreReadAt { store: Arc, @@ -43,6 +52,7 @@ pub struct ObjectStoreReadAt { handle: Handle, allocator: BufferAllocatorRef, concurrency: usize, + file_cache: Option>>>, coalesce_config: Option, } @@ -72,10 +82,20 @@ impl ObjectStoreReadAt { handle, allocator, concurrency: DEFAULT_CONCURRENCY, + file_cache: None, coalesce_config: Some(CoalesceConfig::object_storage()), } } + /// Retain the first local file handle for this reader and its in-flight reads. + /// + /// Use only for immutable files: replacing the path does not change the opened file. + /// Stream-backed object stores continue to issue a GET for each range. + pub fn with_cached_file_handle(mut self) -> Self { + self.file_cache = Some(Arc::new(OnceLock::new())); + self + } + /// Set the concurrency for this source. pub fn with_concurrency(mut self, concurrency: usize) -> Self { self.concurrency = concurrency; @@ -94,6 +114,7 @@ async fn read_object_store_range( path: ObjectPath, io_handle: Handle, allocator: BufferAllocatorRef, + file_cache: Option>>>, request: ReadAtRequest, ) -> VortexResult { let ReadAtRequest { @@ -106,6 +127,20 @@ async fn read_object_store_range( // SAFETY: each return path checks that every byte was initialized. unsafe { buffer.set_len(length) }; + #[cfg(not(target_arch = "wasm32"))] + if let Some(file) = file_cache.as_ref().and_then(|cache| cache.get()).cloned() { + let buffer = io_handle + .spawn_blocking(move || { + read_exact_at(&file.file, buffer.as_mut_slice(), offset)?; + Ok::<_, io::Error>(buffer) + }) + .await + .map_err(io::Error::other)?; + return Ok(BufferHandle::new_host(buffer.freeze())); + } + #[cfg(target_arch = "wasm32")] + let _ = file_cache; + let response = store .get_opts( &path, @@ -118,13 +153,23 @@ async fn read_object_store_range( let buffer = match response.payload { #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => io_handle - .spawn_blocking(move || { - read_exact_at(&file, buffer.as_mut_slice(), range.start)?; - Ok::<_, io::Error>(buffer) - }) - .await - .map_err(io::Error::other)?, + GetResultPayload::File(file, _) => { + let opened = LocalFile { + file, + size: response.meta.size, + }; + let file = match file_cache { + Some(cache) => Arc::clone(cache.get_or_init(|| Arc::new(opened))), + None => Arc::new(opened), + }; + io_handle + .spawn_blocking(move || { + read_exact_at(&file.file, buffer.as_mut_slice(), range.start)?; + Ok::<_, io::Error>(buffer) + }) + .await + .map_err(io::Error::other)? + } #[cfg(target_arch = "wasm32")] GetResultPayload::File(..) => { unreachable!("File payload not supported on wasm32") @@ -173,6 +218,9 @@ impl VortexReadAt for ObjectStoreReadAt { } fn size(&self) -> BoxFuture<'static, VortexResult> { + if let Some(file) = self.file_cache.as_ref().and_then(|cache| cache.get()) { + return futures::future::ready(Ok(file.size)).boxed(); + } let store = Arc::clone(&self.store); let path = self.path.clone(); async move { @@ -202,6 +250,7 @@ impl VortexReadAt for ObjectStoreReadAt { path, io_handle, allocator, + self.file_cache.clone(), ReadAtRequest::new(offset, length, alignment), )) .boxed() @@ -217,6 +266,7 @@ impl VortexReadAt for ObjectStoreReadAt { let handle = self.handle.clone(); let allocator = self.allocator.clone(); let concurrency = self.concurrency.max(1); + let file_cache = self.file_cache.clone(); let (mut send, recv) = mpsc::channel(concurrency); let io_handle = handle.clone(); @@ -229,9 +279,12 @@ impl VortexReadAt for ObjectStoreReadAt { let path = path.clone(); let io_handle = io_handle.clone(); let allocator = allocator.clone(); + let file_cache = file_cache.clone(); async move { - let result = - read_object_store_range(store, path, io_handle, allocator, request).await; + let result = read_object_store_range( + store, path, io_handle, allocator, file_cache, request, + ) + .await; (request, result) } }); @@ -262,6 +315,8 @@ mod tests { use std::sync::atomic::Ordering; use object_store::PutPayload; + #[cfg(unix)] + use object_store::local::LocalFileSystem; use object_store::memory::InMemory; use super::*; @@ -331,6 +386,62 @@ mod tests { Ok(()) } + #[cfg(unix)] + #[tokio::test] + async fn cached_file_handle_survives_path_replacement() -> VortexResult<()> { + let runtime = Arc::new(CountingExecutor::default()) as Arc; + let handle = Handle::new(Arc::downgrade(&runtime)); + let directory = tempfile::tempdir()?; + let path = directory.path().join("test.bin"); + std::fs::write(&path, TEST_DATA)?; + let store = + Arc::new(LocalFileSystem::new_with_prefix(directory.path())?) as Arc; + let reader = ObjectStoreReadAt::new( + Arc::clone(&store), + ObjectPath::from("test.bin"), + handle.clone(), + ) + .with_cached_file_handle(); + let first = reader.read_at(7, 5, Alignment::new(1)).await?; + assert_eq!(first.to_host().await.as_slice(), b"store"); + + let replacement = directory.path().join("replacement.bin"); + std::fs::write(&replacement, b"replacement")?; + std::fs::rename(replacement, path)?; + assert_eq!(reader.size().await?, TEST_DATA.len() as u64); + let requests = Arc::from([ + ReadAtRequest::new(0, 6, Alignment::new(1)), + ReadAtRequest::new(7, 5, Alignment::new(1)), + ]); + let results = reader.read_ranges(requests).collect::>().await; + assert_eq!(results.len(), 2); + for (request, result) in results { + let buffer = result?; + let start = usize::try_from(request.offset)?; + assert_eq!( + buffer.to_host().await.as_slice(), + &TEST_DATA[start..start + request.length] + ); + } + let cache = reader.file_cache.as_ref().map(Arc::downgrade); + let pending = reader.read_at(7, 5, Alignment::new(1)); + drop(reader); + assert_eq!(pending.await?.to_host().await.as_slice(), b"store"); + assert!(cache.is_some_and(|cache| cache.upgrade().is_none())); + + let fresh = ObjectStoreReadAt::new(store, ObjectPath::from("test.bin"), handle); + assert_eq!( + fresh + .read_at(0, 11, Alignment::new(1)) + .await? + .to_host() + .await + .as_slice(), + b"replacement" + ); + Ok(()) + } + #[tokio::test] async fn read_ranges_uses_one_io_task() -> anyhow::Result<()> { let executor = Arc::new(CountingExecutor::default()); diff --git a/vortex-layout/src/plan/physical/operators/mod.rs b/vortex-layout/src/plan/physical/operators/mod.rs index b0f5937d71d..65fdd3bdcb4 100644 --- a/vortex-layout/src/plan/physical/operators/mod.rs +++ b/vortex-layout/src/plan/physical/operators/mod.rs @@ -27,11 +27,16 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::Filter; +use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_err; pub(crate) use whole::*; +use crate::arrays::shared::Shared; use crate::plan::physical::operator::Cx; use crate::plan::physical::pipeline::Inlet; @@ -51,6 +56,23 @@ fn take_rows(inlet: &mut Inlet<'_>, len: usize) -> VortexResult { left -= front.len(); pieces.extend(inlet.take()); } else { + // Misaligned fields can split one lazy filter across several output batches. + // Materialize selected values once when a consumer first needs either slice. + // Extension encodings can push filtering into their component arrays. + let filters_components = front.dtype().is_extension() + && !front.is::() + && front.children_iter().any(|child| child.is::()); + if front.is::() || filters_components { + *front = Shared::new(front.clone()).into_array(); + } else if let Some(dict) = front.as_opt::() + && dict.codes().is::() + { + *front = DictArray::try_new( + Shared::new(dict.codes().clone()).into_array(), + dict.values().clone(), + )? + .into_array(); + } let rest = front.slice(left..front.len())?; pieces.push(std::mem::replace(front, rest).slice(0..left)?); left = 0; diff --git a/vortex-layout/src/plan/physical/operators/take.rs b/vortex-layout/src/plan/physical/operators/take.rs index b844996117d..5b468354a74 100644 --- a/vortex-layout/src/plan/physical/operators/take.rs +++ b/vortex-layout/src/plan/physical/operators/take.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use super::UNBOUNDED; use super::join_inlet; use crate::arrays::shared::Shared; +use crate::arrays::shared::SharedArrayExt; use crate::plan::physical::TakePlan; use crate::plan::physical::operator::Cx; use crate::plan::physical::operator::Input; @@ -85,7 +86,20 @@ impl Operator for TakeSource { } }; Ok(match cx.inlet(CODES).take() { - Some(codes) => Step::Final(DictArray::try_new(codes, values)?.into_array()), + Some(codes) => { + // Expose encoded values to sparse take kernels instead of forcing the shared + // dictionary to decode in full for a small projection. + let values = if (values.dtype().is_utf8() || values.dtype().is_binary()) + && values.len() >= 1024 + && codes.len().saturating_mul(4) < values.len() + && let Some(shared) = values.as_opt::() + { + shared.current_array_ref().clone() + } else { + values + }; + Step::Final(DictArray::try_new(codes, values)?.into_array()) + } None => Step::End, }) } diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 816c3566d7e..00d7f0df050 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -656,8 +656,20 @@ impl Scan { plan: &PlanRef, rows: Range, mask: &Mask, + ) -> VortexResult> { + Self::segments_with_value_cache(session, plan, rows, mask, &ValueCache::default()) + } + + /// Lists reads after accounting for whole values already retained by the file's scans. + pub(crate) fn segments_with_value_cache( + session: &VortexSession, + plan: &PlanRef, + rows: Range, + mask: &Mask, + values: &ValueCache, ) -> VortexResult> { let mut core = Core::new(session.clone(), 0); + core.values = values.clone(); core.compile(plan, rows, mask)?; let mut segments = core.segments; segments.sort_unstable(); diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 804d56553d9..e021744bc64 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -14,16 +14,23 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::lit; +use vortex_array::expr::pack; use vortex_array::expr::root; +use vortex_array::expr::select; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::extension::datetime::Timestamp; use vortex_array::serde::SerializeOptions; use vortex_buffer::Alignment; use vortex_buffer::Buffer; @@ -614,6 +621,37 @@ fn scan_keeps_the_selection_with_or_without_a_filter(#[case] sel: Sel) -> Vortex Ok(()) } +/// Splitting filtered dictionary codes to match another field preserves nulls and row order. +#[rstest] +#[case::alternating(Sel::EveryOther)] +#[case::sparse(Sel::Rows(&[0, 4, 7, 10, 14, 19]))] +fn pack_splits_filtered_dictionary_codes( + #[case] sel: Sel, + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, +) -> VortexResult<()> { + let mut store = Store::default(); + let values = PrimitiveArray::from_option_iter([Some(10_i32), None, Some(30)]).into_array(); + let codes = PrimitiveArray::from_option_iter( + (0..ROWS).map(|row| (row % 5 != 0).then_some((row % 3) as u8)), + ) + .into_array(); + let dict = DictArray::try_new(codes.clone(), values.clone())?.into_array(); + let other = PrimitiveArray::from_iter(0..ROWS as i64).into_array(); + let expected = + StructArray::from_fields(&[("dict", dict), ("other", other.clone())])?.into_array(); + let dictionary = DictLayout::new(store.flat(&values)?, store.flat(&codes)?).into_layout(); + let layout = StructLayout::new( + ROWS, + expected.dtype().clone(), + vec![dictionary, store.chunked(&other, &[7, 5, 8])?], + ) + .into_layout(); + let plan = lower(&layout)?; + let mask = sel.mask(ROWS as usize); + let run = run(&store, &plan, 0..ROWS, mask.clone(), delivery(order))?; + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} + /// A take reads its values over their whole domain and its codes over the selection, including /// when a predicate has been pushed onto the values, and emits nothing before the values are /// whole. Later scans sharing its value cache reuse the values and read only the codes; a scan @@ -621,9 +659,14 @@ fn scan_keeps_the_selection_with_or_without_a_filter(#[case] sel: Sel) -> Vortex #[rstest] #[case::values(false)] #[case::predicate(true)] -fn take_waits_for_whole_values_and_keeps_them(#[case] predicate: bool) -> VortexResult<()> { +fn take_waits_for_whole_values_and_keeps_them( + #[case] predicate: bool, + #[values(3, 1536)] values_len: usize, +) -> VortexResult<()> { let mut store = Store::default(); - let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array(); + let values = + VarBinViewArray::from_iter_str((0..values_len).map(|row| ["a", "b", "c"][row % 3])) + .into_array(); let codes = PrimitiveArray::from_iter((0..ROWS).map(|v| (v % 3) as u8)).into_array(); let layout = DictLayout::new(store.flat(&values)?, store.flat(&codes)?).into_layout(); let mut plan = lower(&layout)?; @@ -730,7 +773,7 @@ fn a_segment_two_readers_need_is_read_by_each() -> VortexResult<()> { .clone(); let plan = PackPlan::try_new( fields, - vortex_array::dtype::Nullability::NonNullable, + Nullability::NonNullable, ROWS, vec![column.clone(), column], None, @@ -921,3 +964,101 @@ fn bound_selection_keeps_whole_dictionary_values() -> VortexResult<()> { ); Ok(()) } + +#[rstest] +#[case::alternating(Sel::EveryOther)] +#[case::sparse(Sel::Rows(&[0, 4, 7, 10, 14, 19]))] +fn pack_splits_filtered_values( + #[case] sel: Sel, + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, +) -> VortexResult<()> { + let mut store = Store::default(); + let values = + PrimitiveArray::from_option_iter((0..ROWS).map(|row| (row % 5 != 0).then_some(row as i32))) + .into_array(); + let other = PrimitiveArray::from_iter(0..ROWS as i64).into_array(); + let expected = + StructArray::from_fields(&[("values", values.clone()), ("other", other.clone())])? + .into_array(); + let layout = StructLayout::new( + ROWS, + expected.dtype().clone(), + vec![store.flat(&values)?, store.chunked(&other, &[7, 5, 8])?], + ) + .into_layout(); + let plan = lower(&layout)?; + let mask = sel.mask(ROWS as usize); + let run = run(&store, &plan, 0..ROWS, mask.clone(), delivery(order))?; + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} + +#[rstest] +#[case::select(false)] +#[case::aliases(true)] +fn direct_projection_preserves_field_order_and_aliases(#[case] aliases: bool) -> VortexResult<()> { + let mut store = Store::default(); + let (plan, _) = two_columns(&mut store)?; + let a = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let b = PrimitiveArray::from_iter((0..ROWS as i64).map(|v| v * 10)).into_array(); + let (expression, expected) = if aliases { + ( + pack( + [ + ("x", get_item("b", root())), + ("y", get_item("a", root())), + ("z", get_item("a", root())), + ], + Nullability::NonNullable, + ), + StructArray::from_fields(&[("x", b), ("y", a.clone()), ("z", a)])?.into_array(), + ) + } else { + ( + select(["b", "a"], root()), + StructArray::from_fields(&[("b", b), ("a", a)])?.into_array(), + ) + }; + let expression = expression.bind(plan.dtype())?; + let plan = EvalPlan::try_new(expression, plan)?.into_plan(); + let rows = 2..18; + let mask = Sel::EveryOther.mask(16); + for plan in [plan.clone(), optimize(plan, &SESSION)?] { + for order in [Delivery::Fifo, Delivery::Lifo] { + let result = run(&store, &plan, rows.clone(), mask.clone(), delivery(order))?; + assert_view(&expected, &rows, &mask, result.arrays)?; + } + } + Ok(()) +} + +#[rstest] +#[case::alternating(Sel::EveryOther)] +#[case::sparse(Sel::Rows(&[0, 4, 7, 10, 14, 19]))] +fn pack_splits_filtered_extension_components( + #[case] sel: Sel, + #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, +) -> VortexResult<()> { + let mut store = Store::default(); + let values = ExtensionArray::new( + Timestamp::new(TimeUnit::Microseconds, Nullability::Nullable).erased(), + PrimitiveArray::from_option_iter( + (0..ROWS).map(|row| (row % 5 != 0).then_some(row as i64 * 1_000_000)), + ) + .into_array(), + ) + .into_array(); + let other = PrimitiveArray::from_iter(0..ROWS as i64).into_array(); + let expected = + StructArray::from_fields(&[("values", values.clone()), ("other", other.clone())])? + .into_array(); + let layout = StructLayout::new( + ROWS, + expected.dtype().clone(), + vec![store.flat(&values)?, store.chunked(&other, &[7, 5, 8])?], + ) + .into_layout(); + let plan = lower(&layout)?; + let mask = sel.mask(ROWS as usize); + let run = run(&store, &plan, 0..ROWS, mask.clone(), delivery(order))?; + assert_view(&expected, &(0..ROWS), &mask, run.arrays) +} diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index d3886de19ad..1ec9871b1ec 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod cache; + use std::borrow::Cow; use std::ops::Range; @@ -11,7 +13,6 @@ use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; @@ -25,13 +26,12 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_expr::descendent_bound_annotations; -use vortex_expr::make_bound_free_field_annotator; -use vortex_expr::transform::partition_bound; use vortex_mask::Mask; +use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use self::cache::PartitionCache; use crate::plan::physical::Eval; use crate::plan::physical::EvalPlan; use crate::plan::physical::Plan; @@ -331,23 +331,38 @@ impl PlanParentReduceRule for ExpressionPackRule { child: &Plan, parent: &Plan, _child_idx: usize, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { if child.dtype().is_nullable() { return Ok(None); } let expression = parent.expression(); + if let Some(projected) = project_fields(child, expression)? { + return Ok(Some(projected)); + } let fields = child.fields(); - let referenced_fields = - descendent_bound_annotations(expression, make_bound_free_field_annotator(fields)) - .get(&ExactBoundExpr(expression.clone())) - .vortex_expect("Bound expression missing free-field annotations") - .clone(); - let expanded_root = expanded_struct_root(child.dtype(), fields)?; - let expanded = expand_struct_root(expression.clone(), &expanded_root, fields)?; - let partitioned = - partition_bound(expanded.clone(), make_bound_free_field_annotator(fields))?; + // Most predicates reference one field. Find it without allocating annotations for + // every expression node; struct-valued expressions use the general partitioner below. + let mut single_field = None; + if collect_single_field(expression, &mut single_field) + && let Some(name) = single_field + { + let index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, index)?; + let lowered = step_into_struct_field(expression.clone(), name, field.dtype().clone())?; + if lowered.is_root_bound_to(field.dtype()) { + return Ok(Some(EvalPlan::try_new(lowered, field)?.into_plan())); + } + } + let cached = session + .get::() + .partition(expression, fields)?; + let expanded = &cached.expanded; + let referenced_fields = &cached.referenced_fields; + let partitioned = &cached.partitions; if partitioned.partition_names.is_empty() { let selected_indices = fields @@ -384,59 +399,31 @@ impl PlanParentReduceRule for ExpressionPackRule { vortex_err!("Struct expression references unknown field '{name}'") })?; let field = field_plan(child, index)?; - let lowered = step_into_struct_field(expanded, name, field.dtype().clone())?; + let lowered = step_into_struct_field(expanded.clone(), name, field.dtype().clone())?; return Ok(Some(EvalPlan::try_new(lowered, field)?.into_plan())); } - let residual = partitioned.root; - let mut collapsed = Vec::with_capacity(partitioned.partitions.len()); - let mut field_expressions = vec![None; fields.nfields()]; - for index in 0..partitioned.partitions.len() { - let name = &partitioned.partition_names[index]; - let partition = &partitioned.partitions[index]; - let field_index = fields.find(name).ok_or_else(|| { - vortex_err!("Struct expression references unknown field '{name}'") - })?; - let field = field_plan(child, field_index)?; - let lowered = if let Some(pack) = partition - .as_scalar() - .and_then(|scalar_fn| scalar_fn.as_opt::()) - && partition.children().len() == 1 - { - let value_name = pack - .names - .get(0) - .ok_or_else(|| vortex_err!("Struct expression partition pack is empty"))?; - collapsed.push((name.clone(), value_name.clone())); - partition.children()[0].clone() - } else { - partition.clone() - }; - let lowered = step_into_struct_field(lowered, name, field.dtype().clone())?; - field_expressions[field_index] = Some(lowered); - } - - let mut fields_changed = partitioned.partition_names.len() != fields.nfields(); - let mut pruned_fields = Vec::with_capacity(partitioned.partition_names.len()); - for (field_index, expression) in field_expressions.into_iter().enumerate() { - let Some(expression) = expression else { - continue; - }; - let field = field_plan(child, field_index)?; - let field = if is_identity_expression(&expression, field.dtype())? { + let lowered = cached + .lowered + .as_ref() + .vortex_expect("Multiple partitions are lowered"); + let fields_changed = lowered.fields_changed; + let mut pruned_fields = Vec::with_capacity(lowered.fields.len()); + for lowered_field in &lowered.fields { + let field = field_plan(child, lowered_field.index)?; + let field = if lowered_field.identity { field } else { - fields_changed = true; - EvalPlan::try_new(expression, field)?.into_plan() + EvalPlan::try_new(lowered_field.expression.clone(), field)?.into_plan() }; - pruned_fields.push((field_name(fields, field_index)?, field)); + pruned_fields.push((field_name(fields, lowered_field.index)?, field)); } let rewritten = if fields_changed { child.with_pruned_fields(pruned_fields)?.into_plan() } else { child.to_plan() }; - let residual = rewrite_partition_root(residual, rewritten.dtype().clone(), &collapsed)?; + let residual = lowered.residual.clone(); if !fields_changed && residual == *expression { return Ok(None); @@ -446,6 +433,93 @@ impl PlanParentReduceRule for ExpressionPackRule { } } +/// Projects direct field references without partitioning an expression tree. +fn project_fields( + child: &Plan, + expression: &BoundExpression, +) -> VortexResult> { + if expression.dtype().is_nullable() { + return Ok(None); + } + let fields = child.fields(); + let inputs = if let Some(selection) = expression.as_opt::() { let names = selection.normalize_to_included_fields(fields.names())?; + let root = &node.children()[0]; let children = names .iter() .map(|name| { - let index = fields - .find(name) - .vortex_expect("normalized selection fields must exist in the root"); - expanded_root.children()[index].clone() + BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()]) }) - .collect(); + .collect::>>()?; return Ok(Transformed { value: bound_pack(names, children)?, changed: true, diff --git a/vortex-layout/src/plan/physical/plans/pack/cache.rs b/vortex-layout/src/plan/physical/plans/pack/cache.rs new file mode 100644 index 00000000000..6d9ed61a2ce --- /dev/null +++ b/vortex-layout/src/plan/physical/plans/pack/cache.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::collections::VecDeque; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use parking_lot::Mutex; +use rustc_hash::FxHasher; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; +use vortex_array::scalar_fn::fns::pack::Pack as PackFn; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_expr::descendent_bound_annotations; +use vortex_expr::make_bound_free_field_annotator; +use vortex_expr::transform::BoundPartitionedExpr; +use vortex_expr::transform::partition_bound_annotations; +use vortex_session::SessionVar; + +use super::expand_struct_root; +use super::is_identity_expression; +use super::rewrite_partition_root; +use super::step_into_struct_field; + +const MAX_ENTRIES: usize = 64; +const MAX_NODES: usize = 256; + +#[derive(Debug)] +pub(super) struct Partitioned { + pub(super) expanded: BoundExpression, + pub(super) referenced_fields: Vec, + pub(super) partitions: BoundPartitionedExpr, + pub(super) lowered: Option, +} + +#[derive(Debug)] +pub(super) struct Lowered { + pub(super) fields: Vec, + pub(super) residual: BoundExpression, + pub(super) fields_changed: bool, +} + +#[derive(Debug)] +pub(super) struct LoweredField { + pub(super) index: usize, + pub(super) expression: BoundExpression, + pub(super) identity: bool, +} + +#[derive(Debug)] +struct Entry { + hash: u64, + expression: BoundExpression, + fields: StructFields, + result: Arc, +} + +/// Pure expression rewrites shared across file scans, without retaining file plans or values. +#[derive(Clone, Debug, Default)] +pub(super) struct PartitionCache(Arc>>); + +impl SessionVar for PartitionCache { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +impl PartitionCache { + pub(super) fn partition( + &self, + expression: &BoundExpression, + fields: &StructFields, + ) -> VortexResult> { + let Some(hash) = fingerprint(expression) else { + return partition(expression, fields).map(Arc::new); + }; + let mut entries = self.0.lock(); + if let Some(entry) = entries.iter().find(|entry| { + entry.hash == hash && entry.fields == *fields && entry.expression == *expression + }) { + return Ok(Arc::clone(&entry.result)); + } + // Partitioning is a pure rewrite. Holding this lock avoids concurrent scans repeating + // the same allocation-heavy work; it never waits for file IO or invokes an optimizer. + let result = Arc::new(partition(expression, fields)?); + if fingerprint(&result.expanded).is_some() { + if entries.len() == MAX_ENTRIES { + entries.pop_front(); + } + entries.push_back(Entry { + hash, + expression: expression.clone(), + fields: fields.clone(), + result: Arc::clone(&result), + }); + } + Ok(result) + } +} + +fn partition(expression: &BoundExpression, fields: &StructFields) -> VortexResult { + let expanded = expand_struct_root(expression.clone(), fields)?; + let annotations = + descendent_bound_annotations(&expanded, make_bound_free_field_annotator(fields)); + let referenced_fields = annotations + .get(&ExactBoundExpr(expanded.clone())) + .vortex_expect("Bound expression missing free-field annotations") + .iter() + .cloned() + .collect(); + let partitions = partition_bound_annotations(expanded.clone(), annotations)?; + let lowered = if partitions.partition_names.len() > 1 { + Some(lower_partitions(&partitions, fields)?) + } else { + None + }; + Ok(Partitioned { + expanded, + referenced_fields, + partitions, + lowered, + }) +} + +fn lower_partitions( + partitioned: &BoundPartitionedExpr, + fields: &StructFields, +) -> VortexResult { + let mut collapsed = Vec::with_capacity(partitioned.partitions.len()); + let mut lowered = Vec::with_capacity(partitioned.partitions.len()); + for (name, partition) in partitioned + .partition_names + .iter() + .zip(partitioned.partitions.iter()) + { + let index = fields + .find(name) + .ok_or_else(|| vortex_err!("Struct expression references unknown field '{name}'"))?; + let dtype = fields.field(name).vortex_expect("Partition field exists"); + let expression = if let Some(pack) = partition + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + && partition.children().len() == 1 + { + let value_name = pack + .names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition pack is empty"))?; + collapsed.push((name.clone(), value_name.clone())); + partition.children()[0].clone() + } else { + partition.clone() + }; + let expression = step_into_struct_field(expression, name, dtype.clone())?; + let identity = is_identity_expression(&expression, &dtype)?; + lowered.push(LoweredField { + index, + expression, + identity, + }); + } + lowered.sort_by_key(|field| field.index); + let fields_changed = + lowered.len() != fields.nfields() || lowered.iter().any(|field| !field.identity); + let root_dtype = DType::Struct( + StructFields::from_iter(lowered.iter().map(|field| { + ( + fields.names()[field.index].clone(), + field.expression.dtype().clone(), + ) + })), + Nullability::NonNullable, + ); + let residual = rewrite_partition_root(partitioned.root.clone(), root_dtype, &collapsed)?; + Ok(Lowered { + fields: lowered, + residual, + fields_changed, + }) +} + +fn fingerprint(expression: &BoundExpression) -> Option { + let mut remaining = MAX_NODES; + let mut hasher = FxHasher::default(); + hash_expression(expression, &mut remaining, &mut hasher)?; + Some(hasher.finish()) +} + +fn hash_expression( + expression: &BoundExpression, + remaining: &mut usize, + hasher: &mut FxHasher, +) -> Option<()> { + *remaining = remaining.checked_sub(1)?; + if let Some(function) = expression.as_scalar() { + hasher.write_u8(1); + function.hash(hasher); + expression.children().len().hash(hasher); + for child in expression.children() { + hash_expression(child, remaining, hasher)?; + } + } else { + hasher.write_u8(0); + } + Some(()) +} + +#[cfg(test)] +mod tests { + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::expr::lit; + use vortex_array::expr::root; + + use super::*; + + fn fields(name: &str) -> StructFields { + StructFields::from_iter([(name, DType::Bool(Nullability::NonNullable))]) + } + + #[test] + fn independent_bindings_share_partitions() -> VortexResult<()> { + let cache = PartitionCache::default(); + let fields = fields("a"); + let dtype = DType::Struct(fields.clone(), Nullability::NonNullable); + let first = cache.partition(&root().bind(&dtype)?, &fields)?; + let second = cache.partition(&root().bind(&dtype)?, &fields)?; + assert!(Arc::ptr_eq(&first, &second)); + Ok(()) + } + + #[test] + fn constant_partitions_keep_scopes_separate() -> VortexResult<()> { + let cache = PartitionCache::default(); + let a = fields("a"); + let b = fields("b"); + let expression = lit(1_i32).bind(&DType::Struct(a.clone(), Nullability::NonNullable))?; + let first = cache.partition(&expression, &a)?; + let second = cache.partition(&expression, &b)?; + assert!(!Arc::ptr_eq(&first, &second)); + Ok(()) + } + + #[test] + fn retained_entries_are_bounded() -> VortexResult<()> { + let cache = PartitionCache::default(); + let fields = fields("a"); + let dtype = DType::Struct(fields.clone(), Nullability::NonNullable); + for value in 0..=MAX_ENTRIES { + cache.partition(&lit(i32::try_from(value)?).bind(&dtype)?, &fields)?; + } + assert_eq!(cache.0.lock().len(), MAX_ENTRIES); + Ok(()) + } +} diff --git a/vortex-layout/src/plan/physical/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs index 5c45d2cae4f..5dde579a927 100644 --- a/vortex-layout/src/plan/physical/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -124,7 +124,28 @@ pub fn plan_row_idx_expression( expression: BoundExpression, child: PlanRef, ) -> VortexResult { - let partitioned = partition_bound(expression.clone(), |node| { + match expression_domains(&expression) { + (true, false) => return Ok(EvalPlan::try_new(expression, child)?.into_plan()), + (false, true) => { + let expression = replace_row_idx(expression)?; + let values = RowIdxPlan::new(child.row_count()).into_plan(); + return Ok(EvalPlan::try_new(expression, values)?.into_plan()); + } + (false, false) => { + let row_domain = PackPlan::try_new( + StructFields::empty(), + Nullability::NonNullable, + child.row_count(), + Vec::new(), + None, + )? + .into_plan(); + return Ok(EvalPlan::try_new(expression, row_domain)?.into_plan()); + } + (true, true) => {} + } + + let partitioned = partition_bound(expression, |node| { if node .as_scalar() .is_some_and(|scalar_fn| scalar_fn.is::()) @@ -137,31 +158,6 @@ pub fn plan_row_idx_expression( } })?; - if partitioned.partition_annotations.is_empty() { - let row_domain = PackPlan::try_new( - StructFields::empty(), - Nullability::NonNullable, - child.row_count(), - Vec::new(), - None, - )? - .into_plan(); - return Ok(EvalPlan::try_new(expression, row_domain)?.into_plan()); - } - - if partitioned.partition_annotations.len() == 1 { - return match partitioned.partition_annotations[0] { - RowIdxExpressionPartition::RowIdx => { - let expression = replace_row_idx(expression)?; - let values = RowIdxPlan::new(child.row_count()).into_plan(); - Ok(EvalPlan::try_new(expression, values)?.into_plan()) - } - RowIdxExpressionPartition::Child => { - Ok(EvalPlan::try_new(expression, child)?.into_plan()) - } - }; - } - vortex_ensure_eq!( partitioned.partition_annotations.len(), 2, @@ -237,6 +233,29 @@ pub fn plan_row_idx_expression( Ok(EvalPlan::try_new(residual, partitions.into_plan())?.into_plan()) } +/// Finds data and row-index references without allocating per-node annotations. +fn expression_domains(expression: &BoundExpression) -> (bool, bool) { + if expression.is_root() { + return (true, false); + } + if expression + .as_scalar() + .is_some_and(|function| function.is::()) + { + return (false, true); + } + let mut domains = (false, false); + for child in expression.children() { + let child_domains = expression_domains(child); + domains.0 |= child_domains.0; + domains.1 |= child_domains.1; + if domains == (true, true) { + break; + } + } + domains +} + fn replace_row_idx(expression: BoundExpression) -> VortexResult { Ok(expression .transform_down(|node| { diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index c86edcf4d30..8570274d069 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -649,14 +649,15 @@ fn empty_projection_prunes_row_idx_child_fields() -> VortexResult<()> { let plan = make_row_idx_plan(projection, make_plan(layout)?)?; let optimized = optimize(plan, &SESSION)?; - let projection = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("optimized plan has no projection expression"))?; - let child = projection.child_plan()?; - let empty_struct = child + let empty_struct = optimized .as_opt::() .ok_or_else(|| vortex_err!("empty projection did not prune the RowIdx child"))?; + assert_eq!(empty_struct.row_count(), 3); + assert_eq!( + empty_struct.dtype(), + &DType::Struct(StructFields::empty(), Nullability::NonNullable) + ); assert_eq!(empty_struct.nfields(), 0); assert_eq!(empty_struct.children().len(), 0); Ok(()) diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 382425d9bb0..2e065983a21 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -7,6 +7,7 @@ use std::ops::Range; use std::sync::Arc; use std::vec::IntoIter; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::request::IoConsumer; use vortex_io::request::IoRequestId; @@ -25,9 +26,11 @@ use super::selection::plan_selection; /// The row grid every chunk boundary lies on. pub(crate) const ROW_BLOCK_SIZE: usize = 8192; /// The rows of a projection chunk. -pub(crate) const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; +pub(crate) const MAX_ROWS_PER_CHUNK: usize = 32 * ROW_BLOCK_SIZE; /// The rows of a filtered chunk: large, so a conjunct's reads coalesce and its mask is cheap. -pub(crate) const FILTERED_ROWS_PER_CHUNK: usize = 16 * MAX_ROWS_PER_CHUNK; +pub(crate) const FILTERED_ROWS_PER_CHUNK: usize = 128 * ROW_BLOCK_SIZE; +// Filtered scans retain smaller work items for selective conjuncts. +const MIN_FILTERED_ROWS_PER_CHUNK: usize = 32 * ROW_BLOCK_SIZE; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; /// Limit estimated bytes per chunk without changing the writer-aligned row granularity. @@ -35,13 +38,42 @@ pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> u if row_count <= MAX_ROWS_PER_CHUNK as u64 { return MAX_ROWS_PER_CHUNK; } - let bytes: u64 = locations + let blocks = blocks_per_target_bytes(total_bytes(locations), row_count); + let blocks = blocks.clamp(1, (MAX_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64); + usize::try_from(blocks).vortex_expect("at most MAX_ROWS_PER_CHUNK blocks") * ROW_BLOCK_SIZE +} + +/// Rows per chunk of a filtered scan over `rows` rows of the file. +/// +/// Large chunks amortize filter evaluation over cheap projections. A projection that reads many +/// bytes per row, or a row range that would otherwise give fewer chunks than threads, gets smaller +/// chunks so the surviving rows still spread across the runtime's threads. +pub(crate) fn filtered_rows_per_chunk( + rows: u64, + row_count: u64, + projected_bytes: u64, + parallelism: usize, +) -> usize { + let max_blocks = (FILTERED_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64; + let min_blocks = (MIN_FILTERED_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64; + let by_bytes = blocks_per_target_bytes(projected_bytes, row_count).max(1); + let by_threads = (rows / parallelism.max(1) as u64 / ROW_BLOCK_SIZE as u64).max(min_blocks); + let blocks = by_bytes.min(by_threads).min(max_blocks); + usize::try_from(blocks).vortex_expect("at most FILTERED_ROWS_PER_CHUNK blocks") * ROW_BLOCK_SIZE +} + +/// The row blocks that hold [`TARGET_BYTES_PER_CHUNK`] when `bytes` spread over `row_count` rows. +fn blocks_per_target_bytes(bytes: u64, row_count: u64) -> u64 { + let bytes_per_row = bytes.div_ceil(row_count.max(1)).max(1); + TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64 +} + +/// The serialized bytes of every segment in `locations`. +pub(crate) fn total_bytes(locations: &[SegmentLocation]) -> u64 { + locations .iter() .map(|location| u64::from(location.length)) - .sum(); - let bytes_per_row = bytes.div_ceil(row_count).max(1); - let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; - blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE + .sum() } /// Fixed boundaries in the file's row domain, clipped to the requested range. diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index d3d7b640892..34863b4c93a 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::BitAnd; use std::ops::Range; use std::sync::Arc; @@ -9,6 +10,8 @@ use rustc_hash::FxHashSet; use vortex_array::VortexSessionExecute; use vortex_array::builtins::ArrayBuiltins; use vortex_array::expr::BoundExpression; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; @@ -33,6 +36,7 @@ use super::adaptor::AdaptorStep; use super::adaptor::Piece; use super::adaptor::ScanAdaptor; use super::projection::ProjectionPlanner; +use crate::plan::physical::Eval; use crate::plan::physical::FilterPlan; use crate::plan::physical::PlanRef; use crate::plan::physical::SegmentScan; @@ -40,6 +44,7 @@ use crate::plan::physical::Zoned; use crate::plan::physical::optimize; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; +use crate::plan::physical::pipeline::ValueCache; use crate::plan::physical::plan_row_idx_expression; use crate::scan::filter::FilterExpr; use crate::segments::SegmentId; @@ -47,11 +52,20 @@ use crate::segments::SegmentId; /// Plans and adaptive conjunct ordering shared by every chunk of a scan. pub struct FilterPlans { expression: FilterExpr, - predicates: Vec, + predicates: Vec, /// A metadata-only pruning plan per conjunct, where the conjunct has one. pruning: Vec>, } +struct PredicatePlan { + plan: PlanRef, + /// An infallible predicate evaluated in the original row domain, without gathering and + /// scattering its boolean result through the current selection. + whole: bool, + /// Dense infallible predicates can intersect masks without gathering their boolean output. + dense_whole: bool, +} + impl FilterPlans { pub(crate) fn new( expression: BoundExpression, @@ -67,7 +81,20 @@ impl FilterPlans { plan_row_idx_expression(expr.clone(), root.clone())?, session, )?; - Ok(FilterPlan::new(plan).into_plan()) + // Sparse filtering through a cross-column OR filters each encoded input. + // Evaluate safe disjunctions first, as the V1 struct filter does. + let whole_predicate = plan.as_opt::().is_some_and(|eval| { + eval.expression() + .as_scalar() + .and_then(|function| function.as_opt::()) + .is_some_and(|operator| *operator == Operator::Or) + }) && infallible(expr); + let dense_whole = infallible(expr); + Ok(PredicatePlan { + plan: FilterPlan::new(plan).into_plan(), + whole: whole_predicate, + dense_whole, + }) }) .collect::>()?; let pruning = expression @@ -101,6 +128,13 @@ impl FilterPlans { } } +fn infallible(expression: &BoundExpression) -> bool { + expression + .as_scalar() + .is_none_or(|function| function.signature().is_infallible()) + && expression.children().iter().all(infallible) +} + /// A pruning plan must not fetch data when its requested statistics are unavailable. fn metadata_only(plan: &PlanRef) -> VortexResult { if plan @@ -130,7 +164,7 @@ pub(super) struct FilterPlanner { selected: Option, filter: Arc, remaining: BitVec, - active: Option<(usize, ScanAdaptor)>, + active: Option<(usize, bool, ScanAdaptor)>, pieces: Vec, announced: bool, /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. @@ -141,13 +175,20 @@ pub(super) struct FilterPlanner { /// them without IO. pub(super) fn segments_read<'a>( session: &VortexSession, + values: &ValueCache, plans: impl Iterator, rows: Range, mask: &Mask, ) -> VortexResult> { let mut segments = Vec::new(); for plan in plans { - segments.extend(Scan::segments(session, plan, rows.clone(), mask)?); + segments.extend(Scan::segments_with_value_cache( + session, + plan, + rows.clone(), + mask, + values, + )?); } segments.sort_unstable(); segments.dedup(); @@ -200,7 +241,7 @@ pub(super) fn piece_mask(session: &VortexSession, piece: Piece) -> VortexResult< impl IoConsumer for FilterPlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { - if let Some((_, adaptor)) = &mut self.active { + if let Some((_, _, adaptor)) = &mut self.active { adaptor.set_io_result(request, result); } } @@ -212,7 +253,7 @@ impl Planner for FilterPlanner { return State::Done; } match &self.active { - Some((_, adaptor)) if adaptor.state() != State::Done => adaptor.state(), + Some((_, _, adaptor)) if adaptor.state() != State::Done => adaptor.state(), _ => State::NeedsCompute, } } @@ -227,12 +268,17 @@ impl Planner for FilterPlanner { let rows = selected.scope.rows.clone(); let predicates = segments_read( session, - self.filter.predicates.iter(), + &self.scan.values, + self.filter + .predicates + .iter() + .map(|predicate| &predicate.plan), rows.clone(), &selected.mask, )?; let projection = segments_read( session, + &self.scan.values, std::iter::once(&self.scan.projection), rows, &selected.mask, @@ -260,7 +306,7 @@ impl Planner for FilterPlanner { return Ok(PlannerOutput::NeedsIO(batch)); } } - if let Some((index, adaptor)) = &mut self.active { + if let Some((index, whole, adaptor)) = &mut self.active { if adaptor.state() != State::Done { return Ok(match adaptor.compute()? { AdaptorStep::Yield => PlannerOutput::Continue, @@ -273,8 +319,13 @@ impl Planner for FilterPlanner { } let input = selected.mask.true_count(); let result = Mask::concat(self.pieces.iter())?; - vortex_ensure_eq!(result.len(), input); - selected.mask = selected.mask.intersect_by_rank(&result); + if *whole { + vortex_ensure_eq!(result.len(), selected.mask.len()); + selected.mask = (&selected.mask).bitand(&result); + } else { + vortex_ensure_eq!(result.len(), input); + selected.mask = selected.mask.intersect_by_rank(&result); + } tracing::debug!( row_start = selected.scope.rows.start, row_end = selected.scope.rows.end, @@ -295,18 +346,28 @@ impl Planner for FilterPlanner { } if let Some(index) = self.filter.expression.next_conjunct(&self.remaining) { self.remaining.set(index, false); + let predicate = &self.filter.predicates[index]; + // Match MaskStage's dense evaluation threshold, but retain the original row + // domain so combining conjuncts needs no boolean gather and rank intersection. + let whole = + predicate.whole || (predicate.dense_whole && selected.mask.density() >= 0.2); let scan = Scan::try_new( self.scan.session.clone(), - self.filter.predicates[index].clone(), + self.filter.predicates[index].plan.clone(), vec![Split { rows: selected.scope.rows.clone(), - mask: selected.mask.clone(), + mask: if whole { + Mask::new_true(selected.mask.len()) + } else { + selected.mask.clone() + }, }], )? .with_row_offset(self.scan.row_offset) .with_value_cache(self.scan.values.clone()); self.active = Some(( index, + whole, ScanAdaptor::new(scan, Arc::clone(&self.scan.locations)), )); return Ok(PlannerOutput::Continue); diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 0185f61d6a9..b4c0c3de15d 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -31,12 +31,13 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; pub use self::chunk::ChunkPlanner; -pub(crate) use self::chunk::FILTERED_ROWS_PER_CHUNK; pub(crate) use self::chunk::MAX_ROWS_PER_CHUNK; pub(crate) use self::chunk::ROW_BLOCK_SIZE; +pub(crate) use self::chunk::filtered_rows_per_chunk; #[cfg(test)] pub(crate) use self::chunk::row_chunks; pub(crate) use self::chunk::rows_per_chunk; +pub(crate) use self::chunk::total_bytes; use self::filter::FilterPlanner; pub use self::filter::FilterPlans; use self::projection::ProjectionPlanner; diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 141b7377c26..f0a8ac7fd0f 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,15 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::VecDeque; use std::sync::Arc; use std::vec::IntoIter; +use rustc_hash::FxHashSet; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; use vortex_io::request::IoRequestId; use vortex_io::request::IoResult; use vortex_mask::Mask; @@ -28,26 +33,56 @@ use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; +use crate::scan::execution::chunk::row_chunks; +use crate::scan::execution::filter::hint; +use crate::scan::execution::filter::segments_read; -const MAX_OUTPUT_ROWS: usize = 65_536; +const MAX_OUTPUT_ROWS: usize = 262_144; +const MAX_FILTERED_OUTPUT_ROWS: usize = 65_536; +/// Pieces smaller than this are grouped with their neighbours into one batch. +const MIN_OUTPUT_ROWS: usize = 4_096; /// Keeps sparse selections together and subdivides dense selections to avoid large copies when /// converting chunked projection results to Arrow. pub(super) struct ProjectionPlanner { scan: Arc, selected: IntoIter, + announced: bool, } impl ProjectionPlanner { /// Creates a planner for `selected`. pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { - let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { + let output_rows = if scan.filter.is_some() { + MAX_FILTERED_OUTPUT_ROWS + } else { + MAX_OUTPUT_ROWS + }; + let selected = if scan.filter.is_none() && scan.chunk_rows > MAX_OUTPUT_ROWS { + row_chunks(selected.scope.rows.clone(), MAX_OUTPUT_ROWS) + .into_iter() + .filter_map(|rows| { + let start = usize::try_from(rows.start - selected.scope.rows.start) + .vortex_expect("selection mask offset fits usize"); + let end = usize::try_from(rows.end - selected.scope.rows.start) + .vortex_expect("selection mask offset fits usize"); + let mask = selected.mask.slice(start..end); + (!mask.all_false()).then_some(SelectedRows { + scope: WorkScope { + file_ordinal: selected.scope.file_ordinal, + rows, + }, + mask, + }) + }) + .collect() + } else if selected.mask.true_count() <= output_rows { vec![selected] } else { (0..selected.mask.len()) - .step_by(MAX_OUTPUT_ROWS) + .step_by(output_rows) .filter_map(|start| { - let end = (start + MAX_OUTPUT_ROWS).min(selected.mask.len()); + let end = (start + output_rows).min(selected.mask.len()); let mask = selected.mask.slice(start..end); if mask.all_false() { return None; @@ -64,6 +99,8 @@ impl ProjectionPlanner { .collect() }; Self { + announced: scan.filter.is_some() + || scan.projection.row_count() <= MAX_OUTPUT_ROWS as u64, scan, selected: selected.into_iter(), } @@ -84,6 +121,46 @@ impl Planner for ProjectionPlanner { } fn compute(&mut self) -> VortexResult { + if !self.announced { + self.announced = true; + let mut batch = Vec::new(); + let mut seen = FxHashSet::default(); + let mut next_hint = u32::MAX; + let mut segments = Vec::new(); + for selected in self.selected.as_slice() { + segments.extend(segments_read( + &self.scan.session, + &self.scan.values, + std::iter::once(&self.scan.projection), + selected.scope.rows.clone(), + &selected.mask, + )?); + } + segments.sort_unstable(); + segments.dedup(); + // Register every range before starting any read, allowing adjacent ranges to + // coalesce. Prefetch is limited to this work chunk's selected projection. + let bytes: u64 = segments + .iter() + .map(|segment| u64::from(self.scan.locations[**segment as usize].length)) + .sum(); + if bytes <= 8 * 1024 * 1024 { + for intent in [IoIntent::Announce, IoIntent::Prefetch] { + seen.clear(); + hint( + &self.scan.locations, + &mut next_hint, + intent, + &segments, + &mut seen, + &mut batch, + ); + } + } + if !batch.is_empty() { + return Ok(PlannerOutput::NeedsIO(batch)); + } + } let Some(selected) = self.selected.next() else { return Ok(PlannerOutput::Done); }; @@ -97,15 +174,19 @@ impl Planner for ProjectionPlanner { } } -/// Runs the projection plan over a split's selected rows and returns them as one batch, in row -/// order. A selection that projects to no rows finishes without a batch. +/// Runs the projection plan over a split's selected rows and returns them in row order, one +/// batch per piece the scan produced. Pieces under [`MIN_OUTPUT_ROWS`] are grouped with their +/// neighbours. Emitting pieces separately avoids concatenating them on conversion to Arrow, +/// which rebuilds nested dictionaries. A selection that projects to no rows finishes without a +/// batch. struct ProjectionMorsel { scan: Arc, dtype: DType, selected: SelectedRows, adaptor: Option, pieces: Vec, - done: bool, + /// Batches ready to emit once the scan has finished, in row order. + batches: Option>, } impl ProjectionMorsel { @@ -117,24 +198,39 @@ impl ProjectionMorsel { selected, adaptor: None, pieces: Vec::new(), - done: false, + batches: None, } } - fn finish(&mut self) -> VortexResult { - self.done = true; + /// Groups the finished pieces into batches of at least [`MIN_OUTPUT_ROWS`] rows where + /// possible. + fn group_pieces(&mut self) -> VortexResult> { self.pieces.sort_by_key(|piece| piece.rows.start); - let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) - .into_iter() - .map(|piece| piece.array) - .filter(|array| !array.is_empty()) - .collect(); - let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), - 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), - }; - Ok(MorselOutput::Batch(array)) + let mut batches = VecDeque::new(); + let mut group = Vec::new(); + let mut group_rows = 0; + for piece in std::mem::take(&mut self.pieces) { + if piece.array.is_empty() { + continue; + } + group_rows += piece.array.len(); + group.push(piece.array); + if group_rows >= MIN_OUTPUT_ROWS { + batches.push_back(self.concat(std::mem::take(&mut group))?); + group_rows = 0; + } + } + if !group.is_empty() { + batches.push_back(self.concat(group)?); + } + Ok(batches) + } + + fn concat(&self, mut arrays: Vec) -> VortexResult { + if arrays.len() == 1 { + return Ok(arrays.remove(0)); + } + Ok(ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array()) } } @@ -148,8 +244,14 @@ impl IoConsumer for ProjectionMorsel { impl Morsel for ProjectionMorsel { fn state(&self) -> State { + if let Some(batches) = &self.batches { + return if batches.is_empty() { + State::Done + } else { + State::NeedsCompute + }; + } match &self.adaptor { - _ if self.done => State::Done, None => State::NeedsCompute, Some(adaptor) => match adaptor.state() { State::Done => State::NeedsCompute, @@ -159,8 +261,11 @@ impl Morsel for ProjectionMorsel { } fn compute(&mut self) -> VortexResult { - if self.done { - vortex_bail!("ProjectionMorsel: compute called after Done"); + if let Some(batches) = &mut self.batches { + return match batches.pop_front() { + Some(batch) => Ok(MorselOutput::Batch(batch)), + None => vortex_bail!("ProjectionMorsel: compute called after Done"), + }; } let Some(adaptor) = self.adaptor.as_mut() else { let scan = Scan::try_new( @@ -181,7 +286,8 @@ impl Morsel for ProjectionMorsel { return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { - return self.finish(); + self.batches = Some(self.group_pieces()?); + return Ok(MorselOutput::Continue); } Ok(match adaptor.compute()? { AdaptorStep::Yield => MorselOutput::Continue, diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 325f5f897d7..ee02deae074 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -111,6 +111,7 @@ impl Planner for PrunePlanner { // result and lists no segment, so only unpruned plans prefetch their zone maps. let segments = segments_read( &self.scan.session, + &self.scan.values, self.pruning.as_slice().iter(), self.rows.clone(), &self.mask, diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index c621bf19289..c8de5a5a64c 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -3,18 +3,23 @@ //! State derived from a [`ScanFile`](super::ScanFile) that repeated V2 scans can share. +use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use std::sync::Weak; use parking_lot::Mutex; use rustc_hash::FxHashMap; use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; +use vortex_error::VortexResult; use super::io::SegmentRanges; +use super::retain::RetainedSegments; +use super::retain::RetainingSegmentSource; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; use crate::scan::execution::FilterPlans; +use crate::segments::SegmentSource; /// Reusable state for repeated V2 scans of one file, the V2 counterpart of a cached /// [`LayoutReader`](crate::LayoutReader) tree. @@ -25,7 +30,8 @@ use crate::scan::execution::FilterPlans; /// scans. The lowered plan keeps the children it has initialized, each distinct projection keeps /// its optimized plan, each distinct filter keeps its conjunct and pruning plans along with the /// zones those plans have already excluded, and dictionary values stay keyed by the cached plan -/// that produced them. +/// that produced them. Recently read segments are retained up to a byte budget, so the chunks of +/// the file's scans that share a stored segment read it once. /// /// A cache describes one file's layout, segment locations, and segment source. Use a new cache /// when any of them changes. Cached plans and values live as long as the cache, one set per @@ -38,15 +44,52 @@ struct Inner { root: Mutex>, /// Optimized plans of non-identity projections. Values are cached by the plan that produced /// them, so a projection must reuse its plan for its values to be found again. - projections: Mutex>, + projections: Mutex>, ranges: Mutex>, /// Plans of filters, keyed by the exact filter expression. A pruning plan remembers the zones it /// excluded, so sharing the plans prunes a file's zones once for every scan of that filter. - filters: Mutex>>, + filters: Mutex>>, + /// Serialized bytes the whole-file projection reads, by projection. + projection_bytes: Mutex>, values: ValueCache, + /// Recently read segments, shared by the chunks of every scan of the file. + retained: Arc, +} + +/// Independently bound expressions from different engine partitions share a cache entry. +/// Dtypes still participate in equality, but hashing them would eagerly decode lazy schemas. +#[derive(Clone, PartialEq, Eq)] +struct ScanExpression(BoundExpression); + +impl Hash for ScanExpression { + fn hash(&self, state: &mut H) { + hash_expression(&self.0, state); + } +} + +fn hash_expression(expression: &BoundExpression, state: &mut H) { + if let Some(function) = expression.as_scalar() { + state.write_u8(1); + function.hash(state); + expression.children().len().hash(state); + for child in expression.children() { + hash_expression(child, state); + } + } else { + state.write_u8(0); + } } impl ScanCache { + /// A cache whose scans retain up to `bytes` of recently read segments, so chunks that share + /// a stored segment read it once. The default retains 64 MiB; zero retains nothing. + pub fn with_retained_bytes(bytes: u64) -> Self { + Self(Arc::new(Inner { + retained: Arc::new(RetainedSegments::new(bytes)), + ..Inner::default() + })) + } + /// The lowered root plan an earlier scan kept. pub(super) fn root(&self) -> Option { self.0.root.lock().clone() @@ -62,7 +105,7 @@ impl ScanCache { self.0 .projections .lock() - .get(&ExactBoundExpr(projection.clone())) + .get(&ScanExpression(projection.clone())) .cloned() } @@ -72,7 +115,7 @@ impl ScanCache { self.0 .projections .lock() - .entry(ExactBoundExpr(projection)) + .entry(ScanExpression(projection)) .or_insert(plan) .clone() } @@ -92,7 +135,7 @@ impl ScanCache { self.0 .filters .lock() - .get(&ExactBoundExpr(filter.clone())) + .get(&ScanExpression(filter.clone())) .cloned() } @@ -107,11 +150,35 @@ impl ScanCache { self.0 .filters .lock() - .entry(ExactBoundExpr(filter)) + .entry(ScanExpression(filter)) .or_insert(plans), ) } + /// The serialized bytes a whole-file scan of `projection` reads, computing them with + /// `compute` on first use. + pub(super) fn projection_bytes( + &self, + projection: &BoundExpression, + compute: impl FnOnce() -> VortexResult, + ) -> VortexResult { + let key = ScanExpression(projection.clone()); + if let Some(bytes) = self.0.projection_bytes.lock().get(&key) { + return Ok(*bytes); + } + let bytes = compute()?; + self.0.projection_bytes.lock().insert(key, bytes); + Ok(bytes) + } + + /// Serves `segments` through the segments the file's scans retain. + pub(super) fn retaining(&self, segments: Arc) -> Arc { + Arc::new(RetainingSegmentSource { + inner: segments, + retained: Arc::clone(&self.0.retained), + }) + } + pub(super) fn values(&self) -> ValueCache { self.0.values.clone() } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index f0ad78034ae..1342d30ad53 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -13,16 +13,18 @@ //! Import this module's [`ScanBuilder`] to prepare a V2 scan. It accepts a [`ScanFile`] directly //! and supports projection, filters, row ranges, and selections. Unsupported layouts return an //! error when the plan is prepared; limits are not part of this builder's API. Work is cut into -//! chunks of up to 65,536 rows, or 1,048,576 rows for filtered scans, independently of the stored -//! layout's chunk boundaries. Full-file projections use smaller 8,192-row multiples when -//! serialized bytes per row are high. +//! chunks of up to 262,144 rows, or up to 1,048,576 rows for filtered scans, independently of the +//! stored layout's chunk boundaries. Filtered scans use smaller chunks when the projection reads +//! many bytes per row or the row range would give fewer chunks than threads. Full-file +//! projections use smaller 8,192-row multiples when serialized bytes per row are high. //! //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and pruned -//! zones, segment index, and dictionary values with every other scan of that file, like a cached -//! layout reader does for V1 scans. +//! zones, segment index, dictionary values, and recently read segments with every other scan of +//! that file, like a cached layout reader does for V1 scans. mod cache; pub(crate) mod io; +mod retain; #[cfg(test)] mod tests; @@ -33,18 +35,20 @@ use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_mask::Mask; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; +use vortex_utils::parallelism::get_available_parallelism; pub use self::cache::ScanCache; pub use self::cache::WeakScanCache; use crate::LayoutRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; +use crate::plan::physical::pipeline::Scan; use crate::plan::physical::plan_row_idx_expression; pub use crate::scan::execution::ChunkPlanner; -use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; pub use crate::scan::execution::FilterPlans; use crate::scan::execution::MAX_ROWS_PER_CHUNK; pub use crate::scan::execution::PreparedScan; @@ -52,13 +56,18 @@ pub use crate::scan::execution::PruneNext; pub use crate::scan::execution::PrunePlanner; use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; +use crate::scan::execution::filtered_rows_per_chunk; use crate::scan::execution::rows_per_chunk; +use crate::scan::execution::total_bytes; use crate::scan::v2::io::segment_ranges; use crate::segments::SegmentSource; +const MIN_PARTITION_ROWS: usize = 8 * ROW_BLOCK_SIZE; + /// The rows a byte range of the file stands for when an engine splits the file by bytes before /// opening it. Consecutive byte ranges give consecutive row ranges that tile the file, with -/// boundaries on the grid's 8,192-row blocks so the executor's chunks stay block aligned. +/// boundaries on 65,536-row blocks. Keeping small files in fewer partitions avoids repeatedly +/// slicing the same encoded buffers into sparse batches. pub fn partition_row_range(bytes: Range, file_size: u64, row_count: u64) -> Range { let row = |byte: u64| -> u64 { if byte >= file_size { @@ -66,7 +75,7 @@ pub fn partition_row_range(bytes: Range, file_size: u64, row_count: u64) -> } let row = u128::from(byte) * u128::from(row_count) / u128::from(file_size); let row = u64::try_from(row).vortex_expect("a row of the file fits u64"); - row / ROW_BLOCK_SIZE as u64 * ROW_BLOCK_SIZE as u64 + row / MIN_PARTITION_ROWS as u64 * MIN_PARTITION_ROWS as u64 }; let start = row(bytes.start); start..row(bytes.end).max(start) @@ -97,6 +106,7 @@ pub struct ScanBuilder { row_range: Option>, selection: Selection, row_offset: u64, + prefetch_batches: usize, } impl ScanBuilder { @@ -110,6 +120,7 @@ impl ScanBuilder { row_range: None, selection: Selection::default(), row_offset: 0, + prefetch_batches: 1, } } @@ -158,6 +169,13 @@ impl ScanBuilder { self } + /// Groups up to four output batches into one unfiltered projection read window. + /// Drivers running one chunk at a time can widen IO discovery without changing output sizes. + pub fn with_prefetch_batches(mut self, batches: usize) -> Self { + self.prefetch_batches = batches.clamp(1, 4); + self + } + /// Lowers the file's layout and plans the filter and the projection, from the file's cache /// where an earlier scan planned them already. Touches no segment. Lowering and optimizing a /// wide file is real CPU, so a driver runs this where it runs compute. @@ -192,13 +210,39 @@ impl ScanBuilder { plan_row_idx_expression(self.projection.clone(), root)?, &self.session, )?; - cache.insert_projection(self.projection, plan) + cache.insert_projection(self.projection.clone(), plan) } }; + let rows = self.row_range.unwrap_or(0..row_count); let chunk_rows = if filter.is_some() { - FILTERED_ROWS_PER_CHUNK + let projected_bytes = if identity { + total_bytes(&file.locations) + } else { + cache.projection_bytes(&self.projection, || { + let segments = Scan::segments( + &self.session, + &projection, + 0..row_count, + &Mask::new_true(usize::try_from(row_count)?), + )?; + Ok(segments + .iter() + .map(|segment| u64::from(file.locations[**segment as usize].length)) + .sum()) + })? + }; + filtered_rows_per_chunk( + rows.end - rows.start, + row_count, + projected_bytes, + get_available_parallelism().unwrap_or(1), + ) } else if identity { rows_per_chunk(row_count, &file.locations) + } else if self.prefetch_batches > 1 && row_count > MAX_ROWS_PER_CHUNK as u64 { + // The engine already supplies parallelism. Group read discovery across several + // output morsels without creating fewer parallel tasks. + self.prefetch_batches * MAX_ROWS_PER_CHUNK } else { MAX_ROWS_PER_CHUNK }; @@ -209,9 +253,9 @@ impl ScanBuilder { let values = cache.values(); Ok(PreparedScan { session: self.session, - rows: self.row_range.unwrap_or(0..row_count), + rows, selection: self.selection, - segments: file.segments, + segments: cache.retaining(file.segments), ranges, cache, locations: file.locations, diff --git a/vortex-layout/src/scan/v2/retain.rs b/vortex-layout/src/scan/v2/retain.rs new file mode 100644 index 00000000000..5f50c4b80df --- /dev/null +++ b/vortex-layout/src/scan/v2/retain.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Keeps recently read segments so the chunks of a scan that share a segment read it once. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use futures::FutureExt; +use futures::TryFutureExt; +use parking_lot::Mutex; +use rustc_hash::FxHashMap; +use vortex_array::buffer::BufferHandle; + +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// The host bytes of recently read segments, evicting the least recently used past a budget. +/// +/// A V2 scan cuts its rows into chunks that each read through a session of their own, and a +/// stored segment that spans a chunk boundary is read by every chunk it overlaps. The file's segment +/// source shares a read only while it is in flight, so a chunk that starts after its neighbour +/// finished would read the segment again. Neighbouring chunks run close together, so a small +/// budget is enough to read each such segment once. +pub(super) struct RetainedSegments { + budget: u64, + state: Mutex, +} + +/// The bytes a file's scans retain by default. +const DEFAULT_RETAINED_BYTES: u64 = 64 << 20; + +impl Default for RetainedSegments { + fn default() -> Self { + Self::new(DEFAULT_RETAINED_BYTES) + } +} + +#[derive(Default)] +struct State { + bytes: u64, + tick: u64, + /// Each retained segment and the tick of its last use. + segments: FxHashMap, + /// Retained segments by the tick of their last use, oldest first. + by_use: BTreeMap, +} + +impl RetainedSegments { + pub(super) fn new(budget: u64) -> Self { + Self { + budget, + state: Mutex::default(), + } + } + + fn get(&self, id: SegmentId) -> Option { + let mut state = self.state.lock(); + state.tick += 1; + let tick = state.tick; + let (buffer, used) = state.segments.get_mut(&id)?; + let buffer = buffer.clone(); + let previous = std::mem::replace(used, tick); + state.by_use.remove(&previous); + state.by_use.insert(tick, id); + Some(buffer) + } + + fn insert(&self, id: SegmentId, buffer: &BufferHandle) { + let len = buffer.len() as u64; + if !buffer.is_on_host() || len > self.budget / 4 { + return; + } + let mut state = self.state.lock(); + if state.segments.contains_key(&id) { + return; + } + state.tick += 1; + let tick = state.tick; + state.segments.insert(id, (buffer.clone(), tick)); + state.by_use.insert(tick, id); + state.bytes += len; + while state.bytes > self.budget { + let Some((_, oldest)) = state.by_use.pop_first() else { + break; + }; + if let Some((evicted, _)) = state.segments.remove(&oldest) { + state.bytes -= evicted.len() as u64; + } + } + } +} + +/// Serves segments from [`RetainedSegments`] before reading them from `inner`, and retains +/// what it reads. +pub(super) struct RetainingSegmentSource { + pub(super) inner: Arc, + pub(super) retained: Arc, +} + +impl SegmentSource for RetainingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + if let Some(buffer) = self.retained.get(id) { + return futures::future::ready(Ok(buffer)).boxed(); + } + let retained = Arc::clone(&self.retained); + self.inner + .request(id) + .inspect_ok(move |buffer| retained.insert(id, buffer)) + .boxed() + } +} + +#[cfg(test)] +mod tests { + use vortex_array::buffer::BufferHandle; + use vortex_buffer::ByteBuffer; + + use super::RetainedSegments; + use crate::segments::SegmentId; + + fn buffer(len: usize) -> BufferHandle { + BufferHandle::new_host(ByteBuffer::from(vec![0u8; len])) + } + + #[test] + fn evicts_least_recently_used_past_budget() { + let retained = RetainedSegments::new(400); + retained.insert(SegmentId::from(0), &buffer(100)); + retained.insert(SegmentId::from(1), &buffer(100)); + retained.insert(SegmentId::from(2), &buffer(100)); + // Using segment 0 makes segment 1 the least recently used. + assert!(retained.get(SegmentId::from(0)).is_some()); + retained.insert(SegmentId::from(3), &buffer(100)); + retained.insert(SegmentId::from(4), &buffer(100)); + + assert!(retained.get(SegmentId::from(1)).is_none()); + for id in [0, 2, 3, 4] { + assert!(retained.get(SegmentId::from(id)).is_some(), "segment {id}"); + } + } + + #[test] + fn skips_segments_over_a_quarter_of_the_budget() { + let retained = RetainedSegments::new(400); + retained.insert(SegmentId::from(0), &buffer(101)); + assert!(retained.get(SegmentId::from(0)).is_none()); + } +} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index eff68dfe996..6edb643ab4e 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -33,6 +33,7 @@ use vortex_array::expr::gt; use vortex_array::expr::is_null; use vortex_array::expr::lit; use vortex_array::expr::lt; +use vortex_array::expr::or; use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::validity::Validity; @@ -60,6 +61,7 @@ use super::ROW_BLOCK_SIZE; use super::ScanBuilder; use super::ScanCache; use super::ScanFile; +use super::filtered_rows_per_chunk; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; @@ -344,6 +346,13 @@ async fn large_filtered_chunks_match_v1(#[case] cutoff: i32) -> VortexResult<()> gt(get_item("ints", root()), lit(995i32)), lt(get_item("ints", root()), lit(2050i32)) ))] +#[case::cross_column_or(and( + gt(get_item("ints", root()), lit(995i32)), + or( + lt(get_item("ints", root()), lit(1010i32)), + eq(get_item("strings", root()), lit("row 2001")) + ) +))] #[case::nullable(eq(get_item("strings", root()), lit("row 1001")))] #[case::nulls(is_null(get_item("strings", root())))] #[case::no_matches(gt(get_item("ints", root()), lit(10_000i32)))] @@ -352,7 +361,7 @@ async fn large_filtered_chunks_match_v1(#[case] cutoff: i32) -> VortexResult<()> async fn filter_matches_v1( #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, #[case] filter: Expression, - #[values(false, true)] sparse: bool, + #[values(1, 2, 7)] selection_step: usize, ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, strategy.get().as_ref()).await?; @@ -365,9 +374,9 @@ async fn filter_matches_v1( .with_some_filter(Some(filter.bind(layout.dtype())?)) .with_row_range(900..3100) .with_row_offset(100); - if sparse { + if selection_step > 1 { scan = scan.with_row_indices(StrictSortedBuffer::try_new(Buffer::from_iter( - (0..4000u64).step_by(7), + (0..4000u64).step_by(selection_step), ))?); } Ok(scan) @@ -480,10 +489,10 @@ fn fixed_row_chunks(#[case] rows: Range, #[case] expected: Vec>) } #[rstest] -#[case::empty(0, 0, 65_536)] -#[case::single_chunk(1000, 100_000_000, 65_536)] -#[case::no_bytes(100_000, 0, 65_536)] -#[case::small_rows(1_000_000, 64_000_000, 65_536)] +#[case::empty(0, 0, 262_144)] +#[case::single_chunk(1000, 100_000_000, 262_144)] +#[case::no_bytes(100_000, 0, 262_144)] +#[case::small_rows(1_000_000, 64_000_000, 131_072)] #[case::large_rows(1_000_000, 384_000_000, 16_384)] #[case::minimum(1_000_000, 2_000_000_000, 8192)] fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: usize) { @@ -495,6 +504,24 @@ fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: us assert_eq!(rows_per_chunk(rows, &locations), expected); } +#[rstest] +#[case::minimum_work(1_000_000, 1_000_000, 0, 32, 262_144)] +#[case::maximum_work(64_000_000, 64_000_000, 0, 32, 1_048_576)] +#[case::wide_projection(1_000_000, 1_000_000, 384_000_000, 32, 16_384)] +#[case::zero_parallelism(1_048_576, 1_048_576, 0, 0, 1_048_576)] +fn filtered_chunk_budget( + #[case] rows: u64, + #[case] row_count: u64, + #[case] projected_bytes: u64, + #[case] parallelism: usize, + #[case] expected: usize, +) { + assert_eq!( + filtered_rows_per_chunk(rows, row_count, projected_bytes, parallelism), + expected + ); +} + #[test] fn smaller_chunks_keep_file_alignment() { assert_eq!( @@ -518,7 +545,7 @@ impl SegmentSource for CountingSegments { #[tokio::test] async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let session = new_session().with_tokio(); - const STORED_CHUNK_ROWS: usize = 70_000; + const STORED_CHUNK_ROWS: usize = 150_000; const STORED_CHUNKS: usize = 3; let segments = Arc::new(TestSegments::default()); let (mut sequence, eof) = SequenceId::root().split(); @@ -550,7 +577,10 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); + let mut file = scan_file(&segments, &layout)?; + // Without retained segments, so a chunk reads every codes chunk it overlaps. + file.cache = Some(ScanCache::with_retained_bytes(0)); + let scan = Arc::new(ScanBuilder::new(session.clone(), file).prepare()?); let chunks = row_chunks(scan.rows().clone(), scan.chunk_rows); assert_eq!( chunks.len(), @@ -583,6 +613,45 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { Ok(()) } +/// Work chunks that each overlap half of a stored segment read it once while the file's cache +/// retains segments, and once per chunk when it retains nothing. +#[rstest] +#[case::retained(None, 1)] +#[case::not_retained(Some(0), 2)] +#[tokio::test] +async fn chunks_sharing_a_segment_read_it_once( + #[case] retained_bytes: Option, + #[case] reads_per_segment: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, chunked_flat().as_ref()).await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments = Arc::clone(&counted) as Arc; + let mut file = scan_file(&segments, &layout)?; + file.cache = + Some(retained_bytes.map_or_else(ScanCache::default, ScanCache::with_retained_bytes)); + let projection = get_item("ints", root()).bind(layout.dtype())?; + let mut scan = ScanBuilder::new(session, file) + .with_projection(projection) + .prepare()?; + let chunk_rows = CHUNK_ROWS / 2; + scan.chunk_rows = chunk_rows; + let scan = Arc::new(scan); + + // One chunk at a time, so no read is still in flight when the next chunk starts. + for rows in row_chunks(scan.rows().clone(), chunk_rows) { + run_rows(&scan, rows).await?; + } + assert_eq!( + counted.reads.load(Ordering::Relaxed), + CHUNKS * reads_per_segment + ); + Ok(()) +} + /// V2 rejects list layouts: their element ranges are only known once their offsets are read. #[tokio::test] async fn list_layouts_return_an_error() -> VortexResult<()> { @@ -704,7 +773,8 @@ async fn scan_cache_reuses_values_between_scans( }); let segments = Arc::clone(&counted) as Arc; let mut file = scan_file(&segments, &layout)?; - file.cache = cached.then(ScanCache::default); + // Without retained segments, so only the values come from the cache. + file.cache = cached.then(|| ScanCache::with_retained_bytes(0)); let mut reads = Vec::new(); for _ in 0..2 { @@ -724,6 +794,39 @@ async fn scan_cache_reuses_values_between_scans( Ok(()) } +#[tokio::test] +async fn dictionary_cache_lives_with_its_prepared_file_scans() -> VortexResult<()> { + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let values = chunked_ints(&session, &segments).await?; + let codes = chunked_ints(&session, &segments).await?; + let layout = DictLayout::new(values, codes).into_layout(); + let segments = segments as Arc; + let mut file = scan_file(&segments, &layout)?; + let cache = ScanCache::with_retained_bytes(0); + let weak = cache.downgrade(); + file.cache = Some(cache); + let first = Arc::new(ScanBuilder::new(session.clone(), file.clone()).prepare()?); + let second = ScanBuilder::new(session.clone(), file).prepare()?; + + let mut other_file = scan_file(&segments, &layout)?; + let other_cache = ScanCache::with_retained_bytes(0); + let other_weak = other_cache.downgrade(); + other_file.cache = Some(other_cache); + let other = ScanBuilder::new(session, other_file).prepare()?; + + // Populate dictionary values before releasing the plans owned by this file's drivers. + assert!(!run_rows(&first, first.rows().clone()).await?.is_empty()); + drop(first); + assert!(weak.upgrade().is_some()); + drop(second); + assert!(weak.upgrade().is_none()); + assert!(other_weak.upgrade().is_some()); + drop(other); + assert!(other_weak.upgrade().is_none()); + Ok(()) +} + #[tokio::test] async fn selected_split_scope_is_bounded_before_projection() -> VortexResult<()> { let session = new_session().with_tokio(); @@ -762,7 +865,8 @@ async fn repeated_projections_reuse_their_values() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let mut file = scan_file(&segments, &layout)?; - file.cache = Some(ScanCache::default()); + // Without retained segments, so only the values come from the cache. + file.cache = Some(ScanCache::with_retained_bytes(0)); let projection = gt(root(), lit(15i32)).bind(layout.dtype())?; let mut reads = Vec::new(); @@ -797,7 +901,7 @@ fn partition_rows_tile_the_file_on_row_blocks() { assert_eq!(pair[0].end, pair[1].start); assert_eq!(pair[0].end % ROW_BLOCK_SIZE as u64, 0); } - assert!(rows.iter().all(|rows| !rows.is_empty())); + assert_eq!(rows, vec![0..0, 0..0, 0..100_000]); assert_eq!( super::partition_row_range(1..2, file_size, row_count), 0..0, @@ -847,7 +951,8 @@ async fn shared_cache_prunes_zones_once() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let mut file = scan_file(&segments, &layout)?; - let cache = ScanCache::default(); + // Without retained segments, so only the proved zones come from the cache. + let cache = ScanCache::with_retained_bytes(0); file.cache = Some(cache.clone()); let filter = gt(get_item("ints", root()), lit(2_999i32)).bind(layout.dtype())?; let projection = get_item("ints", root()).bind(layout.dtype())?; @@ -856,6 +961,9 @@ async fn shared_cache_prunes_zones_once() -> VortexResult<()> { let mut reads = Vec::new(); for _ in 0..2 { let before = counted.reads.load(Ordering::Relaxed); + // Engine partitions bind separate trees for the same logical expressions. + let filter = gt(get_item("ints", root()), lit(2_999i32)).bind(layout.dtype())?; + let projection = get_item("ints", root()).bind(layout.dtype())?; let scan = ScanBuilder::new(session.clone(), file.clone()) .with_projection(projection.clone()) .with_some_filter(Some(filter.clone())); @@ -867,7 +975,37 @@ async fn shared_cache_prunes_zones_once() -> VortexResult<()> { .filter_plans(&filter) .expect("the first scan kept its filter plans"); assert_eq!(plans.pruning().count(), 1); + assert!(cache.projection(&projection).is_some()); // Each scan reads the one chunk pruning keeps; only the first reads the zone map. assert_eq!(reads, vec![2, 1]); Ok(()) } + +#[rstest] +#[case(0)] +#[case(1)] +#[case(65_535)] +#[case(65_536)] +#[case(102_000)] +#[case(28_800_991)] +fn partition_rows_preserve_coverage(#[case] row_count: u64) { + let rows: Vec<_> = (0..8) + .map(|partition| { + super::partition_row_range(partition * 125..(partition + 1) * 125, 1_000, row_count) + }) + .collect(); + assert_eq!(rows[0].start, 0); + assert_eq!(rows[7].end, row_count); + for pair in rows.windows(2) { + assert_eq!(pair[0].end, pair[1].start); + } + assert_eq!( + rows.iter() + .map(|range| range.end - range.start) + .sum::(), + row_count + ); + if row_count == 102_000 { + assert_eq!(rows.iter().filter(|range| !range.is_empty()).count(), 2); + } +}