From dea698435de5adda34c2a38677ebc4f06b2d481a Mon Sep 17 00:00:00 2001 From: bunny <1065339646@qq.com> Date: Thu, 16 Jul 2026 17:44:38 +0800 Subject: [PATCH 1/2] fix(parquet): prevent batch-induced oversized BYTE_ARRAY pages --- bolt/dwio/parquet/arrow/ColumnWriter.cpp | 334 ++++++++++++-- bolt/dwio/parquet/arrow/Encoding.cpp | 14 +- bolt/dwio/parquet/arrow/Encoding.h | 2 +- .../parquet/arrow/tests/ColumnWriterTest.cpp | 79 ++++ .../tests/writer/ParquetWriterBenchmark.cpp | 271 ++++++++++-- .../tests/writer/ParquetWriterTest.cpp | 415 ++++++++++++++++++ 6 files changed, 1054 insertions(+), 61 deletions(-) diff --git a/bolt/dwio/parquet/arrow/ColumnWriter.cpp b/bolt/dwio/parquet/arrow/ColumnWriter.cpp index 919605dd5..6eb26bba1 100644 --- a/bolt/dwio/parquet/arrow/ColumnWriter.cpp +++ b/bolt/dwio/parquet/arrow/ColumnWriter.cpp @@ -36,8 +36,11 @@ #include #include #include +#include #include #include +#include +#include #include #include @@ -96,6 +99,19 @@ using util::CodecOptions; namespace { +constexpr int64_t kMaxPageHeaderSize = std::numeric_limits::max(); + +int32_t checkPageHeaderSize(std::string_view size_name, int64_t size) { + if (size < 0 || size > kMaxPageHeaderSize) { + throw ParquetException( + std::string(size_name), + " page size cannot be represented in a Parquet PageHeader int32 " + "field: ", + size); + } + return static_cast(size); +} + int64_t bufferAllocatedBytes(const std::shared_ptr& buffer) { return buffer != nullptr ? buffer->capacity() : 0; } @@ -410,7 +426,9 @@ class SerializedPageWriter : public PageWriter { int64_t WriteDictionaryPage( std::unique_ptr pagePtr) override { const auto& page = *pagePtr; - int64_t uncompressed_size = page.size(); + const int64_t uncompressed_size = page.buffer()->size(); + const int32_t uncompressed_page_size = + checkPageHeaderSize("Uncompressed dictionary", uncompressed_size); std::shared_ptr compressed_data; if (has_compressor()) { auto buffer = std::static_pointer_cast( @@ -427,25 +445,28 @@ class SerializedPageWriter : public PageWriter { dict_page_header.__set_is_sorted(page.is_sorted()); const uint8_t* output_data_buffer = compressed_data->data(); - int32_t output_data_len = static_cast(compressed_data->size()); + int64_t output_data_len = compressed_data->size(); + int32_t compressed_page_size = + checkPageHeaderSize("Compressed dictionary", output_data_len); if (data_encryptor_.get()) { UpdateEncryption(encryption::kDictionaryPage); PARQUET_THROW_NOT_OK(encryption_buffer_->Resize( - data_encryptor_->CiphertextSizeDelta() + output_data_len, false)); + data_encryptor_->CiphertextSizeDelta() + compressed_page_size, + false)); output_data_len = data_encryptor_->Encrypt( compressed_data->data(), - output_data_len, + compressed_page_size, encryption_buffer_->mutable_data()); output_data_buffer = encryption_buffer_->data(); + compressed_page_size = + checkPageHeaderSize("Encrypted dictionary", output_data_len); } format::PageHeader page_header; page_header.__set_type(format::PageType::DICTIONARY_PAGE); - page_header.__set_uncompressed_page_size( - static_cast(uncompressed_size)); - page_header.__set_compressed_page_size( - static_cast(output_data_len)); + page_header.__set_uncompressed_page_size(uncompressed_page_size); + page_header.__set_compressed_page_size(compressed_page_size); page_header.__set_dictionary_page_header(dict_page_header); if (page_checksum_verification_) { uint32_t crc32 = @@ -531,26 +552,31 @@ class SerializedPageWriter : public PageWriter { int64_t WriteDataPage(std::unique_ptr pagePtr) override { const DataPage& page = *pagePtr; const int64_t uncompressed_size = page.uncompressed_size(); + const int32_t uncompressed_page_size = + checkPageHeaderSize("Uncompressed data", uncompressed_size); std::shared_ptr compressed_data = page.buffer(); const uint8_t* output_data_buffer = compressed_data->data(); - int32_t output_data_len = static_cast(compressed_data->size()); + int64_t output_data_len = compressed_data->size(); + int32_t compressed_page_size = + checkPageHeaderSize("Compressed data", output_data_len); if (data_encryptor_.get()) { PARQUET_THROW_NOT_OK(encryption_buffer_->Resize( - data_encryptor_->CiphertextSizeDelta() + output_data_len, false)); + data_encryptor_->CiphertextSizeDelta() + compressed_page_size, + false)); UpdateEncryption(encryption::kDataPage); output_data_len = data_encryptor_->Encrypt( compressed_data->data(), - output_data_len, + compressed_page_size, encryption_buffer_->mutable_data()); output_data_buffer = encryption_buffer_->data(); + compressed_page_size = + checkPageHeaderSize("Encrypted data", output_data_len); } format::PageHeader page_header; - page_header.__set_uncompressed_page_size( - static_cast(uncompressed_size)); - page_header.__set_compressed_page_size( - static_cast(output_data_len)); + page_header.__set_uncompressed_page_size(uncompressed_page_size); + page_header.__set_compressed_page_size(compressed_page_size); if (page_checksum_verification_) { uint32_t crc32 = @@ -2216,8 +2242,10 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, return; } + const int64_t dictionary_page_byte_limit = + std::clamp(dictionary_pagesize_limit_, 0, kMaxPageHeaderSize); if (current_dict_encoder_->dict_encoded_size() >= - dictionary_pagesize_limit_) { + dictionary_page_byte_limit) { FallbackToPlainEncoding(); } } @@ -2976,18 +3004,109 @@ Status TypedColumnWriterImpl::WriteArrowDense( ARROW_UNSUPPORTED(); } - int64_t value_offset = 0; - auto WriteChunk = [&](int64_t offset, int64_t batch_size, bool check_page) { - int64_t batch_num_values = 0; - int64_t batch_num_spaced_values = 0; - int64_t null_count = 0; + // The configured page sizes are soft flush targets. PageHeader fields are + // the hard format boundary, so cap byte-budget arithmetic at INT32_MAX and + // leave the final encoded/compressed size checks to PageWriter. + auto PageByteLimit = [](int64_t configured_limit) { + return std::clamp(configured_limit, 0, kMaxPageHeaderSize); + }; + const int64_t data_page_byte_limit = PageByteLimit(data_pagesize_); + const int64_t dictionary_page_byte_limit = + PageByteLimit(dictionary_pagesize_limit_); - MaybeCalculateValidityBits( - AddIfNotNull(def_levels, offset), - batch_size, - &batch_num_values, - &batch_num_spaced_values, - &null_count); + // We only need to know whether an encoded byte count fits in a given budget. + // Saturating at limit + 1 avoids overflowing int64_t for malformed or + // exceptionally large LargeBinary offsets. + auto CappedPageBytes = [](int64_t first, int64_t second, int64_t limit) { + if (first < 0 || second < 0) { + throw ParquetException("Invalid negative BYTE_ARRAY encoded size"); + } + if (first > limit || second > limit - first) { + return limit + 1; + } + return first + second; + }; + + auto ValueLength = [&](int64_t index) { + if (::arrow::is_binary_like(array.type_id())) { + return static_cast( + checked_cast(array).value_length(index)); + } + DCHECK(::arrow::is_large_binary_like(array.type_id())); + return static_cast( + checked_cast(array).value_length( + index)); + }; + + auto ValueRangeByteLength = [&](int64_t start, int64_t count) { + if (::arrow::is_binary_like(array.type_id())) { + const auto& binary_array = + checked_cast(array); + return static_cast( + binary_array.value_offset(start + count) - + binary_array.value_offset(start)); + } + DCHECK(::arrow::is_large_binary_like(array.type_id())); + const auto& large_binary_array = + checked_cast(array); + return static_cast( + large_binary_array.value_offset(start + count) - + large_binary_array.value_offset(start)); + }; + + auto HasSpacedValue = [&](int64_t level_index) { + if (def_levels == nullptr || level_info_.def_level == 0) { + return true; + } + return def_levels[level_index] >= level_info_.repeated_ancestor_def_level; + }; + + auto HasNonNullValue = [&](int64_t level_index, int64_t value_index) { + if (def_levels != nullptr && + def_levels[level_index] != level_info_.def_level) { + return false; + } + return array.IsValid(value_index); + }; + + auto PlainEncodedRangeSize = [&](int64_t start, + int64_t spaced_values, + int64_t non_null_values, + int64_t byte_limit) { + const int64_t length_prefix_bytes = + non_null_values > byte_limit / sizeof(uint32_t) + ? byte_limit + 1 + : non_null_values * static_cast(sizeof(uint32_t)); + return CappedPageBytes( + ValueRangeByteLength(start, spaced_values), + length_prefix_bytes, + byte_limit); + }; + + auto IsCurrentDictionaryEncoding = [&]() { + return IsDictionaryEncoding(current_encoder_->encoding()); + }; + + auto CurrentPageByteLimit = [&]() { + return IsCurrentDictionaryEncoding() ? dictionary_page_byte_limit + : data_page_byte_limit; + }; + + auto CurrentPageBytes = [&]() { + if (IsCurrentDictionaryEncoding()) { + DCHECK_NE(current_dict_encoder_, nullptr); + return current_dict_encoder_->dict_encoded_size(); + } + return current_encoder_->EstimatedDataEncodedSize(); + }; + + int64_t value_offset = 0; + auto WritePreparedChunk = [&](int64_t offset, + int64_t batch_size, + int64_t batch_num_values, + int64_t batch_num_spaced_values, + int64_t null_count, + bool check_page) { WriteLevelsSpaced( batch_size, AddIfNotNull(def_levels, offset), @@ -3012,6 +3131,167 @@ Status TypedColumnWriterImpl::WriteArrowDense( value_offset += batch_num_spaced_values; }; + auto PrepareAndWriteChunk = + [&](int64_t offset, int64_t batch_size, bool check_page) { + int64_t batch_num_values = 0; + int64_t batch_num_spaced_values = 0; + int64_t null_count = 0; + MaybeCalculateValidityBits( + AddIfNotNull(def_levels, offset), + batch_size, + &batch_num_values, + &batch_num_spaced_values, + &null_count); + WritePreparedChunk( + offset, + batch_size, + batch_num_values, + batch_num_spaced_values, + null_count, + check_page); + }; + + auto WriteChunk = [&](int64_t offset, int64_t batch_size, bool check_page) { + // Calculate the metadata that the old write path needs before deciding + // whether this chunk requires splitting. This makes the no-split path use + // the same single level scan as before, including for nullable and nested + // columns. + int64_t batch_num_values = 0; + int64_t batch_num_spaced_values = 0; + int64_t null_count = 0; + MaybeCalculateValidityBits( + AddIfNotNull(def_levels, offset), + batch_size, + &batch_num_values, + &batch_num_spaced_values, + &null_count); + + // For PLAIN encoding this is the data-page value size. During dictionary + // encoding it is a conservative upper bound on dictionary growth: duplicate + // values consume no additional dictionary bytes, but treating them as new + // entries keeps a single write batch from causing a large peak allocation. + // The offset range is exact for canonical Arrow binary arrays and a + // conservative upper bound if a null slot retains bytes. + const int64_t page_byte_limit = CurrentPageByteLimit(); + const int64_t batch_encoded_bytes = PlainEncodedRangeSize( + value_offset, + batch_num_spaced_values, + batch_num_values, + page_byte_limit); + if (CappedPageBytes( + CurrentPageBytes(), batch_encoded_bytes, page_byte_limit) <= + page_byte_limit) { + WritePreparedChunk( + offset, + batch_size, + batch_num_values, + batch_num_spaced_values, + null_count, + check_page); + return; + } + + const bool dictionary_encoding = IsCurrentDictionaryEncoding(); + + // A false check_page means this chunk is the final, potentially partial + // nested record. It cannot be split or flushed after writing, but its + // beginning is a known record boundary, so flush an existing page before + // appending an oversized record to it. + if (!dictionary_encoding && !check_page) { + if (num_buffered_values_ > 0) { + AddDataPage(); + } + WritePreparedChunk( + offset, + batch_size, + batch_num_values, + batch_num_spaced_values, + null_count, + check_page); + return; + } + + // Prefer starting a new page over scanning levels merely to fill the + // remainder of the current page. If the whole chunk fits on an empty page, + // this preserves the old one-pass path for nullable and nested columns. + if (!dictionary_encoding && num_buffered_values_ > 0 && + batch_encoded_bytes <= data_page_byte_limit) { + AddDataPage(); + WritePreparedChunk( + offset, + batch_size, + batch_num_values, + batch_num_spaced_values, + null_count, + check_page); + return; + } + + // The chunk itself is larger than a page. Only this actual split path + // inspects individual levels. The PrepareAndWriteChunk lambda recalculates + // validity for each output segment because bits_buffer_ must start at that + // segment's offset. + int64_t local_offset = offset; + int64_t remaining = batch_size; + while (remaining > 0) { + const bool current_dictionary_encoding = IsCurrentDictionaryEncoding(); + const int64_t current_page_byte_limit = CurrentPageByteLimit(); + int64_t current_page_bytes = CurrentPageBytes(); + int64_t subchunk_levels = 0; + int64_t subchunk_spaced_values = 0; + int64_t subchunk_encoded_bytes = 0; + + while (subchunk_levels < remaining) { + const int64_t level_index = local_offset + subchunk_levels; + const bool can_break_before_level = + !pages_change_on_record_boundaries() || rep_levels == nullptr || + rep_levels[level_index] == 0; + int64_t value_bytes = 0; + if (HasSpacedValue(level_index)) { + const int64_t value_index = value_offset + subchunk_spaced_values; + if (HasNonNullValue(level_index, value_index)) { + value_bytes = CappedPageBytes( + static_cast(sizeof(uint32_t)), + ValueLength(value_index), + current_page_byte_limit); + } + } + + if (can_break_before_level && + CappedPageBytes( + CappedPageBytes( + current_page_bytes, + subchunk_encoded_bytes, + current_page_byte_limit), + value_bytes, + current_page_byte_limit) > current_page_byte_limit) { + if (subchunk_levels == 0) { + if (!current_dictionary_encoding && num_buffered_values_ > 0) { + AddDataPage(); + current_page_bytes = 0; + } + } else { + break; + } + } + + if (HasSpacedValue(level_index)) { + ++subchunk_spaced_values; + } + subchunk_encoded_bytes = CappedPageBytes( + subchunk_encoded_bytes, value_bytes, current_page_byte_limit); + ++subchunk_levels; + } + + if (subchunk_levels == 0) { + throw ParquetException("Unable to split BYTE_ARRAY write chunk"); + } + PrepareAndWriteChunk(local_offset, subchunk_levels, check_page); + local_offset += subchunk_levels; + remaining -= subchunk_levels; + } + }; + PARQUET_CATCH_NOT_OK(DoInBatches( def_levels, rep_levels, diff --git a/bolt/dwio/parquet/arrow/Encoding.cpp b/bolt/dwio/parquet/arrow/Encoding.cpp index be2ebd0fa..3a9a56799 100644 --- a/bolt/dwio/parquet/arrow/Encoding.cpp +++ b/bolt/dwio/parquet/arrow/Encoding.cpp @@ -636,7 +636,7 @@ class DictEncoderImpl : public EncoderImpl, virtual public DictEncoder { DCHECK(buffered_indices_.empty()); } - int dict_encoded_size() const override { + int64_t dict_encoded_size() const override { return dict_encoded_size_; } @@ -803,7 +803,7 @@ class DictEncoderImpl : public EncoderImpl, virtual public DictEncoder { throw ParquetException( "Parquet cannot store strings with size 2GB or more"); } - dict_encoded_size_ += static_cast(v.size() + sizeof(uint32_t)); + dict_encoded_size_ += static_cast(v.size()) + sizeof(uint32_t); int32_t unused_memo_index; PARQUET_THROW_NOT_OK(memo_table_.GetOrInsert( v.data(), static_cast(v.size()), &unused_memo_index)); @@ -811,7 +811,7 @@ class DictEncoderImpl : public EncoderImpl, virtual public DictEncoder { } /// The number of bytes needed to encode the dictionary. - int dict_encoded_size_; + int64_t dict_encoded_size_; MemoTableType memo_table_; }; @@ -850,7 +850,7 @@ inline void DictEncoderImpl::Put(const T& v) { // Put() implementation for primitive types auto on_found = [](int32_t memo_index) {}; auto on_not_found = [this](int32_t memo_index) { - dict_encoded_size_ += static_cast(sizeof(T)); + dict_encoded_size_ += sizeof(T); }; int32_t memo_index; @@ -874,7 +874,7 @@ inline void DictEncoderImpl::PutByteArray( auto on_found = [](int32_t memo_index) {}; auto on_not_found = [&](int32_t memo_index) { - dict_encoded_size_ += static_cast(length + sizeof(uint32_t)); + dict_encoded_size_ += static_cast(length) + sizeof(uint32_t); }; DCHECK(ptr != nullptr || length == 0); @@ -989,7 +989,7 @@ void DictEncoderImpl::PutDictionary(const ::arrow::Array& values) { const auto& data = checked_cast(values); dict_encoded_size_ += - static_cast(sizeof(typename DType::c_type) * data.length()); + static_cast(sizeof(typename DType::c_type)) * data.length(); for (int64_t i = 0; i < data.length(); i++) { int32_t unused_memo_index; PARQUET_THROW_NOT_OK( @@ -1004,7 +1004,7 @@ void DictEncoderImpl::PutDictionary(const ::arrow::Array& values) { const auto& data = checked_cast(values); - dict_encoded_size_ += static_cast(type_length_ * data.length()); + dict_encoded_size_ += static_cast(type_length_) * data.length(); for (int64_t i = 0; i < data.length(); i++) { int32_t unused_memo_index; PARQUET_THROW_NOT_OK(memo_table_.GetOrInsert( diff --git a/bolt/dwio/parquet/arrow/Encoding.h b/bolt/dwio/parquet/arrow/Encoding.h index ae85a7509..76d53a333 100644 --- a/bolt/dwio/parquet/arrow/Encoding.h +++ b/bolt/dwio/parquet/arrow/Encoding.h @@ -245,7 +245,7 @@ class DictEncoder : virtual public TypedEncoder { /// bytes. Use EstimatedDataEncodedSize() to size buffer. virtual int WriteIndices(uint8_t* buffer, int buffer_len) = 0; - virtual int dict_encoded_size() const = 0; + virtual int64_t dict_encoded_size() const = 0; virtual int bit_width() const = 0; diff --git a/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp b/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp index 2b0f9a097..5fb870531 100644 --- a/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp +++ b/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp @@ -32,7 +32,9 @@ #include #include +#include #include +#include #include #include @@ -993,6 +995,83 @@ TEST(TestColumnWriter, RepeatedListsUpdateSpacedBug) { writer->Close(); } +class PageWriterSizeGuardTest : public ::testing::Test { + protected: + void SetUp() override { + node_ = std::static_pointer_cast(GroupNode::Make( + "schema", + Repetition::REQUIRED, + {schema::ByteArray("value", Repetition::REQUIRED)})); + schemaDescriptor_.Init(node_); + properties_ = WriterProperties::Builder().disable_dictionary()->build(); + metadata_ = ColumnChunkMetaDataBuilder::Make( + properties_, schemaDescriptor_.Column(0)); + sink_ = CreateOutputStream(); + pageWriter_ = PageWriter::Open( + sink_, + Compression::UNCOMPRESSED, + Codec::UseDefaultCompressionLevel(), + metadata_.get()); + } + + template + void expectSizeError(Callback&& callback, const std::string& expectedText) { + try { + callback(); + FAIL() << "Expected ParquetException containing: " << expectedText; + } catch (const ParquetException& error) { + EXPECT_NE(std::string(error.what()).find(expectedText), std::string::npos) + << error.what(); + } + } + + std::shared_ptr node_; + SchemaDescriptor schemaDescriptor_; + std::shared_ptr properties_; + std::unique_ptr metadata_; + std::shared_ptr<::arrow::io::BufferOutputStream> sink_; + std::unique_ptr pageWriter_; +}; + +TEST_F(PageWriterSizeGuardTest, rejectsUncompressedDataPageSizeOverflow) { + static const uint8_t kDummyData = 0; + auto buffer = std::make_shared<::arrow::Buffer>(&kDummyData, 1); + constexpr int64_t kTooLarge = + static_cast(std::numeric_limits::max()) + 1; + auto page = std::make_unique( + buffer, 1, Encoding::PLAIN, Encoding::RLE, Encoding::RLE, kTooLarge); + + expectSizeError( + [&]() { pageWriter_->WriteDataPage(std::move(page)); }, + "Uncompressed data page size"); +} + +TEST_F(PageWriterSizeGuardTest, rejectsCompressedDataPageSizeOverflow) { + static const uint8_t kDummyData = 0; + constexpr int64_t kTooLarge = + static_cast(std::numeric_limits::max()) + 1; + auto buffer = std::make_shared<::arrow::Buffer>(&kDummyData, kTooLarge); + auto page = std::make_unique( + buffer, 1, Encoding::PLAIN, Encoding::RLE, Encoding::RLE, 1); + + expectSizeError( + [&]() { pageWriter_->WriteDataPage(std::move(page)); }, + "Compressed data page size"); +} + +TEST_F(PageWriterSizeGuardTest, rejectsDictionaryPageSizeBeforeNarrowing) { + static const uint8_t kDummyData = 0; + constexpr int64_t kTooLarge = + static_cast(std::numeric_limits::max()) + 1; + auto buffer = std::make_shared<::arrow::Buffer>(&kDummyData, kTooLarge); + auto page = + std::make_unique(buffer, 1, Encoding::PLAIN, false); + + expectSizeError( + [&]() { pageWriter_->WriteDictionaryPage(std::move(page)); }, + "Uncompressed dictionary page size"); +} + void GenerateLevels( int min_repeat_factor, int max_repeat_factor, diff --git a/bolt/dwio/parquet/tests/writer/ParquetWriterBenchmark.cpp b/bolt/dwio/parquet/tests/writer/ParquetWriterBenchmark.cpp index 2d7fb63e0..492236cff 100644 --- a/bolt/dwio/parquet/tests/writer/ParquetWriterBenchmark.cpp +++ b/bolt/dwio/parquet/tests/writer/ParquetWriterBenchmark.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "bolt/common/file/File.h" #include "bolt/dwio/common/FileSink.h" #include "bolt/dwio/common/Options.h" #include "bolt/dwio/common/Statistics.h" @@ -24,34 +25,57 @@ #include "bolt/dwio/parquet/writer/Writer.h" #include "bolt/exec/tests/utils/TempDirectoryPath.h" +#include #include #include + +#include +#include +#include using namespace bytedance::bolt; using namespace bytedance::bolt::dwio; using namespace bytedance::bolt::dwio::common; using namespace bytedance::bolt::parquet; using namespace bytedance::bolt::test; -const uint32_t kNumBatches = 50; -const uint32_t kNumRowsPerRowGroup = 10000; +constexpr uint32_t kNumBatches = 50; +constexpr uint32_t kNumRowsPerRowGroup = 10000; +constexpr int64_t kNoSplitDataPageSize = 64 * 1024 * 1024; +constexpr int64_t kNoSplitSinkCapacity = 8 * 1024 * 1024; + +struct WriteMetrics { + uint64_t outputSize; + uint64_t dataPageCount; +}; class ParquetWriterBenchmark { public: explicit ParquetWriterBenchmark( bool disableDictionary, - const RowTypePtr& rowType) + const RowTypePtr& rowType, + int64_t dataPageSize, + bool useMemorySink) : disableDictionary_(disableDictionary) { rootPool_ = memory::memoryManager()->addRootPool("ParquetWriterBenchmark"); leafPool_ = rootPool_->addLeafChild("ParquetWriterBenchmark"); dataSetBuilder_ = std::make_unique(*leafPool_, 0); - auto path = fileFolder_->path + "/" + fileName_; - auto localWriteFile = std::make_unique(path, true, false); - auto sink = - std::make_unique(std::move(localWriteFile), path); + std::unique_ptr sink; + if (useMemorySink) { + auto memorySink = std::make_unique( + kNoSplitSinkCapacity, FileSink::Options{.pool = leafPool_.get()}); + sink_ = memorySink.get(); + sink = std::move(memorySink); + } else { + fileFolder_ = bytedance::bolt::exec::test::TempDirectoryPath::create(); + auto path = fileFolder_->path + "/test.parquet"; + auto localWriteFile = std::make_unique(path, true, false); + sink = std::make_unique(std::move(localWriteFile), path); + } bytedance::bolt::parquet::WriterOptions options; options.enableFlushBasedOnBlockSize = true; options.parquetWriteTimestampUnit = TimestampUnit::kNano; options.writeInt96AsTimestamp = true; + options.dataPageSize = dataPageSize; options.dataPageVersion = bytedance::bolt::parquet::arrow::ParquetDataPageVersion::V2; if (disableDictionary_) { @@ -63,9 +87,9 @@ class ParquetWriterBenchmark { std::move(sink), options, rowType); } - ~ParquetWriterBenchmark() {} + ~ParquetWriterBenchmark() = default; - void writeToFile( + void writeToSink( const std::vector& batches, bool /*forRowGroupSkip*/) { for (auto& batch : batches) { @@ -75,46 +99,167 @@ class ParquetWriterBenchmark { writer_->close(); } - void writeSingleColumn( + std::unique_ptr> makeSingleColumnData( const std::string& columnName, const TypePtr& type, uint8_t nullsRateX100, + uint32_t numBatches, uint32_t batchSize) { - folly::BenchmarkSuspender suspender; - auto rowType = ROW({columnName}, {type}); // Generating the data (consider the null rate). - auto batches = dataSetBuilder_->makeDataset(rowType, kNumBatches, batchSize) - .withRowGroupSpecificData(kNumRowsPerRowGroup) - .withNullsForField(Subfield(columnName), nullsRateX100) - .build(); - suspender.dismiss(); - writeToFile(*batches, true); + return dataSetBuilder_->makeDataset(rowType, numBatches, batchSize) + .withRowGroupSpecificData(kNumRowsPerRowGroup) + .withNullsForField(Subfield(columnName), nullsRateX100) + .build(); + } + + WriteMetrics collectMetrics() const { + BOLT_CHECK_NOT_NULL(sink_); + std::string_view data(sink_->data(), sink_->size()); + dwio::common::ReaderOptions readerOptions{leafPool_.get()}; + ParquetReader reader( + std::make_unique( + std::make_shared(data), + readerOptions.getMemoryPool()), + readerOptions); + + uint64_t dataPageCount = 0; + const auto metadata = reader.fileMetaData(); + for (int rowGroupIndex = 0; rowGroupIndex < metadata.numRowGroups(); + ++rowGroupIndex) { + const auto rowGroup = metadata.rowGroup(rowGroupIndex); + for (int columnIndex = 0; columnIndex < rowGroup.numColumns(); + ++columnIndex) { + for (const auto& stats : + rowGroup.columnChunk(columnIndex).pageEncodingStats()) { + if (stats.page_type == thrift::PageType::DATA_PAGE || + stats.page_type == thrift::PageType::DATA_PAGE_V2) { + dataPageCount += stats.count; + } + } + } + } + return { + .outputSize = static_cast(sink_->size()), + .dataPageCount = dataPageCount}; } private: - const std::string fileName_ = "test.parquet"; - const std::shared_ptr - fileFolder_ = bytedance::bolt::exec::test::TempDirectoryPath::create(); const bool disableDictionary_; std::unique_ptr dataSetBuilder_; std::shared_ptr rootPool_; std::shared_ptr leafPool_; + std::shared_ptr fileFolder_; + MemorySink* sink_{nullptr}; std::unique_ptr writer_; }; +void reportLayoutOnce( + const std::string& benchmarkName, + const WriteMetrics& metrics) { + static std::mutex mutex; + static std::unordered_set reportedBenchmarks; + std::lock_guard lock(mutex); + if (reportedBenchmarks.insert(benchmarkName).second) { + std::cerr << fmt::format( + "LAYOUT {} output_bytes={} data_pages={}\n", + benchmarkName, + metrics.outputSize, + metrics.dataPageCount); + } +} + +void runImpl( + uint32_t iterations, + const std::string& columnName, + const TypePtr& type, + uint8_t nullsRateX100, + uint32_t batchSize, + uint32_t numBatches, + bool disableDictionary, + int64_t dataPageSize, + bool reportLayout) { + RowTypePtr rowType = ROW({columnName}, {type}); + for (uint32_t i = 0; i < iterations; ++i) { + // Measure only write/flush/close. Writer construction, input generation, + // and output metadata inspection are deliberately suspended so the + // benchmark remains sensitive to the ColumnWriter hot path. + folly::BenchmarkSuspender suspender; + ParquetWriterBenchmark benchmark( + disableDictionary, rowType, dataPageSize, reportLayout); + const auto batches = benchmark.makeSingleColumnData( + columnName, type, nullsRateX100, numBatches, batchSize); + suspender.dismiss(); + benchmark.writeToSink(*batches, true); + suspender.rehire(); + + if (reportLayout) { + const auto metrics = benchmark.collectMetrics(); + folly::doNotOptimizeAway(metrics.outputSize); + folly::doNotOptimizeAway(metrics.dataPageCount); + BOLT_CHECK_EQ(metrics.dataPageCount, 1); + reportLayoutOnce( + fmt::format( + "{}_batch_{}_null{}", columnName, batchSize, nullsRateX100), + metrics); + } + } +} + void run( - uint32_t, + uint32_t iterations, const std::string& columnName, const TypePtr& type, uint8_t nullsRateX100, uint32_t batchSize, bool disableDictionary) { - RowTypePtr rowType = ROW({columnName}, {type}); - ParquetWriterBenchmark benchmark(disableDictionary, rowType); - BIGINT()->toString(); - benchmark.writeSingleColumn(columnName, type, nullsRateX100, batchSize); + runImpl( + iterations, + columnName, + type, + nullsRateX100, + batchSize, + kNumBatches, + disableDictionary, + bytedance::bolt::parquet::WriterOptions{}.dataPageSize, + false); +} + +void runNoSplit( + uint32_t iterations, + const std::string& columnName, + const TypePtr& type, + uint8_t nullsRateX100, + uint32_t batchSize) { + runImpl( + iterations, + columnName, + type, + nullsRateX100, + batchSize, + 1, + true, + kNoSplitDataPageSize, + true); +} + +void runDictionaryNoSplit( + uint32_t iterations, + const std::string& columnName, + const TypePtr& type, + uint8_t nullsRateX100, + uint32_t batchSize) { + runImpl( + iterations, + columnName, + type, + nullsRateX100, + batchSize, + 1, + false, + kNoSplitDataPageSize, + true); } #define PARQUET_BENCHMARKS_NULLS(_type_, _name_, _null_) \ @@ -131,6 +276,64 @@ void run( #define PARQUET_BENCHMARKS(_type_, _name_) \ PARQUET_BENCHMARKS_NULLS(_type_, _name_, 20) +#define PARQUET_NO_SPLIT_BENCHMARKS_NULLS(_type_, _name_, _null_) \ + BENCHMARK_NAMED_PARAM( \ + runNoSplit, \ + _name_##_batch_4k_null##_null_, \ + #_name_, \ + _type_, \ + _null_, \ + 4096); \ + BENCHMARK_NAMED_PARAM( \ + runNoSplit, \ + _name_##_batch_32k_null##_null_, \ + #_name_, \ + _type_, \ + _null_, \ + 32768); \ + BENCHMARK_NAMED_PARAM( \ + runNoSplit, \ + _name_##_batch_256k_null##_null_, \ + #_name_, \ + _type_, \ + _null_, \ + 262144); \ + BENCHMARK_DRAW_LINE(); + +#define PARQUET_NESTED_NO_SPLIT_BENCHMARKS_NULLS(_type_, _name_, _null_) \ + BENCHMARK_NAMED_PARAM( \ + runNoSplit, \ + _name_##_batch_4k_null##_null_, \ + #_name_, \ + _type_, \ + _null_, \ + 4096); \ + BENCHMARK_NAMED_PARAM( \ + runNoSplit, \ + _name_##_batch_32k_null##_null_, \ + #_name_, \ + _type_, \ + _null_, \ + 32768); \ + BENCHMARK_DRAW_LINE(); + +#define PARQUET_DICTIONARY_NO_SPLIT_BENCHMARKS_NULLS(_type_, _name_, _null_) \ + BENCHMARK_NAMED_PARAM( \ + runDictionaryNoSplit, \ + _name_##_batch_4k_null##_null_, \ + #_name_, \ + _type_, \ + _null_, \ + 4096); \ + BENCHMARK_NAMED_PARAM( \ + runDictionaryNoSplit, \ + _name_##_batch_32k_null##_null_, \ + #_name_, \ + _type_, \ + _null_, \ + 32768); \ + BENCHMARK_DRAW_LINE(); + PARQUET_BENCHMARKS(VARCHAR(), Varchar); PARQUET_BENCHMARKS(BIGINT(), BigInt); PARQUET_BENCHMARKS(DOUBLE(), Double); @@ -139,6 +342,22 @@ PARQUET_BENCHMARKS(DECIMAL(38, 3), LongDecimalType); PARQUET_BENCHMARKS(MAP(BIGINT(), BIGINT()), Map); PARQUET_BENCHMARKS(ARRAY(BIGINT()), List); +// These cases use one batch of short values and a 64 MiB page target. Both the +// baseline and the fixed writer must therefore produce exactly one data page. +// This isolates the no-split BYTE_ARRAY hot path and prints output size/page +// count so page-layout changes cannot be mistaken for CPU improvements. +PARQUET_NO_SPLIT_BENCHMARKS_NULLS(VARCHAR(), VarcharPlainNoSplit, 0); +PARQUET_NO_SPLIT_BENCHMARKS_NULLS(VARCHAR(), VarcharPlainNoSplit, 20); +PARQUET_NO_SPLIT_BENCHMARKS_NULLS(VARBINARY(), VarbinaryPlainNoSplit, 20); +PARQUET_NESTED_NO_SPLIT_BENCHMARKS_NULLS( + ARRAY(VARCHAR()), + VarcharListPlainNoSplit, + 20); +PARQUET_DICTIONARY_NO_SPLIT_BENCHMARKS_NULLS( + VARCHAR(), + VarcharDictionaryNoSplit, + 20); + // TODO: Add all data types int main(int argc, char** argv) { diff --git a/bolt/dwio/parquet/tests/writer/ParquetWriterTest.cpp b/bolt/dwio/parquet/tests/writer/ParquetWriterTest.cpp index cad77a88d..a264a9ac5 100644 --- a/bolt/dwio/parquet/tests/writer/ParquetWriterTest.cpp +++ b/bolt/dwio/parquet/tests/writer/ParquetWriterTest.cpp @@ -28,8 +28,10 @@ * -------------------------------------------------------------------------- */ +#include #include #include +#include #include "arrow/c/abi.h" #include "arrow/c/bridge.h" @@ -40,7 +42,9 @@ #include "bolt/connectors/hive/HiveConnector.h" #include "bolt/dwio/common/tests/utils/BatchMaker.h" #include "bolt/dwio/parquet/RegisterParquetWriter.h" +#include "bolt/dwio/parquet/arrow/ColumnPage.h" #include "bolt/dwio/parquet/arrow/Encoding.h" +#include "bolt/dwio/parquet/arrow/tests/ColumnReader.h" #include "bolt/dwio/parquet/arrow/tests/FileReader.h" #include "bolt/dwio/parquet/tests/ParquetTestBase.h" #include "bolt/type/fbhive/HiveTypeParser.h" @@ -124,6 +128,24 @@ class ParquetWriterTest : public ParquetTestBase { readerOptions); } + template + int64_t forEachDataPage( + const std::string& parquetPath, + int columnIndex, + Callback&& callback) { + auto fileReader = vp::arrow::ParquetFileReader::OpenFile(parquetPath); + auto pageReader = fileReader->RowGroup(0)->GetColumnPageReader(columnIndex); + int64_t dataPageCount = 0; + while (auto page = pageReader->NextPage()) { + if (page->type() == vp::arrow::PageType::DATA_PAGE || + page->type() == vp::arrow::PageType::DATA_PAGE_V2) { + callback(page); + ++dataPageCount; + } + } + return dataPageCount; + } + ::arrow::MemoryPool* getArrowMemoryPool() { if (arrowPool_) { return arrowPool_.get(); @@ -705,6 +727,399 @@ TEST_F(ParquetWriterTest, columnPageSize) { ASSERT_EQ(4, chunk2PageEncodingStats[0].count); // data page num } +TEST_F(ParquetWriterTest, byteArrayPlainEncodingRespectsDataPageSize) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {VARCHAR()}); + const int64_t kRows = 1024; + const int64_t kValueSize = 256; + const int64_t kDataPageSize = 1024; + const int64_t kMaxValuesPerPage = + kDataPageSize / (kValueSize + sizeof(uint32_t)); + std::string parquetPath = tempPath_->path + "/byteArrayPageSize.parquet"; + + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = false; + writerOptions.columnDataPageSizeMap[c0] = kDataPageSize; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + auto data = makeRowVector({makeFlatVector(kRows, [&](auto row) { + return std::string(kValueSize, 'a' + row % 26); + })}); + + writer->write(data); + writer->close(); + assertRead(parquetPath, kRows, schema, data); + + int64_t totalValues = 0; + const auto dataPageCount = forEachDataPage( + parquetPath, 0, [&](const std::shared_ptr& page) { + ASSERT_EQ(vp::arrow::PageType::DATA_PAGE, page->type()); + const auto dataPage = + std::static_pointer_cast(page); + ASSERT_EQ(vp::arrow::Encoding::PLAIN, dataPage->encoding()); + // All values are non-null and have the same length, so this directly + // bounds the PLAIN value bytes in each page without counting level + // encoding overhead. + EXPECT_LE(dataPage->num_values(), kMaxValuesPerPage); + totalValues += dataPage->num_values(); + }); + EXPECT_EQ(kRows, totalValues); + EXPECT_GT(dataPageCount, 1); +} + +TEST_F(ParquetWriterTest, nullableByteArrayPlainEncodingSplitsPages) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {VARBINARY()}); + const int64_t kRows = 1024; + const int64_t kValueSize = 128; + const int64_t kDataPageSize = 1024; + std::string parquetPath = tempPath_->path + "/nullableVarbinaryPages.parquet"; + + auto data = makeRowVector({makeFlatVector( + kRows, + [&](auto row) { return std::string(kValueSize, 'a' + row % 26); }, + [](auto row) { return row % 5 == 0; }, + VARBINARY())}); + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = false; + writerOptions.dataPageVersion = vp::arrow::ParquetDataPageVersion::V2; + writerOptions.columnDataPageSizeMap[c0] = kDataPageSize; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(data); + writer->close(); + + assertRead(parquetPath, kRows, schema, data); + int64_t totalValues = 0; + int64_t totalNulls = 0; + const auto dataPageCount = forEachDataPage( + parquetPath, 0, [&](const std::shared_ptr& page) { + ASSERT_EQ(vp::arrow::PageType::DATA_PAGE_V2, page->type()); + const auto dataPage = + std::static_pointer_cast(page); + ASSERT_EQ(vp::arrow::Encoding::PLAIN, dataPage->encoding()); + const int64_t nonNullValues = + dataPage->num_values() - dataPage->num_nulls(); + EXPECT_LE( + nonNullValues * (kValueSize + sizeof(uint32_t)), kDataPageSize); + totalValues += dataPage->num_values(); + totalNulls += dataPage->num_nulls(); + }); + EXPECT_EQ(kRows, totalValues); + EXPECT_EQ((kRows + 4) / 5, totalNulls); + EXPECT_GT(dataPageCount, 1); +} + +TEST_F(ParquetWriterTest, nestedByteArrayPlainEncodingSplitsAtRecords) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {ARRAY(VARCHAR())}); + const vector_size_t kRows = 32; + const vector_size_t kElementsPerRow = 32; + const int64_t kValueSize = 64; + const int64_t kDataPageSize = 1024; + const int64_t kNullArrayCount = (kRows + 10) / 11; + const int64_t kExpectedLevels = + (kRows - kNullArrayCount) * kElementsPerRow + kNullArrayCount; + std::string parquetPath = tempPath_->path + "/nestedVarcharPages.parquet"; + + auto data = makeRowVector({makeArrayVector( + kRows, + [&](auto /*row*/) { return kElementsPerRow; }, + [&](auto index) { return std::string(kValueSize, 'a' + index % 26); }, + [](auto row) { return row % 11 == 0; }, + [](auto index) { return index % 13 == 0; })}); + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = false; + writerOptions.dataPageVersion = vp::arrow::ParquetDataPageVersion::V2; + writerOptions.columnDataPageSizeMap["c0.list.element"] = kDataPageSize; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(data); + writer->close(); + + assertRead(parquetPath, kRows, schema, data); + int64_t totalRows = 0; + int64_t totalLevels = 0; + const auto dataPageCount = forEachDataPage( + parquetPath, 0, [&](const std::shared_ptr& page) { + ASSERT_EQ(vp::arrow::PageType::DATA_PAGE_V2, page->type()); + const auto dataPage = + std::static_pointer_cast(page); + ASSERT_EQ(vp::arrow::Encoding::PLAIN, dataPage->encoding()); + // Every non-null array is larger than the page target and therefore + // must occupy its own page. A neighboring null array may share that + // page, but no page may contain two non-null array records. + EXPECT_GT(dataPage->num_rows(), 0); + EXPECT_LE(dataPage->num_rows(), 2); + EXPECT_LE( + dataPage->num_values() - dataPage->num_nulls(), kElementsPerRow); + totalRows += dataPage->num_rows(); + totalLevels += dataPage->num_values(); + }); + EXPECT_EQ(kRows, totalRows); + EXPECT_EQ(kExpectedLevels, totalLevels); + EXPECT_GT(dataPageCount, 2); +} + +TEST_F(ParquetWriterTest, nestedFinalOversizedRecordStartsNewPage) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {ARRAY(VARCHAR())}); + const vector_size_t kRows = 2; + const vector_size_t kElementsPerRow = 32; + const int64_t kValueSize = 64; + std::string parquetPath = + tempPath_->path + "/nestedFinalOversizedRecord.parquet"; + + auto data = makeRowVector({makeArrayVector( + kRows, + [&](auto row) { return row == 0 ? 1 : kElementsPerRow; }, + [&](auto index) { return std::string(kValueSize, 'a' + index % 26); })}); + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = false; + writerOptions.dataPageVersion = vp::arrow::ParquetDataPageVersion::V2; + writerOptions.columnDataPageSizeMap["c0.list.element"] = 1024; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(data); + writer->close(); + + assertRead(parquetPath, kRows, schema, data); + int64_t totalRows = 0; + const auto dataPageCount = forEachDataPage( + parquetPath, 0, [&](const std::shared_ptr& page) { + ASSERT_EQ(vp::arrow::PageType::DATA_PAGE_V2, page->type()); + const auto dataPage = + std::static_pointer_cast(page); + EXPECT_EQ(1, dataPage->num_rows()); + totalRows += dataPage->num_rows(); + }); + EXPECT_EQ(2, dataPageCount); + EXPECT_EQ(kRows, totalRows); +} + +TEST_F(ParquetWriterTest, byteArrayDictionaryFallbackSplitsPlainPages) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {VARCHAR()}); + const int64_t kRows = 2048; + const int64_t kValuePayloadSize = 128; + const int64_t kDataPageSize = 1024; + const int64_t kMinEncodedValueSize = kValuePayloadSize + 1 + sizeof(uint32_t); + const int64_t kMaxPlainValuesPerPage = kDataPageSize / kMinEncodedValueSize; + std::string parquetPath = + tempPath_->path + "/dictionaryFallbackPages.parquet"; + + auto data = makeRowVector({makeFlatVector(kRows, [&](auto row) { + return std::to_string(row) + std::string(kValuePayloadSize, 'a' + row % 26); + })}); + vp::WriterOptions writerOptions{}; + writerOptions.dictionaryPageSizeLimit = 512; + writerOptions.columnDataPageSizeMap[c0] = kDataPageSize; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(data); + writer->close(); + + assertRead(parquetPath, kRows, schema, data); + auto reader = createLocalParquetReader(parquetPath); + auto pageEncodingStats = + reader->fileMetaData().rowGroup(0).columnChunk(0).pageEncodingStats(); + ASSERT_EQ(3, pageEncodingStats.size()); + + auto findPageStats = [&](thrift::PageType::type pageType, + thrift::Encoding::type encoding) { + return std::find_if( + pageEncodingStats.begin(), + pageEncodingStats.end(), + [&](const auto& stats) { + return stats.page_type == pageType && stats.encoding == encoding; + }); + }; + const auto dictionaryPage = + findPageStats(thrift::PageType::DICTIONARY_PAGE, thrift::Encoding::PLAIN); + const auto dictionaryDataPage = findPageStats( + thrift::PageType::DATA_PAGE, thrift::Encoding::RLE_DICTIONARY); + const auto plainDataPage = + findPageStats(thrift::PageType::DATA_PAGE, thrift::Encoding::PLAIN); + + ASSERT_NE(dictionaryPage, pageEncodingStats.end()); + ASSERT_NE(dictionaryDataPage, pageEncodingStats.end()); + ASSERT_NE(plainDataPage, pageEncodingStats.end()); + + int64_t totalValues = 0; + int64_t observedPlainDataPages = 0; + forEachDataPage( + parquetPath, 0, [&](const std::shared_ptr& page) { + const auto dataPage = + std::static_pointer_cast(page); + totalValues += dataPage->num_values(); + if (dataPage->encoding() == vp::arrow::Encoding::PLAIN) { + ++observedPlainDataPages; + EXPECT_LE(dataPage->num_values(), kMaxPlainValuesPerPage); + } + }); + EXPECT_EQ(kRows, totalValues); + EXPECT_EQ(plainDataPage->count, observedPlainDataPages); + EXPECT_GT(observedPlainDataPages, 1); +} + +TEST_F(ParquetWriterTest, byteArrayDictionaryPageUsesByteBudgetBeforeFallback) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {VARBINARY()}); + const int64_t kRows = 1024; + const int64_t kValueSize = 256; + const int64_t kDictionaryPageSize = 1024; + const int64_t kEncodedValueSize = kValueSize + sizeof(uint32_t); + std::string parquetPath = + tempPath_->path + "/dictionaryPageByteBudget.parquet"; + + auto data = makeRowVector({makeFlatVector( + kRows, + [&](auto row) { + std::string value = std::to_string(row); + value.resize(kValueSize, 'a' + row % 26); + return value; + }, + [](auto row) { return row % 7 == 0; }, + VARBINARY())}); + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = true; + writerOptions.dictionaryPageSizeLimit = kDictionaryPageSize; + writerOptions.columnDataPageSizeMap[c0] = kDictionaryPageSize; + writerOptions.compression = CompressionKind::CompressionKind_NONE; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(data); + writer->close(); + + assertRead(parquetPath, kRows, schema, data); + + auto fileReader = vp::arrow::ParquetFileReader::OpenFile(parquetPath); + auto pageReader = fileReader->RowGroup(0)->GetColumnPageReader(0); + bool sawDictionaryPage = false; + bool sawPlainDataPage = false; + while (auto page = pageReader->NextPage()) { + if (page->type() == vp::arrow::PageType::DICTIONARY_PAGE) { + sawDictionaryPage = true; + const auto dictionaryPage = + std::static_pointer_cast(page); + // The soft limit may be crossed by one boundary value, but not by the + // entire writer batch. This bounds dictionary peak memory before + // fallback in the same way as the plain data-page path. + EXPECT_LE( + dictionaryPage->size(), kDictionaryPageSize + kEncodedValueSize); + } else if ( + page->type() == vp::arrow::PageType::DATA_PAGE || + page->type() == vp::arrow::PageType::DATA_PAGE_V2) { + const auto dataPage = std::static_pointer_cast(page); + sawPlainDataPage |= dataPage->encoding() == vp::arrow::Encoding::PLAIN; + } + } + EXPECT_TRUE(sawDictionaryPage); + EXPECT_TRUE(sawPlainDataPage); +} + +TEST_F(ParquetWriterTest, repeatedByteArrayValuesDoNotForceDictionaryFallback) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {VARCHAR()}); + const int64_t kRows = 64; + const int64_t kValueSize = 256; + const int64_t kDictionaryPageSize = 1024; + std::string parquetPath = + tempPath_->path + "/repeatedDictionaryValues.parquet"; + + auto data = makeRowVector({makeFlatVector( + kRows, [&](auto /*row*/) { return std::string(kValueSize, 'a'); })}); + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = true; + writerOptions.dictionaryPageSizeLimit = kDictionaryPageSize; + writerOptions.columnDataPageSizeMap[c0] = kDictionaryPageSize; + writerOptions.compression = CompressionKind::CompressionKind_NONE; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(data); + writer->close(); + + assertRead(parquetPath, kRows, schema, data); + + auto reader = createLocalParquetReader(parquetPath); + const auto pageEncodingStats = + reader->fileMetaData().rowGroup(0).columnChunk(0).pageEncodingStats(); + const auto plainDataPage = std::find_if( + pageEncodingStats.begin(), + pageEncodingStats.end(), + [](const auto& stats) { + return stats.page_type == thrift::PageType::DATA_PAGE && + stats.encoding == thrift::Encoding::PLAIN; + }); + EXPECT_EQ(plainDataPage, pageEncodingStats.end()); +} + +TEST_F(ParquetWriterTest, nullableByteArrayWholeChunkStartsNewPage) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {VARBINARY()}); + const vector_size_t kRowsPerChunk = 5; + const int64_t kValueSize = 128; + const int64_t kDataPageSize = 1024; + std::string parquetPath = + tempPath_->path + "/nullableWholeChunkPages.parquet"; + + auto makeChunk = [&](char value) { + return makeRowVector({makeFlatVector( + kRowsPerChunk, + [&](auto /*row*/) { return std::string(kValueSize, value); }, + [](auto row) { return row == 0; }, + VARBINARY())}); + }; + auto first = makeChunk('a'); + auto second = makeChunk('b'); + auto expected = BaseVector::create(schema, 2 * kRowsPerChunk, pool_.get()); + expected->copy(first.get(), 0, 0, first->size()); + expected->copy(second.get(), kRowsPerChunk, 0, second->size()); + + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = false; + writerOptions.dataPageVersion = vp::arrow::ParquetDataPageVersion::V2; + writerOptions.columnDataPageSizeMap[c0] = kDataPageSize; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(first); + writer->write(second); + writer->close(); + + assertRead(parquetPath, 2 * kRowsPerChunk, schema, expected); + int64_t totalValues = 0; + const auto dataPageCount = forEachDataPage( + parquetPath, 0, [&](const std::shared_ptr& page) { + ASSERT_EQ(vp::arrow::PageType::DATA_PAGE_V2, page->type()); + const auto dataPage = + std::static_pointer_cast(page); + EXPECT_EQ(kRowsPerChunk, dataPage->num_values()); + EXPECT_EQ(1, dataPage->num_nulls()); + totalValues += dataPage->num_values(); + }); + EXPECT_EQ(2, dataPageCount); + EXPECT_EQ(2 * kRowsPerChunk, totalValues); +} + +TEST_F(ParquetWriterTest, singleOversizedByteArrayValueMakesProgress) { + std::string c0{"c0"}; + auto schema = ROW({c0}, {VARCHAR()}); + const int64_t kValueSize = 2048; + std::string parquetPath = + tempPath_->path + "/singleOversizedByteArray.parquet"; + + auto data = makeRowVector( + {makeFlatVector({std::string(kValueSize, 'a')})}); + vp::WriterOptions writerOptions{}; + writerOptions.enableDictionary = false; + writerOptions.columnDataPageSizeMap[c0] = 1024; + auto writer = createLocalWriter(parquetPath, schema, writerOptions); + writer->write(data); + writer->close(); + + assertRead(parquetPath, 1, schema, data); + const auto dataPageCount = forEachDataPage( + parquetPath, 0, [&](const std::shared_ptr& page) { + const auto dataPage = + std::static_pointer_cast(page); + EXPECT_EQ(vp::arrow::Encoding::PLAIN, dataPage->encoding()); + EXPECT_EQ(1, dataPage->num_values()); + }); + EXPECT_EQ(1, dataPageCount); +} + TEST_F(ParquetWriterTest, arrowPool) { const size_t kRows = 4 * 1024; auto type = getType(); From f19f73058dd3420b964cec86fd8a543aedffc026 Mon Sep 17 00:00:00 2001 From: bunny <1065339646@qq.com> Date: Wed, 22 Jul 2026 16:30:02 +0800 Subject: [PATCH 2/2] fix(parquet): harden PageHeader size handling, add tests --- bolt/dwio/parquet/arrow/ColumnWriter.cpp | 164 +++++++--- bolt/dwio/parquet/arrow/ColumnWriter.h | 6 + bolt/dwio/parquet/arrow/Encoding.cpp | 18 +- .../parquet/arrow/tests/ColumnWriterTest.cpp | 305 ++++++++++++++++++ 4 files changed, 445 insertions(+), 48 deletions(-) diff --git a/bolt/dwio/parquet/arrow/ColumnWriter.cpp b/bolt/dwio/parquet/arrow/ColumnWriter.cpp index 6eb26bba1..6c4665845 100644 --- a/bolt/dwio/parquet/arrow/ColumnWriter.cpp +++ b/bolt/dwio/parquet/arrow/ColumnWriter.cpp @@ -451,9 +451,12 @@ class SerializedPageWriter : public PageWriter { if (data_encryptor_.get()) { UpdateEncryption(encryption::kDictionaryPage); - PARQUET_THROW_NOT_OK(encryption_buffer_->Resize( - data_encryptor_->CiphertextSizeDelta() + compressed_page_size, - false)); + const int32_t encrypted_buffer_size = checkPageHeaderSize( + "Encrypted dictionary", + static_cast(data_encryptor_->CiphertextSizeDelta()) + + compressed_page_size); + PARQUET_THROW_NOT_OK( + encryption_buffer_->Resize(encrypted_buffer_size, false)); output_data_len = data_encryptor_->Encrypt( compressed_data->data(), compressed_page_size, @@ -561,9 +564,12 @@ class SerializedPageWriter : public PageWriter { checkPageHeaderSize("Compressed data", output_data_len); if (data_encryptor_.get()) { - PARQUET_THROW_NOT_OK(encryption_buffer_->Resize( - data_encryptor_->CiphertextSizeDelta() + compressed_page_size, - false)); + const int32_t encrypted_buffer_size = checkPageHeaderSize( + "Encrypted data", + static_cast(data_encryptor_->CiphertextSizeDelta()) + + compressed_page_size); + PARQUET_THROW_NOT_OK( + encryption_buffer_->Resize(encrypted_buffer_size, false)); UpdateEncryption(encryption::kDataPage); output_data_len = data_encryptor_->Encrypt( compressed_data->data(), @@ -1081,7 +1087,8 @@ class ColumnWriterImpl { const bool use_dictionary, Encoding::type encoding, const WriterProperties* properties, - std::shared_ptr buffered_resources) + std::shared_ptr buffered_resources, + int64_t byte_array_page_size_limit) : metadata_(metadata), descr_(metadata->descr()), level_info_(ComputeLevelInfo(metadata->descr())), @@ -1097,6 +1104,10 @@ class ColumnWriterImpl { dictionary_pagesize_limit_( properties->dictionary_pagesize_limit(descr_->path())), data_pagesize_(properties->data_pagesize(descr_->path())), + byte_array_page_size_limit_(std::clamp( + byte_array_page_size_limit, + 0, + kMaxPageHeaderSize)), num_buffered_values_(0), num_buffered_encoded_values_(0), num_buffered_nulls_(0), @@ -1216,6 +1227,8 @@ class ColumnWriterImpl { const int64_t data_pagesize_; + const int64_t byte_array_page_size_limit_; + // The total number of values stored in the data page. This is the maximum of // the number of encoded definition levels or encoded values. For // non-repeated, required columns, this is equal to the number of encoded @@ -1691,14 +1704,16 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, const bool use_dictionary, Encoding::type encoding, const WriterProperties* properties, - std::shared_ptr buffered_resources) + std::shared_ptr buffered_resources, + int64_t byte_array_page_size_limit) : ColumnWriterImpl( metadata, std::move(pager), use_dictionary, encoding, properties, - std::move(buffered_resources)), + std::move(buffered_resources), + byte_array_page_size_limit), encoder_pool_(std::make_unique<::arrow::ProxyMemoryPool>( properties->memory_pool())) { current_encoder_ = MakeEncoder( @@ -2242,8 +2257,8 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, return; } - const int64_t dictionary_page_byte_limit = - std::clamp(dictionary_pagesize_limit_, 0, kMaxPageHeaderSize); + const int64_t dictionary_page_byte_limit = std::clamp( + dictionary_pagesize_limit_, 0, byte_array_page_size_limit_); if (current_dict_encoder_->dict_encoded_size() >= dictionary_page_byte_limit) { FallbackToPlainEncoding(); @@ -3007,8 +3022,9 @@ Status TypedColumnWriterImpl::WriteArrowDense( // The configured page sizes are soft flush targets. PageHeader fields are // the hard format boundary, so cap byte-budget arithmetic at INT32_MAX and // leave the final encoded/compressed size checks to PageWriter. - auto PageByteLimit = [](int64_t configured_limit) { - return std::clamp(configured_limit, 0, kMaxPageHeaderSize); + auto PageByteLimit = [&](int64_t configured_limit) { + return std::clamp( + configured_limit, 0, byte_array_page_size_limit_); }; const int64_t data_page_byte_limit = PageByteLimit(data_pagesize_); const int64_t dictionary_page_byte_limit = @@ -3166,21 +3182,25 @@ Status TypedColumnWriterImpl::WriteArrowDense( &batch_num_spaced_values, &null_count); - // For PLAIN encoding this is the data-page value size. During dictionary - // encoding it is a conservative upper bound on dictionary growth: duplicate - // values consume no additional dictionary bytes, but treating them as new - // entries keeps a single write batch from causing a large peak allocation. - // The offset range is exact for canonical Arrow binary arrays and a - // conservative upper bound if a null slot retains bytes. + // For PLAIN encoding this is the data-page value size. For dictionary and + // DELTA encodings it bounds the raw bytes passed to one encoder Put(). const int64_t page_byte_limit = CurrentPageByteLimit(); - const int64_t batch_encoded_bytes = PlainEncodedRangeSize( + const int64_t batch_plain_bytes = PlainEncodedRangeSize( value_offset, batch_num_spaced_values, batch_num_values, page_byte_limit); - if (CappedPageBytes( - CurrentPageBytes(), batch_encoded_bytes, page_byte_limit) <= - page_byte_limit) { + + const bool dictionary_encoding = IsCurrentDictionaryEncoding(); + const bool plain_encoding = current_encoder_->encoding() == Encoding::PLAIN; + + // The final nested chunk may end inside a record. Flush before writing it + // when the conservative bound cannot fit because it cannot be split later. + if (!dictionary_encoding && !check_page && num_buffered_values_ > 0 && + CappedPageBytes( + CurrentPageBytes(), batch_plain_bytes, page_byte_limit) > + page_byte_limit) { + AddDataPage(); WritePreparedChunk( offset, batch_size, @@ -3191,7 +3211,23 @@ Status TypedColumnWriterImpl::WriteArrowDense( return; } - const bool dictionary_encoding = IsCurrentDictionaryEncoding(); + // PLAIN uses the exact remaining page budget, and dictionary uses a + // conservative remaining budget. DELTA raw bytes only bound one Put(). + const bool batch_fits = dictionary_encoding || plain_encoding + ? CappedPageBytes( + CurrentPageBytes(), batch_plain_bytes, page_byte_limit) <= + page_byte_limit + : batch_plain_bytes <= page_byte_limit; + if (batch_fits) { + WritePreparedChunk( + offset, + batch_size, + batch_num_values, + batch_num_spaced_values, + null_count, + check_page); + return; + } // A false check_page means this chunk is the final, potentially partial // nested record. It cannot be split or flushed after writing, but its @@ -3214,8 +3250,8 @@ Status TypedColumnWriterImpl::WriteArrowDense( // Prefer starting a new page over scanning levels merely to fill the // remainder of the current page. If the whole chunk fits on an empty page, // this preserves the old one-pass path for nullable and nested columns. - if (!dictionary_encoding && num_buffered_values_ > 0 && - batch_encoded_bytes <= data_page_byte_limit) { + if (plain_encoding && num_buffered_values_ > 0 && + batch_plain_bytes <= data_page_byte_limit) { AddDataPage(); WritePreparedChunk( offset, @@ -3235,11 +3271,16 @@ Status TypedColumnWriterImpl::WriteArrowDense( int64_t remaining = batch_size; while (remaining > 0) { const bool current_dictionary_encoding = IsCurrentDictionaryEncoding(); + const bool current_plain_encoding = + current_encoder_->encoding() == Encoding::PLAIN; const int64_t current_page_byte_limit = CurrentPageByteLimit(); - int64_t current_page_bytes = CurrentPageBytes(); + const bool use_current_page_bytes = + current_dictionary_encoding || current_plain_encoding; + int64_t current_page_bytes = + use_current_page_bytes ? CurrentPageBytes() : 0; int64_t subchunk_levels = 0; int64_t subchunk_spaced_values = 0; - int64_t subchunk_encoded_bytes = 0; + int64_t subchunk_plain_bytes = 0; while (subchunk_levels < remaining) { const int64_t level_index = local_offset + subchunk_levels; @@ -3261,12 +3302,12 @@ Status TypedColumnWriterImpl::WriteArrowDense( CappedPageBytes( CappedPageBytes( current_page_bytes, - subchunk_encoded_bytes, + subchunk_plain_bytes, current_page_byte_limit), value_bytes, current_page_byte_limit) > current_page_byte_limit) { if (subchunk_levels == 0) { - if (!current_dictionary_encoding && num_buffered_values_ > 0) { + if (current_plain_encoding && num_buffered_values_ > 0) { AddDataPage(); current_page_bytes = 0; } @@ -3278,8 +3319,8 @@ Status TypedColumnWriterImpl::WriteArrowDense( if (HasSpacedValue(level_index)) { ++subchunk_spaced_values; } - subchunk_encoded_bytes = CappedPageBytes( - subchunk_encoded_bytes, value_bytes, current_page_byte_limit); + subchunk_plain_bytes = CappedPageBytes( + subchunk_plain_bytes, value_bytes, current_page_byte_limit); ++subchunk_levels; } @@ -3433,11 +3474,14 @@ Status TypedColumnWriterImpl::WriteArrowDense( // ---------------------------------------------------------------------- // Dynamic column writer constructor -std::shared_ptr ColumnWriter::Make( +namespace { + +std::shared_ptr MakeColumnWriter( ColumnChunkMetaDataBuilder* metadata, std::unique_ptr pager, const WriterProperties* properties, - std::shared_ptr buffered_resources) { + std::shared_ptr buffered_resources, + int64_t byte_array_page_size_limit) { const ColumnDescriptor* descr = metadata->descr(); const bool use_dictionary = properties->dictionary_enabled(descr->path()) && descr->physical_type() != Type::BOOLEAN; @@ -3464,7 +3508,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); case Type::INT32: return std::make_shared>( metadata, @@ -3472,7 +3517,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); case Type::INT64: return std::make_shared>( metadata, @@ -3480,7 +3526,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); case Type::INT96: return std::make_shared>( metadata, @@ -3488,7 +3535,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); case Type::FLOAT: return std::make_shared>( metadata, @@ -3496,7 +3544,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); case Type::DOUBLE: return std::make_shared>( metadata, @@ -3504,7 +3553,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); case Type::BYTE_ARRAY: return std::make_shared>( metadata, @@ -3512,7 +3562,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); case Type::FIXED_LEN_BYTE_ARRAY: return std::make_shared>( metadata, @@ -3520,7 +3571,8 @@ std::shared_ptr ColumnWriter::Make( use_dictionary, encoding, properties, - std::move(buffered_resources)); + std::move(buffered_resources), + byte_array_page_size_limit); default: ParquetException::NYI("type reader not implemented"); } @@ -3528,4 +3580,32 @@ std::shared_ptr ColumnWriter::Make( return std::shared_ptr(nullptr); } +} // namespace + +std::shared_ptr ColumnWriter::Make( + ColumnChunkMetaDataBuilder* metadata, + std::unique_ptr pager, + const WriterProperties* properties, + std::shared_ptr buffered_resources) { + return MakeColumnWriter( + metadata, + std::move(pager), + properties, + std::move(buffered_resources), + kMaxPageHeaderSize); +} + +std::shared_ptr ColumnWriter::MakeForTest( + ColumnChunkMetaDataBuilder* metadata, + std::unique_ptr pager, + const WriterProperties* properties, + int64_t byte_array_page_size_limit) { + return MakeColumnWriter( + metadata, + std::move(pager), + properties, + nullptr, + byte_array_page_size_limit); +} + } // namespace bytedance::bolt::parquet::arrow diff --git a/bolt/dwio/parquet/arrow/ColumnWriter.h b/bolt/dwio/parquet/arrow/ColumnWriter.h index d5dbf1fab..8e0b2d5da 100644 --- a/bolt/dwio/parquet/arrow/ColumnWriter.h +++ b/bolt/dwio/parquet/arrow/ColumnWriter.h @@ -225,6 +225,12 @@ class PARQUET_EXPORT ColumnWriter { std::shared_ptr buffered_resources = nullptr); + static std::shared_ptr MakeForTest( + ColumnChunkMetaDataBuilder*, + std::unique_ptr, + const WriterProperties* properties, + int64_t byte_array_page_size_limit); + /// \brief Closes the ColumnWriter, commits any buffered values to pages. /// \return Total size of the column in bytes virtual int64_t Close() = 0; diff --git a/bolt/dwio/parquet/arrow/Encoding.cpp b/bolt/dwio/parquet/arrow/Encoding.cpp index 3a9a56799..44487d176 100644 --- a/bolt/dwio/parquet/arrow/Encoding.cpp +++ b/bolt/dwio/parquet/arrow/Encoding.cpp @@ -3150,13 +3150,12 @@ class DeltaLengthByteArrayEncoder : public EncoderImpl, Encoding::DELTA_LENGTH_BYTE_ARRAY, pool = ::arrow::default_memory_pool()), sink_(pool), - length_encoder_(nullptr, pool), - encoded_size_{0} {} + length_encoder_(nullptr, pool) {} std::shared_ptr<::arrow::Buffer> FlushValues() override; int64_t EstimatedDataEncodedSize() override { - return encoded_size_ + length_encoder_.EstimatedDataEncodedSize(); + return sink_.length() + length_encoder_.EstimatedDataEncodedSize(); } using TypedEncoder::Put; @@ -3182,6 +3181,13 @@ class DeltaLengthByteArrayEncoder : public EncoderImpl, return Status::Invalid( "Parquet cannot store strings with size 2GB or more"); } + if (ARROW_PREDICT_FALSE( + view.size() + sink_.length() > + static_cast( + std::numeric_limits::max()))) { + return Status::Invalid( + "excess expansion in DELTA_LENGTH_BYTE_ARRAY"); + } length_encoder_.Put({static_cast(view.length())}, 1); PARQUET_THROW_NOT_OK(sink_.Append(view.data(), view.length())); return Status::OK(); @@ -3191,7 +3197,6 @@ class DeltaLengthByteArrayEncoder : public EncoderImpl, ::arrow::BufferBuilder sink_; DeltaBitPackEncoder length_encoder_; - uint32_t encoded_size_; }; template @@ -3225,7 +3230,9 @@ void DeltaLengthByteArrayEncoder::Put(const T* src, int num_values) { length_encoder_.Put(lengths.data(), batch_size); } - if (AddWithOverflow(encoded_size_, total_increment_size, &encoded_size_)) { + if (ARROW_PREDICT_FALSE( + sink_.length() + total_increment_size > + std::numeric_limits::max())) { throw ParquetException("excess expansion in DELTA_LENGTH_BYTE_ARRAY"); } PARQUET_THROW_NOT_OK(sink_.Reserve(total_increment_size)); @@ -3267,7 +3274,6 @@ DeltaLengthByteArrayEncoder::FlushValues() { PARQUET_THROW_NOT_OK(sink_.Append(data->data(), data->size())); std::shared_ptr buffer = FinishSink<::arrow::BufferBuilder>(sink_); - encoded_size_ = 0; return buffer; } diff --git a/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp b/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp index 5fb870531..706c0de92 100644 --- a/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp +++ b/bolt/dwio/parquet/arrow/tests/ColumnWriterTest.cpp @@ -40,6 +40,7 @@ #include +#include "arrow/array/builder_binary.h" #include "arrow/buffer.h" #include "arrow/io/buffered.h" #include "arrow/testing/gtest_util.h" @@ -49,6 +50,8 @@ #include "bolt/dwio/parquet/arrow/ColumnWriter.h" #include "bolt/dwio/parquet/arrow/Encoding.h" #include "bolt/dwio/parquet/arrow/EncodingInternal.h" +#include "bolt/dwio/parquet/arrow/EncryptionInternal.h" +#include "bolt/dwio/parquet/arrow/FileEncryptorInternal.h" #include "bolt/dwio/parquet/arrow/FileWriter.h" #include "bolt/dwio/parquet/arrow/Metadata.h" #include "bolt/dwio/parquet/arrow/PageBufferArena.h" @@ -1025,11 +1028,35 @@ class PageWriterSizeGuardTest : public ::testing::Test { } } + void enableDataEncryption() { + aesEncryptor_ = std::make_unique( + ParquetCipher::AES_GCM_V1, + /*key_len=*/16, + /*metadata=*/false); + auto dataEncryptor = std::make_shared( + aesEncryptor_.get(), + std::string(16, 'k'), + "file-aad", + "page-aad", + ::arrow::default_memory_pool()); + pageWriter_ = PageWriter::Open( + sink_, + Compression::UNCOMPRESSED, + metadata_.get(), + /*row_group_ordinal=*/0, + /*column_chunk_ordinal=*/0, + ::arrow::default_memory_pool(), + /*buffered_row_group=*/false, + /*header_encryptor=*/nullptr, + std::move(dataEncryptor)); + } + std::shared_ptr node_; SchemaDescriptor schemaDescriptor_; std::shared_ptr properties_; std::unique_ptr metadata_; std::shared_ptr<::arrow::io::BufferOutputStream> sink_; + std::unique_ptr aesEncryptor_; std::unique_ptr pageWriter_; }; @@ -1072,6 +1099,284 @@ TEST_F(PageWriterSizeGuardTest, rejectsDictionaryPageSizeBeforeNarrowing) { "Uncompressed dictionary page size"); } +TEST_F(PageWriterSizeGuardTest, rejectsEncryptedDataPageSizeOverflow) { + enableDataEncryption(); + static const uint8_t kDummyData = 0; + constexpr int64_t kLargestInt32 = std::numeric_limits::max(); + auto buffer = std::make_shared<::arrow::Buffer>(&kDummyData, kLargestInt32); + auto page = std::make_unique( + buffer, 1, Encoding::PLAIN, Encoding::RLE, Encoding::RLE, 1); + + expectSizeError( + [&]() { pageWriter_->WriteDataPage(std::move(page)); }, + "Encrypted data page size"); +} + +TEST_F(PageWriterSizeGuardTest, rejectsEncryptedDictionaryPageSizeOverflow) { + enableDataEncryption(); + static const uint8_t kDummyData = 0; + constexpr int64_t kLargestInt32 = std::numeric_limits::max(); + auto buffer = std::make_shared<::arrow::Buffer>(&kDummyData, kLargestInt32); + auto page = + std::make_unique(buffer, 1, Encoding::PLAIN, false); + + expectSizeError( + [&]() { pageWriter_->WriteDictionaryPage(std::move(page)); }, + "Encrypted dictionary page size"); +} + +class ByteArraySplitterPageHeaderLimitTest + : public ::testing::TestWithParam {}; + +TEST_P(ByteArraySplitterPageHeaderLimitTest, SplitsBeforeHardLimit) { + const bool dictionary_enabled = GetParam(); + auto node = std::static_pointer_cast(GroupNode::Make( + "schema", + Repetition::REQUIRED, + {schema::ByteArray("value", Repetition::REQUIRED)})); + SchemaDescriptor schema_descriptor; + schema_descriptor.Init(node); + WriterProperties::Builder properties_builder; + properties_builder.data_pagesize(1024)->dictionary_pagesize_limit(1024); + if (dictionary_enabled) { + properties_builder.enable_dictionary(); + } else { + properties_builder.disable_dictionary(); + } + auto properties = properties_builder.build(); + auto metadata = + ColumnChunkMetaDataBuilder::Make(properties, schema_descriptor.Column(0)); + auto sink = CreateOutputStream(); + auto pager = PageWriter::Open( + sink, + Compression::UNCOMPRESSED, + Codec::UseDefaultCompressionLevel(), + metadata.get()); + + constexpr int64_t kPageHeaderSizeLimit = 64; + auto writer = ColumnWriter::MakeForTest( + metadata.get(), std::move(pager), properties.get(), kPageHeaderSizeLimit); + + ::arrow::StringBuilder builder; + constexpr int64_t kNumValues = 4; + const int64_t value_length = dictionary_enabled ? 28 : 40; + for (int64_t i = 0; i < kNumValues; ++i) { + ASSERT_OK(builder.Append(std::string(value_length, 'a' + i))); + } + ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + auto arrow_properties = default_arrow_writer_properties(); + ArrowWriteContext context( + ::arrow::default_memory_pool(), arrow_properties.get()); + + ASSERT_OK(writer->WriteArrow( + nullptr, nullptr, array->length(), *array, &context, false)); + writer->Close(); + + ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + auto page_reader = PageReader::Open( + std::make_shared<::arrow::io::BufferReader>(buffer), + kNumValues, + Compression::UNCOMPRESSED, + default_reader_properties()); + int64_t num_values = 0; + bool saw_dictionary_page = false; + bool saw_plain_page = false; + while (auto page = page_reader->NextPage()) { + EXPECT_LE(page->size(), kPageHeaderSizeLimit); + if (page->type() == PageType::DICTIONARY_PAGE) { + saw_dictionary_page = true; + continue; + } + ASSERT_EQ(PageType::DATA_PAGE, page->type()); + auto data_page = std::static_pointer_cast(page); + EXPECT_LE(data_page->uncompressed_size(), kPageHeaderSizeLimit); + saw_plain_page |= data_page->encoding() == Encoding::PLAIN; + num_values += data_page->num_values(); + } + EXPECT_EQ(kNumValues, num_values); + EXPECT_EQ(dictionary_enabled, saw_dictionary_page); + EXPECT_TRUE(saw_plain_page); +} + +INSTANTIATE_TEST_SUITE_P( + Encoding, + ByteArraySplitterPageHeaderLimitTest, + ::testing::Values(false, true), + [](const ::testing::TestParamInfo& info) { + return info.param ? "Dictionary" : "Plain"; + }); + +class ByteArrayDeltaPageTest : public ::testing::Test { + protected: + void SetUp() override { + node_ = std::static_pointer_cast(GroupNode::Make( + "schema", + Repetition::REQUIRED, + {schema::ByteArray("value", Repetition::REQUIRED)})); + schema_descriptor_.Init(node_); + } + + std::shared_ptr write_arrow_array( + const std::shared_ptr<::arrow::Array>& array, + Encoding::type encoding, + int64_t data_page_size, + int64_t write_batch_size) { + WriterProperties::Builder properties_builder; + properties_builder.disable_dictionary() + ->encoding(encoding) + ->data_pagesize(data_page_size) + ->write_batch_size(write_batch_size); + auto properties = properties_builder.build(); + auto metadata = ColumnChunkMetaDataBuilder::Make( + properties, schema_descriptor_.Column(0)); + auto sink = CreateOutputStream(); + auto pager = PageWriter::Open( + sink, + Compression::UNCOMPRESSED, + Codec::UseDefaultCompressionLevel(), + metadata.get()); + auto writer = + ColumnWriter::Make(metadata.get(), std::move(pager), properties.get()); + auto arrow_properties = default_arrow_writer_properties(); + ArrowWriteContext context( + ::arrow::default_memory_pool(), arrow_properties.get()); + + auto status = writer->WriteArrow( + nullptr, nullptr, array->length(), *array, &context, false); + EXPECT_TRUE(status.ok()) << status.ToString(); + if (!status.ok()) { + return nullptr; + } + writer->Close(); + EXPECT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + return buffer; + } + + std::vector data_page_value_counts( + const std::shared_ptr& buffer, + int64_t num_values, + Encoding::type encoding) { + auto page_reader = PageReader::Open( + std::make_shared<::arrow::io::BufferReader>(buffer), + num_values, + Compression::UNCOMPRESSED, + default_reader_properties()); + std::vector counts; + while (auto page = page_reader->NextPage()) { + if (page->type() != PageType::DATA_PAGE) { + ADD_FAILURE() << "Unexpected non-data page"; + continue; + } + auto data_page = std::static_pointer_cast(page); + EXPECT_EQ(encoding, data_page->encoding()); + counts.push_back(data_page->num_values()); + } + return counts; + } + + void expect_round_trip( + const std::shared_ptr& buffer, + const std::vector& expected) { + auto page_reader = PageReader::Open( + std::make_shared<::arrow::io::BufferReader>(buffer), + expected.size(), + Compression::UNCOMPRESSED, + default_reader_properties()); + auto reader = std::static_pointer_cast>( + ColumnReader::Make( + schema_descriptor_.Column(0), std::move(page_reader))); + std::vector batch(expected.size()); + std::vector actual(expected.size()); + int64_t total_read = 0; + while (total_read < static_cast(expected.size())) { + int64_t values_read = 0; + reader->ReadBatch( + expected.size() - total_read, + nullptr, + nullptr, + batch.data(), + &values_read); + ASSERT_GT(values_read, 0); + for (int64_t i = 0; i < values_read; ++i) { + actual[total_read + i] = std::string_view(batch[i]); + } + total_read += values_read; + } + EXPECT_EQ(expected, actual); + } + + std::shared_ptr node_; + SchemaDescriptor schema_descriptor_; +}; + +TEST_F(ByteArrayDeltaPageTest, DeltaByteArrayUsesActualSizeForPageFlush) { + std::vector values(64, std::string(112, 'a')); + ::arrow::StringBuilder builder; + for (const auto& value : values) { + ASSERT_OK(builder.Append(value)); + } + ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + + auto buffer = write_arrow_array( + array, + Encoding::DELTA_BYTE_ARRAY, + /*data_page_size=*/1024, + /*write_batch_size=*/8); + ASSERT_NE(nullptr, buffer); + + EXPECT_EQ( + std::vector({64}), + data_page_value_counts( + buffer, values.size(), Encoding::DELTA_BYTE_ARRAY)); + expect_round_trip(buffer, values); +} + +TEST_F(ByteArrayDeltaPageTest, DeltaLengthAccountsForArrowPayload) { + std::vector values; + ::arrow::StringBuilder builder; + for (int i = 0; i < 12; ++i) { + values.push_back(std::string(40, static_cast('a' + i))); + ASSERT_OK(builder.Append(values.back())); + } + ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + + auto buffer = write_arrow_array( + array, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + /*data_page_size=*/128, + /*write_batch_size=*/4); + ASSERT_NE(nullptr, buffer); + + EXPECT_EQ( + std::vector({4, 4, 4}), + data_page_value_counts( + buffer, values.size(), Encoding::DELTA_LENGTH_BYTE_ARRAY)); + expect_round_trip(buffer, values); +} + +TEST_F(ByteArrayDeltaPageTest, DeltaByteArrayBoundsIncompressibleBatches) { + std::vector values; + ::arrow::StringBuilder builder; + for (int i = 0; i < 12; ++i) { + values.push_back(std::string(40, static_cast('a' + i))); + ASSERT_OK(builder.Append(values.back())); + } + ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + + auto buffer = write_arrow_array( + array, + Encoding::DELTA_BYTE_ARRAY, + /*data_page_size=*/128, + /*write_batch_size=*/4); + ASSERT_NE(nullptr, buffer); + + EXPECT_EQ( + std::vector({2, 2, 2, 2, 2, 2}), + data_page_value_counts( + buffer, values.size(), Encoding::DELTA_BYTE_ARRAY)); + expect_round_trip(buffer, values); +} + void GenerateLevels( int min_repeat_factor, int max_repeat_factor,