[opt](split) get file splits in batch mode

be/src/common/config.cpp

@@ -244,6 +244,7 @@ DEFINE_Validator(doris_scanner_thread_pool_thread_num, [](const int config) -> b
     }
     return true;
 });
+DEFINE_Int32(remote_split_source_batch_size, "1024");
 DEFINE_Int32(doris_max_remote_scanner_thread_pool_thread_num, "-1");
 // number of olap scanner thread pool queue size
 DEFINE_Int32(doris_scanner_thread_pool_queue_size, "102400");

be/src/common/config.h

@@ -286,6 +286,8 @@ DECLARE_mInt64(doris_blocking_priority_queue_wait_timeout_ms);
 // number of scanner thread pool size for olap table
 // and the min thread num of remote scanner thread pool
 DECLARE_mInt32(doris_scanner_thread_pool_thread_num);
+// number of batch size to fetch the remote split source
+DECLARE_mInt32(remote_split_source_batch_size);
 // max number of remote scanner thread pool size
 // if equal to -1, value is std::max(512, CpuInfo::num_cores() * 10)
 DECLARE_Int32(doris_max_remote_scanner_thread_pool_thread_num);

be/src/pipeline/exec/file_scan_operator.cpp

@@ -31,22 +31,20 @@
 namespace doris::pipeline {
 
 Status FileScanLocalState::_init_scanners(std::list<vectorized::VScannerSPtr>* scanners) {
-    if (_scan_ranges.empty()) {
+    if (_split_source->num_scan_ranges() == 0) {
         _eos = true;
-        _scan_dependency->set_ready();
         return Status::OK();
     }
 
     auto& p = _parent->cast<FileScanOperatorX>();
     size_t shard_num = std::min<size_t>(
             config::doris_scanner_thread_pool_thread_num / state()->query_parallel_instance_num(),
-            _scan_ranges.size());
+            _max_scanners);
     shard_num = std::max(shard_num, (size_t)1);
     _kv_cache.reset(new vectorized::ShardedKVCache(shard_num));
-    for (auto& scan_range : _scan_ranges) {
+    for (int i = 0; i < _max_scanners; ++i) {
         std::unique_ptr<vectorized::VFileScanner> scanner = vectorized::VFileScanner::create_unique(
-                state(), this, p._limit_per_scanner,
-                scan_range.scan_range.ext_scan_range.file_scan_range, _scanner_profile.get(),
+                state(), this, p._limit_per_scanner, _split_source, _scanner_profile.get(),
                 _kv_cache.get());
         RETURN_IF_ERROR(
                 scanner->prepare(_conjuncts, &_colname_to_value_range, &_colname_to_slot_id));
@@ -62,47 +60,25 @@ std::string FileScanLocalState::name_suffix() const {
 
 void FileScanLocalState::set_scan_ranges(RuntimeState* state,
                                          const std::vector<TScanRangeParams>& scan_ranges) {
-    int max_scanners =
+    _max_scanners =
             config::doris_scanner_thread_pool_thread_num / state->query_parallel_instance_num();
-    max_scanners = std::max(std::max(max_scanners, state->parallel_scan_max_scanners_count()), 1);
+    _max_scanners = std::max(std::max(_max_scanners, state->parallel_scan_max_scanners_count()), 1);
     // For select * from table limit 10; should just use one thread.
     if (should_run_serial()) {
-        max_scanners = 1;
+        _max_scanners = 1;
     }
-    if (scan_ranges.size() <= max_scanners) {
-        _scan_ranges = scan_ranges;
-    } else {
-        // There is no need for the number of scanners to exceed the number of threads in thread pool.
-        // scan_ranges is sorted by path(as well as partition path) in FE, so merge scan ranges in order.
-        // In the insert statement, reading data in partition order can reduce the memory usage of BE
-        // and prevent the generation of smaller tables.
-        _scan_ranges.resize(max_scanners);
-        int num_ranges = scan_ranges.size() / max_scanners;
-        int num_add_one = scan_ranges.size() - num_ranges * max_scanners;
-        int scan_index = 0;
-        int range_index = 0;
-        for (int i = 0; i < num_add_one; ++i) {
-            _scan_ranges[scan_index] = scan_ranges[range_index++];
-            auto& ranges =
-                    _scan_ranges[scan_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-            for (int j = 0; j < num_ranges; j++) {
-                auto& merged_ranges =
-                        scan_ranges[range_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-                ranges.insert(ranges.end(), merged_ranges.begin(), merged_ranges.end());
-            }
+    if (scan_ranges.size() == 1) {
+        auto scan_range = scan_ranges[0].scan_range.ext_scan_range.file_scan_range;
+        if (scan_range.__isset.split_source) {
+            auto split_source = scan_range.split_source;
+            _split_source = std::make_shared<vectorized::RemoteSplitSourceConnector>(
+                    state, split_source.split_source_id, split_source.num_splits);
         }
-        for (int i = num_add_one; i < max_scanners; ++i) {
-            _scan_ranges[scan_index] = scan_ranges[range_index++];
-            auto& ranges =
-                    _scan_ranges[scan_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-            for (int j = 0; j < num_ranges - 1; j++) {
-                auto& merged_ranges =
-                        scan_ranges[range_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-                ranges.insert(ranges.end(), merged_ranges.begin(), merged_ranges.end());
-            }
-        }
-        LOG(INFO) << "Merge " << scan_ranges.size() << " scan ranges to " << _scan_ranges.size();
     }
+    if (_split_source == nullptr) {
+        _split_source = std::make_shared<vectorized::LocalSplitSourceConnector>(scan_ranges);
+    }
+    _max_scanners = std::min(_max_scanners, _split_source->num_scan_ranges());
     if (scan_ranges.size() > 0 &&
         scan_ranges[0].scan_range.ext_scan_range.file_scan_range.__isset.params) {
         // for compatibility.

be/src/pipeline/exec/file_scan_operator.h

@@ -26,6 +26,7 @@
 #include "operator.h"
 #include "pipeline/exec/scan_operator.h"
 #include "vec/exec/format/format_common.h"
+#include "vec/exec/scan/split_source_connector.h"
 
 namespace doris {
 namespace vectorized {
@@ -54,7 +55,8 @@ class FileScanLocalState final : public ScanLocalState<FileScanLocalState> {
     std::string name_suffix() const override;
 
 private:
-    std::vector<TScanRangeParams> _scan_ranges;
+    std::shared_ptr<vectorized::SplitSourceConnector> _split_source = nullptr;
+    int _max_scanners;
     // A in memory cache to save some common components
     // of the this scan node. eg:
     // 1. iceberg delete file

be/src/vec/exec/scan/new_file_scan_node.cpp

@@ -26,6 +26,7 @@
 
 #include "common/config.h"
 #include "common/object_pool.h"
+#include "runtime/client_cache.h"
 #include "vec/exec/scan/vfile_scanner.h"
 #include "vec/exec/scan/vscanner.h"
 
@@ -36,6 +37,8 @@ class RuntimeState;
 
 namespace doris::vectorized {
 
+using apache::thrift::transport::TTransportException;
+
 NewFileScanNode::NewFileScanNode(ObjectPool* pool, const TPlanNode& tnode,
                                  const DescriptorTbl& descs)
         : VScanNode(pool, tnode, descs) {
@@ -60,47 +63,25 @@ Status NewFileScanNode::prepare(RuntimeState* state) {
 
 void NewFileScanNode::set_scan_ranges(RuntimeState* state,
                                       const std::vector<TScanRangeParams>& scan_ranges) {
-    int max_scanners =
+    _max_scanners =
             config::doris_scanner_thread_pool_thread_num / state->query_parallel_instance_num();
-    max_scanners = std::max(std::max(max_scanners, state->parallel_scan_max_scanners_count()), 1);
+    _max_scanners = std::max(std::max(_max_scanners, state->parallel_scan_max_scanners_count()), 1);
     // For select * from table limit 10; should just use one thread.
     if (should_run_serial()) {
-        max_scanners = 1;
+        _max_scanners = 1;
     }
-    if (scan_ranges.size() <= max_scanners) {
-        _scan_ranges = scan_ranges;
-    } else {
-        // There is no need for the number of scanners to exceed the number of threads in thread pool.
-        // scan_ranges is sorted by path(as well as partition path) in FE, so merge scan ranges in order.
-        // In the insert statement, reading data in partition order can reduce the memory usage of BE
-        // and prevent the generation of smaller tables.
-        _scan_ranges.resize(max_scanners);
-        int num_ranges = scan_ranges.size() / max_scanners;
-        int num_add_one = scan_ranges.size() - num_ranges * max_scanners;
-        int scan_index = 0;
-        int range_index = 0;
-        for (int i = 0; i < num_add_one; ++i) {
-            _scan_ranges[scan_index] = scan_ranges[range_index++];
-            auto& ranges =
-                    _scan_ranges[scan_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-            for (int j = 0; j < num_ranges; j++) {
-                auto& merged_ranges =
-                        scan_ranges[range_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-                ranges.insert(ranges.end(), merged_ranges.begin(), merged_ranges.end());
-            }
-        }
-        for (int i = num_add_one; i < max_scanners; ++i) {
-            _scan_ranges[scan_index] = scan_ranges[range_index++];
-            auto& ranges =
-                    _scan_ranges[scan_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-            for (int j = 0; j < num_ranges - 1; j++) {
-                auto& merged_ranges =
-                        scan_ranges[range_index++].scan_range.ext_scan_range.file_scan_range.ranges;
-                ranges.insert(ranges.end(), merged_ranges.begin(), merged_ranges.end());
-            }
+    if (scan_ranges.size() == 1) {
+        auto scan_range = scan_ranges[0].scan_range.ext_scan_range.file_scan_range;
+        if (scan_range.__isset.split_source) {
+            auto split_source = scan_range.split_source;
+            _split_source = std::make_shared<RemoteSplitSourceConnector>(
+                    state, split_source.split_source_id, split_source.num_splits);
         }
-        LOG(INFO) << "Merge " << scan_ranges.size() << " scan ranges to " << _scan_ranges.size();
     }
+    if (_split_source == nullptr) {
+        _split_source = std::make_shared<LocalSplitSourceConnector>(scan_ranges);
+    }
+    _max_scanners = std::min(_max_scanners, _split_source->num_scan_ranges());
     if (scan_ranges.size() > 0 &&
         scan_ranges[0].scan_range.ext_scan_range.file_scan_range.__isset.params) {
         // for compatibility.
@@ -125,20 +106,19 @@ Status NewFileScanNode::_process_conjuncts() {
 }
 
 Status NewFileScanNode::_init_scanners(std::list<VScannerSPtr>* scanners) {
-    if (_scan_ranges.empty()) {
+    if (_split_source->num_scan_ranges() == 0) {
         _eos = true;
         return Status::OK();
     }
 
     size_t shard_num = std::min<size_t>(
             config::doris_scanner_thread_pool_thread_num / _state->query_parallel_instance_num(),
-            _scan_ranges.size());
+            _max_scanners);
     shard_num = std::max(shard_num, (size_t)1);
     _kv_cache.reset(new ShardedKVCache(shard_num));
-    for (auto& scan_range : _scan_ranges) {
+    for (int i = 0; i < _max_scanners; ++i) {
         std::unique_ptr<VFileScanner> scanner =
-                VFileScanner::create_unique(_state, this, _limit_per_scanner,
-                                            scan_range.scan_range.ext_scan_range.file_scan_range,
+                VFileScanner::create_unique(_state, this, _limit_per_scanner, _split_source,
                                             runtime_profile(), _kv_cache.get());
         RETURN_IF_ERROR(
                 scanner->prepare(_conjuncts, &_colname_to_value_range, &_colname_to_slot_id));

be/src/vec/exec/scan/new_file_scan_node.h

@@ -25,6 +25,7 @@
 
 #include "common/status.h"
 #include "vec/exec/format/format_common.h"
+#include "vec/exec/scan/split_source_connector.h"
 #include "vec/exec/scan/vscan_node.h"
 
 namespace doris {
@@ -58,7 +59,8 @@ class NewFileScanNode : public VScanNode {
     Status _init_scanners(std::list<VScannerSPtr>* scanners) override;
 
 private:
-    std::vector<TScanRangeParams> _scan_ranges;
+    std::shared_ptr<SplitSourceConnector> _split_source = nullptr;
+    int _max_scanners;
     // A in memory cache to save some common components
     // of the this scan node. eg:
     // 1. iceberg delete file

be/src/vec/exec/scan/split_source_connector.cpp

@@ -0,0 +1,86 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "vec/exec/scan/split_source_connector.h"
+
+#include "runtime/exec_env.h"
+#include "runtime/query_context.h"
+
+namespace doris::vectorized {
+
+using apache::thrift::transport::TTransportException;
+
+Status LocalSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* range) {
+    std::lock_guard<std::mutex> l(_range_lock);
+    *has_next = false;
+    if (_scan_index < _scan_ranges.size()) {
+        auto& ranges = _scan_ranges[_scan_index].scan_range.ext_scan_range.file_scan_range.ranges;
+        if (_range_index < ranges.size()) {
+            *has_next = true;
+            *range = ranges[_range_index++];
+            if (_range_index == ranges.size()) {
+                _scan_index++;
+                _range_index = 0;
+            }
+        }
+    }
+    return Status::OK();
+}
+
+Status RemoteSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* range) {
+    std::lock_guard<std::mutex> l(_range_lock);
+    *has_next = false;
+    if (_scan_index == _scan_ranges.size() && !_last_batch) {
+        Status coord_status;
+        FrontendServiceConnection coord(_state->exec_env()->frontend_client_cache(),
+                                        _state->get_query_ctx()->coord_addr, &coord_status);
+        RETURN_IF_ERROR(coord_status);
+        TFetchSplitBatchRequest request;
+        request.__set_split_source_id(_split_source_id);
+        request.__set_max_num_splits(config::remote_split_source_batch_size);
+        TFetchSplitBatchResult result;
+        try {
+            coord->fetchSplitBatch(result, request);
+        } catch (std::exception& e1) {
+            LOG(WARNING) << "Failed to get batch of split source: {}, try to reopen" << e1.what();
+            RETURN_IF_ERROR(coord.reopen());
+            try {
+                coord->fetchSplitBatch(result, request);
+            } catch (std::exception& e2) {
+                return Status::IOError("Failed to get batch of split source: {}", e2.what());
+            }
+        }
+        _last_batch = result.splits.empty();
+        _scan_ranges = result.splits;
+        _scan_index = 0;
+        _range_index = 0;
+    }
+    if (_scan_index < _scan_ranges.size()) {
+        auto& ranges = _scan_ranges[_scan_index].scan_range.ext_scan_range.file_scan_range.ranges;
+        if (_range_index < ranges.size()) {
+            *has_next = true;
+            *range = ranges[_range_index++];
+            if (_range_index == ranges.size()) {
+                _scan_index++;
+                _range_index = 0;
+            }
+        }
+    }
+    return Status::OK();
+}
+
+} // namespace doris::vectorized