perf: optimize IO path for reading manifest

rust/lance-file/src/page_table.rs

@@ -215,7 +215,7 @@ mod tests {
             .unwrap();
         writer.shutdown().await.unwrap();
 
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         let actual = PageTable::load(
@@ -284,7 +284,7 @@ mod tests {
         let mut writer = tokio::fs::File::create(&path).await.unwrap();
         let res = page_table.write(&mut writer, 0).await.unwrap();
 
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
 

rust/lance-io/src/encodings/binary.rs

@@ -137,7 +137,7 @@ impl<'a, T: ByteArrayType> BinaryDecoder<'a, T> {
     /// use lance_io::{local::LocalObjectReader, encodings::binary::BinaryDecoder, traits::Reader};
     ///
     /// async {
-    ///     let reader = LocalObjectReader::open_local_path("/tmp/foo.lance", 2048).await.unwrap();
+    ///     let reader = LocalObjectReader::open_local_path("/tmp/foo.lance", 2048, None).await.unwrap();
     ///     let string_decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), 100, 1024, true);
     /// };
     /// ```
@@ -494,7 +494,7 @@ mod tests {
 
         let pos = write_test_data(&path, arrs).await.unwrap();
 
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         let read_len = arrs.iter().map(|a| a.len()).sum();
@@ -562,7 +562,7 @@ mod tests {
         let pos = encoder.encode(&[&data]).await.unwrap();
         object_writer.shutdown().await.unwrap();
 
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
@@ -605,7 +605,7 @@ mod tests {
         let path = temp_dir.path().join("foo");
 
         let pos = write_test_data(&path, &[&data]).await.unwrap();
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
@@ -627,7 +627,7 @@ mod tests {
         let temp_dir = tempfile::tempdir().unwrap();
         let path = temp_dir.path().join("foo");
         let pos = write_test_data(&path, &[&data]).await.unwrap();
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
@@ -658,7 +658,7 @@ mod tests {
         let path = temp_dir.path().join("foo");
         let pos = write_test_data(&path, &[&data]).await.unwrap();
 
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
@@ -738,7 +738,7 @@ mod tests {
             pos
         };
 
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         let decoder = BinaryDecoder::<BinaryType>::new(reader.as_ref(), pos, data.len(), true);

rust/lance-io/src/encodings/dictionary.rs

@@ -246,7 +246,7 @@ mod tests {
             object_writer.shutdown().await.unwrap();
         }
 
-        let reader = LocalObjectReader::open_local_path(&path, 2048)
+        let reader = LocalObjectReader::open_local_path(&path, 2048, None)
             .await
             .unwrap();
         let decoder = DictionaryDecoder::new(

rust/lance-io/src/encodings/plain.rs

@@ -581,7 +581,7 @@ mod tests {
             writer.flush().await.unwrap();
         }
 
-        let reader = LocalObjectReader::open_local_path(&path, 1024)
+        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
             .await
             .unwrap();
         assert!(reader.size().await.unwrap() > 0);
@@ -705,7 +705,7 @@ mod tests {
             writer.flush().await.unwrap();
         }
 
-        let reader = LocalObjectReader::open_local_path(&path, 2048)
+        let reader = LocalObjectReader::open_local_path(&path, 2048, None)
             .await
             .unwrap();
         assert!(reader.size().await.unwrap() > 0);
@@ -753,7 +753,7 @@ mod tests {
             writer.shutdown().await.unwrap();
         }
 
-        let reader = LocalObjectReader::open_local_path(&path, 2048)
+        let reader = LocalObjectReader::open_local_path(&path, 2048, None)
             .await
             .unwrap();
         assert!(reader.size().await.unwrap() > 0);

rust/lance-io/src/local.rs

@@ -20,6 +20,7 @@ use lance_core::{Error, Result};
 use object_store::path::Path;
 use snafu::{location, Location};
 use tokio::io::AsyncSeekExt;
+use tokio::sync::OnceCell;
 use tracing::instrument;
 
 use crate::traits::{Reader, Writer};
@@ -55,6 +56,10 @@ pub struct LocalObjectReader {
     /// Fie path.
     path: Path,
 
+    /// Known size of the file. This is either passed in on construction or
+    /// cached on the first metadata call.
+    size: OnceCell<usize>,
+
     /// Block size, in bytes.
     block_size: usize,
 }
@@ -63,23 +68,20 @@ impl LocalObjectReader {
     pub async fn open_local_path(
         path: impl AsRef<std::path::Path>,
         block_size: usize,
+        known_size: Option<usize>,
     ) -> Result<Box<dyn Reader>> {
         let path = path.as_ref().to_owned();
         let object_store_path = Path::from_filesystem_path(&path)?;
-        tokio::task::spawn_blocking(move || {
-            let local_file = File::open(&path)?;
-            Ok(Box::new(Self {
-                file: Arc::new(local_file),
-                path: object_store_path,
-                block_size,
-            }) as Box<dyn Reader>)
-        })
-        .await?
+        Self::open(&object_store_path, block_size, known_size).await
     }
 
     /// Open a local object reader, with default prefetch size.
     #[instrument(level = "debug")]
-    pub async fn open(path: &Path, block_size: usize) -> Result<Box<dyn Reader>> {
+    pub async fn open(
+        path: &Path,
+        block_size: usize,
+        known_size: Option<usize>,
+    ) -> Result<Box<dyn Reader>> {
         let path = path.clone();
         let local_path = to_local_path(&path);
         tokio::task::spawn_blocking(move || {
@@ -90,9 +92,11 @@ impl LocalObjectReader {
                 },
                 _ => e.into(),
             })?;
+            let size = OnceCell::new_with(known_size);
             Ok(Box::new(Self {
                 file: Arc::new(file),
                 block_size,
+                size,
                 path: path.clone(),
             }) as Box<dyn Reader>)
         })
@@ -111,13 +115,26 @@ impl Reader for LocalObjectReader {
     }
 
     /// Returns the file size.
-    async fn size(&self) -> Result<usize> {
-        Ok(self.file.metadata()?.len() as usize)
+    async fn size(&self) -> object_store::Result<usize> {
+        let file = self.file.clone();
+        self.size
+            .get_or_try_init(|| async move {
+                let metadata = tokio::task::spawn_blocking(move || {
+                    file.metadata().map_err(|err| object_store::Error::Generic {
+                        store: "LocalFileSystem",
+                        source: err.into(),
+                    })
+                })
+                .await??;
+                Ok(metadata.len() as usize)
+            })
+            .await
+            .cloned()
     }
 
     /// Reads a range of data.
     #[instrument(level = "debug", skip(self))]
-    async fn get_range(&self, range: Range<usize>) -> Result<Bytes> {
+    async fn get_range(&self, range: Range<usize>) -> object_store::Result<Bytes> {
         let file = self.file.clone();
         tokio::task::spawn_blocking(move || {
             let mut buf = BytesMut::with_capacity(range.len());
@@ -132,6 +149,10 @@ impl Reader for LocalObjectReader {
             Ok(buf.freeze())
         })
         .await?
+        .map_err(|err: std::io::Error| object_store::Error::Generic {
+            store: "LocalFileSystem",
+            source: err.into(),
+        })
     }
 }
 

rust/lance-io/src/object_reader.rs

@@ -9,6 +9,7 @@ use bytes::Bytes;
 use futures::future::BoxFuture;
 use lance_core::Result;
 use object_store::{path::Path, ObjectStore};
+use tokio::sync::OnceCell;
 use tracing::instrument;
 
 use crate::traits::Reader;
@@ -22,16 +23,24 @@ pub struct CloudObjectReader {
     pub object_store: Arc<dyn ObjectStore>,
     // File path
     pub path: Path,
+    // File size, if known.
+    size: OnceCell<usize>,
 
     block_size: usize,
 }
 
 impl CloudObjectReader {
     /// Create an ObjectReader from URI
-    pub fn new(object_store: Arc<dyn ObjectStore>, path: Path, block_size: usize) -> Result<Self> {
+    pub fn new(
+        object_store: Arc<dyn ObjectStore>,
+        path: Path,
+        block_size: usize,
+        known_size: Option<usize>,
+    ) -> Result<Self> {
         Ok(Self {
             object_store,
             path,
+            size: OnceCell::new_with(known_size),
             block_size,
         })
     }
@@ -42,14 +51,14 @@ impl CloudObjectReader {
     async fn do_with_retry<'a, O>(
         &self,
         f: impl Fn() -> BoxFuture<'a, std::result::Result<O, object_store::Error>>,
-    ) -> Result<O> {
+    ) -> object_store::Result<O> {
         let mut retries = 3;
         loop {
             match f().await {
                 Ok(val) => return Ok(val),
                 Err(err) => {
                     if retries == 0 {
-                        return Err(err.into());
+                        return Err(err);
                     }
                     retries -= 1;
                 }
@@ -69,15 +78,20 @@ impl Reader for CloudObjectReader {
     }
 
     /// Object/File Size.
-    async fn size(&self) -> Result<usize> {
-        let meta = self
-            .do_with_retry(|| self.object_store.head(&self.path))
-            .await?;
-        Ok(meta.size)
+    async fn size(&self) -> object_store::Result<usize> {
+        self.size
+            .get_or_try_init(|| async move {
+                let meta = self
+                    .do_with_retry(|| self.object_store.head(&self.path))
+                    .await?;
+                Ok(meta.size)
+            })
+            .await
+            .cloned()
     }
 
     #[instrument(level = "debug", skip(self))]
-    async fn get_range(&self, range: Range<usize>) -> Result<Bytes> {
+    async fn get_range(&self, range: Range<usize>) -> object_store::Result<Bytes> {
         self.do_with_retry(|| self.object_store.get_range(&self.path, range.clone()))
             .await
     }

rust/lance-io/src/object_store.rs

@@ -440,11 +440,29 @@ impl ObjectStore {
     /// - ``path``: Absolute path to the file.
     pub async fn open(&self, path: &Path) -> Result<Box<dyn Reader>> {
         match self.scheme.as_str() {
-            "file" => LocalObjectReader::open(path, self.block_size).await,
+            "file" => LocalObjectReader::open(path, self.block_size, None).await,
             _ => Ok(Box::new(CloudObjectReader::new(
                 self.inner.clone(),
                 path.clone(),
                 self.block_size,
+                None,
+            )?)),
+        }
+    }
+
+    /// Open a reader for a file with known size.
+    ///
+    /// This size may either have been retrieved from a list operation or
+    /// cached metadata. By passing in the known size, we can skip a HEAD / metadata
+    /// call.
+    pub async fn open_with_size(&self, path: &Path, known_size: usize) -> Result<Box<dyn Reader>> {
+        match self.scheme.as_str() {
+            "file" => LocalObjectReader::open(path, self.block_size, Some(known_size)).await,
+            _ => Ok(Box::new(CloudObjectReader::new(
+                self.inner.clone(),
+                path.clone(),
+                self.block_size,
+                Some(known_size),
             )?)),
         }
     }

rust/lance-io/src/scheduler.rs

@@ -91,7 +91,8 @@ impl IoTask {
         let bytes = self
             .reader
             .get_range(self.to_read.start as usize..self.to_read.end as usize)
-            .await;
+            .await
+            .map_err(Error::from);
         (self.when_done)(bytes);
     }
 }

rust/lance-io/src/traits.rs

@@ -86,10 +86,10 @@ pub trait Reader: std::fmt::Debug + Send + Sync {
     fn block_size(&self) -> usize;
 
     /// Object/File Size.
-    async fn size(&self) -> Result<usize>;
+    async fn size(&self) -> object_store::Result<usize>;
 
     /// Read a range of bytes from the object.
     ///
     /// TODO: change to read_at()?
-    async fn get_range(&self, range: Range<usize>) -> Result<Bytes>;
+    async fn get_range(&self, range: Range<usize>) -> object_store::Result<Bytes>;
 }

rust/lance-io/src/utils.rs

@@ -77,14 +77,16 @@ pub async fn read_fixed_stride_array(
 }
 
 /// Read a protobuf message at file position 'pos'.
-// TODO: pub(crate)
+///
+/// We write protobuf by first writing the length of the message as a u32,
+/// followed by the message itself.
 pub async fn read_message<M: Message + Default>(reader: &dyn Reader, pos: usize) -> Result<M> {
     let file_size = reader.size().await?;
     if pos > file_size {
         return Err(Error::io("file size is too small".to_string(), location!()));
     }
 
-    let range = pos..min(pos + 4096, file_size);
+    let range = pos..min(pos + reader.block_size(), file_size);
     let buf = reader.get_range(range.clone()).await?;
     let msg_len = LittleEndian::read_u32(&buf) as usize;
 
@@ -113,7 +115,7 @@ pub async fn read_struct<
     T::try_from(msg)
 }
 
-pub async fn read_last_block(reader: &dyn Reader) -> Result<Bytes> {
+pub async fn read_last_block(reader: &dyn Reader) -> object_store::Result<Bytes> {
     let file_size = reader.size().await?;
     let block_size = reader.block_size();
     let begin = if file_size < block_size {
@@ -225,7 +227,7 @@ mod tests {
         assert_eq!(pos, 0);
         object_writer.shutdown().await.unwrap();
 
-        let object_reader = CloudObjectReader::new(Arc::new(store), path, 1024).unwrap();
+        let object_reader = CloudObjectReader::new(Arc::new(store), path, 1024, None).unwrap();
         let actual: BytesWrapper = read_struct(&object_reader, pos).await.unwrap();
         assert_eq!(some_message, actual);
     }