You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2024/07/01 10:28:11 UTC
(arrow-rs) branch master updated: feat: support reading OPTIONAL column in parquet_derive (#5717)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 62c1615e8f feat: support reading OPTIONAL column in parquet_derive (#5717)
62c1615e8f is described below
commit 62c1615e8ff822f5aa3adbc1e0acb1b5c59ea2a9
Author: double-free <cr...@126.com>
AuthorDate: Mon Jul 1 18:28:05 2024 +0800
feat: support reading OPTIONAL column in parquet_derive (#5717)
* support def_level=1 but non-null column in reader
* update comment, adapt ut to the uuid change
---------
Co-authored-by: Ye Yuan <yu...@qq.com>
---
parquet_derive/src/parquet_field.rs | 59 +++++++++++++++++++++++++++---------
parquet_derive_test/src/lib.rs | 60 +++++++++++++++++++++++++++++++++++++
2 files changed, 105 insertions(+), 14 deletions(-)
diff --git a/parquet_derive/src/parquet_field.rs b/parquet_derive/src/parquet_field.rs
index 9fff76c42d..f99ea3e035 100644
--- a/parquet_derive/src/parquet_field.rs
+++ b/parquet_derive/src/parquet_field.rs
@@ -239,7 +239,8 @@ impl Field {
/// because this parsing logic is not sophisticated enough for definition
/// levels beyond 2.
///
- /// `Option` types and references not supported
+ /// `Option` types and references not supported, but the column itself can be nullable
+ /// (i.e., def_level==1), as long as the values are all valid.
pub fn reader_snippet(&self) -> proc_macro2::TokenStream {
let ident = &self.ident;
let column_reader = self.ty.column_reader();
@@ -248,7 +249,13 @@ impl Field {
let write_batch_expr = quote! {
let mut vals = Vec::new();
if let #column_reader(mut typed) = column_reader {
- typed.read_records(num_records, None, None, &mut vals)?;
+ let mut definition_levels = Vec::new();
+ let (total_num, valid_num, decoded_num) = typed.read_records(
+ num_records, Some(&mut definition_levels), None, &mut vals)?;
+ if valid_num != decoded_num {
+ panic!("Support only valid records, found {} null records in column type {}",
+ decoded_num - valid_num, stringify!{#ident});
+ }
} else {
panic!("Schema and struct disagree on type for {}", stringify!{#ident});
}
@@ -876,15 +883,21 @@ mod test {
snippet,
(quote! {
{
- let mut vals = Vec::new();
- if let ColumnReader::Int64ColumnReader(mut typed) = column_reader {
- typed.read_records(num_records, None, None, &mut vals)?;
- } else {
- panic!("Schema and struct disagree on type for {}", stringify!{ counter });
- }
- for (i, r) in &mut records[..num_records].iter_mut().enumerate() {
- r.counter = vals[i] as usize;
- }
+ let mut vals = Vec::new();
+ if let ColumnReader::Int64ColumnReader(mut typed) = column_reader {
+ let mut definition_levels = Vec::new();
+ let (total_num, valid_num, decoded_num) = typed.read_records(
+ num_records, Some(&mut definition_levels), None, &mut vals)?;
+ if valid_num != decoded_num {
+ panic!("Support only valid records, found {} null records in column type {}",
+ decoded_num - valid_num, stringify!{counter});
+ }
+ } else {
+ panic!("Schema and struct disagree on type for {}", stringify!{counter});
+ }
+ for (i, r) in &mut records[..num_records].iter_mut().enumerate() {
+ r.counter = vals[i] as usize;
+ }
}
})
.to_string()
@@ -1291,7 +1304,13 @@ mod test {
{
let mut vals = Vec::new();
if let ColumnReader::Int64ColumnReader(mut typed) = column_reader {
- typed.read_records(num_records, None, None, &mut vals)?;
+ let mut definition_levels = Vec::new();
+ let (total_num, valid_num, decoded_num) = typed.read_records(
+ num_records, Some(&mut definition_levels), None, &mut vals)?;
+ if valid_num != decoded_num {
+ panic!("Support only valid records, found {} null records in column type {}",
+ decoded_num - valid_num, stringify!{henceforth});
+ }
} else {
panic!("Schema and struct disagree on type for {}", stringify!{ henceforth });
}
@@ -1359,7 +1378,13 @@ mod test {
{
let mut vals = Vec::new();
if let ColumnReader::Int32ColumnReader(mut typed) = column_reader {
- typed.read_records(num_records, None, None, &mut vals)?;
+ let mut definition_levels = Vec::new();
+ let (total_num, valid_num, decoded_num) = typed.read_records(
+ num_records, Some(&mut definition_levels), None, &mut vals)?;
+ if valid_num != decoded_num {
+ panic!("Support only valid records, found {} null records in column type {}",
+ decoded_num - valid_num, stringify!{henceforth});
+ }
} else {
panic!("Schema and struct disagree on type for {}", stringify!{ henceforth });
}
@@ -1427,7 +1452,13 @@ mod test {
{
let mut vals = Vec::new();
if let ColumnReader::FixedLenByteArrayColumnReader(mut typed) = column_reader {
- typed.read_records(num_records, None, None, &mut vals)?;
+ let mut definition_levels = Vec::new();
+ let (total_num, valid_num, decoded_num) = typed.read_records(
+ num_records, Some(&mut definition_levels), None, &mut vals)?;
+ if valid_num != decoded_num {
+ panic!("Support only valid records, found {} null records in column type {}",
+ decoded_num - valid_num, stringify!{unique_id});
+ }
} else {
panic!("Schema and struct disagree on type for {}", stringify!{ unique_id });
}
diff --git a/parquet_derive_test/src/lib.rs b/parquet_derive_test/src/lib.rs
index b8e6c27f53..310a3d8c90 100644
--- a/parquet_derive_test/src/lib.rs
+++ b/parquet_derive_test/src/lib.rs
@@ -66,6 +66,25 @@ struct APartiallyCompleteRecord {
pub byte_vec: Vec<u8>,
}
+// This struct has OPTIONAL columns
+// If these fields are guaranteed to be valid
+// we can load this struct into APartiallyCompleteRecord
+#[derive(PartialEq, ParquetRecordWriter, Debug)]
+struct APartiallyOptionalRecord {
+ pub bool: bool,
+ pub string: String,
+ pub maybe_i16: Option<i16>,
+ pub maybe_i32: Option<i32>,
+ pub maybe_u64: Option<u64>,
+ pub isize: isize,
+ pub float: f32,
+ pub double: f64,
+ pub now: chrono::NaiveDateTime,
+ pub date: chrono::NaiveDate,
+ pub uuid: uuid::Uuid,
+ pub byte_vec: Vec<u8>,
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -218,6 +237,47 @@ mod tests {
assert_eq!(drs[0], out[0]);
}
+ #[test]
+ fn test_parquet_derive_read_optional_but_valid_column() {
+ let file = get_temp_file("test_parquet_derive_read_optional", &[]);
+ let drs: Vec<APartiallyOptionalRecord> = vec![APartiallyOptionalRecord {
+ bool: true,
+ string: "a string".into(),
+ maybe_i16: Some(-45),
+ maybe_i32: Some(456),
+ maybe_u64: Some(4563424),
+ isize: -365,
+ float: 3.5,
+ double: std::f64::NAN,
+ now: chrono::Utc::now().naive_local(),
+ date: chrono::naive::NaiveDate::from_ymd_opt(2015, 3, 14).unwrap(),
+ uuid: uuid::Uuid::new_v4(),
+ byte_vec: vec![0x65, 0x66, 0x67],
+ }];
+
+ let generated_schema = drs.as_slice().schema().unwrap();
+
+ let props = Default::default();
+ let mut writer =
+ SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
+
+ let mut row_group = writer.next_row_group().unwrap();
+ drs.as_slice().write_to_row_group(&mut row_group).unwrap();
+ row_group.close().unwrap();
+ writer.close().unwrap();
+
+ use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
+ let reader = SerializedFileReader::new(file).unwrap();
+ let mut out: Vec<APartiallyCompleteRecord> = Vec::new();
+
+ let mut row_group = reader.get_row_group(0).unwrap();
+ out.read_from_row_group(&mut *row_group, 1).unwrap();
+
+ assert_eq!(drs[0].maybe_i16.unwrap(), out[0].i16);
+ assert_eq!(drs[0].maybe_i32.unwrap(), out[0].i32);
+ assert_eq!(drs[0].maybe_u64.unwrap(), out[0].u64);
+ }
+
/// Returns file handle for a temp file in 'target' directory with a provided content
pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
// build tmp path to a file in "target/debug/testdata"