bug: Add some missing field in row group metadata: ordinal, total co…

parquet/src/file/metadata.rs

@@ -277,6 +277,9 @@ pub struct RowGroupMetaData {
     sorting_columns: Option<Vec<SortingColumn>>,
     total_byte_size: i64,
     schema_descr: SchemaDescPtr,
+    // We can't infer from file offset of first column since there may empty columns in row group.
+    file_offset: Option<i64>,
+    ordinal: Option<i16>,
 }
 
 impl RowGroupMetaData {
@@ -330,6 +333,18 @@ impl RowGroupMetaData {
         self.schema_descr.clone()
     }
 
+    /// Returns ordinal of this row group in file
+    #[inline(always)]
+    pub fn ordinal(&self) -> Option<i16> {
+        self.ordinal
+    }
+
+    /// Returns file offset of this row group in file.
+    #[inline(always)]
+    pub fn file_offset(&self) -> Option<i64> {
+        self.file_offset
+    }
+
     /// Method to convert from Thrift.
     pub fn from_thrift(
         schema_descr: SchemaDescPtr,
@@ -350,6 +365,8 @@ impl RowGroupMetaData {
             sorting_columns,
             total_byte_size,
             schema_descr,
+            file_offset: rg.file_offset,
+            ordinal: rg.ordinal,
         })
     }
 
@@ -360,9 +377,9 @@ impl RowGroupMetaData {
             total_byte_size: self.total_byte_size,
             num_rows: self.num_rows,
             sorting_columns: self.sorting_columns().cloned(),
-            file_offset: None,
-            total_compressed_size: None,
-            ordinal: None,
+            file_offset: self.file_offset(),
+            total_compressed_size: Some(self.compressed_size()),
+            ordinal: self.ordinal,
         }
     }
 
@@ -381,9 +398,11 @@ impl RowGroupMetaDataBuilder {
         Self(RowGroupMetaData {
             columns: Vec::with_capacity(schema_descr.num_columns()),
             schema_descr,
+            file_offset: None,
             num_rows: 0,
             sorting_columns: None,
             total_byte_size: 0,
+            ordinal: None,
         })
     }
 
@@ -411,6 +430,17 @@ impl RowGroupMetaDataBuilder {
         self
     }
 
+    /// Sets ordinal for this row group.
+    pub fn set_ordinal(mut self, value: i16) -> Self {
+        self.0.ordinal = Some(value);
+        self
+    }
+
+    pub fn set_file_offset(mut self, value: i64) -> Self {
+        self.0.file_offset = Some(value);
+        self
+    }
+
     /// Builds row group metadata.
     pub fn build(self) -> Result<RowGroupMetaData> {
         if self.0.schema_descr.num_columns() != self.0.columns.len() {
@@ -966,6 +996,7 @@ mod tests {
             .set_num_rows(1000)
             .set_total_byte_size(2000)
             .set_column_metadata(columns)
+            .set_ordinal(1)
             .build()
             .unwrap();
 

parquet/src/file/writer.rs

@@ -183,6 +183,8 @@ impl<W: Write + Send> SerializedFileWriter<W> {
     /// previous row group must be finalised and closed using `RowGroupWriter::close` method.
     pub fn next_row_group(&mut self) -> Result<SerializedRowGroupWriter<'_, W>> {
         self.assert_previous_writer_closed()?;
+        let ordinal = self.row_group_index;
+
         self.row_group_index += 1;
 
         let row_groups = &mut self.row_groups;
@@ -204,6 +206,7 @@ impl<W: Write + Send> SerializedFileWriter<W> {
             self.descr.clone(),
             self.props.clone(),
             &mut self.buf,
+            ordinal as i16,
             Some(Box::new(on_close)),
         );
         Ok(row_group_writer)
@@ -409,6 +412,8 @@ pub struct SerializedRowGroupWriter<'a, W: Write> {
     bloom_filters: Vec<Option<Sbbf>>,
     column_indexes: Vec<Option<ColumnIndex>>,
     offset_indexes: Vec<Option<OffsetIndex>>,
+    row_group_index: i16,
+    file_offset: i64,
     on_close: Option<OnCloseRowGroup<'a>>,
 }
 
@@ -418,16 +423,22 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> {
     /// - `schema_descr` - the schema to write
     /// - `properties` - writer properties
     /// - `buf` - the buffer to write data to
+    /// - `row_group_index` - row group index in this parquet file.
+    /// - `file_offset` - file offset of this row group in this parquet file.
     /// - `on_close` - an optional callback that will invoked on [`Self::close`]
     pub fn new(
         schema_descr: SchemaDescPtr,
         properties: WriterPropertiesPtr,
         buf: &'a mut TrackedWrite<W>,
+        row_group_index: i16,
         on_close: Option<OnCloseRowGroup<'a>>,
     ) -> Self {
         let num_columns = schema_descr.num_columns();
+        let file_offset = buf.bytes_written() as i64;
         Self {
             buf,
+            row_group_index,
+            file_offset,
             on_close,
             total_rows_written: None,
             descr: schema_descr,
@@ -603,6 +614,8 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> {
                 .set_total_byte_size(self.total_uncompressed_bytes)
                 .set_num_rows(self.total_rows_written.unwrap_or(0) as i64)
                 .set_sorting_columns(self.props.sorting_columns().cloned())
+                .set_ordinal(self.row_group_index)
+                .set_file_offset(self.file_offset)
                 .build()?;
 
             let metadata = Arc::new(row_group_metadata);
@@ -1312,6 +1325,7 @@ mod tests {
         let mut rows: i64 = 0;
 
         for (idx, subset) in data.iter().enumerate() {
+            let row_group_file_offset = file_writer.buf.bytes_written();
             let mut row_group_writer = file_writer.next_row_group().unwrap();
             if let Some(mut writer) = row_group_writer.next_column().unwrap() {
                 rows += writer
@@ -1323,6 +1337,8 @@ mod tests {
             let last_group = row_group_writer.close().unwrap();
             let flushed = file_writer.flushed_row_groups();
             assert_eq!(flushed.len(), idx + 1);
+            assert_eq!(Some(idx as i16), last_group.ordinal());
+            assert_eq!(Some(row_group_file_offset as i64), last_group.file_offset());
             assert_eq!(flushed[idx].as_ref(), last_group.as_ref());
         }
         let file_metadata = file_writer.close().unwrap();