Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Fixed statistics writing flag and correct null_count in dictionaries #1414

Merged
merged 2 commits into from
Feb 22, 2023
Merged

Fixed statistics writing flag and correct null_count in dictionaries #1414

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2556574
fix(parquet): repsect statistics and correct null_count
ritchie46 Feb 22, 2023
bbeb28a
fix invalid test
ritchie46 Feb 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 98 additions & 67 deletions src/io/parquet/write/dictionary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ fn serialize_keys<K: DictionaryKey>(
array: &DictionaryArray<K>,
type_: PrimitiveType,
nested: &[Nested],
statistics: ParquetStatistics,
statistics: Option<ParquetStatistics>,
options: WriteOptions,
) -> Result<Page> {
let mut buffer = vec![];
Expand Down Expand Up @@ -146,7 +146,7 @@ fn serialize_keys<K: DictionaryKey>(
array.null_count(),
repetition_levels_byte_length,
definition_levels_byte_length,
Some(statistics),
statistics,
type_,
options,
Encoding::RleDictionary,
Expand All @@ -159,8 +159,15 @@ macro_rules! dyn_prim {
let values = $array.values().as_any().downcast_ref().unwrap();

let buffer = primitive_encode_plain::<$from, $to>(values, false, vec![]);
let stats = primitive_build_statistics::<$from, $to>(values, $type_.clone());
let stats = serialize_statistics(&stats);

let stats: Option<ParquetStatistics> = if $options.write_statistics {
let mut stats = primitive_build_statistics::<$from, $to>(values, $type_.clone());
stats.null_count = Some($array.null_count() as i64);
let stats = serialize_statistics(&stats);
Some(stats)
} else {
None
};
(DictPage::new(buffer, values.len(), false), stats)
}};
}
Expand All @@ -175,69 +182,93 @@ pub fn array_to_pages<K: DictionaryKey>(
match encoding {
Encoding::PlainDictionary | Encoding::RleDictionary => {
// write DictPage
let (dict_page, statistics) = match array.values().data_type().to_logical_type() {
DataType::Int8 => dyn_prim!(i8, i32, array, options, type_),
DataType::Int16 => dyn_prim!(i16, i32, array, options, type_),
DataType::Int32 | DataType::Date32 | DataType::Time32(_) => {
dyn_prim!(i32, i32, array, options, type_)
}
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
| DataType::Timestamp(_, _)
| DataType::Duration(_) => dyn_prim!(i64, i64, array, options, type_),
DataType::UInt8 => dyn_prim!(u8, i32, array, options, type_),
DataType::UInt16 => dyn_prim!(u16, i32, array, options, type_),
DataType::UInt32 => dyn_prim!(u32, i32, array, options, type_),
DataType::UInt64 => dyn_prim!(u64, i64, array, options, type_),
DataType::Float32 => dyn_prim!(f32, f32, array, options, type_),
DataType::Float64 => dyn_prim!(f64, f64, array, options, type_),
DataType::Utf8 => {
let array = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
utf8_encode_plain::<i32>(array, false, &mut buffer);
let stats = utf8_build_statistics(array, type_.clone());
(DictPage::new(buffer, array.len(), false), stats)
}
DataType::LargeUtf8 => {
let array = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
utf8_encode_plain::<i64>(array, false, &mut buffer);
let stats = utf8_build_statistics(array, type_.clone());
(DictPage::new(buffer, array.len(), false), stats)
}
DataType::Binary => {
let array = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
binary_encode_plain::<i32>(array, false, &mut buffer);
let stats = binary_build_statistics(array, type_.clone());
(DictPage::new(buffer, array.len(), false), stats)
}
DataType::LargeBinary => {
let array = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
binary_encode_plain::<i64>(array, false, &mut buffer);
let stats = binary_build_statistics(array, type_.clone());
(DictPage::new(buffer, array.len(), false), stats)
}
DataType::FixedSizeBinary(_) => {
let mut buffer = vec![];
let array = array.values().as_any().downcast_ref().unwrap();
fixed_binary_encode_plain(array, false, &mut buffer);
let stats = fixed_binary_build_statistics(array, type_.clone());
let stats = serialize_statistics(&stats);
(DictPage::new(buffer, array.len(), false), stats)
}
other => {
return Err(Error::NotYetImplemented(format!(
"Writing dictionary arrays to parquet only support data type {other:?}"
)))
}
};
let (dict_page, statistics): (_, Option<ParquetStatistics>) =
match array.values().data_type().to_logical_type() {
DataType::Int8 => dyn_prim!(i8, i32, array, options, type_),
DataType::Int16 => dyn_prim!(i16, i32, array, options, type_),
DataType::Int32 | DataType::Date32 | DataType::Time32(_) => {
dyn_prim!(i32, i32, array, options, type_)
}
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
| DataType::Timestamp(_, _)
| DataType::Duration(_) => dyn_prim!(i64, i64, array, options, type_),
DataType::UInt8 => dyn_prim!(u8, i32, array, options, type_),
DataType::UInt16 => dyn_prim!(u16, i32, array, options, type_),
DataType::UInt32 => dyn_prim!(u32, i32, array, options, type_),
DataType::UInt64 => dyn_prim!(u64, i64, array, options, type_),
DataType::Float32 => dyn_prim!(f32, f32, array, options, type_),
DataType::Float64 => dyn_prim!(f64, f64, array, options, type_),
DataType::Utf8 => {
let array = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
utf8_encode_plain::<i32>(array, false, &mut buffer);
let stats = if options.write_statistics {
Some(utf8_build_statistics(array, type_.clone()))
} else {
None
};
(DictPage::new(buffer, array.len(), false), stats)
}
DataType::LargeUtf8 => {
let array = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
utf8_encode_plain::<i64>(array, false, &mut buffer);
let stats = if options.write_statistics {
Some(utf8_build_statistics(array, type_.clone()))
} else {
None
};
(DictPage::new(buffer, array.len(), false), stats)
}
DataType::Binary => {
let array = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
binary_encode_plain::<i32>(array, false, &mut buffer);
let stats = if options.write_statistics {
Some(binary_build_statistics(array, type_.clone()))
} else {
None
};
(DictPage::new(buffer, array.len(), false), stats)
}
DataType::LargeBinary => {
let values = array.values().as_any().downcast_ref().unwrap();

let mut buffer = vec![];
binary_encode_plain::<i64>(values, false, &mut buffer);
let stats = if options.write_statistics {
let mut stats = binary_build_statistics(values, type_.clone());
stats.null_count = Some(array.null_count() as i64);
Some(stats)
} else {
None
};
(DictPage::new(buffer, values.len(), false), stats)
}
DataType::FixedSizeBinary(_) => {
let mut buffer = vec![];
let array = array.values().as_any().downcast_ref().unwrap();
fixed_binary_encode_plain(array, false, &mut buffer);
let stats = if options.write_statistics {
let mut stats = fixed_binary_build_statistics(array, type_.clone());
stats.null_count = Some(array.null_count() as i64);
Some(serialize_statistics(&stats))
} else {
None
};
(DictPage::new(buffer, array.len(), false), stats)
}
other => {
return Err(Error::NotYetImplemented(format!(
"Writing dictionary arrays to parquet only support data type {other:?}"
)))
}
};
let dict_page = Page::Dict(dict_page);

// write DataPage pointing to DictPage
Expand Down
2 changes: 1 addition & 1 deletion tests/it/io/parquet/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -591,7 +591,7 @@ pub fn pyarrow_nullable_statistics(column: &str) -> Statistics {

Statistics {
distinct_count: UInt64Array::from([None]).boxed(),
null_count: UInt64Array::from([Some(0)]).boxed(),
null_count: UInt64Array::from([Some(1)]).boxed(),
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test result was invalid. I debugged and clearly found null_count of 1.

min_value: new_dict(Box::new(Int32Array::from_slice([10]))),
max_value: new_dict(Box::new(Int32Array::from_slice([200]))),
}
Expand Down