use try_collect to avoid a double collect

spiraldb · Sep 5, 2024 · 59417b3 · 59417b3
1 parent 2bfb65b
commit 59417b3
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/pyvortex/src/io.rs b/pyvortex/src/io.rs
@@ -1,6 +1,6 @@
 use std::path::Path;
 
-use futures::StreamExt;
+use futures::TryStreamExt;
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::pyfunction;
@@ -156,11 +156,7 @@ pub fn read<'py>(
 
         let dtype = stream.schema().into_dtype();
 
-        let vecs = stream
-            .collect::<Vec<VortexResult<Array>>>()
-            .await
-            .into_iter() // TODO(dk) unclear why I need two collects to pacify the compiler
-            .collect::<VortexResult<Vec<Array>>>()?;
+        let vecs: Vec<Array> = stream.try_collect().await?;
 
         if vecs.len() == 1 {
             Ok(vecs.into_iter().next().unwrap())