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

Commit

Permalink
Added support to write FixedLen dictionary
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao committed Jan 30, 2022
1 parent 6adaccf commit 66906ed
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 11 deletions.
7 changes: 7 additions & 0 deletions src/io/parquet/write/dictionary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use parquet2::{
};

use super::binary::encode_plain as binary_encode_plain;
use super::fixed_len_bytes::encode_plain as fixed_binary_encode_plain;
use super::primitive::encode_plain as primitive_encode_plain;
use super::utf8::encode_plain as utf8_encode_plain;
use crate::array::{Array, DictionaryArray, DictionaryKey, PrimitiveArray};
Expand Down Expand Up @@ -172,6 +173,12 @@ pub fn array_to_pages<K: DictionaryKey>(
binary_encode_plain::<i64>(values, false, &mut buffer);
EncodedDictPage::new(buffer, values.len())
}
DataType::FixedSizeBinary(_) => {
let mut buffer = vec![];
let array = array.values().as_any().downcast_ref().unwrap();
fixed_binary_encode_plain(array, false, &mut buffer);
EncodedDictPage::new(buffer, array.len())
}
other => {
return Err(ArrowError::NotYetImplemented(format!(
"Writing dictionary arrays to parquet only support data type {:?}",
Expand Down
25 changes: 14 additions & 11 deletions src/io/parquet/write/fixed_len_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,19 @@ use crate::{
io::parquet::read::is_type_nullable,
};

pub(crate) fn encode_plain(array: &FixedSizeBinaryArray, is_optional: bool, buffer: &mut Vec<u8>) {
// append the non-null values
if is_optional {
array.iter().for_each(|x| {
if let Some(x) = x {
buffer.extend_from_slice(x);
}
})
} else {
buffer.extend_from_slice(array.values());
}
}

pub fn array_to_page(
array: &FixedSizeBinaryArray,
options: WriteOptions,
Expand All @@ -32,17 +45,7 @@ pub fn array_to_page(

let definition_levels_byte_length = buffer.len();

if is_optional {
// append the non-null values
array.iter().for_each(|x| {
if let Some(x) = x {
buffer.extend_from_slice(x);
}
});
} else {
// append all values
buffer.extend_from_slice(array.values());
}
encode_plain(array, is_optional, &mut buffer);

let statistics = if options.write_statistics {
build_statistics(array, descriptor.clone())
Expand Down

0 comments on commit 66906ed

Please sign in to comment.