Add ChunkReader::get_bytes

parquet/src/arrow/async_reader.rs

@@ -628,10 +628,14 @@ impl ChunkReader for ColumnChunkData {
     type T = bytes::buf::Reader<Bytes>;
 
     fn get_read(&self, start: u64, length: usize) -> Result<Self::T> {
+        Ok(self.get_bytes(start, length)?.reader())
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes> {
         match &self {
             ColumnChunkData::Sparse { data, .. } => data
                 .binary_search_by_key(&start, |(offset, _)| *offset as u64)
-                .map(|idx| data[idx].1.slice(0..length).reader())
+                .map(|idx| data[idx].1.slice(0..length))
                 .map_err(|_| {
                     ParquetError::General(format!(
                         "Invalid offset in sparse column chunk data: {}",
@@ -641,7 +645,7 @@ impl ChunkReader for ColumnChunkData {
             ColumnChunkData::Dense { offset, data } => {
                 let start = start as usize - *offset;
                 let end = start + length;
-                Ok(data.slice(start..end).reader())
+                Ok(data.slice(start..end))
             }
         }
     }

parquet/src/file/footer.rs

@@ -62,19 +62,8 @@ pub fn parse_metadata<R: ChunkReader>(chunk_reader: &R) -> Result<ParquetMetaDat
         ));
     }
 
-    let mut metadata = Vec::with_capacity(metadata_len);
-
-    let read = chunk_reader
-        .get_read(file_size - footer_metadata_len as u64, metadata_len)?
-        .read_to_end(&mut metadata)?;
-
-    if read != metadata_len {
-        return Err(eof_err!(
-            "Expected to read {} bytes of metadata, got {}",
-            metadata_len,
-            read
-        ));
-    }
+    let metadata =
+        chunk_reader.get_bytes(file_size - footer_metadata_len as u64, metadata_len)?;
 
     decode_metadata(&metadata)
 }

parquet/src/file/reader.rs

@@ -18,6 +18,7 @@
 //! Contains file reader API and provides methods to access file metadata, row group
 //! readers to read individual column chunks, or access record iterator.
 
+use bytes::Bytes;
 use std::{boxed::Box, io::Read, sync::Arc};
 
 use crate::column::page::PageIterator;
@@ -48,6 +49,22 @@ pub trait ChunkReader: Length + Send + Sync {
     /// Get a serially readable slice of the current reader
     /// This should fail if the slice exceeds the current bounds
     fn get_read(&self, start: u64, length: usize) -> Result<Self::T>;
+
+    /// Get a range as bytes
+    /// This should fail if the exact number of bytes cannot be read
+    fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes> {
+        let mut buffer = Vec::with_capacity(length);
+        let read = self.get_read(start, length)?.read_to_end(&mut buffer)?;
+
+        if read != length {
+            return Err(eof_err!(
+                "Expected to read {} bytes, read only {}",
+                length,
+                read
+            ));
+        }
+        Ok(buffer.into())
+    }
 }
 
 // ----------------------------------------------------------------------

parquet/src/file/serialized_reader.rs

@@ -79,8 +79,12 @@ impl ChunkReader for Bytes {
     type T = bytes::buf::Reader<Bytes>;
 
     fn get_read(&self, start: u64, length: usize) -> Result<Self::T> {
+        Ok(self.get_bytes(start, length)?.reader())
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes> {
         let start = start as usize;
-        Ok(self.slice(start..start + length).reader())
+        Ok(self.slice(start..start + length))
     }
 }
 
@@ -623,26 +627,13 @@ impl<R: ChunkReader> PageReader for SerializedPageReader<R> {
 
                     let page_len = front.compressed_page_size as usize;
 
-                    // TODO: Add ChunkReader get_bytes to potentially avoid copy
-                    let mut buffer = Vec::with_capacity(page_len);
-                    let read = self
-                        .reader
-                        .get_read(front.offset as u64, page_len)?
-                        .read_to_end(&mut buffer)?;
-
-                    if read != page_len {
-                        return Err(eof_err!(
-                            "Expected to read {} bytes of page, read only {}",
-                            page_len,
-                            read
-                        ));
-                    }
+                    let buffer = self.reader.get_bytes(front.offset as u64, page_len)?;
 
-                    let mut cursor = Cursor::new(buffer);
+                    let mut cursor = Cursor::new(buffer.as_ref());
                     let header = read_page_header(&mut cursor)?;
                     let offset = cursor.position();
 
-                    let bytes = Bytes::from(cursor.into_inner()).slice(offset as usize..);
+                    let bytes = buffer.slice(offset as usize..);
                     decode_page(
                         header,
                         bytes.into(),