diff --git a/bench-orchestrator/bench_orchestrator/runner/executor.py b/bench-orchestrator/bench_orchestrator/runner/executor.py index 1369c164f00..2814aabd9f2 100644 --- a/bench-orchestrator/bench_orchestrator/runner/executor.py +++ b/bench-orchestrator/bench_orchestrator/runner/executor.py @@ -17,6 +17,12 @@ console = Console() +# The benchmarks each engine runs on the planner-driven V2 Vortex scan. +SCAN_V2_BENCHMARKS: dict[Engine, frozenset[Benchmark]] = { + Engine.DATAFUSION: frozenset({Benchmark.TPCH}), + Engine.DUCKDB: frozenset({Benchmark.TPCH, Benchmark.TPCDS, Benchmark.CLICKBENCH}), +} + @final class BenchmarkExecutor: @@ -56,6 +62,8 @@ def build_command( cmd.extend(["--formats", ",".join(fmt.value for fmt in formats)]) if self.backend == Engine.DUCKDB: cmd.append("--delete-duckdb-database") + if benchmark in SCAN_V2_BENCHMARKS.get(self.backend, frozenset()): + cmd.extend(["--scan-v2", "true"]) if query is not None: cmd.extend(["--queries", str(query)]) diff --git a/bench-orchestrator/tests/test_executor.py b/bench-orchestrator/tests/test_executor.py index b8afb066e75..36ef737e766 100644 --- a/bench-orchestrator/tests/test_executor.py +++ b/bench-orchestrator/tests/test_executor.py @@ -5,6 +5,7 @@ import textwrap from pathlib import Path +import pytest from bench_orchestrator.config import Benchmark, Engine, Format from bench_orchestrator.runner.executor import BenchmarkExecutor @@ -33,6 +34,29 @@ def test_build_command_adds_duckdb_cleanup_flag() -> None: assert cmd[cmd.index("--opt") + 1] == "scale-factor=1.0" +@pytest.mark.parametrize( + ("engine", "benchmark", "expected"), + [ + (Engine.DUCKDB, Benchmark.TPCDS, True), + (Engine.DUCKDB, Benchmark.CLICKBENCH, True), + (Engine.DUCKDB, Benchmark.TPCH, True), + (Engine.DUCKDB, Benchmark.FINEWEB, False), + (Engine.DATAFUSION, Benchmark.TPCH, True), + (Engine.DATAFUSION, Benchmark.CLICKBENCH, False), + (Engine.LANCE, Benchmark.TPCH, False), + ], +) +def test_build_command_turns_on_scan_v2(engine: Engine, benchmark: Benchmark, expected: bool) -> None: + executor = BenchmarkExecutor(Path("/tmp/bench"), engine) + + cmd = executor.build_command(benchmark=benchmark, formats=[Format.VORTEX]) + + if expected: + assert cmd[cmd.index("--scan-v2") + 1] == "true" + else: + assert "--scan-v2" not in cmd + + def test_build_command_serializes_vortex_spatial_native_format() -> None: executor = BenchmarkExecutor(Path("/tmp/duckdb-bench"), Engine.DUCKDB) diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 1a83d374fcb..765b34fafcf 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -120,9 +120,9 @@ 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, + /// Use the planner-driven V2 Vortex scan. + #[arg(long, default_value_t = false, action = clap::ArgAction::Set)] + scan_v2: bool, #[arg(long = "opt", value_delimiter = ',', value_parser = value_parser!(Opt))] options: Vec, @@ -136,10 +136,7 @@ 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") - }); + let scan_v2 = args.scan_v2; set_join_set_tracer(get_static_tracer())?; setup_logging_and_tracing(args.verbose, args.tracing)?; diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index a047e275b24..f0b2cede63d 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -103,9 +103,9 @@ struct Args { )] reuse: bool, - /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H). - #[arg(long, action = clap::ArgAction::Set)] - scan_v2: Option, + /// Use the planner-driven V2 Vortex scan. + #[arg(long, default_value_t = false, action = clap::ArgAction::Set)] + scan_v2: bool, } #[global_allocator] @@ -116,11 +116,9 @@ fn main() -> anyhow::Result<()> { let args = Args::parse(); let opts = Opts::from(args.options); - let scan_v2 = args.scan_v2.unwrap_or_else(|| { - matches!(args.benchmark, BenchmarkArg::TpcH) - && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0") - }); - vortex_duckdb::set_scan_v2(scan_v2); + if args.scan_v2 { + vortex_duckdb::set_scan_v2(true); + } setup_logging_and_tracing(args.verbose, args.tracing)?; diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 97b310de84f..af005b4a5d4 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -1410,6 +1410,10 @@ fn query_i64_column(path: &str, query: &str) -> Result> { #[case::expression("SELECT a + b FROM ? WHERE a >= 50000 ORDER BY 1")] #[case::full_scan("SELECT a FROM ? ORDER BY a")] #[case::empty("SELECT a FROM ? WHERE b > 5000")] +#[case::top_n("SELECT a FROM ? ORDER BY b DESC, a LIMIT 7")] +#[case::filtered_top_n("SELECT a FROM ? WHERE b < 500 ORDER BY a DESC LIMIT 10")] +// DuckDB fetches the columns beyond the sort key by row id, as in ClickBench q23. +#[case::late_materialized_top_n("SELECT * FROM ? WHERE b < 700 ORDER BY a DESC LIMIT 5")] fn test_vortex_scan_v2_matches_v1(#[case] query: &str) -> Result<()> { // Many row blocks, so the file is cut into several V2 splits. let rows = 100_000i64; diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index 9f402f0e3ca..d8dc5f23072 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -12,6 +12,7 @@ use url::Url; use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::struct_::StructArrayExt as _; +use vortex::array::scalar_fn::fns::dynamic::DynamicComparison; use vortex::dtype::FieldPath; use vortex::error::VortexExpect; use vortex::error::VortexResult; @@ -28,6 +29,8 @@ use vortex::io::runtime::BlockingRuntime as _; use vortex::io::std_file::StdFileSystem; use vortex::layout::LayoutReaderRef; use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::layout::scan::scan_builder::referenced_field_masks; +use vortex::layout::scan::split_by::SplitBy; use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder; use vortex::layout::scan::v2::ScanFile; use vortex::mask::Mask; @@ -216,7 +219,18 @@ pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> Vor // global lock. if global.scan_v2 { let state = file.scan_state(); - let scan = FileScan::try_new(v2_builder(state, global))?; + let dynamic = global + .filter + .filter + .as_ref() + .map(|filter| filter.contains::()) + .transpose()? + .unwrap_or(false); + let scan = FileScan::try_new( + v2_builder(state, global), + &v2_boundaries(state, global)?, + dynamic, + )?; if state.v2.set(scan).is_err() { vortex_panic!("file initialized twice"); } @@ -257,6 +271,16 @@ fn v2_builder(state: &FileScanState, global: &GlobalState) -> V2ScanBuilder { .with_row_range(rows) } +/// The rows the file's chunks of the columns the scan reads end at, the boundaries V1 splits on. +fn v2_boundaries(state: &FileScanState, global: &GlobalState) -> VortexResult> { + let field_mask = referenced_field_masks(&global.projection, global.filter.filter.as_ref())?; + SplitBy::default().splits( + state.reader.as_ref(), + &(0..state.reader.row_count()), + &field_mask, + ) +} + /// What a thread's split yields next. enum SplitItem { /// The split has no more rows. @@ -301,13 +325,13 @@ fn next_split_item(local: &mut LocalState) -> VortexResult { /// the global lock. It only moves references between the file and the local /// state, and reader_scan drops the retired ones. pub fn reader_try_initialize_scan(file: &mut OpenFileReader, local: &mut LocalState) -> bool { - // A V2 scan's threads take chunks from the file's scan until it has spawned - // them all, so the file has work for another thread until then. + // A V2 scan's threads take chunks from the file's queue until it is empty, + // so the file has work for one more thread per queued chunk at most. let v2 = file .scan_state .as_ref() .and_then(|state| state.v2.get()) - .map(|scan| (!scan.is_exhausted()).then_some(Split::V2)); + .map(|scan| scan.claim().then_some(Split::V2)); let Some(split) = v2.unwrap_or_else(|| file.splits.pop()) else { // DuckDB releases the file after this call, under the global lock. // Give the reference of the file to this thread so it drops outside. diff --git a/vortex-duckdb/src/scan_v2.rs b/vortex-duckdb/src/scan_v2.rs index cee78cdc441..b74ad56b1e2 100644 --- a/vortex-duckdb/src/scan_v2.rs +++ b/vortex-duckdb/src/scan_v2.rs @@ -6,21 +6,37 @@ //! A planner says what it needs: compute, or reads. This driver decides where each runs, and //! DuckDB's own threads run all of it. Each file is prepared once and its planners are composed //! over the prepared scan: with a filter, its pruning plans over the scanned rows first, -//! continuing into the chunks of the rows they keep; without, the chunks directly. The root -//! planner's run is shared by every thread scanning the file. A thread takes the next chunk the -//! root spawns and runs it on a run of its own, inline, waiting only for its reads. Threads keep -//! taking chunks until the root is done, so the file's work spreads over as many threads as -//! DuckDB gives the scan, however its rows divide. +//! continuing into the chunks of the rows they keep; without, the chunks directly. //! -//! The root cuts the scanned rows into parts of whole 65,536-row blocks, a few per thread, and -//! spawns a chunk planner per part. A filtered scan's chunks are up to 1,048,576 rows, so without -//! parts a file of a few million rows would keep only a few threads busy. +//! The thread initializing the file runs the root planner to the end, pruning included, and +//! queues the chunks it spawns. DuckDB initializes a file on one thread while its others scan +//! other files, so no thread waits on another file's pruning reads. DuckDB then gives the file +//! one thread per queued chunk at most. A thread takes the next chunk and runs it on a run of its +//! own, inline, waiting only for its reads, and keeps taking chunks until the queue is empty, so +//! the file's work spreads over as many threads as it has chunks, however its rows divide. +//! +//! A part runs its chunks as children of one root, on the thread that takes it, so they share an +//! IO session: a segment that spans several chunks, such as a column chunk longer than the +//! scan's chunks, is read once. +//! +//! A filter with dynamic comparisons, such as the top-N bound DuckDB tightens during the scan, is +//! pruned again over each part when a thread takes it, as V1 prunes each split when it runs: a +//! part taken once the bound has tightened skips the zones it excludes. The zone maps were read +//! when the file was planned, so pruning a part again reads nothing. +//! +//! Chunks end only on the file's own chunk boundaries, of the columns the scan reads, so no chunk +//! of the file is decoded twice by two chunks of the scan. Spans between boundaries are merged up +//! to the scan's chunk size, but no further than leaves a few chunks per thread: a filtered scan's +//! chunks are up to 1,048,576 rows, so a file of a few million rows would otherwise keep only a few +//! threads busy. The root groups the chunks into parts, a few per thread, and spawns a chunk +//! planner per part. +use std::collections::VecDeque; use std::future; use std::ops::Range; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use parking_lot::Mutex; @@ -52,18 +68,18 @@ use crate::RUNTIME; /// compute. const CHUNKS_PER_THREAD: usize = 1; /// The parts of a file's rows per thread, so threads that finish early find more work. -const PARTS_PER_THREAD: u64 = 4; -/// The row blocks parts are cut on: the rows of a projection chunk, and of the writer's chunks. -const PART_ROW_BLOCK: u64 = 65_536; +const PARTS_PER_THREAD: usize = 4; /// One file's V2 scan, shared by the threads scanning it. pub(crate) struct FileScan { scan: Arc, io: Arc, - /// The run of the root planner, locked while a thread takes the chunk it spawns next. - root: Mutex, - /// Set once the root has spawned its last chunk. - exhausted: AtomicBool, + /// The work the root planned, in row order, that no thread has taken yet. + work: Mutex>, + /// The work the root planned, taken or not. + planned: usize, + /// The threads DuckDB has given the file so far. + claims: AtomicUsize, /// The rows of the chunks taken so far, for DuckDB's progress bar. rows_taken: AtomicU64, } @@ -77,39 +93,60 @@ enum Work { } impl FileScan { - /// Prepares `builder` and composes its planners. Lowering and planning touch no segment. - pub(crate) fn try_new(builder: ScanBuilder) -> VortexResult { + /// Prepares `builder`, composes its planners, cutting chunks on `boundaries`, the sorted rows + /// the file's chunks of the scanned columns end at, and runs the root planner to the end on + /// the calling thread, waiting for its pruning reads. With `dynamic`, each part is pruned + /// again when a thread takes it. + pub(crate) fn try_new( + builder: ScanBuilder, + boundaries: &[u64], + dynamic: bool, + ) -> VortexResult { let scan = Arc::new(builder.prepare()?); let io = scan.io(); let rows = scan.rows().clone(); + let mut parts = get_available_parallelism().unwrap_or(1) * PARTS_PER_THREAD; + let max_rows = (scan.chunk_rows() as u64).min(merged_rows(&rows, parts)); + let mut chunks = chunk_ranges(rows.clone(), boundaries, max_rows); + if scan.sparse_selection() { + // A sparse selection, such as an engine's fetch of a few rows by id, selects rows in + // few chunks: each of them is its own part, and the rest are not planned at all, so + // DuckDB gives the file a thread per chunk that has work, as V1 does with its splits. + let selection = scan.selection(); + chunks.retain(|chunk| selection.row_count_in_range(chunk) > 0); + parts = chunks.len(); + } let mut root = Run::new(); root.admit( - compose(&scan), + compose(&scan, chunks, parts, dynamic), WorkScope { file_ordinal: 0, rows, }, io.session(), ); + let work = plan(&mut root)?; Ok(Self { scan, io, - root: Mutex::new(root), - exhausted: AtomicBool::new(false), + planned: work.len(), + work: Mutex::new(work), + claims: AtomicUsize::new(0), rows_taken: AtomicU64::new(0), }) } - /// Whether the root has spawned its last chunk, so a thread joining the scan finds no work. - pub(crate) fn is_exhausted(&self) -> bool { - self.exhausted.load(Ordering::Acquire) + /// Whether the file has work for one more thread: DuckDB gives it one thread per piece of + /// planned work at most. + pub(crate) fn claim(&self) -> bool { + self.claims.fetch_add(1, Ordering::Relaxed) < self.planned } /// The percentage of the scanned rows whose chunks threads have taken. pub(crate) fn progress(&self) -> f64 { let rows = self.scan.rows(); let total = rows.end - rows.start; - if total == 0 || self.exhausted.load(Ordering::Relaxed) { + if total == 0 || self.work.lock().is_empty() { return 100.0; } 100.0 * self.rows_taken.load(Ordering::Relaxed) as f64 / total as f64 @@ -143,33 +180,29 @@ impl FileScan { } } - /// Advances the root until it spawns a chunk or produces a batch, or `None` once it is done. + /// The next work no thread has taken, or `None` once all is taken. fn next_work(&self) -> VortexResult> { - if self.exhausted.load(Ordering::Acquire) { - return Ok(None); + let work = self.work.lock().pop_front(); + if let Some(Work::Chunk(scope, _)) = &work { + self.rows_taken + .fetch_add(scope.rows.end - scope.rows.start, Ordering::Relaxed); } - let mut root = self.root.lock(); - loop { - if self.exhausted.load(Ordering::Acquire) { - return Ok(None); - } - match root.advance()? { - Progress::Batch(batch) => return Ok(Some(Work::Batch(batch.array))), - Progress::Spawn(scope, chunk) => { - self.rows_taken - .fetch_add(scope.rows.end - scope.rows.start, Ordering::Relaxed); - return Ok(Some(Work::Chunk(scope, chunk))); - } - Progress::RootDone(_) if root.live_roots() > 0 => {} - Progress::RootDone(_) | Progress::Idle => { - self.exhausted.store(true, Ordering::Release); - return Ok(None); - } - // The root waits for the pruning plans' reads. The threads waiting on the lock - // would have nothing to run until they arrive. - Progress::Waiting => { - RUNTIME.block_on(future::poll_fn(|cx| root.poll_completion(cx)))?; - } + Ok(work) + } +} + +/// Runs `root` to the end, waiting for its reads, and returns the chunks it spawns and the +/// batches it produces, in order. +fn plan(root: &mut Run) -> VortexResult> { + let mut work = VecDeque::new(); + loop { + match root.advance()? { + Progress::Batch(batch) => work.push_back(Work::Batch(batch.array)), + Progress::Spawn(scope, chunk) => work.push_back(Work::Chunk(scope, chunk)), + Progress::RootDone(_) if root.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => return Ok(work), + Progress::Waiting => { + RUNTIME.block_on(future::poll_fn(|cx| root.poll_completion(cx)))?; } } } @@ -177,18 +210,23 @@ impl FileScan { /// The planners of a scan, composed: with a filter, its pruning plans first over the scan's rows, /// continuing into the parts of the rows they keep; without, the parts directly. -fn compose(scan: &Arc) -> Box { +fn compose( + scan: &Arc, + chunks: Vec>, + count: usize, + dynamic: bool, +) -> Box { let rows = scan.rows().clone(); let parts = { let scan = Arc::clone(scan); let rows = rows.clone(); move |kept: Option| -> Box { - let parallelism = get_available_parallelism().unwrap_or(1) as u64; Box::new(PartPlanner::new( scan, rows, kept, - parallelism * PARTS_PER_THREAD, + parts(chunks, count), + dynamic, )) } }; @@ -199,8 +237,9 @@ fn compose(scan: &Arc) -> Box { } } -/// Spawns a chunk planner per part of a scan's rows, in row order. Parts end on whole -/// [`PART_ROW_BLOCK`]s, so no chunk of the scan spans two parts. +/// Spawns a chunk planner per part of a scan's rows, in row order, which runs the part's chunks +/// inline, behind the part's own pruning when the filter is dynamic. A part is a run of whole +/// chunks. struct PartPlanner { scan: Arc, selection: Selection, @@ -208,17 +247,27 @@ struct PartPlanner { first_row: u64, /// The rows zone pruning kept, from `first_row`, or `None` when nothing was pruned. kept: Option, - parts: std::vec::IntoIter>, + /// The chunks of each part, which tile the scan's rows. + parts: std::vec::IntoIter>>, + /// Whether each part is pruned again when it runs. + dynamic: bool, } impl PartPlanner { - fn new(scan: Arc, rows: Range, kept: Option, parts: u64) -> Self { + fn new( + scan: Arc, + rows: Range, + kept: Option, + parts: Vec>>, + dynamic: bool, + ) -> Self { Self { selection: scan.selection().clone(), scan, first_row: rows.start, kept, - parts: part_ranges(rows, parts).into_iter(), + parts: parts.into_iter(), + dynamic, } } } @@ -237,7 +286,11 @@ impl Planner for PartPlanner { } fn compute(&mut self) -> VortexResult { - for part in self.parts.by_ref() { + for chunks in self.parts.by_ref() { + let (Some(first), Some(last)) = (chunks.first(), chunks.last()) else { + continue; + }; + let part = first.start..last.end; let kept = match &self.kept { Some(kept) => { let offset = |row: u64| { @@ -251,38 +304,78 @@ impl Planner for PartPlanner { } None => None, }; - let chunks = ChunkPlanner::new( - Arc::clone(&self.scan), - part.clone(), - self.selection.clone(), - kept, - ); let scope = WorkScope { file_ordinal: 0, - rows: part, + rows: part.clone(), }; - return Ok(PlannerOutput::Spawn(scope, Box::new(chunks))); + let scan = Arc::clone(&self.scan); + let selection = self.selection.clone(); + if !self.dynamic { + let chunks = + ChunkPlanner::with_chunks(scan, part, chunks, selection, kept).inline(); + return Ok(PlannerOutput::Spawn(scope, Box::new(chunks))); + } + let next = move |pruned: Option| -> Box { + let kept = match (kept, pruned) { + (Some(kept), Some(pruned)) => Some(kept & &pruned), + (kept, pruned) => kept.or(pruned), + }; + Box::new( + ChunkPlanner::with_chunks(Arc::clone(&scan), part, chunks, selection, kept) + .inline(), + ) + }; + let pruning = + PrunePlanner::new(Arc::clone(&self.scan), scope.rows.clone(), Box::new(next)); + return Ok(PlannerOutput::Spawn(scope, Box::new(pruning))); } Ok(PlannerOutput::Done) } } -/// Cuts `rows` into up to `parts` non-empty ranges of whole [`PART_ROW_BLOCK`]s of the file's -/// row domain that tile it, clipped to `rows`. -fn part_ranges(rows: Range, parts: u64) -> Vec> { - if rows.is_empty() { +/// The most rows spans merge into while `rows` still divides into `parts` chunks. +fn merged_rows(rows: &Range, parts: usize) -> u64 { + (rows.end - rows.start).div_ceil(parts.max(1) as u64).max(1) +} + +/// Cuts `rows` into chunks that end only on `boundaries` or at the end of `rows`, merging the +/// spans between boundaries into chunks of at most `max_rows` rows. A span wider than `max_rows` +/// stays one chunk, so no chunk of the file is split. +fn chunk_ranges(rows: Range, boundaries: &[u64], max_rows: u64) -> Vec> { + let mut chunks = Vec::new(); + let mut start = rows.start; + let mut end = rows.start; + let ends = boundaries + .iter() + .copied() + .filter(|&row| row > rows.start && row < rows.end) + .chain((!rows.is_empty()).then_some(rows.end)); + for row in ends { + if row - start > max_rows && end > start { + chunks.push(start..end); + start = end; + } + end = row; + } + if end > start { + chunks.push(start..end); + } + chunks +} + +/// Groups `chunks` into up to `parts` runs of consecutive chunks, of as even a count as they +/// divide into. +fn parts(chunks: Vec>, parts: usize) -> Vec>> { + if chunks.is_empty() { return Vec::new(); } - let first_block = rows.start / PART_ROW_BLOCK; - let blocks = rows.end.div_ceil(PART_ROW_BLOCK) - first_block; - let blocks_per_part = blocks.div_ceil(parts.max(1)); - (0..blocks.div_ceil(blocks_per_part)) - .map(|part| { - let start = (first_block + part * blocks_per_part) * PART_ROW_BLOCK; - let end = start + blocks_per_part * PART_ROW_BLOCK; - start.max(rows.start)..end.min(rows.end) - }) - .collect() + let per_part = chunks.len().div_ceil(parts.max(1)); + let mut chunks = chunks.into_iter().peekable(); + let mut grouped = Vec::new(); + while chunks.peek().is_some() { + grouped.push(chunks.by_ref().take(per_part).collect()); + } + grouped } #[cfg(test)] @@ -291,23 +384,52 @@ mod tests { use rstest::rstest; - use super::PART_ROW_BLOCK; - use super::part_ranges; - - const B: u64 = PART_ROW_BLOCK; + use super::chunk_ranges; + use super::merged_rows; + use super::parts; #[rstest] - #[case::even(0..4 * B, 2, vec![0..2 * B, 2 * B..4 * B])] - #[case::ragged_end(0..3 * B + 10, 2, vec![0..2 * B, 2 * B..3 * B + 10])] - #[case::more_parts_than_blocks(0..2 * B, 8, vec![0..B, B..2 * B])] - #[case::small(0..100, 4, vec![0..100])] - #[case::unaligned_start(B / 2..3 * B, 2, vec![B / 2..2 * B, 2 * B..3 * B])] - #[case::empty(5..5, 4, vec![])] - fn parts_end_on_blocks( + #[case::merged(0..400, &[100, 200, 300, 400], 200, vec![0..200, 200..400])] + #[case::wide_spans(0..300, &[0, 87, 182, 300], 64, vec![0..87, 87..182, 182..300])] + #[case::clipped(50..250, &[100, 200, 300], 100, vec![50..100, 100..200, 200..250])] + #[case::no_boundaries(0..10, &[], 4, vec![0..10])] + #[case::empty(5..5, &[5], 4, vec![])] + fn chunks_end_on_boundaries( #[case] rows: Range, - #[case] parts: u64, + #[case] boundaries: &[u64], + #[case] max_rows: u64, #[case] expected: Vec>, ) { - assert_eq!(part_ranges(rows, parts), expected); + assert_eq!(chunk_ranges(rows, boundaries, max_rows), expected); + } + + #[rstest] + #[case::few_parts(0..1_000, 4, 250)] + #[case::many_parts(0..10, 384, 1)] + #[case::empty(5..5, 4, 1)] + fn merging_leaves_chunks_for_every_part( + #[case] rows: Range, + #[case] parts: usize, + #[case] expected: u64, + ) { + assert_eq!(merged_rows(&rows, parts), expected); + } + + #[rstest] + #[case::even(4, 2, vec![2, 2])] + #[case::ragged(5, 2, vec![3, 2])] + #[case::more_parts_than_chunks(2, 8, vec![1, 1])] + #[case::empty(0, 4, vec![])] + fn parts_group_whole_chunks( + #[case] chunks: u64, + #[case] count: usize, + #[case] expected: Vec, + ) { + let grouped = parts((0..chunks).map(|i| i..i + 1).collect(), count); + assert_eq!(grouped.iter().map(Vec::len).collect::>(), expected); + assert_eq!( + grouped.concat(), + (0..chunks).map(|i| i..i + 1).collect::>() + ); } } diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index f9641cbd6cb..389cad6059d 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -293,8 +293,8 @@ fn row_idx_dtype() -> DType { DType::Primitive(PType::U64, NonNullable) } -// Returns a SequenceArray representing the row indices for the given row range, -fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { +/// A [`SequenceArray`] of the row indices of `row_range`, offset by `row_offset`. +pub(crate) fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { Sequence::try_new( PValue::U64(row_offset + row_range.start), PValue::U64(1), diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index 043f0676198..d3ed93771bb 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -12,6 +12,7 @@ use vortex_array::expr::BoundExpression; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::pack::Pack; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -20,6 +21,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::plan::physical::Pack as PackVTable; use crate::plan::physical::PackPlan; use crate::plan::physical::Plan; use crate::plan::physical::PlanChildren; @@ -194,7 +196,12 @@ impl PlanParentReduceRule for EvalEvalRule { } } -/// Eval(pack({0: expr})) -> PackPlan(expr) for a non-nullable pack +/// Eval(pack({0: expr})) -> PackPlan(expr) for a non-nullable pack, and +/// Eval(pack({0: path0, 1: path1, ..}), Pack) -> PackPlan(path0, path1, ..) when every field is a +/// field path into a non-nullable [`Pack`] child. +/// +/// Lifting a pack of field paths into the plan lets each field reach its column's plan, so a scan +/// assembles each chunk's struct once instead of packing, renaming and repacking it per chunk. #[derive(Debug)] pub(crate) struct EvalPackRule; impl PlanReduceRule for EvalPackRule { @@ -205,21 +212,52 @@ impl PlanReduceRule for EvalPackRule { if expr.nullability != Nullability::NonNullable { return Ok(None); } - let [child_expr] = plan.expression.children() else { - return Ok(None); - }; - let fields = StructFields::new(expr.names.clone(), vec![child_expr.dtype().clone()]); - + let child_exprs = plan.expression.children(); let child_plan = plan.child(0)?.vortex_expect("no child plan"); - let child_plan = EvalPlan::try_new(child_expr.clone(), child_plan)?.into_plan(); + let child_plans = match child_exprs { + [child_expr] => vec![EvalPlan::try_new(child_expr.clone(), child_plan)?.into_plan()], + // A child read by several fields is shared only where each field steps into it by + // name, so that each field reads the plan of its own column. + _ => { + let Some(pack) = child_plan.as_opt::() else { + return Ok(None); + }; + if child_exprs.is_empty() + || pack.dtype().is_nullable() + || !child_exprs.iter().all(is_field_path) + { + return Ok(None); + } + child_exprs + .iter() + .map(|child| pack.field_path_plan(child)) + .collect::>>()? + } + }; + let fields = StructFields::new( + expr.names.clone(), + child_exprs + .iter() + .map(|child| child.dtype().clone()) + .collect(), + ); let plan = PackPlan::try_new( fields, Nullability::NonNullable, plan.row_count(), - vec![child_plan], + child_plans, None, )?; Ok(Some(plan.into_plan())) } } + +/// Whether `expression` is the root or a chain of field accesses from it, such as `$.a.b`. +pub(super) fn is_field_path(expression: &BoundExpression) -> bool { + if expression.is_root() { + return true; + } + expression.as_opt::().is_some() + && matches!(expression.children(), [child] if is_field_path(child)) +} diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index d3886de19ad..4b82bc6d0f9 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -32,6 +32,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use super::eval::is_field_path; use crate::plan::physical::Eval; use crate::plan::physical::EvalPlan; use crate::plan::physical::Plan; @@ -291,6 +292,30 @@ fn validate_validity_child(expected_row_count: u64, child: &PlanRef) -> VortexRe } impl PackPlan { + /// The plan of `path`, a chain of field accesses such as `$.a.b` over this non-nullable + /// pack, stepped into the plan of its first field. + pub(crate) fn field_path_plan(&self, path: &BoundExpression) -> VortexResult { + let mut first = path; + while let [child] = first.children() + && !child.is_root() + { + first = child; + } + let Some(name) = first.as_opt::() else { + return Ok(EvalPlan::try_new(path.clone(), self.to_plan())?.into_plan()); + }; + let index = self + .fields() + .find(name) + .ok_or_else(|| vortex_err!("Struct expression references unknown field '{name}'"))?; + let field = field_plan(self, index)?; + let lowered = step_into_struct_field(path.clone(), name, field.dtype().clone())?; + if lowered.is_root() { + return Ok(field); + } + Ok(EvalPlan::try_new(lowered, field)?.into_plan()) + } + /// Rebuilds this plan with only `fields`, which must be a subset of the current fields. /// /// Pruning is only sound for a non-nullable struct: dropping a field of a nullable struct @@ -338,6 +363,11 @@ impl PlanParentReduceRule for ExpressionPackRule { } let expression = parent.expression(); + // A field path steps straight into its field's plan, without the annotation of every + // field that a general expression needs. + if !expression.is_root() && is_field_path(expression) { + return Ok(Some(child.field_path_plan(expression)?)); + } let fields = child.fields(); let referenced_fields = descendent_bound_annotations(expression, make_bound_free_field_annotator(fields)) diff --git a/vortex-layout/src/plan/physical/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs index 5c45d2cae4f..e300506b866 100644 --- a/vortex-layout/src/plan/physical/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -17,7 +17,6 @@ use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar_fn::fns::pack::Pack as PackFn; -use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -26,6 +25,7 @@ use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::layouts::row_idx::RowIdx as RowIdxFn; +use crate::layouts::row_idx::idx_array; use crate::plan::physical::EvalPlan; use crate::plan::physical::PackPlan; use crate::plan::physical::Plan; @@ -104,13 +104,10 @@ impl PlanVTable for RowIdx { if mask.all_false() { return Ok(None); } - let offset = compiler.row_offset(); - let indices = Buffer::from_iter(rows.start + offset..rows.end + offset).into_array(); - let indices = if mask.all_true() { - indices - } else { - indices.filter(mask.clone())? - }; + // The sequence and its filter are lazy, so only the rows the mask keeps are produced. + let indices = idx_array(compiler.row_offset(), &rows) + .into_array() + .filter(mask.clone())?; Ok(Some(Chain::new(OnceSource::new(indices)))) } } diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index c86edcf4d30..6791a0003c2 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -630,6 +630,42 @@ fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { Ok(()) } +/// A projection of row ids and renamed columns, as an engine fetching rows by id plans it, +/// reads each column's plan directly beside the row indices. +#[test] +fn field_projection_with_row_idx_packs_column_plans() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), + Nullability::NonNullable, + ), + vec![flat(3, value_dtype.clone(), 0), flat(3, value_dtype, 1)], + ) + .into_layout(); + let projection = pack( + [ + ("id", row_idx()), + ("y", get_item("b", root())), + ("x", get_item("a", root())), + ], + Nullability::NonNullable, + ); + let plan = make_row_idx_plan(projection, make_plan(layout)?)?; + + let optimized = optimize(plan, &SESSION)?; + insta::assert_snapshot!(optimized.display_tree(), @" + root: vortex.plan.pack({id=u64, y=i32, x=i32}, rows=3) + id: vortex.plan.row_idx(u64, rows=3) + y: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) + x: vortex.plan.filter(i32, rows=3) + child: vortex.plan.segment_scan(i32, rows=3) + "); + Ok(()) +} + #[test] fn empty_projection_prunes_row_idx_child_fields() -> VortexResult<()> { let value_dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 382425d9bb0..0a94e18d7eb 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -68,6 +68,8 @@ pub struct ChunkPlanner { selection: Selection, /// The rows zone pruning kept, over `rows`, or `None` when nothing was pruned. kept: Option, + /// Whether the chunks run as children of this planner instead of being spawned. + inline: bool, } impl ChunkPlanner { @@ -79,15 +81,42 @@ impl ChunkPlanner { selection: Selection, kept: Option, ) -> Self { - let chunks = row_chunks(rows.clone(), scan.chunk_rows).into_iter(); + let chunks = row_chunks(rows.clone(), scan.chunk_rows); + Self::with_chunks(scan, rows, chunks, selection, kept) + } + + /// The planner of `chunks`, which tile `rows` in row order, for a driver that cuts the rows + /// itself, such as on the file's own chunk boundaries. + pub fn with_chunks( + scan: Arc, + rows: Range, + chunks: Vec>, + selection: Selection, + kept: Option, + ) -> Self { + debug_assert!( + chunks.first().is_none_or(|chunk| chunk.start == rows.start) + && chunks.last().is_none_or(|chunk| chunk.end == rows.end) + && chunks.windows(2).all(|pair| pair[0].end == pair[1].start), + "chunks must tile the rows in order" + ); Self { scan, rows, - chunks, + chunks: chunks.into_iter(), selection, kept, + inline: false, } } + + /// Runs the chunks as children of this planner, on its run and IO session, instead of + /// spawning them for the driver to place. Chunks that read the same segment, such as a + /// column chunk longer than the scan's chunks, then share the read. + pub fn inline(mut self) -> Self { + self.inline = true; + self + } } impl IoConsumer for ChunkPlanner { @@ -112,7 +141,11 @@ impl Planner for ChunkPlanner { }; if let Some((scope, planner)) = plan_selection(&self.scan, scope, &self.selection, kept) { - return Ok(PlannerOutput::Spawn(scope, planner)); + return Ok(if self.inline { + PlannerOutput::Planner(scope, planner) + } else { + PlannerOutput::Spawn(scope, planner) + }); } } Ok(PlannerOutput::Done) diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 0185f61d6a9..61a785c9568 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -71,6 +71,9 @@ impl SegmentLocation { } } +/// A selection of at most one row in this many of a scan's rows is sparse. +const SPARSE_SELECTION: u64 = 1024; + /// A scan's plans over one file, built once from the file's footer and read by every planner of /// the scan. Nothing here touches a segment, and no planner plans again. pub struct PreparedScan { @@ -118,12 +121,32 @@ impl PreparedScan { self.filter.as_deref() } + /// The most rows of one chunk a [`ChunkPlanner`] cuts the scan's rows into. + pub fn chunk_rows(&self) -> usize { + self.chunk_rows + } + + /// Whether the scan selects at most one row in 1024 of its rows, as an + /// engine fetching a few rows by id does. + pub fn sparse_selection(&self) -> bool { + !matches!(self.selection, Selection::All) + && self + .selection + .row_count_in_range(&self.rows) + .saturating_mul(SPARSE_SELECTION) + <= self.rows.end - self.rows.start + } + /// The service a driver takes IO sessions from: every read a planner of the scan requests is /// served from it by segment. + /// + /// A [sparse selection](Self::sparse_selection) retains its reads for the scan's lifetime, + /// as the chunks that select its few rows read the same wide segments one after another. pub fn io(&self) -> Arc { Arc::new(SegmentScanIo::new( Arc::clone(&self.segments), Arc::clone(&self.ranges), + self.sparse_selection(), )) } diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 0e690ebfde7..1b77b5f578d 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -12,6 +12,7 @@ use futures::StreamExt; use futures::TryFutureExt; use futures::future::BoxFuture; use futures::future::Shared; +use futures::future::WeakShared; use futures::stream::FuturesUnordered; use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; @@ -46,26 +47,89 @@ pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges { } /// Serves splits from any [`SegmentSource`], by segment id. +/// +/// Its sessions share a segment's read while any of them still holds it, so splits that run at +/// once and need the same segment, such as chunks that each cut into one column segment, read +/// it once. A scan that retains its reads keeps every read for its lifetime, so splits that run +/// one after another share them too. pub(crate) struct SegmentScanIo { segments: Arc, ranges: SegmentRanges, + reads: SharedReads, } impl SegmentScanIo { - pub(crate) fn new(segments: Arc, ranges: SegmentRanges) -> Self { - Self { segments, ranges } + /// A scan's IO, which retains its reads when `retain` is set. + /// + /// Retaining suits a sparse selection, such as an engine fetching a few rows by id: its + /// reads are few, but each wide column segment it touches is needed by every chunk that + /// selects a row of it. + pub(crate) fn new( + segments: Arc, + ranges: SegmentRanges, + retain: bool, + ) -> Self { + Self { + segments, + ranges, + reads: Arc::new(Mutex::new(ScanReads { + retain, + reads: HashMap::default(), + })), + } } } impl IoService for SegmentScanIo { fn session(&self) -> Arc { - Arc::new(SegmentIoSource::new( + Arc::new(SegmentIoSource::with_shared_reads( Arc::clone(&self.segments), Arc::clone(&self.ranges), + Arc::clone(&self.reads), )) } } +/// The reads of a scan's sessions, by segment. +type SharedReads = Arc>; + +/// The reads of a scan's sessions, by segment. +#[derive(Default)] +struct ScanReads { + /// Whether reads are held for the scan's lifetime rather than while a session holds them. + retain: bool, + reads: HashMap, +} + +/// A read of a scan, held weakly so that sharing it keeps no segment in memory longer, or +/// retained. +enum ScanRead { + Weak(WeakShared), + Retained(Prefetched), +} + +impl ScanReads { + /// The read of `segment` the scan still has. + fn get(&self, segment: SegmentId) -> Option { + match self.reads.get(&segment)? { + ScanRead::Weak(read) => read.upgrade(), + ScanRead::Retained(read) => Some(read.clone()), + } + } + + fn insert(&mut self, segment: SegmentId, read: &Prefetched) { + let read = if self.retain { + ScanRead::Retained(read.clone()) + } else { + let Some(read) = read.downgrade() else { + return; + }; + ScanRead::Weak(read) + }; + self.reads.insert(segment, read); + } +} + /// Serves a split's range requests from a segment source, without blocking. /// /// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once, @@ -85,13 +149,27 @@ pub(super) struct SegmentIoSource { prefetched: Mutex>, /// Announced segments, requested and never polled. announced: Mutex>>, + /// The reads of the scan's other sessions. + shared: SharedReads, } +/// The read of a segment. +type SegmentRead = BoxFuture<'static, Result>>; + /// A read started ahead of the fetch that needs it. -type Prefetched = Shared>>>; +type Prefetched = Shared; impl SegmentIoSource { + #[cfg(test)] pub(super) fn new(segments: Arc, ranges: SegmentRanges) -> Self { + Self::with_shared_reads(segments, ranges, SharedReads::default()) + } + + fn with_shared_reads( + segments: Arc, + ranges: SegmentRanges, + shared: SharedReads, + ) -> Self { Self { segments, ranges, @@ -99,8 +177,29 @@ impl SegmentIoSource { ready: Mutex::default(), prefetched: Mutex::default(), announced: Mutex::default(), + shared, } } + + /// The read of `segment`: one another session of the scan still holds, else the read of an + /// announcement of it, else a new one. + fn read(&self, segment: SegmentId) -> Prefetched { + let mut shared = self.shared.lock(); + if let Some(read) = shared.get(segment) { + return read; + } + let read = self + .announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared(); + shared.insert(segment, &read); + read + } } impl IoSource for SegmentIoSource { @@ -129,15 +228,7 @@ impl IoSource for SegmentIoSource { } if request.intent == IoIntent::Prefetch { prefetched.entry(segment).or_insert_with(|| { - let read = self - .announced - .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) - .map_err(Arc::new) - .boxed() - .shared(); + let read = self.read(segment); // Polling registers the read as wanted, so the source starts it now. drop( read.clone() @@ -152,16 +243,7 @@ impl IoSource for SegmentIoSource { // first has completed. let mut bytes = prefetched .entry(segment) - .or_insert_with(|| { - self.announced - .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) - .map_err(Arc::new) - .boxed() - .shared() - }) + .or_insert_with(|| self.read(segment)) .clone() .map_err(move |err| vortex_err!("read of segment {segment} failed: {err}")) .boxed(); @@ -249,10 +331,12 @@ mod tests { use vortex_io::request::IoOwnerId; use vortex_io::request::IoRequest; use vortex_io::request::IoRequestId; + use vortex_io::request::IoService; use vortex_io::request::IoSource; use vortex_io::request::IoTarget; use super::SegmentIoSource; + use super::SegmentScanIo; use super::segment_ranges; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentFuture; @@ -370,6 +454,75 @@ mod tests { Ok(()) } + /// Sessions of one scan share a segment's read while one of them holds it, and read it again + /// once none does. + #[tokio::test] + async fn sessions_share_a_held_read() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let scan = SegmentScanIo::new( + Arc::clone(&segments) as _, + segment_ranges(&[location]), + false, + ); + let first = scan.session(); + let second = scan.session(); + for (io, id) in [(&first, 0), (&second, 1)] { + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + assert_eq!(completion.request, IoRequestId(id)); + assert!(completion.result.is_ok()); + } + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + + first.clear(); + second.clear(); + let third = scan.session(); + third.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 2)])?; + assert!( + std::future::poll_fn(|cx| third.poll_completion(cx)) + .await? + .result + .is_ok() + ); + assert_eq!(segments.0.load(Ordering::Relaxed), 2); + Ok(()) + } + + /// A scan that retains its reads shares a read with a later session after the earlier one + /// let it go. + #[tokio::test] + async fn retained_reads_outlive_their_sessions() -> VortexResult<()> { + let segments = Arc::new(CountingSegments::default()); + let location = SegmentLocation { + offset: 0, + length: 4, + alignment: Alignment::none(), + }; + let scan = SegmentScanIo::new( + Arc::clone(&segments) as _, + segment_ranges(&[location]), + true, + ); + for id in 0..2 { + let io = scan.session(); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + assert!( + std::future::poll_fn(|cx| io.poll_completion(cx)) + .await? + .result + .is_ok() + ); + io.clear(); + } + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + Ok(()) + } + #[tokio::test] async fn announcements_keep_independent_consumers() -> VortexResult<()> { let segments = Arc::new(CountingSegments::default());