Use string_view size for string stream memory accounting (facebookincubator#432)

Summary:

The new `NullableContentStringStreamData` chunker was using `sizeof(uint64_t)` (8 bytes) when calculating chunk sizes and memory usage, the batch reader consume string data using `std::string_view` (16 bytes). This mismatch causes memory regressions because chunk size calculations underestimate the actual memory footprint, leading to larger-than-expected chunks being created. This change uses `sizeof(std::string_view)` consistently across the chunking and stream data layers to protect readers that use string views from unexpected memory growth

Reviewed By: helfman

Differential Revision: D91016115

Author: macvincent

dwio/nimble/velox/StreamChunker.h

@@ -659,7 +659,10 @@ class NullableContentStringStreamChunker final : public StreamChunker {
       if (nonNulls[idx]) {
         currentExtraMemory =
             stringLengths[lengthsOffset_ + chunkSize.dataElementCount];
-        currentTotalSize += currentExtraMemory + sizeof(uint64_t);
+        // NOTE: This is not a bug, we use sizeof(std::string_view) instead of
+        // sizeof(uint64_t) to prevent reader regressions with the current
+        // default writer settings.
+        currentTotalSize += currentExtraMemory + sizeof(std::string_view);
         ++currentDataCount;
       }
 

dwio/nimble/velox/StreamData.h

@@ -356,13 +356,16 @@ class NullableContentStringStreamData final : public NullsStreamData {
   }
 
   inline uint64_t memoryUsed() const override {
+    // NOTE: This is not a bug, we multiply lengths_.size() by
+    // sizeof(std::string_view) instead of sizeof(uint64_t) to prevent reader
+    // regressions with the current default writer settings.
     return NullsStreamData::memoryUsed() + buffer_.size() +
-        lengths_.size() * sizeof(uint64_t) +
+        lengths_.size() * sizeof(std::string_view) +
         stringViews_.size() * sizeof(std::string_view);
   }
 
   inline uint64_t bufferSize() const {
-    return lengths_.size() * sizeof(uint64_t) + buffer_.size();
+    return lengths_.size() * sizeof(std::string_view) + buffer_.size();
   }
 
   inline StringBuffer mutableData() {

dwio/nimble/velox/tests/StreamChunkerTests.cpp

@@ -121,7 +121,7 @@ class StreamChunkerTestsBase : public ::testing::Test {
             std::vector<uint64_t>(lengths.begin(), lengths.end()),
             ::testing::ElementsAreArray(expectedLength));
         expectedStreamDataSize = expectedBufferContent.size() +
-            expectedLength.size() * sizeof(uint64_t);
+            expectedLength.size() * sizeof(std::string_view);
       } else {
         NIMBLE_UNREACHABLE(
             "NullableContentStringStreamData can only be called on string data");
@@ -1139,16 +1139,16 @@ TEST_F(StreamChunkerTestsBase, nullableContentStringStreamChunking) {
   populateStringData(data, testData);
 
   // number of strings: 7
-  // size of string buffer length = 7 * 8 = 56 bytes
+  // size of string buffer length = 7 * 16 = 112 bytes
   // size of nulls = 11 * 1 = 11 bytes
   // total number of characters = 36 bytes
-  // total size of stream data = 56 + 11 + 36 = 103 bytes
-  ASSERT_EQ(stream.memoryUsed(), 103);
+  // total size of stream data = 112 + 11 + 36 = 159 bytes
+  ASSERT_EQ(stream.memoryUsed(), 159);
 
   // Test 1: Not last chunk
   {
-    const uint64_t maxChunkSize = 45;
-    const uint64_t minChunkSize = 40;
+    const uint64_t maxChunkSize = 72;
+    const uint64_t minChunkSize = 50;
     auto chunker = getStreamChunker(
         stream,
         {
@@ -1220,7 +1220,8 @@ TEST_F(StreamChunkerTestsBase, nullableContentStringStreamChunking) {
     populateData(smallNonNulls, smallNonNullsData);
 
     // Exactly what is needed to fit just the first entry
-    const uint64_t minChunkSize = smallTestData.at(0).size() + sizeof(uint64_t);
+    const uint64_t minChunkSize =
+        smallTestData.at(0).size() + sizeof(std::string_view);
     const uint64_t maxChunkSize = minChunkSize;
     auto chunker = getStreamChunker(
         stream,
@@ -1254,7 +1255,7 @@ TEST_F(StreamChunkerTestsBase, nullableContentStringStreamChunking) {
 
     const uint64_t minChunkSize = 1;
     // Size of "really really large string", string view, and null minus 1.
-    const uint64_t maxChunkSize = 24 + sizeof(uint64_t) + sizeof(bool);
+    const uint64_t maxChunkSize = 24 + sizeof(std::string_view) + sizeof(bool);
     auto chunker = getStreamChunker(
         stream,
         {

dwio/nimble/velox/tests/StreamDataTest.cpp

@@ -548,7 +548,8 @@ TEST_F(NullableContentStringStreamDataTest, basicOperations) {
   EXPECT_EQ(buffer.size(), 10);
   EXPECT_EQ(
       streamData.memoryUsed(),
-      nonNulls.size() + buffer.size() + lengths.size() * sizeof(uint64_t));
+      nonNulls.size() + buffer.size() +
+          lengths.size() * sizeof(std::string_view));
 
   streamData.materialize();
   const auto* views =

dwio/nimble/velox/tests/VeloxWriterTest.cpp

@@ -30,6 +30,7 @@
 #include "dwio/nimble/velox/ChunkedStream.h"
 #include "dwio/nimble/velox/EncodingLayoutTree.h"
 #include "dwio/nimble/velox/FlushPolicy.h"
+#include "dwio/nimble/velox/SchemaReader.h"
 #include "dwio/nimble/velox/SchemaSerialization.h"
 #include "dwio/nimble/velox/StatsGenerated.h"
 #include "dwio/nimble/velox/VeloxReader.h"
@@ -392,10 +393,27 @@ ChunkSizeResults validateChunkSize(
     nimble::VeloxReader& reader,
     const uint64_t minStreamChunkRawSize,
     const uint64_t maxStreamChunkRawSize) {
-  const double kMaxErrorRate = 0.2;
+  constexpr int kRowCountOffset = 2;
+  constexpr double kMaxErrorRate = 0.2;
   const auto& tablet = reader.tabletReader();
   auto& pool = reader.memoryPool();
 
+  std::unordered_set<uint32_t> stringStreamOffsets;
+  nimble::SchemaReader::traverseSchema(
+      reader.schema(),
+      [&](uint32_t /*level*/,
+          const nimble::Type& type,
+          const nimble::SchemaReader::NodeInfo& /*info*/) {
+        if (type.isScalar()) {
+          auto scalarKind = type.asScalar().scalarDescriptor().scalarKind();
+          if (scalarKind == nimble::ScalarKind::String ||
+              scalarKind == nimble::ScalarKind::Binary) {
+            stringStreamOffsets.insert(
+                type.asScalar().scalarDescriptor().offset());
+          }
+        }
+      });
+
   const uint32_t stripeCount = tablet.stripeCount();
   uint32_t maxChunkCount = 0;
   uint32_t minChunkCount = std::numeric_limits<uint32_t>::max();
@@ -415,21 +433,35 @@ ChunkSizeResults validateChunkSize(
       nimble::InMemoryChunkedStream chunkedStream{
           pool, std::move(streamLoaders[streamId])};
       uint32_t currentStreamChunkCount = 0;
+      const bool isStringStream = stringStreamOffsets.contains(streamId);
       while (chunkedStream.hasNext()) {
         ++currentStreamChunkCount;
         const auto chunk = chunkedStream.nextChunk();
         const uint64_t chunkRawDataSize =
             nimble::test::TestUtils::getRawDataSize(pool, chunk);
-        EXPECT_LE(chunkRawDataSize, maxStreamChunkRawSize)
-            << "Stream " << streamId << " has a chunk with size "
-            << chunkRawDataSize << " which is above max chunk size of "
-            << maxStreamChunkRawSize;
+        const uint32_t rowCount =
+            *reinterpret_cast<const uint32_t*>(chunk.data() + kRowCountOffset);
+        const double stringError =
+            isStringStream ? rowCount * sizeof(std::string_view) : 0;
+
+        // For string streams, a single row may exceed max chunk size if the
+        // string itself is larger than the limit. This is acceptable since we
+        // cannot split a single string value.
+        if (!(isStringStream && rowCount == 1)) {
+          const double maxError =
+              kMaxErrorRate * maxStreamChunkRawSize + stringError;
+          EXPECT_LE(chunkRawDataSize, maxStreamChunkRawSize + maxError)
+              << "Stream " << streamId << " has a chunk with size "
+              << chunkRawDataSize << " which is above max chunk size of "
+              << maxStreamChunkRawSize;
+        }
 
         // Validate min chunk size when not last chunk
         if (chunkedStream.hasNext() &&
             chunkRawDataSize < minStreamChunkRawSize) {
-          uint64_t difference = minStreamChunkRawSize - chunkRawDataSize;
-          EXPECT_LE(difference * 1.0 / minStreamChunkRawSize, kMaxErrorRate)
+          const double minError =
+              kMaxErrorRate * minStreamChunkRawSize + stringError;
+          EXPECT_GE(chunkRawDataSize, minStreamChunkRawSize - minError)
               << "Stream " << streamId << " has a non-last chunk with size "
               << chunkRawDataSize << " which is below min chunk size of "
               << minStreamChunkRawSize;
@@ -2248,7 +2280,8 @@ TEST_F(VeloxWriterTest, fuzzComplex) {
 
     const auto iterations = 20;
     // provide sufficient buffer between min and max chunk size thresholds
-    constexpr uint64_t chunkThresholdBuffer = sizeof(int64_t) + sizeof(bool);
+    constexpr uint64_t chunkThresholdBuffer =
+        sizeof(std::string_view) + sizeof(bool);
     for (auto i = 0; i < iterations; ++i) {
       writerOptions.minStreamChunkRawSize =
           std::uniform_int_distribution<uint64_t>(10, 4096)(rng);