Skip to content

Commit 5deafb8

Browse files
committed
chore:optimize next_coalesced() by tracking removed ids in a HashSet
Signed-off-by: cancaicai <2356672992@qq.com>
1 parent a11336d commit 5deafb8

1 file changed

Lines changed: 15 additions & 7 deletions

File tree

vortex-file/src/read/driver.rs

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
use std::collections::BTreeMap;
55
use std::collections::BTreeSet;
6+
use std::collections::HashSet;
67
use std::pin::Pin;
78
use std::task::Context;
89
use std::task::Poll;
@@ -209,7 +210,10 @@ impl State {
209210
let mut current_end = requests[0].offset + requests[0].length as u64;
210211
let align = *self.coalesced_buffer_alignment as u64;
211212

212-
let mut keys_to_remove = Vec::new();
213+
// Track requests that we've already decided to remove (or that were cancelled) so that
214+
// we don't repeatedly process them during range scans.
215+
let mut keys_to_remove: Vec<(u64, RequestId)> = Vec::new();
216+
let mut ids_to_remove: HashSet<RequestId> = HashSet::new();
213217
let mut found_new_requests = true;
214218

215219
// Keep expanding the window while we can find new requests within constraints
@@ -225,8 +229,8 @@ impl State {
225229
.requests_by_offset
226230
.range((scan_start, RequestId::MIN)..=(scan_end, RequestId::MAX))
227231
{
228-
// Skip if we've already marked this request for removal
229-
if keys_to_remove.iter().any(|&(_, id)| id == req_id) {
232+
// Skip if we've already marked this request for removal.
233+
if ids_to_remove.contains(&req_id) {
230234
continue;
231235
}
232236

@@ -236,13 +240,15 @@ impl State {
236240
.or_else(|| self.requests.get(&req_id))
237241
.vortex_expect("Missing request in requests_by_offset");
238242

239-
// Skip any cancelled requests
243+
// Skip any cancelled requests.
240244
if req.callback.is_closed() {
241-
keys_to_remove.push((req_offset, req_id));
245+
if ids_to_remove.insert(req_id) {
246+
keys_to_remove.push((req_offset, req_id));
247+
}
242248
continue;
243249
}
244250

245-
// Check if this request is within coalescing distance of our current range
251+
// Check if this request is within coalescing distance of our current range.
246252
let req_end = req_offset + req.length as u64;
247253
if (req_offset <= current_end + window.distance && req_end >= current_start)
248254
|| (req_end + window.distance >= current_start && req_offset <= current_end)
@@ -267,7 +273,9 @@ impl State {
267273
.vortex_expect("Missing request in requests_by_offset");
268274

269275
requests.push(req);
270-
keys_to_remove.push((req_offset, req_id));
276+
if ids_to_remove.insert(req_id) {
277+
keys_to_remove.push((req_offset, req_id));
278+
}
271279
found_new_requests = true;
272280
}
273281
}

0 commit comments

Comments
 (0)