branch-4.2: [improvement](fe) Default external scans to bounded split batches (#69683)

be/src/exec/operator/file_scan_operator.cpp

@@ -69,6 +69,22 @@ FileScanLocalState::FileScanLocalState(RuntimeState* state, OperatorXBase* paren
 
 FileScanLocalState::~FileScanLocalState() = default;
 
+Status FileScanLocalState::close(RuntimeState* state) {
+    if (_closed) {
+        return Status::OK();
+    }
+    Status status = Base::close(state);
+    // Base::close stops every scanner and releases ScannerContext's transfer lock before RPC.
+    // A pruned/failed initialization can close before constructing a split source.
+    if (_split_source != nullptr) {
+        Status source_status = _split_source->close();
+        if (status.ok()) {
+            status = source_status;
+        }
+    }
+    return status;
+}
+
 PushDownType FileScanLocalState::_should_push_down_binary_predicate(
         VectorizedFnCall* fn_call, VExprContext* expr_ctx, Field& constant_val,
         const std::set<std::string> fn_name) const {
@@ -336,7 +352,8 @@ void FileScanLocalState::set_scan_ranges(RuntimeState* state,
             _max_scanners = calc_max_scanners(p.parallelism(state));
             _split_source = std::make_shared<RemoteSplitSourceConnector>(
                     state, get_split_timer, split_source.split_source_id, split_source.num_splits,
-                    _max_scanners);
+                    _max_scanners,
+                    split_source.__isset.supports_close && split_source.supports_close);
         }
     }
 

be/src/exec/operator/file_scan_operator.h

@@ -50,6 +50,7 @@ class FileScanLocalState final : public ScanLocalState<FileScanLocalState> {
     ~FileScanLocalState() override;
 
     Status init(RuntimeState* state, LocalStateInfo& info) override;
+    Status close(RuntimeState* state) override;
 
     Status _process_conjuncts(RuntimeState* state) override;
     Status _init_scanners(std::list<ScannerSPtr>* scanners) override;

be/src/exec/scan/file_scanner.cpp

@@ -482,7 +482,7 @@ void FileScanner::_get_slot_ids(VExpr* expr, std::vector<int>* slot_ids) {
 Status FileScanner::_open_impl(RuntimeState* state) {
     RETURN_IF_CANCELLED(state);
     RETURN_IF_ERROR(Scanner::_open_impl(state));
-    RETURN_IF_ERROR(_split_source->get_next(&_first_scan_range, &_current_range));
+    RETURN_IF_ERROR(_split_source->get_next(&_first_scan_range, &_current_range, &_should_stop));
     if (_first_scan_range) {
         RETURN_IF_ERROR(_init_expr_ctxes());
         if (_state->query_options().enable_runtime_filter_partition_prune &&
@@ -1048,7 +1048,7 @@ Status FileScanner::_get_next_reader() {
         _src_block_init = false;
         bool has_next = _first_scan_range;
         if (!_first_scan_range) {
-            RETURN_IF_ERROR(_split_source->get_next(&has_next, &_current_range));
+            RETURN_IF_ERROR(_split_source->get_next(&has_next, &_current_range, &_should_stop));
         }
         _first_scan_range = false;
         if (!has_next || _should_stop) {

be/src/exec/scan/scanner.h

@@ -20,6 +20,7 @@
 #include <stdint.h>
 
 #include <algorithm>
+#include <atomic>
 #include <vector>
 
 #include "common/status.h"
@@ -275,7 +276,7 @@ class Scanner {
     int64_t _per_scanner_timer = 0;
     int64_t _projection_timer = 0;
 
-    bool _should_stop = false;
+    std::atomic_bool _should_stop = false;
 };
 
 using ScannerSPtr = std::shared_ptr<Scanner>;

be/src/exec/scan/split_source_connector.cpp

@@ -19,6 +19,7 @@
 
 #include "runtime/exec_env.h"
 #include "runtime/query_context.h"
+#include "util/thrift_rpc_helper.h"
 
 namespace doris {
 
@@ -130,7 +131,8 @@ void SplitSourceConnector::stop() {
     _split_ready.notify_all();
 }
 
-Status LocalSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* range) {
+Status LocalSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* range,
+                                           const std::atomic_bool* /*scanner_stop*/) {
     std::lock_guard<std::mutex> l(_range_lock);
     *has_next = false;
     if (_scan_index < _scan_ranges.size()) {
@@ -147,31 +149,42 @@ Status LocalSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* range
     return Status::OK();
 }
 
-Status RemoteSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* range) {
+Status RemoteSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* range,
+                                            const std::atomic_bool* scanner_stop) {
     std::lock_guard<std::mutex> l(_range_lock);
     *has_next = false;
+    if (is_stopped() || (scanner_stop && scanner_stop->load(std::memory_order_acquire))) {
+        return Status::OK();
+    }
     if (_scan_index == _scan_ranges.size() && !_last_batch) {
         SCOPED_TIMER(_get_split_timer);
-        Status coord_status;
-        // No need to set timeout because on FE side, there is a max fetch time
-        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);
+        do {
+            RETURN_IF_CANCELLED(_state);
+            // A local LIMIT stops this scanner without cancelling the query. Do not keep a scan
+            // worker in the retry loop while another branch of the same query remains active.
+            if (is_stopped() || (scanner_stop && scanner_stop->load(std::memory_order_acquire))) {
+                return Status::OK();
+            }
+            const auto fetch_status = _fetch_batch(&result, request);
+            if (is_stopped() || (scanner_stop && scanner_stop->load(std::memory_order_acquire))) {
+                return Status::OK();
+            }
+            RETURN_IF_ERROR(fetch_status);
+            if (result.__isset.status && result.status.status_code == TStatusCode::INCOMPLETE) {
+                continue;
+            }
             if (result.__isset.status && result.status.status_code != TStatusCode::OK) {
                 return Status::IOError<false>("Failed to get batch of split source: {}",
                                               result.status.error_msgs.empty()
                                                       ? "unknown error"
                                                       : result.status.error_msgs[0]);
             }
-        } catch (std::exception& e) {
-            return Status::IOError<false>("Failed to get batch of split source: {}", e.what());
-        }
+            break;
+        } while (true);
         _last_batch = result.splits.empty();
         _merge_ranges<TScanRangeLocations>(_scan_ranges, result.splits);
         _scan_index = 0;
@@ -191,4 +204,62 @@ Status RemoteSplitSourceConnector::get_next(bool* has_next, TFileRangeDesc* rang
     return Status::OK();
 }
 
+Status RemoteSplitSourceConnector::_fetch_batch(TFetchSplitBatchResult* result,
+                                                const TFetchSplitBatchRequest& request) {
+    if (request.__isset.close_source && request.close_source) {
+        // FE closure is idempotent, including when the first reply is lost. Reuse the standard
+        // reconnect/retry path only here: an ordinary split fetch consumes queue entries.
+        const auto& address = _state->get_query_ctx()->coord_addr;
+        return ThriftRpcHelper::rpc<FrontendServiceClient>(
+                address.hostname, address.port,
+                [&request, result](FrontendServiceConnection& coord) {
+                    coord->fetchSplitBatch(*result, request);
+                });
+    }
+    Status status;
+    FrontendServiceConnection coord(_state->exec_env()->frontend_client_cache(),
+                                    _state->get_query_ctx()->coord_addr, &status);
+    RETURN_IF_ERROR(status);
+    try {
+        coord->fetchSplitBatch(*result, request);
+    } catch (const std::exception& e) {
+        return Status::IOError<false>("Failed to get batch of split source: {}", e.what());
+    }
+    return Status::OK();
+}
+
+Status RemoteSplitSourceConnector::close() {
+    stop();
+    // Older FEs do not advertise support for the source-close request.
+    if (!_supports_close) {
+        return Status::OK();
+    }
+    std::scoped_lock lock(_close_lock);
+    if (_close_completed) {
+        return Status::OK();
+    }
+    TFetchSplitBatchRequest request;
+    request.__set_split_source_id(_split_source_id);
+    request.__set_max_num_splits(config::remote_split_source_batch_size);
+    request.__set_close_source(true);
+    Status status;
+    // The owner normally closes once. Retry a failed close here, without latching an unacknowledged
+    // request. An exhausted error must cancel the query rather than strand other BEs behind a queue.
+    for (int attempt = 0; attempt < 2; ++attempt) {
+        TFetchSplitBatchResult result;
+        status = _fetch_batch(&result, request);
+        if (status.ok() && result.__isset.status && result.status.status_code != TStatusCode::OK) {
+            status = Status::IOError<false>("Failed to close split source {}: {}", _split_source_id,
+                                            result.status.error_msgs.empty()
+                                                    ? "unknown error"
+                                                    : result.status.error_msgs[0]);
+        }
+        if (status.ok()) {
+            _close_completed = true;
+            return Status::OK();
+        }
+    }
+    return status;
+}
+
 } // namespace doris

be/src/exec/scan/split_source_connector.h

@@ -18,13 +18,16 @@
 #pragma once
 
 #include <algorithm>
+#include <atomic>
 #include <condition_variable>
 #include <deque>
+#include <mutex>
 #include <unordered_set>
 
 #include "common/config.h"
 #include "core/custom_allocator.h"
 #include "format_v2/file_scan_context.h"
+#include "gen_cpp/FrontendService_types.h"
 #include "runtime/runtime_state.h"
 #include "util/client_cache.h"
 
@@ -46,7 +49,8 @@ class SplitSourceConnector {
      * @param has_next whether exists the next scan range
      * @param range the obtained next scan range
      */
-    virtual Status get_next(bool* has_next, TFileRangeDesc* range) = 0;
+    virtual Status get_next(bool* has_next, TFileRangeDesc* range,
+                            const std::atomic_bool* scanner_stop = nullptr) = 0;
 
     // FileScannerV2 may replace one FE split with several physical children after inspecting file
     // metadata. These methods are shared by local and remote sources so neither source can report
@@ -55,6 +59,11 @@ class SplitSourceConnector {
     Status finish_source_split(const FileScanSplit& source_split,
                                std::vector<FileScanSplit> generated_splits);
     void stop();
+    // Called once by the scan operator owner after stopping all shared consumers.
+    virtual Status close() {
+        stop();
+        return Status::OK();
+    }
 
     virtual int num_scan_ranges() = 0;
 
@@ -63,6 +72,8 @@ class SplitSourceConnector {
     virtual bool all_ranges_have_table_level_row_count() const { return false; }
 
 protected:
+    bool is_stopped() const { return _stopped.load(std::memory_order_acquire); }
+
     template <typename T, typename V1 = std::vector<T>, typename V2 = std::vector<T>>
         requires(std::is_same_v<std::remove_cvref_t<V1>,
                                 std::vector<T, typename V1::allocator_type>> &&
@@ -117,7 +128,7 @@ class SplitSourceConnector {
     uint64_t _next_source_split_id = 1;
     bool _source_claim_in_progress = false;
     bool _source_exhausted = false;
-    bool _stopped = false;
+    std::atomic_bool _stopped = false;
 };
 
 /**
@@ -139,7 +150,8 @@ class LocalSplitSourceConnector : public SplitSourceConnector {
         _merge_ranges<TScanRangeParams>(_scan_ranges, scan_ranges);
     }
 
-    Status get_next(bool* has_next, TFileRangeDesc* range) override;
+    Status get_next(bool* has_next, TFileRangeDesc* range,
+                    const std::atomic_bool* scanner_stop = nullptr) override;
 
     int num_scan_ranges() override { return static_cast<int>(_scan_ranges.size()); }
 
@@ -181,6 +193,10 @@ class RemoteSplitSourceConnector : public SplitSourceConnector {
     RuntimeProfile::Counter* _get_split_timer;
     int64_t _split_source_id;
     int _num_splits;
+    const bool _supports_close;
+    // Closing is independent of a possibly blocked fetch, and completes only after FE acknowledges.
+    std::mutex _close_lock;
+    bool _close_completed = false;
 
     DorisVector<TScanRangeLocations> _scan_ranges;
     bool _last_batch = false;
@@ -189,16 +205,25 @@ class RemoteSplitSourceConnector : public SplitSourceConnector {
 
 public:
     RemoteSplitSourceConnector(RuntimeState* state, RuntimeProfile::Counter* get_split_timer,
-                               int64_t split_source_id, int num_splits, int max_scanners)
+                               int64_t split_source_id, int num_splits, int max_scanners,
+                               bool supports_close = false)
             : _state(state),
               _get_split_timer(get_split_timer),
               _split_source_id(split_source_id),
-              _num_splits(num_splits) {
+              _num_splits(num_splits),
+              _supports_close(supports_close) {
         _max_scanners = max_scanners;
     }
 
-    Status get_next(bool* has_next, TFileRangeDesc* range) override;
+    Status get_next(bool* has_next, TFileRangeDesc* range,
+                    const std::atomic_bool* scanner_stop = nullptr) override;
+    Status close() override;
 
+protected:
+    virtual Status _fetch_batch(TFetchSplitBatchResult* result,
+                                const TFetchSplitBatchRequest& request);
+
+public:
     /*
      * Remote split source is fetched in batch mode, and the splits are generated while scanning,
      * so the number of scan ranges may not be accurate.