fix local exchanger low mem mode

jacktengg · Dec 11, 2024 · 57446b7 · 57446b7
1 parent f8af843
commit 57446b7
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 9 deletions.
diff --git a/be/src/pipeline/dependency.h b/be/src/pipeline/dependency.h
@@ -847,8 +847,7 @@ struct LocalExchangeSharedState : public BasicSharedState {
     }
 
     virtual void set_low_memory_mode() {
-        _buffer_mem_limit =
-                std::min<int64_t>(config::local_exchange_buffer_mem_limit, 10 * 1024 * 1024);
+        _buffer_mem_limit = std::min<int64_t>(config::local_exchange_buffer_mem_limit, 512 * 1024);
     }
 };
 
@@ -896,8 +895,7 @@ struct LocalMergeExchangeSharedState : public LocalExchangeSharedState {
     }
 
     void set_low_memory_mode() override {
-        _buffer_mem_limit =
-                std::min<int64_t>(config::local_exchange_buffer_mem_limit, 10 * 1024 * 1024);
+        _buffer_mem_limit = std::min<int64_t>(config::local_exchange_buffer_mem_limit, 512 * 1024);
         DCHECK(!_queues_mem_usage.empty());
         _each_queue_limit =
                 std::max<int64_t>(64 * 1024, _buffer_mem_limit / _queues_mem_usage.size());

diff --git a/be/src/pipeline/local_exchange/local_exchange_sink_operator.cpp b/be/src/pipeline/local_exchange/local_exchange_sink_operator.cpp
@@ -147,6 +147,7 @@ Status LocalExchangeSinkOperatorX::sink(RuntimeState* state, vectorized::Block*
 
     if (state->get_query_ctx()->low_memory_mode()) {
         local_state._shared_state->set_low_memory_mode();
+        local_state._exchanger->set_low_memory_mode();
     }
 
     RETURN_IF_ERROR(local_state._exchanger->sink(state, in_block, eos, local_state));

diff --git a/be/src/pipeline/local_exchange/local_exchanger.cpp b/be/src/pipeline/local_exchange/local_exchanger.cpp
@@ -352,6 +352,16 @@ void ExchangerBase::finalize(LocalExchangeSourceLocalState& local_state) {
         // do nothing
     }
 }
+
+void ExchangerBase::set_low_memory_mode() {
+    _free_block_limit = 0;
+
+    vectorized::Block block;
+    while (_free_blocks.try_dequeue(block)) {
+        // do nothing
+    }
+}
+
 void LocalMergeSortExchanger::finalize(LocalExchangeSourceLocalState& local_state) {
     BlockWrapperSPtr next_block;
     vectorized::Block block;

diff --git a/be/src/pipeline/local_exchange/local_exchanger.h b/be/src/pipeline/local_exchange/local_exchanger.h
@@ -72,6 +72,8 @@ class ExchangerBase {
 
     virtual std::string data_queue_debug_string(int i) = 0;
 
+    void set_low_memory_mode();
+
 protected:
     friend struct LocalExchangeSharedState;
     friend struct BlockWrapper;
@@ -83,7 +85,7 @@ class ExchangerBase {
     const int _num_partitions;
     const int _num_senders;
     const int _num_sources;
-    const int _free_block_limit = 0;
+    int _free_block_limit = 0;
     moodycamel::ConcurrentQueue<vectorized::Block> _free_blocks;
 };
 
@@ -181,10 +183,9 @@ struct BlockWrapper {
         if (ref_count.fetch_sub(1) == 1) {
             DCHECK_GT(allocated_bytes, 0);
             shared_state->sub_total_mem_usage(allocated_bytes, channel_id);
-            if (shared_state->exchanger->_free_block_limit == 0 ||
-                shared_state->exchanger->_free_blocks.size_approx() <
-                        shared_state->exchanger->_free_block_limit *
-                                shared_state->exchanger->_num_sources) {
+            if (shared_state->exchanger->_free_blocks.size_approx() <
+                shared_state->exchanger->_free_block_limit *
+                        shared_state->exchanger->_num_sources) {
                 data_block.clear_column_data();
                 shared_state->exchanger->_free_blocks.enqueue(std::move(data_block));
             }