fix: [Nimble][Chunking] Determine Threshold for Chunking Stages Based on Schema Width (#316)

Summary:
Pull Request resolved: #316

There are two chunking stages. A soft stage where we only chunk streams above a pre-determined max size and a hard stage where we chunk all streams above min size. In D86435350, we introduced an option to determine max size based on schema width. While we use this new max size value to initialize the chunker, we fail to use it to when determining chunking stages. This diff fixes that.

Reviewed By: helfman

Differential Revision: D86924613

fbshipit-source-id: 4c07eefe22d4474aa1c07cfa5227fb92d71a89b4

Author: macvincent

dwio/nimble/velox/VeloxWriter.cpp

@@ -1009,15 +1009,26 @@ bool VeloxWriter::evalauateFlushPolicy() {
     // Relieve memory pressure by chunking streams above max size.
     const auto& streams = context_->streams();
     std::vector<uint32_t> streamIndices;
-    streamIndices.reserve(streams.size());
+    const auto streamCount = streams.size();
+    streamIndices.reserve(streamCount);
+
+    // Determine size threshold for soft chunking based on schema width.
+    const auto& options = context_->options;
+    const auto maxChunkSize = streamCount > options.largeSchemaThreshold
+        ? options.wideSchemaMaxStreamChunkRawSize
+        : options.maxStreamChunkRawSize;
     for (auto streamIndex = 0; streamIndex < streams.size(); ++streamIndex) {
-      if (streams[streamIndex]->memoryUsed() >=
-          context_->options.maxStreamChunkRawSize) {
+      if (streams[streamIndex]->memoryUsed() >= maxChunkSize) {
         streamIndices.push_back(streamIndex);
       }
     }
+
+    // Soft chunking.
     const bool continueChunking =
         batchChunkStreams(streamIndices, /*ensureFullChunks=*/true);
+
+    // Hard chunking when chunking streams above maxChunkSize fails to
+    // relieve memory pressure.
     if (continueChunking) {
       // Relieve memory pressure by chunking small streams.
       // Sort streams for chunking based on raw memory usage.