etseidl commented on code in PR #6000:
URL: https://github.com/apache/arrow-rs/pull/6000#discussion_r1688740876
##########
parquet/src/file/writer.rs:
##########
@@ -815,20 +735,299 @@ impl<'a, W: Write + Send> PageWriter for
SerializedPageWriter<'a, W> {
}
}
+/// Writes `crate::file::metadata` structures to a thrift encdoded byte streams
+///
+/// This structure handles the details of writing the various parts of parquet
+/// metadata into a byte stream. It is used to write the metadata into a
+/// parquet file and can also write metadata into other locations (such as a
+/// store of bytes).
+///
+/// This is somewhat trickey because the metadata is not store as a single
inline
+/// thrift struture. It can have several "out of band" structures such as the
OffsetIndex
+/// and BloomFilters which are stored separately whose locations are stored as
offsets
+struct ThriftMetadataWriter<'a, W: Write> {
+ buf: &'a mut TrackedWrite<W>,
+ schema: &'a TypePtr,
+ schema_descr: &'a SchemaDescPtr,
+ row_groups: Vec<RowGroup>,
+ column_indexes: Option<&'a [Vec<Option<ColumnIndex>>]>,
+ offset_indexes: Option<&'a [Vec<Option<OffsetIndex>>]>,
+ key_value_metadata: Option<Vec<KeyValue>>,
+ created_by: Option<String>,
+ writer_version: i32,
+}
+
+impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
+ /// Serialize all the offset index to the file
+ fn write_offset_indexes(&mut self, offset_indexes:
&[Vec<Option<OffsetIndex>>]) -> Result<()> {
+ // iter row group
+ // iter each column
+ // write offset index to the file
+ for (row_group_idx, row_group) in
self.row_groups.iter_mut().enumerate() {
+ for (column_idx, column_metadata) in
row_group.columns.iter_mut().enumerate() {
+ match &offset_indexes[row_group_idx][column_idx] {
+ Some(offset_index) => {
+ let start_offset = self.buf.bytes_written();
+ let mut protocol = TCompactOutputProtocol::new(&mut
self.buf);
+ offset_index.write_to_out_protocol(&mut protocol)?;
+ let end_offset = self.buf.bytes_written();
+ // set offset and index for offset index
+ column_metadata.offset_index_offset =
Some(start_offset as i64);
+ column_metadata.offset_index_length =
+ Some((end_offset - start_offset) as i32);
+ }
+ None => {}
+ }
+ }
+ }
+ Ok(())
+ }
+
+ /// Serialize all the column index to the file
+ fn write_column_indexes(&mut self, column_indexes:
&[Vec<Option<ColumnIndex>>]) -> Result<()> {
+ // iter row group
+ // iter each column
+ // write column index to the file
+ for (row_group_idx, row_group) in
self.row_groups.iter_mut().enumerate() {
+ for (column_idx, column_metadata) in
row_group.columns.iter_mut().enumerate() {
+ match &column_indexes[row_group_idx][column_idx] {
+ Some(column_index) => {
+ let start_offset = self.buf.bytes_written();
+ let mut protocol = TCompactOutputProtocol::new(&mut
self.buf);
+ column_index.write_to_out_protocol(&mut protocol)?;
+ let end_offset = self.buf.bytes_written();
+ // set offset and index for offset index
+ column_metadata.column_index_offset =
Some(start_offset as i64);
+ column_metadata.column_index_length =
+ Some((end_offset - start_offset) as i32);
+ }
+ None => {}
+ }
+ }
+ }
+ Ok(())
+ }
+
+ /// Assembles and writes the final metadata to self.buf
+ pub fn finish(mut self) -> Result<parquet::FileMetaData> {
+ let num_rows = self.row_groups.iter().map(|x| x.num_rows).sum();
+
+ // Write column indexes and offset indexes
+ if let Some(column_indexes) = self.column_indexes {
+ self.write_column_indexes(column_indexes)?;
+ }
+ if let Some(offset_indexes) = self.offset_indexes {
+ self.write_offset_indexes(offset_indexes)?;
+ }
+
+ // We only include ColumnOrder for leaf nodes.
+ // Currently only supported ColumnOrder is TypeDefinedOrder so we set
this
+ // for all leaf nodes.
+ // Even if the column has an undefined sort order, such as INTERVAL,
this
+ // is still technically the defined TYPEORDER so it should still be
set.
+ let column_orders = (0..self.schema_descr.num_columns())
+ .map(|_| parquet::ColumnOrder::TYPEORDER(parquet::TypeDefinedOrder
{}))
+ .collect();
+ // This field is optional, perhaps in cases where no min/max fields
are set
+ // in any Statistics or ColumnIndex object in the whole file.
+ // But for simplicity we always set this field.
+ let column_orders = Some(column_orders);
+
+ let file_metadata = parquet::FileMetaData {
+ num_rows,
+ row_groups: self.row_groups,
+ key_value_metadata: self.key_value_metadata.clone(),
+ version: self.writer_version,
+ schema: types::to_thrift(self.schema.as_ref())?,
+ created_by: self.created_by.clone(),
+ column_orders,
+ encryption_algorithm: None,
+ footer_signing_key_metadata: None,
+ };
+
+ // Write file metadata
+ let start_pos = self.buf.bytes_written();
+ {
+ let mut protocol = TCompactOutputProtocol::new(&mut self.buf);
+ file_metadata.write_to_out_protocol(&mut protocol)?;
+ }
+ let end_pos = self.buf.bytes_written();
+
+ // Write footer
+ let metadata_len = (end_pos - start_pos) as u32;
+
+ self.buf.write_all(&metadata_len.to_le_bytes())?;
+ self.buf.write_all(&PARQUET_MAGIC)?;
+ Ok(file_metadata)
+ }
+
+ pub(self) fn new(
+ buf: &'a mut TrackedWrite<W>,
+ schema: &'a TypePtr,
+ schema_descr: &'a SchemaDescPtr,
+ row_groups: Vec<RowGroup>,
+ created_by: Option<String>,
+ writer_version: i32,
+ ) -> Self {
+ Self {
+ buf,
+ schema,
+ schema_descr,
+ row_groups,
+ column_indexes: None,
+ offset_indexes: None,
+ key_value_metadata: None,
+ created_by,
+ writer_version,
+ }
+ }
+
+ pub fn with_column_indexes(mut self, column_indexes: &'a
[Vec<Option<ColumnIndex>>]) -> Self {
+ self.column_indexes = Some(column_indexes);
+ self
+ }
+
+ pub fn with_offset_indexes(mut self, offset_indexes: &'a
[Vec<Option<OffsetIndex>>]) -> Self {
+ self.offset_indexes = Some(offset_indexes);
+ self
+ }
+
+ pub fn with_key_value_metadata(mut self, key_value_metadata:
Vec<KeyValue>) -> Self {
+ self.key_value_metadata = Some(key_value_metadata);
+ self
+ }
+}
+
+pub struct ParquetMetadataWriter<'a, W: Write> {
+ buf: TrackedWrite<W>,
+ write_page_index: bool,
+ metadata: &'a ParquetMetaData,
+}
+
+impl<'a, W: Write> ParquetMetadataWriter<'a, W> {
+ pub fn new(buf: W, metadata: &'a ParquetMetaData) -> Self {
+ Self {
+ buf: TrackedWrite::new(buf),
+ write_page_index: true,
+ metadata,
+ }
+ }
+
+ pub fn write_page_index(&mut self, write_page_index: bool) -> &mut Self {
+ self.write_page_index = write_page_index;
+ self
+ }
+
+ pub fn finish(&mut self) -> Result<()> {
+ let file_metadata = self.metadata.file_metadata();
+
+ let schema = Arc::new(file_metadata.schema().clone());
+ let schema_descr = Arc::new(SchemaDescriptor::new(schema.clone()));
+ let created_by = file_metadata.created_by().map(str::to_string);
+
+ let row_groups = self
+ .metadata
+ .row_groups()
+ .iter()
+ .map(|rg| rg.to_thrift())
+ .collect::<Vec<_>>();
+
+ let key_value_metadata = file_metadata.key_value_metadata().cloned();
+
+ let column_indexes = self.convert_column_indexes();
+ let offset_indexes = self.convert_offset_index();
+
+ let mut encoder = ThriftMetadataWriter::new(
+ &mut self.buf,
+ &schema,
+ &schema_descr,
+ row_groups,
+ created_by,
+ file_metadata.version(),
+ );
+ encoder = encoder.with_column_indexes(&column_indexes);
+ encoder = encoder.with_offset_indexes(&offset_indexes);
+ if let Some(key_value_metadata) = key_value_metadata {
+ encoder = encoder.with_key_value_metadata(key_value_metadata);
+ }
+ encoder.finish()?;
+
+ Ok(())
+ }
+
+ fn convert_column_indexes(&self) -> Vec<Vec<Option<ColumnIndex>>> {
+ if let Some(row_group_column_indexes) = self.metadata.column_index() {
+ (0..self.metadata.row_groups().len())
+ .map(|rg_idx| {
+ let column_indexes = &row_group_column_indexes[rg_idx];
+ column_indexes
+ .iter()
+ .map(|column_index| match column_index {
+ Index::NONE => None,
+ Index::BOOLEAN(column_index) =>
Some(column_index.to_thrift()),
+ Index::BYTE_ARRAY(column_index) =>
Some(column_index.to_thrift()),
+ Index::DOUBLE(column_index) =>
Some(column_index.to_thrift()),
+ Index::FIXED_LEN_BYTE_ARRAY(column_index) => {
+ Some(column_index.to_thrift())
+ }
+ Index::FLOAT(column_index) =>
Some(column_index.to_thrift()),
+ Index::INT32(column_index) =>
Some(column_index.to_thrift()),
+ Index::INT64(column_index) =>
Some(column_index.to_thrift()),
+ Index::INT96(column_index) =>
Some(column_index.to_thrift()),
+ })
+ .collect()
+ })
+ .collect()
+ } else {
+ // make a None for each row group, for each column
+ self.metadata
+ .row_groups()
+ .iter()
+ .map(|rg|
std::iter::repeat(None).take(rg.columns().len()).collect())
+ .collect()
+ }
+ }
+
+ fn convert_offset_index(&self) -> Vec<Vec<Option<OffsetIndex>>> {
+ if let Some(row_group_offset_indexes) = self.metadata.offset_index() {
+ (0..self.metadata.row_groups().len())
+ .map(|rg_idx| {
+ let offset_indexes = &row_group_offset_indexes[rg_idx];
+ offset_indexes
+ .iter()
+ .map(|column_index|
Some(OffsetIndex::new(column_index.clone(), None)))
Review Comment:
After merging with the latest 53.0.0-dev, this can then become:
```suggestion
.map(|offset_index| Some(offset_index.to_thrift()))
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
