etseidl commented on code in PR #11911: URL: https://github.com/apache/arrow-rs/pull/11911#discussion_r4416365542
########## parquet/src/arrow/arrow_reader/statistics/page_index.rs: ########## @@ -0,0 +1,1135 @@ +// 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. + +//! Reads the page statistics stored in a Parquet `ColumnIndex` and puts them +//! straight into Arrow arrays. +//! +//! The usual route first builds a [`ColumnIndexMetaData`] and then copies its +//! values into Arrow arrays one by one. This module skips the middle step: it +//! reads the stored bytes once and writes each value directly into the memory +//! that will become the final Arrow array. +//! +//! Only `INT32` columns are supported. +//! +//! [`ColumnIndexMetaData`]: crate::file::page_index::column_index::ColumnIndexMetaData + +use super::DataPageStatistics; +use crate::basic::BoundaryOrder; +use crate::errors::{ParquetError, Result}; +use crate::parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftSliceInputProtocol, + validate_list_type, +}; +use arrow_array::types::{ + Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Int8Type, + Int16Type, Time32MillisecondType, Time32SecondType, UInt8Type, UInt16Type, UInt32Type, +}; +use arrow_array::{ArrayRef, Int32Array, UInt64Array, new_null_array}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer, i256}; +use arrow_schema::{DataType, TimeUnit}; +use std::sync::Arc; + +/// Adds `n` default values to `v`. +fn pad<T: Default + Clone>(v: &mut Vec<T>, n: usize) { + v.resize(v.len() + n, T::default()); +} + +/// Reads one `INT32` value per page from the list of binary values at the +/// start of `buf`, and returns how many bytes the list took. +/// +/// `has_min_max` holds one entry for each page in the list. Pages marked +/// `false` get a filler value, and their stored bytes are not looked at. +#[inline(never)] +fn read_list(mut buf: &[u8], out: &mut Vec<i32>, has_min_max: &[bool]) -> Result<usize> { + // `buf` is a local copy of the unread bytes, so the compiler can keep + // the read position in a register rather than in memory. This makes + // this loop much faster. + let start = buf.len(); + // Fill in filler values first and then overwrite them, which saves + // checking the capacity for every value. + let first = out.len(); + out.resize(first + has_min_max.len(), 0); + for (slot, &has_value) in out[first..].iter_mut().zip(has_min_max) { + // Writers store each value as a one byte length of 4 followed by + // the 4 bytes, so handle that directly. Anything else, such as the + // empty value of a null page, takes the general route below. + if has_value + && let Some((&4, rest)) = buf.split_first() + && let Some((value, rest)) = rest.split_first_chunk::<4>() + { + *slot = i32::from_le_bytes(*value); + buf = rest; + continue; + } + let (bytes, rest) = split_binary(buf)?; + buf = rest; + if has_value { + *slot = i32::from_le_bytes(first_bytes(bytes)?); + } + } + Ok(start - buf.len()) +} + +/// Splits one Thrift binary value off the front of `buf`, returning the +/// value and the bytes after it. A binary value is its length followed by +/// its bytes. +#[inline(always)] +fn split_binary(buf: &[u8]) -> Result<(&[u8], &[u8])> { + let (len, rest) = split_varint(buf)?; + let len = len as usize; + if rest.len() < len { + return Err(eof_err!("Unexpected EOF")); + } + Ok(rest.split_at(len)) +} + +/// Splits one stored whole number off the front of `buf`, returning the +/// number and the bytes after it. +/// +/// Numbers below 128 take a single byte, which covers nearly every length +/// and count in a page index, so that case is handled right here. Longer +/// numbers use the normal reader. +#[inline(always)] +fn split_varint(buf: &[u8]) -> Result<(u64, &[u8])> { + match buf.split_first() { + Some((&n, rest)) if n < 0x80 => Ok((n as u64, rest)), + Some(_) => { + let mut prot = ThriftSliceInputProtocol::new(buf); + Ok((prot.read_vlq()?, prot.as_slice())) + } + None => Err(eof_err!("Unexpected EOF")), + } +} + +/// Takes the first `N` bytes of `bytes`. +/// +/// This behaves like the older decoder: extra bytes are ignored and too few +/// bytes is an error, with the same message. +fn first_bytes<const N: usize>(bytes: &[u8]) -> Result<[u8; N]> { + match bytes.get(..N) { + Some(b) => Ok(b.try_into().unwrap()), + None => Err(general_err!( + "error converting value, expected {} bytes got {}", + N, + bytes.len() + )), + } +} + +/// The null or NaN counts of every page. +#[derive(Debug)] +struct Counts { + values: Vec<u64>, + /// One entry per page: `true` if its row group stored these counts. + known: BooleanBufferBuilder, +} + +impl Counts { + fn with_capacity(capacity: usize) -> Self { + Self { + values: Vec::with_capacity(capacity), + known: BooleanBufferBuilder::new(capacity), + } + } + + /// Adds `n` pages whose count is not known. + fn append_nulls(&mut self, n: usize) { + pad(&mut self.values, n); + self.known.append_n(n, false); + } + + /// Reads a list of counts, refusing negative numbers. `what` names the + /// counts in errors. + fn read(&mut self, prot: &mut ThriftSliceInputProtocol, what: &str) -> Result<()> { + let size = read_list_size(prot, ElementType::I64)?; + let mut buf = prot.as_slice(); + self.values.reserve(size); + let mut remaining = size; + while remaining > 0 { + // Counts from 0 to 63 take one byte each, with the top bit and + // the sign bit (the lowest) clear. When the next 8 counts are all + // like that, which is usual, convert them together. + if remaining >= 8 + && let Some((chunk, rest)) = buf.split_first_chunk::<8>() + && u64::from_le_bytes(*chunk) & 0x8181_8181_8181_8181 == 0 + { + self.values.extend(chunk.map(|b| u64::from(b >> 1))); + buf = rest; + remaining -= 8; + continue; + } + remaining -= 1; + // Numbers are stored so that small values of either sign are short: + // 0, -1, 1, -2, 2, ... are stored as 0, 1, 2, 3, 4, ... + let (stored, rest) = split_varint(buf)?; + buf = rest; + let count = (stored >> 1) as i64 ^ -((stored & 1) as i64); + let count = u64::try_from(count) + .map_err(|_| general_err!("ColumnIndex {what} count is negative {count}"))?; + self.values.push(count); + } + let used = prot.as_slice().len() - buf.len(); + Ok(prot.skip_bytes(used)?) + } + + /// Ends a row group of `len` pages whose counts start at `first`: checks + /// that it gave one count per page, or marks every page as unknown if it + /// gave none. + fn finish_row_group( + &mut self, + first: usize, + present: bool, + len: usize, + name: &str, + ) -> Result<()> { + if !present { + self.append_nulls(len); + return Ok(()); + } + let got = self.values.len() - first; + if got != len { + return Err(general_err!( + "ColumnIndex {name} length mismatch: expected {len}, got {got}" + )); + } + self.known.append_n(len, true); + Ok(()) + } + + fn finish(mut self) -> UInt64Array { + let nulls = NullBuffer::new(self.known.finish()); + let nulls = (nulls.null_count() > 0).then_some(nulls); + UInt64Array::new(ScalarBuffer::from(self.values), nulls) + } +} + +/// The names of the min and max fields, in the order of +/// [`ColumnIndexDecoder::min_max`]. +const MIN_MAX_NAMES: [&str; 2] = ["min_values", "max_values"]; + +/// Collects the page statistics of one column, across any number of row +/// groups, while reading the stored `ColumnIndex` bytes. +#[derive(Debug)] +pub(super) struct ColumnIndexDecoder { + /// The Arrow type the mins and maxes are turned into. A dictionary type + /// is replaced by its value type, as statistics come out as the values. + data_type: DataType, + /// One entry per page: `true` if the page has a min and max, `false` if + /// every value in the page is null (or the row group has no index). + has_min_max: Vec<bool>, + /// The mins, then the maxes. + min_max: [Vec<i32>; 2], + null_counts: Counts, + nan_counts: Counts, +} + +impl ColumnIndexDecoder { + /// `data_type` is the Arrow type the statistics will be turned into. + pub(super) fn new(data_type: &DataType, capacity: usize) -> Self { + let data_type = match data_type { + DataType::Dictionary(_, value_type) => value_type.as_ref(), + data_type => data_type, + }; + Self { + data_type: data_type.clone(), + has_min_max: Vec::with_capacity(capacity), + min_max: std::array::from_fn(|_| Vec::with_capacity(capacity)), + null_counts: Counts::with_capacity(capacity), + nan_counts: Counts::with_capacity(capacity), + } + } + + /// Adds `n` pages with no statistics, for a row group that has no index. + pub(super) fn append_nulls(&mut self, n: usize) { + self.has_min_max.resize(self.has_min_max.len() + n, false); + for values in &mut self.min_max { + pad(values, n); + } + self.null_counts.append_nulls(n); + self.nan_counts.append_nulls(n); + } + + /// Reads one stored `ColumnIndex` (the statistics for one column in one + /// row group) and adds its pages. + /// + /// If this returns an error, the decoder is left half filled and must not + /// be used again. + pub(super) fn append(&mut self, data: &[u8]) -> Result<()> { + let mut prot = ThriftSliceInputProtocol::new(data); + + // Where this row group's pages start in the shared buffers. + let first_page = self.has_min_max.len(); + let first_null_count = self.null_counts.values.len(); + let first_nan_count = self.nan_counts.values.len(); + + let mut num_pages: Option<usize> = None; + // The number of mins and maxes, in the order of `self.min_max`. + let mut num_min_max: [Option<usize>; 2] = [None; 2]; + let mut has_boundary_order = false; + let mut has_null_counts = false; + let mut has_nan_counts = false; + // Writers put the list of null pages first, so we normally know which + // pages to skip while reading the mins and maxes. If a file stores the + // mins or maxes before that list, we note where they are and read them + // at the end. + let mut deferred: [Option<&[u8]>; 2] = [None; 2]; + + let mut last_field_id = 0i16; + loop { + let field = prot.read_field_begin(last_field_id)?; + if field.field_type == FieldType::Stop { + break; + } + match (field.id, field.field_type) { + // null_pages + (1, FieldType::List) => { + if num_pages.is_some() { + return Err(general_err!("ColumnIndex has more than one null_pages")); + } + let size = read_list_size(&mut prot, ElementType::Bool)?; + // Each item is one byte. `read_list_size` checked they are there. + let flags = &prot.as_slice()[..size]; + // `true` is stored as 1, and `false` as 0 or 2. Checking + // all of them first, then converting, lets the compiler + // handle many at a time. + if flags.iter().fold(0, |max, &flag| max.max(flag)) > 2 { + let flag = flags.iter().find(|&&flag| flag > 2).unwrap(); + return Err(general_err!("cannot convert {} into bool", flag)); + } + // stored as "is this page all null", kept as the opposite + self.has_min_max.extend(flags.iter().map(|&flag| flag != 1)); + prot.skip_bytes(size)?; + num_pages = Some(size); + } + // min_values and max_values + (2 | 3, FieldType::List) => { + let i = (field.id - 2) as usize; + if num_min_max[i].is_some() { + return Err(general_err!( + "ColumnIndex has more than one {}", + MIN_MAX_NAMES[i] + )); + } + let size = read_list_size(&mut prot, ElementType::Binary)?; + if num_pages == Some(size) { + let has_min_max = &self.has_min_max[first_page..]; + let used = read_list(prot.as_slice(), &mut self.min_max[i], has_min_max)?; + prot.skip_bytes(used)?; + } else { + // The null pages are not known yet, or the counts do + // not match and the check after the loop reports the + // error. Step over the values for now. + deferred[i] = Some(prot.as_slice()); + for _ in 0..size { + prot.skip(FieldType::Binary)?; + } + } + num_min_max[i] = Some(size); + } + // boundary_order: not needed, but read so that a bad value is caught + (4, FieldType::I32) => { + BoundaryOrder::read_thrift(&mut prot)?; + has_boundary_order = true; + } + // null_counts + (5, FieldType::List) => { + if has_null_counts { + return Err(general_err!("ColumnIndex has more than one null_counts")); + } + self.null_counts.read(&mut prot, "null")?; + has_null_counts = true; + } + // nan_counts + (8, FieldType::List) => { + if has_nan_counts { + return Err(general_err!("ColumnIndex has more than one nan_counts")); + } + self.nan_counts.read(&mut prot, "NaN")?; + has_nan_counts = true; + } + // Anything else, including the level histograms (fields 6 and + // 7) and fields added to the format later, is not needed. + (_, FieldType::List) => skip_list(&mut prot)?, + _ => prot.skip(field.field_type)?, + } + last_field_id = field.id; + } + + let Some(len) = num_pages else { + return Err(general_err!("Required field null_pages is missing")); + }; + let [Some(num_mins), Some(num_maxes)] = num_min_max else { + let name = MIN_MAX_NAMES[usize::from(num_min_max[0].is_some())]; + return Err(general_err!("Required field {name} is missing")); + }; + if !has_boundary_order { + return Err(general_err!("Required field boundary_order is missing")); + } + if num_mins != len || num_maxes != len { + return Err(general_err!( + "ColumnIndex min/max length mismatch: expected {len}, got min={num_mins} max={num_maxes}" + )); + } + for (values, list) in self.min_max.iter_mut().zip(deferred) { + if let Some(list) = list { + read_list(list, values, &self.has_min_max[first_page..])?; + } + } + + self.null_counts + .finish_row_group(first_null_count, has_null_counts, len, "null_counts")?; + self.nan_counts + .finish_row_group(first_nan_count, has_nan_counts, len, "nan_counts")?; + Ok(()) + } + + /// Returns the statistics, with the mins and maxes turned into the Arrow + /// type given to [`Self::new`]. + pub(super) fn finish(self) -> Result<DataPageStatistics> { + // `collect_bool` packs 64 pages at a time, which is much faster than + // converting the list one page at a time + let len = self.has_min_max.len(); + let has_min_max = &self.has_min_max; + let nulls = NullBuffer::new(BooleanBuffer::collect_bool(len, |i| has_min_max[i])); + let nulls = (nulls.null_count() > 0).then_some(nulls); + let [mins, maxes] = self.min_max; + let to_array = |values: Vec<i32>, nulls| { + let array = Int32Array::new(values.into(), nulls); + // Every other type gives nulls, as the old code does. + Ok::<_, ParquetError>( + int32_to_logical(&array, &self.data_type)? + .unwrap_or_else(|| new_null_array(&self.data_type, len)), + ) + }; + Ok(DataPageStatistics { + mins: to_array(mins, nulls.clone())?, + maxes: to_array(maxes, nulls)?, + null_counts: self.null_counts.finish(), + nan_counts: self.nan_counts.finish(), + }) + } +} + +/// Reads the start of a list and checks that it holds the expected kind of item. +fn read_list_size(prot: &mut ThriftSliceInputProtocol, expected: ElementType) -> Result<usize> { + let list = prot.read_list_begin()?; + let size = list.size as usize; + // Some writers mark an empty list with a made up item kind, so only + // check the kind when there are items. + if size > 0 { + validate_list_type(expected, &list)?; + } + // Every item takes at least one byte, so a list claiming more items than + // there are bytes left is broken. Checking this first avoids reserving a + // huge amount of memory for a bad file. + let remaining = prot.as_slice().len(); + if size > remaining { + return Err(general_err!( + "Thrift list size {size} exceeds remaining input length {remaining}" + )); + } + Ok(size) +} + +/// Steps over a list we do not need. +/// +/// The level histograms are long lists of whole numbers, one or more per +/// page. The general purpose skip looks at each number on its own, which +/// made skipping them take a quarter of the total time. Whole numbers are +/// stored so that every byte except the last has its top bit set, so here we +/// just count the bytes whose top bit is clear. +fn skip_list(prot: &mut ThriftSliceInputProtocol) -> Result<()> { + let list = prot.read_list_begin()?; + match list.element_type { + ElementType::I16 | ElementType::I32 | ElementType::I64 => { + let count = list.size as usize; + let buf = prot.as_slice(); + let mut seen = 0; + let mut len = 0; + // While at least 8 numbers are left, the next 8 bytes all belong + // to the list, so count them together. + while count - seen >= 8 + && let Some(chunk) = buf.get(len..len + 8) + { + let word = u64::from_le_bytes(chunk.try_into().unwrap()); + seen += (!word & 0x8080_8080_8080_8080).count_ones() as usize; + len += 8; + } + while seen < count { + let byte = *buf.get(len).ok_or_else(|| eof_err!("Unexpected EOF"))?; + len += 1; + seen += usize::from(byte & 0x80 == 0); + } + prot.skip_bytes(len)?; + } + element_type => { Review Comment: ```suggestion ElementType::Bool => prot.skip_bytes(list.size as usize)?, element_type => { ``` see #11926/#11923 Not really an issue today, but future proofing in case another list<bool> is added down the road. -- 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]
