rluvaton commented on code in PR #9086:
URL: https://github.com/apache/arrow-rs/pull/9086#discussion_r2681931811
##########
arrow-json/src/reader/struct_array.rs:
##########
@@ -38,131 +43,162 @@ impl StructArrayDecoder {
is_nullable: bool,
struct_mode: StructMode,
) -> Result<Self, ArrowError> {
- let decoders = struct_fields(&data_type)
- .iter()
- .map(|f| {
- // If this struct nullable, need to permit nullability in
child array
- // StructArrayDecoder::decode verifies that if the child is
not nullable
- // it doesn't contain any nulls not masked by its parent
- let nullable = f.is_nullable() || is_nullable;
- make_decoder(
- f.data_type().clone(),
- coerce_primitive,
- strict_mode,
- nullable,
- struct_mode,
- )
- })
- .collect::<Result<Vec<_>, ArrowError>>()?;
+ let (decoders, field_name_to_index) = {
+ let fields = struct_fields(&data_type);
+ let decoders = fields
+ .iter()
+ .map(|f| {
+ // If this struct nullable, need to permit nullability in
child array
+ // StructArrayDecoder::decode verifies that if the child
is not nullable
+ // it doesn't contain any nulls not masked by its parent
+ let nullable = f.is_nullable() || is_nullable;
+ make_decoder(
+ f.data_type().clone(),
+ coerce_primitive,
+ strict_mode,
+ nullable,
+ struct_mode,
+ )
+ })
+ .collect::<Result<Vec<_>, ArrowError>>()?;
+ let field_name_to_index = if struct_mode == StructMode::ObjectOnly
{
+ build_field_index(fields)
+ } else {
+ None
+ };
+ (decoders, field_name_to_index)
+ };
Ok(Self {
data_type,
decoders,
strict_mode,
is_nullable,
struct_mode,
+ field_name_to_index,
+ field_tape_positions: Vec::new(),
})
}
}
impl ArrayDecoder for StructArrayDecoder {
fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayData,
ArrowError> {
let fields = struct_fields(&self.data_type);
- let mut child_pos: Vec<_> = (0..fields.len()).map(|_| vec![0;
pos.len()]).collect();
-
+ let row_count = pos.len();
+ let field_count = fields.len();
+ let total_len = field_count.checked_mul(row_count).ok_or_else(|| {
+ ArrowError::JsonError(format!(
+ "StructArrayDecoder child position buffer size overflow for
rows={row_count} fields={field_count}"
+ ))
+ })?;
+ self.field_tape_positions.clear();
+ self.field_tape_positions.resize(total_len, 0);
let mut nulls = self
.is_nullable
.then(|| BooleanBufferBuilder::new(pos.len()));
- // We avoid having the match on self.struct_mode inside the hot loop
for performance
- // TODO: Investigate how to extract duplicated logic.
- match self.struct_mode {
- StructMode::ObjectOnly => {
- for (row, p) in pos.iter().enumerate() {
- let end_idx = match (tape.get(*p), nulls.as_mut()) {
- (TapeElement::StartObject(end_idx), None) => end_idx,
- (TapeElement::StartObject(end_idx), Some(nulls)) => {
- nulls.append(true);
- end_idx
- }
- (TapeElement::Null, Some(nulls)) => {
- nulls.append(false);
- continue;
- }
- (_, _) => return Err(tape.error(*p, "{")),
- };
-
- let mut cur_idx = *p + 1;
- while cur_idx < end_idx {
- // Read field name
- let field_name = match tape.get(cur_idx) {
- TapeElement::String(s) => tape.get_string(s),
- _ => return Err(tape.error(cur_idx, "field name")),
+ {
+ let child_pos = self.field_tape_positions.as_mut_slice();
+ // We avoid having the match on self.struct_mode inside the hot
loop for performance
+ // TODO: Investigate how to extract duplicated logic.
+ match self.struct_mode {
+ StructMode::ObjectOnly => {
+ for (row, p) in pos.iter().enumerate() {
+ let end_idx = match (tape.get(*p), nulls.as_mut()) {
+ (TapeElement::StartObject(end_idx), None) =>
end_idx,
+ (TapeElement::StartObject(end_idx), Some(nulls))
=> {
+ nulls.append(true);
+ end_idx
+ }
+ (TapeElement::Null, Some(nulls)) => {
+ nulls.append(false);
+ continue;
+ }
+ (_, _) => return Err(tape.error(*p, "{")),
};
- // Update child pos if match found
- match fields.iter().position(|x| x.name() ==
field_name) {
- Some(field_idx) => child_pos[field_idx][row] =
cur_idx + 1,
- None => {
- if self.strict_mode {
- return Err(ArrowError::JsonError(format!(
- "column '{field_name}' missing from
schema",
- )));
+ let mut cur_idx = *p + 1;
+ while cur_idx < end_idx {
+ // Read field name
+ let field_name = match tape.get(cur_idx) {
+ TapeElement::String(s) => tape.get_string(s),
+ _ => return Err(tape.error(cur_idx, "field
name")),
+ };
+
+ // Update child pos if match found
+ let field_idx = match &self.field_name_to_index {
+ Some(map) => map.get(field_name).copied(),
+ None => fields.iter().position(|x| x.name() ==
field_name),
+ };
+ match field_idx {
+ Some(field_idx) => {
+ child_pos[field_idx * row_count + row] =
cur_idx + 1;
+ }
+ None => {
+ if self.strict_mode {
+ return
Err(ArrowError::JsonError(format!(
+ "column '{field_name}' missing
from schema",
+ )));
+ }
}
}
+ // Advance to next field
+ cur_idx = tape.next(cur_idx + 1, "field value")?;
}
- // Advance to next field
- cur_idx = tape.next(cur_idx + 1, "field value")?;
}
}
- }
- StructMode::ListOnly => {
- for (row, p) in pos.iter().enumerate() {
- let end_idx = match (tape.get(*p), nulls.as_mut()) {
- (TapeElement::StartList(end_idx), None) => end_idx,
- (TapeElement::StartList(end_idx), Some(nulls)) => {
- nulls.append(true);
- end_idx
- }
- (TapeElement::Null, Some(nulls)) => {
- nulls.append(false);
- continue;
- }
- (_, _) => return Err(tape.error(*p, "[")),
- };
+ StructMode::ListOnly => {
+ for (row, p) in pos.iter().enumerate() {
+ let end_idx = match (tape.get(*p), nulls.as_mut()) {
+ (TapeElement::StartList(end_idx), None) => end_idx,
+ (TapeElement::StartList(end_idx), Some(nulls)) => {
+ nulls.append(true);
+ end_idx
+ }
+ (TapeElement::Null, Some(nulls)) => {
+ nulls.append(false);
+ continue;
+ }
+ (_, _) => return Err(tape.error(*p, "[")),
+ };
- let mut cur_idx = *p + 1;
- let mut entry_idx = 0;
- while cur_idx < end_idx {
- if entry_idx >= fields.len() {
+ let mut cur_idx = *p + 1;
+ let mut entry_idx = 0;
+ while cur_idx < end_idx {
+ if entry_idx >= fields.len() {
+ return Err(ArrowError::JsonError(format!(
+ "found extra columns for {} fields",
+ fields.len()
+ )));
+ }
+ child_pos[entry_idx * row_count + row] = cur_idx;
Review Comment:
This can also be part of the dedicated struct
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
