Make the VortexReadAt::size method return an io::Result (#1471)

Author: joseph-isaacs

pyvortex/src/dataset.rs

@@ -61,7 +61,7 @@ pub async fn read_array_from_reader<T: VortexReadAt + Unpin + 'static>(
 }
 
 pub async fn read_dtype_from_reader<T: VortexReadAt + Unpin>(reader: T) -> VortexResult<DType> {
-    let initial_read = read_initial_bytes(&reader, reader.size().await).await?;
+    let initial_read = read_initial_bytes(&reader, reader.size().await?).await?;
     initial_read.lazy_dtype()?.value().cloned()
 }
 

vortex-file/src/read/builder/mod.rs

@@ -116,7 +116,7 @@ impl<R: VortexReadAt> VortexReadBuilder<R> {
 
     pub async fn build(self) -> VortexResult<VortexFileArrayStream<R>> {
         // we do a large enough initial read to get footer, layout, and schema
-        let initial_read = read_initial_bytes(&self.read_at, self.size().await).await?;
+        let initial_read = read_initial_bytes(&self.read_at, self.size().await?).await?;
 
         let layout = initial_read.fb_layout()?;
 
@@ -181,10 +181,10 @@ impl<R: VortexReadAt> VortexReadBuilder<R> {
         ))
     }
 
-    async fn size(&self) -> u64 {
-        match self.size {
+    async fn size(&self) -> VortexResult<u64> {
+        Ok(match self.size {
             Some(s) => s,
-            None => self.read_at.size().await,
-        }
+            None => self.read_at.size().await?,
+        })
     }
 }

vortex-io/src/compio.rs

@@ -6,7 +6,6 @@ use compio::buf::{IoBuf, IoBufMut, SetBufInit};
 use compio::fs::File;
 use compio::io::AsyncReadAtExt;
 use compio::BufResult;
-use vortex_error::vortex_panic;
 
 use crate::aligned::{AlignedBytesMut, PowerOfTwo};
 use crate::{VortexReadAt, ALIGNMENT};
@@ -67,14 +66,9 @@ impl VortexReadAt for File {
         }
     }
 
-    fn size(&self) -> impl Future<Output = u64> + 'static {
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
         let this = self.clone();
-        async move {
-            this.metadata()
-                .await
-                .map(|metadata| metadata.len())
-                .unwrap_or_else(|e| vortex_panic!("compio File::size: {e}"))
-        }
+        async move { this.metadata().await.map(|metadata| metadata.len()) }
     }
 }
 

vortex-io/src/object_store.rs

@@ -9,7 +9,7 @@ use object_store::path::Path;
 use object_store::{GetOptions, GetRange, ObjectStore, WriteMultipart};
 use vortex_buffer::io_buf::IoBuf;
 use vortex_buffer::Buffer;
-use vortex_error::{vortex_panic, VortexError, VortexResult};
+use vortex_error::VortexResult;
 
 use crate::aligned::AlignedBytesMut;
 use crate::{VortexBufReader, VortexReadAt, VortexWrite, ALIGNMENT};
@@ -103,19 +103,16 @@ impl VortexReadAt for ObjectStoreReadAt {
     }
 
     #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
-    fn size(&self) -> impl Future<Output = u64> + 'static {
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
         let object_store = self.object_store.clone();
         let location = self.location.clone();
 
         Box::pin(async move {
             object_store
                 .head(&location)
                 .await
-                .map_err(VortexError::ObjectStore)
-                .unwrap_or_else(|err| {
-                    vortex_panic!(err, "Failed to get size of object at location {}", location)
-                })
-                .size as u64
+                .map(|obj| obj.size as u64)
+                .map_err(io::Error::other)
         })
     }
 }

vortex-io/src/offset.rs

@@ -1,4 +1,5 @@
 use std::future::Future;
+use std::io;
 
 use bytes::Bytes;
 use futures::FutureExt;
@@ -34,16 +35,16 @@ impl<R: VortexReadAt> VortexReadAt for OffsetReadAt<R> {
         &self,
         pos: u64,
         len: u64,
-    ) -> impl Future<Output = std::io::Result<Bytes>> + 'static {
+    ) -> impl Future<Output = io::Result<Bytes>> + 'static {
         self.read.read_byte_range(pos + self.offset, len)
     }
 
     fn performance_hint(&self) -> usize {
         self.read.performance_hint()
     }
 
-    fn size(&self) -> impl Future<Output = u64> + 'static {
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
         let offset = self.offset;
-        self.read.size().map(move |len| len - offset)
+        self.read.size().map(move |len| len.map(|len| len - offset))
     }
 }

vortex-io/src/read.rs

@@ -38,7 +38,7 @@ pub trait VortexReadAt: Send + Sync + Clone + 'static {
     ///
     /// For a file it will be the size in bytes, for an object in an
     /// `ObjectStore` it will be the `ObjectMeta::size`.
-    fn size(&self) -> impl Future<Output = u64> + 'static;
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static;
 }
 
 impl<T: VortexReadAt> VortexReadAt for Arc<T> {
@@ -54,7 +54,7 @@ impl<T: VortexReadAt> VortexReadAt for Arc<T> {
         T::performance_hint(self)
     }
 
-    fn size(&self) -> impl Future<Output = u64> + 'static {
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
         T::size(self)
     }
 }
@@ -80,8 +80,8 @@ impl VortexReadAt for Buffer {
         }
     }
 
-    fn size(&self) -> impl Future<Output = u64> + 'static {
-        future::ready(self.len() as u64)
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
+        future::ready(Ok(self.len() as u64))
     }
 }
 
@@ -102,8 +102,8 @@ impl VortexReadAt for Bytes {
         }
     }
 
-    fn size(&self) -> impl Future<Output = u64> + 'static {
-        let len = self.len() as u64;
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
+        let len = Ok(self.len() as u64);
         future::ready(len)
     }
 }

vortex-io/src/tokio.rs

@@ -9,7 +9,6 @@ use std::sync::Arc;
 use bytes::Bytes;
 use tokio::io::{AsyncWrite, AsyncWriteExt};
 use vortex_buffer::io_buf::IoBuf;
-use vortex_error::vortex_panic;
 
 use crate::aligned::AlignedBytesMut;
 use crate::{VortexReadAt, VortexWrite, ALIGNMENT};
@@ -81,19 +80,10 @@ impl VortexReadAt for TokioFile {
     }
 
     #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
-    fn size(&self) -> impl Future<Output = u64> + 'static {
+    fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
         let this = self.clone();
 
-        async move {
-            let res = tokio::task::spawn_blocking(move || {
-                this.metadata()
-                    .unwrap_or_else(|e| vortex_panic!("access TokioFile metadata: {e}"))
-                    .len()
-            })
-            .await;
-
-            res.unwrap_or_else(|e| vortex_panic!("Joining spawn_blocking: size: {e}"))
-        }
+        async move { this.metadata().map(|metadata| metadata.len()) }
     }
 }