bytedance · lizhen-0710 · Jul 30, 2026
diff --git 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<StringView>(dictionary_.numValues, &pool_);
       auto numBytes = uncompressedLen;
       auto values = dictionary_.values->asMutable<StringView>();
-      dictionary_.strings = AlignedBuffer::allocate<char>(numBytes, &pool_);
-      auto strings = dictionary_.strings->asMutable<char>();
-      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<char>(numBytes, &pool_);
+        auto strings = dictionary_.strings->asMutable<char>();
+        if (pageData_) {
+          memcpy(strings, pageData_, numBytes);
+        } else {
+          dwio::common::readBytes(
+              numBytes, inputStream_.get(), strings, bufferStart_, bufferEnd_);
+        }
       }
+      auto strings = dictionary_.strings->asMutable<char>();
       auto header = strings;
       for (auto i = 0; i < dictionary_.numValues; ++i) {
         auto length = *reinterpret_cast<const int32_t*>(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<char>()) {
+    decompressedData_->setSize(size);
+    return std::move(decompressedData_);
+  }
+  if (decryptionBuffer_ && data == decryptionBuffer_->as<char>()) {
+    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<uint8_t*>(decryptionBuffer_->as<uint8_t>()),
       decryptionBuffer_->size());
+  if (pageBuffer_ && pageData_ == pageBuffer_->as<char>()) {
+    pageBuffer_.reset();
+  }
   pageData_ = decryptionBuffer_->as<char>();
 }
 

diff --git 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<int16_t>(-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
@@ -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<char>(reader.decompressedData_, reader, payload);
+    return reader.takeOwnedPageBuffer(
+        reader.decompressedData_->as<char>(), payload.size());
+  }
+
+  static BufferPtr setAndTakeDecryptionBuffer(
+      PageReader& reader,
+      const std::string& payload) {
+    setBuffer<uint8_t>(reader.decryptionBuffer_, reader, payload);
+    return reader.takeOwnedPageBuffer(
+        reader.decryptionBuffer_->as<char>(), 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 <typename T>
+  static void
+  setBuffer(BufferPtr& buffer, PageReader& reader, const std::string& payload) {
+    buffer = AlignedBuffer::allocate<T>(payload.size(), &reader.pool_);
+    memcpy(buffer->asMutable<char>(), payload.data(), payload.size());
+    buffer->setSize(payload.size());
+  }
+};
+} // namespace bytedance::bolt::parquet
+
 class ParquetReaderTest : public ParquetTestBase {
  public:
   std::unique_ptr<dwio::common::RowReader> 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<char>(), 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<char>(decryptedPayload.size(), leafPool_.get());
+  EXPECT_EQ(
+      PageReaderTestPeer::takeOwnedPageBuffer(
+          reader, external->as<char>(), 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<BufferPtr> stringBuffers;
+  stringBuffers.reserve(kDictionaryValues);
+  std::vector<StringView> dictionary;
+  dictionary.reserve(kDictionaryValues);
+  for (int32_t i = 0; i < kDictionaryValues; ++i) {
+    auto buffer = AlignedBuffer::allocate<char>(kStringBytes, dataPool.get());
+    auto* raw = buffer->asMutable<char>();
+    memset(raw, static_cast<int>('a' + (i % 26)), kStringBytes);
+    raw[kStringBytes - 1] = static_cast<char>('A' + (i % 26));
+    buffer->setSize(kStringBytes);
+    dictionary.emplace_back(raw, kStringBytes);
+    stringBuffers.push_back(std::move(buffer));
+  }
+
+  auto values = AlignedBuffer::allocate<StringView>(kRows, dataPool.get());
+  auto* rawValues = values->asMutable<StringView>();
+  for (int32_t i = 0; i < kRows; ++i) {
+    rawValues[i] = dictionary[i % dictionary.size()];
+  }
+  auto input = std::make_shared<RowVector>(
+      dataPool.get(),
+      rowType,
+      nullptr,
+      kRows,
+      std::vector<VectorPtr>{std::make_shared<FlatVector<StringView>>(
+          dataPool.get(),
+          VARCHAR(),
+          nullptr,
+          kRows,
+          values,
+          std::move(stringBuffers))});
+
+  auto tempFile = exec::test::TempFilePath::create();
+  {
+    auto writeFile =
+        std::make_unique<LocalWriteFile>(tempFile->getPath(), true, false);
+    auto sink = std::make_unique<dwio::common::WriteFileSink>(
+        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<bytedance::bolt::parquet::Writer>(
+        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"));