sdf-jkl commented on code in PR #9372: URL: https://github.com/apache/arrow-rs/pull/9372#discussion_r3716871094
########## parquet/src/encodings/decoding/alp_decoder.rs: ########## @@ -0,0 +1,1408 @@ +// 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. + +use std::ops::Range; + +use bytes::Bytes; + +use crate::basic::Encoding; +use crate::data_type::DataType; +use crate::encodings::alp::{ + ALP_COMPRESSION_MODE, ALP_DEFAULT_LOG_VECTOR_SIZE, ALP_HEADER_SIZE, + ALP_INTEGER_ENCODING_FOR_BIT_PACK, ALP_MAX_EXPONENT_F32, ALP_MAX_EXPONENT_F64, + ALP_MAX_LOG_VECTOR_SIZE, ALP_MIN_LOG_VECTOR_SIZE, AlpExact, AlpFloat, AlpHeader, AlpInfo, + ForInfo, +}; +use crate::encodings::decoding::Decoder; +use crate::errors::{ParquetError, Result}; +use crate::util::bit_util::BitReader; + +/// Parsed view of one vector's metadata and data sections. +/// +/// Each data section is described by its start offset into the page body; the +/// section bytes themselves stay in the body and are decoded lazily when the +/// vector is decoded. Section lengths are fully determined by the fixed-size +/// metadata at the front of the vector (`bit_width` for `packed_values`, +/// `num_exceptions` for both exception sections), so only the start offset is +/// stored. +#[derive(Debug, Clone, Copy)] +struct AlpEncodedVectorView<Exact: AlpExact> { + num_elements: u16, + alp_info: AlpInfo, + for_info: ForInfo<Exact>, + packed_values: usize, + exception_positions: usize, + exception_values: usize, +} + +impl<Exact: AlpExact> AlpEncodedVectorView<Exact> { + fn expected_stored_size(&self) -> usize { + AlpInfo::STORED_SIZE + + ForInfo::<Exact>::stored_size() + + self + .for_info + .get_data_stored_size(self.num_elements, self.alp_info.num_exceptions) + } + + /// Byte range of the bit-packed values section in the page body. + fn packed_values_range(&self) -> Range<usize> { + let len = self.for_info.get_bit_packed_size(self.num_elements); + self.packed_values..self.packed_values + len + } + + /// Byte range of the exception positions section (`u16` each) in the page body. + fn exception_positions_range(&self) -> Range<usize> { + let len = self.alp_info.num_exceptions as usize * std::mem::size_of::<u16>(); + self.exception_positions..self.exception_positions + len + } + + /// Byte range of the exception values section (`Exact::WIDTH` each) in the page body. + fn exception_values_range(&self) -> Range<usize> { + let len = self.alp_info.num_exceptions as usize * Exact::WIDTH; + self.exception_values..self.exception_values + len + } +} + +/// Parse and validate the 7-byte ALP page header: compression mode, integer +/// encoding, and vector-size range. +fn parse_alp_page_header(data: &[u8]) -> Result<AlpHeader> { + let header = AlpHeader::deserialize(data)?; + + if header.compression_mode != ALP_COMPRESSION_MODE { + return Err(general_err!( + "Invalid ALP page: unsupported compression mode {}", + header.compression_mode + )); + } + if header.integer_encoding != ALP_INTEGER_ENCODING_FOR_BIT_PACK { + return Err(general_err!( + "Invalid ALP page: unsupported integer encoding {}", + header.integer_encoding + )); + } + if header.vector_size < (1usize << ALP_MIN_LOG_VECTOR_SIZE) { + return Err(general_err!( + "Invalid ALP page: log_vector_size {} below min {}", + header.vector_size.trailing_zeros(), + ALP_MIN_LOG_VECTOR_SIZE + )); + } + if header.vector_size > (1usize << ALP_MAX_LOG_VECTOR_SIZE) { + return Err(general_err!( + "Invalid ALP page: log_vector_size {} exceeds max {}", + header.vector_size.trailing_zeros(), + ALP_MAX_LOG_VECTOR_SIZE + )); + } + + Ok(header) +} + +/// Read the little-endian `u32` vector offset at index `idx` from the offsets +/// section at the start of the page body. +fn read_offset(body: &[u8], idx: usize) -> Result<usize> { + let start = idx * std::mem::size_of::<u32>(); + let bytes = body + .get(start..start + std::mem::size_of::<u32>()) + .ok_or_else(|| general_err!("Invalid ALP page: offset index {} out of bounds", idx))?; + Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize) +} + +/// Parse a single vector section: +/// `[AlpInfo][ForInfo][PackedValues][ExceptionPositions][ExceptionValues]`. +fn parse_vector_view<Exact: AlpExact>( + body: &[u8], + vector_start: usize, + vector_end: usize, + num_elements: u16, +) -> Result<AlpEncodedVectorView<Exact>> { + let vector_bytes = &body[vector_start..vector_end]; + + let metadata_size = AlpInfo::STORED_SIZE + ForInfo::<Exact>::stored_size(); + if vector_bytes.len() < metadata_size { + return Err(general_err!( + "Invalid ALP page: vector metadata too short, expected at least {} bytes, got {}", + metadata_size, + vector_bytes.len() + )); + } + + let alp_info = AlpInfo { + exponent: vector_bytes[0], + factor: vector_bytes[1], + num_exceptions: u16::from_le_bytes([vector_bytes[2], vector_bytes[3]]), + }; + + let max_exponent = if Exact::WIDTH == 4 { + ALP_MAX_EXPONENT_F32 + } else { + ALP_MAX_EXPONENT_F64 + }; + + if alp_info.exponent > max_exponent { + return Err(general_err!( + "Invalid ALP page: exponent {} exceeds max {}", + alp_info.exponent, + max_exponent + )); + } + + if alp_info.factor > alp_info.exponent { + return Err(general_err!( + "Invalid ALP page: factor {} exceeds exponent {}", + alp_info.factor, + alp_info.exponent + )); + } + + if alp_info.num_exceptions > num_elements { + return Err(general_err!( + "Invalid ALP page: num_exceptions {} exceeds vector num_elements {}", + alp_info.num_exceptions, + num_elements + )); + } + + let for_start = AlpInfo::STORED_SIZE; + let for_end = for_start + Exact::WIDTH; + let frame_of_reference = Exact::from_le_slice(&vector_bytes[for_start..for_end]); + let bit_width = vector_bytes[for_end]; + + if bit_width as usize > Exact::WIDTH * 8 { + return Err(general_err!( + "Invalid ALP page: bit width {} exceeds {}", + bit_width, + Exact::WIDTH * 8 + )); + } + + let for_info = ForInfo::<Exact> { + frame_of_reference, + bit_width, + }; + + let data_size = for_info.get_data_stored_size(num_elements, alp_info.num_exceptions); + let expected_size = metadata_size + data_size; + if vector_bytes.len() < expected_size { + return Err(general_err!( + "Invalid ALP page: vector data too short, expected at least {} bytes, got {}", + expected_size, + vector_bytes.len() + )); + } + if vector_bytes.len() > expected_size { + return Err(general_err!( + "Invalid ALP page: vector data too long, expected {} bytes, got {}", + expected_size, + vector_bytes.len() + )); + } + + let data = &vector_bytes[metadata_size..expected_size]; + let packed_size = for_info.get_bit_packed_size(num_elements); + let positions_size = alp_info.num_exceptions as usize * std::mem::size_of::<u16>(); + + // Section offsets relative to the start of the data section: packed values + // first, then exception positions, then exception values. + let positions_start = packed_size; + let values_start = positions_start + positions_size; + + // Validate exception positions without materializing them. They are decoded + // straight from the body when the vector is decoded; here we only enforce + // that every position is in range so the whole page is validated up front. + for chunk in data[positions_start..values_start].chunks_exact(2) { + let position = u16::from_le_bytes([chunk[0], chunk[1]]); + if position >= num_elements { + return Err(general_err!( + "Invalid ALP page: exception position {} out of bounds for vector length {}", + position, + num_elements + )); + } + } + + // Store each section's start offset into the page body. Lengths are derived + // from the vector metadata at decode time, so no end offset is stored. + let data_start = vector_start + metadata_size; + let packed_values = data_start; + let exception_positions = data_start + positions_start; + let exception_values = data_start + values_start; + + Ok(AlpEncodedVectorView { + num_elements, + alp_info, + for_info, + packed_values, + exception_positions, + exception_values, + }) +} + +/// Live decode state for the one vector currently being consumed. +/// +/// Holds the bit position inside that vector's packed values plus the +/// vector-level constants needed to turn each packed integer back into a float. +/// `delivered` is the vector-local index of the next element to produce, so +/// exception patches (which use vector-local positions) land in the right place +/// even when a vector is split across several `get`/`skip` calls. +struct CurrentVector<Value: AlpFloat> { + reader: BitReader, + bit_width: u8, + frame_of_reference: Value::Exact, + scale: Value::Scale, + /// Number of this vector's elements not yet delivered or skipped. + remaining: usize, + /// Vector-local index of the next element to produce. + delivered: usize, + exception_positions: Bytes, + exception_values: Bytes, +} + +/// Largest slice decoded in one unpack-then-decode pass: the canonical ALP +/// vector size - 1024. +/// +/// The unpack scratch is sized to `min(vector_size, this)`, so vectors at the +/// default size or smaller are decoded whole, while larger (non-default) vectors +/// are decoded in canonical-vector-sized tiles. +/// +/// Bounding the tile to one canonical vector keeps the scratch L1-resident, which +/// is what makes the staged unpack-then-decode beat an in-place decode. +const DECODE_TILE_CAP: usize = 1 << ALP_DEFAULT_LOG_VECTOR_SIZE; + +/// Decode the next `out.len()` elements of the current vector into `out`, +/// patching any exceptions whose vector-local position falls in the +/// just-produced sub-range. +/// +/// Deltas are bulk-unpacked a tile at a time into the caller-provided `scratch` +/// via `get_batch` (which dispatches to the SIMD-friendly fixed-width `unpack` +/// kernels), then the inverse FOR and decimal decode run as one branchless, +/// state-free loop over that contiguous tile so the compiler can autovectorize +/// it. +fn decode_range<Value: AlpFloat>( + cur: &mut CurrentVector<Value>, + scratch: &mut [Value::Exact], + out: &mut [Value], +) -> Result<()> { + let frame_of_reference = cur.frame_of_reference; + if cur.bit_width == 0 { + // Every packed delta is zero, so all values share `frame_of_reference`. + let signed = frame_of_reference.reinterpret_as_signed(); + out.fill(Value::decode_value(signed, cur.scale)); + } else { + let bit_width = cur.bit_width as usize; + let scale = cur.scale; + for chunk in out.chunks_mut(scratch.len()) { + let deltas = &mut scratch[..chunk.len()]; + let unpacked = cur.reader.get_batch::<Value::Exact>(deltas, bit_width); + if unpacked != chunk.len() { + return Err(general_err!( + "Invalid ALP page: not enough packed bits to decode vector" + )); + } + for (slot, &delta) in chunk.iter_mut().zip(deltas.iter()) { + let signed = delta + .wrapping_add(frame_of_reference) + .reinterpret_as_signed(); + *slot = Value::decode_value(signed, scale); + } + } + } + + // Patch exceptions landing in `[delivered, delivered + out.len())`. Positions + // were validated in bounds when the vector was parsed, and patching is a + // positional overwrite, so it is independent of exception ordering. + let lo = cur.delivered; + let hi = cur.delivered + out.len(); + for (pos_chunk, value_chunk) in cur Review Comment: added here https://github.com/apache/arrow-rs/pull/9372/commits/199a50ac6882b8b0cd5b59f0d61554bbf06f6cda -- 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]
