add local exchange partition buffer size system config

Author: kagamiori

presto-docs/src/main/sphinx/presto_cpp/properties.rst

@@ -218,6 +218,14 @@ avoid exceeding memory limits for the query.
 When ``spill_enabled`` is ``true``, this determines whether Presto will try spilling memory to disk for order by to
 avoid exceeding memory limits for the query.
 
+``local-exchange.max-partition-buffer-size``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``integer``
+* **Default value:** ``65536`` (64KB)
+
+  Specifies the maximum size in bytes to accumulate for a single partition of a local exchange before flushing.
+
 
 ``shared-arbitrator.reserved-capacity``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

presto-native-execution/presto_cpp/main/QueryContextManager.cpp

@@ -57,7 +57,9 @@ void updateFromSystemConfigs(
           {core::QueryConfig::kAggregationSpillEnabled,
           std::string(SystemConfig::kAggregationSpillEnabled)},
           {core::QueryConfig::kRequestDataSizesMaxWaitSec,
-          std::string(SystemConfig::kRequestDataSizesMaxWaitSec)}};
+          std::string(SystemConfig::kRequestDataSizesMaxWaitSec)},
+          {core::QueryConfig::kMaxLocalExchangePartitionBufferSize,
+          std::string(SystemConfig::kMaxLocalExchangePartitionBufferSize)}};
   for (const auto& configNameEntry : sessionSystemConfigMapping) {
     const auto& sessionName = configNameEntry.first;
     const auto& systemConfigName = configNameEntry.second;

presto-native-execution/presto_cpp/main/common/Configs.cpp

@@ -263,6 +263,7 @@ SystemConfig::SystemConfig() {
           STR_PROP(kPluginDir, ""),
           NUM_PROP(kExchangeIoEvbViolationThresholdMs, 1000),
           NUM_PROP(kHttpSrvIoEvbViolationThresholdMs, 1000),
+          NUM_PROP(kMaxLocalExchangePartitionBufferSize, 65536),
       };
 }
 
@@ -909,6 +910,10 @@ int32_t SystemConfig::httpSrvIoEvbViolationThresholdMs() const {
       .value();
 }
 
+uint64_t SystemConfig::maxLocalExchangePartitionBufferSize() const {
+  return optionalProperty<uint64_t>(kMaxLocalExchangePartitionBufferSize).value();
+}
+
 NodeConfig::NodeConfig() {
   registeredProps_ =
       std::unordered_map<std::string, folly::Optional<std::string>>{

presto-native-execution/presto_cpp/main/common/Configs.h

@@ -755,6 +755,9 @@ class SystemConfig : public ConfigBase {
   static constexpr std::string_view kHttpSrvIoEvbViolationThresholdMs{
       "http-server.io-evb-violation-threshold-ms"};
 
+  static constexpr std::string_view kMaxLocalExchangePartitionBufferSize{
+      "local-exchange.max-partition-buffer-size"};
+
   SystemConfig();
 
   virtual ~SystemConfig() = default;
@@ -1035,6 +1038,8 @@ class SystemConfig : public ConfigBase {
   int32_t exchangeIoEvbViolationThresholdMs() const;
 
   int32_t httpSrvIoEvbViolationThresholdMs() const;
+
+  uint64_t maxLocalExchangePartitionBufferSize() const;
 };
 
 /// Provides access to node properties defined in node.properties file.

presto-native-execution/presto_cpp/main/tests/QueryContextManagerTest.cpp

@@ -127,6 +127,8 @@ TEST_F(QueryContextManagerTest, defaultSessionProperties) {
       queryConfig.spillWriteBufferSize(), defaultQC->spillWriteBufferSize());
   EXPECT_EQ(
       queryConfig.requestDataSizesMaxWaitSec(), defaultQC->requestDataSizesMaxWaitSec());
+  EXPECT_EQ(
+      queryConfig.maxLocalExchangePartitionBufferSize(), defaultQC->maxLocalExchangePartitionBufferSize());
 }
 
 TEST_F(QueryContextManagerTest, overridingSessionProperties) {
@@ -162,6 +164,9 @@ TEST_F(QueryContextManagerTest, overridingSessionProperties) {
     EXPECT_EQ(
         queryCtx->queryConfig().requestDataSizesMaxWaitSec(),
         systemConfig->requestDataSizesMaxWaitSec());
+    EXPECT_EQ(
+        queryCtx->queryConfig().maxLocalExchangePartitionBufferSize(),
+        systemConfig->maxLocalExchangePartitionBufferSize());
   }
   {
     protocol::SessionRepresentation session{
@@ -171,7 +176,8 @@ TEST_F(QueryContextManagerTest, overridingSessionProperties) {
             {"spill_enabled", "true"},
             {"aggregation_spill_enabled", "false"},
             {"join_spill_enabled", "true"},
-            {"request_data_sizes_max_wait_sec", "12"}}};
+            {"request_data_sizes_max_wait_sec", "12"},
+            {"local-exchange.max-partition-buffer-size", "1024"}}};
     protocol::TaskUpdateRequest updateRequest;
     updateRequest.session = session;
     auto queryCtx =
@@ -194,6 +200,11 @@ TEST_F(QueryContextManagerTest, overridingSessionProperties) {
     EXPECT_NE(
         queryCtx->queryConfig().aggregationSpillEnabled(),
         systemConfig->aggregationSpillEnabled());
+    // Override with different value
+    EXPECT_EQ(queryCtx->queryConfig().maxLocalExchangePartitionBufferSize(), 1024);
+    EXPECT_NE(
+        queryCtx->queryConfig().maxLocalExchangePartitionBufferSize(),
+        systemConfig->maxLocalExchangePartitionBufferSize());
     // Override with same value
     EXPECT_EQ(
         queryCtx->queryConfig().joinSpillEnabled(), true);