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"));