Update more parquet APIs to use u64

Author: alamb

parquet/src/arrow/arrow_writer/mod.rs

@@ -431,7 +431,7 @@ impl ChunkReader for ArrowColumnChunkData {
         ))
     }
 
-    fn get_bytes(&self, _start: u64, _length: usize) -> Result<Bytes> {
+    fn get_bytes(&self, _start: u64, _length: u64`) -> Result<Bytes> {
         unimplemented!()
     }
 }

parquet/src/arrow/async_reader/metadata.rs

@@ -19,7 +19,7 @@ use crate::arrow::async_reader::AsyncFileReader;
 use crate::errors::{ParquetError, Result};
 use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 use crate::file::page_index::index::Index;
-use crate::file::page_index::index_reader::{acc_range, decode_column_index, decode_offset_index};
+use crate::file::page_index::index_reader::{acc_range_usize, decode_column_index, decode_offset_index};
 use crate::file::FOOTER_SIZE;
 use bytes::Bytes;
 use futures::future::BoxFuture;
@@ -125,7 +125,7 @@ impl<F: MetadataFetch> MetadataLoader<F> {
         footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]);
 
         let footer = ParquetMetaDataReader::decode_footer_tail(&footer)?;
-        let length = footer.metadata_length();
+        let length = footer.metadata_length() as usize;
 
         if file_size < length + FOOTER_SIZE {
             return Err(ParquetError::EOF(format!(
@@ -181,8 +181,8 @@ impl<F: MetadataFetch> MetadataLoader<F> {
 
         let mut range = None;
         for c in self.metadata.row_groups().iter().flat_map(|r| r.columns()) {
-            range = acc_range(range, c.column_index_range());
-            range = acc_range(range, c.offset_index_range());
+            range = acc_range_usize(range, c.column_index_range());
+            range = acc_range_usize(range, c.offset_index_range());
         }
         let range = match range {
             None => return Ok(()),
@@ -294,6 +294,8 @@ where
     F: FnMut(Range<usize>) -> Fut + Send,
     Fut: Future<Output = Result<Bytes>> + Send,
 {
+    let prefetch = prefetch.map(|v| v as u64);
+    let file_size = file_size as u64;
     let fetch = MetadataFetchFn(fetch);
     ParquetMetaDataReader::new()
         .with_prefetch_hint(prefetch)

parquet/src/arrow/async_reader/mod.rs

@@ -138,7 +138,7 @@ impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
 }
 
 impl<T: AsyncFileReader + MetadataFetch + AsyncRead + AsyncSeek + Unpin> MetadataSuffixFetch for T {
-    fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result<Bytes>> {
+    fn fetch_suffix(&mut self, suffix: u64) -> BoxFuture<'_, Result<Bytes>> {
         async move {
             self.seek(SeekFrom::End(-(suffix as i64))).await?;
             let mut buf = Vec::with_capacity(suffix);
@@ -1058,6 +1058,7 @@ impl ChunkReader for ColumnChunkData {
     }
 
     fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes> {
+        let length: usize = length.try_into()?;
         Ok(self.get(start)?.slice(..length))
     }
 }

parquet/src/errors.rs

@@ -46,12 +46,12 @@ pub enum ParquetError {
     ArrowError(String),
     /// Error when the requested column index is more than the
     /// number of columns in the row group
-    IndexOutOfBound(usize, usize),
+    IndexOutOfBound(u64, u64),
     /// An external error variant
     External(Box<dyn Error + Send + Sync>),
-    /// Returned when a function needs more data to complete properly. The `usize` field indicates
+    /// Returned when a function needs more data to complete properly. The `u64` field indicates
     /// the total number of bytes required, not the number of additional bytes.
-    NeedMoreData(usize),
+    NeedMoreData(u64),
 }
 
 impl std::fmt::Display for ParquetError {

parquet/src/file/footer.rs

@@ -76,6 +76,6 @@ pub fn decode_metadata(buf: &[u8]) -> Result<ParquetMetaData> {
     since = "53.1.0",
     note = "Use ParquetMetaDataReader::decode_footer_tail"
 )]
-pub fn decode_footer(slice: &[u8; FOOTER_SIZE]) -> Result<usize> {
+pub fn decode_footer(slice: &[u8; FOOTER_SIZE as usize]) -> Result<u64> {
     ParquetMetaDataReader::decode_footer_tail(slice).map(|f| f.metadata_length())
 }

parquet/src/file/metadata/reader.rs

@@ -72,10 +72,10 @@ pub struct ParquetMetaDataReader {
     metadata: Option<ParquetMetaData>,
     column_index: bool,
     offset_index: bool,
-    prefetch_hint: Option<usize>,
-    // Size of the serialized thrift metadata plus the 8 byte footer. Only set if
-    // `self.parse_metadata` is called.
-    metadata_size: Option<usize>,
+    prefetch_hint: Option<u64>,
+    /// Size of the serialized thrift metadata plus the 8 byte footer. Only set if
+    /// `self.parse_metadata` is called.
+    metadata_size: Option<u64>,
     #[cfg(feature = "encryption")]
     file_decryption_properties: Option<FileDecryptionProperties>,
 }
@@ -84,13 +84,13 @@ pub struct ParquetMetaDataReader {
 ///
 /// This is parsed from the last 8 bytes of the Parquet file
 pub struct FooterTail {
-    metadata_length: usize,
+    metadata_length: u64,
     encrypted_footer: bool,
 }
 
 impl FooterTail {
     /// The length of the footer metadata in bytes
-    pub fn metadata_length(&self) -> usize {
+    pub fn metadata_length(&self) -> u64 {
         self.metadata_length
     }
 
@@ -151,7 +151,7 @@ impl ParquetMetaDataReader {
     /// to fully decode the [`ParquetMetaData`], which can reduce the number of fetch requests and
     /// reduce latency. Setting `prefetch` too small will not trigger an error, but will result
     /// in extra fetches being performed.
-    pub fn with_prefetch_hint(mut self, prefetch: Option<usize>) -> Self {
+    pub fn with_prefetch_hint(mut self, prefetch: Option<u64>) -> Self {
         self.prefetch_hint = prefetch;
         self
     }
@@ -209,7 +209,7 @@ impl ParquetMetaDataReader {
     /// the request, and must include the Parquet footer. If page indexes are desired, the buffer
     /// must contain the entire file, or [`Self::try_parse_sized()`] should be used.
     pub fn try_parse<R: ChunkReader>(&mut self, reader: &R) -> Result<()> {
-        self.try_parse_sized(reader, reader.len() as usize)
+        self.try_parse_sized(reader, reader.len())
     }
 
     /// Same as [`Self::try_parse()`], but provide the original file size in the case that `reader`
@@ -284,13 +284,13 @@ impl ParquetMetaDataReader {
     /// }
     /// let metadata = reader.finish().unwrap();
     /// ```
-    pub fn try_parse_sized<R: ChunkReader>(&mut self, reader: &R, file_size: usize) -> Result<()> {
+    pub fn try_parse_sized<R: ChunkReader>(&mut self, reader: &R, file_size: u64) -> Result<()> {
         self.metadata = match self.parse_metadata(reader) {
             Ok(metadata) => Some(metadata),
             Err(ParquetError::NeedMoreData(needed)) => {
                 // If reader is the same length as `file_size` then presumably there is no more to
                 // read, so return an EOF error.
-                if file_size == reader.len() as usize || needed > file_size {
+                if file_size == reader.len() || needed > file_size {
                     return Err(eof_err!(
                         "Parquet file too small. Size is {} but need {}",
                         file_size,
@@ -315,7 +315,7 @@ impl ParquetMetaDataReader {
     /// Read the page index structures when a [`ParquetMetaData`] has already been obtained.
     /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`].
     pub fn read_page_indexes<R: ChunkReader>(&mut self, reader: &R) -> Result<()> {
-        self.read_page_indexes_sized(reader, reader.len() as usize)
+        self.read_page_indexes_sized(reader, reader.len())
     }
 
     /// Read the page index structures when a [`ParquetMetaData`] has already been obtained.
@@ -326,7 +326,7 @@ impl ParquetMetaDataReader {
     pub fn read_page_indexes_sized<R: ChunkReader>(
         &mut self,
         reader: &R,
-        file_size: usize,
+        file_size: u64,
     ) -> Result<()> {
         if self.metadata.is_none() {
             return Err(general_err!(
@@ -350,7 +350,7 @@ impl ParquetMetaDataReader {
 
         // Check to see if needed range is within `file_range`. Checking `range.end` seems
         // redundant, but it guards against `range_for_page_index()` returning garbage.
-        let file_range = file_size.saturating_sub(reader.len() as usize)..file_size;
+        let file_range = file_size.saturating_sub(reader.len())..file_size;
         if !(file_range.contains(&range.start) && file_range.contains(&range.end)) {
             // Requested range starts beyond EOF
             if range.end > file_size {
@@ -378,7 +378,7 @@ impl ParquetMetaDataReader {
         }
 
         let bytes_needed = range.end - range.start;
-        let bytes = reader.get_bytes((range.start - file_range.start) as u64, bytes_needed)?;
+        let bytes = reader.get_bytes(range.start - file_range.start, bytes_needed)?;
         let offset = range.start;
 
         self.parse_column_index(&bytes, offset)?;
@@ -397,7 +397,7 @@ impl ParquetMetaDataReader {
     pub async fn load_and_finish<F: MetadataFetch>(
         mut self,
         fetch: F,
-        file_size: usize,
+        file_size: u64,
     ) -> Result<ParquetMetaData> {
         self.try_load(fetch, file_size).await?;
         self.finish()
@@ -426,7 +426,7 @@ impl ParquetMetaDataReader {
     pub async fn try_load<F: MetadataFetch>(
         &mut self,
         mut fetch: F,
-        file_size: usize,
+        file_size: u64,
     ) -> Result<()> {
         let (metadata, remainder) = self.load_metadata(&mut fetch, file_size).await?;
 
@@ -473,7 +473,7 @@ impl ParquetMetaDataReader {
     async fn load_page_index_with_remainder<F: MetadataFetch>(
         &mut self,
         mut fetch: F,
-        remainder: Option<(usize, Bytes)>,
+        remainder: Option<(u64, Bytes)>,
     ) -> Result<()> {
         if self.metadata.is_none() {
             return Err(general_err!("Footer metadata is not present"));
@@ -507,7 +507,8 @@ impl ParquetMetaDataReader {
         Ok(())
     }
 
-    fn parse_column_index(&mut self, bytes: &Bytes, start_offset: usize) -> Result<()> {
+    fn parse_column_index(&mut self, bytes: &Bytes, start_offset: u64) -> Result<()> {
+        let start_offset: usize = start_offset.try_into()?;
         let metadata = self.metadata.as_mut().unwrap();
         if self.column_index {
             let index = metadata
@@ -531,7 +532,8 @@ impl ParquetMetaDataReader {
         Ok(())
     }
 
-    fn parse_offset_index(&mut self, bytes: &Bytes, start_offset: usize) -> Result<()> {
+    fn parse_offset_index(&mut self, bytes: &Bytes, start_offset: u64) -> Result<()> {
+        let start_offset: usize = start_offset.try_into()?;
         let metadata = self.metadata.as_mut().unwrap();
         if self.offset_index {
             let index = metadata
@@ -555,7 +557,7 @@ impl ParquetMetaDataReader {
         Ok(())
     }
 
-    fn range_for_page_index(&self) -> Option<Range<usize>> {
+    fn range_for_page_index(&self) -> Option<Range<u64>> {
         // sanity check
         self.metadata.as_ref()?;
 
@@ -592,7 +594,7 @@ impl ParquetMetaDataReader {
         let footer_metadata_len = FOOTER_SIZE + metadata_len;
         self.metadata_size = Some(footer_metadata_len);
 
-        if footer_metadata_len > file_size as usize {
+        if footer_metadata_len > file_size {
             return Err(ParquetError::NeedMoreData(footer_metadata_len));
         }
 
@@ -607,7 +609,7 @@ impl ParquetMetaDataReader {
     /// been provided, then return that value if it is larger than the size of the Parquet
     /// file footer (8 bytes). Otherwise returns `8`.
     #[cfg(all(feature = "async", feature = "arrow"))]
-    fn get_prefetch_size(&self) -> usize {
+    fn get_prefetch_size(&self) -> u64 {
         if let Some(prefetch) = self.prefetch_hint {
             if prefetch > FOOTER_SIZE {
                 return prefetch;
@@ -620,8 +622,8 @@ impl ParquetMetaDataReader {
     async fn load_metadata<F: MetadataFetch>(
         &self,
         fetch: &mut F,
-        file_size: usize,
-    ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> {
+        file_size: u64,
+    ) -> Result<(ParquetMetaData, Option<(u64, Bytes)>)> {
         let prefetch = self.get_prefetch_size();
 
         if file_size < FOOTER_SIZE {
@@ -679,7 +681,7 @@ impl ParquetMetaDataReader {
     async fn load_metadata_via_suffix<F: MetadataSuffixFetch>(
         &self,
         fetch: &mut F,
-    ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> {
+    ) -> Result<(ParquetMetaData, Option<(u64, Bytes)>)> {
         let prefetch = self.get_prefetch_size();
 
         let suffix = fetch.fetch_suffix(prefetch as _).await?;
@@ -747,18 +749,18 @@ impl ParquetMetaDataReader {
         } else {
             return Err(general_err!("Invalid Parquet file. Corrupt footer"));
         };
-        // get the metadata length from the footer
+        // get the metadata length from the footer (infallible)
         let metadata_len = u32::from_le_bytes(slice[..4].try_into().unwrap());
+        let metadata_length = metadata_len as u64;
         Ok(FooterTail {
-            // u32 won't be larger than usize in most cases
-            metadata_length: metadata_len as usize,
+            metadata_length,
             encrypted_footer,
         })
     }
 
     /// Decodes the Parquet footer, returning the metadata length in bytes
     #[deprecated(note = "use decode_footer_tail instead")]
-    pub fn decode_footer(slice: &[u8; FOOTER_SIZE]) -> Result<usize> {
+    pub fn decode_footer(slice: &[u8; FOOTER_SIZE]) -> Result<u64> {
         Self::decode_footer_tail(slice).map(|f| f.metadata_length)
     }
 
@@ -1091,7 +1093,7 @@ mod tests {
     #[test]
     fn test_try_parse() {
         let file = get_test_file("alltypes_tiny_pages.parquet");
-        let len = file.len() as usize;
+        let len = file.len();
 
         let mut reader = ParquetMetaDataReader::new().with_page_indexes(true);
 

parquet/src/file/page_index/index_reader.rs

@@ -31,7 +31,15 @@ use std::ops::Range;
 /// Computes the covering range of two optional ranges
 ///
 /// For example `acc_range(Some(7..9), Some(1..3)) = Some(1..9)`
-pub(crate) fn acc_range(a: Option<Range<usize>>, b: Option<Range<usize>>) -> Option<Range<usize>> {
+pub(crate) fn acc_range(a: Option<Range<u64>>, b: Option<Range<u64>>) -> Option<Range<u64>> {
+    match (a, b) {
+        (Some(a), Some(b)) => Some(a.start.min(b.start)..a.end.max(b.end)),
+        (None, x) | (x, None) => x,
+    }
+}
+
+/// Computes the covering range of two optional ranges of `usize`
+pub(crate) fn acc_range_usize(a: Option<Range<usize>>, b: Option<Range<usize>>) -> Option<Range<usize>> {
     match (a, b) {
         (Some(a), Some(b)) => Some(a.start.min(b.start)..a.end.max(b.end)),
         (None, x) | (x, None) => x,

parquet/src/file/reader.rs

@@ -23,7 +23,6 @@ use bytes::{Buf, Bytes};
 use std::fs::File;
 use std::io::{BufReader, Seek, SeekFrom};
 use std::{io::Read, sync::Arc};
-
 use crate::bloom_filter::Sbbf;
 use crate::column::page::PageIterator;
 use crate::column::{page::PageReader, reader::ColumnReader};
@@ -77,7 +76,7 @@ pub trait ChunkReader: Length + Send + Sync {
     ///
     /// Similarly to [`Self::get_read`], this method may have side-effects on
     /// previously returned readers.
-    fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes>;
+    fn get_bytes(&self, start: u64, length: u64) -> Result<Bytes>;
 }
 
 impl Length for File {
@@ -95,11 +94,13 @@ impl ChunkReader for File {
         Ok(BufReader::new(self.try_clone()?))
     }
 
-    fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes> {
-        let mut buffer = Vec::with_capacity(length);
+    fn get_bytes(&self, start: u64, length: u64) -> Result<Bytes> {
+        let cap: usize = length.try_into()?;
+        let mut buffer = Vec::with_capacity(cap);
         let mut reader = self.try_clone()?;
         reader.seek(SeekFrom::Start(start))?;
         let read = reader.take(length as _).read_to_end(&mut buffer)?;
+        let read = read as u64;
 
         if read != length {
             return Err(eof_err!(
@@ -126,8 +127,9 @@ impl ChunkReader for Bytes {
         Ok(self.slice(start..).reader())
     }
 
-    fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes> {
-        let start = start as usize;
+    fn get_bytes(&self, start: u64, length: u64) -> Result<Bytes> {
+        let start: usize = start.try_into()?;
+        let length: usize = length.try_into()?;
         Ok(self.slice(start..start + length))
     }
 }