diff --git a/benchmarks/datafusion-bench/src/lib.rs b/benchmarks/datafusion-bench/src/lib.rs index 44299c87bac..5f0b1a6800e 100644 --- a/benchmarks/datafusion-bench/src/lib.rs +++ b/benchmarks/datafusion-bench/src/lib.rs @@ -98,13 +98,15 @@ pub fn make_object_store( } } -pub fn format_to_df_format(format: Format) -> anyhow::Result> { +pub fn format_to_df_format(format: Format, scan_v2: bool) -> anyhow::Result> { + let mut options = vortex_table_options(); + options.scan_v2 = scan_v2; Ok(match format { Format::Csv => Arc::new(CsvFormat::default()) as _, Format::Parquet => Arc::new(ParquetFormat::new()), - Format::OnDiskVortex | Format::VortexCompact | Format::VortexSpatialNative => Arc::new( - VortexFormat::new_with_options(SESSION.clone(), vortex_table_options()), - ), + Format::OnDiskVortex | Format::VortexCompact | Format::VortexSpatialNative => { + Arc::new(VortexFormat::new_with_options(SESSION.clone(), options)) + } Format::ArrowIpc | Format::OnDiskDuckDB | Format::Lance | Format::LitData => { anyhow::bail!("Format {format} cannot be turned into a DataFusion `FileFormat`") } diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index f65d7b64973..1a83d374fcb 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -120,6 +120,10 @@ 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). + #[arg(long, action = clap::ArgAction::Set)] + scan_v2: Option, + #[arg(long = "opt", value_delimiter = ',', value_parser = value_parser!(Opt))] options: Vec, } @@ -132,6 +136,10 @@ async fn main() -> anyhow::Result<()> { vortex_bench::configure_allocator(); 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") + }); set_join_set_tracer(get_static_tracer())?; setup_logging_and_tracing(args.verbose, args.tracing)?; @@ -191,7 +199,7 @@ async fn main() -> anyhow::Result<()> { session.sql(&sql).await?.collect().await?; } datafusion_bench::make_object_store(&session, benchmark.data_url())?; - register_benchmark_tables(&session, benchmark, format).await?; + register_benchmark_tables(&session, benchmark, format, scan_v2).await?; Ok((session, format)) } }, @@ -255,12 +263,14 @@ async fn register_benchmark_tables( session: &SessionContext, benchmark: &B, format: Format, + scan_v2: bool, ) -> anyhow::Result<()> { - if use_scan_api() && matches!(format, Format::OnDiskVortex | Format::VortexCompact) { + if !scan_v2 && use_scan_api() && matches!(format, Format::OnDiskVortex | Format::VortexCompact) + { register_v2_tables(session, benchmark, format).await } else { let benchmark_base = benchmark.data_url().join(&format!("{}/", format.name()))?; - let file_format = format_to_df_format(format)?; + let file_format = format_to_df_format(format, scan_v2)?; for table in benchmark.table_specs().iter() { let pattern = benchmark.pattern(table.name, format); diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 50f2bbf8f84..7e69693400b 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -197,6 +197,13 @@ extensions_options! { /// [`ParquetFormatFactory::new_with_options`]: https://docs.rs/datafusion/latest/datafusion/datasource/file_format/parquet/struct.ParquetFormatFactory.html#method.new_with_options /// [`JsonFormatFactory::new_with_options`]: https://docs.rs/datafusion/latest/datafusion/datasource/file_format/json/struct.JsonFormatFactory.html#method.new_with_options pub struct VortexTableOptions { + /// Use the planner-driven V2 scan executor for supported layouts. + /// + /// Disabled by default. Select V2 before creating a table with + /// `SET vortex.scan_v2 = true` (with the `vortex` options extension registered), + /// or for one table with `OPTIONS(scan_v2 'true')`. In Rust, set + /// `VortexTableOptions { scan_v2: true, ..Default::default() }` on the factory. + pub scan_v2: bool, default = false /// The number of bytes to read when parsing a file footer. /// /// Values smaller than `MAX_POSTSCRIPT_SIZE + EOF_SIZE` will be clamped to that minimum @@ -1033,17 +1040,18 @@ mod tests { Ok(()) } + #[rstest] #[tokio::test] - async fn configure_format_source() -> anyhow::Result<()> { + async fn configure_format_source(#[values(false, true)] scan_v2: bool) -> anyhow::Result<()> { let ctx = TestSessionContext::default(); ctx.session .sql( - "CREATE EXTERNAL TABLE my_tbl \ + &format!("CREATE EXTERNAL TABLE my_tbl \ (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \ STORED AS vortex \ LOCATION 'table/' \ - OPTIONS( footer_initial_read_size_bytes '12345', predicate_pushdown 'false', scan_concurrency '3' );", + OPTIONS( footer_initial_read_size_bytes '12345', predicate_pushdown 'false', scan_concurrency '3', scan_v2 '{scan_v2}' );"), ) .await? .collect() @@ -1061,6 +1069,35 @@ mod tests { assert_eq!(format.options().footer_initial_read_size_bytes, 12345); assert!(!format.options().predicate_pushdown); assert_eq!(format.options().scan_concurrency, Some(3)); + assert_eq!(format.options().scan_v2, scan_v2); + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn session_scan_version(#[values(false, true)] scan_v2: bool) -> anyhow::Result<()> { + let factory = Arc::new(VortexFormatFactory::new()); + let ctx = TestSessionContext::new_with_factory(Arc::clone(&factory)); + ctx.session + .state_ref() + .write() + .config_mut() + .options_mut() + .extensions + .insert(VortexTableOptions::default()); + + for enabled in [scan_v2, !scan_v2] { + ctx.session + .sql(&format!("SET vortex.scan_v2 = {enabled}")) + .await? + .collect() + .await?; + let format = factory.create(&ctx.session.state(), &Default::default())?; + let format = format + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("expected VortexFormat"))?; + assert_eq!(format.options().scan_v2, enabled); + } Ok(()) } @@ -1076,6 +1113,7 @@ mod tests { #[test] fn format_plumbs_source_options() -> anyhow::Result<()> { let opts = VortexTableOptions { + scan_v2: true, projection_pushdown: true, predicate_pushdown: false, scan_concurrency: Some(3), @@ -1095,6 +1133,7 @@ mod tests { ); assert_eq!(source.options().predicate_pushdown, opts.predicate_pushdown); assert_eq!(source.options().scan_concurrency, opts.scan_concurrency); + assert_eq!(source.options().scan_v2, opts.scan_v2); Ok(()) } diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index a2fd75a86f3..71dfb22ffc1 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -34,9 +34,12 @@ use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_adapter::replace_columns_with_literals; use datafusion_physical_plan::filter::batch_filter; +use datafusion_physical_plan::limit::LimitStream; +use datafusion_physical_plan::metrics::BaselineMetrics; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::metrics::MetricBuilder; use datafusion_physical_plan::metrics::MetricCategory; +use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_pruning::FilePruner; use futures::FutureExt; use futures::StreamExt; @@ -52,6 +55,9 @@ use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; +use vortex::layout::scan::v2::partition_row_range; +use vortex::layout::segments::SegmentSource; use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; @@ -100,7 +106,13 @@ pub(crate) struct VortexOpener { /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache /// a file reader the first time we read a file. pub layout_readers: Arc>>, - /// Shared full-file natural splits keyed by file path. + /// Shared V2 segment sources keyed by file path. + /// + /// 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>>, + /// 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. pub file_schemas: Arc>, @@ -112,6 +124,7 @@ pub(crate) struct VortexOpener { /// Whether to enable expression pushdown into the underlying Vortex scan. pub projection_pushdown: bool, pub scan_concurrency: Option, + pub scan_v2: bool, } impl FileOpener for VortexOpener { @@ -144,10 +157,12 @@ impl FileOpener for VortexOpener { let unified_file_schema = Arc::clone(self.table_schema.file_schema()); let limit = self.limit; let layout_readers = Arc::clone(&self.layout_readers); + let segment_sources = Arc::clone(&self.segment_sources); let natural_splits = Arc::clone(&self.natural_splits); let file_schemas = Arc::clone(&self.file_schemas); let has_output_ordering = self.has_output_ordering; let scan_concurrency = self.scan_concurrency; + let scan_v2 = self.scan_v2; let expr_converter = Arc::clone(&self.expression_converter); let projection_pushdown = self.projection_pushdown; @@ -327,7 +342,11 @@ impl FileOpener for VortexOpener { leftover_projection, } = if let Some(residual) = &residual_filter { let mut indices = projection.column_indices(); - indices.extend(collect_columns(residual).iter().map(|column| column.index())); + indices.extend( + collect_columns(residual) + .iter() + .map(|column| column.index()), + ); indices.sort_unstable(); indices.dedup(); let (scan_projection, scan_reference_schema) = @@ -367,136 +386,182 @@ impl FileOpener for VortexOpener { .transpose()?; let projector = leftover_projection.make_projector(&stream_schema)?; - // We share our layout readers with others partitions in the scan, so we can only need to read each layout in each file once. - let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { - Entry::Occupied(mut occupied_entry) => { - if let Some(reader) = occupied_entry.get().upgrade() { - tracing::trace!("reusing layout reader for {}", occupied_entry.key()); - reader - } else { - tracing::trace!("creating layout reader for {}", occupied_entry.key()); + 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 + .filter(|_| scan_v2 && residual_filter.is_none()) + .map(|limit| (limit, Arc::clone(&output_schema))); + let map_batch = move |chunk| { + let mut ctx = session.create_execution_ctx(); + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }; + let 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, + ); + let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file) + .with_projection(scan_projection) + .with_some_filter(native_filter) + .with_ordered(has_output_ordering); + 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()); + } + builder = builder.with_row_range(rows); + } + if let Some(selection) = file + .extensions + .get::() + .and_then(VortexAccessPlan::selection) + { + builder = builder.with_selection(selection.clone()); + } + if let Some(concurrency) = scan_concurrency { + builder = builder.with_concurrency(concurrency); + } + builder.map(map_batch).into_stream() + } else { + // Share V1 layout readers across partitions of the file. + let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { + Entry::Occupied(mut occupied_entry) => { + if let Some(reader) = occupied_entry.get().upgrade() { + tracing::trace!("reusing layout reader for {}", occupied_entry.key()); + reader + } else { + tracing::trace!("creating layout reader for {}", occupied_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create layout reader: {e}" + )) + })?; + occupied_entry.insert(Arc::downgrade(&reader)); + reader + } + } + Entry::Vacant(vacant_entry) => { + tracing::trace!("creating layout reader for {}", vacant_entry.key()); let reader = vxf.layout_reader().map_err(|e| { DataFusionError::Execution(format!( "Failed to create layout reader: {e}" )) })?; - occupied_entry.insert(Arc::downgrade(&reader)); + vacant_entry.insert(Arc::downgrade(&reader)); + reader } - } - Entry::Vacant(vacant_entry) => { - tracing::trace!("creating layout reader for {}", vacant_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!("Failed to create layout reader: {e}")) - })?; - vacant_entry.insert(Arc::downgrade(&reader)); + }; - reader + let mut scan_builder = ScanBuilder::new(vxf.session().clone(), layout_reader) + .with_projection(scan_projection) + .with_some_filter(native_filter); + if let Some(vortex_plan) = file.extensions.get::() { + scan_builder = vortex_plan.apply_to_builder(scan_builder); } - }; - - let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); - - if let Some(vortex_plan) = file.extensions.get::() { - scan_builder = vortex_plan.apply_to_builder(scan_builder); - } - - // Native scans apply limits after filtering. Residual filters still need the - // full scan output, with the limit enforced by DataFusion's file stream. - if let Some(limit) = limit - && residual_filter.is_none() - { - scan_builder = scan_builder.with_limit(limit); - } - - if let Some(concurrency) = scan_concurrency { - scan_builder = scan_builder.with_concurrency(concurrency); - } - - // Set before the byte-range translation below, which computes natural splits for - // the fields the scan's projection and filter reference. - scan_builder = scan_builder - .with_projection(scan_projection) - .with_some_filter(native_filter); - - if let Some(file_range) = &file.range { - let byte_range = Range { - start: u64::try_from(file_range.start) - .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, - end: u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, - }; - if byte_range.start != 0 || byte_range.end != file.object_meta.size { - // Full-file scans already cover every natural split. Only translate the - // byte range back into row boundaries when DataFusion has trimmed the file. - let natural_splits = natural_splits_for_file( + // Residual filters need the full scan output before DataFusion applies the limit. + if let Some(limit) = limit + && residual_filter.is_none() + { + scan_builder = scan_builder.with_limit(limit); + } + // V1 requires positive concurrency; V2 also accepts zero. + if let Some(concurrency) = scan_concurrency.filter(|&concurrency| concurrency > 0) { + scan_builder = scan_builder.with_concurrency(concurrency); + } + if let Some(bytes) = byte_range + && (bytes.start != 0 || bytes.end != file.object_meta.size) + { + let splits = natural_splits_for_file( natural_splits.as_ref(), &file.object_meta.location, &scan_builder, file.object_meta.size, )?; - - let Some(row_range) = - split_aligned_row_range(byte_range, natural_splits.as_ref()) - else { + let Some(rows) = split_aligned_row_range(bytes, splits.as_ref()) else { return Ok(stream::empty().boxed()); }; - scan_builder = scan_builder - .with_row_range(row_range) - // Hand the shared full-file boundaries back to the scan so prepare() - // skips its own layout walk. - .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + .with_row_range(rows) + .with_natural_splits(Arc::clone(&splits.row_boundaries)); } + scan_builder + .with_metrics_registry(metrics_registry) + .with_ordered(has_output_ordering) + .map(map_batch) + .into_stream() + .map(StreamExt::boxed) } - - let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); - let residual_path = file.path().clone(); - let stream = scan_builder - .with_metrics_registry(metrics_registry) - .with_ordered(has_output_ordering) - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) - .into_stream() - .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {}", - file.object_meta.location - )))) - }) - .map(move |batch| -> DFResult> { - let mut batch = batch?; - if let Some(residual) = &residual_filter { - batch = batch_filter(&batch, residual).map_err(|e| { - exec_datafusion_err!( - "Failed to evaluate residual filter {residual} in {residual_path}: {e}" - ) - })?; - // Do not evaluate the projection on batches the residual filter emptied. - if batch.num_rows() == 0 { - return Ok(None); - } + .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? + .map_err(move |e: VortexError| { + DataFusionError::External(Box::new(e.with_context(format!( + "Failed to read Vortex file: {}", + file.object_meta.location + )))) + }) + .map(move |batch| -> DFResult> { + let mut batch = batch?; + if let Some(residual) = &residual_filter { + batch = batch_filter(&batch, residual).map_err(|e| { + exec_datafusion_err!( + "Failed to evaluate residual filter {residual} in {residual_path}: {e}" + ) + })?; + // Do not evaluate the projection on batches the residual filter emptied. + if batch.num_rows() == 0 { + return Ok(None); } - let batch = projector.project_batch(&batch)?; - - let (_, columns, row_count) = batch.into_parts(); - Ok(Some(RecordBatch::try_new_with_options( - Arc::clone(&output_schema), - columns, - &RecordBatchOptions::new().with_row_count(Some(row_count)), - )?)) - }) - .filter_map(|batch| ready(batch.transpose())) - .boxed(); + } + let batch = projector.project_batch(&batch)?; + + let (_, columns, row_count) = batch.into_parts(); + Ok(Some(RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + )?)) + }) + .filter_map(|batch| ready(batch.transpose())) + .boxed(); + + let stream = if let Some((limit, schema)) = native_limit { + LimitStream::new( + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + 0, + Some( + usize::try_from(limit) + .map_err(|e| exec_datafusion_err!("Invalid scan limit: {e}"))?, + ), + BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + ) + .boxed() + } else { + stream + }; if let Some(file_pruner) = file_pruner && file_pruner.is_watching() @@ -578,6 +643,30 @@ fn contains_lambda(projection: &ProjectionExprs) -> DFResult { Ok(false) } +/// Returns the segment source other partitions of `path` are reading 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, + fresh: Arc, +) -> Arc { + match segment_sources.entry(path.clone()) { + Entry::Occupied(mut entry) => { + if let Some(shared) = entry.get().upgrade() { + tracing::trace!("reusing segment source for {path}"); + shared + } else { + entry.insert(Arc::downgrade(&fresh)); + fresh + } + } + Entry::Vacant(entry) => { + entry.insert(Arc::downgrade(&fresh)); + fresh + } + } +} + /// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. fn natural_splits_for_file( natural_splits: &DashMap>, @@ -695,7 +784,11 @@ mod tests { use datafusion_common::arrow::compute::concat_batches; use datafusion_common::assert_batches_eq; use datafusion_common::stats::Precision; + use datafusion_datasource::file::FileSource; + use datafusion_datasource::file_groups::FileGroup; + use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_execution::cache::default_cache::DefaultCache; + use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::Operator; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions as df_expr; @@ -718,6 +811,7 @@ mod tests { use super::*; use crate::VortexAccessPlan; + use crate::VortexSource; use crate::convert::exprs::DefaultExpressionConverter; use crate::persistent::reader::DefaultVortexReaderFactory; @@ -860,6 +954,88 @@ mod tests { assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4)); } + /// Byte partitions read every row exactly once, independently of footer cache state. + #[rstest] + #[tokio::test] + async fn test_repartitioned_rows_tile_the_file( + #[values(false, true)] scan_v2: bool, + #[values(false, true)] warm_cache: bool, + #[values(3, 8)] partitions: usize, + ) -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "rows/file.vortex"; + let values: Vec> = (0..30_000).map(Some).collect(); + let batch = record_batch!(("a", Int32, values)).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let file = PartitionedFile::new(file_path.to_string(), data_size); + let table_schema = TableSchema::from(batch.schema()); + + let cache: Arc = Arc::new( + DefaultCache::::new(64 * 1024 * 1024), + ); + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.file_metadata_cache = Some(Arc::clone(&cache)); + opener.scan_v2 = scan_v2; + opener.has_output_ordering = true; + if warm_cache { + opener + .open(file.clone())? + .await? + .try_collect::>() + .await?; + } + + let source = + VortexSource::new(table_schema, SESSION.clone()).with_file_metadata_cache(cache); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(source.clone()), + ) + .with_file_group(FileGroup::new(vec![file])) + .build(); + let partitioned = source + .repartitioned(partitions, 0, None, &config)? + .ok_or_else(|| anyhow::anyhow!("file was not repartitioned"))?; + let files: Vec<_> = partitioned + .file_groups + .iter() + .flat_map(FileGroup::iter) + .collect(); + assert_eq!(files.len(), partitions); + + let mut data = Vec::new(); + for file in files { + let batches = opener + .open(file.clone())? + .await? + .try_collect::>() + .await?; + if scan_v2 { + let bytes = file.range.as_ref().expect("partition has a byte range"); + let rows = partition_row_range( + u64::try_from(bytes.start)?..u64::try_from(bytes.end)?, + data_size, + batch.num_rows() as u64, + ); + assert_eq!( + concat_batches(&batch.schema(), &batches)?, + batch.slice( + usize::try_from(rows.start)?, + usize::try_from(rows.end - rows.start)?, + ), + ); + } + data.extend(batches); + } + assert_eq!(concat_batches(&batch.schema(), &data)?, batch); + if scan_v2 { + assert!(opener.layout_readers.is_empty()); + assert!(opener.natural_splits.is_empty()); + } + Ok(()) + } + async fn write_arrow_to_vortex( object_store: Arc, path: &str, @@ -897,6 +1073,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), + segment_sources: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -904,11 +1081,13 @@ mod tests { file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, + scan_v2: false, } } + #[rstest] #[tokio::test] - async fn test_open() -> anyhow::Result<()> { + async fn test_open(#[values(false, true)] scan_v2: bool) -> anyhow::Result<()> { let object_store = Arc::new(InMemory::new()) as Arc; let file_path = "part=1/file.vortex"; let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); @@ -927,11 +1106,12 @@ mod tests { let filter = col("part").eq(lit(1)); let filter = logical2physical(&filter, table_schema.table_schema()); - let opener = make_opener( + let mut opener = make_opener( Arc::clone(&object_store), table_schema.clone(), Some(filter), ); + opener.scan_v2 = scan_v2; let stream = opener.open(file.clone()).unwrap().await.unwrap(); let data = stream.try_collect::>().await?; @@ -945,11 +1125,12 @@ mod tests { let filter = col("part").eq(lit(2)); let filter = logical2physical(&filter, table_schema.table_schema()); - let opener = make_opener( + let mut opener = make_opener( Arc::clone(&object_store), table_schema.clone(), Some(filter), ); + opener.scan_v2 = scan_v2; let stream = opener.open(file.clone()).unwrap().await.unwrap(); let data = stream.try_collect::>().await?; @@ -965,6 +1146,7 @@ mod tests { async fn test_residual_filter_unprojected_column( #[values(false, true)] projection_pushdown: bool, #[values(false, true)] reordered_schema: bool, + #[values(false, true)] scan_v2: bool, ) -> anyhow::Result<()> { let store: Arc = Arc::new(InMemory::new()); let batch = record_batch!( @@ -992,6 +1174,7 @@ mod tests { let mut opener = make_opener(store, TableSchema::from(Arc::clone(&schema)), Some(filter)); opener.projection = ProjectionExprs::from_indices(&[schema.index_of("a")?], &schema); opener.projection_pushdown = projection_pushdown; + opener.scan_v2 = scan_v2; opener.limit = Some(1); let batches = opener .open(PartitionedFile::new("residual.vortex", size))? @@ -1130,8 +1313,11 @@ mod tests { Ok(()) } + #[rstest] #[tokio::test] - async fn test_open_applies_limit_after_filtering() -> anyhow::Result<()> { + async fn test_open_applies_limit_after_filtering( + #[values(false, true)] scan_v2: bool, + ) -> anyhow::Result<()> { let object_store = Arc::new(InMemory::new()) as Arc; let file_path = "filtered-limit/file.vortex"; let batch = record_batch!(( @@ -1161,6 +1347,7 @@ mod tests { let mut opener = make_opener(object_store, table_schema, Some(filter)); opener.limit = Some(3); + opener.scan_v2 = scan_v2; let batches = opener.open(file)?.await?.try_collect::>().await?; let values = batches @@ -1294,6 +1481,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), + segment_sources: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1301,6 +1489,7 @@ mod tests { file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, + scan_v2: false, }; let filter = col("a").lt(lit(100_i32)); @@ -1382,6 +1571,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), + segment_sources: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1389,6 +1579,7 @@ mod tests { file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, + scan_v2: false, }; let stream = opener.open(file)?.await?; @@ -1561,6 +1752,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), + segment_sources: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1568,6 +1760,7 @@ mod tests { file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, + scan_v2: false, }; // This should succeed and return the correctly projected and cast data @@ -1622,6 +1815,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), + segment_sources: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1629,6 +1823,7 @@ mod tests { file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, + scan_v2: false, } } @@ -1834,6 +2029,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), + segment_sources: Default::default(), natural_splits: Default::default(), file_schemas: Default::default(), has_output_ordering: false, @@ -1841,6 +2037,7 @@ mod tests { file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, + scan_v2: false, }; let file = PartitionedFile::new(file_path.to_string(), data_size); diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 22fb4659155..5a86abee50c 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -11,10 +11,12 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_groups::FileGroupPartitioner; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::file_stream::FileOpener; use datafusion_execution::cache::cache_manager::FileMetadataCache; use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::conjunction; @@ -32,6 +34,7 @@ use object_store::ObjectStore; use object_store::path::Path; use vortex::file::VORTEX_FILE_EXTENSION; use vortex::layout::LayoutReader; +use vortex::layout::segments::SegmentSource; use vortex::metrics::DefaultMetricsRegistry; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; @@ -198,6 +201,8 @@ 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>>, /// Shared full-file natural splits keyed by path. natural_splits: Arc>>, /// Shared physical Arrow schemas of each file, keyed by path. @@ -234,6 +239,7 @@ impl VortexSource { vortex_predicate: None, df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), + segment_sources: Arc::new(DashMap::default()), natural_splits: Arc::new(DashMap::default()), file_schemas: Arc::new(DashMap::default()), expression_converter, @@ -359,6 +365,7 @@ 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), 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, @@ -366,6 +373,7 @@ impl VortexSource { file_metadata_cache: self.file_metadata_cache.clone(), projection_pushdown: self.options.projection_pushdown, scan_concurrency: self.options.scan_concurrency, + scan_v2: self.options.scan_v2, }; Ok(opener) @@ -445,6 +453,29 @@ impl FileSource for VortexSource { true } + fn repartitioned( + &self, + target_partitions: usize, + repartition_file_min_size: usize, + output_ordering: Option, + config: &FileScanConfig, + ) -> DFResult> { + if config.file_compression_type.is_compressed() { + return Ok(None); + } + let Some(file_groups) = FileGroupPartitioner::new() + .with_target_partitions(target_partitions) + .with_repartition_file_min_size(repartition_file_min_size) + .with_preserve_order_within_groups(output_ordering.is_some()) + .repartition_file_groups(&config.file_groups) + else { + return Ok(None); + }; + let mut repartitioned = config.clone(); + repartitioned.file_groups = file_groups; + Ok(Some(repartitioned)) + } + fn try_pushdown_filters( &self, filters: Vec>, diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 30ef93453ce..5d91242c91d 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -61,8 +61,14 @@ use crate::common_tests::TestSessionContext; fn make_session( object_store: Arc, repartition_file_scans: bool, + scan_v2: bool, ) -> SessionContext { - let factory = Arc::new(VortexFormatFactory::new()); + let factory = Arc::new( + VortexFormatFactory::new().with_options(crate::VortexTableOptions { + scan_v2, + ..Default::default() + }), + ); let config = SessionConfig::new() .with_target_partitions(4) @@ -471,9 +477,11 @@ async fn doc_example() -> anyhow::Result<()> { Ok(()) } +#[rstest] #[tokio::test] -async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits() -> anyhow::Result<()> -{ +async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits( + #[values(false, true)] scan_v2: bool, +) -> anyhow::Result<()> { let store = Arc::new(InMemory::new()) as _; let session = VortexSession::default(); let path = object_store::path::Path::parse("/split-aligned-repartition.vortex")?; @@ -544,8 +552,8 @@ async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits() - .any(|window| window[0] != window[1]) ); - let serial_ctx = make_session(Arc::clone(&store), false); - let repartitioned_ctx = make_session(Arc::clone(&store), true); + let serial_ctx = make_session(Arc::clone(&store), false, scan_v2); + let repartitioned_ctx = make_session(Arc::clone(&store), true, scan_v2); let repartitioned_partitions = count_query_partitions( &repartitioned_ctx, "SELECT value FROM '/split-aligned-repartition.vortex'", @@ -744,8 +752,10 @@ async fn test_predicate_memtable_oracle( #[case] projection: &str, #[case] predicate: &str, #[values(false, true)] pushdown: bool, + #[values(false, true)] scan_v2: bool, ) -> anyhow::Result<()> { let options = crate::VortexTableOptions { + scan_v2, projection_pushdown: pushdown, predicate_pushdown: pushdown, ..Default::default() diff --git a/vortex-layout/src/plan/physical/operators/scan.rs b/vortex-layout/src/plan/physical/operators/scan.rs index b6d2c463a81..842bd33ee87 100644 --- a/vortex-layout/src/plan/physical/operators/scan.rs +++ b/vortex-layout/src/plan/physical/operators/scan.rs @@ -115,28 +115,3 @@ impl Operator for ScanSource { )?)) } } - -/// Narrows a whole decoded segment, the one batch of a share's port, to a reader's rows. -pub(crate) struct SelectStage { - slice: Option>, - filter: Option, -} - -impl SelectStage { - pub(crate) fn new(slice: Option>, filter: Option) -> Self { - Self { slice, filter } - } -} - -impl Operator for SelectStage { - fn compute(&mut self, input: Input, _cx: &mut Cx<'_>) -> VortexResult { - let Input::Batch(batch) = input else { - return Ok(Step::End); - }; - Ok(Step::Final(select( - batch, - self.slice.as_ref(), - self.filter.as_ref(), - )?)) - } -} diff --git a/vortex-layout/src/plan/physical/optimize.rs b/vortex-layout/src/plan/physical/optimize.rs index 26f13c00021..d52f3da241f 100644 --- a/vortex-layout/src/plan/physical/optimize.rs +++ b/vortex-layout/src/plan/physical/optimize.rs @@ -8,6 +8,7 @@ use vortex_error::VortexResult; +use crate::plan::physical::Concat; use crate::plan::physical::PlanRef; use crate::plan::physical::optimizer::reduce_parent; use crate::plan::physical::optimizer::reduce_plan; @@ -26,6 +27,13 @@ fn reduce(plan: &PlanRef) -> VortexResult> { /// Optimizes `plan`, preserving its dtype and row domain. pub fn optimize(plan: PlanRef) -> VortexResult { + // Concat has no local rewrites. Optimizing its independent chunks eagerly traverses the + // entire file even when a partition only scans a small row range. + if let Some(concat) = plan.as_opt::() { + return Ok(concat + .map_chunks(plan.dtype().clone(), optimize) + .into_plan()); + } if let Some(rewritten) = reduce(&plan)? { return optimize(rewritten); } diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index 74f5296cd19..0e18ebe4e4c 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -6,8 +6,6 @@ use std::ops::Range; -use rustc_hash::FxHashMap; -use smallvec::SmallVec; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -22,10 +20,7 @@ use crate::plan::physical::PlanRef; use crate::plan::physical::SegmentScanPlan; use crate::plan::physical::operator::Operator; use crate::plan::physical::operator::Source; -use crate::plan::physical::operators::ConcatSource; use crate::plan::physical::operators::ScanSource; -use crate::plan::physical::operators::SelectStage; -use crate::segments::SegmentId; /// A pipeline under construction: a source and the stages after it, not yet given outlets, so a /// parent plan can add stages before it becomes a pipeline. @@ -46,14 +41,6 @@ impl Chain { } /// A chain passing the batches of `port` through. - pub(crate) fn port(port: PortId) -> Self { - Self { - source: Box::new(ConcatSource::new(1)), - stages: Vec::new(), - inlets: vec![port], - } - } - /// This chain with `stage` after its last stage. pub fn with(mut self, stage: impl Operator + 'static) -> Self { self.stages.push(Box::new(stage)); @@ -161,217 +148,12 @@ impl Compiler<'_> { } else { Some(usize::try_from(rows.start)?..usize::try_from(rows.end)?) }; - let split = self.core.split; - Ok(Some(match self.claim(scan, split) { - None => Chain::new(ScanSource::new(scan.clone(), slice, filter)), - Some(port) => { - let chain = Chain::port(port); - if slice.is_none() && filter.is_none() { - chain - } else { - chain.with(SelectStage::new(slice, filter)) - } - } - })) - } - - /// Claims split `split`'s port for `scan`'s segment when the segment is shared, building - /// the pipeline that decodes it for the first claimant with one port per reader in every - /// split not yet finished, or `None` when this split's reader is its only one. - fn claim(&mut self, scan: &SegmentScanPlan, split: usize) -> Option { - if split == usize::MAX { - return None; - } - let segment = scan.segment_id(); - let shares = &mut self.core.shares; - let ports = &mut shares.ports[split]; - if let Some(waiting) = ports.get_mut(&segment) { - let port = waiting.pop(); - if waiting.is_empty() { - ports.remove(&segment); - } - return port; - } - let ranges = shares.shared.remove(&segment)?; - let mut claimants: SmallVec<[usize; 4]> = SmallVec::new(); - for rows in &ranges { - claimants.extend( - shares - .overlapping(rows) - .filter(|&split| !shares.finished[split] && shares.overlaps(split, rows)), - ); - } - if claimants.len() <= 1 { - return None; - } - let mut outlets: SmallVec<[PortId; 1]> = SmallVec::with_capacity(claimants.len()); - let mut mine = None; - for claimant in claimants { - let port = self - .core - .arena - .create(usize::MAX, None, SourceEnd::Unclaimed); - outlets.push(port); - if claimant == split && mine.is_none() { - mine = Some(port); - } else { - self.core.shares.ports[claimant] - .entry(segment) - .or_default() - .push(port); - } - } - self.core.add_pipeline( - Chain::new(ScanSource::new(scan.clone(), None, None)), - outlets, - ); - mine - } -} - -/// Where each segment read more than once is read from, and the ports of the segments whose -/// decoding pipeline is built. -#[derive(Default)] -pub(crate) struct Shares { - /// The splits' row ranges, in split order. - splits: Vec>, - /// Whether the splits are sorted and disjoint, so the splits overlapping a range are found - /// by binary search. - ordered: bool, - /// Splits that have finished, and so will claim nothing. - finished: Vec, - /// Segment readers collected before the scan starts. A split may read only a few of a - /// file's segments, so recording a reader must not allocate by the file-wide segment id. - reaches: Vec<(SegmentId, Range)>, - /// The row ranges of segments read more than once, until their decoding pipeline is built. - shared: FxHashMap; 1]>>, - /// By split: the unclaimed ports of each shared segment the split reads. - ports: Vec>>, -} - -impl Shares { - pub(crate) fn new(splits: Vec>) -> Self { - let ordered = splits.windows(2).all(|pair| pair[0].end <= pair[1].start); - Self { - finished: vec![false; splits.len()], - ports: (0..splits.len()).map(|_| FxHashMap::default()).collect(), - splits, - ordered, - reaches: Vec::new(), - shared: FxHashMap::default(), - } - } - - /// The splits overlapping `rows`. - fn overlapping(&self, rows: &Range) -> Range { - if self.ordered { - let first = self.splits.partition_point(|split| split.end <= rows.start); - let end = self.splits.partition_point(|split| split.start < rows.end); - return first..end.max(first); - } - let mut overlapping = self - .splits - .iter() - .enumerate() - .filter(|(_, split)| split.start < rows.end && rows.start < split.end) - .map(|(index, _)| index); - match overlapping.next() { - Some(first) => first..overlapping.next_back().unwrap_or(first) + 1, - None => 0..0, - } - } - - fn overlaps(&self, split: usize, rows: &Range) -> bool { - let split = &self.splits[split]; - split.start < rows.end && rows.start < split.end - } - - /// Records the ranges `plan` reads each of its segments for, read over `rows` by every - /// split overlapping them. Call once per plan a split stage runs, before the scan starts, - /// then [`retain_shared`](Self::retain_shared). - pub(crate) fn add(&mut self, plan: &PlanRef, rows: Range) -> VortexResult<()> { - let reaches = &mut self.reaches; - plan.reach(rows, &Reach::Offset(0), &mut |segment, rows| { - reaches.push((segment, rows)) - }) - } - - /// Keeps only the segments more than one reader reads. The rest are read by the one - /// pipeline that needs them. - pub(crate) fn retain_shared(&mut self) { - let mut reaches = std::mem::take(&mut self.reaches); - reaches.sort_unstable_by_key(|(segment, _)| *segment); - for readers in reaches.chunk_by(|(left, _), (right, _)| left == right) { - let (segment, rows) = &readers[0]; - // One reader per split overlapping the range: shared when the range crosses a - // split boundary, or when multiple plan nodes read the segment. - if readers.len() == 1 - && self - .overlapping(rows) - .filter(|&split| self.overlaps(split, rows)) - .nth(1) - .is_none() - { - continue; - } - self.shared.insert( - *segment, - readers.iter().map(|(_, rows)| rows.clone()).collect(), - ); - } - } - - /// Every port made for a reader that has not claimed it yet. - #[cfg(debug_assertions)] - pub(crate) fn waiting(&self) -> impl Iterator + '_ { - self.ports - .iter() - .flat_map(|ports| ports.values().flatten().copied()) - } - - /// Marks split `split` finished and returns the ports it never claimed: those of its - /// readers that never came, as a chunk its mask ruled out or a stage it never ran. Each is - /// still an outlet of the pipeline that shares it, which must stop pushing into it: a port - /// has a consumer until it is closed and drained, so the caller detaches each port from - /// its sink before discarding it. See [`Core::finish_split`]. - pub(crate) fn finish(&mut self, split: usize) -> SmallVec<[PortId; 4]> { - self.finished[split] = true; - std::mem::take(&mut self.ports[split]) - .into_values() - .flatten() - .collect() - } -} - -/// How a plan's rows map to the rows of the plan a scan's splits range over. -#[derive(Clone, Debug)] -pub enum Reach { - /// Row `r` is scanned row `r + offset`. - Offset(u64), - /// Every row is read for these scanned rows, as a dictionary's values are for its codes. - Fixed(Range), -} - -impl Reach { - /// The scanned rows `rows` are read for. - pub fn root(&self, rows: &Range) -> Range { - match self { - Reach::Offset(offset) => rows.start + offset..rows.end + offset, - Reach::Fixed(root) => root.clone(), - } - } - - /// The mapping of a child whose row zero is this plan's row `by`. - pub fn shift(&self, by: u64) -> Self { - match self { - Reach::Offset(offset) => Reach::Offset(offset + by), - Reach::Fixed(root) => Reach::Fixed(root.clone()), - } - } - - /// The mapping of a child read whole for `rows` of this plan, as dictionary values are. - pub fn fixed(&self, rows: &Range) -> Self { - Reach::Fixed(self.root(rows)) + self.core.segments.push(scan.segment_id()); + Ok(Some(Chain::new(ScanSource::new( + scan.clone(), + slice, + filter, + )))) } } @@ -398,15 +180,3 @@ pub(crate) fn overlapping<'a>( (local.start < local.end).then_some((index, start, local)) }) } - -/// Unclaimed shared ports, for tests. -#[cfg(test)] -impl Shares { - pub(crate) fn pending(&self) -> usize { - self.ports - .iter() - .flat_map(|ports| ports.values()) - .map(SmallVec::len) - .sum() - } -} diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs index b6037c9e6ea..869babeaf8d 100644 --- a/vortex-layout/src/plan/physical/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -46,7 +46,6 @@ mod values; pub use self::compile::Chain; pub use self::compile::Compiler; -pub use self::compile::Reach; pub(crate) use self::compile::overlapping; pub(super) use self::port::Arena; pub use self::port::Inlet; diff --git a/vortex-layout/src/plan/physical/pipeline/port.rs b/vortex-layout/src/plan/physical/pipeline/port.rs index e089a9f0f76..1c149248727 100644 --- a/vortex-layout/src/plan/physical/pipeline/port.rs +++ b/vortex-layout/src/plan/physical/pipeline/port.rs @@ -88,11 +88,6 @@ impl Queue { usize::from(self.front.is_some()) + self.batches.len() } - #[cfg(debug_assertions)] - pub(crate) fn is_free(&self) -> bool { - self.free - } - pub(crate) fn take_changed(&mut self) -> bool { std::mem::take(&mut self.changed) } @@ -293,21 +288,6 @@ impl Arena { self.maybe_free(id); } - /// Frees a port nobody will read, once it is no longer an outlet of any pipeline: a share's - /// port whose claimant never came. Closes it on the sink's behalf if the sink has not. - pub(crate) fn discard(&mut self, id: PortId) { - let queue = &mut self.queues[id]; - debug_assert!( - queue.sink.is_none() && queue.source == SourceEnd::Unclaimed, - "discarding port {id} still attached to a pipeline" - ); - queue.closed = true; - queue.source_gone = true; - queue.front = None; - queue.batches.clear(); - self.maybe_free(id); - } - fn maybe_free(&mut self, id: PortId) { let queue = &mut self.queues[id]; if queue.closed && queue.source_gone && !queue.free { @@ -318,12 +298,6 @@ impl Arena { } } - /// Every port slot, free or not, with its id. - #[cfg(debug_assertions)] - pub(crate) fn iter(&self) -> impl Iterator { - self.queues.iter().enumerate() - } - /// Ports in use. #[cfg(test)] pub(crate) fn live(&self) -> usize { diff --git a/vortex-layout/src/plan/physical/pipeline/scan.rs b/vortex-layout/src/plan/physical/pipeline/scan.rs index 9d650ebbd26..816c3566d7e 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan.rs @@ -23,7 +23,6 @@ use vortex_session::VortexSession; use super::DEFAULT_CAPACITY; use super::compile::Chain; -use super::compile::Shares; use super::port::Arena; use super::port::PipelineId; use super::port::PortId; @@ -285,16 +284,14 @@ pub(crate) struct Core { /// queued once. dirty: Vec, dirty_flags: Vec, - pub(crate) shares: Shares, - /// The split being compiled, whose shared ports a compile claims. `usize::MAX` when no - /// split is, as while the pipeline a share decodes with is built: it reads for no split. - pub(crate) split: usize, + /// The segment of every scan source compiled so far, in compile order. + pub(crate) segments: Vec, /// Whole values kept for later scans of the same split. pub(crate) values: ValueCache, } impl Core { - fn new(session: VortexSession, row_offset: u64, splits: Vec>) -> Self { + fn new(session: VortexSession, row_offset: u64) -> Self { Self { exec: session.create_execution_ctx(), session, @@ -308,8 +305,7 @@ impl Core { new_reads: VecDeque::new(), dirty: Vec::new(), dirty_flags: Vec::new(), - shares: Shares::new(splits), - split: usize::MAX, + segments: Vec::new(), values: ValueCache::default(), } } @@ -360,12 +356,9 @@ impl Core { plan: &PlanRef, rows: Range, mask: &Mask, - (slot, split): (usize, usize), + slot: usize, ) -> VortexResult> { - self.split = split; - let chain = self.compile(plan, rows, mask); - self.split = usize::MAX; - let Some(chain) = chain? else { + let Some(chain) = self.compile(plan, rows, mask)? else { return Ok(None); }; let port = self @@ -495,29 +488,6 @@ impl Core { .all(|&outlet| self.arena.has_room(outlet)) } - /// Finishes split `index`: its unclaimed share ports are detached from the pipelines - /// sharing them, which then push to the claimants alone, and freed. - pub(crate) fn finish_split(&mut self, index: usize) { - for port in self.shares.finish(index) { - // A sink that has finished already closed the port, and so left it. - if let Some(sink) = self.arena.get(port).sink { - let pipe = self.pipeline_mut(sink); - pipe.outlets.retain(|outlet| *outlet != port); - // The claimant that built the shared pipeline reads its own port to the end, - // so no split's leaving takes a shared pipeline's last outlet. - debug_assert!( - !pipe.outlets.is_empty(), - "pipeline {sink} shares its output with no split" - ); - self.arena.get_mut(port).sink = None; - if self.sink_has_room(sink) { - self.make_runnable(sink); - } - } - self.arena.discard(port); - } - } - /// Drops a finished pipeline. Its outlets are closed; its inlets lose their source; the reads /// it asked for and never took are cancelled, so their bytes are dropped on arrival instead /// of reaching a pipeline that reuses its id. @@ -666,22 +636,7 @@ impl Scan { ); } let root = Root::Plan(plan); - let mut core = Core::new( - session, - 0, - splits.iter().map(|split| split.rows.clone()).collect(), - ); - // Only the rows some split reads can share a segment between readers. - let start = splits - .iter() - .map(|split| split.rows.start) - .min() - .unwrap_or(0); - let end = splits.iter().map(|split| split.rows.end).max().unwrap_or(0); - match &root { - Root::Plan(plan) => core.shares.add(plan, start..end)?, - } - core.shares.retain_shared(); + let mut core = Core::new(session, 0); core.dirty_flags = vec![false; 1]; Ok(Self { core, @@ -694,6 +649,22 @@ impl Scan { }) } + /// The segments compiling `plan` over `rows` under `mask` reads, each once, in id order. + /// Compiling does no IO, so the segments a stage will read are known before it starts. + pub fn segments( + session: &VortexSession, + plan: &PlanRef, + rows: Range, + mask: &Mask, + ) -> VortexResult> { + let mut core = Core::new(session.clone(), 0); + core.compile(plan, rows, mask)?; + let mut segments = core.segments; + segments.sort_unstable(); + segments.dedup(); + Ok(segments) + } + /// Runs up to `max_active` splits at once. pub fn with_max_active(mut self, max_active: usize) -> Self { self.active.resize_with(max_active.max(1), || None); @@ -773,7 +744,7 @@ impl Scan { { Root::Plan(plan) => { self.core - .compile_stage(plan, split.rows.clone(), &split.mask, (slot, index))? + .compile_stage(plan, split.rows.clone(), &split.mask, slot)? } }; if self.next == self.splits.len() { @@ -781,9 +752,7 @@ impl Scan { // lowering tree or the source layouts retained by its lazy children. self.root = None; } - if output.is_none() { - self.core.finish_split(index); - } else { + if output.is_some() { self.active[slot] = Some(Active { index, output }); self.live += 1; } @@ -807,7 +776,6 @@ impl Scan { self.core.arena.drop_source(port); active.output = None; } - self.core.finish_split(active.index); self.active[slot] = None; self.live -= 1; Ok(()) @@ -828,10 +796,4 @@ impl Scan { .filter(|(owner, _)| owner.is_some()) .count() } - - /// Shared ports not yet claimed or dropped. - #[cfg(test)] - pub(crate) fn pending_shares(&self) -> usize { - self.core.shares.pending() - } } diff --git a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs index df404bd34ea..58174ba42c0 100644 --- a/vortex-layout/src/plan/physical/pipeline/scan/validate.rs +++ b/vortex-layout/src/plan/physical/pipeline/scan/validate.rs @@ -44,18 +44,6 @@ impl Core { ); } } - let waiting: Vec<_> = self.shares.waiting().collect(); - for (port, queue) in self.arena.iter() { - if queue.is_free() || queue.source != SourceEnd::Unclaimed { - continue; - } - assert!( - waiting.contains(&port), - "port {port} has no reader and no split will claim it, so its sink {:?} \ - pushes into it for nothing", - queue.sink - ); - } } /// Checks what one run of pipeline `id` reported: a source waits only on a condition some diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 5a8daa39665..da9af0ab96c 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -194,8 +194,8 @@ struct Run { /// The arrays of each split, in order. splits: Vec>, events: Vec, - /// Ports and shared readers still held once the scan finished. - leftover: (usize, usize), + /// Ports still held once the scan finished. + leftover: usize, } /// Drives a scan the way an owner does, completing reads only while the scan waits. @@ -248,7 +248,7 @@ fn drive( 0, "scan finished with a pipeline waiting for a read" ); - let leftover = (scan.live_ports(), scan.pending_shares()); + let leftover = scan.live_ports(); Ok(Run { arrays, splits, @@ -280,7 +280,7 @@ fn run_with_values( let scan = Scan::try_new(SESSION.clone(), plan.clone(), vec![Split { rows, mask }])? .with_value_cache(values); let run = drive(store, scan, pick)?; - assert_eq!(run.leftover, (0, 0), "a finished scan holds no port"); + assert_eq!(run.leftover, 0, "a finished scan holds no port"); Ok(run) } @@ -401,9 +401,9 @@ fn views_of_one_plan( } /// Splits cutting every column at different places, run as one scan, each produce their own -/// rows, and a segment several splits read is read and decoded once, then dropped. +/// rows, and between them read every segment. #[rstest] -fn splits_of_one_scan_read_each_segment_once( +fn splits_of_one_scan_read_the_segments_they_overlap( #[values([0, 4, 11, 20], [0, 9, 13, 20], [0, 1, 2, 20])] cuts: [u64; 4], #[values(1, 3)] active: usize, #[values(Delivery::Fifo, Delivery::Lifo)] order: Delivery, @@ -419,11 +419,10 @@ fn splits_of_one_scan_read_each_segment_once( .collect(); let scan = Scan::try_new(SESSION.clone(), plan, splits.clone())?.with_max_active(active); let run = drive(&store, scan, delivery(order))?; - assert_eq!(run.leftover, (0, 0), "every shared segment was dropped"); - assert_eq!( - segments_read(&run.events), - (0..store.segments.len() as u32).collect::>() - ); + assert_eq!(run.leftover, 0, "a finished scan holds no port"); + let mut read = segments_read(&run.events); + read.dedup(); + assert_eq!(read, (0..store.segments.len() as u32).collect::>()); for (index, split) in splits.iter().enumerate() { assert_view( &expected, @@ -717,9 +716,9 @@ fn eval_applies_expression_to_selected_rows() -> VortexResult<()> { Ok(()) } -/// A column read by two fields of one struct is read and decoded once, and both fields get it. +/// A column read by two fields of one struct is read by each field, and both get it. #[test] -fn a_segment_two_readers_need_is_read_once() -> VortexResult<()> { +fn a_segment_two_readers_need_is_read_by_each() -> VortexResult<()> { let mut store = Store::default(); let a = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); let column = lower(&store.chunked(&a, &[7, 13])?)?; @@ -747,7 +746,7 @@ fn a_segment_two_readers_need_is_read_once() -> VortexResult<()> { mask.clone(), delivery(Delivery::Lifo), )?; - assert_eq!(segments_read(&run.events), [0, 1]); + assert_eq!(segments_read(&run.events), [0, 0, 1, 1]); assert_view(&expected, &rows, &mask, run.arrays) } diff --git a/vortex-layout/src/plan/physical/plans/concat.rs b/vortex-layout/src/plan/physical/plans/concat.rs index c4c6b41dc11..0d74d705f4b 100644 --- a/vortex-layout/src/plan/physical/plans/concat.rs +++ b/vortex-layout/src/plan/physical/plans/concat.rs @@ -9,6 +9,8 @@ use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::CachedId; @@ -24,10 +26,8 @@ use crate::plan::physical::operators::ConcatSource; use crate::plan::physical::optimizer::PlanParentReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; use crate::plan::physical::pipeline::overlapping; use crate::plan::physical::selection::selected_slice; -use crate::segments::SegmentId; /// Concatenates its children row-wise. #[derive(Clone, Debug)] @@ -93,6 +93,34 @@ impl ConcatPlan { pub fn row_offsets(&self) -> &[u64] { &self.data().row_offsets } + + /// Defers independent child rewrites until a scan reaches the child. Row offsets remain + /// valid because each materialized result must preserve its child's row count. + pub(crate) fn map_chunks( + &self, + dtype: DType, + map: impl Fn(PlanRef) -> VortexResult + Send + Sync + 'static, + ) -> Self { + let children = self.children().clone(); + let child_dtype = dtype.clone(); + PlanParts { + vtable: Concat, + dtype, + row_count: self.row_count(), + data: self.data().clone(), + children: PlanChildren::lazy(self.children().len(), move |index| { + let child = children + .get(index)? + .ok_or_else(|| vortex_err!("Missing chunk {index}"))?; + let rows = child.row_count(); + let mapped = map(child)?; + vortex_ensure_eq!(mapped.dtype(), &child_dtype); + vortex_ensure_eq!(mapped.row_count(), rows); + Ok(mapped) + }), + } + .into_typed() + } } impl PlanVTable for Concat { @@ -188,22 +216,6 @@ impl PlanVTable for Concat { count => Some(compiler.join(chains, ConcatSource::new(count))), }) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - for (index, start, local) in overlapping(plan, &rows) { - plan.child_required(index)?.reach( - local.start - start..local.end - start, - &at.shift(start), - visit, - )?; - } - Ok(()) - } } /// Pushes an expression into every chunk of a [`Concat`]. @@ -219,14 +231,13 @@ impl PlanParentReduceRule for ExpressionConcatRule { parent: &Plan, _child_idx: usize, ) -> VortexResult> { - let expression = parent.expression(); - let chunks = child - .children() - .iter() - .map(|chunk| Ok(EvalPlan::try_new(expression.clone(), chunk?)?.into_plan())) - .collect::>>()?; + let expression = parent.expression().clone(); Ok(Some( - ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(), + child + .map_chunks(expression.dtype().clone(), move |chunk| { + Ok(EvalPlan::try_new(expression.clone(), chunk)?.into_plan()) + }) + .into_plan(), )) } } diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index 3de19590419..791e9715d8d 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -23,8 +23,6 @@ use crate::plan::physical::operators::EvalStage; use crate::plan::physical::optimizer::PlanReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; /// Applies an expression to the output of its child. #[derive(Clone, Debug)] @@ -130,15 +128,6 @@ impl PlanVTable for Eval { let chain = compiler.compile(&plan.child_plan()?, rows, mask)?; Ok(chain.map(|chain| chain.with(EvalStage::new(plan.expression().clone())))) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - plan.child_plan()?.reach(rows, at, visit) - } } fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> VortexResult<()> { diff --git a/vortex-layout/src/plan/physical/plans/filter.rs b/vortex-layout/src/plan/physical/plans/filter.rs index 6e8f032845a..857a69a376e 100644 --- a/vortex-layout/src/plan/physical/plans/filter.rs +++ b/vortex-layout/src/plan/physical/plans/filter.rs @@ -22,8 +22,6 @@ use crate::plan::physical::check_child_count; use crate::plan::physical::operators::MaskStage; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; /// Keeps only the selected rows of its child. /// @@ -121,13 +119,4 @@ impl PlanVTable for Filter { } })) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - plan.child_plan()?.reach(rows, at, visit) - } } diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index a005ecc0c83..7eb6bec85ff 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/vortex-layout/src/plan/physical/plans/legacy_stats.rs @@ -29,8 +29,6 @@ use crate::plan::physical::optimizer::PlanParentReduceRule; use crate::plan::physical::optimizer::PlanReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; const DATA: usize = 0; const ZONES: usize = 1; @@ -146,15 +144,6 @@ impl PlanVTable for LegacyStats { ) -> VortexResult> { compiler.compile(&plan.data_plan()?, rows, mask) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - plan.data_plan()?.reach(rows, at, visit) - } } /// Rewrites a legacy stats plan into a [`Zoned`] plan over its zone table converted to the diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index 25f12380493..3ead88f8892 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -46,8 +46,6 @@ use crate::plan::physical::operators::empty_struct; use crate::plan::physical::optimizer::PlanParentReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; /// Assembles a struct from one child per field, plus an optional trailing validity child. #[derive(Clone, Debug)] @@ -251,18 +249,6 @@ impl PlanVTable for Pack { let source = PackSource::new(plan.fields().clone(), nullable, count); Ok(Some(compiler.join(chains, source))) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - for child in plan.children().iter_refs() { - child?.reach(rows.clone(), at, visit)?; - } - Ok(()) - } } fn validate_field_child( diff --git a/vortex-layout/src/plan/physical/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs index 48ee399bb8f..c7970ba8a12 100644 --- a/vortex-layout/src/plan/physical/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -38,9 +38,7 @@ use crate::plan::physical::check_child_count; use crate::plan::physical::operators::OnceSource; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; use crate::plan::physical::plans::pack::rewrite_partition_root; -use crate::segments::SegmentId; const ROW_IDX_PARTITION_NAME: &str = "row_idx"; const CHILD_PARTITION_NAME: &str = "child"; @@ -115,16 +113,6 @@ impl PlanVTable for RowIdx { }; Ok(Some(Chain::new(OnceSource::new(indices)))) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - let _ = (plan, rows, at, visit); - Ok(()) - } } /// Plans an expression over a data source and its global row-index domain. diff --git a/vortex-layout/src/plan/physical/plans/segment_scan.rs b/vortex-layout/src/plan/physical/plans/segment_scan.rs index e75181f53cc..a07ff0e8b4d 100644 --- a/vortex-layout/src/plan/physical/plans/segment_scan.rs +++ b/vortex-layout/src/plan/physical/plans/segment_scan.rs @@ -19,7 +19,6 @@ use crate::plan::physical::PlanVTable; use crate::plan::physical::check_child_count; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; use crate::segments::SegmentId; /// Reads one serialized array segment. @@ -108,14 +107,4 @@ impl PlanVTable for SegmentScan { // A scan keeps the selected rows itself, as every plan produces only those. compiler.scan(plan, rows, Some(mask.clone())) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - visit(plan.segment_id(), at.root(&rows)); - Ok(()) - } } diff --git a/vortex-layout/src/plan/physical/plans/share.rs b/vortex-layout/src/plan/physical/plans/share.rs index e5ec641f342..96166e7c81f 100644 --- a/vortex-layout/src/plan/physical/plans/share.rs +++ b/vortex-layout/src/plan/physical/plans/share.rs @@ -26,8 +26,6 @@ use crate::plan::physical::operators::UNBOUNDED; use crate::plan::physical::operators::WholeSource; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; /// Produces its child's whole value once per split, as a /// [`SharedArray`](crate::arrays::shared::SharedArray) that the split's later scans reuse. @@ -133,16 +131,6 @@ impl PlanVTable for Share { }); Ok(Some(compiler.join_with(vec![chain], |_| UNBOUNDED, source))) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - plan.child_plan()? - .reach(0..plan.row_count(), &at.fixed(&rows), visit) - } } /// The rows of `rows` that `mask` selects, from a whole shared value. diff --git a/vortex-layout/src/plan/physical/plans/take.rs b/vortex-layout/src/plan/physical/plans/take.rs index 69171e10eef..acef148774a 100644 --- a/vortex-layout/src/plan/physical/plans/take.rs +++ b/vortex-layout/src/plan/physical/plans/take.rs @@ -33,8 +33,6 @@ use crate::plan::physical::operators::TakeSource; use crate::plan::physical::optimizer::PlanParentReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; const CODES: usize = 0; const VALUES: usize = 1; @@ -172,18 +170,6 @@ impl PlanVTable for Take { source, ))) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - let values = plan.values()?; - let len = values.row_count(); - values.reach(0..len, &at.fixed(&rows), visit)?; - plan.codes()?.reach(rows, at, visit) - } } /// Pushes an expression, or the part of it that reads the values, onto the dictionary values of diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index e5913298bab..d1a17702ebf 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -48,8 +48,6 @@ use crate::plan::physical::operators::WholeSource; use crate::plan::physical::optimizer::PlanParentReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; const DATA: usize = 0; const ZONES: usize = 1; @@ -373,22 +371,6 @@ impl PlanVTable for Zoned { }); Ok(Some(compiler.join_with(vec![chain], |_| UNBOUNDED, source))) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - if !plan.is_pruning() { - return match plan.data_plan()? { - Some(data) => data.reach(rows, at, visit), - None => Ok(()), - }; - } - let zones = plan.zones_plan()?; - zones.reach(0..zones.row_count(), &at.fixed(&rows), visit) - } } /// The per-zone results kept on a pruning `plan`. diff --git a/vortex-layout/src/plan/physical/selection.rs b/vortex-layout/src/plan/physical/selection.rs index 177f800afcd..1087ad61538 100644 --- a/vortex-layout/src/plan/physical/selection.rs +++ b/vortex-layout/src/plan/physical/selection.rs @@ -11,22 +11,13 @@ use vortex_error::vortex_ensure; use vortex_mask::Mask; use vortex_session::registry::CachedId; -use super::Concat; -use super::Eval; -use super::LegacyStats; -use super::Pack; use super::Plan; use super::PlanId; use super::PlanParts; use super::PlanRef; use super::PlanVTable; -use super::Take; -use super::Zoned; use super::pipeline::Chain; use super::pipeline::Compiler; -use super::pipeline::Reach; -use super::pipeline::overlapping; -use crate::segments::SegmentId; #[derive(Clone, Debug)] struct Selected; @@ -80,16 +71,6 @@ impl PlanVTable for Selected { } compiler.compile(&plan.child_required(0)?, rows, &mask) } - - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - let mask = plan.data().mask(&rows)?; - reach_selected(&plan.child_required(0)?, rows, &mask, at, visit) - } } impl SelectedData { @@ -135,62 +116,3 @@ pub(crate) fn selected_slice(mask: &Mask, rows: Range) -> Option<(Range, - mask: &Mask, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), -) -> VortexResult<()> { - if mask.all_false() || rows.is_empty() { - return Ok(()); - } - if mask.all_true() { - return plan.reach(rows, at, visit); - } - if let Some(concat) = plan.as_opt::() { - for (index, start, local) in overlapping(concat, &rows) { - let Some((span, mask)) = selected_slice( - mask, - usize::try_from(local.start - rows.start)? - ..usize::try_from(local.end - rows.start)?, - ) else { - continue; - }; - reach_selected( - &concat.child_required(index)?, - rows.start + span.start as u64 - start..rows.start + span.end as u64 - start, - &mask, - &at.shift(start), - visit, - )?; - } - return Ok(()); - } - if plan.is::() || plan.is::() { - for child in plan.children().iter() { - reach_selected(&child?, rows.clone(), mask, at, visit)?; - } - return Ok(()); - } - if let Some(take) = plan.as_opt::() { - let values = take.values()?; - values.reach(0..values.row_count(), &at.fixed(&rows), visit)?; - return reach_selected(&take.codes()?, rows, mask, at, visit); - } - if let Some(zoned) = plan.as_opt::() - && !zoned.is_pruning() - { - return match zoned.data_plan()? { - Some(data) => reach_selected(&data, rows, mask, at, visit), - None => Ok(()), - }; - } - if let Some(stats) = plan.as_opt::() { - return reach_selected(&stats.data_plan()?, rows, mask, at, visit); - } - // Filters over non-segment children consume every input row. Share, pruning, and unknown - // operators may also read whole dependencies under a sparse mask. - plan.reach(rows, at, visit) -} diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index da6fa56ccb9..44705753b3e 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -150,6 +150,28 @@ fn chunked_plan_exposes_chunks() -> VortexResult<()> { Ok(()) } +#[test] +fn optimize_preserves_unvisited_concat_children() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( + 4, + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![ + flat(2, dtype.clone(), 0), + unsupported(2, dtype), + ]), + ) + .into_layout(); + let plan = optimize(make_plan(layout)?)?; + assert_eq!(child_of(&plan, 0)?.row_count(), 2); + assert_unsupported( + child_of(&plan, 1) + .err() + .ok_or_else(|| vortex_err!("unsupported child unexpectedly lowered"))?, + ); + Ok(()) +} + #[test] fn chunked_plan_lowers_each_chunk_on_access() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-layout/src/plan/physical/typed.rs b/vortex-layout/src/plan/physical/typed.rs index 1c6d99fd93c..c783dc8137b 100644 --- a/vortex-layout/src/plan/physical/typed.rs +++ b/vortex-layout/src/plan/physical/typed.rs @@ -28,8 +28,6 @@ use crate::plan::physical::PlanVTable; use crate::plan::physical::display::PlanTreeDisplay; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; -use crate::segments::SegmentId; /// The combined allocation behind [`PlanRef`]. /// @@ -156,20 +154,6 @@ impl PlanRef { self.dyn_plan().dyn_compile(self, rows, mask, compiler) } - /// Visits the segments compiling this plan over `rows` could read. See - /// [`PlanVTable::reach`]. - pub fn reach( - &self, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - if rows.start >= rows.end { - return Ok(()); - } - self.dyn_plan().dyn_reach(self, rows, at, visit) - } - /// Displays this plan and its descendants with the default plan extractors. pub fn display_tree(&self) -> PlanTreeDisplay<'_> { PlanTreeDisplay::default_display(self) @@ -412,15 +396,6 @@ pub trait DynPlan: 'static + Send + Sync + Debug { mask: &Mask, compiler: &mut Compiler<'_>, ) -> VortexResult>; - - /// Visits this operator's segments. See [`PlanVTable::reach`]. - fn dyn_reach( - &self, - plan: &PlanRef, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()>; } impl DynPlan for PlanData { @@ -462,14 +437,4 @@ impl DynPlan for PlanData { ) -> VortexResult> { V::compile(plan.as_::(), rows, mask, compiler) } - - fn dyn_reach( - &self, - plan: &PlanRef, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - V::reach(plan.as_::(), rows, at, visit) - } } diff --git a/vortex-layout/src/plan/physical/vtable.rs b/vortex-layout/src/plan/physical/vtable.rs index 328c7436a6e..75f3e4139c8 100644 --- a/vortex-layout/src/plan/physical/vtable.rs +++ b/vortex-layout/src/plan/physical/vtable.rs @@ -16,9 +16,7 @@ use vortex_session::registry::Id; use crate::plan::physical::PlanChildren; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; -use crate::plan::physical::pipeline::Reach; use crate::plan::physical::typed::Plan; -use crate::segments::SegmentId; /// A unique identifier for a plan operator. pub type PlanId = Id; @@ -92,20 +90,4 @@ pub trait PlanVTable: 'static + Clone + Sized + Send + Sync + Debug { let _ = (rows, mask, compiler); vortex_bail!("Plan {} has no pipeline implementation", plan.vtable().id()) } - - /// Visits every segment compiling this plan over `rows` could read, under any mask, with - /// the rows of the scanned plan it is read for, so a segment more than one reader reads is - /// decoded once. `at` maps this plan's rows to the scanned plan's. - /// - /// The default visits nothing, so the segments under a plan that does not override it are - /// read by each reader alone. - fn reach( - plan: &Plan, - rows: Range, - at: &Reach, - visit: &mut dyn FnMut(SegmentId, Range), - ) -> VortexResult<()> { - let _ = (plan, rows, at, visit); - Ok(()) - } } diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs new file mode 100644 index 00000000000..2ce5f588f02 --- /dev/null +++ b/vortex-layout/src/scan/execution/filter.rs @@ -0,0 +1,264 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use bit_vec::BitVec; +use vortex_array::VortexSessionExecute; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::transform::find_between; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_session::VortexSession; + +use super::ScanPlans; +use super::SelectedRows; +use super::adaptor::AdaptorStep; +use super::adaptor::ScanAdaptor; +use super::projection::ProjectionPlanner; +use crate::plan::physical::FilterPlan; +use crate::plan::physical::PlanRef; +use crate::plan::physical::SegmentScan; +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::plan_row_idx_expression; +use crate::scan::filter::FilterExpr; + +/// Plans and adaptive conjunct ordering shared by all splits of a scan. +pub(in crate::scan) struct FilterPlans { + expression: FilterExpr, + predicates: Vec, + pruning: Vec>, +} + +impl FilterPlans { + pub(in crate::scan) fn new( + expression: BoundExpression, + root: &PlanRef, + session: &VortexSession, + ) -> VortexResult { + let expression = FilterExpr::new(find_between(expression)?); + let predicates = expression + .conjuncts() + .iter() + .map(|expr| { + let plan = optimize(plan_row_idx_expression(expr.clone(), root.clone())?)?; + Ok(FilterPlan::new(plan).into_plan()) + }) + .collect::>()?; + let pruning = expression + .conjuncts() + .iter() + .map(|expr| { + expr.falsify(session)? + .map(|proof| optimize(plan_row_idx_expression(proof, root.clone())?)) + .transpose() + }) + .map(|plan| { + let plan = plan?; + match plan { + Some(plan) if metadata_only(&plan)? => Ok(Some(plan)), + _ => Ok(None), + } + }) + .collect::>()?; + Ok(Self { + expression, + predicates, + pruning, + }) + } +} + +/// A pruning proof must not fetch data when its requested statistics are unavailable. +fn metadata_only(plan: &PlanRef) -> VortexResult { + if plan + .as_opt::() + .is_some_and(|zoned| zoned.is_pruning()) + { + return Ok(true); + } + if plan.is::() { + return Ok(false); + } + for child in plan.children().iter() { + if !metadata_only(&child?)? { + return Ok(false); + } + } + Ok(true) +} + +/// Proves zone exclusions, evaluates conjuncts over survivors, then hands the mask to projection. +pub(super) struct FilterPlanner { + plans: ScanPlans, + selected: Option, + filter: Arc, + remaining: BitVec, + prune_index: usize, + active: Option<(usize, bool, ScanAdaptor)>, + pieces: Vec, + announced: bool, +} + +impl FilterPlanner { + pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + let filter = Arc::clone(plans.filter.as_ref().vortex_expect("filter plans exist")); + Self { + remaining: BitVec::from_elem(filter.predicates.len(), true), + plans, + selected: Some(selected), + filter, + prune_index: 0, + active: None, + pieces: Vec::new(), + announced: false, + } + } +} + +impl IoConsumer for FilterPlanner { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some((_, _, adaptor)) = &mut self.active { + adaptor.set_io_result(request, result); + } + } +} + +impl Planner for FilterPlanner { + fn state(&self) -> State { + if self.selected.is_none() { + return State::Done; + } + match &self.active { + Some((_, _, adaptor)) if adaptor.state() != State::Done => adaptor.state(), + _ => State::NeedsCompute, + } + } + + fn compute(&mut self) -> VortexResult { + let Some(selected) = &mut self.selected else { + return Ok(PlannerOutput::Done); + }; + if !self.announced { + self.announced = true; + let mut batch = Vec::new(); + for plan in self + .filter + .pruning + .iter() + .flatten() + .chain(self.filter.predicates.iter()) + .chain(std::iter::once(&self.plans.projection)) + { + let segments = Scan::segments( + &self.plans.session, + plan, + selected.scope.rows.clone(), + &selected.mask, + )?; + for segment in segments { + batch.push(IoRequest { + intent: IoIntent::Announce, + request: IoRequestId(u32::MAX - u32::try_from(batch.len())?), + target: self.plans.locations[*segment as usize].target(), + }); + } + } + if !batch.is_empty() { + return Ok(PlannerOutput::NeedsIO(batch)); + } + } + if let Some((index, pruning, adaptor)) = &mut self.active { + if adaptor.state() != State::Done { + return Ok(match adaptor.compute()? { + AdaptorStep::Yield => PlannerOutput::Continue, + AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + AdaptorStep::Piece(piece) => { + let mut ctx = self.plans.session.create_execution_ctx(); + let mask = piece.array.fill_null(false)?.execute::(&mut ctx)?; + self.pieces.push(mask); + PlannerOutput::Continue + } + }); + } + let input = selected.mask.true_count(); + let result = Mask::concat(self.pieces.iter())?; + vortex_ensure_eq!(result.len(), input); + let result = if *pruning { !result } else { result }; + selected.mask = selected.mask.intersect_by_rank(&result); + tracing::debug!( + row_start = selected.scope.rows.start, + row_end = selected.scope.rows.end, + input_rows = input, + output_rows = selected.mask.true_count(), + pruning = *pruning, + conjunct = *index, + "v2 filter conjunct evaluated" + ); + if !*pruning { + self.filter + .expression + .report_selectivity(*index, selected.mask.true_count() as f64 / input as f64); + } + self.active = None; + self.pieces.clear(); + } + if selected.mask.all_false() { + self.selected = None; + return Ok(PlannerOutput::Done); + } + + let mut next = None; + while self.prune_index < self.filter.pruning.len() { + let index = self.prune_index; + self.prune_index += 1; + if let Some(plan) = &self.filter.pruning[index] { + next = Some((index, true, plan.clone())); + break; + } + } + if next.is_none() + && let Some(index) = self.filter.expression.next_conjunct(&self.remaining) + { + self.remaining.set(index, false); + next = Some((index, false, self.filter.predicates[index].clone())); + } + if let Some((index, pruning, plan)) = next { + let scan = Scan::try_new( + self.plans.session.clone(), + plan, + vec![Split { + rows: selected.scope.rows.clone(), + mask: selected.mask.clone(), + }], + )? + .with_row_offset(self.plans.row_offset) + .with_value_cache(self.plans.values.clone()); + self.active = Some(( + index, + pruning, + ScanAdaptor::new(scan, Arc::clone(&self.plans.locations)), + )); + return Ok(PlannerOutput::Continue); + } + let selected = self.selected.take().vortex_expect("selected rows exist"); + Ok(PlannerOutput::Planner( + selected.scope.clone(), + Box::new(ProjectionPlanner::new(self.plans.clone(), selected)), + )) + } +} diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index f5fdafd3345..53cc685bccb 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -3,12 +3,13 @@ //! Scan execution over physical plans through the planner I/O and compute protocol. //! -//! A `ProjectionPlanner` emits one `ProjectionMorsel` for each fixed-size work chunk. The morsel -//! runs the physical projection plan over the chunk's selected rows through +//! A `ProjectionPlanner` emits bounded `ProjectionMorsel`s for a work chunk's selected rows. Each +//! morsel runs the physical projection plan through //! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published //! through a `ScanAdaptor`, so the planning driver can advance other work while reads wait. mod adaptor; +mod filter; mod projection; mod selection; @@ -20,6 +21,7 @@ use vortex_mask::Mask; use vortex_scan::planning::planner::WorkScope; use vortex_session::VortexSession; +pub(super) use self::filter::FilterPlans; pub(super) use self::selection::plan_selection; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; @@ -55,6 +57,7 @@ pub(super) struct ScanPlans { pub(super) locations: Arc<[SegmentLocation]>, /// The optimized projection plan, over the scan's root row domain. pub(super) projection: PlanRef, + pub(super) filter: Option>, /// The global row index of the root row domain's first row. pub(super) row_offset: u64, /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 4941cacb0a7..00c0a0614f6 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::vec::IntoIter; + use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; @@ -16,6 +18,7 @@ use vortex_scan::planning::morsel::MorselOutput; use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; @@ -26,16 +29,43 @@ use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; -/// Hands one fixed-size work chunk to a projection morsel. +const MAX_OUTPUT_ROWS: usize = 65_536; + +/// Keeps sparse selections together and subdivides dense selections to avoid large copies when +/// converting chunked projection results to Arrow. pub(super) struct ProjectionPlanner { - work: Option<(ScanPlans, SelectedRows)>, + plans: ScanPlans, + selected: IntoIter, } impl ProjectionPlanner { /// Creates a planner for `selected`. pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { + vec![selected] + } else { + (0..selected.mask.len()) + .step_by(MAX_OUTPUT_ROWS) + .filter_map(|start| { + let end = (start + MAX_OUTPUT_ROWS).min(selected.mask.len()); + let mask = selected.mask.slice(start..end); + if mask.all_false() { + return None; + } + Some(SelectedRows { + scope: WorkScope { + file_ordinal: selected.scope.file_ordinal, + rows: selected.scope.rows.start + start as u64 + ..selected.scope.rows.start + end as u64, + }, + mask, + }) + }) + .collect() + }; Self { - work: Some((plans, selected)), + plans, + selected: selected.into_iter(), } } } @@ -46,7 +76,7 @@ impl IoConsumer for ProjectionPlanner { impl Planner for ProjectionPlanner { fn state(&self) -> State { - if self.work.is_some() { + if !self.selected.as_slice().is_empty() { State::NeedsCompute } else { State::Done @@ -54,7 +84,7 @@ impl Planner for ProjectionPlanner { } fn compute(&mut self) -> VortexResult { - let Some((plans, selected)) = self.work.take() else { + let Some(selected) = self.selected.next() else { return Ok(PlannerOutput::Done); }; if selected.mask.all_false() { @@ -62,7 +92,7 @@ impl Planner for ProjectionPlanner { } Ok(PlannerOutput::Morsel( selected.scope.clone(), - Box::new(ProjectionMorsel::new(plans, selected)), + Box::new(ProjectionMorsel::new(self.plans.clone(), selected)), )) } } diff --git a/vortex-layout/src/scan/execution/selection.rs b/vortex-layout/src/scan/execution/selection.rs index b699eaaf98f..fe751b48887 100644 --- a/vortex-layout/src/scan/execution/selection.rs +++ b/vortex-layout/src/scan/execution/selection.rs @@ -11,6 +11,7 @@ use vortex_scan::selection::Selection; use super::ScanPlans; use super::SelectedRows; +use super::filter::FilterPlanner; use super::projection::ProjectionPlanner; /// Plans only the selected span, retaining one projection planner and IO session per work chunk. @@ -20,10 +21,13 @@ pub(crate) fn plan_selection( selection: &Selection, ) -> Option<(WorkScope, Box)> { let selected = select_rows(scope, selection)?; - Some(( - selected.scope.clone(), - Box::new(ProjectionPlanner::new(plans.clone(), selected)), - )) + let scope = selected.scope.clone(); + let planner: Box = if plans.filter.is_some() { + Box::new(FilterPlanner::new(plans.clone(), selected)) + } else { + Box::new(ProjectionPlanner::new(plans.clone(), selected)) + }; + Some((scope, planner)) } fn select_rows(mut scope: WorkScope, selection: &Selection) -> Option { diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 706de41515c..c6b359fa328 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -9,9 +9,10 @@ //! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. //! //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] -//! directly and supports projection, row ranges, and selections. Unsupported layouts return an -//! error; filters and limits are not part of this builder's API. Work is partitioned into fixed -//! chunks of up to 65,536 rows, independently of the stored layout's chunk boundaries. +//! directly and supports projection, filters, row ranges, and selections. Unsupported layouts return an +//! error; limits are not part of this builder's API. Work is partitioned into fixed +//! 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. //! //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, segment index, and dictionary @@ -33,12 +34,14 @@ use arrow_schema::SchemaRef; use futures::Stream; use futures::StreamExt; use futures::TryStreamExt; +use futures::channel::mpsc; use futures::stream; use futures::stream::BoxStream; use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; @@ -54,6 +57,7 @@ use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; +use crate::scan::execution::FilterPlans; use crate::scan::execution::ScanPlans; pub use crate::scan::execution::SegmentLocation; use crate::scan::execution::plan_selection; @@ -64,6 +68,7 @@ use crate::segments::SegmentSource; const ROW_BLOCK_SIZE: usize = 8192; const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; +const FILTERED_ROWS_PER_CHUNK: usize = 16 * MAX_ROWS_PER_CHUNK; const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; /// Limit estimated bytes per task without changing the writer-aligned row granularity. @@ -94,6 +99,22 @@ fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { .collect() } +/// 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. +pub fn partition_row_range(bytes: Range, file_size: u64, row_count: u64) -> Range { + let row = |byte: u64| -> u64 { + if byte >= file_size { + return row_count; + } + 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 + }; + let start = row(bytes.start); + start..row(bytes.end).max(start) +} + /// The parts of a file the planner-driven executor reads through. #[derive(Clone)] pub struct ScanFile { @@ -110,11 +131,12 @@ pub struct ScanFile { /// Builds a projection scan on the planner-driven V2 executor. /// /// This builder reads the file's layout and segments directly. Unsupported layouts return an -/// error when the stream is polled. Filters and row limits are not supported. +/// error when the stream is polled. Row limits are not supported. pub struct ScanBuilder { session: VortexSession, file: ScanFile, projection: BoundExpression, + filter: Option, ordered: bool, row_range: Option>, selection: Selection, @@ -128,6 +150,7 @@ impl ScanBuilder { pub fn new(session: VortexSession, file: ScanFile) -> Self { Self { projection: BoundExpression::new_root(file.layout.dtype().clone()), + filter: None, session, file, ordered: true, @@ -168,6 +191,12 @@ impl ScanBuilder { self } + /// Evaluates a bound predicate before projecting the surviving rows. + pub fn with_some_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + /// Configures whether batches are yielded in file order. pub fn with_ordered(mut self, ordered: bool) -> Self { self.ordered = ordered; @@ -205,7 +234,7 @@ impl ScanBuilder { } /// Maps each output batch on its split's task. - fn map( + pub fn map( self, map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, ) -> ScanBuilder { @@ -214,6 +243,7 @@ impl ScanBuilder { session: self.session, file: self.file, projection: self.projection, + filter: self.filter, ordered: self.ordered, row_range: self.row_range, selection: self.selection, @@ -262,6 +292,10 @@ impl PreparedScan { None => cache.insert_root(lower(&file.layout)?), }; let row_count = root.row_count(); + let filter = builder + .filter + .map(|filter| FilterPlans::new(filter, &root, &builder.session).map(Arc::new)) + .transpose()?; // An identity projection needs no rewrites. Keep the lowered children lazy, including // statistics that a projection-only scan never reads. let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype(); @@ -277,7 +311,9 @@ impl PreparedScan { } } }; - let chunk_rows = if identity { + let chunk_rows = if filter.is_some() { + FILTERED_ROWS_PER_CHUNK + } else if identity { rows_per_chunk(row_count, &file.locations) } else { MAX_ROWS_PER_CHUNK @@ -294,6 +330,7 @@ impl PreparedScan { session: builder.session, locations: file.locations, projection, + filter, row_offset: builder.row_offset, values: cache.values(), }, @@ -334,6 +371,21 @@ impl PreparedScan { let split_plans = self.split_plans(); let io = self.io; let map_fn = self.map_fn; + if !self.ordered + && split_plans.iter().any(|split| { + split.scope.rows.end - split.scope.rows.start > MAX_ROWS_PER_CHUNK as u64 + }) + { + let streams = stream::iter(split_plans).map(move |split| { + let (send, recv) = mpsc::channel(1); + let task = + handle.spawn(split.run_unordered(Arc::clone(&io), Arc::clone(&map_fn), send)); + let errors = stream::once(task) + .filter_map(|result| std::future::ready(result.err().map(Err))); + stream::select(recv, errors) + }); + return Ok(streams.flatten_unordered(concurrency).boxed()); + } // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. let tasks = stream::iter(split_plans).map(move |split| { let io = Arc::clone(&io); diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs index b9371735d39..624a775ad61 100644 --- a/vortex-layout/src/scan/v2/split.rs +++ b/vortex-layout/src/scan/v2/split.rs @@ -4,6 +4,8 @@ use std::future; use std::sync::Arc; +use futures::SinkExt; +use futures::channel::mpsc; use vortex_array::ArrayRef; use vortex_error::VortexResult; use vortex_io::request::IoService; @@ -24,6 +26,31 @@ pub(super) struct SplitPlan { } impl SplitPlan { + /// Publishes mapped batches as they finish, bounding how far projection runs ahead of its + /// consumer. Unordered scans do not need to retain the whole split to sort its output. + pub(super) async fn run_unordered( + self, + io: Arc, + map: Arc VortexResult + Send + Sync>, + mut send: mpsc::Sender>, + ) -> VortexResult<()> { + let mut run = Run::new(); + run.admit(self.root, self.scope, io.session()); + loop { + match run.advance()? { + Progress::Batch(batch) => { + let batch = map(batch.array); + let failed = batch.is_err(); + if send.send(batch).await.is_err() || failed { + return Ok(()); + } + } + Progress::RootDone(_) | Progress::Idle => return Ok(()), + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + } + /// Runs the split on a planning run of its own, reading through a session of `io`, and /// returns its batches in row order. /// diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 0f1f261eb94..ba29376bafb 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -6,10 +6,14 @@ use std::ops::Range; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::time::Duration; use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; +use futures::FutureExt; use futures::TryStreamExt; +use futures::channel::oneshot; +use futures::future::Shared; use futures::stream; use rstest::rstest; use vortex_array::ArrayContext; @@ -24,9 +28,14 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; +use vortex_array::expr::Expression; +use vortex_array::expr::and; +use vortex_array::expr::eq; use vortex_array::expr::get_item; 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::pack; use vortex_array::expr::root; use vortex_array::validity::Validity; @@ -34,12 +43,14 @@ use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use super::MAX_ROWS_PER_CHUNK; use super::PreparedScan; +use super::ROW_BLOCK_SIZE; use super::ScanBuilder; use super::ScanCache; use super::ScanFile; @@ -200,6 +211,7 @@ async fn assert_same_as_v1( )?; let mut v1 = V1ScanBuilder::new(session.clone(), reader) .with_projection(builder.projection.clone()) + .with_some_filter(builder.filter.clone()) .with_selection(builder.selection.clone()) .with_row_offset(builder.row_offset) .with_ordered(builder.ordered) @@ -220,6 +232,74 @@ enum Strategy { ZonedTable, } +#[rstest] +#[case::dense(-1)] +#[case::sparse(200_000)] +#[case::empty(300_000)] +#[tokio::test] +async fn large_filtered_chunks_match_v1(#[case] cutoff: i32) -> VortexResult<()> { + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let values = PrimitiveArray::from_iter(0i32..210_000).into_array(); + let dtype = values.dtype().clone(); + let layout = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new( + dtype.clone(), + stream::iter([Ok((sequence.advance(), values))]), + ) + .sendable(), + eof, + &session, + ) + .await?; + let segments: Arc = segments; + assert_same_as_v1(&session, || { + Ok(builder(&session, &segments, &layout)? + .with_some_filter(Some(gt(root(), lit(cutoff)).bind(&dtype)?))) + }) + .await +} + +#[rstest] +#[case::range(and( + gt(get_item("ints", root()), lit(995i32)), + lt(get_item("ints", root()), lit(2050i32)) +))] +#[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)))] +#[case::row_index(gt(row_idx(), lit(1500u64)))] +#[tokio::test] +async fn filter_matches_v1( + #[values(Strategy::Flat, Strategy::Table, Strategy::ZonedTable)] strategy: Strategy, + #[case] filter: Expression, + #[values(false, true)] sparse: bool, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, strategy.get().as_ref()).await?; + assert_same_as_v1(&session, || { + let mut scan = builder(&session, &segments, &layout)? + .with_projection( + pack([("strings", get_item("strings", root()))], false.into()) + .bind(layout.dtype())?, + ) + .with_some_filter(Some(filter.bind(layout.dtype())?)) + .with_row_range(900..3100) + .with_row_offset(100); + if sparse { + scan = scan.with_row_indices(StrictSortedBuffer::try_new(Buffer::from_iter( + (0..4000u64).step_by(7), + ))?); + } + Ok(scan) + }) + .await +} + impl Strategy { fn get(self) -> Arc { match self { @@ -402,12 +482,14 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { let reads = counted.reads.load(Ordering::Relaxed) - before; let code_chunks = (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; + // The first split reads every values chunk once for the whole scan, plus its codes. + let code_chunks = usize::try_from(code_chunks)?; assert_eq!( reads, if index == 0 { - STORED_CHUNKS + STORED_CHUNKS + code_chunks } else { - usize::try_from(code_chunks)? + code_chunks } ); let expected = @@ -437,6 +519,18 @@ async fn list_layouts_return_an_error() -> VortexResult<()> { Ok(()) } +#[tokio::test] +async fn projection_skips_unsupported_unselected_layouts() -> VortexResult<()> { + let session = new_session().with_tokio(); + let strategy = TableStrategy::new(chunked_flat(), chunked_flat()).with_list_layout(); + let (segments, layout) = write(&session, &strategy).await?; + assert_same_as_v1(&session, || { + Ok(builder(&session, &segments, &layout)? + .with_projection(get_item("ints", root()).bind(layout.dtype())?)) + }) + .await +} + #[tokio::test] async fn unordered_returns_every_row() -> VortexResult<()> { let session = new_session().with_tokio(); @@ -448,6 +542,88 @@ async fn unordered_returns_every_row() -> VortexResult<()> { Ok(()) } +struct GatedSegments { + inner: Arc, + blocked: SegmentId, + release: Shared>, +} + +impl SegmentSource for GatedSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + let bytes = self.inner.request(id); + if id != self.blocked { + return bytes; + } + let release = self.release.clone(); + async move { + let _ = release.await; + bytes.await + } + .boxed() + } +} + +#[rstest] +#[timeout(Duration::from_secs(5))] +#[tokio::test] +async fn unordered_yields_before_the_split_finishes() -> VortexResult<()> { + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let arrays = (0..2) + .map(|index| { + let start = i32::try_from(index * MAX_ROWS_PER_CHUNK)?; + Ok(( + sequence.advance(), + PrimitiveArray::from_iter(start..start + i32::try_from(MAX_ROWS_PER_CHUNK)?) + .into_array(), + )) + }) + .collect::>>()?; + let dtype = arrays[0].1.dtype().clone(); + let layout = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))) + .sendable(), + eof, + &session, + ) + .await?; + let blocked = *layout + .slot(1)? + .ok_or_else(|| vortex_err!("second chunk is missing"))? + .segment_ids() + .first() + .ok_or_else(|| vortex_err!("second chunk has no segment"))?; + let (release, gate) = oneshot::channel(); + let segments: Arc = Arc::new(GatedSegments { + inner: segments, + blocked, + release: gate.shared(), + }); + let mut scan = + PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; + // Put both projection morsels in one split, as a dense filtered scan would. + scan.splits = std::iter::once(0..layout.row_count()).collect(); + let mut batches = scan.execute_stream()?; + let first = batches.try_next().await; + let _ = release.send(()); + let first = first?.ok_or_else(|| vortex_err!("first batch is missing"))?; + assert_arrays_eq!( + first, + PrimitiveArray::from_iter(0i32..i32::try_from(MAX_ROWS_PER_CHUNK)?), + &mut session.create_execution_ctx() + ); + let rest: Vec<_> = batches.try_collect().await?; + assert_eq!( + rest.iter().map(|array| array.len()).sum::(), + MAX_ROWS_PER_CHUNK + ); + Ok(()) +} + #[tokio::test] async fn primitive_column() -> VortexResult<()> { let session = new_session().with_tokio(); @@ -603,3 +779,25 @@ async fn repeated_projections_reuse_their_values() -> VortexResult<()> { assert_eq!(reads, vec![4, 1, 1]); Ok(()) } + +#[test] +fn partition_rows_tile_the_file_on_row_blocks() { + let (file_size, row_count) = (1_000u64, 100_000u64); + let bytes = [0..300, 300..600, 600..1_000]; + let rows: Vec<_> = bytes + .iter() + .map(|bytes| super::partition_row_range(bytes.clone(), file_size, row_count)) + .collect(); + assert_eq!(rows[0].start, 0); + assert_eq!(rows[2].end, row_count); + for pair in rows.windows(2) { + 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!( + super::partition_row_range(1..2, file_size, row_count), + 0..0, + "a range inside one block owns no rows" + ); +}