diff --git a/cpp/src/arrow/util/alp/alp_codec.cc b/cpp/src/arrow/util/alp/alp_codec.cc index bbdc5d829287..89fd64e1297a 100644 --- a/cpp/src/arrow/util/alp/alp_codec.cc +++ b/cpp/src/arrow/util/alp/alp_codec.cc @@ -153,61 +153,45 @@ Result EncodePage(std::span inputs, const AlpEncodingPreset& p std::span body = output.subspan(AlpHeader::kSize); const int64_t element_count = static_cast(inputs.size()); - // Phase 1: compress all vectors and collect them. - std::vector> encoded_vectors; const int64_t num_vectors = ::arrow::bit_util::CeilDiv(element_count, vector_size); - encoded_vectors.reserve(num_vectors); - - int64_t input_offset = 0; - for (int64_t remaining_elements = element_count; remaining_elements > 0;) { - const int64_t elements_to_encode = - std::min(static_cast(vector_size), remaining_elements); - encoded_vectors.push_back(AlpCompression::Compress( - inputs.subspan(static_cast(input_offset), - static_cast(elements_to_encode)), - preset)); - input_offset += elements_to_encode; - remaining_elements -= elements_to_encode; - } - - // Phase 2: calculate sizes and offsets. const int64_t per_vector_metadata_size = AlpInfo::kStoredSize + GetIntegerEncodingMetadataSize(preset.integer_encoding); const int64_t offsets_section_size = num_vectors * static_cast(sizeof(AlpFormatConstants::OffsetType)); - - std::vector vector_offsets; - vector_offsets.reserve(num_vectors); + if (offsets_section_size > static_cast(body.size())) { + return Status::Invalid("ALP output buffer too small for offsets: ", body.size(), + " < ", offsets_section_size); + } int64_t current_offset = offsets_section_size; - for (const auto& vec : encoded_vectors) { + int64_t input_offset = 0; + for (int64_t vector_index = 0; vector_index < num_vectors; ++vector_index) { + const int64_t remaining_elements = element_count - input_offset; + const int64_t elements_to_encode = + std::min(static_cast(vector_size), remaining_elements); + const auto vec = AlpCompression::Compress( + inputs.subspan(static_cast(input_offset), + static_cast(elements_to_encode)), + preset); + const int64_t stored_size = per_vector_metadata_size + vec.GetDataStoredSize(); if (current_offset > static_cast( std::numeric_limits::max())) { return Status::Invalid("ALP encoded data exceeds the uint32 offset range"); } - vector_offsets.push_back(static_cast(current_offset)); - current_offset += per_vector_metadata_size + vec.GetDataStoredSize(); + if (stored_size > static_cast(body.size()) - current_offset) { + return Status::Invalid("ALP output buffer too small: ", body.size(), " < ", + current_offset + stored_size); + } + const auto offset = static_cast(current_offset); + util::SafeStore(body.data() + vector_index * sizeof(AlpFormatConstants::OffsetType), + bit_util::ToLittleEndian(offset)); + vec.Store(body.subspan(static_cast(current_offset), + static_cast(stored_size))); + current_offset += stored_size; + input_offset += elements_to_encode; } const int64_t body_size = current_offset; - if (body_size > static_cast(body.size())) { - return Status::Invalid("ALP output buffer too small: ", body.size(), " < ", - body_size); - } - - // Phase 3: write the offset table. - uint8_t* offset_ptr = body.data(); - for (const auto& offset : vector_offsets) { - util::SafeStore(offset_ptr, bit_util::ToLittleEndian(offset)); - offset_ptr += sizeof(AlpFormatConstants::OffsetType); - } - - // Phase 4: write interleaved vectors [AlpInfo | ForInfo | Data]. - for (size_t i = 0; i < encoded_vectors.size(); ++i) { - const auto& vec = encoded_vectors[i]; - const int64_t stored_size = per_vector_metadata_size + vec.GetDataStoredSize(); - vec.Store(body.subspan(vector_offsets[i], static_cast(stored_size))); - } AlpHeader header{}; header.compression_mode = static_cast(AlpMode::kAlp); header.integer_encoding = static_cast(AlpIntegerEncoding::kForBitPack); @@ -228,8 +212,11 @@ Result EncodePage(std::span inputs, const AlpEncodingPreset& p template Result AlpCodec::MakePreset(std::span inputs) { RETURN_NOT_OK(ValidateElementCount(inputs.size())); + if (inputs.empty()) { + return AlpEncodingPreset::MakeDefault(); + } - AlpSampler sampler; + AlpSampler sampler(static_cast(inputs.size())); sampler.AddSample(inputs); return sampler.MakePreset(); } diff --git a/cpp/src/arrow/util/alp/alp_compression.cc b/cpp/src/arrow/util/alp/alp_compression.cc index 3ba243c7dd4f..73917ac5ebb7 100644 --- a/cpp/src/arrow/util/alp/alp_compression.cc +++ b/cpp/src/arrow/util/alp/alp_compression.cc @@ -18,6 +18,7 @@ #include "arrow/util/alp/alp_compression_internal.h" #include +#include #include #include #include @@ -58,6 +59,20 @@ class AlpInlines { using EncodedUnsigned = typename Constants::EncodedUnsigned; using EncodedSigned = typename Constants::EncodedSigned; + struct Multipliers { + T encode_exponent; + T encode_factor; + T decode_factor; + T decode_exponent; + }; + + static Multipliers GetMultipliers(const AlpExponentAndFactor exponent_and_factor) { + return {Constants::GetExponent(exponent_and_factor.exponent), + Constants::GetFactor(exponent_and_factor.factor), + Constants::GetExponent(exponent_and_factor.factor), + Constants::GetFactor(exponent_and_factor.exponent)}; + } + // Returns true when the scaled value cannot be encoded. static bool IsImpossibleToEncode(const T n) { // The two limit comparisons already cover both infinities: the encoding @@ -89,22 +104,15 @@ class AlpInlines { return FastRound(n); } - // Encodes one value using the exponent/factor pair. - static EncodedSigned EncodeValue(const T value, - const AlpExponentAndFactor exponent_and_factor) { - const T tmp_encoded_value = value * - Constants::GetExponent(exponent_and_factor.exponent) * - Constants::GetFactor(exponent_and_factor.factor); - return NumberToInt(tmp_encoded_value); + static EncodedSigned EncodeValue(const T value, const Multipliers& multipliers) { + return NumberToInt(value * multipliers.encode_exponent * multipliers.encode_factor); } - // Decodes one encoded integer using the exponent/factor pair. static T DecodeValue(const EncodedSigned encoded_value, - const AlpExponentAndFactor exponent_and_factor) { + const Multipliers& multipliers) { // The cast to T is needed to prevent a signed integer overflow. - return static_cast(encoded_value) * - Constants::GetExponent(exponent_and_factor.factor) * - Constants::GetFactor(exponent_and_factor.exponent); + return static_cast(encoded_value) * multipliers.decode_factor * + multipliers.decode_exponent; } }; @@ -149,6 +157,7 @@ std::optional EstimateCompressedSizeBits( static constexpr int kBatchSize = 8; const size_t n = inputs.size(); int64_t num_exceptions = 0; + const auto multipliers = AlpInlines::GetMultipliers(exponent_and_factor); size_t i = 0; for (; i + kBatchSize <= n; i += kBatchSize) { @@ -157,12 +166,11 @@ std::optional EstimateCompressedSizeBits( // Pass 1: Encode (vectorizable: multiply + magic-number round) for (int j = 0; j < kBatchSize; j++) { - encoded_values[j] = AlpInlines::EncodeValue(inputs[i + j], exponent_and_factor); + encoded_values[j] = AlpInlines::EncodeValue(inputs[i + j], multipliers); } // Pass 2: Decode (vectorizable: int->float cast + multiply) for (int j = 0; j < kBatchSize; j++) { - decoded_values[j] = - AlpInlines::DecodeValue(encoded_values[j], exponent_and_factor); + decoded_values[j] = AlpInlines::DecodeValue(encoded_values[j], multipliers); } // Pass 3: Compare and accumulate min/max/exceptions for (int j = 0; j < kBatchSize; j++) { @@ -177,9 +185,8 @@ std::optional EstimateCompressedSizeBits( // Scalar tail for remaining elements for (; i < n; i++) { const EncodedSigned encoded_value = - AlpInlines::EncodeValue(inputs[i], exponent_and_factor); - const T decoded_value = - AlpInlines::DecodeValue(encoded_value, exponent_and_factor); + AlpInlines::EncodeValue(inputs[i], multipliers); + const T decoded_value = AlpInlines::DecodeValue(encoded_value, multipliers); if (decoded_value == inputs[i]) { max_encoded_value = std::max(encoded_value, max_encoded_value); min_encoded_value = std::min(encoded_value, min_encoded_value); @@ -304,21 +311,6 @@ AlpEncodingPreset AlpCompression::MakePreset( return AlpEncodingPreset{std::move(combinations), estimated_compressed_size_bytes}; } -template -std::vector AlpCompression::CreateSample(std::span inputs) { - // Pick an equidistant stride so the sample count is capped at kSamplesPerChunk. - const int32_t idx_increments = std::max( - 1, static_cast(std::ceil(static_cast(inputs.size()) / - AlpSamplingOptions::kSamplesPerChunk))); - const size_t increment = static_cast(idx_increments); - std::vector sample; - sample.reserve((inputs.size() + increment - 1) / increment); - for (size_t i = 0; i < inputs.size(); i += idx_increments) { - sample.push_back(inputs[i]); - } - return sample; -} - template AlpExponentAndFactor AlpCompression::FindBestExponentAndFactor( std::span inputs, const std::vector& combinations) { @@ -331,7 +323,17 @@ AlpExponentAndFactor AlpCompression::FindBestExponentAndFactor( return combinations.front(); } - const std::vector sample = CreateSample(inputs); + // The sample is capped at 256 values. Keep it on the stack to avoid a heap + // allocation for every vector parameter search. + std::array sample_storage; + const int32_t idx_increments = std::max( + 1, static_cast(std::ceil(static_cast(inputs.size()) / + AlpSamplingOptions::kSamplesPerChunk))); + size_t sample_count = 0; + for (size_t i = 0; i < inputs.size(); i += idx_increments) { + sample_storage[sample_count++] = inputs[i]; + } + const std::span sample(sample_storage.data(), sample_count); AlpExponentAndFactor best_exponent_and_factor; int64_t best_total_bits = std::numeric_limits::max(); @@ -375,6 +377,7 @@ typename AlpCompression::EncodingResult AlpCompression::EncodeVector( for_deltas.reserve(inputs.size()); std::vector exceptions; std::vector exception_positions; + const auto multipliers = AlpInlines::GetMultipliers(exponent_and_factor); EncodedSigned min_encoded_value = std::numeric_limits::max(); EncodedSigned max_encoded_value = std::numeric_limits::min(); @@ -383,10 +386,8 @@ typename AlpCompression::EncodingResult AlpCompression::EncodeVector( // while the remaining values define the FOR range. int64_t input_offset = 0; for (const T input : inputs) { - const EncodedSigned encoded_value = - AlpInlines::EncodeValue(input, exponent_and_factor); - const T decoded_value = - AlpInlines::DecodeValue(encoded_value, exponent_and_factor); + const EncodedSigned encoded_value = AlpInlines::EncodeValue(input, multipliers); + const T decoded_value = AlpInlines::DecodeValue(encoded_value, multipliers); for_deltas.push_back(static_cast(encoded_value)); if (decoded_value != input) { @@ -524,16 +525,25 @@ template void AlpCompression::PatchExceptions( std::span exceptions, std::span exception_positions, - std::span outputs) { + const bool positions_strictly_increasing, std::span outputs) { ARROW_DCHECK_EQ(exceptions.size(), exception_positions.size()); - // `ivdep` assumes the positions are distinct. The encoder writes them in - // order, and loading a page rejects positions that do not increase. - int64_t exception_idx = 0; - ARROW_ALP_UNROLL_AND_ASSUME_INDEPENDENT - for (const AlpFormatConstants::PositionType exception_position : exception_positions) { - ARROW_DCHECK_LT(exception_position, outputs.size()); - outputs[exception_position] = static_cast(exceptions[exception_idx]); - exception_idx++; + const size_t num_exceptions = exception_positions.size(); + const auto patch_exception = [&](size_t i) { + const auto position = exception_positions[i]; + ARROW_DCHECK_LT(position, outputs.size()); + outputs[position] = static_cast(exceptions[i]); + }; + if (positions_strictly_increasing) { + // The encoder emits strictly increasing positions. Use the vectorized + // fast path for that common case. + ARROW_ALP_UNROLL_AND_ASSUME_INDEPENDENT + for (size_t i = 0; i < num_exceptions; ++i) { + patch_exception(i); + } + } else { + for (size_t i = 0; i < num_exceptions; ++i) { + patch_exception(i); + } } } @@ -563,17 +573,18 @@ void AlpCompression::Decompress(const AlpEncodedVectorView& encoded_view, // Read the exponent/factor once. AlpInfo stores them as uint8_t, which may // alias the output stores, so reading them inside the loop blocks vectorization. const AlpExponentAndFactor exponent_and_factor = alp_info.GetExponentAndFactor(); + const auto multipliers = AlpInlines::GetMultipliers(exponent_and_factor); for (size_t i = 0; i < output.size(); ++i) { const EncodedUnsigned unfored_value = scratch[i] + frame_of_ref; const EncodedSigned signed_value = util::SafeCopy(unfored_value); - output[i] = AlpInlines::DecodeValue(signed_value, exponent_and_factor); + output[i] = AlpInlines::DecodeValue(signed_value, multipliers); } PatchExceptions( {encoded_view.exceptions().data(), encoded_view.exceptions().size()}, {encoded_view.exception_positions().data(), encoded_view.exception_positions().size()}, - output); + encoded_view.exception_positions_strictly_increasing(), output); } // Template instantiations diff --git a/cpp/src/arrow/util/alp/alp_compression_internal.h b/cpp/src/arrow/util/alp/alp_compression_internal.h index 99464785f2e2..d437760027fd 100644 --- a/cpp/src/arrow/util/alp/alp_compression_internal.h +++ b/cpp/src/arrow/util/alp/alp_compression_internal.h @@ -79,8 +79,6 @@ class ARROW_EXPORT AlpCompression { std::span integer_scratch); private: - static std::vector CreateSample(std::span inputs); - static AlpExponentAndFactor FindBestExponentAndFactor( std::span inputs, const std::vector& combinations); @@ -113,7 +111,7 @@ class ARROW_EXPORT AlpCompression { static void PatchExceptions( std::span exceptions, std::span exception_positions, - std::span outputs); + bool positions_strictly_increasing, std::span outputs); }; // The marker on a class template does not reach its member templates: each is diff --git a/cpp/src/arrow/util/alp/alp_constants_internal.h b/cpp/src/arrow/util/alp/alp_constants_internal.h index 53e64d38727f..d89c24f823bb 100644 --- a/cpp/src/arrow/util/alp/alp_constants_internal.h +++ b/cpp/src/arrow/util/alp/alp_constants_internal.h @@ -38,9 +38,6 @@ enum class AlpMode : uint8_t { kAlp = 0 }; enum class AlpIntegerEncoding : uint8_t { kForBitPack = 0 }; struct AlpFormatConstants { - /// Recommended default vector size. - static constexpr int32_t kDefaultVectorSize = 1024; - /// Minimum log2(vector size), corresponding to 8 elements. static constexpr uint8_t kMinLogVectorSize = 3; @@ -63,14 +60,14 @@ struct AlpSamplingOptions { /// Input elements per sampling chunk. static constexpr int64_t kChunkSize = 4096; - /// Elements per sampling interval. - static constexpr int64_t kIntervalSize = 122880; + /// Maximum prefix length examined in each sampled chunk. + static constexpr int64_t kSampleWindowSize = 1024; /// Samples collected per chunk. static constexpr int64_t kSamplesPerChunk = 256; - /// Chunks sampled per interval. - static constexpr int64_t kSampleChunksPerInterval = 8; + /// Target number of sampled chunks per window. + static constexpr int64_t kSampleChunksPerWindow = 8; /// Consecutive non-improving combinations before an early exit. static constexpr uint8_t kEarlyExitThreshold = 4; diff --git a/cpp/src/arrow/util/alp/alp_metadata.cc b/cpp/src/arrow/util/alp/alp_metadata.cc index 3d507b3118dd..a7230fa7d3f9 100644 --- a/cpp/src/arrow/util/alp/alp_metadata.cc +++ b/cpp/src/arrow/util/alp/alp_metadata.cc @@ -102,24 +102,23 @@ Status ValidateVectorInfo(const AlpInfo& alp_info, const AlpForInfo& for_info return Status::OK(); } -// Exception positions index the output and the encoder writes them in order, so -// they must be increasing and inside the vector. A repeated position would make -// the patch step write the same slot twice and leave another slot untouched. -Status ValidateExceptionPositions( +// Exception positions index the output and must be inside the vector. The wire +// format does not require positions to be sorted or unique. +Result ValidateExceptionPositions( std::span exception_positions, int32_t num_elements) { + bool strictly_increasing = true; for (size_t i = 0; i < exception_positions.size(); ++i) { const AlpFormatConstants::PositionType position = exception_positions[i]; if (position >= num_elements) { return Status::Invalid("ALP exception position ", position, " is outside a vector of ", num_elements, " elements"); } - if (i > 0 && position <= exception_positions[i - 1]) { - return Status::Invalid("ALP exception positions must increase, but got ", - exception_positions[i - 1], " then ", position); + if (i > 0 && exception_positions[i] <= exception_positions[i - 1]) { + strictly_increasing = false; } } - return Status::OK(); + return strictly_increasing; } template @@ -378,7 +377,8 @@ Status AlpEncodedVectorView::Reset(std::span input_buffer, input_buffer.subspan(static_cast(input_offset), exception_position_size), std::span(exception_positions_)); input_offset += exception_position_size; - RETURN_NOT_OK(ValidateExceptionPositions(exception_positions_, num_elements)); + ARROW_ASSIGN_OR_RAISE(exception_positions_strictly_increasing_, + ValidateExceptionPositions(exception_positions_, num_elements)); exceptions_.resize(alp_info.num_exceptions()); LoadLittleEndianArray( @@ -397,6 +397,7 @@ void AlpEncodedVectorView::Clear() { packed_values_ = {}; exception_positions_.clear(); exceptions_.clear(); + exception_positions_strictly_increasing_ = true; } template class AlpEncodedVectorView; diff --git a/cpp/src/arrow/util/alp/alp_metadata_internal.h b/cpp/src/arrow/util/alp/alp_metadata_internal.h index 13317fa6b883..09fcf5fbaeb2 100644 --- a/cpp/src/arrow/util/alp/alp_metadata_internal.h +++ b/cpp/src/arrow/util/alp/alp_metadata_internal.h @@ -253,6 +253,9 @@ class ARROW_EXPORT AlpEncodedVectorView { const AlpVector& exception_positions() const { return exception_positions_; } + bool exception_positions_strictly_increasing() const { + return exception_positions_strictly_increasing_; + } const AlpVector& exceptions() const { return exceptions_; } @@ -279,6 +282,7 @@ class ARROW_EXPORT AlpEncodedVectorView { std::span packed_values_; AlpVector exception_positions_; AlpVector exceptions_; + bool exception_positions_strictly_increasing_{true}; }; /// Returns the metadata size for one per-vector integer encoding. diff --git a/cpp/src/arrow/util/alp/alp_sampler.cc b/cpp/src/arrow/util/alp/alp_sampler.cc index f1f2ddc9d1b2..007a35e3ab6b 100644 --- a/cpp/src/arrow/util/alp/alp_sampler.cc +++ b/cpp/src/arrow/util/alp/alp_sampler.cc @@ -37,26 +37,19 @@ struct AlpSamplingParameters { int64_t sample_stride; }; -// The interval holds 30 4096-element chunks. Round up the jump so about eight -// chunks per interval are sampled instead of ten under integer truncation. -constexpr int64_t kSampledChunkStride = bit_util::CeilDiv( - bit_util::CeilDiv(AlpSamplingOptions::kIntervalSize, AlpSamplingOptions::kChunkSize), - AlpSamplingOptions::kSampleChunksPerInterval); -static_assert(kSampledChunkStride > 0, "sample jump must be positive"); - // Caps the lookup window and chooses a stride that yields at most one sample per // kSamplesPerChunk values. AlpSamplingParameters GetSamplingParameters(int64_t chunk_size) { const int64_t values_to_examine = - std::min(chunk_size, static_cast(AlpFormatConstants::kDefaultVectorSize)); + std::min(chunk_size, AlpSamplingOptions::kSampleWindowSize); const int64_t sample_stride = std::max( 1, bit_util::CeilDiv(values_to_examine, AlpSamplingOptions::kSamplesPerChunk)); return AlpSamplingParameters{values_to_examine, sample_stride}; } bool ShouldSkipChunk(const int64_t chunk_index, const int64_t sampled_chunk_count, - const int64_t chunk_size) { - if ((chunk_index % kSampledChunkStride) != 0) { + const int64_t chunk_size, const int64_t sampled_chunk_stride) { + if ((chunk_index % sampled_chunk_stride) != 0) { return true; } @@ -67,6 +60,14 @@ bool ShouldSkipChunk(const int64_t chunk_index, const int64_t sampled_chunk_coun } // namespace +template +AlpSampler::AlpSampler(int64_t sampling_window_size) + : sampled_chunk_stride_(bit_util::CeilDiv( + bit_util::CeilDiv(sampling_window_size, AlpSamplingOptions::kChunkSize), + AlpSamplingOptions::kSampleChunksPerWindow)) { + ARROW_DCHECK_GT(sampling_window_size, 0); +} + template void AlpSampler::AddSample(std::span input) { const int64_t input_size = static_cast(input.size()); @@ -83,8 +84,8 @@ void AlpSampler::AddSampleChunk(std::span input) { } const int64_t input_size = static_cast(input.size()); - const bool should_skip = - ShouldSkipChunk(chunks_processed_, chunks_sampled_, input_size); + const bool should_skip = ShouldSkipChunk(chunks_processed_, chunks_sampled_, input_size, + sampled_chunk_stride_); chunks_processed_ += 1; values_processed_ += input_size; diff --git a/cpp/src/arrow/util/alp/alp_sampler_internal.h b/cpp/src/arrow/util/alp/alp_sampler_internal.h index b7e428e965f2..ecbe4cb46cb9 100644 --- a/cpp/src/arrow/util/alp/alp_sampler_internal.h +++ b/cpp/src/arrow/util/alp/alp_sampler_internal.h @@ -34,7 +34,7 @@ namespace arrow::util::alp { template class ARROW_EXPORT AlpSampler { public: - AlpSampler() = default; + explicit AlpSampler(int64_t sampling_window_size); /// Adds input values to the sample set. void AddSample(std::span input); @@ -46,9 +46,11 @@ class ARROW_EXPORT AlpSampler { private: /// Adds one sampling chunk. Only the first - /// AlpFormatConstants::kDefaultVectorSize elements are examined. + /// AlpSamplingOptions::kSampleWindowSize elements are examined. void AddSampleChunk(std::span input); + /// Stride between sampled chunks. + const int64_t sampled_chunk_stride_; /// Number of chunks that have contributed a sample. int64_t chunks_sampled_{0}; /// Number of chunks processed. diff --git a/cpp/src/arrow/util/alp/alp_test.cc b/cpp/src/arrow/util/alp/alp_test.cc index ede13b1082d1..087099e6b059 100644 --- a/cpp/src/arrow/util/alp/alp_test.cc +++ b/cpp/src/arrow/util/alp/alp_test.cc @@ -42,6 +42,8 @@ namespace arrow::util::alp { +constexpr int32_t kVectorSize = 1024; + // Test helpers // Compares two ranges by bit pattern. ALP must be lossless, and `==` cannot show @@ -109,14 +111,12 @@ std::vector EncodeSmallPage(std::vector* input) { for (size_t i = 0; i < input->size(); ++i) { (*input)[i] = static_cast(i) * static_cast(0.1); } - EXPECT_OK_AND_ASSIGN( - int64_t max_comp_size, - AlpCodec::GetMaxCompressedSize(static_cast(input->size()), - AlpFormatConstants::kDefaultVectorSize)); + EXPECT_OK_AND_ASSIGN(int64_t max_comp_size, + AlpCodec::GetMaxCompressedSize( + static_cast(input->size()), kVectorSize)); std::vector comp_buffer(max_comp_size); - EXPECT_OK_AND_ASSIGN( - const int64_t comp_size, - AlpCodec::Encode(*input, AlpFormatConstants::kDefaultVectorSize, comp_buffer)); + EXPECT_OK_AND_ASSIGN(const int64_t comp_size, + AlpCodec::Encode(*input, kVectorSize, comp_buffer)); comp_buffer.resize(comp_size); return comp_buffer; } @@ -167,8 +167,7 @@ TEST(AlpEncodingPresetTest, RejectsInvalidPreset) { SCOPED_TRACE(expected_message); EXPECT_RAISES_WITH_MESSAGE_THAT( Invalid, ::testing::HasSubstr(expected_message), - AlpCodec::Encode(input, preset, AlpFormatConstants::kDefaultVectorSize, - output)); + AlpCodec::Encode(input, preset, kVectorSize, output)); } } // AlpInfo Serialization Tests @@ -757,12 +756,11 @@ class AlpSamplerTest : public ::testing::Test {}; TYPED_TEST_SUITE(AlpSamplerTest, FloatingTestTypes); TYPED_TEST(AlpSamplerTest, PresetGenerationDecimalData) { - AlpSampler sampler; - std::vector data(10000); for (size_t i = 0; i < data.size(); ++i) { data[i] = static_cast(100.0 + i * 0.01); } + AlpSampler sampler(static_cast(data.size())); sampler.AddSample(data); constexpr size_t kNumValues = 1024; @@ -776,11 +774,34 @@ TYPED_TEST(AlpSamplerTest, PresetGenerationDecimalData) { output, std::vector(data.begin(), data.begin() + kNumValues))); } +TYPED_TEST(AlpSamplerTest, SamplingWindowSize) { + constexpr int64_t kChunkSize = AlpSamplingOptions::kChunkSize; + const std::vector nan_chunk(kChunkSize, + std::numeric_limits::quiet_NaN()); + const std::vector decimal_chunk(kChunkSize, static_cast(1.25)); + const auto decimal_preset = AlpCompression::MakePreset({decimal_chunk}); + + // The sampler targets eight chunks per window. An eight-chunk window samples + // both inputs, but only decimal_chunk can contribute usable parameters. + AlpSampler every_chunk(/*sampling_window_size=*/8 * kChunkSize); + every_chunk.AddSample(nan_chunk); + every_chunk.AddSample(decimal_chunk); + EXPECT_EQ(decimal_preset.combinations, every_chunk.MakePreset().combinations); + + // A sixteen-chunk window samples every other chunk, skipping decimal_chunk. + // With only NaNs sampled, the preset falls back to identity. + AlpSampler every_other_chunk(/*sampling_window_size=*/16 * kChunkSize); + every_other_chunk.AddSample(nan_chunk); + every_other_chunk.AddSample(decimal_chunk); + EXPECT_EQ(AlpEncodingPreset::MakeDefault().combinations, + every_other_chunk.MakePreset().combinations); +} + TYPED_TEST(AlpSamplerTest, NoUsableSamplesFallBackToDefault) { const std::vector identity = AlpEncodingPreset::MakeDefault().combinations; - AlpSampler empty_sampler; + AlpSampler empty_sampler(/*sampling_window_size=*/1); EXPECT_EQ(empty_sampler.MakePreset().combinations, identity); const std::vector unusable(8, std::numeric_limits::quiet_NaN()); @@ -860,12 +881,10 @@ TEST(AlpRobustnessTest, TruncatedData) { const int64_t num_elements = static_cast(input.size()); ASSERT_OK_AND_ASSIGN(int64_t max_size, - AlpCodec::GetMaxCompressedSize( - num_elements, AlpFormatConstants::kDefaultVectorSize)); + AlpCodec::GetMaxCompressedSize(num_elements, kVectorSize)); std::vector buffer(max_size); - ASSERT_OK_AND_ASSIGN( - const int64_t comp_size, - AlpCodec::Encode(input, AlpFormatConstants::kDefaultVectorSize, buffer)); + ASSERT_OK_AND_ASSIGN(const int64_t comp_size, + AlpCodec::Encode(input, kVectorSize, buffer)); const auto page = std::span(buffer).first(static_cast(comp_size)); ASSERT_OK(AlpVectorReader::Open(page, arrow::default_memory_pool())); @@ -890,12 +909,10 @@ TEST(AlpRobustnessTest, HeaderElementCountMismatch) { } const int64_t num_elements = static_cast(input.size()); ASSERT_OK_AND_ASSIGN(int64_t max_size, - AlpCodec::GetMaxCompressedSize( - num_elements, AlpFormatConstants::kDefaultVectorSize)); + AlpCodec::GetMaxCompressedSize(num_elements, kVectorSize)); std::vector buffer(max_size); - ASSERT_OK_AND_ASSIGN( - const int64_t comp_size, - AlpCodec::Encode(input, AlpFormatConstants::kDefaultVectorSize, buffer)); + ASSERT_OK_AND_ASSIGN(const int64_t comp_size, + AlpCodec::Encode(input, kVectorSize, buffer)); const auto page = std::span(buffer).first(static_cast(comp_size)); ASSERT_OK_AND_ASSIGN(auto reader, @@ -915,18 +932,16 @@ TEST(AlpRobustnessTest, HeaderElementCountMismatch) { TEST(AlpRobustnessTest, CorruptedOffsetChain) { // The spec fixes every offset, so a duplicate, backward or gapped offset has to // be rejected rather than used to read the wrong bytes. - constexpr int32_t kNumElements = 4 * AlpFormatConstants::kDefaultVectorSize; + constexpr int32_t kNumElements = 4 * kVectorSize; std::vector input(kNumElements); for (size_t i = 0; i < input.size(); ++i) { input[i] = static_cast(i) * 0.123; } ASSERT_OK_AND_ASSIGN(int64_t max_size, - AlpCodec::GetMaxCompressedSize( - kNumElements, AlpFormatConstants::kDefaultVectorSize)); + AlpCodec::GetMaxCompressedSize(kNumElements, kVectorSize)); std::vector buffer(max_size); - ASSERT_OK_AND_ASSIGN( - const int64_t comp_size, - AlpCodec::Encode(input, AlpFormatConstants::kDefaultVectorSize, buffer)); + ASSERT_OK_AND_ASSIGN(const int64_t comp_size, + AlpCodec::Encode(input, kVectorSize, buffer)); const auto page = std::span(buffer).first(static_cast(comp_size)); ASSERT_OK(AlpVectorReader::Open(page, arrow::default_memory_pool())); @@ -971,14 +986,13 @@ TEST(AlpRobustnessTest, ElementCountAboveInt32Max) { double one_value = 1.0; uint8_t output_byte = 0; - EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("exceeds INT32_MAX"), - AlpCodec::GetMaxCompressedSize( - kTooMany, AlpFormatConstants::kDefaultVectorSize)); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("exceeds INT32_MAX"), + AlpCodec::GetMaxCompressedSize(kTooMany, kVectorSize)); // The count is rejected before the input span is read. EXPECT_RAISES_WITH_MESSAGE_THAT( Invalid, ::testing::HasSubstr("exceeds INT32_MAX"), - AlpCodec::Encode({&one_value, static_cast(kTooMany)}, - AlpFormatConstants::kDefaultVectorSize, + AlpCodec::Encode({&one_value, static_cast(kTooMany)}, kVectorSize, {&output_byte, 1})); } @@ -990,14 +1004,12 @@ TYPED_TEST(AlpRobustnessTest, RejectsOutOfRangeLogVectorSizeInHeader) { input[i] = static_cast(i) * static_cast(0.1); } - ASSERT_OK_AND_ASSIGN( - int64_t max_comp_size, - AlpCodec::GetMaxCompressedSize(static_cast(input.size()), - AlpFormatConstants::kDefaultVectorSize)); + ASSERT_OK_AND_ASSIGN(int64_t max_comp_size, + AlpCodec::GetMaxCompressedSize( + static_cast(input.size()), kVectorSize)); std::vector comp_buffer(max_comp_size); ASSERT_OK_AND_ASSIGN(const int64_t comp_size, - AlpCodec::Encode( - input, AlpFormatConstants::kDefaultVectorSize, comp_buffer)); + AlpCodec::Encode(input, kVectorSize, comp_buffer)); for (const uint8_t bad : {uint8_t{0}, uint8_t{1}, uint8_t{2}, uint8_t{16}}) { SCOPED_TRACE("log_vector_size=" + std::to_string(bad)); @@ -1136,43 +1148,23 @@ TYPED_TEST(AlpRobustnessTest, RejectsNumExceptionsAboveVectorLength) { arrow::default_memory_pool())); } -// The patch step writes output[position] once per exception, so repeated -// positions have to be rejected. -TYPED_TEST(AlpRobustnessTest, RejectsDuplicateExceptionPositions) { +// Exception values must stay paired with their positions, even when unsorted. +TYPED_TEST(AlpRobustnessTest, AcceptsUnsortedExceptionPositions) { constexpr int32_t kNumElements = 64; - std::vector input(kNumElements); - for (int32_t i = 0; i < kNumElements; ++i) { - input[i] = static_cast(i); - } - // Whole numbers encode exactly, so only these two NaNs are exceptions. - input[5] = std::numeric_limits::quiet_NaN(); - input[40] = std::numeric_limits::quiet_NaN(); + const auto infinity = std::numeric_limits::infinity(); + const auto encoded = AlpEncodedVector::Make( + AlpInfo(/*exponent=*/0, /*factor=*/0, /*num_exceptions=*/2), + AlpForInfo(/*frame_of_reference=*/0, /*bit_width=*/0), kNumElements, + /*packed_values=*/{}, /*exception_positions=*/{40, 5}, + /*exceptions=*/{-infinity, infinity}); - const auto encoded = - AlpCompression::Compress(input, AlpEncodingPreset::MakeDefault()); - std::vector comp_buffer(static_cast(encoded.GetStoredSize())); - encoded.Store(comp_buffer); + std::vector expected(kNumElements, 0); + expected[40] = -infinity; + expected[5] = infinity; - ASSERT_OK_AND_ASSIGN(const AlpInfo alp_info, AlpInfo::Load(comp_buffer)); - ASSERT_GE(alp_info.num_exceptions(), 2); - - constexpr int64_t kForInfoSize = AlpForInfo::kStoredSize; - const uint8_t bit_width = comp_buffer[AlpInfo::kStoredSize + kForInfoSize - 1]; - const int64_t position_pos = AlpInfo::kStoredSize + kForInfoSize + - bit_util::BytesForBits(int64_t{kNumElements} * bit_width); - constexpr int64_t kPositionSize = - static_cast(sizeof(AlpFormatConstants::PositionType)); - ASSERT_LE(position_pos + 2 * kPositionSize, static_cast(comp_buffer.size())); - - // Copy the second position over the first, so both name the same slot. - std::vector corrupted = comp_buffer; - std::memcpy(corrupted.data() + position_pos, - corrupted.data() + position_pos + kPositionSize, kPositionSize); - - EXPECT_RAISES_WITH_MESSAGE_THAT( - Invalid, ::testing::HasSubstr("exception positions must increase"), - AlpEncodedVectorView::Load({corrupted.data(), corrupted.size()}, - kNumElements, arrow::default_memory_pool())); + std::vector output; + ASSERT_NO_FATAL_FAILURE(DecodeEncodedVector(encoded, &output)); + EXPECT_TRUE(IsBitwiseEqual(output, expected)); } // The patch step writes output[position], so a position at or past the end of diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index c899f5a9847a..bdf50b77c222 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -1293,7 +1293,8 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, : ColumnWriterImpl(metadata, std::move(pager), use_dictionary, encoding, properties, content_defined_chunker) { current_encoder_ = MakeEncoder(ParquetType::type_num, encoding, use_dictionary, - descr_, properties->memory_pool()); + descr_, properties->memory_pool(), + properties->encoding_options(descr_->path())); // We have to dynamic_cast as some compilers don't want to static_cast // through virtual inheritance. current_value_encoder_ = @@ -1876,7 +1877,8 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, fallback_ = true; // Only PLAIN encoding is supported for fallback current_encoder_ = MakeEncoder(ParquetType::type_num, Encoding::PLAIN, false, - descr_, properties_->memory_pool()); + descr_, properties_->memory_pool(), + properties_->encoding_options(descr_->path())); current_value_encoder_ = dynamic_cast(current_encoder_.get()); current_dict_encoder_ = nullptr; // not using dict encoding_ = Encoding::PLAIN; diff --git a/cpp/src/parquet/encoder.cc b/cpp/src/parquet/encoder.cc index 30341bbe25e1..22e7bb146620 100644 --- a/cpp/src/parquet/encoder.cc +++ b/cpp/src/parquet/encoder.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -33,7 +34,7 @@ #include "arrow/stl_allocator.h" #include "arrow/type_traits.h" #include "arrow/util/alp/alp_codec_internal.h" -#include "arrow/util/alp/alp_constants_internal.h" +#include "arrow/util/alp/alp_sampler_internal.h" #include "arrow/util/bit_stream_utils_internal.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" @@ -70,6 +71,30 @@ template using ArrowPoolVector = std::vector>; namespace parquet { + +void AlpOptions::Validate() const { + if (vector_size < 8 || vector_size > 32768 || + !::arrow::bit_util::HasSingleBit(static_cast(vector_size))) { + throw ParquetException("ALP vector size must be a power of two in [8, 32768]"); + } + if (sampling_window_size <= 0) { + throw ParquetException("ALP sampling window size must be positive"); + } +} + +void EncodingOptions::Validate() const { + for (const auto& [encoding, value] : options) { + switch (encoding) { + case Encoding::ALP: + Get(encoding).Validate(); + break; + default: + throw ParquetException("Encoding options are not supported for encoding ", + EncodingToString(encoding)); + } + } +} + namespace { // The Parquet spec isn't very clear whether ByteArray lengths are signed or @@ -1013,40 +1038,56 @@ class AlpEncoder : public EncoderImpl, virtual public TypedEncoder { public: using T = typename DType::c_type; using ArrowType = typename EncodingTraits::ArrowType; + using Codec = ::arrow::util::alp::AlpCodec; + using EncodingPreset = ::arrow::util::alp::AlpEncodingPreset; + using Sampler = ::arrow::util::alp::AlpSampler; using TypedEncoder::Put; static_assert(std::is_same_v || std::is_same_v, "ALP only supports float and double types"); - explicit AlpEncoder(const ColumnDescriptor* descr, - ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) - : EncoderImpl(descr, Encoding::ALP, pool), sink_{pool} {} + explicit AlpEncoder(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, + const AlpOptions& alp_options) + : EncoderImpl(descr, Encoding::ALP, pool), + sink_{pool}, + alp_options_(alp_options), + sampler_(alp_options.sampling_window_size) {} // TODO(GH-48701): use a ratio estimate instead of reporting the raw buffer size int64_t EstimatedDataEncodedSize() override { return sink_.length(); } std::shared_ptr FlushValues() override { - // TODO(GH-48701): allow configuring the vector size. - constexpr auto kVectorSize = - ::arrow::util::alp::AlpFormatConstants::kDefaultVectorSize; - - // An all-null optional page still has to produce a header-only page. const int64_t num_elements = sink_.length() / static_cast(sizeof(T)); + const std::span input{reinterpret_cast(sink_.data()), + static_cast(num_elements)}; + + if (num_elements > 0 && !preset_.has_value()) { + sampler_.AddSample(input); + sampled_values_ += num_elements; + if (sampled_values_ >= alp_options_.sampling_window_size) { + preset_ = sampler_.MakePreset(); + } + } + + EncodingPreset page_preset; + if (!preset_.has_value()) { + // Use a page-specific sampler to avoid affecting the global preset. + Sampler page_sampler(alp_options_.sampling_window_size); + page_sampler.AddSample(input); + page_preset = page_sampler.MakePreset(); + } + const auto& preset = preset_.has_value() ? *preset_ : page_preset; + PARQUET_ASSIGN_OR_THROW( int64_t max_comp_size, - ::arrow::util::alp::AlpCodec::GetMaxCompressedSize(num_elements, kVectorSize)); - + Codec::GetMaxCompressedSize(num_elements, alp_options_.vector_size)); PARQUET_ASSIGN_OR_THROW( - std::shared_ptr compressed_buffer, + auto compressed_buffer, ::arrow::AllocateResizableBuffer(max_comp_size, this->memory_pool())); - - PARQUET_ASSIGN_OR_THROW( - const int64_t compressed_size, - ::arrow::util::alp::AlpCodec::Encode( - {reinterpret_cast(sink_.data()), static_cast(num_elements)}, - kVectorSize, - {compressed_buffer->mutable_data(), static_cast(max_comp_size)})); - + PARQUET_ASSIGN_OR_THROW(const int64_t compressed_size, + Codec::Encode(input, preset, alp_options_.vector_size, + {compressed_buffer->mutable_data(), + static_cast(max_comp_size)})); PARQUET_THROW_NOT_OK(compressed_buffer->Resize(compressed_size)); sink_.Reset(); @@ -1062,13 +1103,23 @@ class AlpEncoder : public EncoderImpl, virtual public TypedEncoder { void PutSpaced(const T* src, int num_values, const uint8_t* valid_bits, int64_t valid_bits_offset) override { + if (num_values == 0) { + return; + } if (valid_bits != NULLPTR) { - PARQUET_ASSIGN_OR_THROW(auto buffer, ::arrow::AllocateBuffer(num_values * sizeof(T), - this->memory_pool())); + const int num_valid_values = static_cast(::arrow::internal::CountSetBits( + valid_bits, valid_bits_offset, static_cast(num_values))); + if (num_valid_values == 0) { + return; + } + PARQUET_ASSIGN_OR_THROW( + auto buffer, + ::arrow::AllocateBuffer(num_valid_values * sizeof(T), this->memory_pool())); T* data = buffer->template mutable_data_as(); - const int num_valid_values = ::arrow::util::internal::SpacedCompress( + const int compacted_values = ::arrow::util::internal::SpacedCompress( src, num_values, valid_bits, valid_bits_offset, data); - Put(data, num_valid_values); + ARROW_DCHECK_EQ(compacted_values, num_valid_values); + Put(data, compacted_values); } else { Put(src, num_values); } @@ -1087,6 +1138,12 @@ class AlpEncoder : public EncoderImpl, virtual public TypedEncoder { private: ::arrow::BufferBuilder sink_; + const AlpOptions alp_options_; + // Sampling state spans pages in one column chunk. + Sampler sampler_; + int64_t sampled_values_{0}; + // Reuse the sampled exponent/factor pairs across pages. + std::optional preset_; }; // ---------------------------------------------------------------------- @@ -1852,7 +1909,7 @@ std::shared_ptr RleBooleanEncoder::FlushValues() { std::unique_ptr MakeEncoder(Type::type type_num, Encoding::type encoding, bool use_dictionary, const ColumnDescriptor* descr, - MemoryPool* pool) { + MemoryPool* pool, const EncodingOptions& options) { if (use_dictionary) { switch (type_num) { case Type::INT32: @@ -1913,11 +1970,12 @@ std::unique_ptr MakeEncoder(Type::type type_num, Encoding::type encodin "and FIXED_LEN_BYTE_ARRAY"); } } else if (encoding == Encoding::ALP) { + const auto& alp_options = options.Get(encoding); switch (type_num) { case Type::FLOAT: - return std::make_unique>(descr, pool); + return std::make_unique>(descr, pool, alp_options); case Type::DOUBLE: - return std::make_unique>(descr, pool); + return std::make_unique>(descr, pool, alp_options); default: throw ParquetException("ALP encoding only supports FLOAT and DOUBLE"); } diff --git a/cpp/src/parquet/encoding.h b/cpp/src/parquet/encoding.h index 157d67264c56..d9bc858be31d 100644 --- a/cpp/src/parquet/encoding.h +++ b/cpp/src/parquet/encoding.h @@ -17,9 +17,11 @@ #pragma once +#include #include #include #include +#include #include #include "arrow/type_fwd.h" @@ -35,6 +37,46 @@ class Dictionary32Builder; namespace parquet { +/// Configuration for ALP encoding. +struct PARQUET_EXPORT AlpOptions { + static constexpr int32_t kDefaultVectorSize = 1024; + static constexpr int64_t kDefaultSamplingWindowSize = 122880; + + /// Values per encoded vector. Must be a power of two in [8, 32768]. + int32_t vector_size{kDefaultVectorSize}; + /// Target number of non-null input values examined before caching the preset. + /// Must be positive. + int64_t sampling_window_size{kDefaultSamplingWindowSize}; + + /// Throws ParquetException if a configuration value is invalid. + void Validate() const; +}; + +/// Options for the encodings that a column writer or encoder may use. +struct PARQUET_EXPORT EncodingOptions { + /// Encoding-specific options. An ALP entry must contain AlpOptions. + std::unordered_map options; + + /// Return options for an encoding, or default value if absent. + template + const T& Get(Encoding::type encoding) const { + const auto it = options.find(encoding); + if (it == options.end()) { + static const T default_options{}; + return default_options; + } + const auto* encoding_options = std::any_cast(&it->second); + if (!encoding_options) { + throw ParquetException("Invalid options type for encoding ", + EncodingToString(encoding)); + } + return *encoding_options; + } + + /// Throws ParquetException for invalid or unsupported options. + void Validate() const; +}; + template class TypedEncoder; @@ -433,16 +475,18 @@ PARQUET_EXPORT std::unique_ptr MakeEncoder( Type::type type_num, Encoding::type encoding, bool use_dictionary = false, const ColumnDescriptor* descr = NULLPTR, - ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool(), + const EncodingOptions& options = {}); template std::unique_ptr::Encoder> MakeTypedEncoder( Encoding::type encoding, bool use_dictionary = false, const ColumnDescriptor* descr = NULLPTR, - ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) { + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool(), + const EncodingOptions& options = {}) { using OutType = typename EncodingTraits::Encoder; std::unique_ptr base = - MakeEncoder(DType::type_num, encoding, use_dictionary, descr, pool); + MakeEncoder(DType::type_num, encoding, use_dictionary, descr, pool, options); return std::unique_ptr(dynamic_cast(base.release())); } diff --git a/cpp/src/parquet/encoding_test.cc b/cpp/src/parquet/encoding_test.cc index 672e55a900dc..0bfc480312bf 100644 --- a/cpp/src/parquet/encoding_test.cc +++ b/cpp/src/parquet/encoding_test.cc @@ -3002,28 +3002,43 @@ TYPED_TEST(TestAlpEncoding, BatchedDecode) { } } -// Flush twice on one encoder. The two pages must not mix. TYPED_TEST(TestAlpEncoding, ReusesEncoderAcrossFlushes) { using c_type = typename TypeParam::c_type; - auto encoder = MakeTypedEncoder(Encoding::ALP, false, this->descr_.get()); - auto decoder = MakeTypedDecoder(Encoding::ALP, this->descr_.get()); - const std::vector> pages = { - {static_cast(1.25), static_cast(2.5), static_cast(3.75)}, - // Two vectors, so a bad offset table would be caught. - std::vector(1500, static_cast(7.5)), - }; - - for (const auto& page : pages) { - encoder->Put(page.data(), static_cast(page.size())); - auto buffer = encoder->FlushValues(); + for (const AlpOptions& alp_options : + {AlpOptions{}, AlpOptions{/*vector_size=*/8, /*sampling_window_size=*/16}, + AlpOptions{/*vector_size=*/4096, /*sampling_window_size=*/8192}}) { + SCOPED_TRACE(::testing::Message() + << "vector_size=" << alp_options.vector_size + << " sampling_window_size=" << alp_options.sampling_window_size); + EncodingOptions options; + options.options.emplace(Encoding::ALP, alp_options); + auto encoder = MakeTypedEncoder( + /*encoding=*/Encoding::ALP, /*use_dictionary=*/false, this->descr_.get(), + ::arrow::default_memory_pool(), options); + auto decoder = MakeTypedDecoder(Encoding::ALP, this->descr_.get()); - std::vector output(page.size()); - decoder->SetData(static_cast(page.size()), buffer->data(), - static_cast(buffer->size())); - ASSERT_EQ(page.size(), static_cast(decoder->Decode( - output.data(), static_cast(page.size())))); - ASSERT_THAT(output, ::testing::ElementsAreArray(page)); + for (c_type value : {static_cast(1.25), static_cast(7.5), + static_cast(0.001)}) { + SCOPED_TRACE(value); + // Each page contains one full vector and one trailing value. + const std::vector page(alp_options.vector_size + 1, value); + encoder->Put(page.data(), static_cast(page.size())); + auto buffer = encoder->FlushValues(); + + ASSERT_OK_AND_ASSIGN(auto reader, + ::arrow::util::alp::AlpVectorReader::Open( + {buffer->data(), static_cast(buffer->size())}, + ::arrow::default_memory_pool())); + EXPECT_EQ(alp_options.vector_size, reader.vector_size()); + + std::vector output(page.size()); + decoder->SetData(static_cast(page.size()), buffer->data(), + static_cast(buffer->size())); + ASSERT_EQ(page.size(), static_cast(decoder->Decode( + output.data(), static_cast(page.size())))); + EXPECT_EQ(0, std::memcmp(output.data(), page.data(), page.size() * sizeof(c_type))); + } } } diff --git a/cpp/src/parquet/properties.cc b/cpp/src/parquet/properties.cc index 6e3cfddd7bb8..fbfa92541a1e 100644 --- a/cpp/src/parquet/properties.cc +++ b/cpp/src/parquet/properties.cc @@ -61,6 +61,24 @@ std::shared_ptr default_arrow_writer_properties() { return default_writer_properties; } +void ColumnProperties::set_encoding_options(const EncodingOptions& options) { + options.Validate(); + encoding_options_ = options; +} + +WriterProperties::Builder* WriterProperties::Builder::encoding_options( + const EncodingOptions& options) { + default_column_properties_.set_encoding_options(options); + return this; +} + +WriterProperties::Builder* WriterProperties::Builder::encoding_options( + const std::string& path, const EncodingOptions& options) { + options.Validate(); + encoding_options_[path] = options; + return this; +} + void WriterProperties::Builder::CopyColumnSpecificProperties( const WriterProperties& properties) { for (const auto& [col_path, col_props] : properties.column_properties_) { @@ -106,6 +124,10 @@ void WriterProperties::Builder::CopyColumnSpecificProperties( if (col_props.bloom_filter_options().has_value()) { this->enable_bloom_filter(col_path, col_props.bloom_filter_options().value()); } + + if (!col_props.encoding_options().options.empty()) { + this->encoding_options(col_path, col_props.encoding_options()); + } } } diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index 251448213896..15b6a809e86c 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -29,6 +29,7 @@ #include "arrow/type_fwd.h" #include "arrow/util/compression.h" #include "arrow/util/type_fwd.h" +#include "parquet/encoding.h" #include "parquet/encryption/encryption.h" #include "parquet/exception.h" #include "parquet/parquet_version.h" @@ -257,6 +258,8 @@ class PARQUET_EXPORT ColumnProperties { void set_encoding(Encoding::type encoding) { encoding_ = encoding; } + void set_encoding_options(const EncodingOptions& options); + void set_compression(Compression::type codec) { codec_ = codec; } void set_dictionary_enabled(bool dictionary_enabled) { @@ -301,6 +304,8 @@ class PARQUET_EXPORT ColumnProperties { Encoding::type encoding() const { return encoding_; } + const EncodingOptions& encoding_options() const { return encoding_options_; } + Compression::type compression() const { return codec_; } bool dictionary_enabled() const { return dictionary_enabled_; } @@ -335,6 +340,7 @@ class PARQUET_EXPORT ColumnProperties { std::shared_ptr codec_options_; bool page_index_enabled_; std::optional bloom_filter_options_; + EncodingOptions encoding_options_; }; // EXPERIMENTAL: Options for content-defined chunking. @@ -517,6 +523,18 @@ class PARQUET_EXPORT WriterProperties { return this; } + /// Set the encoding-specific options for all columns. + Builder* encoding_options(const EncodingOptions& options); + + /// Set encoding-specific options for column specified by `path`. + Builder* encoding_options(const std::string& path, const EncodingOptions& options); + + /// Set encoding-specific options for column specified by `path`. + Builder* encoding_options(const std::shared_ptr& path, + const EncodingOptions& options) { + return encoding_options(path->ToDotString(), options); + } + /// Specify the maximum number of rows per data page. /// Default 20K rows. Builder* max_rows_per_page(int64_t max_rows) { @@ -914,6 +932,9 @@ class PARQUET_EXPORT WriterProperties { get(item.first).set_codec_options(item.second); for (const auto& item : dictionary_enabled_) get(item.first).set_dictionary_enabled(item.second); + for (const auto& [path, options] : encoding_options_) { + get(path).set_encoding_options(options); + } for (const auto& item : statistics_enabled_) get(item.first).set_statistics_enabled(item.second); for (const auto& item : page_index_enabled_) @@ -970,6 +991,7 @@ class PARQUET_EXPORT WriterProperties { std::unordered_map statistics_enabled_; std::unordered_map page_index_enabled_; std::unordered_map bloom_filter_options_; + std::unordered_map encoding_options_; bool content_defined_chunking_enabled_; CdcOptions content_defined_chunking_options_; @@ -1028,6 +1050,15 @@ class PARQUET_EXPORT WriterProperties { } } + const EncodingOptions& encoding_options() const { + return default_column_properties_.encoding_options(); + } + + const EncodingOptions& encoding_options( + const std::shared_ptr& path) const { + return column_properties(path).encoding_options(); + } + const ColumnProperties& column_properties( const std::shared_ptr& path) const { auto it = column_properties_.find(path->ToDotString()); diff --git a/cpp/src/parquet/properties_test.cc b/cpp/src/parquet/properties_test.cc index 0dfebaf7a4f5..c9d2bf1eb1b9 100644 --- a/cpp/src/parquet/properties_test.cc +++ b/cpp/src/parquet/properties_test.cc @@ -71,6 +71,40 @@ TEST(TestWriterProperties, Basics) { ASSERT_FALSE(props->page_checksum_enabled()); } +TEST(TestWriterProperties, AlpOptions) { + WriterProperties::Builder builder; + const AlpOptions alp_options{/*vector_size=*/4096, /*sampling_window_size=*/8192}; + EncodingOptions options; + options.options.emplace(Encoding::ALP, alp_options); + auto properties = builder.encoding_options(options)->build(); + const auto& stored = properties->encoding_options().Get(Encoding::ALP); + EXPECT_EQ(alp_options.vector_size, stored.vector_size); + EXPECT_EQ(alp_options.sampling_window_size, stored.sampling_window_size); + + auto copy = WriterProperties::Builder(*properties).build(); + const auto& copied = copy->encoding_options().Get(Encoding::ALP); + EXPECT_EQ(alp_options.vector_size, copied.vector_size); + EXPECT_EQ(alp_options.sampling_window_size, copied.sampling_window_size); + + for (int32_t invalid_size : {0, 7, 1000, 65536}) { + EncodingOptions invalid_options; + invalid_options.options.emplace( + Encoding::ALP, AlpOptions{invalid_size, alp_options.sampling_window_size}); + EXPECT_THROW(builder.encoding_options(invalid_options), ParquetException); + } + for (int64_t invalid_size : {0, -1}) { + EncodingOptions invalid_options; + invalid_options.options.emplace(Encoding::ALP, + AlpOptions{alp_options.vector_size, invalid_size}); + EXPECT_THROW(builder.encoding_options(invalid_options), ParquetException); + } + for (const std::any& value : {std::any{}, std::any{1024}}) { + EncodingOptions invalid_options; + invalid_options.options.emplace(Encoding::ALP, value); + EXPECT_THROW(builder.encoding_options(invalid_options), ParquetException); + } +} + TEST(TestWriterProperties, DefaultCompression) { std::shared_ptr props = WriterProperties::Builder().build();