Parquet/async: Swap default MetadataLoader::load behaviour to suffix requests, swap get_bytes to take a GetRange (equivalent to object-store's)

Author: H-Plus-Time

parquet/src/arrow/async_reader/metadata.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::arrow::async_reader::AsyncFileReader;
+use crate::arrow::async_reader::{AsyncFileReader, GetRange};
 use crate::errors::{ParquetError, Result};
 use crate::file::footer::{decode_footer, decode_metadata};
 use crate::file::metadata::ParquetMetaData;
@@ -25,15 +25,14 @@ use bytes::Bytes;
 use futures::future::BoxFuture;
 use futures::FutureExt;
 use std::future::Future;
-use std::ops::Range;
 
 /// A data source that can be used with [`MetadataLoader`] to load [`ParquetMetaData`]
 pub trait MetadataFetch {
-    fn fetch(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>>;
+    fn fetch(&mut self, range: GetRange) -> BoxFuture<'_, Result<Bytes>>;
 }
 
 impl<'a, T: AsyncFileReader> MetadataFetch for &'a mut T {
-    fn fetch(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
+    fn fetch(&mut self, range: GetRange) -> BoxFuture<'_, Result<Bytes>> {
         self.get_bytes(range)
     }
 }
@@ -52,49 +51,30 @@ impl<F: MetadataFetch> MetadataLoader<F> {
     /// Create a new [`MetadataLoader`] by reading the footer information
     ///
     /// See [`fetch_parquet_metadata`] for the meaning of the individual parameters
-    pub async fn load(mut fetch: F, file_size: usize, prefetch: Option<usize>) -> Result<Self> {
-        if file_size < 8 {
-            return Err(ParquetError::EOF(format!(
-                "file size of {file_size} is less than footer"
-            )));
-        }
+    pub async fn load(mut fetch: F, prefetch: Option<usize>) -> Result<Self> {
 
-        // If a size hint is provided, read more than the minimum size
-        // to try and avoid a second fetch.
-        let footer_start = if let Some(size_hint) = prefetch {
-            file_size.saturating_sub(size_hint)
-        } else {
-            file_size - 8
-        };
-
-        let suffix = fetch.fetch(footer_start..file_size).await?;
+        let suffix = fetch.fetch(GetRange::Suffix(prefetch.unwrap_or(8))).await?;
         let suffix_len = suffix.len();
 
         let mut footer = [0; 8];
         footer.copy_from_slice(&suffix[suffix_len - 8..suffix_len]);
 
         let length = decode_footer(&footer)?;
 
-        if file_size < length + 8 {
-            return Err(ParquetError::EOF(format!(
-                "file size of {} is less than footer + metadata {}",
-                file_size,
-                length + 8
-            )));
-        }
-
         // Did not fetch the entire file metadata in the initial read, need to make a second request
         let (metadata, remainder) = if length > suffix_len - 8 {
-            let metadata_start = file_size - length - 8;
-            let meta = fetch.fetch(metadata_start..file_size - 8).await?;
-            (decode_metadata(&meta)?, None)
+            let metadata_offset = length + 8;
+            let meta = fetch.fetch(GetRange::Suffix(metadata_offset)).await?;
+            let slice = &meta[0..length];
+            (decode_metadata(&slice)?, None)
         } else {
-            let metadata_start = file_size - length - 8 - footer_start;
+            let metadata_offset = length + 8;
+            let metadata_start = suffix_len - metadata_offset;
 
             let slice = &suffix[metadata_start..suffix_len - 8];
             (
                 decode_metadata(slice)?,
-                Some((footer_start, suffix.slice(..metadata_start))),
+                Some((0, suffix.slice(..metadata_start))),
             )
         };
 
@@ -105,6 +85,10 @@ impl<F: MetadataFetch> MetadataLoader<F> {
         })
     }
 
+    pub async fn load_absolute(mut fetch: F, file_size: usize, prefetch: Option<usize>) -> Result<Self> {
+        todo!()
+    }
+
     /// Create a new [`MetadataLoader`] from an existing [`ParquetMetaData`]
     pub fn new(fetch: F, metadata: ParquetMetaData) -> Self {
         Self {
@@ -133,13 +117,15 @@ impl<F: MetadataFetch> MetadataLoader<F> {
             Some(range) => range,
         };
 
+        let page_index_len = range.end - range.start;
+        // TODO: determine if _remainder_start is needed even in the non-suffix request case
         let data = match &self.remainder {
-            Some((remainder_start, remainder)) if *remainder_start <= range.start => {
-                let offset = range.start - *remainder_start;
-                remainder.slice(offset..range.end - *remainder_start + offset)
+            Some((_remainder_start, remainder)) if remainder.len() >= page_index_len => {
+                let offset = remainder.len() - page_index_len;
+                remainder.slice(offset..)
             }
             // Note: this will potentially fetch data already in remainder, this keeps things simple
-            _ => self.fetch.fetch(range.start..range.end).await?,
+            _ => self.fetch.fetch((range.start..range.end).into()).await?,
         };
 
         // Sanity check
@@ -200,10 +186,10 @@ struct MetadataFetchFn<F>(F);
 
 impl<F, Fut> MetadataFetch for MetadataFetchFn<F>
 where
-    F: FnMut(Range<usize>) -> Fut + Send,
+    F: FnMut(GetRange) -> Fut + Send,
     Fut: Future<Output = Result<Bytes>> + Send,
 {
-    fn fetch(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
+    fn fetch(&mut self, range: GetRange) -> BoxFuture<'_, Result<Bytes>> {
         async move { self.0(range).await }.boxed()
     }
 }
@@ -226,15 +212,18 @@ where
 /// significantly reduce the number of `fetch` requests, and consequently latency
 pub async fn fetch_parquet_metadata<F, Fut>(
     fetch: F,
-    file_size: usize,
+    file_size: Option<usize>,
     prefetch: Option<usize>,
 ) -> Result<ParquetMetaData>
 where
-    F: FnMut(Range<usize>) -> Fut + Send,
+    F: FnMut(GetRange) -> Fut + Send,
     Fut: Future<Output = Result<Bytes>> + Send,
 {
     let fetch = MetadataFetchFn(fetch);
-    let loader = MetadataLoader::load(fetch, file_size, prefetch).await?;
+    let loader = match file_size {
+        Some(file_size) => MetadataLoader::load_absolute(fetch, file_size, prefetch).await?,
+        None => MetadataLoader::load(fetch, prefetch).await?
+    };
     Ok(loader.finish())
 }
 
@@ -247,7 +236,13 @@ mod tests {
     use std::io::{Read, Seek, SeekFrom};
     use std::sync::atomic::{AtomicUsize, Ordering};
 
-    fn read_range(file: &mut File, range: Range<usize>) -> Result<Bytes> {
+    fn read_range(file: &mut File, range: GetRange) -> Result<Bytes> {
+        let file_size = file.len().try_into().unwrap();
+        let range = match range {
+            GetRange::Bounded(range) => range,
+            GetRange::Offset(offset) => offset..file_size,
+            GetRange::Suffix(end_offset) => (file_size.saturating_sub(end_offset.try_into().unwrap())..file_size)
+        };
         file.seek(SeekFrom::Start(range.start as _))?;
         let len = range.end - range.start;
         let mut buf = Vec::with_capacity(len);
@@ -269,41 +264,41 @@ mod tests {
             futures::future::ready(read_range(&mut file, range))
         };
 
-        let actual = fetch_parquet_metadata(&mut fetch, len, None).await.unwrap();
+        let actual = fetch_parquet_metadata(&mut fetch, None, None).await.unwrap();
         assert_eq!(actual.file_metadata().schema(), expected);
         assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
 
         // Metadata hint too small
         fetch_count.store(0, Ordering::SeqCst);
-        let actual = fetch_parquet_metadata(&mut fetch, len, Some(10))
+        let actual = fetch_parquet_metadata(&mut fetch, None, Some(10))
             .await
             .unwrap();
         assert_eq!(actual.file_metadata().schema(), expected);
         assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
 
         // Metadata hint too large
         fetch_count.store(0, Ordering::SeqCst);
-        let actual = fetch_parquet_metadata(&mut fetch, len, Some(500))
+        let actual = fetch_parquet_metadata(&mut fetch, None, Some(500))
             .await
             .unwrap();
         assert_eq!(actual.file_metadata().schema(), expected);
         assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
 
         // Metadata hint exactly correct
         fetch_count.store(0, Ordering::SeqCst);
-        let actual = fetch_parquet_metadata(&mut fetch, len, Some(428))
+        let actual = fetch_parquet_metadata(&mut fetch, None, Some(428))
             .await
             .unwrap();
         assert_eq!(actual.file_metadata().schema(), expected);
         assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
 
-        let err = fetch_parquet_metadata(&mut fetch, 4, None)
+        let err = fetch_parquet_metadata(&mut fetch, Some(4), None)
             .await
             .unwrap_err()
             .to_string();
         assert_eq!(err, "EOF: file size of 4 is less than footer");
 
-        let err = fetch_parquet_metadata(&mut fetch, 20, None)
+        let err = fetch_parquet_metadata(&mut fetch, Some(20), None)
             .await
             .unwrap_err()
             .to_string();
@@ -321,7 +316,7 @@ mod tests {
         };
 
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, None).await.unwrap();
+        let mut loader = MetadataLoader::load(f, None).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
         loader.load_page_index(true, true).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 3);
@@ -331,7 +326,7 @@ mod tests {
         // Prefetch just footer exactly
         fetch_count.store(0, Ordering::SeqCst);
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, Some(1729)).await.unwrap();
+        let mut loader = MetadataLoader::load(f, Some(1729)).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
         loader.load_page_index(true, true).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
@@ -341,7 +336,7 @@ mod tests {
         // Prefetch more than footer but not enough
         fetch_count.store(0, Ordering::SeqCst);
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, Some(130649)).await.unwrap();
+        let mut loader = MetadataLoader::load(f, Some(130649)).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
         loader.load_page_index(true, true).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
@@ -351,7 +346,17 @@ mod tests {
         // Prefetch exactly enough
         fetch_count.store(0, Ordering::SeqCst);
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, Some(130650)).await.unwrap();
+        let mut loader = MetadataLoader::load(f, Some(130650)).await.unwrap();
+        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
+        loader.load_page_index(true, true).await.unwrap();
+        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
+        let metadata = loader.finish();
+        assert!(metadata.offset_index().is_some() && metadata.column_index().is_some());
+        
+        // Prefetch more than enough
+        fetch_count.store(0, Ordering::SeqCst);
+        let f = MetadataFetchFn(&mut fetch);
+        let mut loader = MetadataLoader::load(f, Some(131651)).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
         loader.load_page_index(true, true).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 1);