This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11741-915a1b498c95e6edc5c1b689ed2315aa7635d393 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 68f3f54ebc7171c566a198305a65485087d5f8d4 Author: Bhaskar Gurram <[email protected]> AuthorDate: Wed Oct 28 22:22:07 2026 +0000 Parquet: move `write_leaf`'s INT32/INT64/FLBA conversions into shared visitors (#11741) # Which issue does this PR close? - Part of #5550. This is the write-path refactor from #11549, split out as the reviewers asked so it can be reviewed and benchmarked on its own. #11549 will then only add `ArrowSbbf` on top of it. # Rationale for this change `ArrowSbbf` in #11549 checks Arrow values against a bloom filter written by `ArrowWriter`, so it has to apply the same conversions the writer applies before hashing (`Int8` to `INT32`, `Date64` to days, decimals to `INT32`, `INT64` or `FIXED_LEN_BYTE_ARRAY`, and so on). Instead of a second copy of those conversions that could drift from the writer, I moved them out of `write_leaf` into functions both can use. # What changes are included in this PR? - The `INT32`, `INT64` and `FIXED_LEN_BYTE_ARRAY` arms of `write_leaf` move into crate-private `visit_i32`, `visit_i64` and `visit_fixed_len_byte_array`. These pass the values to a visitor: - `PrimitiveVisitor` gets either the Arrow values as they are (`slice`) or the values with a per-value conversion (`map`). - `FixedLenByteArrayVisitor` gets the value length and a closure that fills each value. - The writer implements these with `WritePrimitive` and `WriteFixedLenByteArray`. Types that were written straight from the Arrow buffer still are. Types that went through `unary` are collected into a `Vec` with the same conversion. FLBA values are still built with `chunk_array_slice`. The one difference is `Decimal32` to `INT32`, which now borrows the buffer instead of copying it through an identity `unary`. - The `get_*_array_slice` helpers for intervals, decimals, `Float16` and `FixedSizeBinary` are folded into `visit_fixed_len_byte_array` and `visit_decimal`. The match arms, the error messages and the decimal length `assert!` are unchanged. The parts only `ArrowSbbf` needs (its own error for unsupported types, a `Decimal32` arm in `visit_i64`, and sign extension for decimal columns wider than the Arrow type) stay in #11549. # Are these changes tested? There is no change in behaviour, so I didn't add tests. The existing `arrow_writer` tests cover every arm. adriangb benchmarked these write paths on #11549 and they were within noise of main. Could someone run `run benchmark arrow_writer` on this PR as well? # Are there any significant user-facing changes? No. The new functions and traits are crate-private. --------- Co-authored-by: Adrian Garcia Badaracco <[email protected]> Co-authored-by: Ed Seidl <[email protected]> --- parquet/Cargo.toml | 4 +- parquet/src/arrow/arrow_writer/mod.rs | 623 ++++++++++++++++++---------------- parquet/src/util/interner.rs | 7 + 3 files changed, 340 insertions(+), 294 deletions(-) diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index b895fcc485..7e66a21e70 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -67,7 +67,9 @@ serde_json = { version = "1.0", default-features = false, features = ["std"], op seq-macro = { version = "0.3", default-features = false } futures = { version = "0.3", default-features = false, features = ["std"], optional = true } tokio = { version = "1.0", optional = true, default-features = false, features = ["macros", "rt", "io-util"] } -hashbrown = { version = "0.17", default-features = false } +# `inline-more` marks `HashTable::entry` `#[inline]`, so the dictionary encoder's +# hash probe is compiled into `Interner::intern` (see the note there). +hashbrown = { version = "0.17", default-features = false, features = ["inline-more"] } twox-hash = { version = "2.0", default-features = false, features = ["xxhash64"] } half = { version = "2.1", default-features = false, features = ["num-traits"] } crc32fast = { version = "1.4.2", optional = true, default-features = false } diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 257ec463c7..7e6264c398 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -44,9 +44,9 @@ use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter}; use crate::column::page_encryption::PageEncryptor; use crate::column::writer::encoder::ColumnValueEncoder; use crate::column::writer::{ - ColumnCloseResult, ColumnWriter, GenericColumnWriter, get_column_writer, + ColumnCloseResult, ColumnWriter, ColumnWriterImpl, GenericColumnWriter, get_column_writer, }; -use crate::data_type::{ByteArray, FixedLenByteArray}; +use crate::data_type::{ByteArray, FixedLenByteArray, FixedLenByteArrayType}; use std::collections::HashSet; type DistinctValuesSet = HashSet<u64>; #[cfg(feature = "encryption")] @@ -1639,87 +1639,13 @@ fn write_leaf( match writer { // Note: this should match the contents of arrow_to_parquet_type - ColumnWriter::Int32ColumnWriter(typed) => { - match column.data_type() { - ArrowDataType::Null => { - let array = Int32Array::new_null(column.len()); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Int8 => { - let array: Int32Array = column.as_primitive::<Int8Type>().unary(|x| x as i32); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Int16 => { - let array: Int32Array = column.as_primitive::<Int16Type>().unary(|x| x as i32); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Int32 => { - write_primitive(typed, column.as_primitive::<Int32Type>().values(), levels) - } - ArrowDataType::UInt8 => { - let array: Int32Array = column.as_primitive::<UInt8Type>().unary(|x| x as i32); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::UInt16 => { - let array: Int32Array = column.as_primitive::<UInt16Type>().unary(|x| x as i32); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::UInt32 => { - // follow C++ implementation and use overflow/reinterpret cast from u32 to i32 which will map - // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0` - let array = column.as_primitive::<UInt32Type>(); - write_primitive(typed, array.values().inner().typed_data(), levels) - } - ArrowDataType::Date32 => { - let array = column.as_primitive::<Date32Type>(); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Time32(TimeUnit::Second) => { - let array = column.as_primitive::<Time32SecondType>(); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Time32(TimeUnit::Millisecond) => { - let array = column.as_primitive::<Time32MillisecondType>(); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Date64 => { - // If the column is a Date64, we truncate it - let array: Int32Array = column - .as_primitive::<Date64Type>() - .unary(|x| (x / 86_400_000) as _); - - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Decimal32(_, _) => { - let array = column - .as_primitive::<Decimal32Type>() - .unary::<_, Int32Type>(|v| v); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Decimal64(_, _) => { - // use the int32 to represent the decimal with low precision - let array = column - .as_primitive::<Decimal64Type>() - .unary::<_, Int32Type>(|v| v as i32); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Decimal128(_, _) => { - // use the int32 to represent the decimal with low precision - let array = column - .as_primitive::<Decimal128Type>() - .unary::<_, Int32Type>(|v| v as i32); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Decimal256(_, _) => { - // use the int32 to represent the decimal with low precision - let array = column - .as_primitive::<Decimal256Type>() - .unary::<_, Int32Type>(|v| v.as_i128() as i32); - write_primitive(typed, array.values(), levels) - } - d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))), - } - } + ColumnWriter::Int32ColumnWriter(typed) => visit_i32( + column, + WritePrimitive { + writer: typed, + levels, + }, + ), ColumnWriter::BoolColumnWriter(typed) => { let array = column.as_boolean(); let values = get_bool_array_slice(array, indices.iter().copied()); @@ -1733,93 +1659,13 @@ fn write_leaf( None, ) } - ColumnWriter::Int64ColumnWriter(typed) => { - match column.data_type() { - ArrowDataType::Date64 => { - let array = column - .as_primitive::<Date64Type>() - .reinterpret_cast::<Int64Type>(); - - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Int64 => { - let array = column.as_primitive::<Int64Type>(); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::UInt64 => { - let values = column.as_primitive::<UInt64Type>().values(); - // follow C++ implementation and use overflow/reinterpret cast from u64 to i64 which will map - // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0` - let array = values.inner().typed_data::<i64>(); - write_primitive(typed, array, levels) - } - ArrowDataType::Time64(TimeUnit::Microsecond) => { - let array = column.as_primitive::<Time64MicrosecondType>(); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Time64(TimeUnit::Nanosecond) => { - let array = column.as_primitive::<Time64NanosecondType>(); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Timestamp(unit, _) => match unit { - TimeUnit::Second => { - let array = column.as_primitive::<TimestampSecondType>(); - write_primitive(typed, array.values(), levels) - } - TimeUnit::Millisecond => { - let array = column.as_primitive::<TimestampMillisecondType>(); - write_primitive(typed, array.values(), levels) - } - TimeUnit::Microsecond => { - let array = column.as_primitive::<TimestampMicrosecondType>(); - write_primitive(typed, array.values(), levels) - } - TimeUnit::Nanosecond => { - let array = column.as_primitive::<TimestampNanosecondType>(); - write_primitive(typed, array.values(), levels) - } - }, - ArrowDataType::Duration(unit) => match unit { - TimeUnit::Second => { - let array = column.as_primitive::<DurationSecondType>(); - write_primitive(typed, array.values(), levels) - } - TimeUnit::Millisecond => { - let array = column.as_primitive::<DurationMillisecondType>(); - write_primitive(typed, array.values(), levels) - } - TimeUnit::Microsecond => { - let array = column.as_primitive::<DurationMicrosecondType>(); - write_primitive(typed, array.values(), levels) - } - TimeUnit::Nanosecond => { - let array = column.as_primitive::<DurationNanosecondType>(); - write_primitive(typed, array.values(), levels) - } - }, - ArrowDataType::Decimal64(_, _) => { - let array = column - .as_primitive::<Decimal64Type>() - .reinterpret_cast::<Int64Type>(); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Decimal128(_, _) => { - // use the int64 to represent the decimal with low precision - let array = column - .as_primitive::<Decimal128Type>() - .unary::<_, Int64Type>(|v| v as i64); - write_primitive(typed, array.values(), levels) - } - ArrowDataType::Decimal256(_, _) => { - // use the int64 to represent the decimal with low precision - let array = column - .as_primitive::<Decimal256Type>() - .unary::<_, Int64Type>(|v| v.as_i128() as i64); - write_primitive(typed, array.values(), levels) - } - d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))), - } - } + ColumnWriter::Int64ColumnWriter(typed) => visit_i64( + column, + WritePrimitive { + writer: typed, + levels, + }, + ), ColumnWriter::Int96ColumnWriter(_typed) => { unreachable!("Currently unreachable because data type not supported") } @@ -1834,63 +1680,321 @@ fn write_leaf( ColumnWriter::ByteArrayColumnWriter(_) => { unreachable!("should use ByteArrayWriter") } - ColumnWriter::FixedLenByteArrayColumnWriter(typed) => { - let bytes = match column.data_type() { - ArrowDataType::Interval(interval_unit) => match interval_unit { - IntervalUnit::YearMonth => { - let array = column.as_primitive::<IntervalYearMonthType>(); - get_interval_ym_array_slice(array, indices.iter().copied()) - } - IntervalUnit::DayTime => { - let array = column.as_primitive::<IntervalDayTimeType>(); - get_interval_dt_array_slice(array, indices.iter().copied()) - } - IntervalUnit::MonthDayNano => { - return Err(ParquetError::NYI(format!( - "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented" - ))); - } - }, - ArrowDataType::FixedSizeBinary(_) => { - let array = column.as_fixed_size_binary(); - get_fsb_array_slice(array, indices.iter().copied()) - } - ArrowDataType::Decimal32(_, _) => { - let array = column.as_primitive::<Decimal32Type>(); - get_decimal_array_slice(array, indices.iter().copied()) - } - ArrowDataType::Decimal64(_, _) => { - let array = column.as_primitive::<Decimal64Type>(); - get_decimal_array_slice(array, indices.iter().copied()) - } - ArrowDataType::Decimal128(_, _) => { - let array = column.as_primitive::<Decimal128Type>(); - get_decimal_array_slice(array, indices.iter().copied()) - } - ArrowDataType::Decimal256(_, _) => { - let array = column.as_primitive::<Decimal256Type>(); - get_decimal_array_slice(array, indices.iter().copied()) - } - ArrowDataType::Float16 => { - let array = column.as_primitive::<Float16Type>(); - get_float_16_array_slice(array, indices.iter().copied()) - } - _ => { - return Err(ParquetError::NYI( - "Attempting to write an Arrow type that is not yet implemented".to_string(), - )); - } - }; - typed.write_batch_internal( - bytes.as_slice(), - None, - levels.def_level_data().as_ref(), - levels.rep_level_data().as_ref(), - None, - None, - None, - ) + ColumnWriter::FixedLenByteArrayColumnWriter(typed) => visit_fixed_len_byte_array( + column, + WriteFixedLenByteArray { + writer: typed, + levels, + }, + ), + } +} + +/// Receives the values [`ArrowWriter`] stores for an Arrow array in an `INT32` +/// (`T = i32`) or `INT64` (`T = i64`) column, which are also the values hashed +/// into the column's bloom filter +/// +/// [`visit_i32`] and [`visit_i64`] hold the conversion from each Arrow type, and +/// call [`Self::slice`] when the Arrow values are stored as they are, or +/// [`Self::map`] when each value is converted. This lets the writer borrow the +/// Arrow buffer when no conversion is needed. +pub(crate) trait PrimitiveVisitor<T> { + type Output; + + /// Receives values that are stored as they are + fn slice(self, values: &[T]) -> Result<Self::Output>; + + /// Receives values that are each stored as `convert(value)` + fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<Self::Output>; +} + +/// Passes the values [`ArrowWriter`] stores for `column` in an `INT32` column to `visitor` +pub(crate) fn visit_i32<V: PrimitiveVisitor<i32>>( + column: &dyn arrow_array::Array, + visitor: V, +) -> Result<V::Output> { + // Note: this should match the contents of arrow_to_parquet_type + match column.data_type() { + ArrowDataType::Null => visitor.slice(Int32Array::new_null(column.len()).values()), + ArrowDataType::Int8 => { + visitor.map(column.as_primitive::<Int8Type>().values(), |x| x as i32) + } + ArrowDataType::Int16 => { + visitor.map(column.as_primitive::<Int16Type>().values(), |x| x as i32) + } + ArrowDataType::Int32 => visitor.slice(column.as_primitive::<Int32Type>().values()), + ArrowDataType::UInt8 => { + visitor.map(column.as_primitive::<UInt8Type>().values(), |x| x as i32) + } + ArrowDataType::UInt16 => { + visitor.map(column.as_primitive::<UInt16Type>().values(), |x| x as i32) + } + ArrowDataType::UInt32 => { + // follow C++ implementation and use overflow/reinterpret cast from u32 to i32 which will map + // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0` + let values = column.as_primitive::<UInt32Type>().values(); + visitor.slice(values.inner().typed_data()) + } + ArrowDataType::Date32 => visitor.slice(column.as_primitive::<Date32Type>().values()), + ArrowDataType::Time32(TimeUnit::Second) => { + visitor.slice(column.as_primitive::<Time32SecondType>().values()) + } + ArrowDataType::Time32(TimeUnit::Millisecond) => { + visitor.slice(column.as_primitive::<Time32MillisecondType>().values()) + } + ArrowDataType::Date64 => { + // If the column is a Date64, we truncate it + visitor.map(column.as_primitive::<Date64Type>().values(), |x| { + (x / 86_400_000) as i32 + }) + } + ArrowDataType::Decimal32(_, _) => { + visitor.slice(column.as_primitive::<Decimal32Type>().values()) + } + ArrowDataType::Decimal64(_, _) => { + // use the int32 to represent the decimal with low precision + visitor.map(column.as_primitive::<Decimal64Type>().values(), |v| { + v as i32 + }) + } + ArrowDataType::Decimal128(_, _) => { + // use the int32 to represent the decimal with low precision + visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| { + v as i32 + }) + } + ArrowDataType::Decimal256(_, _) => { + // use the int32 to represent the decimal with low precision + visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| { + v.as_i128() as i32 + }) + } + d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))), + } +} + +/// Passes the values [`ArrowWriter`] stores for `column` in an `INT64` column to `visitor` +pub(crate) fn visit_i64<V: PrimitiveVisitor<i64>>( + column: &dyn arrow_array::Array, + visitor: V, +) -> Result<V::Output> { + // Note: this should match the contents of arrow_to_parquet_type + match column.data_type() { + ArrowDataType::Date64 => visitor.slice(column.as_primitive::<Date64Type>().values()), + ArrowDataType::Int64 => visitor.slice(column.as_primitive::<Int64Type>().values()), + ArrowDataType::UInt64 => { + // follow C++ implementation and use overflow/reinterpret cast from u64 to i64 which will map + // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0` + let values = column.as_primitive::<UInt64Type>().values(); + visitor.slice(values.inner().typed_data()) + } + ArrowDataType::Time64(TimeUnit::Microsecond) => { + visitor.slice(column.as_primitive::<Time64MicrosecondType>().values()) + } + ArrowDataType::Time64(TimeUnit::Nanosecond) => { + visitor.slice(column.as_primitive::<Time64NanosecondType>().values()) + } + ArrowDataType::Timestamp(unit, _) => match unit { + TimeUnit::Second => { + visitor.slice(column.as_primitive::<TimestampSecondType>().values()) + } + TimeUnit::Millisecond => { + visitor.slice(column.as_primitive::<TimestampMillisecondType>().values()) + } + TimeUnit::Microsecond => { + visitor.slice(column.as_primitive::<TimestampMicrosecondType>().values()) + } + TimeUnit::Nanosecond => { + visitor.slice(column.as_primitive::<TimestampNanosecondType>().values()) + } + }, + ArrowDataType::Duration(unit) => match unit { + TimeUnit::Second => visitor.slice(column.as_primitive::<DurationSecondType>().values()), + TimeUnit::Millisecond => { + visitor.slice(column.as_primitive::<DurationMillisecondType>().values()) + } + TimeUnit::Microsecond => { + visitor.slice(column.as_primitive::<DurationMicrosecondType>().values()) + } + TimeUnit::Nanosecond => { + visitor.slice(column.as_primitive::<DurationNanosecondType>().values()) + } + }, + ArrowDataType::Decimal64(_, _) => { + visitor.slice(column.as_primitive::<Decimal64Type>().values()) } + ArrowDataType::Decimal128(_, _) => { + // use the int64 to represent the decimal with low precision + visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| { + v as i64 + }) + } + ArrowDataType::Decimal256(_, _) => { + // use the int64 to represent the decimal with low precision + visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| { + v.as_i128() as i64 + }) + } + d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))), + } +} + +/// Writes the values passed to it to an `INT32` or `INT64` column +struct WritePrimitive<'a, 'b, E: ColumnValueEncoder> { + writer: &'a mut GenericColumnWriter<'b, E>, + levels: &'a ArrayLevels, +} + +impl<T, E> PrimitiveVisitor<T> for WritePrimitive<'_, '_, E> +where + E: ColumnValueEncoder<Values = [T]>, +{ + type Output = usize; + + fn slice(self, values: &[T]) -> Result<usize> { + write_primitive(self.writer, values, self.levels) + } + + fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<usize> { + let values: Vec<T> = values.iter().map(|v| convert(*v)).collect(); + write_primitive(self.writer, &values, self.levels) + } +} + +/// Receives the values [`ArrowWriter`] stores for an Arrow array in a +/// `FIXED_LEN_BYTE_ARRAY` column, which are also the values hashed into the +/// column's bloom filter +/// +/// [`visit_fixed_len_byte_array`] holds the conversion from each Arrow type +pub(crate) trait FixedLenByteArrayVisitor { + type Output; + + /// Returns the length of the values that decimals of `precision` are stored as + /// + /// `native_length` is the length of the Arrow decimal type, for example 16 + /// for `Decimal128`. + fn decimal_length(&self, precision: u8, native_length: usize) -> usize; + + /// Receives values that are `length` bytes long + /// + /// `fill(i, value)` writes the value at index `i` of the array to `value`, + /// which is `length` zeroed bytes. + fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<Self::Output>; +} + +/// Passes the values [`ArrowWriter`] stores for `column` in a `FIXED_LEN_BYTE_ARRAY` +/// column to `visitor` +pub(crate) fn visit_fixed_len_byte_array<V: FixedLenByteArrayVisitor>( + column: &dyn arrow_array::Array, + visitor: V, +) -> Result<V::Output> { + match column.data_type() { + ArrowDataType::Interval(interval_unit) => match interval_unit { + IntervalUnit::YearMonth => { + // 12 bytes of months, days and milliseconds (4 bytes each). An Arrow + // YearMonth interval only stores months, so only the first 4 bytes are set. + let array = column.as_primitive::<IntervalYearMonthType>(); + visitor.values(12, |i, value| { + value[0..4].copy_from_slice(&array.value(i).to_le_bytes()); + }) + } + IntervalUnit::DayTime => { + // 12 bytes of months, days and milliseconds (4 bytes each). An Arrow + // DayTime interval only stores days and milliseconds, so the first 4 + // bytes are not set. + let array = column.as_primitive::<IntervalDayTimeType>(); + visitor.values(12, |i, value| { + let interval = array.value(i); + value[4..8].copy_from_slice(&interval.days.to_le_bytes()); + value[8..12].copy_from_slice(&interval.milliseconds.to_le_bytes()); + }) + } + IntervalUnit::MonthDayNano => Err(ParquetError::NYI(format!( + "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented" + ))), + }, + ArrowDataType::FixedSizeBinary(_) => { + let array = column.as_fixed_size_binary(); + visitor.values(array.value_size(), |i, value| { + value.copy_from_slice(array.value(i)); + }) + } + ArrowDataType::Decimal32(_, _) => { + visit_decimal(column.as_primitive::<Decimal32Type>(), visitor) + } + ArrowDataType::Decimal64(_, _) => { + visit_decimal(column.as_primitive::<Decimal64Type>(), visitor) + } + ArrowDataType::Decimal128(_, _) => { + visit_decimal(column.as_primitive::<Decimal128Type>(), visitor) + } + ArrowDataType::Decimal256(_, _) => { + visit_decimal(column.as_primitive::<Decimal256Type>(), visitor) + } + ArrowDataType::Float16 => { + let array = column.as_primitive::<Float16Type>(); + visitor.values(2, |i, value| { + value.copy_from_slice(&array.value(i).to_le_bytes()); + }) + } + _ => Err(ParquetError::NYI( + "Attempting to write an Arrow type that is not yet implemented".to_string(), + )), + } +} + +/// Passes the big-endian two's complement bytes of the values of `array` to `visitor` +/// +/// The values are truncated to their low [`FixedLenByteArrayVisitor::decimal_length`] +/// bytes, which must not be more than the length of the Arrow decimal type. A value +/// too large for the length is truncated, as [`ArrowWriter`] has always done, since +/// Arrow does not validate precision. +fn visit_decimal<T: NativeDecimalType, V: FixedLenByteArrayVisitor>( + array: &PrimitiveArray<T>, + visitor: V, +) -> Result<V::Output> { + let length = visitor.decimal_length(array.precision(), T::BYTE_LENGTH); + let be_bytes = |i| T::to_be_bytes(array.value(i)); + if length == T::BYTE_LENGTH { + // Special-case that allows inlining memcpy. + visitor.values(length, |i, value| { + value.copy_from_slice(be_bytes(i).as_ref()); + }) + } else { + visitor.values(length, |i, value| { + value.copy_from_slice(&be_bytes(i).as_ref()[T::BYTE_LENGTH - length..]); + }) + } +} + +/// Writes the values passed to it to a `FIXED_LEN_BYTE_ARRAY` column +struct WriteFixedLenByteArray<'a, 'b> { + writer: &'a mut ColumnWriterImpl<'b, FixedLenByteArrayType>, + levels: &'a ArrayLevels, +} + +impl FixedLenByteArrayVisitor for WriteFixedLenByteArray<'_, '_> { + type Output = usize; + + fn decimal_length(&self, precision: u8, native_length: usize) -> usize { + let length = decimal_length_from_precision(precision); + // The precision of a valid Arrow decimal type fits in its native type + assert!(length <= native_length); + length + } + + fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<usize> { + let indices = self.levels.non_null_indices().iter().copied(); + let values = chunk_array_slice(length, indices, fill); + self.writer.write_batch_internal( + values.as_slice(), + None, + self.levels.def_level_data().as_ref(), + self.levels.rep_level_data().as_ref(), + None, + None, + None, + ) } } @@ -1921,31 +2025,6 @@ fn get_bool_array_slice( values } -/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each). -/// An Arrow YearMonth interval only stores months, thus only the first 4 bytes are populated. -fn get_interval_ym_array_slice( - array: &arrow_array::IntervalYearMonthArray, - indices: impl ExactSizeIterator<Item = usize>, -) -> Vec<FixedLenByteArray> { - chunk_array_slice(12, indices, move |i, chunk| { - let value = array.value(i); - chunk[0..4].copy_from_slice(&value.to_le_bytes()); - }) -} - -/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each). -/// An Arrow DayTime interval only stores days and millis, thus the first 4 bytes are not populated. -fn get_interval_dt_array_slice( - array: &arrow_array::IntervalDayTimeArray, - indices: impl ExactSizeIterator<Item = usize>, -) -> Vec<FixedLenByteArray> { - chunk_array_slice(12, indices, move |i, chunk| { - let value = array.value(i); - chunk[4..8].copy_from_slice(&value.days.to_le_bytes()); - chunk[8..12].copy_from_slice(&value.milliseconds.to_le_bytes()); - }) -} - trait NativeDecimalType: DecimalType { type NativeBytes: AsRef<[u8]>; @@ -1980,48 +2059,6 @@ impl NativeDecimalType for Decimal256Type { } } -fn get_decimal_array_slice<T: NativeDecimalType>( - array: &PrimitiveArray<T>, - indices: impl ExactSizeIterator<Item = usize>, -) -> Vec<FixedLenByteArray> { - let chunk_size = decimal_length_from_precision(array.precision()); - assert!(chunk_size <= T::BYTE_LENGTH); - - if chunk_size == T::BYTE_LENGTH { - // Special-case that allows inlining memcpy. - chunk_array_slice(chunk_size, indices, move |i, chunk| { - let as_be_bytes = T::to_be_bytes(array.value(i)); - chunk.copy_from_slice(as_be_bytes.as_ref()); - }) - } else { - chunk_array_slice(chunk_size, indices, move |i, chunk| { - let as_be_bytes = T::to_be_bytes(array.value(i)); - let resized_value = &as_be_bytes.as_ref()[(T::BYTE_LENGTH - chunk.len())..]; - chunk.copy_from_slice(resized_value); - }) - } -} - -fn get_float_16_array_slice( - array: &arrow_array::Float16Array, - indices: impl ExactSizeIterator<Item = usize>, -) -> Vec<FixedLenByteArray> { - chunk_array_slice(2, indices, move |i, chunk| { - let value = array.value(i).to_le_bytes(); - chunk.copy_from_slice(&value); - }) -} - -fn get_fsb_array_slice( - array: &arrow_array::FixedSizeBinaryArray, - indices: impl ExactSizeIterator<Item = usize>, -) -> Vec<FixedLenByteArray> { - chunk_array_slice(array.value_size(), indices, move |i, chunk| { - let value = array.value(i); - chunk.copy_from_slice(value); - }) -} - #[inline] fn chunk_array_slice( chunk_size: usize, diff --git a/parquet/src/util/interner.rs b/parquet/src/util/interner.rs index f156541c8f..f4e5e170c9 100644 --- a/parquet/src/util/interner.rs +++ b/parquet/src/util/interner.rs @@ -58,6 +58,13 @@ impl<S: Storage> Interner<S> { } /// Intern the value, returning the interned key, and if this was a new value + /// + /// The dictionary encoder calls this once per value, so it and the hash probe + /// of `HashTable::entry` must be inlined into the caller's loop. Without + /// `#[inline]` (and hashbrown's `inline-more` feature) each has a single copy, + /// and whether it is inlined depends on how rustc splits the crate into + /// codegen units, so unrelated changes can turn it into a call per value. + #[inline] pub fn intern(&mut self, value: &S::Value) -> S::Key { let hash = self.state.hash_one(value.as_bytes());
