Add a C abi safe function to dest api to get bytes written so far

Author: aykut-bozkurt

src/lib.rs

@@ -19,6 +19,8 @@ mod type_compat;
 pub use crate::arrow_parquet::compression::PgParquetCompression;
 #[allow(unused_imports)]
 pub use crate::parquet_copy_hook::copy_to_split_dest_receiver::create_copy_to_parquet_split_dest_receiver;
+#[allow(unused_imports)]
+pub use crate::parquet_copy_hook::copy_to_split_dest_receiver::parquet_split_dest_receiver_bytes_written;
 
 pgrx::pg_module_magic!();
 

src/parquet_copy_hook/copy_to_dest_receiver.rs

@@ -91,6 +91,10 @@ impl CopyToParquetDestReceiver {
     }
 
     pub(crate) fn collected_bytes(&self) -> usize {
+        if self.parquet_writer_context.is_null() {
+            return 0;
+        }
+
         let current_parquet_writer_context = unsafe {
             self.parquet_writer_context
                 .as_ref()

src/parquet_copy_hook/copy_to_split_dest_receiver.rs

@@ -27,6 +27,7 @@ struct CopyToParquetSplitDestReceiver {
     options: CopyToParquetOptions,
     current_child_id: i64,
     current_child_receiver: *mut CopyToParquetDestReceiver,
+    bytes_written: u64,
 }
 
 #[repr(C)]
@@ -90,6 +91,11 @@ impl CopyToParquetSplitDestReceiver {
             }
         }
 
+        self.bytes_written += unsafe {
+            PgBox::<CopyToParquetDestReceiver>::from_pg(self.current_child_receiver as _)
+                .collected_bytes() as u64
+        };
+
         self.current_child_receiver = std::ptr::null_mut();
     }
 
@@ -140,6 +146,18 @@ impl CopyToParquetSplitDestReceiver {
 
         child_uri.to_str().expect("invalid uri").as_pg_cstr()
     }
+
+    fn bytes_written(&self) -> u64 {
+        if self.current_child_receiver.is_null() {
+            return self.bytes_written;
+        }
+
+        let child_parquet_dest = unsafe {
+            PgBox::<CopyToParquetDestReceiver>::from_pg(self.current_child_receiver as _)
+        };
+
+        self.bytes_written + child_parquet_dest.collected_bytes() as u64
+    }
 }
 
 #[pg_guard]
@@ -265,6 +283,21 @@ pub extern "C" fn create_copy_to_parquet_split_dest_receiver(
     split_dest.operation = -1;
     split_dest.options = options;
     split_dest.current_child_id = 0;
+    split_dest.bytes_written = 0;
 
     unsafe { std::mem::transmute(split_dest) }
 }
+
+// parquet_split_dest_receiver_bytes_written returns the total number of bytes written to parquet files
+// so far. This includes all child receivers that have been created and flushed.
+#[pg_guard]
+#[no_mangle]
+pub extern "C" fn parquet_split_dest_receiver_bytes_written(dest: *const DestReceiver) -> u64 {
+    let split_parquet_dest = unsafe {
+        (dest as *const CopyToParquetSplitDestReceiver)
+            .as_ref()
+            .expect("invalid split parquet dest receiver ptr")
+    };
+
+    split_parquet_dest.bytes_written()
+}