This is an automated email from the ASF dual-hosted git repository.
Jefffrey 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 891711ec1d fix: ndv calculation for dictionary arrays (#11173)
891711ec1d is described below
commit 891711ec1dbec2dcf4924e92258dd66bfc7b96f3
Author: RIchard Baah <[email protected]>
AuthorDate: Thu Sep 24 01:33:01 2026 -0400
fix: ndv calculation for dictionary arrays (#11173)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- Closes #11172.
# Rationale for this change
Hashing key indices for dictionary arrays was incorrect across batches:
key 0 in one batch and key 0 in another can refer to completely
different values, causing the NDV to be undercounted.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
Fix iterates over the values array; O(distinct values) instead of
O(rows).
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
yes, added a regression test that was mentioned in the tagged issue
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
correct ndv field
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
parquet/src/arrow/arrow_writer/mod.rs | 146 ++++++++++++++++++++++++++++------
1 file changed, 120 insertions(+), 26 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 4bee015d9c..29b413531c 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -1145,21 +1145,16 @@ impl ArrowColumnWriter {
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]));
- }
- }
+ // Hash referenced values, not key indices: keys can map
to different
+ // values across batches, and unreferenced values must not
count toward NDV.
+ let values = dict.values();
+ let keys = dict.normalized_keys();
+ let referenced_value_indices: Vec<usize> = non_null
+ .iter()
+ .map(|&pos| keys[pos])
+ .filter(|&val_idx| values.is_valid(val_idx))
+ .collect();
+ update_distinct_values_seen(values.as_ref(),
&referenced_value_indices, seen);
}
// For plain arrays, hash the actual values directly.
None => update_distinct_values_seen(array.as_ref(), non_null,
seen),
@@ -1982,17 +1977,6 @@ 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::*;
@@ -6566,6 +6550,116 @@ mod tests {
assert!(count.is_none());
}
+ #[test]
+ fn test_dictionary_ndv_single_batch() {
+ // Dictionary array with 3 distinct string values repeated many times.
+ // NDV must equal the number of distinct values in the dictionary (3),
+ // not the number of rows.
+ let keys = Int32Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog",
"bird"]));
+ let dict: ArrayRef =
Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "x",
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
+ false,
+ )]));
+ let batch = RecordBatch::try_new(schema, vec![dict]).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");
+ assert_eq!(count, 3);
+ }
+
+ #[test]
+ fn test_dictionary_ndv_excludes_unreferenced_values() {
+ // Keys only reference indices 0 and 1; value at index 2
("unreferenced") should not
+ // count toward NDV even though it appears in the dictionary's values
array.
+ let keys = Int32Array::from(vec![0, 1, 0, 1]);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog",
"unreferenced"]));
+ let dict: ArrayRef =
Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "x",
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
+ false,
+ )]));
+ let batch = RecordBatch::try_new(schema, vec![dict]).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");
+ assert_eq!(
+ count, 2,
+ "unreferenced dictionary values must not count toward NDV"
+ );
+ }
+
+ #[test]
+ fn test_dictionary_ndv_across_batches_regression() {
+ // Regression test for https://github.com/apache/arrow-rs/issues/11172.
+ let make_dict_batch = |a: &str, b: &str| -> RecordBatch {
+ let keys = Int32Array::from(vec![0, 1, 0, 1]);
+ let values: ArrayRef = Arc::new(StringArray::from(vec![a, b]));
+ let dict: ArrayRef =
+ Arc::new(DictionaryArray::<Int32Type>::try_new(keys,
values).unwrap());
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "x",
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
+ false,
+ )]));
+ RecordBatch::try_new(schema, vec![dict]).unwrap()
+ };
+
+ // batch1: dict = ["cat", "dog"], batch2: dict = ["fish", "cat"]
+ // Distinct values across both batches: "cat", "dog", "fish" NDV = 3
+ let batch1 = make_dict_batch("cat", "dog");
+ let batch2 = make_dict_batch("fish", "cat");
+
+ 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, batch1.schema(),
Some(props)).unwrap();
+ writer.write(&batch1).unwrap();
+ writer.write(&batch2).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");
+ assert_eq!(
+ count, 3,
+ "NDV should count distinct values, not distinct key indices"
+ );
+ }
+
#[test]
fn ree_struct_with_ree_child() {
// Struct with a REE string field and a REE int field — confirms
