add support for max_page_header_size for pyarrow

Author: ayushbansal07

cpp/src/arrow/dataset/file_parquet.cc

@@ -76,6 +76,9 @@ parquet::ReaderProperties MakeReaderProperties(
     properties.disable_buffered_stream();
   }
   properties.set_buffer_size(parquet_scan_options->reader_properties->buffer_size());
+  if (format.reader_options.max_page_header_size != 0) {
+    properties.set_max_page_header_size(format.reader_options.max_page_header_size);
+  }
 
   auto file_decryption_prop =
       parquet_scan_options->reader_properties->file_decryption_properties();

cpp/src/arrow/dataset/file_parquet.h

@@ -92,6 +92,7 @@ class ARROW_DS_EXPORT ParquetFileFormat : public FileFormat {
     arrow::TimeUnit::type coerce_int96_timestamp_unit = arrow::TimeUnit::NANO;
     Type::type binary_type = Type::BINARY;
     Type::type list_type = Type::LIST;
+    uint32_t max_page_header_size = 0;
     /// @}
   } reader_options;
 

cpp/src/parquet/column_reader.cc

@@ -231,7 +231,6 @@ class SerializedPageReader : public PageReader {
       crypto_ctx_ = *crypto_ctx;
       InitDecryption();
     }
-    max_page_header_size_ = kDefaultMaxPageHeaderSize;
     decompressor_ = GetCodec(codec);
     always_compressed_ = always_compressed;
   }
@@ -244,7 +243,9 @@ class SerializedPageReader : public PageReader {
   // called then the content of previous page might be invalidated.
   std::shared_ptr<Page> NextPage() override;
 
-  void set_max_page_header_size(uint32_t size) override { max_page_header_size_ = size; }
+  void set_max_page_header_size(uint32_t size) override {
+    properties_.set_max_page_header_size(size);
+  }
 
  private:
   void UpdateDecryption(Decryptor* decryptor, int8_t module_type, std::string* page_aad);
@@ -260,7 +261,7 @@ class SerializedPageReader : public PageReader {
   // Fills in data_page_statistics.
   bool ShouldSkipPage(EncodedStatistics* data_page_statistics);
 
-  const ReaderProperties properties_;
+  ReaderProperties properties_;
   std::shared_ptr<ArrowInputStream> stream_;
 
   format::PageHeader current_page_header_;
@@ -288,9 +289,6 @@ class SerializedPageReader : public PageReader {
   // The ordinal fields in the context below are used for AAD suffix calculation.
   int32_t page_ordinal_;  // page ordinal does not count the dictionary page
 
-  // Maximum allowed page size
-  uint32_t max_page_header_size_;
-
   // Number of values read in data pages so far
   int64_t seen_num_values_;
 
@@ -422,7 +420,7 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() {
         std::stringstream ss;
         ss << e.what();
         allowed_page_size *= 2;
-        if (allowed_page_size > max_page_header_size_) {
+        if (allowed_page_size > properties_.max_page_header_size()) {
           ss << "Deserializing page header failed.\n";
           throw ParquetException(ss.str());
         }

cpp/src/parquet/column_reader.h

@@ -50,9 +50,6 @@ namespace parquet {
 class Decryptor;
 class Page;
 
-// 16 MB is the default maximum page header size
-static constexpr uint32_t kDefaultMaxPageHeaderSize = 16 * 1024 * 1024;
-
 // 16 KB is the default expected page header size
 static constexpr uint32_t kDefaultPageHeaderSize = 16 * 1024;
 

cpp/src/parquet/properties.h

@@ -57,6 +57,9 @@ enum class SizeStatisticsLevel : uint8_t {
   PageAndColumnChunk
 };
 
+// 16 MB is the default maximum page header size
+static constexpr uint32_t kDefaultMaxPageHeaderSize = 16 * 1024 * 1024;
+
 /// Align the default buffer size to a small multiple of a page size.
 constexpr int64_t kDefaultBufferSize = 4096 * 4;
 
@@ -101,6 +104,11 @@ class PARQUET_EXPORT ReaderProperties {
   /// Set the size of the buffered stream buffer in bytes.
   void set_buffer_size(int64_t size) { buffer_size_ = size; }
 
+  /// Return the size of the buffered stream buffer. 0 means default
+  uint32_t max_page_header_size() const { return max_page_header_size_; }
+  /// Set the size of the buffered stream buffer in bytes. 0 means default
+  void set_max_page_header_size(uint32_t size) { max_page_header_size_ = size; }
+
   /// \brief Return the size limit on thrift strings.
   ///
   /// This limit helps prevent space and time bombs in files, but may need to
@@ -142,6 +150,7 @@ class PARQUET_EXPORT ReaderProperties {
  private:
   MemoryPool* pool_;
   int64_t buffer_size_ = kDefaultBufferSize;
+  uint32_t max_page_header_size_ = kDefaultMaxPageHeaderSize;
   int32_t thrift_string_size_limit_ = kDefaultThriftStringSizeLimit;
   int32_t thrift_container_size_limit_ = kDefaultThriftContainerSizeLimit;
   bool buffered_stream_enabled_ = false;

python/pyarrow/_parquet.pyx

@@ -1571,7 +1571,8 @@ cdef class ParquetReader(_Weakrefable):
              thrift_string_size_limit=None,
              thrift_container_size_limit=None,
              page_checksum_verification=False,
-             arrow_extensions_enabled=False):
+             arrow_extensions_enabled=False,
+             max_page_header_size=None):
         """
         Open a parquet file for reading.
 
@@ -1591,6 +1592,7 @@ cdef class ParquetReader(_Weakrefable):
         thrift_container_size_limit : int, optional
         page_checksum_verification : bool, default False
         arrow_extensions_enabled : bool, default False
+        max_page_header_size : int, default None
         """
         cdef:
             shared_ptr[CFileMetaData] c_metadata
@@ -1624,6 +1626,11 @@ cdef class ParquetReader(_Weakrefable):
                                  "must be larger than zero")
             properties.set_thrift_container_size_limit(
                 thrift_container_size_limit)
+        if max_page_header_size is not None:
+            if max_page_header_size <= 0:
+                raise ValueError("max_page_header_size "
+                                 "must be larger than zero")
+            properties.set_max_page_header_size(max_page_header_size)
 
         if decryption_properties is not None:
             properties.file_decryption_properties(

python/pyarrow/includes/libparquet.pxd

@@ -422,6 +422,9 @@ cdef extern from "parquet/api/reader.h" namespace "parquet" nogil:
         void set_buffer_size(int64_t buf_size)
         int64_t buffer_size() const
 
+        void set_max_page_header_size(uint32_t size)
+        uint32_t max_page_header_size() const
+
         void set_thrift_string_size_limit(int32_t size)
         int32_t thrift_string_size_limit() const
 

python/pyarrow/parquet/core.py

@@ -265,6 +265,9 @@ class ParquetFile:
         If True, read Parquet logical types as Arrow extension types where possible,
         (e.g., read JSON as the canonical `arrow.json` extension type or UUID as
         the canonical `arrow.uuid` extension type).
+    max_page_header_size : int, default None
+        If not None, override the maximum size of a page header.
+        Deafults to 16MB, which should be sufficient for most Parquet files.
 
     Examples
     --------
@@ -314,7 +317,8 @@ def __init__(self, source, *, metadata=None, common_metadata=None,
                  coerce_int96_timestamp_unit=None,
                  decryption_properties=None, thrift_string_size_limit=None,
                  thrift_container_size_limit=None, filesystem=None,
-                 page_checksum_verification=False, arrow_extensions_enabled=True):
+                 page_checksum_verification=False, arrow_extensions_enabled=True,
+                 max_page_header_size=None):
 
         self._close_source = getattr(source, 'closed', True)
 
@@ -336,6 +340,7 @@ def __init__(self, source, *, metadata=None, common_metadata=None,
             thrift_container_size_limit=thrift_container_size_limit,
             page_checksum_verification=page_checksum_verification,
             arrow_extensions_enabled=arrow_extensions_enabled,
+            max_page_header_size=max_page_header_size,
         )
         self.common_metadata = common_metadata
         self._nested_paths_by_prefix = self._build_nested_paths()