Change Parquet API interaction to use u64 (support files larger than 4GB in WASM) (#7371)

* Change AsyncFileReader trait for u64

Signed-off-by: Arpit Bandejiya <[email protected]>

* update metadatafetch trait

Signed-off-by: Arpit Bandejiya <[email protected]>

* Fix lint issue

* fix tests for latest main

* Address comments by @mbrobbel from #7252

* Fix compile

* Revert suffix length back to usize

* Use `u64` for `file_size`

* address comments

* fix calculation of metadata_start

* change file_size type to u64

* change _sized functions to take u64 file_size

* clippy

* remove some potential panics

* use u64 for page index ranges

* Revert change to deprecated method

---------

Signed-off-by: Arpit Bandejiya <[email protected]>
Co-authored-by: Arpit Bandejiya <[email protected]>
Co-authored-by: Ed Seidl <[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>

Author: 4 people

Diff for: parquet/examples/external_metadata.rs

@@ -112,7 +112,7 @@ async fn get_metadata_from_remote_parquet_file(
     // tell the reader to read the page index
     ParquetMetaDataReader::new()
         .with_page_indexes(true)
-        .load_and_finish(remote_file, file_size as usize)
+        .load_and_finish(remote_file, file_size)
         .await
         .unwrap()
 }

Diff for: parquet/examples/read_with_rowgroup.rs

@@ -166,9 +166,7 @@ impl InMemoryRowGroup {
         for (leaf_idx, meta) in self.metadata.columns().iter().enumerate() {
             if self.mask.leaf_included(leaf_idx) {
                 let (start, len) = meta.byte_range();
-                let data = reader
-                    .get_bytes(start as usize..(start + len) as usize)
-                    .await?;
+                let data = reader.get_bytes(start..(start + len)).await?;
 
                 vs[leaf_idx] = Some(Arc::new(ColumnChunkData {
                     offset: start as usize,

Diff for: parquet/src/arrow/arrow_reader/selection.rs

@@ -162,8 +162,8 @@ impl RowSelection {
     /// Note: this method does not make any effort to combine consecutive ranges, nor coalesce
     /// ranges that are close together. This is instead delegated to the IO subsystem to optimise,
     /// e.g. [`ObjectStore::get_ranges`](object_store::ObjectStore::get_ranges)
-    pub fn scan_ranges(&self, page_locations: &[crate::format::PageLocation]) -> Vec<Range<usize>> {
-        let mut ranges = vec![];
+    pub fn scan_ranges(&self, page_locations: &[crate::format::PageLocation]) -> Vec<Range<u64>> {
+        let mut ranges: Vec<Range<u64>> = vec![];
         let mut row_offset = 0;
 
         let mut pages = page_locations.iter().peekable();
@@ -175,8 +175,8 @@ impl RowSelection {
 
         while let Some((selector, page)) = current_selector.as_mut().zip(current_page) {
             if !(selector.skip || current_page_included) {
-                let start = page.offset as usize;
-                let end = start + page.compressed_page_size as usize;
+                let start = page.offset as u64;
+                let end = start + page.compressed_page_size as u64;
                 ranges.push(start..end);
                 current_page_included = true;
             }
@@ -200,8 +200,8 @@ impl RowSelection {
                 }
             } else {
                 if !(selector.skip || current_page_included) {
-                    let start = page.offset as usize;
-                    let end = start + page.compressed_page_size as usize;
+                    let start = page.offset as u64;
+                    let end = start + page.compressed_page_size as u64;
                     ranges.push(start..end);
                 }
                 current_selector = selectors.next()

Diff for: parquet/src/arrow/async_reader/metadata.rs

@@ -48,12 +48,13 @@ use std::ops::Range;
 ///     file: tokio::fs::File,
 /// }
 /// impl MetadataFetch for TokioFileMetadata {
-///     fn fetch(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
+///     fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
 ///         // return a future that fetches data in range
 ///         async move {
-///             let mut buf = vec![0; range.len()]; // target buffer
+///             let len = (range.end - range.start).try_into().unwrap();
+///             let mut buf = vec![0; len]; // target buffer
 ///             // seek to the start of the range and read the data
-///             self.file.seek(SeekFrom::Start(range.start as u64)).await?;
+///             self.file.seek(SeekFrom::Start(range.start)).await?;
 ///             self.file.read_exact(&mut buf).await?;
 ///             Ok(Bytes::from(buf)) // convert to Bytes
 ///         }
@@ -66,11 +67,11 @@ pub trait MetadataFetch {
     ///
     /// Note the returned type is a boxed future, often created by
     /// [FutureExt::boxed]. See the trait documentation for an example
-    fn fetch(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>>;
+    fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>>;
 }
 
 impl<T: AsyncFileReader> MetadataFetch for &mut T {
-    fn fetch(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
+    fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
         self.get_bytes(range)
     }
 }
@@ -117,7 +118,7 @@ impl<F: MetadataFetch> MetadataLoader<F> {
             file_size - FOOTER_SIZE
         };
 
-        let suffix = fetch.fetch(footer_start..file_size).await?;
+        let suffix = fetch.fetch(footer_start as u64..file_size as u64).await?;
         let suffix_len = suffix.len();
 
         let mut footer = [0; FOOTER_SIZE];
@@ -137,7 +138,9 @@ impl<F: MetadataFetch> MetadataLoader<F> {
         // Did not fetch the entire file metadata in the initial read, need to make a second request
         let (metadata, remainder) = if length > suffix_len - FOOTER_SIZE {
             let metadata_start = file_size - length - FOOTER_SIZE;
-            let meta = fetch.fetch(metadata_start..file_size - FOOTER_SIZE).await?;
+            let meta = fetch
+                .fetch(metadata_start as u64..(file_size - FOOTER_SIZE) as u64)
+                .await?;
             (ParquetMetaDataReader::decode_metadata(&meta)?, None)
         } else {
             let metadata_start = file_size - length - FOOTER_SIZE - footer_start;
@@ -187,16 +190,18 @@ impl<F: MetadataFetch> MetadataLoader<F> {
         };
 
         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_start as u64 <= range.start => {
+                let remainder_start = *remainder_start as u64;
+                let range_start = usize::try_from(range.start - remainder_start)?;
+                let range_end = usize::try_from(range.end - remainder_start)?;
+                remainder.slice(range_start..range_end)
             }
             // Note: this will potentially fetch data already in remainder, this keeps things simple
             _ => self.fetch.fetch(range.start..range.end).await?,
         };
 
         // Sanity check
-        assert_eq!(data.len(), range.end - range.start);
+        assert_eq!(data.len(), (range.end - range.start) as usize);
         let offset = range.start;
 
         if column_index {
@@ -208,10 +213,11 @@ impl<F: MetadataFetch> MetadataLoader<F> {
                     x.columns()
                         .iter()
                         .map(|c| match c.column_index_range() {
-                            Some(r) => decode_column_index(
-                                &data[r.start - offset..r.end - offset],
-                                c.column_type(),
-                            ),
+                            Some(r) => {
+                                let r_start = usize::try_from(r.start - offset)?;
+                                let r_end = usize::try_from(r.end - offset)?;
+                                decode_column_index(&data[r_start..r_end], c.column_type())
+                            }
                             None => Ok(Index::NONE),
                         })
                         .collect::<Result<Vec<_>>>()
@@ -230,7 +236,11 @@ impl<F: MetadataFetch> MetadataLoader<F> {
                     x.columns()
                         .iter()
                         .map(|c| match c.offset_index_range() {
-                            Some(r) => decode_offset_index(&data[r.start - offset..r.end - offset]),
+                            Some(r) => {
+                                let r_start = usize::try_from(r.start - offset)?;
+                                let r_end = usize::try_from(r.end - offset)?;
+                                decode_offset_index(&data[r_start..r_end])
+                            }
                             None => Err(general_err!("missing offset index")),
                         })
                         .collect::<Result<Vec<_>>>()
@@ -256,8 +266,8 @@ where
     F: FnMut(Range<usize>) -> Fut + Send,
     Fut: Future<Output = Result<Bytes>> + Send,
 {
-    fn fetch(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
-        async move { self.0(range).await }.boxed()
+    fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
+        async move { self.0(range.start.try_into()?..range.end.try_into()?).await }.boxed()
     }
 }
 
@@ -287,6 +297,7 @@ where
     F: FnMut(Range<usize>) -> Fut + Send,
     Fut: Future<Output = Result<Bytes>> + Send,
 {
+    let file_size = u64::try_from(file_size)?;
     let fetch = MetadataFetchFn(fetch);
     ParquetMetaDataReader::new()
         .with_prefetch_hint(prefetch)

Diff for: parquet/src/arrow/async_reader/mod.rs

@@ -80,10 +80,10 @@ pub use store::*;
 /// [`tokio::fs::File`]: https://docs.rs/tokio/latest/tokio/fs/struct.File.html
 pub trait AsyncFileReader: Send {
     /// Retrieve the bytes in `range`
-    fn get_bytes(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>>;
+    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>>;
 
     /// Retrieve multiple byte ranges. The default implementation will call `get_bytes` sequentially
-    fn get_byte_ranges(&mut self, ranges: Vec<Range<usize>>) -> BoxFuture<'_, Result<Vec<Bytes>>> {
+    fn get_byte_ranges(&mut self, ranges: Vec<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>> {
         async move {
             let mut result = Vec::with_capacity(ranges.len());
 
@@ -121,11 +121,11 @@ pub trait AsyncFileReader: Send {
 
 /// This allows Box<dyn AsyncFileReader + '_> to be used as an AsyncFileReader,
 impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
-    fn get_bytes(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
+    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
         self.as_mut().get_bytes(range)
     }
 
-    fn get_byte_ranges(&mut self, ranges: Vec<Range<usize>>) -> BoxFuture<'_, Result<Vec<Bytes>>> {
+    fn get_byte_ranges(&mut self, ranges: Vec<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>> {
         self.as_mut().get_byte_ranges(ranges)
     }
 
@@ -150,14 +150,14 @@ impl<T: AsyncFileReader + MetadataFetch + AsyncRead + AsyncSeek + Unpin> Metadat
 }
 
 impl<T: AsyncRead + AsyncSeek + Unpin + Send> AsyncFileReader for T {
-    fn get_bytes(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
+    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
         async move {
-            self.seek(SeekFrom::Start(range.start as u64)).await?;
+            self.seek(SeekFrom::Start(range.start)).await?;
 
             let to_read = range.end - range.start;
-            let mut buffer = Vec::with_capacity(to_read);
-            let read = self.take(to_read as u64).read_to_end(&mut buffer).await?;
-            if read != to_read {
+            let mut buffer = Vec::with_capacity(to_read.try_into()?);
+            let read = self.take(to_read).read_to_end(&mut buffer).await?;
+            if read as u64 != to_read {
                 return Err(eof_err!("expected to read {} bytes, got {}", to_read, read));
             }
 
@@ -424,7 +424,7 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> {
         let metadata = self.metadata.row_group(row_group_idx);
         let column_metadata = metadata.column(column_idx);
 
-        let offset: usize = if let Some(offset) = column_metadata.bloom_filter_offset() {
+        let offset: u64 = if let Some(offset) = column_metadata.bloom_filter_offset() {
             offset
                 .try_into()
                 .map_err(|_| ParquetError::General("Bloom filter offset is invalid".to_string()))?
@@ -433,16 +433,16 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> {
         };
 
         let buffer = match column_metadata.bloom_filter_length() {
-            Some(length) => self.input.0.get_bytes(offset..offset + length as usize),
+            Some(length) => self.input.0.get_bytes(offset..offset + length as u64),
             None => self
                 .input
                 .0
-                .get_bytes(offset..offset + SBBF_HEADER_SIZE_ESTIMATE),
+                .get_bytes(offset..offset + SBBF_HEADER_SIZE_ESTIMATE as u64),
         }
         .await?;
 
         let (header, bitset_offset) =
-            chunk_read_bloom_filter_header_and_offset(offset as u64, buffer.clone())?;
+            chunk_read_bloom_filter_header_and_offset(offset, buffer.clone())?;
 
         match header.algorithm {
             BloomFilterAlgorithm::BLOCK(_) => {
@@ -461,14 +461,17 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> {
         }
 
         let bitset = match column_metadata.bloom_filter_length() {
-            Some(_) => buffer.slice((bitset_offset as usize - offset)..),
+            Some(_) => buffer.slice(
+                (TryInto::<usize>::try_into(bitset_offset).unwrap()
+                    - TryInto::<usize>::try_into(offset).unwrap())..,
+            ),
             None => {
-                let bitset_length: usize = header.num_bytes.try_into().map_err(|_| {
+                let bitset_length: u64 = header.num_bytes.try_into().map_err(|_| {
                     ParquetError::General("Bloom filter length is invalid".to_string())
                 })?;
                 self.input
                     .0
-                    .get_bytes(bitset_offset as usize..bitset_offset as usize + bitset_length)
+                    .get_bytes(bitset_offset..bitset_offset + bitset_length)
                     .await?
             }
         };
@@ -880,7 +883,7 @@ impl InMemoryRowGroup<'_> {
         if let Some((selection, offset_index)) = selection.zip(self.offset_index) {
             // If we have a `RowSelection` and an `OffsetIndex` then only fetch pages required for the
             // `RowSelection`
-            let mut page_start_offsets: Vec<Vec<usize>> = vec![];
+            let mut page_start_offsets: Vec<Vec<u64>> = vec![];
 
             let fetch_ranges = self
                 .column_chunks
@@ -893,11 +896,11 @@ impl InMemoryRowGroup<'_> {
                 .flat_map(|(idx, (_chunk, chunk_meta))| {
                     // If the first page does not start at the beginning of the column,
                     // then we need to also fetch a dictionary page.
-                    let mut ranges = vec![];
+                    let mut ranges: Vec<Range<u64>> = vec![];
                     let (start, _len) = chunk_meta.byte_range();
                     match offset_index[idx].page_locations.first() {
                         Some(first) if first.offset as u64 != start => {
-                            ranges.push(start as usize..first.offset as usize);
+                            ranges.push(start..first.offset as u64);
                         }
                         _ => (),
                     }
@@ -925,7 +928,11 @@ impl InMemoryRowGroup<'_> {
 
                     *chunk = Some(Arc::new(ColumnChunkData::Sparse {
                         length: metadata.column(idx).byte_range().1 as usize,
-                        data: offsets.into_iter().zip(chunks.into_iter()).collect(),
+                        data: offsets
+                            .into_iter()
+                            .map(|x| x as usize)
+                            .zip(chunks.into_iter())
+                            .collect(),
                     }))
                 }
             }
@@ -938,7 +945,7 @@ impl InMemoryRowGroup<'_> {
                 .map(|(idx, _chunk)| {
                     let column = metadata.column(idx);
                     let (start, length) = column.byte_range();
-                    start as usize..(start + length) as usize
+                    start..(start + length)
                 })
                 .collect();
 
@@ -1108,9 +1115,16 @@ mod tests {
     }
 
     impl AsyncFileReader for TestReader {
-        fn get_bytes(&mut self, range: Range<usize>) -> BoxFuture<'_, Result<Bytes>> {
-            self.requests.lock().unwrap().push(range.clone());
-            futures::future::ready(Ok(self.data.slice(range))).boxed()
+        fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
+            let range = range.clone();
+            self.requests
+                .lock()
+                .unwrap()
+                .push(range.start as usize..range.end as usize);
+            futures::future::ready(Ok(self
+                .data
+                .slice(range.start as usize..range.end as usize)))
+            .boxed()
         }
 
         fn get_metadata<'a>(
@@ -2238,7 +2252,7 @@ mod tests {
         let file_size = file.metadata().await.unwrap().len();
         let mut metadata = ParquetMetaDataReader::new()
             .with_page_indexes(true)
-            .load_and_finish(&mut file, file_size as usize)
+            .load_and_finish(&mut file, file_size)
             .await
             .unwrap();
 
@@ -2263,7 +2277,7 @@ mod tests {
         let file_size = file.metadata().await.unwrap().len();
         let metadata = ParquetMetaDataReader::new()
             .with_page_indexes(true)
-            .load_and_finish(&mut file, file_size as usize)
+            .load_and_finish(&mut file, file_size)
             .await
             .unwrap();
 
@@ -2309,7 +2323,7 @@ mod tests {
         let file_size = file.metadata().await.unwrap().len();
         let metadata = ParquetMetaDataReader::new()
             .with_page_indexes(true)
-            .load_and_finish(&mut file, file_size as usize)
+            .load_and_finish(&mut file, file_size)
             .await
             .unwrap();