From 30426e2be69ad10c2f955a384fe4876a7c32bd40 Mon Sep 17 00:00:00 2001 From: lizhen-0710 Date: Thu, 30 Jul 2026 19:11:24 +0800 Subject: [PATCH] misc: reuse owned parquet string dictionary page buffers Avoid keeping two resident copies of large string dictionary pages after decoding. For BYTE_ARRAY dictionary pages, the decoded string data can already live in PageReader-owned storage: the decompressed page buffer for compressed pages or the decrypted buffer for encrypted uncompressed pages. When the current page data points at one of those owned buffers, move that buffer directly into dictionary_.strings instead of allocating a second buffer and copying the bytes. Also clear dictionary-page scratch buffers after dictionary materialization so non-adopted page storage does not linger on the PageReader. Add targeted tests for decompressed and decrypted buffer adoption, plus an end-to-end compressed string dictionary read under a constrained memory pool to catch regressions back to the duplicate-copy behavior. Verified with: CCACHE_DISABLE=1 _build/Release/bolt/dwio/parquet/tests/reader/bolt_dwio_parquet_reader_test --gtest_filter='ParquetReaderTest.stringDictionaryTakesOwnedPageBuffers:ParquetReaderTest.compressedStringDictionaryDoesNotKeepDuplicatePageCopy' Verified with: CCACHE_DISABLE=1 ctest --test-dir _build/Release -R 'parquet|Parquet' --output-on-failure Co-authored-by: TRAE CLI --- bolt/dwio/parquet/reader/PageReader.cpp | 47 +++++- bolt/dwio/parquet/reader/PageReader.h | 4 + .../tests/reader/ParquetReaderTest.cpp | 159 ++++++++++++++++++ 3 files changed, 204 insertions(+), 6 deletions(-) diff --git a/bolt/dwio/parquet/reader/PageReader.cpp b/bolt/dwio/parquet/reader/PageReader.cpp index bbec00ebb..72e5962e9 100644 --- a/bolt/dwio/parquet/reader/PageReader.cpp +++ b/bolt/dwio/parquet/reader/PageReader.cpp @@ -531,6 +531,7 @@ void PageReader::prepareDataPageV2( } void PageReader::prepareDictionary(const PageHeader& pageHeader) { + pageData_ = nullptr; dictionary_.numValues = pageHeader.dictionary_page_header.num_values; dictionaryEncoding_ = pageHeader.dictionary_page_header.encoding; dictionary_.sorted = pageHeader.dictionary_page_header.__isset.is_sorted && @@ -542,13 +543,18 @@ void PageReader::prepareDictionary(const PageHeader& pageHeader) { cryptoCtx_.startDecryptWithDictionaryPage = false; int32_t compressedLen = pageHeader.compressed_page_size; int32_t uncompressedLen = pageHeader.uncompressed_page_size; + bool decryptedPage = false; if (codec_ != thrift::CompressionCodec::UNCOMPRESSED) { pageData_ = readBytes(compressedLen, pageBuffer_); if (cryptoCtx_.dataDecryptor != nullptr) { decryptPageData(compressedLen); + decryptedPage = true; } pageData_ = decompressData( pageData_, compressedLen, pageHeader.uncompressed_page_size); + if (decryptedPage) { + decryptionBuffer_.reset(); + } } else { if (cryptoCtx_.dataDecryptor != nullptr) { pageData_ = readBytes(uncompressedLen, pageBuffer_); @@ -652,14 +658,19 @@ void PageReader::prepareDictionary(const PageHeader& pageHeader) { AlignedBuffer::allocate(dictionary_.numValues, &pool_); auto numBytes = uncompressedLen; auto values = dictionary_.values->asMutable(); - dictionary_.strings = AlignedBuffer::allocate(numBytes, &pool_); - auto strings = dictionary_.strings->asMutable(); - if (pageData_) { - memcpy(strings, pageData_, numBytes); + if (auto ownedPageBuffer = takeOwnedPageBuffer(pageData_, numBytes)) { + dictionary_.strings = std::move(ownedPageBuffer); } else { - dwio::common::readBytes( - numBytes, inputStream_.get(), strings, bufferStart_, bufferEnd_); + dictionary_.strings = AlignedBuffer::allocate(numBytes, &pool_); + auto strings = dictionary_.strings->asMutable(); + if (pageData_) { + memcpy(strings, pageData_, numBytes); + } else { + dwio::common::readBytes( + numBytes, inputStream_.get(), strings, bufferStart_, bufferEnd_); + } } + auto strings = dictionary_.strings->asMutable(); auto header = strings; for (auto i = 0; i < dictionary_.numValues; ++i) { auto length = *reinterpret_cast(header); @@ -741,6 +752,27 @@ void PageReader::prepareDictionary(const PageHeader& pageHeader) { BOLT_UNSUPPORTED( "Parquet type {} not supported for dictionary", parquetType); } + pageData_ = nullptr; + pageBuffer_.reset(); + decompressedData_.reset(); + decryptionBuffer_.reset(); +} + +BufferPtr PageReader::takeOwnedPageBuffer( + const char* FOLLY_NULLABLE data, + size_t size) { + if (data == nullptr) { + return nullptr; + } + if (decompressedData_ && data == decompressedData_->as()) { + decompressedData_->setSize(size); + return std::move(decompressedData_); + } + if (decryptionBuffer_ && data == decryptionBuffer_->as()) { + decryptionBuffer_->setSize(size); + return std::move(decryptionBuffer_); + } + return nullptr; } void PageReader::makeFilterCache(dwio::common::ScanState& state) { @@ -1570,6 +1602,9 @@ void PageReader::decryptPageData(int32_t& compressedLen) { compressedLen, const_cast(decryptionBuffer_->as()), decryptionBuffer_->size()); + if (pageBuffer_ && pageData_ == pageBuffer_->as()) { + pageBuffer_.reset(); + } pageData_ = decryptionBuffer_->as(); } diff --git a/bolt/dwio/parquet/reader/PageReader.h b/bolt/dwio/parquet/reader/PageReader.h index 51038529e..c304c39c7 100644 --- a/bolt/dwio/parquet/reader/PageReader.h +++ b/bolt/dwio/parquet/reader/PageReader.h @@ -49,6 +49,7 @@ namespace bytedance::bolt::parquet { constexpr int16_t kNonPageOrdinal = static_cast(-1); constexpr uint32_t kDefaultMaxPageHeaderSize = 16 * 1024 * 1024; +struct PageReaderTestPeer; struct CryptoContext { CryptoContext( @@ -302,6 +303,7 @@ class PageReader { int64_t row, const bool keepRepDefRawData); void makeDecoder(); + BufferPtr takeOwnedPageBuffer(const char* FOLLY_NULLABLE data, size_t size); // For a non-top level leaf, reads the defs and sets 'leafNulls_' and // 'numRowsInPage_' accordingly. This is used for non-top level leaves when @@ -634,6 +636,8 @@ class PageReader { // Tracks output count for the current physical page. -1 means there is no // active page being accounted. int32_t currentPageNumValues_{-1}; + + friend struct PageReaderTestPeer; }; FOLLY_ALWAYS_INLINE dwio::common::compression::CompressionOptions diff --git a/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp b/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp index fe50b23f6..d7ca7687e 100644 --- a/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp +++ b/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp @@ -50,6 +50,8 @@ #include "bolt/functions/sparksql/VariantEncoding.h" #include "bolt/functions/sparksql/VariantFunctions.h" #include "bolt/vector/BaseVector.h" +#include "bolt/vector/DictionaryVector.h" +#include "bolt/vector/LazyVector.h" #include "bolt/vector/VariantVector.h" #include "bolt/vector/tests/utils/VectorMaker.h" @@ -92,6 +94,48 @@ std::string getVariantFixturePath(const std::string& fileName) { } } // namespace +namespace bytedance::bolt::parquet { +struct PageReaderTestPeer { + static BufferPtr setAndTakeDecompressedBuffer( + PageReader& reader, + const std::string& payload) { + setBuffer(reader.decompressedData_, reader, payload); + return reader.takeOwnedPageBuffer( + reader.decompressedData_->as(), payload.size()); + } + + static BufferPtr setAndTakeDecryptionBuffer( + PageReader& reader, + const std::string& payload) { + setBuffer(reader.decryptionBuffer_, reader, payload); + return reader.takeOwnedPageBuffer( + reader.decryptionBuffer_->as(), payload.size()); + } + + static bool hasDecompressedBuffer(const PageReader& reader) { + return reader.decompressedData_ != nullptr; + } + + static bool hasDecryptionBuffer(const PageReader& reader) { + return reader.decryptionBuffer_ != nullptr; + } + + static BufferPtr + takeOwnedPageBuffer(PageReader& reader, const char* data, size_t size) { + return reader.takeOwnedPageBuffer(data, size); + } + + private: + template + static void + setBuffer(BufferPtr& buffer, PageReader& reader, const std::string& payload) { + buffer = AlignedBuffer::allocate(payload.size(), &reader.pool_); + memcpy(buffer->asMutable(), payload.data(), payload.size()); + buffer->setSize(payload.size()); + } +}; +} // namespace bytedance::bolt::parquet + class ParquetReaderTest : public ParquetTestBase { public: std::unique_ptr createRowReader( @@ -1459,6 +1503,121 @@ TEST_F(ParquetReaderTest, prefetchRowGroups) { } } +TEST_F(ParquetReaderTest, stringDictionaryTakesOwnedPageBuffers) { + auto reader = PageReader( + nullptr, *leafPool_, thrift::CompressionCodec::UNCOMPRESSED, 0); + + auto assertBufferContent = [](const BufferPtr& buffer, + const std::string& expected) { + ASSERT_NE(buffer, nullptr); + ASSERT_EQ(buffer->size(), expected.size()); + ASSERT_EQ(std::string(buffer->as(), buffer->size()), expected); + }; + + const std::string decompressedPayload{ + "\x03\x00\x00\x00one\x03\x00\x00\x00two", 2 * sizeof(int32_t) + 6}; + auto decompressed = PageReaderTestPeer::setAndTakeDecompressedBuffer( + reader, decompressedPayload); + assertBufferContent(decompressed, decompressedPayload); + ASSERT_FALSE(PageReaderTestPeer::hasDecompressedBuffer(reader)); + + const std::string decryptedPayload{ + "\x05\x00\x00\x00three", sizeof(int32_t) + 5}; + auto decrypted = + PageReaderTestPeer::setAndTakeDecryptionBuffer(reader, decryptedPayload); + assertBufferContent(decrypted, decryptedPayload); + ASSERT_FALSE(PageReaderTestPeer::hasDecryptionBuffer(reader)); + + auto external = + AlignedBuffer::allocate(decryptedPayload.size(), leafPool_.get()); + EXPECT_EQ( + PageReaderTestPeer::takeOwnedPageBuffer( + reader, external->as(), external->size()), + nullptr); +} + +TEST_F( + ParquetReaderTest, + compressedStringDictionaryDoesNotKeepDuplicatePageCopy) { + constexpr int32_t kRows = 512; + constexpr int32_t kDictionaryValues = 40; + constexpr int32_t kStringBytes = 256 * 1024; + constexpr int64_t kReadMemoryLimit = 18 << 20; + constexpr int64_t kDictionaryPageBytes = + int64_t{kDictionaryValues} * (kStringBytes + sizeof(int32_t)); + auto rowType = ROW({"payload"}, {VARCHAR()}); + + auto dataPool = rootPool_->addLeafChild("dictionaryReuseData"); + std::vector stringBuffers; + stringBuffers.reserve(kDictionaryValues); + std::vector dictionary; + dictionary.reserve(kDictionaryValues); + for (int32_t i = 0; i < kDictionaryValues; ++i) { + auto buffer = AlignedBuffer::allocate(kStringBytes, dataPool.get()); + auto* raw = buffer->asMutable(); + memset(raw, static_cast('a' + (i % 26)), kStringBytes); + raw[kStringBytes - 1] = static_cast('A' + (i % 26)); + buffer->setSize(kStringBytes); + dictionary.emplace_back(raw, kStringBytes); + stringBuffers.push_back(std::move(buffer)); + } + + auto values = AlignedBuffer::allocate(kRows, dataPool.get()); + auto* rawValues = values->asMutable(); + for (int32_t i = 0; i < kRows; ++i) { + rawValues[i] = dictionary[i % dictionary.size()]; + } + auto input = std::make_shared( + dataPool.get(), + rowType, + nullptr, + kRows, + std::vector{std::make_shared>( + dataPool.get(), + VARCHAR(), + nullptr, + kRows, + values, + std::move(stringBuffers))}); + + auto tempFile = exec::test::TempFilePath::create(); + { + auto writeFile = + std::make_unique(tempFile->getPath(), true, false); + auto sink = std::make_unique( + std::move(writeFile), tempFile->getPath()); + bytedance::bolt::parquet::WriterOptions writerOptions; + writerOptions.memoryPool = rootPool_.get(); + writerOptions.enableDictionary = true; + writerOptions.dictionaryPageSizeLimit = kDictionaryPageBytes * 2; + writerOptions.dataPageSize = 1 << 20; + writerOptions.compression = common::CompressionKind_SNAPPY; + auto writer = std::make_unique( + std::move(sink), writerOptions, rowType); + writer->write(input); + writer->close(); + } + + auto readRootPool = memory::memoryManager()->addRootPool( + "compressedDictionaryReuseRead", kReadMemoryLimit); + auto readPool = readRootPool->addLeafChild("leaf"); + dwio::common::ReaderOptions readerOptions{readPool.get()}; + readerOptions.setFilePreloadThreshold(0); + readerOptions.setPrefetchRowGroups(0); + + auto reader = createReader(tempFile->getPath(), readerOptions); + RowReaderOptions rowReaderOpts; + rowReaderOpts.setScanSpec(makeScanSpec(rowType)); + auto rowReader = reader->createRowReader(rowReaderOpts); + VectorPtr result = BaseVector::create(rowType, 0, readPool.get()); + + ASSERT_EQ(rowReader->next(kRows, result), kRows); + LazyVector::ensureLoadedRows(result, SelectivityVector(kRows)); + test::assertEqualVectors(input, result); + EXPECT_EQ(rowReader->next(kRows, result), 0); + EXPECT_LT(readPool->peakBytes(), kReadMemoryLimit); +} + TEST_F(ParquetReaderTest, testEmptyRowGroups) { // empty_row_groups.parquet contains empty row groups const std::string sample(getExampleFilePath("empty_row_groups.parquet"));