[native] Add support shuffle compression for Prestissimo

prestodb · Dec 23, 2024 · 1b76b44 · 1b76b44
1 parent 1a7fd4e
commit 1b76b44
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -277,6 +277,7 @@ void PrestoServer::run() {
   registerPrestoToVeloxConnector(
       std::make_unique<SystemPrestoToVeloxConnector>("$system@system"));
 
+  velox::exec::OutputBufferManager::initialize({});
   initializeVeloxMemory();
   initializeThreadPools();
 

diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
@@ -224,6 +224,17 @@ QueryContextManager::toVeloxConfigs(
       traceFragmentId = it.second;
     } else if (it.first == SessionProperties::kQueryTraceShardId) {
       traceShardId = it.second;
+    } else if (it.first == SessionProperties::kShuffleCompressionEnabled) {
+      if (it.second == "true") {
+        // NOTE: Presto java only support lz4 compression so configure the same
+        // compression kind on velox.
+        configs[core::QueryConfig::kShuffleCompressionKind] = std::to_string(
+            static_cast<uint32_t>(velox::common::CompressionKind_LZ4));
+      } else {
+        VELOX_USER_CHECK_EQ(it.second, "false");
+        configs[core::QueryConfig::kShuffleCompressionKind] = std::to_string(
+            static_cast<uint32_t>(velox::common::CompressionKind_NONE));
+      }
     } else {
       configs[sessionProperties_.toVeloxConfig(it.first)] = it.second;
       sessionProperties_.updateVeloxConfig(it.first, it.second);

diff --git a/presto-native-execution/presto_cpp/main/SessionProperties.h b/presto-native-execution/presto_cpp/main/SessionProperties.h
@@ -283,6 +283,10 @@ class SessionProperties {
   static constexpr const char* kPrefixSortMinRows =
       "native_prefixsort_min_rows";
 
+  /// If true, enable the shuffle compression.
+  static constexpr const char* kShuffleCompressionEnabled =
+      "exchange_compression";
+
   SessionProperties();
 
   const std::unordered_map<std::string, std::shared_ptr<SessionProperty>>&