[Parquet] Add row group distinct counts to StatisticsConverter

parquet/src/arrow/arrow_reader/statistics.rs

@@ -1821,6 +1821,28 @@ impl<'a> StatisticsConverter<'a> {
         Ok(UInt64Array::from_iter(nan_counts))
     }
 
+    /// Extract the distinct counts from row group statistics in [`RowGroupMetaData`]
+    ///
+    /// See docs on [`Self::row_group_mins`] for details
+    pub fn row_group_distinct_counts<I>(&self, metadatas: I) -> Result<UInt64Array>
+    where
+        I: IntoIterator<Item = &'a RowGroupMetaData>,
+    {
+        let Some(parquet_index) = self.parquet_column_index else {
+            let num_row_groups = metadatas.into_iter().count();
+            return Ok(UInt64Array::from_iter(std::iter::repeat_n(
+                None,
+                num_row_groups,
+            )));
+        };
+
+        let distinct_counts = metadatas
+            .into_iter()
+            .map(|x| x.column(parquet_index).statistics())
+            .map(|s| s.and_then(|s| s.distinct_count_opt()));
+        Ok(UInt64Array::from_iter(distinct_counts))
+    }
+
     /// Extract the minimum values from Data Page statistics.
     ///
     /// In Parquet files, in addition to the Column Chunk level statistics

parquet/tests/arrow_reader/statistics.rs

@@ -2630,10 +2630,13 @@ mod test {
         Int32Array, Int64Array, RecordBatch, StringArray, StructArray, TimestampNanosecondArray,
         new_empty_array,
     };
-    use arrow_schema::{DataType, SchemaRef, TimeUnit};
+    use arrow_schema::{DataType, Field, SchemaRef, TimeUnit};
     use bytes::Bytes;
     use parquet::arrow::parquet_column;
+    use parquet::data_type::{ByteArray, ByteArrayType, Int32Type};
     use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData};
+    use parquet::file::writer::SerializedFileWriter;
+    use parquet::schema::parser::parse_message_type;
     use std::path::PathBuf;
     use std::sync::Arc;
     // TODO error cases (with parquet statistics that are mismatched in expected type)
@@ -3222,4 +3225,131 @@ mod test {
             .collect();
         Arc::new(array)
     }
+
+    // Verifies that distinct_count is correctly read back from UTF-8 column statistics.
+    // Uses the low-level writer to inject a known distinct_count into the parquet footer
+    // since ArrowWriter does not yet write this field.
+    #[test]
+    fn test_row_group_distinct_counts_utf8_roundtrip() {
+        let unique_string_values: Vec<ByteArray> = (0..10)
+            .map(|index| ByteArray::from(format!("value_{index}").into_bytes()))
+            .collect();
+
+        let parquet_schema = Arc::new(
+            parse_message_type("message schema { REQUIRED BYTE_ARRAY col (UTF8); }").unwrap(),
+        );
+        let writer_properties = Arc::new(
+            WriterProperties::builder()
+                .set_statistics_enabled(EnabledStatistics::Chunk)
+                .build(),
+        );
+
+        let mut file_buffer: Vec<u8> = Vec::new();
+        let mut file_writer =
+            SerializedFileWriter::new(&mut file_buffer, parquet_schema, writer_properties).unwrap();
+
+        let mut row_group_writer = file_writer.next_row_group().unwrap();
+        let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
+        let min_value = unique_string_values.first().unwrap().clone();
+        let max_value = unique_string_values.last().unwrap().clone();
+        let expected_distinct_count = unique_string_values.len() as u64;
+        column_writer
+            .typed::<ByteArrayType>()
+            .write_batch_with_statistics(
+                &unique_string_values,
+                None,
+                None,
+                Some(&min_value),
+                Some(&max_value),
+                Some(expected_distinct_count),
+            )
+            .unwrap();
+        column_writer.close().unwrap();
+        row_group_writer.close().unwrap();
+        file_writer.close().unwrap();
+
+        let parquet_bytes = Bytes::from(file_buffer);
+        let reader_builder = ParquetRecordBatchReaderBuilder::try_new(parquet_bytes).unwrap();
+        let file_metadata = reader_builder.metadata().clone();
+        let arrow_schema = reader_builder.schema().clone();
+        let parquet_schema_descriptor = file_metadata.file_metadata().schema_descr();
+
+        let statistics_converter =
+            StatisticsConverter::try_new("col", &arrow_schema, parquet_schema_descriptor).unwrap();
+        let distinct_counts = statistics_converter
+            .row_group_distinct_counts(file_metadata.row_groups().iter())
+            .unwrap();
+
+        assert_eq!(
+            distinct_counts,
+            UInt64Array::from(vec![Some(expected_distinct_count)]),
+            "expected distinct_count of 10 unique string values"
+        );
+    }
+
+    // Verifies that a missing distinct_count in one row group does not affect the others.
+    #[test]
+    fn test_row_group_distinct_counts_absent() {
+        let parquet_schema =
+            Arc::new(parse_message_type("message schema { REQUIRED INT32 col; }").unwrap());
+        let writer_properties = Arc::new(
+            WriterProperties::builder()
+                .set_statistics_enabled(EnabledStatistics::Chunk)
+                .build(),
+        );
+
+        let mut file_buffer: Vec<u8> = Vec::new();
+        let mut file_writer =
+            SerializedFileWriter::new(&mut file_buffer, parquet_schema, writer_properties).unwrap();
+
+        // row group 0: distinct_count present
+        let mut row_group_writer = file_writer.next_row_group().unwrap();
+        let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
+        column_writer
+            .typed::<Int32Type>()
+            .write_batch_with_statistics(&[1, 2], None, None, Some(&1), Some(&2), Some(2))
+            .unwrap();
+        column_writer.close().unwrap();
+        row_group_writer.close().unwrap();
+
+        // row group 1: distinct_count absent — should appear as null in output
+        let mut row_group_writer = file_writer.next_row_group().unwrap();
+        let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
+        column_writer
+            .typed::<Int32Type>()
+            .write_batch_with_statistics(&[3, 4], None, None, Some(&3), Some(&4), None)
+            .unwrap();
+        column_writer.close().unwrap();
+        row_group_writer.close().unwrap();
+
+        // row group 2: distinct_count present — iteration must reach here despite row group 1
+        let mut row_group_writer = file_writer.next_row_group().unwrap();
+        let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
+        column_writer
+            .typed::<Int32Type>()
+            .write_batch_with_statistics(&[5, 6, 7, 8, 9], None, None, Some(&5), Some(&9), Some(5))
+            .unwrap();
+        column_writer.close().unwrap();
+        row_group_writer.close().unwrap();
+
+        file_writer.close().unwrap();
+
+        let parquet_bytes = Bytes::from(file_buffer);
+        let reader_builder = ParquetRecordBatchReaderBuilder::try_new(parquet_bytes).unwrap();
+        let file_metadata = reader_builder.metadata().clone();
+        let arrow_schema = reader_builder.schema().clone();
+        let parquet_schema_descriptor = file_metadata.file_metadata().schema_descr();
+
+        let statistics_converter =
+            StatisticsConverter::try_new("col", &arrow_schema, parquet_schema_descriptor).unwrap();
+        let distinct_counts = statistics_converter
+            .row_group_distinct_counts(file_metadata.row_groups().iter())
+            .unwrap();
+
+        assert_eq!(
+            distinct_counts,
+            UInt64Array::from(vec![Some(2), None, Some(5)]),
+            "a missing distinct_count in one row group should not affect the others"
+        );
+    }
 }