Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Read nested Parquet 2-level lists correctly #6757

Merged
merged 5 commits into from
Nov 27, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
52 changes: 51 additions & 1 deletion parquet/src/arrow/arrow_reader/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -937,7 +937,7 @@ mod tests {
};
use arrow_array::*;
use arrow_buffer::{i256, ArrowNativeType, Buffer, IntervalDayTime};
use arrow_data::ArrayDataBuilder;
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{
ArrowError, DataType as ArrowDataType, Field, Fields, Schema, SchemaRef, TimeUnit,
};
Expand Down Expand Up @@ -4065,4 +4065,54 @@ mod tests {
}
}
}

#[test]
fn test_read_old_nested_list() {
use arrow::datatypes::DataType;
use arrow::datatypes::ToByteSlice;

let testdata = arrow::util::test_util::parquet_test_data();
// message my_record {

Choose a reason for hiding this comment

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

Nice, thank you for putting a comment with the test data. I like this :)

// REQUIRED group a (LIST) {
// REPEATED group array (LIST) {
// REPEATED INT32 array;
// }
// }
// }
// should be read as list<list<int32>>
let path = format!("{testdata}/old_list_structure.parquet");
let test_file = File::open(path).unwrap();

// create expected ListArray
let a_values = Int32Array::from(vec![1, 2, 3, 4]);

// Construct a buffer for value offsets, for the nested array: [[1, 2], [3, 4]]
let a_value_offsets = arrow::buffer::Buffer::from([0, 2, 4].to_byte_slice());

// Construct a list array from the above two
let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new(
"array",
DataType::Int32,
false,
))))
.len(2)
.add_buffer(a_value_offsets)
.add_child_data(a_values.into_data())
.build()
.unwrap();
let a = ListArray::from(a_list_data);

let builder = ParquetRecordBatchReaderBuilder::try_new(test_file).unwrap();
let mut reader = builder.build().unwrap();
let out = reader.next().unwrap().unwrap();
assert_eq!(out.num_rows(), 1);
assert_eq!(out.num_columns(), 1);
// grab first column
let c0 = out.column(0);
let c0arr = c0.as_any().downcast_ref::<ListArray>().unwrap();
// get first row: [[1, 2], [3, 4]]
let r0 = c0arr.value(0);
let r0arr = r0.as_any().downcast_ref::<ListArray>().unwrap();
assert_eq!(r0arr, &a);
}
}
16 changes: 11 additions & 5 deletions parquet/src/arrow/schema/complex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -448,15 +448,21 @@ impl Visitor {
};
}

// test to see if the repeated field is a struct or one-tuple
let items = repeated_field.get_fields();
if items.len() != 1
|| repeated_field.name() == "array"
|| repeated_field.name() == format!("{}_tuple", list_type.name())
|| (!repeated_field.is_list()
Copy link
Contributor

Choose a reason for hiding this comment

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

Why is this check necessary, given we're in visit_list?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

We're testing if the child of the list we're processing is also LIST annotated.

optional group my_list (LIST) {  <---- this is `list_type`
    repeated group array (LIST) { <---- this is `repeated_field`
        repeated int32 array;
    };
}

&& !repeated_field.has_single_repeated_child()
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
&& !repeated_field.has_single_repeated_child()
&& items[0].get_basic_info().repetition() != Repetition::REPEATED

Copy link
Contributor Author

@etseidl etseidl Nov 22, 2024

Choose a reason for hiding this comment

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

I didn't do this because repetition() can theoretically panic, even though by this point in the schema every node should have a repetition. Perhaps panicking here is warranted if the schema is invalid. I'll change if you'd prefer the panic (or return an error à la #6738).

&& (repeated_field.name() == "array"
|| repeated_field.name() == format!("{}_tuple", list_type.name())))
{
// If the repeated field is a group with multiple fields, then its type is the element type and elements are required.
// If the repeated field is a group with multiple fields, then its type is the element
// type and elements are required.
//
// If the repeated field is a group with one field and is named either array or uses the LIST-annotated group's name
// with _tuple appended then the repeated type is the element type and elements are required.
// If the repeated field is a group with one field and is named either array or uses
// the LIST-annotated group's name with _tuple appended then the repeated type is the
// element type and elements are required. But this rule only applies if the
// repeated field is not annotated, and the single child field is not `repeated`.
let context = VisitorContext {
rep_level: context.rep_level,
def_level,
Expand Down
30 changes: 30 additions & 0 deletions parquet/src/record/reader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,19 @@ impl Reader {
/// <https://github.com/apache/parquet-format/blob/master/LogicalTypes.md>
/// #backward-compatibility-rules
fn is_element_type(repeated_type: &Type) -> bool {
// For legacy 2-level list types whose element type is a 2-level list
//
// // ARRAY<ARRAY<INT>> (nullable list, non-null elements)
// optional group my_list (LIST) {
// repeated group array (LIST) {
// repeated int32 array;
// };
// }
//
if repeated_type.is_list() || repeated_type.has_single_repeated_child() {
return false;
}

// For legacy 2-level list types with primitive element type, e.g.:
//
// // ARRAY<INT> (nullable list, non-null elements)
Expand Down Expand Up @@ -1839,4 +1852,21 @@ mod tests {
let iter = row_group_reader.get_row_iter(schema)?;
Ok(iter.map(|row| row.unwrap()).collect())
}

#[test]
fn test_read_old_nested_list() {
let rows = test_file_reader_rows("old_list_structure.parquet", None).unwrap();
let expected_rows = vec![row![(
"a".to_string(),
Field::ListInternal(make_list(
[
make_list([1, 2].map(Field::Int).to_vec()),
make_list([3, 4].map(Field::Int).to_vec())
]
.map(Field::ListInternal)
.to_vec()
))
),]];
assert_eq!(rows, expected_rows);
}
}
23 changes: 23 additions & 0 deletions parquet/src/schema/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,29 @@ impl Type {
self.get_basic_info().has_repetition()
&& self.get_basic_info().repetition() != Repetition::REQUIRED
}

/// Returns `true` if this type is annotated as a list.
pub fn is_list(&self) -> bool {
etseidl marked this conversation as resolved.
Show resolved Hide resolved
if self.is_group() {
let basic_info = self.get_basic_info();
if let Some(logical_type) = basic_info.logical_type() {
return logical_type == LogicalType::List;
}
return basic_info.converted_type() == ConvertedType::LIST;
}
false
}

/// Returns `true` if this type is a group with a single child field that is `repeated`.
pub fn has_single_repeated_child(&self) -> bool {
etseidl marked this conversation as resolved.
Show resolved Hide resolved
if self.is_group() {
let children = self.get_fields();
return children.len() == 1
&& children[0].get_basic_info().has_repetition()
&& children[0].get_basic_info().repetition() == Repetition::REPEATED;
}
false
}
}

/// A builder for primitive types. All attributes are optional
Expand Down
Loading