do arrow-parquet conversions per batch size

aykut-bozkurt · aykut-bozkurt · commit 5b5725db5cc6 · 2025-03-06T12:32:29.000+03:00
diff --git a/src/arrow_parquet/parquet_writer.rs b/src/arrow_parquet/parquet_writer.rs
@@ -11,13 +11,14 @@ use pgrx::{heap_tuple::PgHeapTuple, AllocatedByRust, PgTupleDesc};
 
 use crate::{
     arrow_parquet::{
-        compression::{PgParquetCompression, PgParquetCompressionWithLevel},
+        compression::PgParquetCompressionWithLevel,
         pg_to_arrow::context::collect_pg_to_arrow_attribute_contexts,
         schema_parser::{
             parquet_schema_string_from_attributes, parse_arrow_schema_from_attributes,
         },
         uri_utils::parquet_writer_from_uri,
     },
+    parquet_copy_hook::copy_to_split_dest_receiver::CopyToParquetOptions,
     pgrx_utils::{collect_attributes_for, CollectAttributesFor},
     type_compat::{
         geometry::{geoparquet_metadata_json_from_tupledesc, reset_postgis_context},
@@ -38,13 +39,13 @@ pub(crate) struct ParquetWriterContext {
     parquet_writer: AsyncArrowWriter<ParquetObjectWriter>,
     schema: SchemaRef,
     attribute_contexts: Vec<PgToArrowAttributeContext>,
+    options: CopyToParquetOptions,
 }
 
 impl ParquetWriterContext {
     pub(crate) fn new(
         uri_info: ParsedUriInfo,
-        compression: PgParquetCompression,
-        compression_level: i32,
+        options: CopyToParquetOptions,
         tupledesc: &PgTupleDesc,
     ) -> ParquetWriterContext {
         // Postgis and Map contexts are used throughout writing the parquet file.
@@ -62,7 +63,7 @@ impl ParquetWriterContext {
         let schema = parse_arrow_schema_from_attributes(&attributes);
         let schema = Arc::new(schema);
 
-        let writer_props = Self::writer_props(tupledesc, compression, compression_level);
+        let writer_props = Self::writer_props(tupledesc, options);
 
         let parquet_writer = parquet_writer_from_uri(uri_info, schema.clone(), writer_props);
 
@@ -73,22 +74,20 @@ impl ParquetWriterContext {
             parquet_writer,
             schema,
             attribute_contexts,
+            options,
         }
     }
 
-    fn writer_props(
-        tupledesc: &PgTupleDesc,
-        compression: PgParquetCompression,
-        compression_level: i32,
-    ) -> WriterProperties {
+    fn writer_props(tupledesc: &PgTupleDesc, options: CopyToParquetOptions) -> WriterProperties {
         let compression = PgParquetCompressionWithLevel {
-            compression,
-            compression_level,
+            compression: options.compression,
+            compression_level: options.compression_level,
         };
 
         let mut writer_props_builder = WriterProperties::builder()
             .set_statistics_enabled(EnabledStatistics::Page)
             .set_compression(compression.into())
+            .set_max_row_group_size(options.row_group_size as usize)
             .set_created_by("pg_parquet".to_string());
 
         let geometry_columns_metadata_value = geoparquet_metadata_json_from_tupledesc(tupledesc);
@@ -103,10 +102,9 @@ impl ParquetWriterContext {
         writer_props_builder.build()
     }
 
-    pub(crate) fn write_new_row_group(
-        &mut self,
-        tuples: Vec<Option<PgHeapTuple<AllocatedByRust>>>,
-    ) {
+    // write_tuples writes the tuples to the parquet file. It flushes the in progress rows to a new row group
+    // if the row group size is reached.
+    pub(crate) fn write_tuples(&mut self, tuples: Vec<Option<PgHeapTuple<AllocatedByRust>>>) {
         let record_batch =
             Self::pg_tuples_to_record_batch(tuples, &self.attribute_contexts, self.schema.clone());
 
@@ -116,9 +114,20 @@ impl ParquetWriterContext {
             .block_on(parquet_writer.write(&record_batch))
             .unwrap_or_else(|e| panic!("failed to write record batch: {}", e));
 
+        if parquet_writer.in_progress_rows() >= self.options.row_group_size as _
+            || parquet_writer.in_progress_size() >= self.options.row_group_size_bytes as _
+        {
+            PG_BACKEND_TOKIO_RUNTIME
+                .block_on(parquet_writer.flush())
+                .unwrap_or_else(|e| panic!("failed to flush record batch: {}", e));
+        }
+    }
+
+    // finalize flushes the in progress rows to a new row group and finally writes metadata to the file.
+    fn finalize(&mut self) {
         PG_BACKEND_TOKIO_RUNTIME
-            .block_on(parquet_writer.flush())
-            .unwrap_or_else(|e| panic!("failed to flush record batch: {}", e));
+            .block_on(self.parquet_writer.finish())
+            .unwrap_or_else(|e| panic!("failed to finish parquet writer: {}", e));
     }
 
     pub(crate) fn bytes_written(&self) -> usize {
@@ -144,10 +153,6 @@ impl ParquetWriterContext {
 
 impl Drop for ParquetWriterContext {
     fn drop(&mut self) {
-        PG_BACKEND_TOKIO_RUNTIME
-            .block_on(self.parquet_writer.finish())
-            .unwrap_or_else(|e| {
-                panic!("failed to close parquet writer: {}", e);
-            });
+        self.finalize();
     }
 }
diff --git a/src/arrow_parquet/uri_utils.rs b/src/arrow_parquet/uri_utils.rs
@@ -19,7 +19,6 @@ use pgrx::{
 use url::Url;
 
 use crate::{
-    arrow_parquet::parquet_writer::DEFAULT_ROW_GROUP_SIZE,
     object_store::{
         aws::parse_s3_bucket, azure::parse_azure_blob_container,
         object_store_cache::get_or_create_object_store,
@@ -139,6 +138,9 @@ pub(crate) fn parquet_metadata_from_uri(uri_info: ParsedUriInfo) -> Arc<ParquetM
     })
 }
 
+// default # of records per batch during arrow-parquet conversions (RecordBatch api)
+pub(crate) const RECORD_BATCH_SIZE: i64 = 1024;
+
 pub(crate) fn parquet_reader_from_uri(
     uri_info: ParsedUriInfo,
 ) -> ParquetRecordBatchStream<ParquetObjectReader> {
@@ -165,13 +167,36 @@ pub(crate) fn parquet_reader_from_uri(
 
         pgrx::debug2!("Converted arrow schema is: {}", builder.schema());
 
+        let batch_size = calculate_reader_batch_size(builder.metadata());
+
         builder
-            .with_batch_size(DEFAULT_ROW_GROUP_SIZE as usize)
+            .with_batch_size(batch_size)
             .build()
             .unwrap_or_else(|e| panic!("{}", e))
     })
 }
 
+fn calculate_reader_batch_size(metadata: &Arc<ParquetMetaData>) -> usize {
+    const MAX_ARROW_ARRAY_SIZE: i64 = i32::MAX as _;
+
+    for row_group in metadata.row_groups() {
+        for column in row_group.columns() {
+            // try our best to get the size of the column
+            let column_size = column
+                .unencoded_byte_array_data_bytes()
+                .unwrap_or(column.uncompressed_size());
+
+            if column_size > MAX_ARROW_ARRAY_SIZE {
+                // to prevent decoding large arrays into memory, process one row at a time
+                return 1;
+            }
+        }
+    }
+
+    // default batch size
+    RECORD_BATCH_SIZE as _
+}
+
 pub(crate) fn parquet_writer_from_uri(
     uri_info: ParsedUriInfo,
     arrow_schema: SchemaRef,
diff --git a/src/parquet_copy_hook/copy_to_dest_receiver.rs b/src/parquet_copy_hook/copy_to_dest_receiver.rs
@@ -9,7 +9,10 @@ use pg_sys::{
 };
 use pgrx::{prelude::*, FromDatum, PgList, PgMemoryContexts, PgTupleDesc};
 
-use crate::arrow_parquet::{parquet_writer::ParquetWriterContext, uri_utils::parse_uri};
+use crate::arrow_parquet::{
+    parquet_writer::ParquetWriterContext,
+    uri_utils::{ParsedUriInfo, RECORD_BATCH_SIZE},
+};
 
 use super::copy_to_split_dest_receiver::CopyToParquetOptions;
 
@@ -22,6 +25,7 @@ pub(crate) struct CopyToParquetDestReceiver {
     collected_tuple_count: i64,
     collected_tuple_size: i64,
     collected_tuple_column_sizes: *mut i64,
+    target_batch_size: i64,
     uri: *const c_char,
     copy_options: CopyToParquetOptions,
     per_copy_context: MemoryContext,
@@ -64,14 +68,6 @@ impl CopyToParquetDestReceiver {
         };
     }
 
-    fn collected_tuples_exceeds_row_group_size(&self) -> bool {
-        self.collected_tuple_count >= self.copy_options.row_group_size
-    }
-
-    fn collected_tuples_exceeds_row_group_size_bytes(&self) -> bool {
-        self.collected_tuple_size >= self.copy_options.row_group_size_bytes
-    }
-
     fn collected_tuples_exceeds_max_col_size(&self, tuple_column_sizes: &[i32]) -> bool {
         const MAX_ARROW_ARRAY_SIZE: i64 = i32::MAX as _;
 
@@ -118,7 +114,7 @@ impl CopyToParquetDestReceiver {
                 .as_mut()
                 .expect("parquet writer context is not found")
         };
-        current_parquet_writer_context.write_new_row_group(tuples);
+        current_parquet_writer_context.write_tuples(tuples);
 
         self.reset_collected_tuples();
     }
@@ -171,6 +167,13 @@ pub(crate) extern "C" fn copy_startup(
     };
     parquet_dest.natts = tupledesc.len();
 
+    parquet_dest.target_batch_size = if parquet_dest.copy_options.row_group_size < RECORD_BATCH_SIZE
+    {
+        parquet_dest.copy_options.row_group_size
+    } else {
+        RECORD_BATCH_SIZE
+    };
+
     let uri = unsafe { CStr::from_ptr(parquet_dest.uri) }
         .to_str()
         .expect("uri is not a valid C string");
@@ -179,13 +182,9 @@ pub(crate) extern "C" fn copy_startup(
         panic!("{}", e.to_string());
     });
 
-    let compression = parquet_dest.copy_options.compression;
-
-    let compression_level = parquet_dest.copy_options.compression_level;
-
     // leak the parquet writer context since it will be used during the COPY operation
     let parquet_writer_context =
-        ParquetWriterContext::new(uri_info, compression, compression_level, &tupledesc);
+        ParquetWriterContext::new(uri_info, parquet_dest.copy_options, &tupledesc);
     parquet_dest.parquet_writer_context = Box::into_raw(Box::new(parquet_writer_context));
 }
 
@@ -220,6 +219,9 @@ pub(crate) extern "C" fn copy_receive(slot: *mut TupleTableSlot, dest: *mut Dest
 
             let column_sizes = tuple_column_sizes(&datums, &tupledesc);
 
+            // we use arrow arrays as intermediate format when writing to parquet.
+            // To not hit into arrow array size limit, write the tuples before
+            // collecting new one into the batch.
             if parquet_dest.collected_tuples_exceeds_max_col_size(&column_sizes) {
                 parquet_dest.write_tuples_to_parquet();
             }
@@ -229,9 +231,7 @@ pub(crate) extern "C" fn copy_receive(slot: *mut TupleTableSlot, dest: *mut Dest
 
             parquet_dest.collect_tuple(heap_tuple, column_sizes);
 
-            if parquet_dest.collected_tuples_exceeds_row_group_size()
-                || parquet_dest.collected_tuples_exceeds_row_group_size_bytes()
-            {
+            if parquet_dest.collected_tuple_count == parquet_dest.target_batch_size {
                 parquet_dest.write_tuples_to_parquet();
             }
         });
@@ -329,6 +329,7 @@ pub(crate) fn create_copy_to_parquet_dest_receiver(
     parquet_dest.collected_tuple_count = 0;
     parquet_dest.collected_tuples = std::ptr::null_mut();
     parquet_dest.collected_tuple_column_sizes = std::ptr::null_mut();
+    parquet_dest.target_batch_size = 0;
     parquet_dest.copy_options = options;
     parquet_dest.per_copy_context = per_copy_context;
 
diff --git a/src/pgrx_tests/copy_options.rs b/src/pgrx_tests/copy_options.rs
@@ -331,7 +331,7 @@ mod tests {
             results
         });
 
-        assert_eq!(result_metadata, vec![2]);
+        assert_eq!(result_metadata, vec![1]);
     }
 
     #[pg_test]
@@ -381,7 +381,7 @@ mod tests {
 
         let id_bytes = 4;
         let name_bytes = 1;
-        let total_rows_size_bytes = (id_bytes + name_bytes) * 1_000_000;
+        let total_rows_size_bytes = (id_bytes + name_bytes) * 1024 * 1024;
 
         let row_group_size_bytes = total_rows_size_bytes / 10;
 
@@ -409,7 +409,7 @@ mod tests {
             results
         });
 
-        assert_eq!(result_metadata, vec![10]);
+        assert_eq!(result_metadata, vec![9]);
     }
 
     #[pg_test]
