Added data_pagesize_limit to write parquet pages

arrow-parquet-integration-testing/src/main.rs

@@ -172,6 +172,7 @@ fn main() -> Result<()> {
         write_statistics: true,
         compression: args.compression.into(),
         version: args.version.into(),
+        data_pagesize_limit: None,
     };
 
     let encodings = schema

benches/write_parquet.rs

@@ -17,6 +17,7 @@ fn write(array: &dyn Array, encoding: Encoding) -> Result<()> {
         write_statistics: false,
         compression: CompressionOptions::Uncompressed,
         version: Version::V1,
+        data_pagesize_limit: None,
     };
 
     let row_groups = RowGroupIterator::try_new(

examples/parquet_write.rs

@@ -16,6 +16,7 @@ fn write_chunk(path: &str, schema: Schema, chunk: Chunk<Box<dyn Array>>) -> Resu
         write_statistics: true,
         compression: CompressionOptions::Uncompressed,
         version: Version::V2,
+        data_pagesize_limit: None,
     };
 
     let iter = vec![Ok(chunk)];

examples/parquet_write_async.rs

@@ -17,6 +17,7 @@ async fn write_batch(path: &str, schema: Schema, columns: Chunk<Box<dyn Array>>)
         write_statistics: true,
         compression: CompressionOptions::Uncompressed,
         version: Version::V2,
+        data_pagesize_limit: None,
     };
 
     let mut stream = futures::stream::iter(vec![Ok(columns)].into_iter());

examples/parquet_write_parallel/src/main.rs

@@ -47,6 +47,7 @@ fn parallel_write(path: &str, schema: Schema, chunks: &[Chunk]) -> Result<()> {
         write_statistics: true,
         compression: CompressionOptions::Snappy,
         version: Version::V2,
+        data_pagesize_limit: None,
     };
 
     let encoding_map = |data_type: &DataType| {

src/doc/lib.md

@@ -42,6 +42,7 @@ fn main() -> Result<()> {
         write_statistics: true,
         compression: CompressionOptions::Snappy,
         version: Version::V1,
+        data_pagesize_limit: None,
     };
 
     let row_groups = RowGroupIterator::try_new(

src/io/parquet/write/mod.rs

@@ -56,6 +56,8 @@ pub struct WriteOptions {
     pub version: Version,
     /// The compression to apply to every page
     pub compression: CompressionOptions,
+    /// The size to flush a page, defaults to 1024 * 1024 if zero
+    pub data_pagesize_limit: Option<usize>,
 }
 
 use crate::compute::aggregate::estimated_bytes_size;
@@ -141,6 +143,7 @@ pub fn array_to_pages(
     // we also check for an array.len > 3 to prevent infinite recursion
     // still have to figure out how to deal with values that are i32::MAX size, such as very large
     // strings or a list column with many elements
+
     if (estimated_bytes_size(array)) >= (2u32.pow(31) - 2u32.pow(25)) as usize && array.len() > 3 {
         let split_at = array.len() / 2;
         let left = array.slice(0, split_at);
@@ -163,8 +166,31 @@ pub fn array_to_pages(
                     )
                 })
             }
-            _ => array_to_page(array, type_, nested, options, encoding)
-                .map(|page| DynIter::new(std::iter::once(Ok(page)))),
+            _ => {
+                const DEFAULT_PAGE_SIZE: usize = 1024 * 1024;
+                let page_size = options.data_pagesize_limit.unwrap_or(DEFAULT_PAGE_SIZE);
+
+                let bytes_per_row = estimated_bytes_size(array.slice(0, 1).as_ref()).max(1);
+                let row_per_page = page_size / bytes_per_row + 1;
+
+                if array.len() <= row_per_page {
+                    array_to_page(array, type_, nested, options, encoding)
+                        .map(|page| DynIter::new(std::iter::once(Ok(page))))
+                } else {
+                    let first = array.slice(0, row_per_page);
+                    let other = array.slice(row_per_page, array.len() - row_per_page);
+                    Ok(DynIter::new(
+                        array_to_pages(first.as_ref(), type_.clone(), nested, options, encoding)?
+                            .chain(array_to_pages(
+                                other.as_ref(),
+                                type_,
+                                nested,
+                                options,
+                                encoding,
+                            )?),
+                    ))
+                }
+            }
         }
     }
 }

tests/it/io/parquet/mod.rs

@@ -1128,6 +1128,7 @@ fn integration_write(schema: &Schema, chunks: &[Chunk<Box<dyn Array>>]) -> Resul
         write_statistics: true,
         compression: CompressionOptions::Uncompressed,
         version: Version::V1,
+        data_pagesize_limit: None,
     };
 
     let encodings = schema

tests/it/io/parquet/read_indexes.rs

@@ -30,6 +30,7 @@ fn pages(
         write_statistics: true,
         compression: CompressionOptions::Uncompressed,
         version: Version::V1,
+        data_pagesize_limit: None,
     };
 
     let pages1 = [array11, array12, array13]
@@ -79,6 +80,7 @@ fn read_with_indexes(
         write_statistics: true,
         compression: CompressionOptions::Uncompressed,
         version: Version::V1,
+        data_pagesize_limit: None,
     };
 
     let to_compressed = |pages: Vec<EncodedPage>| {

tests/it/io/parquet/write.rs

@@ -36,6 +36,7 @@ fn round_trip(
         write_statistics: true,
         compression,
         version,
+        data_pagesize_limit: None,
     };
 
     let iter = vec![Chunk::try_new(vec![array.clone()])];

tests/it/io/parquet/write_async.rs

@@ -31,6 +31,7 @@ async fn test_parquet_async_roundtrip() {
         write_statistics: true,
         compression: CompressionOptions::Uncompressed,
         version: Version::V2,
+        data_pagesize_limit: None,
     };
 
     let mut buffer = Cursor::new(Vec::new());