diff --git a/parquet/src/arrow/arrow_reader/statistics.rs b/parquet/src/arrow/arrow_reader/statistics.rs index 7ff4a3a4e102..e01fe14e1af2 100644 --- a/parquet/src/arrow/arrow_reader/statistics.rs +++ b/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(&self, metadatas: I) -> Result + where + I: IntoIterator, + { + 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 diff --git a/parquet/tests/arrow_reader/statistics.rs b/parquet/tests/arrow_reader/statistics.rs index 0bc61ba108e5..e0740d38e93e 100644 --- a/parquet/tests/arrow_reader/statistics.rs +++ b/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 = (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 = 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::() + .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 = 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::() + .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::() + .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::() + .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" + ); + } }