nevi-me commented on a change in pull request #8792: URL: https://github.com/apache/arrow/pull/8792#discussion_r532209572
########## File path: rust/parquet/src/arrow/levels.rs ########## @@ -0,0 +1,692 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains the logic for computing definition and repetition levels + +#[derive(Debug, Eq, PartialEq, Clone)] +pub(crate) struct LevelInfo { + /// Array's definition levels + pub definition: Vec<i16>, + /// Array's optional repetition levels + pub repetition: Option<Vec<i16>>, + /// Definition mask, to indicate null ListArray slots that should be skipped + pub definition_mask: Vec<(bool, i16)>, + /// Array's offsets, 64-bit is used to accommodate large offset arrays + pub array_offsets: Vec<i64>, + /// Array's validity mask + pub array_mask: Vec<bool>, + /// The maximum definition at this level, 0 at the root (record batch) [TODO: the 0 might be inaccurate] + pub max_definition: i16, + /// Whether this array or any of its parents is a list + pub is_list: bool, + /// Whether the array is nullable (affects definition levels) + pub is_nullable: bool, +} + +impl LevelInfo { + fn calculate_child_levels( + &self, + array_offsets: Vec<i64>, + array_mask: Vec<bool>, + is_list: bool, + is_nullable: bool, + current_def_level: i16, + ) -> Self { + let mut definition = vec![]; + let mut repetition = vec![]; + let mut definition_mask = vec![]; + let has_repetition = self.is_list || is_list; + + // keep track of parent definition nulls seen through the definition_mask + let mut nulls_seen = 0; + + // push any initial array slots that are null + while !self.definition_mask[nulls_seen].0 + && self.definition_mask[nulls_seen].1 + 2 < current_def_level + { + definition_mask.push(self.definition_mask[nulls_seen]); + definition.push(self.definition[nulls_seen]); + repetition.push(0); // TODO is it always 0? + nulls_seen += 1; + println!("Definition length e: {}", definition.len()); + } + + // we use this index to determine if a repetition should be populated based + // on its definition at the index. It needs to be outside of the loop + let mut def_index = 0; + + self.array_offsets.windows(2).for_each(|w| { + // the parent's index allows us to iterate through its offsets and the child's + let from = w[0] as usize; + let to = w[1] as usize; + // dbg!((from, to)); + // if the parent slot is empty, fill it once to show the nullness + if from == to { + definition.push(self.max_definition - 1); + repetition.push(0); + definition_mask.push((false, self.max_definition - 1)); + println!("Definition length d: {}", definition.len()); + } + + (from..to).for_each(|index| { + println!( + "Array level: {}, parent offset: {}", + current_def_level, index + ); + let parent_mask = &self.definition_mask[index + nulls_seen]; + // TODO: this might need to be < instead of ==, but we generate duplicates in that case + if !parent_mask.0 && parent_mask.1 == current_def_level { + println!("Parent mask c: {:?}", parent_mask); + nulls_seen += 1; + definition.push(self.max_definition); + repetition.push(1); + definition_mask.push(*parent_mask); + println!("Definition length c: {}", definition.len()); + } + let mask = array_mask[index]; + let array_from = array_offsets[index]; + let array_to = array_offsets[index + 1]; + + let parent_def_level = &self.definition[index + nulls_seen]; + + // if array_len == 0, the child is null + let array_len = array_to - array_from; + + // compute the definition level + // what happens if array's len is 0? + if array_len == 0 { + definition.push(self.max_definition); + repetition.push(0); // TODO: validate that this is 0 for deeply nested lists + definition_mask.push((false, current_def_level)); + println!("Definition length b: {}", definition.len()); + } + (array_from..array_to).for_each(|_| { + definition.push(if *parent_def_level == self.max_definition { + // TODO: haven't validated this in deeply-nested lists + self.max_definition + mask as i16 + } else { + *parent_def_level + }); + definition_mask.push((true, current_def_level)); + println!("Definition length a: {}", definition.len()); + }); + + // 11-11-2020 (23:57GMT) + // we are pushing defined repetitions even if a definition is < max + // I had initially separated the repetition logic here so that I + // don't perform a `has_repetition` check on each loop. + // The downside's that I now need to index into `definitions` so I + // can check if a value is defined or not. + + if has_repetition && array_len > 0 { + // compute the repetition level + + // dbg!(&definition); + // dbg!(current_def_level, parent_level.max_definition); + // dbg!(&parent_level.repetition); + match &self.repetition { + Some(rep) => { + let parent_rep = rep[index]; + // TODO(11/11/2020) need correct variable to mask repetitions correctly + if definition[def_index] == current_def_level { + repetition.push(parent_rep); + println!("* Index {} definition is {}, and repetition is {}. Current def: {}", def_index, definition[def_index], parent_rep, current_def_level); + dbg!(&repetition); + def_index += 1; + (1..array_len).for_each(|_| { + println!("* Index {} definition is {}, and repetition is {}. Current def: {}", def_index, definition[def_index], parent_rep, current_def_level); + repetition.push(current_def_level); // was parent_rep + 1 + def_index += 1; + }); + } else { + (0..array_len).for_each(|_| { + println!("* Index {} definition is {}, and repetition is {}. Current def: {}", def_index, definition[def_index], parent_rep, current_def_level); + repetition.push(0); // TODO: should it be anything else? + // TODO: use an append instead of pushes + def_index += 1; + }); + } + } + None => { + println!("+ Index {} definition is {}, and repetition is 0. Current def: {}", def_index, definition[def_index], current_def_level); + // if definition[def_index] == current_def_level { + repetition.push(0); + def_index += 1; + (1..array_len).for_each(|_| { + repetition.push(1); // TODO: is it always 0 and 1? + def_index += 1; + }); + // } else { + // (0..array_len).for_each(|_| { + // repetition.push(0); // TODO: should it be anything else? + // // TODO: use an append instead of pushes + // def_index += 1; + // }); + // } + } + } + } + }); + }); + + let lev = LevelInfo { + definition, + repetition: if !has_repetition { + None + } else { + Some(repetition) + }, + definition_mask, + array_mask, + array_offsets, + is_list: has_repetition, + max_definition: current_def_level, + is_nullable, + }; + + println!("done"); + + lev + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_calculate_array_levels_twitter_example() { + // based on the example at https://blog.twitter.com/engineering/en_us/a/2013/dremel-made-simple-with-parquet.html Review comment: I've removed this for now, stashed them somewhere, as these are relevant for list tests. So the `levels.rs` is currently indirectly tested by the roundtrip tests. ---------------------------------------------------------------- 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. For queries about this service, please contact Infrastructure at: us...@infra.apache.org
