Use FuturesUnordered when order of try_join_all is not relevant (#2026)

Author: robert3005

Cargo.toml

@@ -59,6 +59,7 @@ arrow-ord = "53.0.0"
 arrow-schema = "53.0.0"
 arrow-select = "53.0.0"
 arrow-string = "53.0.0"
+async-once-cell = "0.5.4"
 async-trait = "0.1"
 backtrace = "0.3.74"
 bytes = "1.6.0"
@@ -107,7 +108,6 @@ parquet = "53.0.0"
 paste = "1.0.14"
 pin-project = "1.1.5"
 pin-project-lite = "0.2.15"
-tabled = { version = "0.17.0", default-features = false }
 prost = "0.13.0"
 prost-build = "0.13.0"
 prost-types = "0.13.0"
@@ -123,6 +123,7 @@ serde_json = "1.0.116"
 serde_test = "1.0.176"
 simplelog = { version = "0.12.2", features = ["paris"] }
 static_assertions = "1"
+tabled = { version = "0.17.0", default-features = false }
 tar = "0.4"
 tempfile = "3"
 thiserror = "2.0.0"

vortex-file/src/io/file.rs

@@ -4,8 +4,8 @@ use std::sync::Arc;
 
 use futures::channel::oneshot;
 use futures::Stream;
-use futures_util::future::try_join_all;
-use futures_util::{stream, StreamExt};
+use futures_util::stream::FuturesUnordered;
+use futures_util::{stream, StreamExt, TryStreamExt};
 use vortex_buffer::ByteBuffer;
 use vortex_error::{vortex_err, vortex_panic, VortexExpect, VortexResult};
 use vortex_io::VortexReadAt;
@@ -213,7 +213,9 @@ async fn evaluate<R: VortexReadAt>(
     }
 
     // Populate the cache
-    try_join_all(cache_futures).await?;
+    FuturesUnordered::from_iter(cache_futures)
+        .try_collect::<()>()
+        .await?;
 
     Ok(())
 }

vortex-file/src/open/mod.rs

@@ -5,6 +5,8 @@ use std::sync::Arc;
 
 pub use exec::*;
 use flatbuffers::root;
+use futures_util::stream::FuturesUnordered;
+use futures_util::{stream, StreamExt, TryStreamExt};
 use itertools::Itertools;
 pub use split_by::*;
 use vortex_array::ContextRef;
@@ -331,17 +333,25 @@ impl VortexOpenOptions {
         file_layout: &FileLayout,
         segments: &dyn SegmentCache,
     ) -> VortexResult<()> {
-        for (idx, segment) in file_layout.segment_map().iter().enumerate() {
-            if segment.offset < initial_offset {
-                // Skip segments that aren't in the initial read.
-                continue;
-            }
-            let segment_id = SegmentId::from(u32::try_from(idx)?);
-            let offset = usize::try_from(segment.offset - initial_offset)?;
-            let buffer = initial_read.slice(offset..offset + (segment.length as usize));
-
-            segments.put(segment_id, buffer).await?;
-        }
-        Ok(())
+        stream::iter(
+            file_layout
+                .segment_map()
+                .iter()
+                .enumerate()
+                .filter(|(_, segment)| segment.offset > initial_offset)
+                .map(|(idx, segment)| async move {
+                    let segment_id = SegmentId::from(u32::try_from(idx)?);
+                    let offset = usize::try_from(segment.offset - initial_offset)?;
+                    let buffer = initial_read
+                        .slice(offset..offset + (segment.length as usize))
+                        .aligned(segment.alignment);
+
+                    segments.put(segment_id, buffer).await
+                }),
+        )
+        .collect::<FuturesUnordered<_>>()
+        .await
+        .try_collect::<()>()
+        .await
     }
 }

vortex-layout/Cargo.toml

@@ -15,7 +15,7 @@ categories.workspace = true
 
 [dependencies]
 arrow-buffer = { workspace = true }
-async-once-cell = "0.5.4"
+async-once-cell = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 flatbuffers = { workspace = true }