Skip to content

Commit 0736530

Browse files
committed
Move the scan IO and chunk planner changes to the DuckDB driver PR
They are only used by the DuckDB V2 driver, so they move to its own PR, leaving this one with the projection planning changes only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UeniBCK9ajvBLbjtkwq3cn Signed-off-by: Claude <noreply@anthropic.com>
1 parent afd9463 commit 0736530

3 files changed

Lines changed: 26 additions & 235 deletions

File tree

‎vortex-layout/src/scan/execution/chunk.rs‎

Lines changed: 3 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,6 @@ pub struct ChunkPlanner {
6868
selection: Selection,
6969
/// The rows zone pruning kept, over `rows`, or `None` when nothing was pruned.
7070
kept: Option<Mask>,
71-
/// Whether the chunks run as children of this planner instead of being spawned.
72-
inline: bool,
7371
}
7472

7573
impl ChunkPlanner {
@@ -81,42 +79,15 @@ impl ChunkPlanner {
8179
selection: Selection,
8280
kept: Option<Mask>,
8381
) -> Self {
84-
let chunks = row_chunks(rows.clone(), scan.chunk_rows);
85-
Self::with_chunks(scan, rows, chunks, selection, kept)
86-
}
87-
88-
/// The planner of `chunks`, which tile `rows` in row order, for a driver that cuts the rows
89-
/// itself, such as on the file's own chunk boundaries.
90-
pub fn with_chunks(
91-
scan: Arc<PreparedScan>,
92-
rows: Range<u64>,
93-
chunks: Vec<Range<u64>>,
94-
selection: Selection,
95-
kept: Option<Mask>,
96-
) -> Self {
97-
debug_assert!(
98-
chunks.first().is_none_or(|chunk| chunk.start == rows.start)
99-
&& chunks.last().is_none_or(|chunk| chunk.end == rows.end)
100-
&& chunks.windows(2).all(|pair| pair[0].end == pair[1].start),
101-
"chunks must tile the rows in order"
102-
);
82+
let chunks = row_chunks(rows.clone(), scan.chunk_rows).into_iter();
10383
Self {
10484
scan,
10585
rows,
106-
chunks: chunks.into_iter(),
86+
chunks,
10787
selection,
10888
kept,
109-
inline: false,
11089
}
11190
}
112-
113-
/// Runs the chunks as children of this planner, on its run and IO session, instead of
114-
/// spawning them for the driver to place. Chunks that read the same segment, such as a
115-
/// column chunk longer than the scan's chunks, then share the read.
116-
pub fn inline(mut self) -> Self {
117-
self.inline = true;
118-
self
119-
}
12091
}
12192

12293
impl IoConsumer for ChunkPlanner {
@@ -141,11 +112,7 @@ impl Planner for ChunkPlanner {
141112
};
142113
if let Some((scope, planner)) = plan_selection(&self.scan, scope, &self.selection, kept)
143114
{
144-
return Ok(if self.inline {
145-
PlannerOutput::Planner(scope, planner)
146-
} else {
147-
PlannerOutput::Spawn(scope, planner)
148-
});
115+
return Ok(PlannerOutput::Spawn(scope, planner));
149116
}
150117
}
151118
Ok(PlannerOutput::Done)

‎vortex-layout/src/scan/execution/mod.rs‎

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,6 @@ impl SegmentLocation {
7171
}
7272
}
7373

74-
/// A selection of at most one row in this many of a scan's rows is sparse.
75-
const SPARSE_SELECTION: u64 = 1024;
76-
7774
/// A scan's plans over one file, built once from the file's footer and read by every planner of
7875
/// the scan. Nothing here touches a segment, and no planner plans again.
7976
pub struct PreparedScan {
@@ -121,32 +118,12 @@ impl PreparedScan {
121118
self.filter.as_deref()
122119
}
123120

124-
/// The most rows of one chunk a [`ChunkPlanner`] cuts the scan's rows into.
125-
pub fn chunk_rows(&self) -> usize {
126-
self.chunk_rows
127-
}
128-
129-
/// Whether the scan selects at most one row in 1024 of its rows, as an
130-
/// engine fetching a few rows by id does.
131-
pub fn sparse_selection(&self) -> bool {
132-
!matches!(self.selection, Selection::All)
133-
&& self
134-
.selection
135-
.row_count_in_range(&self.rows)
136-
.saturating_mul(SPARSE_SELECTION)
137-
<= self.rows.end - self.rows.start
138-
}
139-
140121
/// The service a driver takes IO sessions from: every read a planner of the scan requests is
141122
/// served from it by segment.
142-
///
143-
/// A [sparse selection](Self::sparse_selection) retains its reads for the scan's lifetime,
144-
/// as the chunks that select its few rows read the same wide segments one after another.
145123
pub fn io(&self) -> Arc<dyn IoService> {
146124
Arc::new(SegmentScanIo::new(
147125
Arc::clone(&self.segments),
148126
Arc::clone(&self.ranges),
149-
self.sparse_selection(),
150127
))
151128
}
152129

‎vortex-layout/src/scan/v2/io.rs‎

Lines changed: 23 additions & 176 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ use futures::StreamExt;
1212
use futures::TryFutureExt;
1313
use futures::future::BoxFuture;
1414
use futures::future::Shared;
15-
use futures::future::WeakShared;
1615
use futures::stream::FuturesUnordered;
1716
use parking_lot::Mutex;
1817
use vortex_array::buffer::BufferHandle;
@@ -47,89 +46,26 @@ pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges {
4746
}
4847

4948
/// Serves splits from any [`SegmentSource`], by segment id.
50-
///
51-
/// Its sessions share a segment's read while any of them still holds it, so splits that run at
52-
/// once and need the same segment, such as chunks that each cut into one column segment, read
53-
/// it once. A scan that retains its reads keeps every read for its lifetime, so splits that run
54-
/// one after another share them too.
5549
pub(crate) struct SegmentScanIo {
5650
segments: Arc<dyn SegmentSource>,
5751
ranges: SegmentRanges,
58-
reads: SharedReads,
5952
}
6053

6154
impl SegmentScanIo {
62-
/// A scan's IO, which retains its reads when `retain` is set.
63-
///
64-
/// Retaining suits a sparse selection, such as an engine fetching a few rows by id: its
65-
/// reads are few, but each wide column segment it touches is needed by every chunk that
66-
/// selects a row of it.
67-
pub(crate) fn new(
68-
segments: Arc<dyn SegmentSource>,
69-
ranges: SegmentRanges,
70-
retain: bool,
71-
) -> Self {
72-
Self {
73-
segments,
74-
ranges,
75-
reads: Arc::new(Mutex::new(ScanReads {
76-
retain,
77-
reads: HashMap::default(),
78-
})),
79-
}
55+
pub(crate) fn new(segments: Arc<dyn SegmentSource>, ranges: SegmentRanges) -> Self {
56+
Self { segments, ranges }
8057
}
8158
}
8259

8360
impl IoService for SegmentScanIo {
8461
fn session(&self) -> Arc<dyn IoSource> {
85-
Arc::new(SegmentIoSource::with_shared_reads(
62+
Arc::new(SegmentIoSource::new(
8663
Arc::clone(&self.segments),
8764
Arc::clone(&self.ranges),
88-
Arc::clone(&self.reads),
8965
))
9066
}
9167
}
9268

93-
/// The reads of a scan's sessions, by segment.
94-
type SharedReads = Arc<Mutex<ScanReads>>;
95-
96-
/// The reads of a scan's sessions, by segment.
97-
#[derive(Default)]
98-
struct ScanReads {
99-
/// Whether reads are held for the scan's lifetime rather than while a session holds them.
100-
retain: bool,
101-
reads: HashMap<SegmentId, ScanRead>,
102-
}
103-
104-
/// A read of a scan, held weakly so that sharing it keeps no segment in memory longer, or
105-
/// retained.
106-
enum ScanRead {
107-
Weak(WeakShared<SegmentRead>),
108-
Retained(Prefetched),
109-
}
110-
111-
impl ScanReads {
112-
/// The read of `segment` the scan still has.
113-
fn get(&self, segment: SegmentId) -> Option<Prefetched> {
114-
match self.reads.get(&segment)? {
115-
ScanRead::Weak(read) => read.upgrade(),
116-
ScanRead::Retained(read) => Some(read.clone()),
117-
}
118-
}
119-
120-
fn insert(&mut self, segment: SegmentId, read: &Prefetched) {
121-
let read = if self.retain {
122-
ScanRead::Retained(read.clone())
123-
} else {
124-
let Some(read) = read.downgrade() else {
125-
return;
126-
};
127-
ScanRead::Weak(read)
128-
};
129-
self.reads.insert(segment, read);
130-
}
131-
}
132-
13369
/// Serves a split's range requests from a segment source, without blocking.
13470
///
13571
/// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once,
@@ -149,57 +85,22 @@ pub(super) struct SegmentIoSource {
14985
prefetched: Mutex<HashMap<SegmentId, Prefetched>>,
15086
/// Announced segments, requested and never polled.
15187
announced: Mutex<HashMap<SegmentId, VecDeque<SegmentFuture>>>,
152-
/// The reads of the scan's other sessions.
153-
shared: SharedReads,
15488
}
15589

156-
/// The read of a segment.
157-
type SegmentRead = BoxFuture<'static, Result<BufferHandle, Arc<VortexError>>>;
158-
15990
/// A read started ahead of the fetch that needs it.
160-
type Prefetched = Shared<SegmentRead>;
91+
type Prefetched = Shared<BoxFuture<'static, Result<BufferHandle, Arc<VortexError>>>>;
16192

16293
impl SegmentIoSource {
163-
#[cfg(test)]
16494
pub(super) fn new(segments: Arc<dyn SegmentSource>, ranges: SegmentRanges) -> Self {
165-
Self::with_shared_reads(segments, ranges, SharedReads::default())
166-
}
167-
168-
fn with_shared_reads(
169-
segments: Arc<dyn SegmentSource>,
170-
ranges: SegmentRanges,
171-
shared: SharedReads,
172-
) -> Self {
17395
Self {
17496
segments,
17597
ranges,
17698
reads: Mutex::new(FuturesUnordered::new()),
17799
ready: Mutex::default(),
178100
prefetched: Mutex::default(),
179101
announced: Mutex::default(),
180-
shared,
181102
}
182103
}
183-
184-
/// The read of `segment`: one another session of the scan still holds, else the read of an
185-
/// announcement of it, else a new one.
186-
fn read(&self, segment: SegmentId) -> Prefetched {
187-
let mut shared = self.shared.lock();
188-
if let Some(read) = shared.get(segment) {
189-
return read;
190-
}
191-
let read = self
192-
.announced
193-
.lock()
194-
.get_mut(&segment)
195-
.and_then(VecDeque::pop_front)
196-
.unwrap_or_else(|| self.segments.request(segment))
197-
.map_err(Arc::new)
198-
.boxed()
199-
.shared();
200-
shared.insert(segment, &read);
201-
read
202-
}
203104
}
204105

205106
impl IoSource for SegmentIoSource {
@@ -228,7 +129,15 @@ impl IoSource for SegmentIoSource {
228129
}
229130
if request.intent == IoIntent::Prefetch {
230131
prefetched.entry(segment).or_insert_with(|| {
231-
let read = self.read(segment);
132+
let read = self
133+
.announced
134+
.lock()
135+
.get_mut(&segment)
136+
.and_then(VecDeque::pop_front)
137+
.unwrap_or_else(|| self.segments.request(segment))
138+
.map_err(Arc::new)
139+
.boxed()
140+
.shared();
232141
// Polling registers the read as wanted, so the source starts it now.
233142
drop(
234143
read.clone()
@@ -243,7 +152,16 @@ impl IoSource for SegmentIoSource {
243152
// first has completed.
244153
let mut bytes = prefetched
245154
.entry(segment)
246-
.or_insert_with(|| self.read(segment))
155+
.or_insert_with(|| {
156+
self.announced
157+
.lock()
158+
.get_mut(&segment)
159+
.and_then(VecDeque::pop_front)
160+
.unwrap_or_else(|| self.segments.request(segment))
161+
.map_err(Arc::new)
162+
.boxed()
163+
.shared()
164+
})
247165
.clone()
248166
.map_err(move |err| vortex_err!("read of segment {segment} failed: {err}"))
249167
.boxed();
@@ -331,12 +249,10 @@ mod tests {
331249
use vortex_io::request::IoOwnerId;
332250
use vortex_io::request::IoRequest;
333251
use vortex_io::request::IoRequestId;
334-
use vortex_io::request::IoService;
335252
use vortex_io::request::IoSource;
336253
use vortex_io::request::IoTarget;
337254

338255
use super::SegmentIoSource;
339-
use super::SegmentScanIo;
340256
use super::segment_ranges;
341257
use crate::scan::execution::SegmentLocation;
342258
use crate::segments::SegmentFuture;
@@ -454,75 +370,6 @@ mod tests {
454370
Ok(())
455371
}
456372

457-
/// Sessions of one scan share a segment's read while one of them holds it, and read it again
458-
/// once none does.
459-
#[tokio::test]
460-
async fn sessions_share_a_held_read() -> VortexResult<()> {
461-
let segments = Arc::new(CountingSegments::default());
462-
let location = SegmentLocation {
463-
offset: 0,
464-
length: 4,
465-
alignment: Alignment::none(),
466-
};
467-
let scan = SegmentScanIo::new(
468-
Arc::clone(&segments) as _,
469-
segment_ranges(&[location]),
470-
false,
471-
);
472-
let first = scan.session();
473-
let second = scan.session();
474-
for (io, id) in [(&first, 0), (&second, 1)] {
475-
io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?;
476-
let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?;
477-
assert_eq!(completion.request, IoRequestId(id));
478-
assert!(completion.result.is_ok());
479-
}
480-
assert_eq!(segments.0.load(Ordering::Relaxed), 1);
481-
482-
first.clear();
483-
second.clear();
484-
let third = scan.session();
485-
third.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 2)])?;
486-
assert!(
487-
std::future::poll_fn(|cx| third.poll_completion(cx))
488-
.await?
489-
.result
490-
.is_ok()
491-
);
492-
assert_eq!(segments.0.load(Ordering::Relaxed), 2);
493-
Ok(())
494-
}
495-
496-
/// A scan that retains its reads shares a read with a later session after the earlier one
497-
/// let it go.
498-
#[tokio::test]
499-
async fn retained_reads_outlive_their_sessions() -> VortexResult<()> {
500-
let segments = Arc::new(CountingSegments::default());
501-
let location = SegmentLocation {
502-
offset: 0,
503-
length: 4,
504-
alignment: Alignment::none(),
505-
};
506-
let scan = SegmentScanIo::new(
507-
Arc::clone(&segments) as _,
508-
segment_ranges(&[location]),
509-
true,
510-
);
511-
for id in 0..2 {
512-
let io = scan.session();
513-
io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?;
514-
assert!(
515-
std::future::poll_fn(|cx| io.poll_completion(cx))
516-
.await?
517-
.result
518-
.is_ok()
519-
);
520-
io.clear();
521-
}
522-
assert_eq!(segments.0.load(Ordering::Relaxed), 1);
523-
Ok(())
524-
}
525-
526373
#[tokio::test]
527374
async fn announcements_keep_independent_consumers() -> VortexResult<()> {
528375
let segments = Arc::new(CountingSegments::default());

0 commit comments

Comments
 (0)