(arrow-rs) branch main updated: Parquet: Avoid page-size overflows i32 (#8264)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/main by this push:
     new 30a779878a Parquet: Avoid page-size overflows i32 (#8264)
30a779878a is described below

commit 30a779878a617116afe73e48ac05abfff0040c6f
Author: mwish <ma...@gmail.com>
AuthorDate: Sat Sep 13 03:54:07 2025 +0800

    Parquet: Avoid page-size overflows i32 (#8264)
    
    # Which issue does this PR close?
    
    - Closes #8263
    
    # Rationale for this change
    
    See https://github.com/apache/arrow-rs/issues/8263
    
    # What changes are included in this PR?
    
    Avoid page overflows
    
    # Are these changes tested?
    
    * [x] Will add
    
    
    # Are there any user-facing changes?
    
    No
---
 parquet/src/arrow/arrow_writer/mod.rs |  2 +-
 parquet/src/column/page.rs            | 40 +++++++++++++++++++++++++++++++++--
 parquet/src/file/writer.rs            |  2 +-
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs
index c6b0b426f9..864c1bf2da 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -571,7 +571,7 @@ impl PageWriter for ArrowPageWriter {
             None => page,
         };
 
-        let page_header = page.to_thrift_header();
+        let page_header = page.to_thrift_header()?;
         let header = {
             let mut header = Vec::with_capacity(1024);
 
diff --git a/parquet/src/column/page.rs b/parquet/src/column/page.rs
index 1dabe6794f..a2f683d71f 100644
--- a/parquet/src/column/page.rs
+++ b/parquet/src/column/page.rs
@@ -196,9 +196,21 @@ impl CompressedPage {
     }
 
     /// Returns the thrift page header
-    pub(crate) fn to_thrift_header(&self) -> PageHeader {
+    pub(crate) fn to_thrift_header(&self) -> Result<PageHeader> {
         let uncompressed_size = self.uncompressed_size();
         let compressed_size = self.compressed_size();
+        if uncompressed_size > i32::MAX as usize {
+            return Err(general_err!(
+                "Page uncompressed size overflow: {}",
+                uncompressed_size
+            ));
+        }
+        if compressed_size > i32::MAX as usize {
+            return Err(general_err!(
+                "Page compressed size overflow: {}",
+                compressed_size
+            ));
+        }
         let num_values = self.num_values();
         let encoding = self.encoding();
         let page_type = self.page_type();
@@ -261,7 +273,7 @@ impl CompressedPage {
                 page_header.dictionary_page_header = Some(dictionary_page_header);
             }
         }
-        page_header
+        Ok(page_header)
     }
 
     /// Update the compressed buffer for a page.
@@ -491,4 +503,28 @@ mod tests {
         assert_eq!(cpage.encoding(), Encoding::PLAIN);
         assert_eq!(cpage.data(), &[0, 1, 2]);
     }
+
+    #[test]
+    fn test_compressed_page_uncompressed_size_overflow() {
+        // Test that to_thrift_header fails when uncompressed size exceeds i32::MAX
+        let data_page = Page::DataPage {
+            buf: Bytes::from(vec![0, 1, 2]),
+            num_values: 10,
+            encoding: Encoding::PLAIN,
+            def_level_encoding: Encoding::RLE,
+            rep_level_encoding: Encoding::RLE,
+            statistics: None,
+        };
+
+        // Create a CompressedPage with uncompressed size larger than i32::MAX
+        let uncompressed_size = (i32::MAX as usize) + 1;
+        let cpage = CompressedPage::new(data_page, uncompressed_size);
+
+        // Verify that to_thrift_header returns an error
+        let result = cpage.to_thrift_header();
+        assert!(result.is_err());
+
+        let error_msg = result.unwrap_err().to_string();
+        assert!(error_msg.contains("Page uncompressed size overflow"));
+    }
 }
diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs
index 690efb36f2..9adf67e68b 100644
--- a/parquet/src/file/writer.rs
+++ b/parquet/src/file/writer.rs
@@ -958,7 +958,7 @@ impl<W: Write + Send> PageWriter for SerializedPageWriter<'_, W> {
         let page_type = page.page_type();
         let start_pos = self.sink.bytes_written() as u64;
 
-        let page_header = page.to_thrift_header();
+        let page_header = page.to_thrift_header()?;
         let header_size = self.serialize_page_header(page_header)?;
 
         self.sink.write_all(page.data())?;