This is an automated email from the ASF dual-hosted git repository.
etseidl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 3320517b91 [Parquet] Implement num distinct values for parquet writer
(#10654)
3320517b91 is described below
commit 3320517b9130fc8d07c064bc48bfa873d4fec020
Author: RIchard Baah <[email protected]>
AuthorDate: Wed Aug 19 17:31:53 2026 -0400
[Parquet] Implement num distinct values for parquet writer (#10654)
# Which issue does this PR close?
- Closes #8608.
- Closes #10650
# Rationale for this change
see #8608, #10650 & https://github.com/apache/datafusion/issues/24114
# What changes are included in this PR?
Adds `set_write_row_group_number_distinct_values` to WriterProperties,
which when enabled causes the ArrowWriter to track the exact number of
distinct non-null values per column across the full row group and write
it into the column chunk statistics footer as `distinct_count`. Tracking
is implemented by hashing each non-null value using XxHash64 into a
per-column HashSet<u64> that persists across batch writes and is
finalized in `close()`. **The flag defaults to false so there is no
impact on existing writers.**
# Are these changes tested?
yes, see test.
# Are there any user-facing changes?
yes, users will not be able to write `distinct_count` meta data to
parquet files.
---
parquet/benches/arrow_writer.rs | 5 +
parquet/src/arrow/arrow_writer/mod.rs | 209 +++++++++++++++++++++++++++++++++-
parquet/src/column/writer/mod.rs | 28 ++++-
parquet/src/file/properties.rs | 43 +++++++
4 files changed, 282 insertions(+), 3 deletions(-)
diff --git a/parquet/benches/arrow_writer.rs b/parquet/benches/arrow_writer.rs
index 73dc6fae79..54d15284d7 100644
--- a/parquet/benches/arrow_writer.rs
+++ b/parquet/benches/arrow_writer.rs
@@ -785,6 +785,11 @@ fn create_writer_props() -> Vec<(&'static str,
WriterProperties)> {
.build();
props.push(("cdc", prop));
+ let prop = WriterProperties::builder()
+ .set_write_row_group_number_distinct_values(true)
+ .build();
+ props.push(("number_distinct_values", prop));
+
props
}
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 956ccf6f15..947f07a627 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -44,6 +44,8 @@ use crate::column::writer::{
ColumnCloseResult, ColumnWriter, GenericColumnWriter, get_column_writer,
};
use crate::data_type::{ByteArray, FixedLenByteArray};
+use std::collections::HashSet;
+type DistinctValuesSet = HashSet<u64>;
#[cfg(feature = "encryption")]
use crate::encryption::encrypt::FileEncryptor;
use crate::errors::{ParquetError, Result};
@@ -1068,6 +1070,9 @@ impl ArrowColumnChunk {
pub struct ArrowColumnWriter {
writer: ArrowColumnWriterImpl,
chunk: SharedColumnChunk,
+ /// Non-null value hashes accumulated across all writes for this column's
row group.
+ /// `None` when tracking is disabled via
[`WriterProperties::write_row_group_number_distinct_values`].
+ distinct_values_seen: Option<DistinctValuesSet>,
}
impl std::fmt::Debug for ArrowColumnWriter {
@@ -1117,6 +1122,32 @@ impl ArrowColumnWriter {
}
fn write_internal(&mut self, levels: &ArrayLevels) -> Result<()> {
+ if let Some(seen) = &mut self.distinct_values_seen {
+ let array = levels.array();
+ let non_null = levels.non_null_indices();
+ match array.as_any_dictionary_opt() {
+ Some(dict) => {
+ // For dictionary arrays, hash the integer keys rather
than the actual values.
+ // Key cardinality equals value cardinality, so
distinct-value counting stays
+ // correct while avoiding the cost of hashing
arbitrary-length values.
+ let keys = dict.keys();
+ let key_data = keys.to_data();
+ let offset = key_data.offset();
+ let width = arrow_key_byte_width(keys.data_type());
+ if width > 0 {
+ let buffer = key_data.buffers()[0].as_slice();
+ // Only visit non-null rows to avoid counting nulls as
a distinct value.
+ for &row in non_null {
+ let pos = (offset + row) * width;
+ seen.insert(hash_bytes(&buffer[pos..pos + width]));
+ }
+ }
+ }
+ // For plain arrays, hash the actual values directly.
+ None => update_distinct_values_seen(array.as_ref(), non_null,
seen),
+ }
+ }
+
match &mut self.writer {
ArrowColumnWriterImpl::Column(c) => {
let leaf = levels.array();
@@ -1138,9 +1169,24 @@ impl ArrowColumnWriter {
/// Close this column returning the written [`ArrowColumnChunk`]
pub fn close(self) -> Result<ArrowColumnChunk> {
+ let distinct_count = self
+ .distinct_values_seen
+ .as_ref()
+ .filter(|s| !s.is_empty())
+ .map(|s| s.len() as u64);
let close = match self.writer {
- ArrowColumnWriterImpl::ByteArray(c) => c.close()?,
- ArrowColumnWriterImpl::Column(c) => c.close()?,
+ ArrowColumnWriterImpl::ByteArray(mut c) => {
+ if let Some(count) = distinct_count {
+ c.set_distinct_count_override(count);
+ }
+ c.close()?
+ }
+ ArrowColumnWriterImpl::Column(mut c) => {
+ if let Some(count) = distinct_count {
+ c.set_distinct_count_override(count);
+ }
+ c.close()?
+ }
};
let chunk = Arc::try_unwrap(self.chunk).ok().unwrap();
let data = chunk.into_inner().unwrap();
@@ -1405,6 +1451,8 @@ impl ArrowColumnWriterFactory {
leaves: &mut Iter<'_, ColumnDescPtr>,
out: &mut Vec<ArrowColumnWriter>,
) -> Result<()> {
+ let write_distinct_values =
props.write_row_group_number_distinct_values();
+
// Instantiate writers for normal columns
let col = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
let page_writer = self.create_page_writer(desc, out.len())?;
@@ -1413,6 +1461,7 @@ impl ArrowColumnWriterFactory {
Ok(ArrowColumnWriter {
chunk,
writer: ArrowColumnWriterImpl::Column(writer),
+ distinct_values_seen: write_distinct_values.then(HashSet::new),
})
};
@@ -1424,6 +1473,7 @@ impl ArrowColumnWriterFactory {
Ok(ArrowColumnWriter {
chunk,
writer: ArrowColumnWriterImpl::ByteArray(writer),
+ distinct_values_seen: write_distinct_values.then(HashSet::new),
})
};
@@ -1899,6 +1949,109 @@ fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size:
usize) -> Vec<FixedLenByteAr
values
}
+/// Hash a byte slice to a u64 for NDV tracking.
+#[inline]
+fn hash_bytes(bytes: &[u8]) -> u64 {
+ twox_hash::XxHash64::oneshot(0, bytes)
+}
+
+/// Returns the byte width of an Arrow dictionary key type, or 0 if
unsupported.
+fn arrow_key_byte_width(dt: &ArrowDataType) -> usize {
+ match dt {
+ ArrowDataType::Int8 | ArrowDataType::UInt8 => 1,
+ ArrowDataType::Int16 | ArrowDataType::UInt16 => 2,
+ ArrowDataType::Int32 | ArrowDataType::UInt32 => 4,
+ ArrowDataType::Int64 | ArrowDataType::UInt64 => 8,
+ _ => 0,
+ }
+}
+
+/// Returns the fixed byte width for primitive Arrow types, or `None` for
variable-length types.
+fn fixed_byte_width(dt: &ArrowDataType) -> Option<usize> {
+ use ArrowDataType::*;
+ match dt {
+ Int8 | UInt8 => Some(1),
+ Int16 | UInt16 | Float16 => Some(2),
+ Int32 | UInt32 | Float32 | Date32 | Time32(_) | Decimal32(_, _) =>
Some(4),
+ Int64
+ | UInt64
+ | Float64
+ | Date64
+ | Time64(_)
+ | Timestamp(_, _)
+ | Duration(_)
+ | Decimal64(_, _) => Some(8),
+ Interval(IntervalUnit::YearMonth) => Some(4),
+ Interval(IntervalUnit::DayTime) => Some(8),
+ Interval(IntervalUnit::MonthDayNano) => Some(16),
+ Decimal128(_, _) => Some(16),
+ Decimal256(_, _) => Some(32),
+ _ => None,
+ }
+}
+
+/// Hash the non-null values in `array` (at `non_null_indices`) into `seen`.
+///
+/// Handles primitive, boolean, fixed-size-binary, and variable-length
(Utf8/Binary)
+/// arrays. Unsupported types are silently skipped, leaving `seen` unchanged
for
+/// those values (NDV is best-effort).
+fn update_distinct_values_seen(
+ array: &dyn arrow_array::Array,
+ non_null_indices: &[usize],
+ seen: &mut DistinctValuesSet,
+) {
+ let data = array.to_data();
+ let offset = data.offset();
+
+ match array.data_type() {
+ ArrowDataType::Boolean => {
+ let arr = array
+ .as_any()
+ .downcast_ref::<arrow_array::BooleanArray>()
+ .unwrap();
+ for &row in non_null_indices {
+ seen.insert(arr.value(row) as u64);
+ }
+ }
+ ArrowDataType::Utf8 | ArrowDataType::Binary => {
+ let offsets = data.buffers()[0].typed_data::<i32>();
+ let values = data.buffers()[1].as_slice();
+ for &row in non_null_indices {
+ let start = offsets[offset + row] as usize;
+ let end = offsets[offset + row + 1] as usize;
+ seen.insert(hash_bytes(&values[start..end]));
+ }
+ }
+ ArrowDataType::LargeUtf8 | ArrowDataType::LargeBinary => {
+ let offsets = data.buffers()[0].typed_data::<i64>();
+ let values = data.buffers()[1].as_slice();
+ for &row in non_null_indices {
+ let start = offsets[offset + row] as usize;
+ let end = offsets[offset + row + 1] as usize;
+ seen.insert(hash_bytes(&values[start..end]));
+ }
+ }
+ ArrowDataType::FixedSizeBinary(byte_width) => {
+ let byte_width = *byte_width as usize;
+ let buffer = data.buffers()[0].as_slice();
+ for &row in non_null_indices {
+ let start = (offset + row) * byte_width;
+ seen.insert(hash_bytes(&buffer[start..start + byte_width]));
+ }
+ }
+ data_type => {
+ if let Some(width) = fixed_byte_width(data_type) {
+ let buffer = data.buffers()[0].as_slice();
+ for &row in non_null_indices {
+ let pos = (offset + row) * width;
+ seen.insert(hash_bytes(&buffer[pos..pos + width]));
+ }
+ }
+ // Utf8View, BinaryView, nested types: skip
+ }
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -6026,6 +6179,58 @@ mod tests {
ree_write_read_roundtrip(sliced, flat);
}
+ #[test]
+ fn test_number_distinct_values_exact_count() {
+ // 50 distinct Int32 values repeated across 100k rows, with every 7th
row null.
+ // Nulls must not be counted as a distinct value.
+ let cardinality = 50u32;
+ let array: ArrayRef =
Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
+ if i % 7 == 0 {
+ None
+ } else {
+ Some((i % cardinality) as i32)
+ }
+ })));
+ let schema = Arc::new(Schema::new(vec![Field::new("x",
DataType::Int32, true)]));
+ let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
+
+ let props = WriterProperties::builder()
+ .set_write_row_group_number_distinct_values(true)
+ .build();
+ let mut buf = Vec::new();
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ let metadata = writer.close().unwrap();
+
+ let count = metadata
+ .row_group(0)
+ .column(0)
+ .statistics()
+ .and_then(|s| s.distinct_count_opt())
+ .expect("distinct_count should be set");
+ // Must equal cardinality exactly; nulls must not inflate the count.
+ assert_eq!(count, cardinality as u64);
+ }
+
+ #[test]
+ fn test_number_distinct_values_not_written_by_default() {
+ let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
+ let schema = Arc::new(Schema::new(vec![Field::new("x",
DataType::Int32, false)]));
+ let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
+
+ let mut buf = Vec::new();
+ let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(),
None).unwrap();
+ writer.write(&batch).unwrap();
+ let metadata = writer.close().unwrap();
+
+ let count = metadata
+ .row_group(0)
+ .column(0)
+ .statistics()
+ .and_then(|s| s.distinct_count_opt());
+ assert!(count.is_none());
+ }
+
#[test]
fn ree_struct_with_ree_child() {
// Struct with a REE string field and a REE int field — confirms
diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs
index d49a1e4897..6bda82d646 100644
--- a/parquet/src/column/writer/mod.rs
+++ b/parquet/src/column/writer/mod.rs
@@ -114,6 +114,14 @@ impl ColumnWriter<'_> {
downcast_writer!(self, typed, typed.add_data_page())
}
+ /// Sets a pre-computed distinct count on this column writer.
+ ///
+ /// See [`GenericColumnWriter::set_distinct_count_override`] for details.
+ #[cfg(feature = "arrow")]
+ pub(crate) fn set_distinct_count_override(&mut self, count: u64) {
+ downcast_writer!(self, typed, typed.set_distinct_count_override(count))
+ }
+
/// Close this [`ColumnWriter`], returning the metadata for the column
chunk.
pub fn close(self) -> Result<ColumnCloseResult> {
downcast_writer!(self, typed, typed.close())
@@ -453,6 +461,10 @@ pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> {
// Metrics per column writer
column_metrics: ColumnMetrics<E::T>,
+ /// Pre-computed distinct count to write into column chunk statistics.
+ /// When set, takes precedence over `column_metrics.column_distinct_count`.
+ distinct_count_override: Option<u64>,
+
/// The order of encodings within the generated metadata does not impact
its meaning,
/// but we use a BTreeSet so that the output is deterministic
encodings: BTreeSet<Encoding>,
@@ -529,6 +541,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
data_pages: VecDeque::new(),
page_metrics,
column_metrics,
+ distinct_count_override: None,
column_index_builder,
offset_index_builder,
encodings,
@@ -539,6 +552,16 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E>
{
}
}
+ /// Sets a pre-computed distinct count to write into column chunk
statistics.
+ ///
+ /// When set, this value is written as `distinct_count` in the row group
statistics
+ /// footer. It takes precedence over any `distinct_count` passed through
+ /// [`Self::write_batch_with_statistics`].
+ #[cfg(feature = "arrow")]
+ pub(crate) fn set_distinct_count_override(&mut self, count: u64) {
+ self.distinct_count_override = Some(count);
+ }
+
#[expect(clippy::too_many_arguments)]
pub(crate) fn write_batch_internal(
&mut self,
@@ -1527,10 +1550,13 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a,
E> {
if self.statistics_enabled != EnabledStatistics::None {
let backwards_compatible_min_max =
self.descr.sort_order().is_signed();
+ let distinct_count = self
+ .distinct_count_override
+ .or(self.column_metrics.column_distinct_count);
let statistics = ValueStatistics::<E::T>::new(
self.column_metrics.min_column_value.clone(),
self.column_metrics.max_column_value.clone(),
- self.column_metrics.column_distinct_count,
+ distinct_count,
Some(self.column_metrics.num_column_nulls),
false,
)
diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs
index 174c783377..c7a4f550b8 100644
--- a/parquet/src/file/properties.rs
+++ b/parquet/src/file/properties.rs
@@ -68,6 +68,8 @@ pub const DEFAULT_STATISTICS_TRUNCATE_LENGTH: Option<usize> =
Some(64);
pub const DEFAULT_OFFSET_INDEX_DISABLED: bool = false;
/// Default values for [`WriterProperties::coerce_types`]
pub const DEFAULT_COERCE_TYPES: bool = false;
+/// Default value for
[`WriterProperties::write_row_group_number_distinct_values`]
+pub const DEFAULT_WRITE_ROW_GROUP_NUMBER_DISTINCT_VALUES: bool = false;
/// Default value for
[`WriterProperties::data_page_v2_compression_ratio_threshold`]
pub const DEFAULT_DATA_PAGE_V2_COMPRESSION_RATIO_THRESHOLD: f64 = 1.0;
/// Default value for [`WriterProperties::write_path_in_schema`]
@@ -254,6 +256,7 @@ pub struct WriterProperties {
column_index_truncate_length: Option<usize>,
statistics_truncate_length: Option<usize>,
coerce_types: bool,
+ write_row_group_number_distinct_values: bool,
content_defined_chunking: Option<CdcOptions>,
write_path_in_schema: bool,
#[cfg(feature = "encryption")]
@@ -433,6 +436,14 @@ impl WriterProperties {
self.coerce_types
}
+ /// Returns `true` if the writer should compute and store the distinct
count
+ /// (`num_distinct_values`) in row group column chunk statistics.
+ ///
+ /// For more details see
[`WriterPropertiesBuilder::set_write_row_group_number_distinct_values`]
+ pub fn write_row_group_number_distinct_values(&self) -> bool {
+ self.write_row_group_number_distinct_values
+ }
+
/// Returns `true` if the `path_in_schema` field of the `ColumnMetaData`
Thrift struct
/// should be written.
///
@@ -595,6 +606,7 @@ pub struct WriterPropertiesBuilder {
column_index_truncate_length: Option<usize>,
statistics_truncate_length: Option<usize>,
coerce_types: bool,
+ write_row_group_number_distinct_values: bool,
content_defined_chunking: Option<CdcOptions>,
write_path_in_schema: bool,
#[cfg(feature = "encryption")]
@@ -620,6 +632,7 @@ impl Default for WriterPropertiesBuilder {
column_index_truncate_length: DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH,
statistics_truncate_length: DEFAULT_STATISTICS_TRUNCATE_LENGTH,
coerce_types: DEFAULT_COERCE_TYPES,
+ write_row_group_number_distinct_values:
DEFAULT_WRITE_ROW_GROUP_NUMBER_DISTINCT_VALUES,
content_defined_chunking: None,
write_path_in_schema: DEFAULT_WRITE_PATH_IN_SCHEMA,
#[cfg(feature = "encryption")]
@@ -675,6 +688,7 @@ impl WriterPropertiesBuilder {
column_index_truncate_length: self.column_index_truncate_length,
statistics_truncate_length: self.statistics_truncate_length,
coerce_types: self.coerce_types,
+ write_row_group_number_distinct_values:
self.write_row_group_number_distinct_values,
content_defined_chunking: self.content_defined_chunking,
write_path_in_schema: self.write_path_in_schema,
#[cfg(feature = "encryption")]
@@ -896,6 +910,34 @@ impl WriterPropertiesBuilder {
self
}
+ /// Enable or disable writing the distinct value count
(`num_distinct_values`) into
+ /// row group column chunk statistics (defaults to `false` via
+ /// [`DEFAULT_WRITE_ROW_GROUP_NUMBER_DISTINCT_VALUES`]).
+ ///
+ /// When enabled, the [`ArrowWriter`] scans each column's values before
encoding
+ /// and stores the number of distinct non-null values in the row group
statistics
+ /// footer.
+ ///
+ /// # Compatibility
+ ///
+ /// This setting only takes effect when using [`ArrowWriter`]. The
row-based
+ /// [`SerializedFileWriter`] / [`SerializedRowGroupWriter`] APIs do not
populate
+ /// `num_distinct_values` and will ignore this flag.
+ ///
+ /// # Performance
+ ///
+ /// Computing the distinct count requires hashing every non-null value in
the column.
+ /// For large row groups or columns with many values this adds significant
overhead.
+ /// Benchmark your workload before enabling this globally.
+ ///
+ /// [`ArrowWriter`]: crate::arrow::ArrowWriter
+ /// [`SerializedFileWriter`]: crate::file::writer::SerializedFileWriter
+ /// [`SerializedRowGroupWriter`]:
crate::file::writer::SerializedRowGroupWriter
+ pub fn set_write_row_group_number_distinct_values(mut self, value: bool)
-> Self {
+ self.write_row_group_number_distinct_values = value;
+ self
+ }
+
/// EXPERIMENTAL: Should the writer emit the `path_in_schema` element of
the
/// `ColumnMetaData` Thrift struct. Defaults to `true` via
[`DEFAULT_WRITE_PATH_IN_SCHEMA`].
///
@@ -1342,6 +1384,7 @@ impl From<WriterProperties> for WriterPropertiesBuilder {
column_index_truncate_length: props.column_index_truncate_length,
statistics_truncate_length: props.statistics_truncate_length,
coerce_types: props.coerce_types,
+ write_row_group_number_distinct_values:
props.write_row_group_number_distinct_values,
content_defined_chunking: props.content_defined_chunking,
write_path_in_schema: props.write_path_in_schema,
#[cfg(feature = "encryption")]
