etseidl commented on code in PR #6000:
URL: https://github.com/apache/arrow-rs/pull/6000#discussion_r1676979247
##########
parquet/src/file/writer.rs:
##########
@@ -791,23 +710,274 @@ impl<'a, W: Write + Send> PageWriter for
SerializedPageWriter<'a, W> {
}
}
+struct ThriftMetadataWriter<'a, W: Write> {
+ buf: &'a mut TrackedWrite<W>,
+ schema: &'a TypePtr,
+ schema_descr: &'a SchemaDescPtr,
+ row_groups: Vec<RowGroup>,
+ column_indexes: &'a [Vec<Option<ColumnIndex>>],
+ offset_indexes: &'a [Vec<Option<OffsetIndex>>],
+ key_value_metadata: &'a 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) -> 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 &self.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) -> 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 &self.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 metadata at the end of the file.
+ 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
+ self.write_column_indexes()?;
+ self.write_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>,
+ column_indexes: &'a [Vec<Option<ColumnIndex>>],
+ offset_indexes: &'a [Vec<Option<OffsetIndex>>],
+ key_value_metadata: &'a Option<Vec<KeyValue>>,
+ created_by: Option<String>,
+ writer_version: i32,
+ ) -> Self {
+ Self {
+ buf,
+ schema,
+ schema_descr,
+ row_groups,
+ column_indexes,
+ offset_indexes,
+ key_value_metadata,
+ created_by,
+ writer_version,
+ }
+ }
+}
+
+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 encoder = ThriftMetadataWriter::new(
+ &mut self.buf,
+ &schema,
+ &schema_descr,
+ row_groups,
+ &column_indexes,
+ &offset_indexes,
+ &key_value_metadata,
+ created_by,
+ file_metadata.version(),
+ );
+ 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_column_index()),
+ Index::BYTE_ARRAY(column_index) =>
Some(column_index.to_column_index()),
+ Index::DOUBLE(column_index) =>
Some(column_index.to_column_index()),
+ Index::FIXED_LEN_BYTE_ARRAY(column_index) => {
+ Some(column_index.to_column_index())
+ }
+ Index::FLOAT(column_index) =>
Some(column_index.to_column_index()),
+ Index::INT32(column_index) =>
Some(column_index.to_column_index()),
+ Index::INT64(column_index) =>
Some(column_index.to_column_index()),
+ Index::INT96(column_index) =>
Some(column_index.to_column_index()),
+ })
+ .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 column_indexes = &row_group_offset_indexes[rg_idx];
Review Comment:
Minor nit: could this be named `offset_indexes`?
--
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]
