Support file_size_bytes option

COPY TO parquet now supports a new option, called `file_size_bytes`, which lets you
generate parquet files with target size = `file_size_bytes`.

When a parquet file exceeds the target size, it will be flushed and a new parquet file
will be generated under a parent directory. (parent directory will be the path without
the parquet extension)

e.g.

```sql
COPY (select 'hellooooo' || i from generate_series(1, 1000000) i) to '/tmp/test.parquet' with (file_size_bytes 1048576);
```

```bash
> ls -alh /tmp/test/
1.4M data_0.parquet
1.4M data_1.parquet
1.4M data_2.parquet
1.4M data_3.parquet
114K data_4.parquet
```

Author: aykut-bozkurt

README.md

@@ -242,6 +242,7 @@ Supported authorization methods' priority order is shown below:
 ## Copy Options
 `pg_parquet` supports the following options in the `COPY TO` command:
 - `format parquet`: you need to specify this option to read or write Parquet files which does not end with `.parquet[.<compression>]` extension,
+- `file_size_bytes <int>`: the total byte size per Parquet file. When set, the parquet files, with target size, are created under parent directory (named the same as file name without file extension). By default, when not specified, a single file is generated without creating a parent folder.
 - `row_group_size <int>`: the number of rows in each row group while writing Parquet files. The default row group size is `122880`,
 - `row_group_size_bytes <int>`: the total byte size of rows in each row group while writing Parquet files. The default row group size bytes is `row_group_size * 1024`,
 - `compression <string>`: the compression format to use while writing Parquet files. The supported compression formats are `uncompressed`, `snappy`, `gzip`, `brotli`, `lz4`, `lz4raw` and `zstd`. The default compression format is `snappy`. If not specified, the compression format is determined by the file extension,

src/arrow_parquet/parquet_writer.rs

@@ -119,6 +119,10 @@ impl ParquetWriterContext {
             .unwrap_or_else(|e| panic!("failed to flush record batch: {}", e));
     }
 
+    pub(crate) fn bytes_written(&self) -> usize {
+        self.parquet_writer.bytes_written()
+    }
+
     fn pg_tuples_to_record_batch(
         tuples: Vec<Option<PgHeapTuple<AllocatedByRust>>>,
         attribute_contexts: &[PgToArrowAttributeContext],

src/lib.rs

@@ -18,7 +18,7 @@ mod type_compat;
 #[allow(unused_imports)]
 pub use crate::arrow_parquet::compression::PgParquetCompression;
 #[allow(unused_imports)]
-pub use crate::parquet_copy_hook::copy_to_dest_receiver::create_copy_to_parquet_dest_receiver;
+pub use crate::parquet_copy_hook::copy_to_split_dest_receiver::create_copy_to_parquet_split_dest_receiver;
 
 pgrx::pg_module_magic!();
 

src/object_store/local_file.rs

@@ -13,6 +13,10 @@ pub(crate) fn create_local_file_object_store(
     let path = uri_as_string(uri);
 
     if !copy_from {
+        // create parent folder if it doesn't exist
+        let parent = std::path::Path::new(&path).parent().unwrap();
+        std::fs::create_dir_all(parent).unwrap_or_else(|e| panic!("{}", e));
+
         // create or overwrite the local file
         std::fs::OpenOptions::new()
             .write(true)

src/parquet_copy_hook.rs

@@ -1,6 +1,7 @@
 pub(crate) mod copy_from;
 pub(crate) mod copy_to;
 pub(crate) mod copy_to_dest_receiver;
+pub(crate) mod copy_to_split_dest_receiver;
 pub(crate) mod copy_utils;
 pub(crate) mod hook;
 pub(crate) mod pg_compat;

src/parquet_copy_hook/copy_to_dest_receiver.rs

@@ -9,22 +9,12 @@ use pg_sys::{
 };
 use pgrx::{prelude::*, FromDatum, PgList, PgMemoryContexts, PgTupleDesc};
 
-use crate::arrow_parquet::{
-    compression::{PgParquetCompression, INVALID_COMPRESSION_LEVEL},
-    parquet_writer::{ParquetWriterContext, DEFAULT_ROW_GROUP_SIZE, DEFAULT_ROW_GROUP_SIZE_BYTES},
-    uri_utils::parse_uri,
-};
+use crate::arrow_parquet::{parquet_writer::ParquetWriterContext, uri_utils::parse_uri};
 
-#[repr(C)]
-struct CopyToParquetOptions {
-    pub row_group_size: i64,
-    pub row_group_size_bytes: i64,
-    pub compression: PgParquetCompression,
-    pub compression_level: i32,
-}
+use super::copy_to_split_dest_receiver::CopyToParquetOptions;
 
 #[repr(C)]
-struct CopyToParquetDestReceiver {
+pub(crate) struct CopyToParquetDestReceiver {
     dest: DestReceiver,
     natts: usize,
     tupledesc: TupleDesc,
@@ -95,6 +85,16 @@ impl CopyToParquetDestReceiver {
             .any(|size| size > MAX_ARROW_ARRAY_SIZE)
     }
 
+    pub(crate) fn collected_bytes(&self) -> usize {
+        let current_parquet_writer_context = unsafe {
+            self.parquet_writer_context
+                .as_ref()
+                .expect("parquet writer context is not found")
+        };
+
+        current_parquet_writer_context.bytes_written()
+    }
+
     fn write_tuples_to_parquet(&mut self) {
         debug_assert!(!self.tupledesc.is_null());
 
@@ -143,7 +143,11 @@ impl CopyToParquetDestReceiver {
 }
 
 #[pg_guard]
-extern "C" fn copy_startup(dest: *mut DestReceiver, _operation: i32, tupledesc: TupleDesc) {
+pub(crate) extern "C" fn copy_startup(
+    dest: *mut DestReceiver,
+    _operation: i32,
+    tupledesc: TupleDesc,
+) {
     let parquet_dest = unsafe {
         (dest as *mut CopyToParquetDestReceiver)
             .as_mut()
@@ -184,7 +188,7 @@ extern "C" fn copy_startup(dest: *mut DestReceiver, _operation: i32, tupledesc:
 }
 
 #[pg_guard]
-extern "C" fn copy_receive(slot: *mut TupleTableSlot, dest: *mut DestReceiver) -> bool {
+pub(crate) extern "C" fn copy_receive(slot: *mut TupleTableSlot, dest: *mut DestReceiver) -> bool {
     let parquet_dest = unsafe {
         (dest as *mut CopyToParquetDestReceiver)
             .as_mut()
@@ -235,7 +239,7 @@ extern "C" fn copy_receive(slot: *mut TupleTableSlot, dest: *mut DestReceiver) -
 }
 
 #[pg_guard]
-extern "C" fn copy_shutdown(dest: *mut DestReceiver) {
+pub(crate) extern "C" fn copy_shutdown(dest: *mut DestReceiver) {
     let parquet_dest = unsafe {
         (dest as *mut CopyToParquetDestReceiver)
             .as_mut()
@@ -250,79 +254,7 @@ extern "C" fn copy_shutdown(dest: *mut DestReceiver) {
 }
 
 #[pg_guard]
-extern "C" fn copy_destroy(_dest: *mut DestReceiver) {}
-
-// create_copy_to_parquet_dest_receiver creates a new CopyToParquetDestReceiver that can be
-// used as a destination receiver for COPY TO command. All arguments, except "uri", are optional
-// and have default values if not provided.
-#[pg_guard]
-#[no_mangle]
-pub extern "C" fn create_copy_to_parquet_dest_receiver(
-    uri: *const c_char,
-    row_group_size: *const i64,
-    row_group_size_bytes: *const i64,
-    compression: *const PgParquetCompression,
-    compression_level: *const i32,
-) -> *mut DestReceiver {
-    let per_copy_context = unsafe {
-        AllocSetContextCreateExtended(
-            CurrentMemoryContext as _,
-            "ParquetCopyDestReceiver".as_pg_cstr(),
-            ALLOCSET_DEFAULT_MINSIZE as _,
-            ALLOCSET_DEFAULT_INITSIZE as _,
-            ALLOCSET_DEFAULT_MAXSIZE as _,
-        )
-    };
-
-    let row_group_size = if row_group_size.is_null() {
-        DEFAULT_ROW_GROUP_SIZE
-    } else {
-        unsafe { *row_group_size }
-    };
-
-    let row_group_size_bytes = if row_group_size_bytes.is_null() {
-        DEFAULT_ROW_GROUP_SIZE_BYTES
-    } else {
-        unsafe { *row_group_size_bytes }
-    };
-
-    let compression = if compression.is_null() {
-        PgParquetCompression::default()
-    } else {
-        unsafe { *compression }
-    };
-
-    let compression_level = if compression_level.is_null() {
-        compression
-            .default_compression_level()
-            .unwrap_or(INVALID_COMPRESSION_LEVEL)
-    } else {
-        unsafe { *compression_level }
-    };
-
-    let mut parquet_dest =
-        unsafe { PgBox::<CopyToParquetDestReceiver, AllocatedByPostgres>::alloc0() };
-
-    parquet_dest.dest.receiveSlot = Some(copy_receive);
-    parquet_dest.dest.rStartup = Some(copy_startup);
-    parquet_dest.dest.rShutdown = Some(copy_shutdown);
-    parquet_dest.dest.rDestroy = Some(copy_destroy);
-    parquet_dest.dest.mydest = CommandDest::DestCopyOut;
-    parquet_dest.uri = uri;
-    parquet_dest.tupledesc = std::ptr::null_mut();
-    parquet_dest.parquet_writer_context = std::ptr::null_mut();
-    parquet_dest.natts = 0;
-    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.copy_options.row_group_size = row_group_size;
-    parquet_dest.copy_options.row_group_size_bytes = row_group_size_bytes;
-    parquet_dest.copy_options.compression = compression;
-    parquet_dest.copy_options.compression_level = compression_level;
-    parquet_dest.per_copy_context = per_copy_context;
-
-    unsafe { std::mem::transmute(parquet_dest) }
-}
+pub(crate) extern "C" fn copy_destroy(_dest: *mut DestReceiver) {}
 
 fn tuple_column_sizes(tuple_datums: &[Option<Datum>], tupledesc: &PgTupleDesc) -> Vec<i32> {
     let mut column_sizes = vec![];
@@ -361,3 +293,42 @@ fn tuple_column_sizes(tuple_datums: &[Option<Datum>], tupledesc: &PgTupleDesc) -
 
     column_sizes
 }
+
+// create_copy_to_parquet_dest_receiver creates a new CopyToParquetDestReceiver that can be
+// used as a destination receiver for COPY TO command. All arguments, except "uri", are optional
+// and have default values if not provided.
+#[pg_guard]
+pub(crate) fn create_copy_to_parquet_dest_receiver(
+    uri: *const c_char,
+    options: CopyToParquetOptions,
+) -> *mut CopyToParquetDestReceiver {
+    let per_copy_context = unsafe {
+        AllocSetContextCreateExtended(
+            CurrentMemoryContext as _,
+            "ParquetCopyDestReceiver".as_pg_cstr(),
+            ALLOCSET_DEFAULT_MINSIZE as _,
+            ALLOCSET_DEFAULT_INITSIZE as _,
+            ALLOCSET_DEFAULT_MAXSIZE as _,
+        )
+    };
+
+    let mut parquet_dest =
+        unsafe { PgBox::<CopyToParquetDestReceiver, AllocatedByPostgres>::alloc0() };
+
+    parquet_dest.dest.receiveSlot = Some(copy_receive);
+    parquet_dest.dest.rStartup = Some(copy_startup);
+    parquet_dest.dest.rShutdown = Some(copy_shutdown);
+    parquet_dest.dest.rDestroy = Some(copy_destroy);
+    parquet_dest.dest.mydest = CommandDest::DestCopyOut;
+    parquet_dest.uri = uri;
+    parquet_dest.tupledesc = std::ptr::null_mut();
+    parquet_dest.parquet_writer_context = std::ptr::null_mut();
+    parquet_dest.natts = 0;
+    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.copy_options = options;
+    parquet_dest.per_copy_context = per_copy_context;
+
+    parquet_dest.into_pg()
+}