Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: add append_key_value_metadata #3367

Merged
merged 7 commits into from
Dec 20, 2022
Merged

feat: add append_key_value_metadata #3367

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b8d01a4
Add update_key_value_metadata
jiacai2050 Dec 19, 2022
ceb6bdd
Add comments
jiacai2050 Dec 19, 2022
f5c5a37
Address review
jiacai2050 Dec 19, 2022
6070937
fix clippy
jiacai2050 Dec 19, 2022
3f1cad4
Update parquet/src/arrow/arrow_writer/mod.rs
jiacai2050 Dec 20, 2022
36bb6f5
Fix reviews
jiacai2050 Dec 20, 2022
a9313a7
Test and fix
tustvold Dec 20, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
9 changes: 8 additions & 1 deletion parquet/src/arrow/arrow_writer/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -158,6 +158,13 @@ impl<W: Write> ArrowWriter<W> {
self.flush_rows(self.buffered_rows)
}

/// Update key value metadata.
jiacai2050 marked this conversation as resolved.
Show resolved Hide resolved
///
/// 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 {
Expand Down
12 changes: 11 additions & 1 deletion parquet/src/file/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand All @@ -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(),
})
}

Expand Down Expand Up @@ -298,7 +301,10 @@ 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(|mut v| {
v.extend(self.kv_metadatas.clone());
v
}),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
key_value_metadata: self.props.key_value_metadata().cloned().map(|mut v| {
v.extend(self.kv_metadatas.clone());
v
}),
key_value_metadata: self.props.key_value_metadata().iter().chain(&self.kv_metadatas).cloned().collect(),

Or something, perhaps?

created_by: Some(self.props.created_by().to_owned()),
column_orders: None,
encryption_algorithm: None,
Expand Down Expand Up @@ -331,6 +337,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()?;
Expand Down