alamb commented on code in PR #9372: URL: https://github.com/apache/arrow-rs/pull/9372#discussion_r3696575823
########## parquet/src/encodings/alp.rs: ########## @@ -0,0 +1,584 @@ +// 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. + +//! ALP (Adaptive Lossless floating-Point) Encoding +//! +//! Spec: <https://github.com/apache/parquet-format/blob/master/Encodings.md#adaptive-lossless-floating-point-alp--10> +//! +//! # Page layout +//! +//! An ALP-encoded page consists of a fixed-size header, an offset array +//! locating each vector inside the body, and the vector data itself: +//! +//! ```text +//! +-------------+-----------------------------+--------------------------------------+ +//! | Header | Offset Array | Vector Data | +//! | (7 bytes) | (num_vectors * 4 bytes) | (variable) | +//! +-------------+------+------+-----+---------+----------+----------+-----+----------+ +//! | Page Header | off0 | off1 | ... | off N-1 | Vector 0 | Vector 1 | ... | Vec N-1 | +//! | (7 bytes) | (4B) | (4B) | | (4B) |(variable)|(variable)| |(variable)| +//! +-------------+------+------+-----+---------+----------+----------+-----+----------+ +//! ``` +//! +//! Each vector entry has the form +//! `[AlpInfo][ForInfo][PackedValues][ExceptionPositions][ExceptionValues]`. + +use crate::errors::{ParquetError, Result}; +use crate::util::bit_util::{FromBitpacked, FromBytes}; + +pub(crate) const ALP_HEADER_SIZE: usize = 7; +pub(crate) const ALP_COMPRESSION_MODE: u8 = 0; +pub(crate) const ALP_INTEGER_ENCODING_FOR_BIT_PACK: u8 = 0; +pub(crate) const ALP_MIN_LOG_VECTOR_SIZE: u8 = 3; +pub(crate) const ALP_MAX_LOG_VECTOR_SIZE: u8 = 15; +/// Spec-recommended default `log_vector_size`: 1024-value vectors. +pub(crate) const ALP_DEFAULT_LOG_VECTOR_SIZE: u8 = 10; +pub(crate) const ALP_MAX_EXPONENT_F32: u8 = 10; +pub(crate) const ALP_MAX_EXPONENT_F64: u8 = 18; + +/// Page-level ALP header (7 bytes). +/// +/// ```text +/// Byte: 0 1 2 3 4 5 6 +/// +----------------+---------------+--------------+----+----+----+----+ +/// | compression | integer | log_vector | num_elements | +/// | _mode | _encoding | _size | (int32 LE) | +/// +----------------+---------------+--------------+----+----+----+----+ +/// ``` +/// +/// Layout in bytes: +/// - `[0]` `compression_mode` +/// - `[1]` `integer_encoding` +/// - `[2]` `log_vector_size` +/// - `[3..7]` `num_elements` (little-endian `i32`) +/// +/// The fields hold the *decoded* values used throughout the decoder, not the +/// raw on-disk encoding: +/// - `num_elements` is stored on disk as an `i32`, kept in memory as a `usize`. +/// - vector size is stored on disk as a `u8` `log_vector_size`, kept in memory +/// as the actual `vector_size` (`1 << log_vector_size`) `usize`. +/// +/// Each conversion happens once, in [`AlpHeader::deserialize`] and +/// [`AlpHeader::serialize`], so the rest of the decoder computes offsets and +/// sizes in `usize`. Those methods reject only what the target type cannot +/// represent; spec-level validity, such as the allowed vector-size range, is +/// enforced by the page parser. +#[derive(Debug, Clone, Copy)] +pub(crate) struct AlpHeader { + pub(crate) compression_mode: u8, + pub(crate) integer_encoding: u8, + pub(crate) vector_size: usize, + pub(crate) num_elements: usize, +} + +impl AlpHeader { + /// Parse a 7-byte page header from its little-endian on-disk form, + /// converting each field to its in-memory type: + /// - `log_vector_size` (`u8`) is expanded to `vector_size` with an + /// overflow-checked shift. + /// - `num_elements` (`i32`) is checked for non-negativity. + pub(crate) fn deserialize(bytes: &[u8]) -> Result<Self> { + if bytes.len() < ALP_HEADER_SIZE { + return Err(general_err!( + "Invalid ALP page: expected at least {} bytes for header, got {}", + ALP_HEADER_SIZE, + bytes.len() + )); + } + + let log_vector_size = bytes[2]; + let vector_size = 1usize + .checked_shl(u32::from(log_vector_size)) + .ok_or_else(|| { + general_err!( + "Invalid ALP page: log_vector_size {} too large to represent a vector size", + log_vector_size + ) + })?; + + let num_elements_i32 = i32::from_le_bytes([bytes[3], bytes[4], bytes[5], bytes[6]]); + let num_elements = usize::try_from(num_elements_i32).map_err(|_| { + general_err!( + "Invalid ALP page: num_elements {} must be >= 0", + num_elements_i32 + ) + })?; + + Ok(Self { + compression_mode: bytes[0], + integer_encoding: bytes[1], + vector_size, + num_elements, + }) + } + + /// Serialize this header into its 7-byte little-endian on-disk form. + /// + /// Converts the in-memory values back to the on-disk encoding, rejecting + /// what cannot be represented: `vector_size` must be a power of two (its log + /// is the on-disk field), and `num_elements` must fit in an `i32`. + /// Counterpart to [`AlpHeader::deserialize`]; consumed by the ALP encoder. + pub(crate) fn serialize(&self) -> Result<[u8; ALP_HEADER_SIZE]> { + if !self.vector_size.is_power_of_two() { + return Err(general_err!( + "Invalid ALP page: vector_size {} is not a power of two", + self.vector_size + )); + } + let log_vector_size = self.vector_size.trailing_zeros() as u8; + + let num_elements = i32::try_from(self.num_elements).map_err(|_| { + general_err!( + "Invalid ALP page: num_elements {} exceeds i32::MAX", + self.num_elements + ) + })?; + + let mut out = [0u8; ALP_HEADER_SIZE]; + out[0] = self.compression_mode; + out[1] = self.integer_encoding; + out[2] = log_vector_size; + out[3..7].copy_from_slice(&num_elements.to_le_bytes()); + Ok(out) + } + + /// `vector_size` is always `1 << log_vector_size` (see [`AlpHeader::deserialize`]), + /// so the division a `div_ceil` would emit is a shift. `vector_size` is a + /// runtime value, so the compiler cannot see that on its own. + pub(crate) fn num_vectors(&self) -> usize { + debug_assert!(self.vector_size.is_power_of_two()); + (self.num_elements + self.vector_size - 1) >> self.vector_size.trailing_zeros() + } + + /// Number of elements in vector `vector_index`: a full vector, the short + /// trailing remainder, or zero past the end of the page. + pub(crate) fn vector_num_elements(&self, vector_index: usize) -> u16 { + let start = vector_index.saturating_mul(self.vector_size); + let remaining = self.num_elements.saturating_sub(start); + remaining.min(self.vector_size) as u16 + } +} + +/// Per-vector ALP metadata (4 bytes). +/// +/// ```text +/// Byte: 0 1 2 3 +/// +----------+----------+---------+---------+ +/// | exponent | factor | num_exceptions | +/// | (uint8) | (uint8) | (uint16 LE) | +/// +----------+----------+---------+---------+ +/// ``` +#[derive(Debug, Clone, Copy)] +pub(crate) struct AlpInfo { + pub(crate) exponent: u8, + pub(crate) factor: u8, + pub(crate) num_exceptions: u16, +} + +impl AlpInfo { + pub(crate) const STORED_SIZE: usize = 4; + + /// Append this vector's ALP metadata in its on-disk little-endian form. + pub(crate) fn extend_serialized(&self, out: &mut Vec<u8>) { + out.push(self.exponent); + out.push(self.factor); + out.extend_from_slice(&self.num_exceptions.to_le_bytes()); + } +} + +/// Per-vector FOR (frame of reference) metadata: 5 bytes for `f32`, 9 for `f64`. +/// +/// ```text +/// +--------------------+-----------+ +/// | frame_of_reference | bit_width | +/// | (Exact::WIDTH, LE) | (uint8) | +/// +--------------------+-----------+ +/// ``` +#[derive(Debug, Clone, Copy)] +pub(crate) struct ForInfo<Exact: AlpExact> { + pub(crate) frame_of_reference: Exact, + pub(crate) bit_width: u8, +} + +impl<Exact: AlpExact> ForInfo<Exact> { + pub(crate) fn stored_size() -> usize { + Exact::WIDTH + 1 + } + + /// Append this vector's FOR metadata in its on-disk little-endian form. + pub(crate) fn extend_serialized(&self, out: &mut Vec<u8>) { + self.frame_of_reference.extend_le_bytes(out); + out.push(self.bit_width); + } + + pub(crate) fn get_bit_packed_size(&self, num_elements: u16) -> usize { + (self.bit_width as usize * num_elements as usize).div_ceil(8) + } + + pub(crate) fn get_data_stored_size(&self, num_elements: u16, num_exceptions: u16) -> usize { + let bit_packed_size = self.get_bit_packed_size(num_elements); + bit_packed_size + + num_exceptions as usize * std::mem::size_of::<u16>() + + num_exceptions as usize * Exact::WIDTH + } +} + +/// Exact integer type used by FOR reconstruction: `u32` for `f32`, `u64` for +/// `f64`. +/// +/// Why unsigned (not `i32`/`i64`)? The spec computes and stores deltas in +/// unsigned wrapping arithmetic: this avoids signed overflow when a vector's +/// range exceeds the signed maximum, and unpacking needs no sign extension. +/// Signed interpretation is applied later during decimal reconstruction. +pub(crate) trait AlpExact: + Copy + std::fmt::Debug + PartialEq + FromBitpacked + Default +{ + const WIDTH: usize; + type Signed: Copy + Ord + std::fmt::Debug + Send; + fn from_le_slice(slice: &[u8]) -> Self; + fn wrapping_add(self, rhs: Self) -> Self; + fn wrapping_sub(self, rhs: Self) -> Self; + fn reinterpret_as_signed(self) -> Self::Signed; + fn reinterpret_from_signed(signed: Self::Signed) -> Self; + /// Widen to `u64` for bit-packing, which is `u64`-oriented throughout. + fn to_u64(self) -> u64; + fn extend_le_bytes(self, out: &mut Vec<u8>); +} + +impl AlpExact for u32 { + const WIDTH: usize = 4; + type Signed = i32; + + fn from_le_slice(slice: &[u8]) -> Self { + u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]) + } + + fn wrapping_add(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + fn wrapping_sub(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + fn reinterpret_as_signed(self) -> Self::Signed { + i32::from_ne_bytes(self.to_ne_bytes()) + } + + fn reinterpret_from_signed(signed: Self::Signed) -> Self { + u32::from_ne_bytes(signed.to_ne_bytes()) + } + + fn to_u64(self) -> u64 { + u64::from(self) + } + + fn extend_le_bytes(self, out: &mut Vec<u8>) { + out.extend_from_slice(&self.to_le_bytes()); + } +} + +impl AlpExact for u64 { + const WIDTH: usize = 8; + type Signed = i64; + + fn from_le_slice(slice: &[u8]) -> Self { + u64::from_le_bytes([ + slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7], + ]) + } + + fn wrapping_add(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + fn wrapping_sub(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + fn reinterpret_as_signed(self) -> Self::Signed { + i64::from_ne_bytes(self.to_ne_bytes()) + } + + fn reinterpret_from_signed(signed: Self::Signed) -> Self { + u64::from_ne_bytes(signed.to_ne_bytes()) + } + + fn to_u64(self) -> u64 { + self + } + + fn extend_le_bytes(self, out: &mut Vec<u8>) { + out.extend_from_slice(&self.to_le_bytes()); + } +} +pub(crate) const ALP_POW10_F32: [f32; 11] = [ + 1.0, + 10.0, + 100.0, + 1000.0, + 10000.0, + 100000.0, + 1000000.0, + 10000000.0, + 100000000.0, + 1000000000.0, + 10000000000.0, +]; + +pub(crate) const ALP_POW10_F64: [f64; 19] = [ + 1.0, + 10.0, + 100.0, + 1000.0, + 10000.0, + 100000.0, + 1000000.0, + 10000000.0, + 100000000.0, + 1000000000.0, + 10000000000.0, + 100000000000.0, + 1000000000000.0, + 10000000000000.0, + 100000000000000.0, + 1000000000000000.0, + 10000000000000000.0, + 100000000000000000.0, + 1000000000000000000.0, +]; + +pub(crate) const ALP_NEG_POW10_F32: [f32; 11] = [ + 1.0, + 0.1, + 0.01, + 0.001, + 0.0001, + 0.00001, + 0.000001, + 0.0000001, + 0.00000001, + 0.000000001, + 0.0000000001, +]; + +pub(crate) const ALP_NEG_POW10_F64: [f64; 19] = [ + 1.0, + 0.1, + 0.01, + 0.001, + 0.0001, + 0.00001, + 0.000001, + 0.0000001, + 0.00000001, + 0.000000001, + 0.0000000001, + 0.00000000001, + 0.000000000001, + 0.0000000000001, + 0.00000000000001, + 0.000000000000001, + 0.0000000000000001, + 0.00000000000000001, + 0.000000000000000001, +]; + +pub(crate) trait AlpFloat: Review Comment: I think we should add some comments on the AlpFloat trait explaining that it represents the floating point representation and the there is a corresponding AlpExact for each AlpFloat and cross link it. Otherwise I was confused at first trying to understand what the difference was ########## parquet/src/encodings/alp.rs: ########## @@ -0,0 +1,584 @@ +// 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. + +//! ALP (Adaptive Lossless floating-Point) Encoding +//! +//! Spec: <https://github.com/apache/parquet-format/blob/master/Encodings.md#adaptive-lossless-floating-point-alp--10> +//! +//! # Page layout +//! +//! An ALP-encoded page consists of a fixed-size header, an offset array +//! locating each vector inside the body, and the vector data itself: +//! +//! ```text +//! +-------------+-----------------------------+--------------------------------------+ +//! | Header | Offset Array | Vector Data | +//! | (7 bytes) | (num_vectors * 4 bytes) | (variable) | +//! +-------------+------+------+-----+---------+----------+----------+-----+----------+ +//! | Page Header | off0 | off1 | ... | off N-1 | Vector 0 | Vector 1 | ... | Vec N-1 | +//! | (7 bytes) | (4B) | (4B) | | (4B) |(variable)|(variable)| |(variable)| +//! +-------------+------+------+-----+---------+----------+----------+-----+----------+ +//! ``` +//! +//! Each vector entry has the form +//! `[AlpInfo][ForInfo][PackedValues][ExceptionPositions][ExceptionValues]`. + +use crate::errors::{ParquetError, Result}; +use crate::util::bit_util::{FromBitpacked, FromBytes}; + +pub(crate) const ALP_HEADER_SIZE: usize = 7; +pub(crate) const ALP_COMPRESSION_MODE: u8 = 0; +pub(crate) const ALP_INTEGER_ENCODING_FOR_BIT_PACK: u8 = 0; +pub(crate) const ALP_MIN_LOG_VECTOR_SIZE: u8 = 3; +pub(crate) const ALP_MAX_LOG_VECTOR_SIZE: u8 = 15; +/// Spec-recommended default `log_vector_size`: 1024-value vectors. +pub(crate) const ALP_DEFAULT_LOG_VECTOR_SIZE: u8 = 10; +pub(crate) const ALP_MAX_EXPONENT_F32: u8 = 10; +pub(crate) const ALP_MAX_EXPONENT_F64: u8 = 18; + +/// Page-level ALP header (7 bytes). +/// +/// ```text +/// Byte: 0 1 2 3 4 5 6 +/// +----------------+---------------+--------------+----+----+----+----+ +/// | compression | integer | log_vector | num_elements | +/// | _mode | _encoding | _size | (int32 LE) | +/// +----------------+---------------+--------------+----+----+----+----+ +/// ``` +/// +/// Layout in bytes: +/// - `[0]` `compression_mode` +/// - `[1]` `integer_encoding` +/// - `[2]` `log_vector_size` +/// - `[3..7]` `num_elements` (little-endian `i32`) +/// +/// The fields hold the *decoded* values used throughout the decoder, not the +/// raw on-disk encoding: +/// - `num_elements` is stored on disk as an `i32`, kept in memory as a `usize`. +/// - vector size is stored on disk as a `u8` `log_vector_size`, kept in memory +/// as the actual `vector_size` (`1 << log_vector_size`) `usize`. +/// +/// Each conversion happens once, in [`AlpHeader::deserialize`] and +/// [`AlpHeader::serialize`], so the rest of the decoder computes offsets and +/// sizes in `usize`. Those methods reject only what the target type cannot +/// represent; spec-level validity, such as the allowed vector-size range, is +/// enforced by the page parser. +#[derive(Debug, Clone, Copy)] +pub(crate) struct AlpHeader { + pub(crate) compression_mode: u8, + pub(crate) integer_encoding: u8, + pub(crate) vector_size: usize, + pub(crate) num_elements: usize, +} + +impl AlpHeader { + /// Parse a 7-byte page header from its little-endian on-disk form, + /// converting each field to its in-memory type: + /// - `log_vector_size` (`u8`) is expanded to `vector_size` with an + /// overflow-checked shift. + /// - `num_elements` (`i32`) is checked for non-negativity. + pub(crate) fn deserialize(bytes: &[u8]) -> Result<Self> { + if bytes.len() < ALP_HEADER_SIZE { + return Err(general_err!( + "Invalid ALP page: expected at least {} bytes for header, got {}", + ALP_HEADER_SIZE, + bytes.len() + )); + } + + let log_vector_size = bytes[2]; + let vector_size = 1usize + .checked_shl(u32::from(log_vector_size)) + .ok_or_else(|| { + general_err!( + "Invalid ALP page: log_vector_size {} too large to represent a vector size", + log_vector_size + ) + })?; + + let num_elements_i32 = i32::from_le_bytes([bytes[3], bytes[4], bytes[5], bytes[6]]); + let num_elements = usize::try_from(num_elements_i32).map_err(|_| { + general_err!( + "Invalid ALP page: num_elements {} must be >= 0", + num_elements_i32 + ) + })?; + + Ok(Self { + compression_mode: bytes[0], + integer_encoding: bytes[1], + vector_size, + num_elements, + }) + } + + /// Serialize this header into its 7-byte little-endian on-disk form. + /// + /// Converts the in-memory values back to the on-disk encoding, rejecting + /// what cannot be represented: `vector_size` must be a power of two (its log + /// is the on-disk field), and `num_elements` must fit in an `i32`. + /// Counterpart to [`AlpHeader::deserialize`]; consumed by the ALP encoder. + pub(crate) fn serialize(&self) -> Result<[u8; ALP_HEADER_SIZE]> { + if !self.vector_size.is_power_of_two() { + return Err(general_err!( + "Invalid ALP page: vector_size {} is not a power of two", + self.vector_size + )); + } + let log_vector_size = self.vector_size.trailing_zeros() as u8; + + let num_elements = i32::try_from(self.num_elements).map_err(|_| { + general_err!( + "Invalid ALP page: num_elements {} exceeds i32::MAX", + self.num_elements + ) + })?; + + let mut out = [0u8; ALP_HEADER_SIZE]; + out[0] = self.compression_mode; + out[1] = self.integer_encoding; + out[2] = log_vector_size; + out[3..7].copy_from_slice(&num_elements.to_le_bytes()); + Ok(out) + } + + /// `vector_size` is always `1 << log_vector_size` (see [`AlpHeader::deserialize`]), + /// so the division a `div_ceil` would emit is a shift. `vector_size` is a + /// runtime value, so the compiler cannot see that on its own. + pub(crate) fn num_vectors(&self) -> usize { + debug_assert!(self.vector_size.is_power_of_two()); + (self.num_elements + self.vector_size - 1) >> self.vector_size.trailing_zeros() + } + + /// Number of elements in vector `vector_index`: a full vector, the short + /// trailing remainder, or zero past the end of the page. + pub(crate) fn vector_num_elements(&self, vector_index: usize) -> u16 { + let start = vector_index.saturating_mul(self.vector_size); + let remaining = self.num_elements.saturating_sub(start); + remaining.min(self.vector_size) as u16 + } +} + +/// Per-vector ALP metadata (4 bytes). +/// +/// ```text +/// Byte: 0 1 2 3 +/// +----------+----------+---------+---------+ +/// | exponent | factor | num_exceptions | +/// | (uint8) | (uint8) | (uint16 LE) | +/// +----------+----------+---------+---------+ +/// ``` +#[derive(Debug, Clone, Copy)] +pub(crate) struct AlpInfo { + pub(crate) exponent: u8, + pub(crate) factor: u8, + pub(crate) num_exceptions: u16, +} + +impl AlpInfo { + pub(crate) const STORED_SIZE: usize = 4; + + /// Append this vector's ALP metadata in its on-disk little-endian form. + pub(crate) fn extend_serialized(&self, out: &mut Vec<u8>) { + out.push(self.exponent); + out.push(self.factor); + out.extend_from_slice(&self.num_exceptions.to_le_bytes()); + } +} + +/// Per-vector FOR (frame of reference) metadata: 5 bytes for `f32`, 9 for `f64`. +/// +/// ```text +/// +--------------------+-----------+ +/// | frame_of_reference | bit_width | +/// | (Exact::WIDTH, LE) | (uint8) | +/// +--------------------+-----------+ +/// ``` +#[derive(Debug, Clone, Copy)] +pub(crate) struct ForInfo<Exact: AlpExact> { + pub(crate) frame_of_reference: Exact, + pub(crate) bit_width: u8, +} + +impl<Exact: AlpExact> ForInfo<Exact> { + pub(crate) fn stored_size() -> usize { + Exact::WIDTH + 1 + } + + /// Append this vector's FOR metadata in its on-disk little-endian form. + pub(crate) fn extend_serialized(&self, out: &mut Vec<u8>) { + self.frame_of_reference.extend_le_bytes(out); + out.push(self.bit_width); + } + + pub(crate) fn get_bit_packed_size(&self, num_elements: u16) -> usize { + (self.bit_width as usize * num_elements as usize).div_ceil(8) + } + + pub(crate) fn get_data_stored_size(&self, num_elements: u16, num_exceptions: u16) -> usize { + let bit_packed_size = self.get_bit_packed_size(num_elements); + bit_packed_size + + num_exceptions as usize * std::mem::size_of::<u16>() + + num_exceptions as usize * Exact::WIDTH + } +} + +/// Exact integer type used by FOR reconstruction: `u32` for `f32`, `u64` for +/// `f64`. +/// +/// Why unsigned (not `i32`/`i64`)? The spec computes and stores deltas in +/// unsigned wrapping arithmetic: this avoids signed overflow when a vector's +/// range exceeds the signed maximum, and unpacking needs no sign extension. +/// Signed interpretation is applied later during decimal reconstruction. +pub(crate) trait AlpExact: + Copy + std::fmt::Debug + PartialEq + FromBitpacked + Default +{ + const WIDTH: usize; + type Signed: Copy + Ord + std::fmt::Debug + Send; + fn from_le_slice(slice: &[u8]) -> Self; + fn wrapping_add(self, rhs: Self) -> Self; + fn wrapping_sub(self, rhs: Self) -> Self; + fn reinterpret_as_signed(self) -> Self::Signed; + fn reinterpret_from_signed(signed: Self::Signed) -> Self; + /// Widen to `u64` for bit-packing, which is `u64`-oriented throughout. + fn to_u64(self) -> u64; + fn extend_le_bytes(self, out: &mut Vec<u8>); +} + +impl AlpExact for u32 { + const WIDTH: usize = 4; + type Signed = i32; + + fn from_le_slice(slice: &[u8]) -> Self { + u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]) + } + + fn wrapping_add(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + fn wrapping_sub(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + fn reinterpret_as_signed(self) -> Self::Signed { + i32::from_ne_bytes(self.to_ne_bytes()) + } + + fn reinterpret_from_signed(signed: Self::Signed) -> Self { + u32::from_ne_bytes(signed.to_ne_bytes()) + } + + fn to_u64(self) -> u64 { + u64::from(self) + } + + fn extend_le_bytes(self, out: &mut Vec<u8>) { + out.extend_from_slice(&self.to_le_bytes()); + } +} + +impl AlpExact for u64 { + const WIDTH: usize = 8; + type Signed = i64; + + fn from_le_slice(slice: &[u8]) -> Self { + u64::from_le_bytes([ + slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7], + ]) + } + + fn wrapping_add(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + fn wrapping_sub(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + fn reinterpret_as_signed(self) -> Self::Signed { + i64::from_ne_bytes(self.to_ne_bytes()) + } + + fn reinterpret_from_signed(signed: Self::Signed) -> Self { + u64::from_ne_bytes(signed.to_ne_bytes()) + } + + fn to_u64(self) -> u64 { + self + } + + fn extend_le_bytes(self, out: &mut Vec<u8>) { + out.extend_from_slice(&self.to_le_bytes()); + } +} +pub(crate) const ALP_POW10_F32: [f32; 11] = [ + 1.0, + 10.0, + 100.0, + 1000.0, + 10000.0, + 100000.0, + 1000000.0, + 10000000.0, + 100000000.0, + 1000000000.0, + 10000000000.0, +]; + +pub(crate) const ALP_POW10_F64: [f64; 19] = [ + 1.0, + 10.0, + 100.0, + 1000.0, + 10000.0, + 100000.0, + 1000000.0, + 10000000.0, + 100000000.0, + 1000000000.0, + 10000000000.0, + 100000000000.0, + 1000000000000.0, + 10000000000000.0, + 100000000000000.0, + 1000000000000000.0, + 10000000000000000.0, + 100000000000000000.0, + 1000000000000000000.0, +]; + +pub(crate) const ALP_NEG_POW10_F32: [f32; 11] = [ + 1.0, + 0.1, + 0.01, + 0.001, + 0.0001, + 0.00001, + 0.000001, + 0.0000001, + 0.00000001, + 0.000000001, + 0.0000000001, +]; + +pub(crate) const ALP_NEG_POW10_F64: [f64; 19] = [ + 1.0, + 0.1, + 0.01, + 0.001, + 0.0001, + 0.00001, + 0.000001, + 0.0000001, + 0.00000001, + 0.000000001, + 0.0000000001, + 0.00000000001, + 0.000000000001, + 0.0000000000001, + 0.00000000000001, + 0.000000000000001, + 0.0000000000000001, + 0.00000000000000001, + 0.000000000000000001, +]; + +pub(crate) trait AlpFloat: + Copy + Default + PartialEq + std::ops::Mul<Output = Self> +{ + type Exact: AlpExact + FromBytes; + type Scale: Copy + Send; + + /// Largest `exponent` this type admits: 10 for `f32`, 18 for `f64`. + const MAX_EXPONENT: u8; + + /// Rounding magic number: `2^22 + 2^23` (`f32`) or `2^51 + 2^52` (`f64`). + const MAGIC_NUMBER: Self; + + /// Bounds outside which the scaled value cannot reach the exact integer + /// type: `i32` for `f32`, `i64` for `f64`. + const ENCODING_UPPER_LIMIT: Self; + const ENCODING_LOWER_LIMIT: Self; + + /// [`AlpFloat::ENCODING_UPPER_LIMIT`] as the exact signed integer. Stands in + /// for values ALP cannot represent, so that the round-trip check that + /// follows fails and the value is recorded as an exception. + const ENCODING_SENTINEL: <Self::Exact as AlpExact>::Signed; + + /// Precompute vector-level ALP decimal scale constants for: + /// `value = (encoded * 10^(factor)) * 10^(-exponent)`. + /// + /// Preconditions are validated during page parse. + fn decode_scale(exponent: u8, factor: u8) -> Self::Scale; + + /// Decode one signed exact integer using a precomputed two-step scale. + fn decode_value(signed_encoded: <Self::Exact as AlpExact>::Signed, scale: Self::Scale) -> Self; + + fn from_exact_bits(bits: Self::Exact) -> Self; + + fn to_exact_bits(self) -> Self::Exact; + + /// Precompute vector-level ALP decimal scale constants for the encode + /// direction: `encoded = fast_round((value * 10^(exponent)) * 10^(-factor))`. + fn encode_scale(exponent: u8, factor: u8) -> Self::Scale; + + /// Apply a scale as the same two separate multiplications the decode side + /// uses. Two steps rather than one multiplication by a combined constant: + /// the spec requires this on the normative decode path, and encoding with + /// the same arithmetic maximizes the values that round-trip. + fn apply_scale(self, scale: Self::Scale) -> Self; + + /// True for values ALP cannot turn into an exact integer: NaN, the + /// infinities, anything scaled past the exact integer type, and `-0.0` + /// (which would come back as `+0.0` and lose its sign). + fn is_impossible_to_encode(self) -> bool; + + /// Round to nearest by the "magic number" technique. Adding `magic` pushes + /// `x` into the binade where floats are spaced exactly 1.0 apart, so the + /// add itself snaps to the nearest integer, and subtracting `magic` back + /// is exact. `magic = 1.5 * 2^mantissa_bits` keeps the sum in that binade + /// for negative `x` too. Values large enough to be mis-rounded just fail + /// the caller's round-trip check and become exceptions. The add/sub must + /// not be simplified away: it *is* the rounding. + fn fast_round(self) -> <Self::Exact as AlpExact>::Signed; Review Comment: I think we could probbaly provide some more rationale about using `fast_round` before getting into the binade details. For example, perhaps we can explains that it is a way to round quickly using floating point math in a way that maximizes chance that multplying / dividing by factor and exponent will result in the same value (and thus avoiding exceptions) ########## 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: In this code, would it make sense to check / assert that the chunks_exact has no remainder (aka debug assert that `remainder` is null?)? ########## parquet/tests/arrow_reader/alp.rs: ########## @@ -0,0 +1,258 @@ +// 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 arrow::compute::concat_batches; +use arrow::util::test_util::parquet_test_data; +use arrow_array::cast::as_primitive_array; +use arrow_array::types::{Float32Type, Float64Type}; +use arrow_array::{Array, Float32Array, Float64Array, RecordBatch}; +use arrow_csv::ReaderBuilder as CsvReaderBuilder; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ArrowReaderBuilder; +use parquet::basic::Encoding; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; + +#[test] +fn test_read_f32_alp() { + let data_dir = PathBuf::from(parquet_test_data()); + let parquet_path = data_dir.join("alp_float_arade.parquet"); Review Comment: Also, I think we should test reading all the example ALP files (not just alp_float_arade.parquet) ########## parquet/tests/arrow_reader/alp.rs: ########## @@ -0,0 +1,258 @@ +// 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 arrow::compute::concat_batches; +use arrow::util::test_util::parquet_test_data; +use arrow_array::cast::as_primitive_array; +use arrow_array::types::{Float32Type, Float64Type}; +use arrow_array::{Array, Float32Array, Float64Array, RecordBatch}; +use arrow_csv::ReaderBuilder as CsvReaderBuilder; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ArrowReaderBuilder; +use parquet::basic::Encoding; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; + +#[test] +fn test_read_f32_alp() { + let data_dir = PathBuf::from(parquet_test_data()); + let parquet_path = data_dir.join("alp_float_arade.parquet"); + let expected_csv_path = data_dir.join("alp_arade_expect.csv"); + + let expected = read_expected_csv_batch(&expected_csv_path); + let actual = read_parquet_batch(&parquet_path); + + assert_eq!(actual.schema(), expected.schema(), "schema mismatch"); + assert_eq!( + actual.num_columns(), + expected.num_columns(), + "column mismatch" + ); + assert_eq!(actual.num_rows(), expected.num_rows(), "row count mismatch"); + + for col_idx in 0..actual.num_columns() { + let col_name = actual.schema().field(col_idx).name().clone(); + let actual_col = as_primitive_array::<Float32Type>(actual.column(col_idx).as_ref()); + let expected_col = as_primitive_array::<Float32Type>(expected.column(col_idx).as_ref()); + + for row_idx in 0..actual.num_rows() { + assert_eq!( + actual_col.is_valid(row_idx), + expected_col.is_valid(row_idx), + "null mismatch at column {col_name} row {row_idx}" + ); + if actual_col.is_valid(row_idx) { + let actual_value = actual_col.value(row_idx); + let expected_value = expected_col.value(row_idx); + assert!( + actual_value.to_bits() == expected_value.to_bits(), + "bit mismatch at column {col_name} row {row_idx}: expected={expected_value} actual={actual_value}" + ); + } + } + } +} + +/// Write the arade values with the ALP encoder and read them back, over real +/// float data rather than synthetic decimals. +/// +/// This checks losslessness only, not compression: these are `f32` values whose +/// encoded integers need 31 bits and which except 6.4% of the time, so ALP is +/// larger than PLAIN here, a property of the data rather than the encoder. +/// ALP's win on arade in the paper is on the `f64` version of the dataset. +#[test] +fn test_write_f32_alp_roundtrip() { Review Comment: As this is writing data, I think it would make make sense to move it to the arrow_writer tests either in parquet/tests/arrow_writer.rs or in parquet/src/arrow/arrow_writer/mod.rs. `parquet/src/arrow/arrow_writer/mod.rs` is probably the most consistent but it is already a massive module 🤔 In general I like the test pattern of round tripping the data used for alp encoding (aka read the csv and then round trip it through parquet with ALP enabled) to increase coverage. We could both: 1. Verify correctness 2. Compare compression (data page size) and that it was comparable to the checked in results ########## parquet/tests/arrow_reader/alp.rs: ########## @@ -0,0 +1,258 @@ +// 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 arrow::compute::concat_batches; +use arrow::util::test_util::parquet_test_data; +use arrow_array::cast::as_primitive_array; +use arrow_array::types::{Float32Type, Float64Type}; +use arrow_array::{Array, Float32Array, Float64Array, RecordBatch}; +use arrow_csv::ReaderBuilder as CsvReaderBuilder; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ArrowReaderBuilder; +use parquet::basic::Encoding; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; + +#[test] +fn test_read_f32_alp() { + let data_dir = PathBuf::from(parquet_test_data()); + let parquet_path = data_dir.join("alp_float_arade.parquet"); + let expected_csv_path = data_dir.join("alp_arade_expect.csv"); + + let expected = read_expected_csv_batch(&expected_csv_path); + let actual = read_parquet_batch(&parquet_path); + + assert_eq!(actual.schema(), expected.schema(), "schema mismatch"); + assert_eq!( + actual.num_columns(), + expected.num_columns(), + "column mismatch" + ); + assert_eq!(actual.num_rows(), expected.num_rows(), "row count mismatch"); + + for col_idx in 0..actual.num_columns() { + let col_name = actual.schema().field(col_idx).name().clone(); + let actual_col = as_primitive_array::<Float32Type>(actual.column(col_idx).as_ref()); + let expected_col = as_primitive_array::<Float32Type>(expected.column(col_idx).as_ref()); + + for row_idx in 0..actual.num_rows() { + assert_eq!( + actual_col.is_valid(row_idx), + expected_col.is_valid(row_idx), + "null mismatch at column {col_name} row {row_idx}" + ); + if actual_col.is_valid(row_idx) { + let actual_value = actual_col.value(row_idx); + let expected_value = expected_col.value(row_idx); + assert!( + actual_value.to_bits() == expected_value.to_bits(), + "bit mismatch at column {col_name} row {row_idx}: expected={expected_value} actual={actual_value}" + ); + } + } + } +} + +/// Write the arade values with the ALP encoder and read them back, over real +/// float data rather than synthetic decimals. +/// +/// This checks losslessness only, not compression: these are `f32` values whose +/// encoded integers need 31 bits and which except 6.4% of the time, so ALP is +/// larger than PLAIN here, a property of the data rather than the encoder. +/// ALP's win on arade in the paper is on the `f64` version of the dataset. +#[test] +fn test_write_f32_alp_roundtrip() { + let data_dir = PathBuf::from(parquet_test_data()); + let expected = read_expected_csv_batch(&data_dir.join("alp_arade_expect.csv")); + + let alp_bytes = write_batch(&expected, Encoding::ALP); + let actual = read_parquet_bytes(alp_bytes); + assert_eq!(actual.num_rows(), expected.num_rows(), "row count mismatch"); + + for col_idx in 0..expected.num_columns() { + let col_name = expected.schema().field(col_idx).name().clone(); + let actual_col = as_primitive_array::<Float32Type>(actual.column(col_idx).as_ref()); + let expected_col = as_primitive_array::<Float32Type>(expected.column(col_idx).as_ref()); + + for row_idx in 0..expected.num_rows() { + assert_eq!( + actual_col.is_valid(row_idx), + expected_col.is_valid(row_idx), + "null mismatch at column {col_name} row {row_idx}" + ); + if expected_col.is_valid(row_idx) { + // Bitwise, so that NaN and -0.0 are held to the same standard. + assert_eq!( + actual_col.value(row_idx).to_bits(), + expected_col.value(row_idx).to_bits(), + "bit mismatch at column {col_name} row {row_idx}" + ); + } + } + } +} + +/// Round-trip a nullable column under both data page versions. +/// +/// The versions differ in what the reader can tell the value decoder. A v2 page +/// header carries `num_nulls`, so the decoder is handed the exact count of +/// encoded values; a v1 header carries only `num_values`, which counts nulls, so +/// the decoder receives the level count instead. Only non-null values are +/// encoded either way, which is why the ALP header's element count - not the +/// count the reader passes in - is what governs the decode. +/// +/// The exception values are here deliberately: they make the encoded and +/// unencoded value counts differ from the level count in two different ways at +/// once. +#[test] +fn test_alp_roundtrip_page_versions_with_nulls() { Review Comment: Likewise this appears to a roudnteip writer test, so it probably makes sense to move it with the other round trip writer tests ########## parquet/tests/arrow_reader/alp.rs: ########## @@ -0,0 +1,258 @@ +// 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 arrow::compute::concat_batches; +use arrow::util::test_util::parquet_test_data; +use arrow_array::cast::as_primitive_array; +use arrow_array::types::{Float32Type, Float64Type}; +use arrow_array::{Array, Float32Array, Float64Array, RecordBatch}; +use arrow_csv::ReaderBuilder as CsvReaderBuilder; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ArrowReaderBuilder; +use parquet::basic::Encoding; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; + +#[test] +fn test_read_f32_alp() { + let data_dir = PathBuf::from(parquet_test_data()); + let parquet_path = data_dir.join("alp_float_arade.parquet"); + let expected_csv_path = data_dir.join("alp_arade_expect.csv"); + + let expected = read_expected_csv_batch(&expected_csv_path); + let actual = read_parquet_batch(&parquet_path); + + assert_eq!(actual.schema(), expected.schema(), "schema mismatch"); + assert_eq!( + actual.num_columns(), + expected.num_columns(), + "column mismatch" + ); + assert_eq!(actual.num_rows(), expected.num_rows(), "row count mismatch"); + + for col_idx in 0..actual.num_columns() { Review Comment: This is a lot of code to compare record batches and it will be much slower than the normal kernels. I recommend using normal record batch comparison here rather than explicit row-by-row (aka `assert_eq!(expected == actual)` ########## parquet/src/encodings/encoding/mod.rs: ########## @@ -84,26 +86,90 @@ pub fn get_encoder<T: DataType>( encoding: Encoding, descr: &ColumnDescPtr, ) -> Result<Box<dyn Encoder<T>>> { - let encoder: Box<dyn Encoder<T>> = match encoding { - Encoding::PLAIN => Box::new(PlainEncoder::new()), - Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => { - return Err(general_err!( - "Cannot initialize this encoding through this function" - )); + <T::T as private::GetEncoder>::get_encoder(descr, encoding) +} + +pub(crate) mod private { + use super::*; + + /// A trait that allows getting an [`Encoder`] implementation for a [`DataType`] + /// with the corresponding [`ParquetValueType`]. This is necessary to support + /// [`Encoder`] implementations that may not be applicable for all [`DataType`] + /// and by extension all [`ParquetValueType`], such as ALP, which encodes only + /// floating-point columns. + /// + /// [`ParquetValueType`]: crate::data_type::private::ParquetValueType + pub trait GetEncoder { Review Comment: Why do we need a new `GetEncoder` trait? I think the ColumnDescPtr physical_type.logical_type For example the fallback encoder depends on type https://github.com/apache/arrow-rs/blob/8042ea288e084107b602f9e25a850314942567b6/parquet/src/column/writer/mod.rs#L1735-L1734 Alternately, maybe we could pass a `&DataType` as a parameter into the system ########## parquet/tests/arrow_reader/alp.rs: ########## @@ -0,0 +1,258 @@ +// 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 arrow::compute::concat_batches; +use arrow::util::test_util::parquet_test_data; +use arrow_array::cast::as_primitive_array; +use arrow_array::types::{Float32Type, Float64Type}; +use arrow_array::{Array, Float32Array, Float64Array, RecordBatch}; +use arrow_csv::ReaderBuilder as CsvReaderBuilder; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ArrowReaderBuilder; +use parquet::basic::Encoding; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; + +#[test] +fn test_read_f32_alp() { + let data_dir = PathBuf::from(parquet_test_data()); + let parquet_path = data_dir.join("alp_float_arade.parquet"); + let expected_csv_path = data_dir.join("alp_arade_expect.csv"); + + let expected = read_expected_csv_batch(&expected_csv_path); + let actual = read_parquet_batch(&parquet_path); + + assert_eq!(actual.schema(), expected.schema(), "schema mismatch"); + assert_eq!( + actual.num_columns(), + expected.num_columns(), + "column mismatch" + ); + assert_eq!(actual.num_rows(), expected.num_rows(), "row count mismatch"); + + for col_idx in 0..actual.num_columns() { + let col_name = actual.schema().field(col_idx).name().clone(); + let actual_col = as_primitive_array::<Float32Type>(actual.column(col_idx).as_ref()); + let expected_col = as_primitive_array::<Float32Type>(expected.column(col_idx).as_ref()); + + for row_idx in 0..actual.num_rows() { + assert_eq!( + actual_col.is_valid(row_idx), + expected_col.is_valid(row_idx), + "null mismatch at column {col_name} row {row_idx}" + ); + if actual_col.is_valid(row_idx) { + let actual_value = actual_col.value(row_idx); + let expected_value = expected_col.value(row_idx); + assert!( + actual_value.to_bits() == expected_value.to_bits(), + "bit mismatch at column {col_name} row {row_idx}: expected={expected_value} actual={actual_value}" + ); + } + } + } +} + +/// Write the arade values with the ALP encoder and read them back, over real +/// float data rather than synthetic decimals. +/// +/// This checks losslessness only, not compression: these are `f32` values whose +/// encoded integers need 31 bits and which except 6.4% of the time, so ALP is +/// larger than PLAIN here, a property of the data rather than the encoder. +/// ALP's win on arade in the paper is on the `f64` version of the dataset. +#[test] +fn test_write_f32_alp_roundtrip() { + let data_dir = PathBuf::from(parquet_test_data()); + let expected = read_expected_csv_batch(&data_dir.join("alp_arade_expect.csv")); + + let alp_bytes = write_batch(&expected, Encoding::ALP); + let actual = read_parquet_bytes(alp_bytes); + assert_eq!(actual.num_rows(), expected.num_rows(), "row count mismatch"); + + for col_idx in 0..expected.num_columns() { Review Comment: this comparison logic I think can be avoided if we just compare the record batches -- 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]
