Access metadata of flushed row groups on write (apache#1691)

tustvold · Jun 2, 2022 · 4cad205 · 4cad205
1 parent c1a91dc
commit 4cad205
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/parquet/src/arrow/arrow_writer.rs b/parquet/src/arrow/arrow_writer.rs
@@ -36,6 +36,7 @@ use super::schema::{
 use crate::arrow::levels::calculate_array_levels;
 use crate::column::writer::ColumnWriter;
 use crate::errors::{ParquetError, Result};
+use crate::file::metadata::RowGroupMetaDataPtr;
 use crate::file::properties::WriterProperties;
 use crate::file::writer::{SerializedColumnWriter, SerializedRowGroupWriter};
 use crate::{data_type::*, file::writer::SerializedFileWriter};
@@ -95,6 +96,11 @@ impl<W: Write> ArrowWriter<W> {
         })
     }
 
+    /// Returns metadata for any flushed row groups
+    pub fn flushed_row_groups(&self) -> &[RowGroupMetaDataPtr] {
+        self.writer.flushed_row_groups()
+    }
+
     /// Enqueues the provided `RecordBatch` to be written
     ///
     /// If following this there are more than `max_row_group_size` rows buffered,

diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs
@@ -153,6 +153,11 @@ impl<W: Write> SerializedFileWriter<W> {
         Ok(row_group_writer)
     }
 
+    /// Returns metadata for any flushed row groups
+    pub fn flushed_row_groups(&self) -> &[RowGroupMetaDataPtr] {
+        &self.row_groups
+    }
+
     /// Closes and finalises file writer, returning the file metadata.
     ///
     /// All row groups must be appended before this method is called.