This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 68f3f54ebc Parquet: move `write_leaf`'s INT32/INT64/FLBA conversions
into shared visitors (#11741)
68f3f54ebc is described below
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());
