diff --git a/parquet/src/arrow/arrow_writer/levels.rs b/parquet/src/arrow/arrow_writer/levels.rs index d23873278ed3..1f66de349df0 100644 --- a/parquet/src/arrow/arrow_writer/levels.rs +++ b/parquet/src/arrow/arrow_writer/levels.rs @@ -41,6 +41,7 @@ //! \[1\] [parquet-format#nested-encoding](https://github.com/apache/parquet-format#nested-encoding) use crate::column::chunker::CdcChunk; +use crate::column::writer::LevelDataRef; use crate::errors::{ParquetError, Result}; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, OffsetSizeTrait}; @@ -316,7 +317,7 @@ impl LevelInfoBuilder { |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| { child.write(start_idx..end_idx); child.visit_leaves(|leaf| { - let rep_levels = leaf.rep_levels.as_mut().unwrap(); + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); let mut rev = rep_levels.iter_mut().rev(); let mut remaining = end_idx - start_idx; @@ -336,17 +337,18 @@ impl LevelInfoBuilder { }) }; + // In a list column, each row falls into one of three categories: + // - "null": the list slot is absent (!is_valid), encoded at def_level - 2 + // - "empty": the list slot is present but has zero elements + // (offsets[i] == offsets[i+1]), encoded at def_level - 1 + // - non-empty: the list slot has child values, which are recursed into + // + // Consecutive runs of null or empty rows are batched and written together. let write_null_run = |child: &mut LevelInfoBuilder, count: usize| { if count > 0 { child.visit_leaves(|leaf| { - leaf.rep_levels - .as_mut() - .unwrap() - .extend(std::iter::repeat_n(ctx.rep_level - 1, count)); - leaf.def_levels - .as_mut() - .unwrap() - .extend(std::iter::repeat_n(ctx.def_level - 2, count)); + leaf.append_rep_level_run(ctx.rep_level - 1, count); + leaf.append_def_level_run(ctx.def_level - 2, count); }); } }; @@ -354,14 +356,8 @@ impl LevelInfoBuilder { let write_empty_run = |child: &mut LevelInfoBuilder, count: usize| { if count > 0 { child.visit_leaves(|leaf| { - leaf.rep_levels - .as_mut() - .unwrap() - .extend(std::iter::repeat_n(ctx.rep_level - 1, count)); - leaf.def_levels - .as_mut() - .unwrap() - .extend(std::iter::repeat_n(ctx.def_level - 1, count)); + leaf.append_rep_level_run(ctx.rep_level - 1, count); + leaf.append_def_level_run(ctx.def_level - 1, count); }); } }; @@ -431,7 +427,7 @@ impl LevelInfoBuilder { |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| { child.write(start_idx..end_idx); child.visit_leaves(|leaf| { - let rep_levels = leaf.rep_levels.as_mut().unwrap(); + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); let mut rev = rep_levels.iter_mut().rev(); let mut remaining = end_idx - start_idx; @@ -453,19 +449,15 @@ impl LevelInfoBuilder { let write_empty_slice = |child: &mut LevelInfoBuilder| { child.visit_leaves(|leaf| { - let rep_levels = leaf.rep_levels.as_mut().unwrap(); - rep_levels.push(ctx.rep_level - 1); - let def_levels = leaf.def_levels.as_mut().unwrap(); - def_levels.push(ctx.def_level - 1); + leaf.append_rep_level_run(ctx.rep_level - 1, 1); + leaf.append_def_level_run(ctx.def_level - 1, 1); }) }; let write_null_slice = |child: &mut LevelInfoBuilder| { child.visit_leaves(|leaf| { - let rep_levels = leaf.rep_levels.as_mut().unwrap(); - rep_levels.push(ctx.rep_level - 1); - let def_levels = leaf.def_levels.as_mut().unwrap(); - def_levels.push(ctx.def_level - 2); + leaf.append_rep_level_run(ctx.rep_level - 1, 1); + leaf.append_def_level_run(ctx.def_level - 2, 1); }) }; @@ -513,13 +505,8 @@ impl LevelInfoBuilder { for child in children { child.visit_leaves(|info| { let len = range.end - range.start; - - let def_levels = info.def_levels.as_mut().unwrap(); - def_levels.extend(std::iter::repeat_n(ctx.def_level - 1, len)); - - if let Some(rep_levels) = info.rep_levels.as_mut() { - rep_levels.extend(std::iter::repeat_n(ctx.rep_level, len)); - } + info.append_def_level_run(ctx.def_level - 1, len); + info.append_rep_level_run(ctx.rep_level, len); }) } }; @@ -579,7 +566,7 @@ impl LevelInfoBuilder { child.write(values_start..values_end); child.visit_leaves(|leaf| { - let rep_levels = leaf.rep_levels.as_mut().unwrap(); + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); let row_indices = (0..fixed_size) .rev() @@ -605,10 +592,8 @@ impl LevelInfoBuilder { let write_empty = |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| { let len = end_idx - start_idx; child.visit_leaves(|leaf| { - let rep_levels = leaf.rep_levels.as_mut().unwrap(); - rep_levels.extend(std::iter::repeat_n(ctx.rep_level - 1, len)); - let def_levels = leaf.def_levels.as_mut().unwrap(); - def_levels.extend(std::iter::repeat_n(ctx.def_level - 1, len)); + leaf.append_rep_level_run(ctx.rep_level - 1, len); + leaf.append_def_level_run(ctx.def_level - 1, len); }) }; @@ -634,10 +619,8 @@ impl LevelInfoBuilder { } // Add null row child.visit_leaves(|leaf| { - let rep_levels = leaf.rep_levels.as_mut().unwrap(); - rep_levels.push(ctx.rep_level - 1); - let def_levels = leaf.def_levels.as_mut().unwrap(); - def_levels.push(ctx.def_level - 2); + leaf.append_rep_level_run(ctx.rep_level - 1, 1); + leaf.append_def_level_run(ctx.def_level - 2, 1); }) } } @@ -655,37 +638,34 @@ impl LevelInfoBuilder { fn write_leaf(info: &mut ArrayLevels, range: Range) { let len = range.end - range.start; - match &mut info.def_levels { - Some(def_levels) => { - def_levels.reserve(len); - info.non_null_indices.reserve(len); - - match &info.logical_nulls { - Some(nulls) => { - assert!(range.end <= nulls.len()); - let nulls = nulls.inner(); - def_levels.extend(range.clone().map(|i| { - // Safety: range.end was asserted to be in bounds earlier - let valid = unsafe { nulls.value_unchecked(i) }; - info.max_def_level - (!valid as i16) - })); - info.non_null_indices.extend( - BitIndexIterator::new(nulls.inner(), nulls.offset() + range.start, len) - .map(|i| i + range.start), - ); - } - None => { - let iter = std::iter::repeat_n(info.max_def_level, len); - def_levels.extend(iter); - info.non_null_indices.extend(range); - } + if matches!(info.def_levels, LevelData::Absent) { + info.non_null_indices.extend(range.clone()); + } else { + let max_def_level = info.max_def_level; + match &info.logical_nulls { + Some(nulls) => { + assert!(range.end <= nulls.len()); + let nulls = nulls.inner(); + info.def_levels.extend_from_iter(range.clone().map(|i| { + // Safety: range.end was asserted to be in bounds earlier + let valid = unsafe { nulls.value_unchecked(i) }; + max_def_level - (!valid as i16) + })); + info.non_null_indices.reserve(len); + info.non_null_indices.extend( + BitIndexIterator::new(nulls.inner(), nulls.offset() + range.start, len) + .map(|i| i + range.start), + ); + } + None => { + info.append_def_level_run(max_def_level, len); + info.non_null_indices.extend(range.clone()); } } - None => info.non_null_indices.extend(range), } - if let Some(rep_levels) = &mut info.rep_levels { - rep_levels.extend(std::iter::repeat_n(info.max_rep_level, len)) + if !matches!(info.rep_levels, LevelData::Absent) { + info.append_rep_level_run(info.max_rep_level, len); } } @@ -754,17 +734,142 @@ impl LevelInfoBuilder { /// The data necessary to write a primitive Arrow array to parquet, taking into account /// any non-primitive parents it may have in the arrow representation +#[derive(Debug, Clone)] +pub(crate) enum LevelData { + Absent, + Materialized(Vec), + Uniform { value: i16, count: usize }, +} + +// Compare logical level contents rather than physical representation, so a +// uniform run compares equal to the equivalent materialized buffer. +impl PartialEq for LevelData { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Absent, Self::Absent) => true, + (Self::Materialized(a), Self::Materialized(b)) => a == b, + (Self::Uniform { value: v, count: n }, Self::Materialized(b)) + | (Self::Materialized(b), Self::Uniform { value: v, count: n }) => { + b.len() == *n && b.iter().all(|x| x == v) + } + ( + Self::Uniform { + value: v1, + count: n1, + }, + Self::Uniform { + value: v2, + count: n2, + }, + ) => v1 == v2 && n1 == n2, + _ => false, + } + } +} + +impl Eq for LevelData {} + +impl LevelData { + fn new(present: bool) -> Self { + match present { + true => Self::Materialized(Vec::new()), + false => Self::Absent, + } + } + + pub(crate) fn as_ref(&self) -> LevelDataRef<'_> { + match self { + Self::Absent => LevelDataRef::Absent, + Self::Materialized(values) => LevelDataRef::Materialized(values), + Self::Uniform { value, count } => LevelDataRef::Uniform { + value: *value, + count: *count, + }, + } + } + + pub(crate) fn slice(&self, offset: usize, len: usize) -> Self { + match self { + Self::Absent => Self::Absent, + Self::Materialized(values) => Self::Materialized(values[offset..offset + len].to_vec()), + Self::Uniform { value, .. } => Self::Uniform { + value: *value, + count: len, + }, + } + } + + fn append_run(&mut self, value: i16, count: usize) { + if count == 0 { + return; + } + + match self { + // No physical level stream exists for this schema. Higher-level + // traversal may still append implicit levels, so this remains a no-op. + Self::Absent => {} + // Start compact: the first appended run can be represented without + // allocating a level buffer. + Self::Materialized(values) if values.is_empty() => { + *self = Self::Uniform { value, count }; + } + // Already materialized, so preserve the buffer representation and append. + Self::Materialized(values) => values.extend(std::iter::repeat_n(value, count)), + // Preserve the compact representation while the appended run has + // the same value. + Self::Uniform { + value: uniform_value, + count: uniform_count, + } if *uniform_value == value => { + *uniform_count += count; + } + // A different value breaks the uniform representation. Materialize + // the existing run, then append the new run to the buffer. + Self::Uniform { .. } => { + let values = self.materialize_mut().unwrap(); + values.extend(std::iter::repeat_n(value, count)); + } + } + } + + fn extend_from_iter(&mut self, iter: I) + where + I: IntoIterator, + { + if let Some(values) = self.materialize_mut() { + values.extend(iter); + } + } + + /// Convert a uniform run into a materialized buffer if needed, then return + /// the mutable level buffer. Returns `None` when no physical level stream exists. + fn materialize_mut(&mut self) -> Option<&mut Vec> { + match self { + Self::Absent => None, + Self::Materialized(values) => Some(values), + Self::Uniform { value, count } => { + let values = vec![*value; *count]; + *self = Self::Materialized(values); + match self { + Self::Materialized(values) => Some(values), + _ => unreachable!(), + } + } + } + } +} + #[derive(Debug, Clone)] pub(crate) struct ArrayLevels { /// Array's definition levels /// /// Present if `max_def_level != 0` - def_levels: Option>, + def_levels: LevelData, /// Array's optional repetition levels /// /// Present if `max_rep_level != 0` - rep_levels: Option>, + rep_levels: LevelData, /// The corresponding array identifying non-null slices of data /// from the primitive array @@ -807,8 +912,8 @@ impl ArrayLevels { let logical_nulls = array.logical_nulls(); Self { - def_levels: (max_def_level != 0).then(Vec::new), - rep_levels: (max_rep_level != 0).then(Vec::new), + def_levels: LevelData::new(max_def_level != 0), + rep_levels: LevelData::new(max_rep_level != 0), non_null_indices: vec![], max_def_level, max_rep_level, @@ -821,12 +926,12 @@ impl ArrayLevels { &self.array } - pub fn def_levels(&self) -> Option<&[i16]> { - self.def_levels.as_deref() + pub(crate) fn def_level_data(&self) -> &LevelData { + &self.def_levels } - pub fn rep_levels(&self) -> Option<&[i16]> { - self.rep_levels.as_deref() + pub(crate) fn rep_level_data(&self) -> &LevelData { + &self.rep_levels } pub fn non_null_indices(&self) -> &[usize] { @@ -839,12 +944,8 @@ impl ArrayLevels { /// `non_null_indices`. The array is sliced to the range covered by /// those indices, and they are shifted to be relative to the slice. pub(crate) fn slice_for_chunk(&self, chunk: &CdcChunk) -> Self { - let def_levels = self.def_levels.as_ref().map(|levels| { - levels[chunk.level_offset..chunk.level_offset + chunk.num_levels].to_vec() - }); - let rep_levels = self.rep_levels.as_ref().map(|levels| { - levels[chunk.level_offset..chunk.level_offset + chunk.num_levels].to_vec() - }); + let def_levels = self.def_levels.slice(chunk.level_offset, chunk.num_levels); + let rep_levels = self.rep_levels.slice(chunk.level_offset, chunk.num_levels); // Select the non-null indices for this chunk. let nni = &self.non_null_indices[chunk.value_offset..chunk.value_offset + chunk.num_values]; @@ -870,6 +971,14 @@ impl ArrayLevels { logical_nulls, } } + + fn append_def_level_run(&mut self, value: i16, count: usize) { + self.def_levels.append_run(value, count); + } + + fn append_rep_level_run(&mut self, value: i16, count: usize) { + self.rep_levels.append_run(value, count); + } } #[cfg(test)] @@ -920,8 +1029,8 @@ mod tests { assert_eq!(levels.len(), 1); let expected = ArrayLevels { - def_levels: Some(vec![2; 10]), - rep_levels: Some(vec![0, 2, 2, 1, 2, 2, 2, 0, 1, 2]), + def_levels: LevelData::Materialized(vec![2; 10]), + rep_levels: LevelData::Materialized(vec![0, 2, 2, 1, 2, 2, 2, 0, 1, 2]), non_null_indices: vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], max_def_level: 2, max_rep_level: 2, @@ -941,8 +1050,8 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: None, - rep_levels: None, + def_levels: LevelData::Absent, + rep_levels: LevelData::Absent, non_null_indices: (0..10).collect(), max_def_level: 0, max_rep_level: 0, @@ -969,8 +1078,8 @@ mod tests { let logical_nulls = array.logical_nulls(); let expected_levels = ArrayLevels { - def_levels: Some(vec![1, 0, 1, 1, 0]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![1, 0, 1, 1, 0]), + rep_levels: LevelData::Absent, non_null_indices: vec![0, 2, 3], max_def_level: 1, max_rep_level: 0, @@ -1004,8 +1113,8 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: Some(vec![1; 5]), - rep_levels: Some(vec![0; 5]), + def_levels: LevelData::Materialized(vec![1; 5]), + rep_levels: LevelData::Materialized(vec![0; 5]), non_null_indices: (0..5).collect(), max_def_level: 1, max_rep_level: 1, @@ -1038,8 +1147,8 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: Some(vec![2, 2, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2]), - rep_levels: Some(vec![0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1]), + def_levels: LevelData::Materialized(vec![2, 2, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1]), non_null_indices: (0..11).collect(), max_def_level: 2, max_rep_level: 1, @@ -1088,8 +1197,8 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: Some(vec![0, 2, 0, 3, 3, 3, 3, 3, 3, 3]), - rep_levels: Some(vec![0, 0, 0, 0, 1, 1, 1, 0, 1, 1]), + def_levels: LevelData::Materialized(vec![0, 2, 0, 3, 3, 3, 3, 3, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 1, 1, 1, 0, 1, 1]), non_null_indices: (4..11).collect(), max_def_level: 3, max_rep_level: 1, @@ -1135,10 +1244,10 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: Some(vec![ + def_levels: LevelData::Materialized(vec![ 5, 5, 5, 5, 1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ]), - rep_levels: Some(vec![ + rep_levels: LevelData::Materialized(vec![ 0, 2, 1, 2, 0, 0, 2, 1, 2, 0, 2, 1, 2, 1, 2, 1, 2, 0, 2, 1, 2, 1, 2, ]), non_null_indices: (0..22).collect(), @@ -1177,8 +1286,8 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: Some(vec![1; 4]), - rep_levels: Some(vec![0; 4]), + def_levels: LevelData::Materialized(vec![1; 4]), + rep_levels: LevelData::Materialized(vec![0; 4]), non_null_indices: (0..4).collect(), max_def_level: 1, max_rep_level: 1, @@ -1210,8 +1319,8 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: Some(vec![1, 3, 3, 3, 3, 3, 3, 3]), - rep_levels: Some(vec![0, 0, 1, 1, 0, 1, 0, 1]), + def_levels: LevelData::Materialized(vec![1, 3, 3, 3, 3, 3, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 1, 1, 0, 1, 0, 1]), non_null_indices: (0..7).collect(), max_def_level: 3, max_rep_level: 1, @@ -1259,8 +1368,12 @@ mod tests { assert_eq!(levels.len(), 1); let expected_levels = ArrayLevels { - def_levels: Some(vec![1, 5, 5, 5, 4, 5, 5, 5, 5, 5, 5, 5, 4, 5, 5, 5, 5, 5]), - rep_levels: Some(vec![0, 0, 1, 2, 1, 0, 2, 2, 1, 2, 2, 2, 0, 1, 2, 2, 2, 2]), + def_levels: LevelData::Materialized(vec![ + 1, 5, 5, 5, 4, 5, 5, 5, 5, 5, 5, 5, 4, 5, 5, 5, 5, 5, + ]), + rep_levels: LevelData::Materialized(vec![ + 0, 0, 1, 2, 1, 0, 2, 2, 1, 2, 2, 2, 0, 1, 2, 2, 2, 2, + ]), non_null_indices: (0..15).collect(), max_def_level: 5, max_rep_level: 2, @@ -1300,8 +1413,8 @@ mod tests { let logical_nulls = leaf.logical_nulls(); let expected_levels = ArrayLevels { - def_levels: Some(vec![3, 2, 3, 1, 0, 3]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![3, 2, 3, 1, 0, 3]), + rep_levels: LevelData::Absent, non_null_indices: vec![0, 2, 5], max_def_level: 3, max_rep_level: 0, @@ -1340,8 +1453,8 @@ mod tests { let list_level = &levels[0]; let expected_level = ArrayLevels { - def_levels: Some(vec![0, 3, 3, 3]), - rep_levels: Some(vec![0, 0, 1, 1]), + def_levels: LevelData::Materialized(vec![0, 3, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 1, 1]), non_null_indices: vec![3, 4, 5], max_def_level: 3, max_rep_level: 1, @@ -1433,8 +1546,8 @@ mod tests { let list_level = &levels[0]; let expected_level = ArrayLevels { - def_levels: None, - rep_levels: None, + def_levels: LevelData::Absent, + rep_levels: LevelData::Absent, non_null_indices: vec![0, 1, 2, 3, 4], max_def_level: 0, max_rep_level: 0, @@ -1448,8 +1561,8 @@ mod tests { let b_logical_nulls = b.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![1, 0, 0, 1, 1]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![1, 0, 0, 1, 1]), + rep_levels: LevelData::Absent, non_null_indices: vec![0, 3, 4], max_def_level: 1, max_rep_level: 0, @@ -1463,8 +1576,8 @@ mod tests { let d_logical_nulls = d.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![1, 1, 1, 2, 1]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![1, 1, 1, 2, 1]), + rep_levels: LevelData::Absent, non_null_indices: vec![3], max_def_level: 2, max_rep_level: 0, @@ -1478,8 +1591,8 @@ mod tests { let f_logical_nulls = f.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![3, 2, 3, 2, 3]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![3, 2, 3, 2, 3]), + rep_levels: LevelData::Absent, non_null_indices: vec![0, 2, 4], max_def_level: 3, max_rep_level: 0, @@ -1586,8 +1699,8 @@ mod tests { let list_level = &levels[0]; let expected_level = ArrayLevels { - def_levels: Some(vec![1; 7]), - rep_levels: Some(vec![0, 1, 0, 1, 0, 1, 1]), + def_levels: LevelData::Materialized(vec![1; 7]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 1, 0, 1, 1]), non_null_indices: vec![0, 1, 2, 3, 4, 5, 6], max_def_level: 1, max_rep_level: 1, @@ -1601,8 +1714,8 @@ mod tests { let map_values_logical_nulls = map.values().logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![2, 2, 2, 1, 2, 1, 2]), - rep_levels: Some(vec![0, 1, 0, 1, 0, 1, 1]), + def_levels: LevelData::Materialized(vec![2, 2, 2, 1, 2, 1, 2]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 1, 0, 1, 1]), non_null_indices: vec![0, 1, 2, 4, 6], max_def_level: 2, max_rep_level: 1, @@ -1688,8 +1801,8 @@ mod tests { let logical_nulls = values.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![4, 1, 0, 2, 2, 3, 4]), - rep_levels: Some(vec![0, 0, 0, 0, 1, 0, 0]), + def_levels: LevelData::Materialized(vec![4, 1, 0, 2, 2, 3, 4]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 1, 0, 0]), non_null_indices: vec![0, 4], max_def_level: 4, max_rep_level: 1, @@ -1730,8 +1843,8 @@ mod tests { let logical_nulls = values.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![4, 4, 3, 2, 0, 4, 4, 0, 1]), - rep_levels: Some(vec![0, 1, 0, 0, 0, 0, 1, 0, 0]), + def_levels: LevelData::Materialized(vec![4, 4, 3, 2, 0, 4, 4, 0, 1]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 0, 0, 1, 0, 0]), non_null_indices: vec![0, 1, 5, 6], max_def_level: 4, max_rep_level: 1, @@ -1817,8 +1930,8 @@ mod tests { let a1_logical_nulls = a1_values.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![0, 0, 1, 6, 5, 2, 3, 1]), - rep_levels: Some(vec![0, 0, 0, 0, 2, 0, 1, 0]), + def_levels: LevelData::Materialized(vec![0, 0, 1, 6, 5, 2, 3, 1]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 2, 0, 1, 0]), non_null_indices: vec![1], max_def_level: 6, max_rep_level: 2, @@ -1830,8 +1943,8 @@ mod tests { let a2_logical_nulls = a2_values.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![0, 0, 1, 3, 2, 4, 1]), - rep_levels: Some(vec![0, 0, 0, 0, 0, 1, 0]), + def_levels: LevelData::Materialized(vec![0, 0, 1, 3, 2, 4, 1]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 0, 1, 0]), non_null_indices: vec![4], max_def_level: 4, max_rep_level: 1, @@ -1870,8 +1983,8 @@ mod tests { let logical_nulls = values.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![0, 0, 3, 3]), - rep_levels: Some(vec![0, 0, 0, 1]), + def_levels: LevelData::Materialized(vec![0, 0, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 1]), non_null_indices: vec![6, 7], max_def_level: 3, max_rep_level: 1, @@ -2022,8 +2135,8 @@ mod tests { // [[{a: 1}, null], null, [null, null], [{a: null}, {a: 2}]] let values_a_logical_nulls = values_a.logical_nulls(); let expected_a = ArrayLevels { - def_levels: Some(vec![4, 2, 0, 2, 2, 3, 4]), - rep_levels: Some(vec![0, 1, 0, 0, 1, 0, 1]), + def_levels: LevelData::Materialized(vec![4, 2, 0, 2, 2, 3, 4]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1]), non_null_indices: vec![0, 7], max_def_level: 4, max_rep_level: 1, @@ -2033,8 +2146,8 @@ mod tests { // [[{b: 2}, null], null, [null, null], [{b: 3}, {b: 4}]] let values_b_logical_nulls = values_b.logical_nulls(); let expected_b = ArrayLevels { - def_levels: Some(vec![3, 2, 0, 2, 2, 3, 3]), - rep_levels: Some(vec![0, 1, 0, 0, 1, 0, 1]), + def_levels: LevelData::Materialized(vec![3, 2, 0, 2, 2, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1]), non_null_indices: vec![0, 6, 7], max_def_level: 3, max_rep_level: 1, @@ -2066,8 +2179,8 @@ mod tests { let logical_nulls = values.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![1, 0, 1]), - rep_levels: Some(vec![0, 0, 0]), + def_levels: LevelData::Materialized(vec![1, 0, 1]), + rep_levels: LevelData::Materialized(vec![0, 0, 0]), non_null_indices: vec![], max_def_level: 3, max_rep_level: 1, @@ -2103,8 +2216,8 @@ mod tests { let logical_nulls = values.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![5, 4, 5, 2, 5, 3, 5, 5, 4, 4, 0]), - rep_levels: Some(vec![0, 2, 2, 1, 0, 1, 0, 2, 1, 2, 0]), + def_levels: LevelData::Materialized(vec![5, 4, 5, 2, 5, 3, 5, 5, 4, 4, 0]), + rep_levels: LevelData::Materialized(vec![0, 2, 2, 1, 0, 1, 0, 2, 1, 2, 0]), non_null_indices: vec![0, 2, 3, 4, 5], max_def_level: 5, max_rep_level: 2, @@ -2136,8 +2249,8 @@ mod tests { let logical_nulls = dict.logical_nulls(); let expected_level = ArrayLevels { - def_levels: Some(vec![0, 0, 1, 1]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![0, 0, 1, 1]), + rep_levels: LevelData::Absent, non_null_indices: vec![2, 3], max_def_level: 1, max_rep_level: 0, @@ -2176,8 +2289,8 @@ mod tests { let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])); let logical_nulls = array.logical_nulls(); let levels = ArrayLevels { - def_levels: None, - rep_levels: None, + def_levels: LevelData::Absent, + rep_levels: LevelData::Absent, non_null_indices: vec![0, 1, 2, 3, 4, 5], max_def_level: 0, max_rep_level: 0, @@ -2190,8 +2303,8 @@ mod tests { value_offset: 2, num_values: 3, }); - assert!(sliced.def_levels.is_none()); - assert!(sliced.rep_levels.is_none()); + assert!(matches!(sliced.def_levels, LevelData::Absent)); + assert!(matches!(sliced.rep_levels, LevelData::Absent)); assert_eq!(sliced.non_null_indices, vec![0, 1, 2]); assert_eq!(sliced.array.len(), 3); @@ -2210,8 +2323,8 @@ mod tests { ])); let logical_nulls = array.logical_nulls(); let levels = ArrayLevels { - def_levels: Some(vec![1, 0, 1, 0, 1, 1]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![1, 0, 1, 0, 1, 1]), + rep_levels: LevelData::Absent, non_null_indices: vec![0, 2, 4, 5], max_def_level: 1, max_rep_level: 0, @@ -2224,8 +2337,8 @@ mod tests { value_offset: 1, num_values: 1, }); - assert_eq!(sliced.def_levels, Some(vec![0, 1, 0])); - assert!(sliced.rep_levels.is_none()); + assert_eq!(sliced.def_levels, LevelData::Materialized(vec![0, 1, 0])); + assert!(matches!(sliced.rep_levels, LevelData::Absent)); assert_eq!(sliced.non_null_indices, vec![0]); // [2] shifted by -2 (nni[0]) assert_eq!(sliced.array.len(), 1); } @@ -2264,8 +2377,8 @@ mod tests { ])); let logical_nulls = array.logical_nulls(); let levels = ArrayLevels { - def_levels: Some(vec![3, 0, 3, 2, 0, 3, 3]), - rep_levels: Some(vec![0, 0, 0, 1, 0, 0, 1]), + def_levels: LevelData::Materialized(vec![3, 0, 3, 2, 0, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 1, 0, 0, 1]), non_null_indices: vec![0, 3, 8, 9], max_def_level: 3, max_rep_level: 1, @@ -2310,8 +2423,8 @@ mod tests { let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, None, Some(4)])); let logical_nulls = array.logical_nulls(); let levels = ArrayLevels { - def_levels: Some(vec![1, 0, 0, 1]), - rep_levels: None, + def_levels: LevelData::Materialized(vec![1, 0, 0, 1]), + rep_levels: LevelData::Absent, non_null_indices: vec![0, 3], max_def_level: 1, max_rep_level: 0, @@ -2325,7 +2438,7 @@ mod tests { value_offset: 1, num_values: 0, }); - assert_eq!(sliced.def_levels, Some(vec![0, 0])); + assert_eq!(sliced.def_levels, LevelData::Materialized(vec![0, 0])); assert_eq!(sliced.non_null_indices, Vec::::new()); assert_eq!(sliced.array.len(), 0); } diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 641b81257a91..b93e17440120 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -920,8 +920,11 @@ impl ArrowColumnWriter { chunker: &mut ContentDefinedChunker, ) -> Result<()> { let levels = &col.0; - let chunks = - chunker.get_arrow_chunks(levels.def_levels(), levels.rep_levels(), levels.array())?; + let chunks = chunker.get_arrow_chunks( + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), + levels.array(), + )?; let num_chunks = chunks.len(); for (i, chunk) in chunks.iter().enumerate() { @@ -1384,10 +1387,15 @@ fn write_leaf( } ColumnWriter::BoolColumnWriter(typed) => { let array = column.as_boolean(); - typed.write_batch( - get_bool_array_slice(array, indices).as_slice(), - levels.def_levels(), - levels.rep_levels(), + let values = get_bool_array_slice(array, indices); + typed.write_batch_internal( + values.as_slice(), + None, + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), + None, + None, + None, ) } ColumnWriter::Int64ColumnWriter(typed) => { @@ -1538,7 +1546,15 @@ fn write_leaf( )); } }; - typed.write_batch(bytes.as_slice(), levels.def_levels(), levels.rep_levels()) + typed.write_batch_internal( + bytes.as_slice(), + None, + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), + None, + None, + None, + ) } } } @@ -1551,8 +1567,8 @@ fn write_primitive( writer.write_batch_internal( values, Some(levels.non_null_indices()), - levels.def_levels(), - levels.rep_levels(), + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), None, None, None, diff --git a/parquet/src/column/chunker/cdc.rs b/parquet/src/column/chunker/cdc.rs index cb4328536df3..b40dd74a8d83 100644 --- a/parquet/src/column/chunker/cdc.rs +++ b/parquet/src/column/chunker/cdc.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#[cfg(feature = "arrow")] +use crate::column::writer::LevelDataRef; use crate::errors::{ParquetError, Result}; use crate::file::properties::CdcOptions; use crate::schema::types::ColumnDescriptor; @@ -275,8 +277,8 @@ impl ContentDefinedChunker { /// evaluate if we need to create a new chunk. fn calculate( &mut self, - def_levels: Option<&[i16]>, - rep_levels: Option<&[i16]>, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, num_levels: usize, mut roll_value: F, ) -> Vec @@ -322,10 +324,11 @@ impl ContentDefinedChunker { // def_levels: [1, 0, 1, 0, 1] // level: 0 1 2 3 4 // value_offset: 0 1 2 (only increments on def==1) - let def_levels = def_levels.expect("def_levels required when max_def_level > 0"); #[allow(clippy::needless_range_loop)] for offset in 0..num_levels { - let def_level = def_levels[offset]; + let def_level = def_levels + .value_at(offset) + .expect("def_levels required when max_def_level > 0"); self.roll_level(def_level); if def_level == self.max_def_level { // For non-nested data, the leaf array has one slot per @@ -380,13 +383,15 @@ impl ContentDefinedChunker { // Using value_offset=1 would index position 1 (the null slot). // // Using value_offset for roll_value would hash the wrong array slot. - let def_levels = def_levels.expect("def_levels required for nested data"); - let rep_levels = rep_levels.expect("rep_levels required for nested data"); let mut leaf_offset: usize = 0; for offset in 0..num_levels { - let def_level = def_levels[offset]; - let rep_level = rep_levels[offset]; + let def_level = def_levels + .value_at(offset) + .expect("def_levels required for nested data"); + let rep_level = rep_levels + .value_at(offset) + .expect("rep_levels required for nested data"); self.roll_level(def_level); self.roll_level(rep_level); @@ -439,16 +444,20 @@ impl ContentDefinedChunker { #[cfg(feature = "arrow")] pub(crate) fn get_arrow_chunks( &mut self, - def_levels: Option<&[i16]>, - rep_levels: Option<&[i16]>, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, array: &dyn arrow_array::Array, ) -> Result> { use arrow_array::cast::AsArray; use arrow_schema::DataType; - let num_levels = match def_levels { - Some(def_levels) => def_levels.len(), - None => array.len(), + // For nested (list) data, null list entries can own non-zero child + // ranges in the leaf array, so `array.len()` may exceed the number of + // levels. Always drive the loop by the level count; fall back to the + // array length only when there are no levels at all. + let num_levels = match (def_levels.len(), rep_levels.len()) { + (0, 0) => array.len(), + (d, r) => d.max(r), }; macro_rules! fixed_width { @@ -566,6 +575,8 @@ impl ContentDefinedChunker { mod tests { use super::*; use crate::basic::Type as PhysicalType; + #[cfg(feature = "arrow")] + use crate::column::writer::LevelDataRef; use crate::schema::types::{ColumnPath, Type}; use std::sync::Arc; @@ -615,9 +626,14 @@ mod tests { // Write a small amount of data — should produce exactly 1 chunk. let num_values = 4; - let chunks = chunker.calculate(None, None, num_values, |c, i| { - c.roll_fixed::<4>(&(i as i32).to_le_bytes()); - }); + let chunks = chunker.calculate( + LevelDataRef::Absent, + LevelDataRef::Absent, + num_values, + |c, i| { + c.roll_fixed::<4>(&(i as i32).to_le_bytes()); + }, + ); assert_eq!(chunks.len(), 1); assert_eq!(chunks[0].level_offset, 0); assert_eq!(chunks[0].value_offset, 0); @@ -636,9 +652,14 @@ mod tests { // Write enough data to exceed max_chunk_size multiple times. // Each i32 = 4 bytes, max_chunk_size=1024, so ~256 values per chunk max. let num_values = 2000; - let chunks = chunker.calculate(None, None, num_values, |c, i| { - c.roll_fixed::<4>(&(i as i32).to_le_bytes()); - }); + let chunks = chunker.calculate( + LevelDataRef::Absent, + LevelDataRef::Absent, + num_values, + |c, i| { + c.roll_fixed::<4>(&(i as i32).to_le_bytes()); + }, + ); // Should have multiple chunks assert!(chunks.len() > 1); @@ -668,10 +689,10 @@ mod tests { }; let mut chunker1 = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap(); - let chunks1 = chunker1.calculate(None, None, 200, roll); + let chunks1 = chunker1.calculate(LevelDataRef::Absent, LevelDataRef::Absent, 200, roll); let mut chunker2 = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap(); - let chunks2 = chunker2.calculate(None, None, 200, roll); + let chunks2 = chunker2.calculate(LevelDataRef::Absent, LevelDataRef::Absent, 200, roll); assert_eq!(chunks1.len(), chunks2.len()); for (a, b) in chunks1.iter().zip(chunks2.iter()) { @@ -699,9 +720,14 @@ mod tests { .collect(); let expected_non_null: usize = def_levels.iter().filter(|&&d| d == 1).count(); - let chunks = chunker.calculate(Some(&def_levels), None, num_levels, |c, i| { - c.roll_fixed::<4>(&(i as i32).to_le_bytes()); - }); + let chunks = chunker.calculate( + LevelDataRef::Materialized(&def_levels), + LevelDataRef::Absent, + num_levels, + |c, i| { + c.roll_fixed::<4>(&(i as i32).to_le_bytes()); + }, + ); assert!(!chunks.is_empty()); let total_levels: usize = chunks.iter().map(|c| c.num_levels).sum(); @@ -729,6 +755,7 @@ mod arrow_tests { use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use crate::arrow::arrow_writer::ArrowWriter; + use crate::column::writer::LevelDataRef; use crate::file::properties::{CdcOptions, WriterProperties}; use crate::file::reader::{FileReader, SerializedFileReader}; @@ -1109,6 +1136,27 @@ mod arrow_tests { buf } + #[test] + fn cdc_all_null_arrow_column_writes_data_pages() { + let array = Arc::new(Int32Array::from(vec![None::; 4096])) as ArrayRef; + let schema = Arc::new(Schema::new(vec![Field::new("f0", DataType::Int32, true)])); + let batch = RecordBatch::try_new(schema, vec![array.clone()]).unwrap(); + + let data = write_with_cdc_options(&[&batch], 64, 256, Some(4096), false); + let info = get_column_info(&data, 0); + + assert_eq!(info.len(), 1); + assert!( + !info[0].page_lengths.is_empty(), + "all-null CDC write should still emit data pages" + ); + assert_eq!( + info[0].page_lengths.iter().sum::(), + array.len() as i64, + "all-null CDC pages should account for every input row" + ); + } + fn read_batches(data: &[u8]) -> Vec { let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(data.to_vec())) .unwrap() @@ -2190,11 +2238,15 @@ mod arrow_tests { let array: Int32Array = (0..n).map(|i| test_hash(0, i as u64) as i32).collect(); let mut chunker = super::ContentDefinedChunker::new(&desc, &options).unwrap(); - let chunks = chunker.get_arrow_chunks(None, None, &array).unwrap(); + let chunks = chunker + .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &array) + .unwrap(); let sliced = array.slice(offset, n - offset); let mut chunker2 = super::ContentDefinedChunker::new(&desc, &options).unwrap(); - let chunks2 = chunker2.get_arrow_chunks(None, None, &sliced).unwrap(); + let chunks2 = chunker2 + .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &sliced) + .unwrap(); let values: Vec = chunks.iter().map(|c| c.num_values).collect(); let values2: Vec = chunks2.iter().map(|c| c.num_values).collect(); diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index f755beed55c4..5d14ac6856f9 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -314,6 +314,67 @@ impl ColumnMetrics { } } +/// Borrowed view of level data, analogous to `&str` for `LevelData`'s `String`. +/// +/// `LevelDataRef` can be constructed from `LevelData` and directly from an existing +/// `&[i16]` without allocating. +/// +/// The variants are different physical representations of the same logical +/// sequence of levels. +#[derive(Debug, Clone, Copy)] +pub(crate) enum LevelDataRef<'a> { + Absent, + Materialized(&'a [i16]), + Uniform { value: i16, count: usize }, +} + +impl<'a> From<&'a [i16]> for LevelDataRef<'a> { + fn from(levels: &'a [i16]) -> Self { + Self::Materialized(levels) + } +} + +impl<'a> From> for LevelDataRef<'a> { + fn from(levels: Option<&'a [i16]>) -> Self { + levels.map_or(Self::Absent, Self::from) + } +} + +impl<'a> LevelDataRef<'a> { + pub(crate) fn len(self) -> usize { + match self { + Self::Absent => 0, + Self::Materialized(values) => values.len(), + Self::Uniform { count, .. } => count, + } + } + + pub(crate) fn first(self) -> Option { + match self { + Self::Absent => None, + Self::Materialized(values) => values.first().copied(), + Self::Uniform { value, count } => (count > 0).then_some(value), + } + } + + #[cfg(feature = "arrow")] + pub(crate) fn value_at(self, idx: usize) -> Option { + match self { + Self::Absent => None, + Self::Materialized(values) => values.get(idx).copied(), + Self::Uniform { value, count } => (idx < count).then_some(value), + } + } + + pub(crate) fn slice(self, offset: usize, len: usize) -> Self { + match self { + Self::Absent => Self::Absent, + Self::Materialized(values) => Self::Materialized(&values[offset..offset + len]), + Self::Uniform { value, .. } => Self::Uniform { value, count: len }, + } + } +} + /// Typed column writer for a primitive column. pub type ColumnWriterImpl<'a, T> = GenericColumnWriter<'a, ColumnValueEncoderImpl>; @@ -424,21 +485,19 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { &mut self, values: &E::Values, value_indices: Option<&[usize]>, - def_levels: Option<&[i16]>, - rep_levels: Option<&[i16]>, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, min: Option<&E::T>, max: Option<&E::T>, distinct_count: Option, ) -> Result { // Check if number of definition levels is the same as number of repetition levels. - if let (Some(def), Some(rep)) = (def_levels, rep_levels) { - if def.len() != rep.len() { - return Err(general_err!( - "Inconsistent length of definition and repetition levels: {} != {}", - def.len(), - rep.len() - )); - } + if def_levels.len() != 0 && rep_levels.len() != 0 && def_levels.len() != rep_levels.len() { + return Err(general_err!( + "Inconsistent length of definition and repetition levels: {} != {}", + def_levels.len(), + rep_levels.len() + )); } // We check for DataPage limits only after we have inserted the values. If a user @@ -450,10 +509,11 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { // TODO: find out why we don't account for size of levels when we estimate page // size. - - let num_levels = match def_levels { - Some(def_levels) => def_levels.len(), - None => values.len(), + let num_levels = def_levels.len().max(rep_levels.len()); + let num_levels = if num_levels > 0 { + num_levels + } else { + value_indices.map_or(values.len(), |i| i.len()) }; if let Some(min) = min { @@ -472,13 +532,24 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { let mut values_offset = 0; let mut levels_offset = 0; - let base_batch_size = self.props.write_batch_size(); + let both_levels_compact = !matches!(def_levels, LevelDataRef::Materialized(_)) + && !matches!(rep_levels, LevelDataRef::Materialized(_)); + let has_levels = !matches!(def_levels, LevelDataRef::Absent) + || !matches!(rep_levels, LevelDataRef::Absent); + // When both level vectors are compact (Uniform or Absent), there is no + // materialized slice to split and the per-mini-batch work is O(1), so we + // can safely use a much larger batch size. + let base_batch_size = if both_levels_compact && has_levels { + self.props.data_page_row_count_limit() + } else { + self.props.write_batch_size() + }; while levels_offset < num_levels { let mut end_offset = num_levels.min(levels_offset + base_batch_size); // Split at record boundary - if let Some(r) = rep_levels { - while end_offset < r.len() && r[end_offset] != 0 { + if let LevelDataRef::Materialized(levels) = rep_levels { + while end_offset < levels.len() && levels[end_offset] != 0 { end_offset += 1; } } @@ -488,8 +559,8 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { values_offset, value_indices, end_offset - levels_offset, - def_levels.map(|lv| &lv[levels_offset..end_offset]), - rep_levels.map(|lv| &lv[levels_offset..end_offset]), + def_levels.slice(levels_offset, end_offset - levels_offset), + rep_levels.slice(levels_offset, end_offset - levels_offset), )?; levels_offset = end_offset; } @@ -516,7 +587,15 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { def_levels: Option<&[i16]>, rep_levels: Option<&[i16]>, ) -> Result { - self.write_batch_internal(values, None, def_levels, rep_levels, None, None, None) + self.write_batch_internal( + values, + None, + LevelDataRef::from(def_levels), + LevelDataRef::from(rep_levels), + None, + None, + None, + ) } /// Writer may optionally provide pre-calculated statistics for use when computing @@ -538,8 +617,8 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { self.write_batch_internal( values, None, - def_levels, - rep_levels, + LevelDataRef::from(def_levels), + LevelDataRef::from(rep_levels), min, max, distinct_count, @@ -650,32 +729,53 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { values_offset: usize, value_indices: Option<&[usize]>, num_levels: usize, - def_levels: Option<&[i16]>, - rep_levels: Option<&[i16]>, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, ) -> Result { // Process definition levels and determine how many values to write. let values_to_write = if self.descr.max_def_level() > 0 { - let levels = def_levels.ok_or_else(|| { - general_err!( - "Definition levels are required, because max definition level = {}", - self.descr.max_def_level() - ) - })?; - - let mut values_to_write = 0usize; let max_def = self.descr.max_def_level(); - let encoder = &mut self.def_levels_encoder; - match self.page_metrics.definition_level_histogram.as_mut() { - Some(histogram) => encoder.put_with_observer(levels, |level, count| { - values_to_write += count * (level == max_def) as usize; - histogram.increment_by(level, count as i64); - }), - None => encoder.put_with_observer(levels, |level, count| { - values_to_write += count * (level == max_def) as usize; - }), - }; - self.page_metrics.num_page_nulls += (levels.len() - values_to_write) as u64; - values_to_write + match def_levels { + LevelDataRef::Absent => { + return Err(general_err!( + "Definition levels are required, because max definition level = {}", + self.descr.max_def_level() + )); + } + LevelDataRef::Materialized(levels) => { + // General path for caller-provided or already-materialized + // level buffers. + let mut values_to_write = 0usize; + let encoder = &mut self.def_levels_encoder; + match self.page_metrics.definition_level_histogram.as_mut() { + Some(histogram) => encoder.put_with_observer(levels, |level, count| { + values_to_write += count * (level == max_def) as usize; + histogram.increment_by(level, count as i64); + }), + None => encoder.put_with_observer(levels, |level, count| { + values_to_write += count * (level == max_def) as usize; + }), + }; + self.page_metrics.num_page_nulls += (levels.len() - values_to_write) as u64; + values_to_write + } + LevelDataRef::Uniform { value, count } => { + // Fast path for all-null, all-valid, or otherwise uniform + // definition levels without materializing a level buffer. + let encoder = &mut self.def_levels_encoder; + match self.page_metrics.definition_level_histogram.as_mut() { + Some(histogram) => { + encoder.put_n_with_observer(value, count, |level, run_len| { + histogram.increment_by(level, run_len as i64); + }) + } + None => encoder.put_n_with_observer(value, count, |_, _| {}), + }; + let values_to_write = count * (value == max_def) as usize; + self.page_metrics.num_page_nulls += (count - values_to_write) as u64; + values_to_write + } + } } else { num_levels }; @@ -683,31 +783,50 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { // Process repetition levels and determine how many rows we are about to process. if self.descr.max_rep_level() > 0 { // A row could contain more than one value. - let levels = rep_levels.ok_or_else(|| { + let first_level = rep_levels.first().ok_or_else(|| { general_err!( "Repetition levels are required, because max repetition level = {}", self.descr.max_rep_level() ) })?; - if !levels.is_empty() && levels[0] != 0 { + if first_level != 0 { return Err(general_err!( "Write must start at a record boundary, got non-zero repetition level of {}", - levels[0] + first_level )); } let mut new_rows = 0u32; - let encoder = &mut self.rep_levels_encoder; - match self.page_metrics.repetition_level_histogram.as_mut() { - Some(histogram) => encoder.put_with_observer(levels, |level, count| { - new_rows += (count as u32) * (level == 0) as u32; - histogram.increment_by(level, count as i64); - }), - None => encoder.put_with_observer(levels, |level, count| { - new_rows += (count as u32) * (level == 0) as u32; - }), - }; + match rep_levels { + LevelDataRef::Absent => unreachable!(), + LevelDataRef::Materialized(levels) => { + let encoder = &mut self.rep_levels_encoder; + match self.page_metrics.repetition_level_histogram.as_mut() { + Some(histogram) => encoder.put_with_observer(levels, |level, count| { + new_rows += (count as u32) * (level == 0) as u32; + histogram.increment_by(level, count as i64); + }), + None => encoder.put_with_observer(levels, |level, count| { + new_rows += (count as u32) * (level == 0) as u32; + }), + }; + } + LevelDataRef::Uniform { value, count } => { + let encoder = &mut self.rep_levels_encoder; + match self.page_metrics.repetition_level_histogram.as_mut() { + Some(histogram) => { + encoder.put_n_with_observer(value, count, |level, run_len| { + new_rows += (run_len as u32) * (level == 0) as u32; + histogram.increment_by(level, run_len as i64); + }) + } + None => encoder.put_n_with_observer(value, count, |level, run_len| { + new_rows += (run_len as u32) * (level == 0) as u32; + }), + }; + } + } self.page_metrics.num_buffered_rows += new_rows; } else { // Each value is exactly one row. @@ -4451,4 +4570,221 @@ mod tests { result.metadata.compressed_size() ); } + + struct ColumnRoundTripUniform<'a, T: DataType> { + props: WriterProperties, + values: &'a [T::T], + def_levels: LevelDataRef<'a>, + rep_levels: LevelDataRef<'a>, + max_def_level: i16, + max_rep_level: i16, + expected_values: &'a [T::T], + expected_def_levels: Option<&'a [i16]>, + expected_rep_levels: Option<&'a [i16]>, + } + + impl<'a, T: DataType> ColumnRoundTripUniform<'a, T> + where + T::T: PartialEq + std::fmt::Debug, + { + fn new() -> Self { + Self { + props: Default::default(), + values: &[], + def_levels: LevelDataRef::Absent, + rep_levels: LevelDataRef::Absent, + max_def_level: 0, + max_rep_level: 0, + expected_values: &[], + expected_def_levels: None, + expected_rep_levels: None, + } + } + + fn with_props(mut self, props: WriterProperties) -> Self { + self.props = props; + self + } + + fn with_values(mut self, values: &'a [T::T]) -> Self { + self.values = values; + self + } + + fn with_def_levels(mut self, def_levels: LevelDataRef<'a>) -> Self { + self.def_levels = def_levels; + self + } + + fn with_rep_levels(mut self, rep_levels: LevelDataRef<'a>) -> Self { + self.rep_levels = rep_levels; + self + } + + fn with_max_def_level(mut self, max_def_level: i16) -> Self { + self.max_def_level = max_def_level; + self + } + + fn with_max_rep_level(mut self, max_rep_level: i16) -> Self { + self.max_rep_level = max_rep_level; + self + } + + fn with_expected_values(mut self, expected_values: &'a [T::T]) -> Self { + self.expected_values = expected_values; + self + } + + fn with_expected_def_levels(mut self, expected_def_levels: &'a [i16]) -> Self { + self.expected_def_levels = Some(expected_def_levels); + self + } + + fn with_expected_rep_levels(mut self, expected_rep_levels: &'a [i16]) -> Self { + self.expected_rep_levels = Some(expected_rep_levels); + self + } + + /// Write-then-read roundtrip using `write_batch_internal` with the given + /// [`LevelDataRef`] variants, and assert the read-back matches `expected_*`. + fn run(self) { + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let mut writer = get_test_column_writer::( + page_writer, + self.max_def_level, + self.max_rep_level, + Arc::new(self.props), + ); + + writer + .write_batch_internal( + self.values, + None, + self.def_levels, + self.rep_levels, + None, + None, + None, + ) + .unwrap(); + let result = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &result.metadata, + result.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + let mut reader = + get_test_column_reader::(page_reader, self.max_def_level, self.max_rep_level); + + let batch_size = self + .expected_def_levels + .map_or(self.expected_values.len(), |l| l.len()); + let mut actual_values = Vec::with_capacity(batch_size); + let mut actual_def = self + .expected_def_levels + .map(|_| Vec::with_capacity(batch_size)); + let mut actual_rep = self + .expected_rep_levels + .map(|_| Vec::with_capacity(batch_size)); + + let (_, values_read, levels_read) = reader + .read_records( + batch_size, + actual_def.as_mut(), + actual_rep.as_mut(), + &mut actual_values, + ) + .unwrap(); + + assert_eq!(&actual_values[..values_read], self.expected_values); + if let Some(ref v) = actual_def { + assert_eq!(&v[..levels_read], self.expected_def_levels.unwrap()); + } + if let Some(ref v) = actual_rep { + assert_eq!(&v[..levels_read], self.expected_rep_levels.unwrap()); + } + } + } + + #[test] + fn test_uniform_def_levels_all_null() { + // All-null column: def_level=0 (null) for every slot, no values written. + let max_def_level = 1; + let count = 100; + let expected_def_levels = vec![0i16; count]; + ColumnRoundTripUniform::::new() + .with_def_levels(LevelDataRef::Uniform { value: 0, count }) + .with_max_def_level(max_def_level) + .with_expected_def_levels(&expected_def_levels) + .run(); + } + + #[test] + fn test_uniform_def_levels_all_valid() { + // All-valid column: def_level=max for every slot, all values written. + let max_def_level = 1; + let values: Vec = (0..50).collect(); + let expected_def_levels = vec![max_def_level; values.len()]; + ColumnRoundTripUniform::::new() + .with_values(&values) + .with_def_levels(LevelDataRef::Uniform { + value: max_def_level, + count: values.len(), + }) + .with_max_def_level(max_def_level) + .with_expected_values(&values) + .with_expected_def_levels(&expected_def_levels) + .run(); + } + + #[test] + fn test_uniform_def_and_rep_levels() { + // Simulates a list column where every row is null: + // def=0, rep=0 for each row (one row = one entry with no child values). + let max_def_level = 2; + let max_rep_level = 1; + let count = 200; + let expected_def_levels = vec![0i16; count]; + let expected_rep_levels = vec![0i16; count]; + ColumnRoundTripUniform::::new() + .with_def_levels(LevelDataRef::Uniform { value: 0, count }) + .with_rep_levels(LevelDataRef::Uniform { value: 0, count }) + .with_max_def_level(max_def_level) + .with_max_rep_level(max_rep_level) + .with_expected_def_levels(&expected_def_levels) + .with_expected_rep_levels(&expected_rep_levels) + .run(); + } + + #[test] + fn test_uniform_levels_v1_and_v2() { + // Verify uniform levels work identically for both Parquet writer versions. + for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + let props = WriterProperties::builder() + .set_writer_version(version) + .build(); + let max_def = 1; + let count = 100; + let expected_def_levels = vec![0i16; count]; + ColumnRoundTripUniform::::new() + .with_props(props) + .with_def_levels(LevelDataRef::Uniform { value: 0, count }) + .with_max_def_level(max_def) + .with_expected_def_levels(&expected_def_levels) + .run(); + } + } } diff --git a/parquet/src/encodings/levels.rs b/parquet/src/encodings/levels.rs index 6b6de8b0e1c6..841afd8b3186 100644 --- a/parquet/src/encodings/levels.rs +++ b/parquet/src/encodings/levels.rs @@ -92,6 +92,37 @@ impl LevelEncoder { } } + /// Encode `count` repetitions of a single level value, calling + /// `observer(value, count)` exactly once. + /// + /// This can be used to encode uniform runs without allocating a level + /// buffer. Like [`Self::put_with_observer`], but specialized for a single + /// repeated level value. + /// + /// This is O(1) amortized for RLE-based encoders (after a small warmup). + #[inline] + pub fn put_n_with_observer(&mut self, value: i16, count: usize, mut observer: F) + where + F: FnMut(i16, usize), + { + let encoder = match *self { + LevelEncoder::Rle(ref mut encoder) | LevelEncoder::RleV2(ref mut encoder) => encoder, + }; + + // Feed values individually until the encoder enters RLE accumulation + // mode for this value, or until we've encoded everything. + let mut remaining = count; + while remaining > 0 && !encoder.is_accumulating_rle(value as u64) { + encoder.put(value as u64); + remaining -= 1; + } + // If we're now in accumulation mode, bulk-extend the rest. + if remaining > 0 { + encoder.extend_run(remaining); + } + observer(value, count); + } + /// Finalizes level encoder, flush all intermediate buffers and return resulting /// encoded buffer. Returned buffer is already truncated to encoded bytes only. #[inline] @@ -139,3 +170,116 @@ impl LevelEncoder { result } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Encode `count` repetitions of `value` using `put_with_observer` and + /// return the raw encoded bytes. + fn reference_encode(max_level: i16, values: &[i16]) -> Vec { + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_with_observer(values, |_, _| {}); + enc.consume() + } + + #[test] + fn test_put_n_with_observer_large_run() { + // Large count exercises the bulk extend_run path (past the 8-value warmup). + let max_level = 3; + let count = 10_000; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(2, count, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &vec![2; count])); + } + + #[test] + fn test_put_n_with_observer_small_count() { + // Count smaller than the RLE warmup threshold — only the per-element loop runs. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(1, 5, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[1; 5])); + } + + #[test] + fn test_put_n_with_observer_exact_threshold() { + // Exactly 8 values: the warmup loop completes and extend_run gets 0. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(3, 8, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[3; 8])); + } + + #[test] + fn test_put_n_with_observer_single_value() { + let max_level = 1; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(1, 1, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[1])); + } + + #[test] + fn test_put_n_with_observer_zero_count() { + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(2, 0, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[])); + } + + #[test] + fn test_put_n_with_observer_calls_observer_exactly_once() { + let mut enc = LevelEncoder::v2_streaming(3); + let mut calls: Vec<(i16, usize)> = Vec::new(); + enc.put_n_with_observer(2, 500, |val, cnt| calls.push((val, cnt))); + assert_eq!(calls, vec![(2, 500)]); + } + + #[test] + fn test_put_n_with_observer_zero_count_calls_observer() { + let mut enc = LevelEncoder::v2_streaming(3); + let mut calls: Vec<(i16, usize)> = Vec::new(); + enc.put_n_with_observer(1, 0, |val, cnt| calls.push((val, cnt))); + assert_eq!(calls, vec![(1, 0)]); + } + + #[test] + fn test_put_n_with_observer_followed_by_different_value() { + // Two consecutive put_n calls with different values — verifies that + // the encoder correctly transitions between runs. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + let mut calls: Vec<(i16, usize)> = Vec::new(); + enc.put_n_with_observer(1, 100, |v, c| calls.push((v, c))); + enc.put_n_with_observer(3, 200, |v, c| calls.push((v, c))); + assert_eq!(calls, vec![(1, 100), (3, 200)]); + + let reference = reference_encode(max_level, &[&[1i16; 100][..], &[3i16; 200]].concat()); + assert_eq!(enc.consume(), reference); + } + + #[test] + fn test_put_n_with_observer_interleaved_with_put_with_observer() { + // Mix put_n_with_observer and put_with_observer to verify they compose. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(2, 50, |_, _| {}); + enc.put_with_observer(&[0, 0, 1, 1, 3], |_, _| {}); + enc.put_n_with_observer(2, 50, |_, _| {}); + + let input = [&[2i16; 50][..], &[0, 0, 1, 1, 3], &[2i16; 50]].concat(); + assert_eq!(enc.consume(), reference_encode(max_level, &input)); + } + + #[test] + fn test_put_n_with_observer_v1_roundtrip() { + // Also verify V1 (Rle variant with length header) works correctly. + let max_level = 3; + let mut enc = LevelEncoder::v1_streaming(max_level); + enc.put_n_with_observer(2, 1000, |_, _| {}); + + let mut ref_enc = LevelEncoder::v1_streaming(max_level); + ref_enc.put_with_observer(&[2; 1000], |_, _| {}); + assert_eq!(enc.consume(), ref_enc.consume()); + } +}