From 3fe1afef93eeed91af58d92b61475529f68c570c Mon Sep 17 00:00:00 2001
From: cancaicai <2356672992@qq.com>
Date: Sun, 8 Feb 2026 12:36:01 +0800
Subject: [PATCH] chore: optimize next_coalesced() by tracking removed ids in a
 HashSet

Signed-off-by: cancaicai <2356672992@qq.com>
---
 vortex-file/src/read/driver.rs | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs
index b2d7e6cc890..3ff21000cc9 100644
--- a/vortex-file/src/read/driver.rs
+++ b/vortex-file/src/read/driver.rs
@@ -3,6 +3,7 @@
 
 use std::collections::BTreeMap;
 use std::collections::BTreeSet;
+use std::collections::HashSet;
 use std::pin::Pin;
 use std::task::Context;
 use std::task::Poll;
@@ -209,7 +210,10 @@ impl State {
         let mut current_end = requests[0].offset + requests[0].length as u64;
         let align = *self.coalesced_buffer_alignment as u64;
 
-        let mut keys_to_remove = Vec::new();
+        // Track requests that we've already decided to remove (or that were cancelled) so that
+        // we don't repeatedly process them during range scans.
+        let mut keys_to_remove: Vec<(u64, RequestId)> = Vec::new();
+        let mut ids_to_remove: HashSet<RequestId> = HashSet::new();
         let mut found_new_requests = true;
 
         // Keep expanding the window while we can find new requests within constraints
@@ -225,8 +229,8 @@ impl State {
                 .requests_by_offset
                 .range((scan_start, RequestId::MIN)..=(scan_end, RequestId::MAX))
             {
-                // Skip if we've already marked this request for removal
-                if keys_to_remove.iter().any(|&(_, id)| id == req_id) {
+                // Skip if we've already marked this request for removal.
+                if ids_to_remove.contains(&req_id) {
                     continue;
                 }
 
@@ -236,13 +240,15 @@ impl State {
                     .or_else(|| self.requests.get(&req_id))
                     .vortex_expect("Missing request in requests_by_offset");
 
-                // Skip any cancelled requests
+                // Skip any cancelled requests.
                 if req.callback.is_closed() {
-                    keys_to_remove.push((req_offset, req_id));
+                    if ids_to_remove.insert(req_id) {
+                        keys_to_remove.push((req_offset, req_id));
+                    }
                     continue;
                 }
 
-                // Check if this request is within coalescing distance of our current range
+                // Check if this request is within coalescing distance of our current range.
                 let req_end = req_offset + req.length as u64;
                 if (req_offset <= current_end + window.distance && req_end >= current_start)
                     || (req_end + window.distance >= current_start && req_offset <= current_end)
@@ -267,7 +273,9 @@ impl State {
                         .vortex_expect("Missing request in requests_by_offset");
 
                     requests.push(req);
-                    keys_to_remove.push((req_offset, req_id));
+                    if ids_to_remove.insert(req_id) {
+                        keys_to_remove.push((req_offset, req_id));
+                    }
                     found_new_requests = true;
                 }
             }