Add support for nested list arrays from parquet to arrow arrays (#993) #1588
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -38,7 +38,7 @@ use crate::data_type::{ | |
| Int96Type, | ||
| }; | ||
| use crate::errors::ParquetError::ArrowError; | ||
| use crate::errors::{ParquetError, Result}; | ||
| use crate::schema::types::{ColumnDescriptor, ColumnPath, SchemaDescPtr, Type, TypePtr}; | ||
| use crate::schema::visitor::TypeVisitor; | ||
|
|
||
|
|
@@ -129,9 +129,10 @@ impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext | |
|
|
||
| let null_mask_only = match cur_type.get_basic_info().repetition() { | ||
| Repetition::REPEATED => { | ||
| return Err(ArrowError(format!( | ||
| "Reading repeated primitive ({:?}) is not supported yet!", | ||
| cur_type.name() | ||
| ))); | ||
| } | ||
| Repetition::OPTIONAL => { | ||
| new_context.def_level += 1; | ||
|
|
@@ -143,19 +144,12 @@ impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext | |
| }; | ||
|
|
||
| let reader = self.build_for_primitive_type_inner( | ||
| cur_type, | ||
| &new_context, | ||
| null_mask_only, | ||
| )?; | ||
|
|
||
| Ok(Some(reader)) | ||
| } else { | ||
| Ok(None) | ||
| } | ||
|
|
@@ -173,30 +167,19 @@ impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext | |
| if cur_type.get_basic_info().has_repetition() { | ||
| match cur_type.get_basic_info().repetition() { | ||
| Repetition::REPEATED => { | ||
| return Err(ArrowError(format!( | ||
| "Reading repeated struct ({:?}) is not supported yet!", | ||
| cur_type.name(), | ||
| ))) | ||
| } | ||
| Repetition::OPTIONAL => { | ||
| new_context.def_level += 1; | ||
| } | ||
| Repetition::REQUIRED => {} | ||
| } | ||
| } | ||
|
|
||
| self.build_for_struct_type_inner(&cur_type, &new_context) | ||
| } | ||
|
|
||
| /// Build array reader for map type. | ||
|
|
@@ -208,42 +191,61 @@ impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext | |
| // Add map type to context | ||
| let mut new_context = context.clone(); | ||
| new_context.path.append(vec![map_type.name().to_string()]); | ||
|
|
||
| match map_type.get_basic_info().repetition() { | ||
| Repetition::REQUIRED => {} | ||
| Repetition::OPTIONAL => { | ||
| new_context.def_level += 1; | ||
| } | ||
| Repetition::REPEATED => { | ||
| return Err(ArrowError("Map cannot be repeated".to_string())) | ||
| } | ||
| } | ||
|
|
||
| if map_type.get_fields().len() != 1 { | ||
| return Err(ArrowError(format!( | ||
| "Map field must have exactly one key_value child, found {}", | ||
| map_type.get_fields().len() | ||
| ))); | ||
| } | ||
|
|
||
| // Add map entry (key_value) to context | ||
| let map_key_value = &map_type.get_fields()[0]; | ||
| if map_key_value.get_basic_info().repetition() != Repetition::REPEATED { | ||
| return Err(ArrowError( | ||
| "Child of map field must be repeated".to_string(), | ||
| )); | ||
| } | ||
|
|
||
| new_context | ||
| .path | ||
| .append(vec![map_key_value.name().to_string()]); | ||
|
|
||
| new_context.rep_level += 1; | ||
| new_context.def_level += 1; | ||
|
|
||
| if map_key_value.get_fields().len() != 2 { | ||
| // According to the specification the values are optional (#1642) | ||
| return Err(ArrowError(format!( | ||
| "Child of map field must have two children, found {}", | ||
| map_key_value.get_fields().len() | ||
| ))); | ||
| } | ||
|
|
||
| // Get key and value, and create context for each | ||
|
There was a problem hiding this comment. This is a drive-by fix, the context is the context of the parent, not the value being dispatched. This would result in map_key appearing twice |
||
| let map_key = &map_key_value.get_fields()[0]; | ||
| let map_value = &map_key_value.get_fields()[1]; | ||
|
|
||
| if map_key.get_basic_info().repetition() != Repetition::REQUIRED { | ||
| return Err(ArrowError("Map keys must be required".to_string())); | ||
| } | ||
|
|
||
| if map_value.get_basic_info().repetition() == Repetition::REPEATED { | ||
| return Err(ArrowError("Map values cannot be repeated".to_string())); | ||
| } | ||
|
|
||
| let key_reader = self.dispatch(map_key.clone(), &new_context)?.unwrap(); | ||
| let value_reader = self.dispatch(map_value.clone(), &new_context)?.unwrap(); | ||
|
|
||
| let arrow_type = self | ||
| .arrow_schema | ||
|
|
@@ -295,96 +297,80 @@ impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext | |
| .first() | ||
| .ok_or_else(|| ArrowError("List field must have a child.".to_string()))? | ||
| .clone(); | ||
|
|
||
| // If the list can contain nulls | ||
|
There was a problem hiding this comment. I recommend reading this with whitespace disabled https://github.com/apache/arrow-rs/pull/1588/files?w=1 |
||
| let nullable = match list_type.get_basic_info().repetition() { | ||
| Repetition::REQUIRED => false, | ||
| Repetition::OPTIONAL => true, | ||
| Repetition::REPEATED => { | ||
| return Err(general_err!("List type cannot be repeated")) | ||
| } | ||
| }; | ||
|
|
||
| let mut new_context = context.clone(); | ||
| new_context.path.append(vec![list_type.name().to_string()]); | ||
|
|
||
| // The repeated field | ||
| new_context.rep_level += 1; | ||
| new_context.def_level += 1; | ||
|
|
||
| // If the list's root is nullable | ||
| if nullable { | ||
| new_context.def_level += 1; | ||
| } | ||
|
|
||
| match self.dispatch(item_type, &new_context) { | ||
| Ok(Some(item_reader)) => { | ||
| let item_type = item_reader.get_data_type().clone(); | ||
|
|
||
| // a list is a group type with a single child. The list child's | ||
| // name comes from the child's field name. | ||
| // if the child's name is "list" and it has a child, then use this child | ||
| if list_child.name() == "list" && !list_child.get_fields().is_empty() { | ||
| list_child = list_child.get_fields().first().unwrap(); | ||
| } | ||
|
|
||
|
There was a problem hiding this comment. This is the error that would previously be returned on a nested list |
||
| let arrow_type = self | ||
| .arrow_schema | ||
| .field_with_name(list_type.name()) | ||
| .ok() | ||
| .map(|f| f.data_type().to_owned()) | ||
| .unwrap_or_else(|| { | ||
| ArrowType::List(Box::new(Field::new( | ||
| list_child.name(), | ||
| item_type.clone(), | ||
| list_child.is_optional(), | ||
| ))) | ||
| }); | ||
|
|
||
| let list_array_reader: Box<dyn ArrayReader> = match arrow_type { | ||
| ArrowType::List(_) => Box::new(ListArrayReader::<i32>::new( | ||
| item_reader, | ||
| arrow_type, | ||
| item_type, | ||
| new_context.def_level, | ||
| new_context.rep_level, | ||
| nullable, | ||
| )), | ||
| ArrowType::LargeList(_) => Box::new(ListArrayReader::<i64>::new( | ||
| item_reader, | ||
| arrow_type, | ||
| item_type, | ||
| new_context.def_level, | ||
| new_context.rep_level, | ||
| nullable, | ||
| )), | ||
| _ => { | ||
| return Err(ArrowError(format!( | ||
| "creating ListArrayReader with type {:?} should be unreachable", | ||
| arrow_type | ||
| ))) | ||
| } | ||
| }; | ||
|
|
||
| Ok(Some(list_array_reader)) | ||
| } | ||
| result => result, | ||
| } | ||
| } | ||
| } | ||
|
|
@@ -637,10 +623,10 @@ impl<'a> ArrayReaderBuilder { | |
| let mut children_reader = Vec::with_capacity(cur_type.get_fields().len()); | ||
|
|
||
| for child in cur_type.get_fields() { | ||
| if let Some(child_reader) = self.dispatch(child.clone(), context)? { | ||
| // TODO: this results in calling get_arrow_field twice, it could be reused | ||
| // from child_reader above, by making child_reader carry its `Field` | ||
| let mut struct_context = context.clone(); | ||
| struct_context.path.append(vec![child.name().to_string()]); | ||
| let field = match self.get_arrow_field(child, &struct_context) { | ||
| Some(f) => f.clone(), | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This just moves the error to earlier, there is no point continuing if this is not supported anyway 😁