[arrow-rs] branch master updated: Fix panix in coalesce (#2554)

[tu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 89f2e4502 Fix panix in coalesce (#2554)
89f2e4502 is described below

commit 89f2e45023fef15a6c37406c1e33bb8a2eca1084
Author: Dan Harris <13...@users.noreply.github.com>
AuthorDate: Tue Aug 23 05:34:41 2022 -0400

    Fix panix in coalesce (#2554)
---
 object_store/src/util.rs | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/object_store/src/util.rs b/object_store/src/util.rs
index 46e9e9ed8..f548ed48a 100644
--- a/object_store/src/util.rs
+++ b/object_store/src/util.rs
@@ -87,11 +87,17 @@ where
     F: Send + FnMut(std::ops::Range<usize>) -> Fut,
     Fut: std::future::Future<Output = Result<Bytes>> + Send,
 {
+    if ranges.is_empty() {
+        return Ok(vec![]);
+    }
+
     let mut ret = Vec::with_capacity(ranges.len());
     let mut start_idx = 0;
     let mut end_idx = 1;
 
     while start_idx != ranges.len() {
+        let mut range_end = ranges[start_idx].end;
+
         while end_idx != ranges.len()
             && ranges[end_idx]
                 .start
@@ -99,12 +105,14 @@ where
                 .map(|delta| delta <= coalesce)
                 .unwrap_or(false)
         {
+            if ranges[end_idx].end > range_end {
+                range_end = ranges[end_idx].end;
+            }
             end_idx += 1;
         }
 
         let start = ranges[start_idx].start;
-        let end = ranges[end_idx - 1].end;
-        let bytes = fetch(start..end).await?;
+        let bytes = fetch(start..range_end).await?;
         for range in ranges.iter().take(end_idx).skip(start_idx) {
             ret.push(bytes.slice(range.start - start..range.end - start))
         }
@@ -164,5 +172,11 @@ mod tests {
 
         let fetches = do_fetch(vec![0..1, 5..6, 7..9, 2..3, 4..6], 1).await;
         assert_eq!(fetches, vec![0..1, 5..9, 2..6]);
+
+        let fetches = do_fetch(vec![0..1, 5..6, 7..9, 2..3, 4..6], 1).await;
+        assert_eq!(fetches, vec![0..1, 5..9, 2..6]);
+
+        let fetches = do_fetch(vec![0..1, 6..7, 8..9, 10..14, 9..10], 4).await;
+        assert_eq!(fetches, vec![0..1, 6..14]);
     }
 }