feat: add append_key_value_metadata

parquet/src/arrow/arrow_writer/mod.rs

@@ -32,7 +32,7 @@ use super::schema::{
 use crate::arrow::arrow_writer::byte_array::ByteArrayWriter;
 use crate::column::writer::{ColumnWriter, ColumnWriterImpl};
 use crate::errors::{ParquetError, Result};
-use crate::file::metadata::RowGroupMetaDataPtr;
+use crate::file::metadata::{KeyValue, RowGroupMetaDataPtr};
 use crate::file::properties::WriterProperties;
 use crate::file::writer::SerializedRowGroupWriter;
 use crate::{data_type::*, file::writer::SerializedFileWriter};
@@ -158,6 +158,13 @@ impl<W: Write> ArrowWriter<W> {
         self.flush_rows(self.buffered_rows)
     }
 
+    /// Additional [`KeyValue`] metadata to be written in addition to those from [`WriterProperties`]
+    ///
+    /// This method provide a way to append kv_metadata after write RecordBatch
+    pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
+        self.writer.append_key_value_metadata(kv_metadata)
+    }
+
     /// Flushes `num_rows` from the buffer into a new row group
     fn flush_rows(&mut self, num_rows: usize) -> Result<()> {
         if num_rows == 0 {

parquet/src/file/writer.rs

@@ -122,6 +122,8 @@ pub struct SerializedFileWriter<W: Write> {
     column_indexes: Vec<Vec<Option<ColumnIndex>>>,
     offset_indexes: Vec<Vec<Option<OffsetIndex>>>,
     row_group_index: usize,
+    // kv_metadatas will be appended to `props` when `write_metadata`
+    kv_metadatas: Vec<KeyValue>,
 }
 
 impl<W: Write> SerializedFileWriter<W> {
@@ -139,6 +141,7 @@ impl<W: Write> SerializedFileWriter<W> {
             column_indexes: Vec::new(),
             offset_indexes: Vec::new(),
             row_group_index: 0,
+            kv_metadatas: Vec::new(),
         })
     }
 
@@ -298,7 +301,11 @@ impl<W: Write> SerializedFileWriter<W> {
             row_groups,
             version: self.props.writer_version().as_num(),
             schema: types::to_thrift(self.schema.as_ref())?,
-            key_value_metadata: self.props.key_value_metadata().cloned(),
+            key_value_metadata: self.props.key_value_metadata().cloned().map(|kv| {
+                kv.into_iter()
+                    .chain(self.kv_metadatas.clone().into_iter())
+                    .collect()
+            }),
             created_by: Some(self.props.created_by().to_owned()),
             column_orders: None,
             encryption_algorithm: None,
@@ -331,6 +338,10 @@ impl<W: Write> SerializedFileWriter<W> {
         }
     }
 
+    pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
+        self.kv_metadatas.push(kv_metadata);
+    }
+
     /// Writes the file footer and returns the underlying writer.
     pub fn into_inner(mut self) -> Result<W> {
         self.assert_previous_writer_closed()?;