This is an automated email from the ASF dual-hosted git repository.
yiguolei pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new fa07b7c85d8 branch-4.2: [improvement](fe) Default external scans to
bounded split batches (#69683) (#69991)
fa07b7c85d8 is described below
commit fa07b7c85d802e5d322387d56d565f3f620a6c76
Author: 924060929 <[email protected]>
AuthorDate: Tue Nov 3 11:47:09 2026 +0800
branch-4.2: [improvement](fe) Default external scans to bounded split
batches (#69683) (#69991)
cherry pick from #69683
---
be/src/exec/operator/file_scan_operator.cpp | 19 +-
be/src/exec/operator/file_scan_operator.h | 1 +
be/src/exec/scan/file_scanner.cpp | 4 +-
be/src/exec/scan/scanner.h | 3 +-
be/src/exec/scan/split_source_connector.cpp | 95 +++++++--
be/src/exec/scan/split_source_connector.h | 37 +++-
be/test/exec/scan/file_scanner_v2_test.cpp | 220 ++++++++++++++++++++-
be/test/exec/scan/vfile_scanner_exception_test.cpp | 3 +-
be/test/format/wal/wal_manager_test.cpp | 3 +-
.../apache/doris/blockrule/SqlBlockRuleMgr.java | 24 +++
.../apache/doris/datasource/FileQueryScanNode.java | 31 ++-
.../apache/doris/datasource/SplitAssignment.java | 85 ++++++--
.../org/apache/doris/datasource/SplitSource.java | 62 +++---
.../doris/datasource/hive/source/HiveScanNode.java | 39 ++--
.../doris/datasource/hudi/source/HudiScanNode.java | 3 +
.../datasource/iceberg/source/IcebergScanNode.java | 2 +-
.../maxcompute/source/MaxComputeScanNode.java | 108 +++++-----
.../java/org/apache/doris/qe/SessionVariable.java | 36 +++-
.../apache/doris/service/FrontendServiceImpl.java | 12 +-
.../doris/blockrule/SqlBlockRuleMgrTest.java | 39 ++++
.../doris/datasource/BoundedSplitDeliveryTest.java | 148 ++++++++++++++
.../doris/datasource/FileQueryScanNodeTest.java | 16 ++
.../doris/datasource/SplitAssignmentTest.java | 7 +-
.../doris/datasource/SplitConsumerCloseTest.java | 136 +++++++++++++
.../datasource/hive/source/HiveScanNodeTest.java | 77 ++++++++
.../hive/source/HiveSplitEstimateTest.java | 216 ++++++++++++++++++++
.../iceberg/source/IcebergScanNodeTest.java | 88 +++++++++
.../source/MaxComputePartitionBatchTaskTest.java | 94 +++++++++
gensrc/thrift/FrontendService.thrift | 2 +
gensrc/thrift/PlanNodes.thrift | 2 +
.../hive/test_hive_partitions.groovy | 9 +-
.../iceberg/test_iceberg_variant_read.groovy | 2 +
.../write/test_iceberg_write_uuid_inputs.groovy | 1 +
33 files changed, 1484 insertions(+), 140 deletions(-)
diff --git a/be/src/exec/operator/file_scan_operator.cpp
b/be/src/exec/operator/file_scan_operator.cpp
index b6762c23fe6..bb078b509cd 100644
--- a/be/src/exec/operator/file_scan_operator.cpp
+++ b/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);
}
}
diff --git a/be/src/exec/operator/file_scan_operator.h
b/be/src/exec/operator/file_scan_operator.h
index 03c0c980d93..564ee16c012 100644
--- a/be/src/exec/operator/file_scan_operator.h
+++ b/be/src/exec/operator/file_scan_operator.h
@@ -50,6 +50,7 @@ public:
~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;
diff --git a/be/src/exec/scan/file_scanner.cpp
b/be/src/exec/scan/file_scanner.cpp
index 0a172da4216..b273ec7aad9 100644
--- a/be/src/exec/scan/file_scanner.cpp
+++ b/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) {
diff --git a/be/src/exec/scan/scanner.h b/be/src/exec/scan/scanner.h
index 6aaa1ec6940..9d5775bc163 100644
--- a/be/src/exec/scan/scanner.h
+++ b/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 @@ protected:
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>;
diff --git a/be/src/exec/scan/split_source_connector.cpp
b/be/src/exec/scan/split_source_connector.cpp
index 8642e57aa12..5c177915ac8 100644
--- a/be/src/exec/scan/split_source_connector.cpp
+++ b/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
diff --git a/be/src/exec/scan/split_source_connector.h
b/be/src/exec/scan/split_source_connector.h
index 583d6b9eebe..3ddff03c847 100644
--- a/be/src/exec/scan/split_source_connector.h
+++ b/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 @@ public:
* @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 @@ public:
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 @@ public:
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 @@ private:
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 @@ public:
_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 @@ private:
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 @@ private:
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.
diff --git a/be/test/exec/scan/file_scanner_v2_test.cpp
b/be/test/exec/scan/file_scanner_v2_test.cpp
index d7d11e0acbe..01817ac9c24 100644
--- a/be/test/exec/scan/file_scanner_v2_test.cpp
+++ b/be/test/exec/scan/file_scanner_v2_test.cpp
@@ -23,6 +23,7 @@
#include <atomic>
#include <chrono>
#include <condition_variable>
+#include <functional>
#include <future>
#include <memory>
#include <mutex>
@@ -570,7 +571,8 @@ public:
explicit RemoteStyleSplitSourceConnector(std::vector<TFileRangeDesc>
ranges)
: _ranges(std::move(ranges)) {}
- Status get_next(bool* has_next, TFileRangeDesc* range) override {
+ Status get_next(bool* has_next, TFileRangeDesc* range,
+ const std::atomic_bool* /*scanner_stop*/) override {
std::lock_guard lock(_lock);
++_get_next_calls;
*has_next = _next < _ranges.size();
@@ -599,7 +601,8 @@ public:
explicit BlockingRemoteStyleSplitSourceConnector(TFileRangeDesc range)
: _range(std::move(range)) {}
- Status get_next(bool* has_next, TFileRangeDesc* range) override {
+ Status get_next(bool* has_next, TFileRangeDesc* range,
+ const std::atomic_bool* /*scanner_stop*/) override {
std::unique_lock lock(_lock);
if (!_source_returned) {
_source_returned = true;
@@ -1587,4 +1590,217 @@ TEST(FileScannerTest,
PartitionPruningStopsAtUnsafePredicate) {
EXPECT_EQ(partition_conjuncts[0], conjuncts[0]);
}
+TEST(FileScannerV2Test, RemoteSplitSourceStopsLocallyWithoutQueryCancellation)
{
+ RuntimeState state {TQueryOptions(), TQueryGlobals()};
+ RemoteSplitSourceConnector source(&state, nullptr, 1, 1, 1);
+ std::atomic_bool scanner_stop {true};
+ bool has_next = true;
+ TFileRangeDesc range;
+
+ ASSERT_FALSE(state.is_cancelled());
+ ASSERT_TRUE(source.get_next(&has_next, &range, &scanner_stop).ok());
+ EXPECT_FALSE(has_next);
+}
+
+class ScriptedRemoteSplitSource final : public RemoteSplitSourceConnector {
+public:
+ explicit ScriptedRemoteSplitSource(RuntimeState* state, bool
supports_close = true)
+ : RemoteSplitSourceConnector(state, nullptr, 7, 1, 1,
supports_close) {}
+
+ std::function<Status(TFetchSplitBatchResult*, const
TFetchSplitBatchRequest&)> fetch;
+
+protected:
+ Status _fetch_batch(TFetchSplitBatchResult* result,
+ const TFetchSplitBatchRequest& request) override {
+ return fetch(result, request);
+ }
+};
+
+TEST(FileScannerV2Test,
RemoteIncompleteRetriesAndOwnerClosesTheSharedSourceOnce) {
+ RuntimeState state {TQueryOptions(), TQueryGlobals()};
+ ScriptedRemoteSplitSource source(&state);
+ int fetched = 0;
+ int closed = 0;
+ source.fetch = [&](TFetchSplitBatchResult* result, const
TFetchSplitBatchRequest& request) {
+ if (request.__isset.close_source && request.close_source) {
+ ++closed;
+ EXPECT_EQ(request.split_source_id, 7);
+ EXPECT_GT(request.max_num_splits, 0);
+ TStatus thrift_status;
+ thrift_status.__set_status_code(TStatusCode::OK);
+ result->__set_status(thrift_status);
+ return Status::OK();
+ }
+ if (++fetched == 1) {
+ TStatus thrift_status;
+ thrift_status.__set_status_code(TStatusCode::INCOMPLETE);
+ result->__set_status(thrift_status);
+ } else {
+ TScanRangeLocations locations;
+ locations.scan_range =
scan_range_with_path("file.parquet").scan_range;
+
locations.scan_range.ext_scan_range.file_scan_range.ranges.push_back(
+
locations.scan_range.ext_scan_range.file_scan_range.ranges.front());
+ result->__set_splits({locations});
+ TStatus thrift_status;
+ thrift_status.__set_status_code(TStatusCode::OK);
+ result->__set_status(thrift_status);
+ }
+ return Status::OK();
+ };
+ bool has_next = false;
+ TFileRangeDesc range;
+ ASSERT_TRUE(source.get_next(&has_next, &range).ok());
+ EXPECT_TRUE(has_next);
+ EXPECT_EQ(fetched, 2);
+ EXPECT_EQ(range.path, "file.parquet");
+ ASSERT_TRUE(source.close().ok());
+ ASSERT_TRUE(source.close().ok());
+ EXPECT_EQ(closed, 1);
+ ASSERT_TRUE(source.get_next(&has_next, &range).ok());
+ EXPECT_FALSE(has_next);
+ EXPECT_EQ(fetched, 2);
+}
+
+TEST(FileScannerV2Test,
OldFrontendDescriptorDoesNotSendAnUnsupportedCloseRequest) {
+ RuntimeState state {TQueryOptions(), TQueryGlobals()};
+ ScriptedRemoteSplitSource source(&state, false);
+ source.fetch = [](TFetchSplitBatchResult*, const TFetchSplitBatchRequest&)
{
+ ADD_FAILURE() << "an old FE did not advertise source closure";
+ return Status::InternalError("unexpected close");
+ };
+ EXPECT_TRUE(source.close().ok());
+}
+
+TEST(FileScannerV2Test, OldFrontendTimeoutRemainsAnErrorWithoutRetry) {
+ RuntimeState state {TQueryOptions(), TQueryGlobals()};
+ ScriptedRemoteSplitSource source(&state, false);
+ int fetched = 0;
+ source.fetch = [&](TFetchSplitBatchResult* result, const
TFetchSplitBatchRequest& request) {
+ EXPECT_FALSE(request.__isset.close_source);
+ if (++fetched > 1) {
+ return Status::InternalError("the old FE timeout must not be
retried");
+ }
+ TStatus thrift_status;
+ thrift_status.__set_status_code(TStatusCode::INTERNAL_ERROR);
+ thrift_status.__set_error_msgs({"Timeout. Max wait time(ms): 1000"});
+ result->__set_status(thrift_status);
+ return Status::OK();
+ };
+ bool has_next = true;
+ TFileRangeDesc range;
+ Status status = source.get_next(&has_next, &range);
+ EXPECT_FALSE(status.ok());
+ EXPECT_NE(status.to_string().find("Timeout. Max wait time(ms): 1000"),
std::string::npos);
+ EXPECT_FALSE(has_next);
+ EXPECT_EQ(fetched, 1);
+ EXPECT_TRUE(source.close().ok());
+ EXPECT_EQ(fetched, 1);
+}
+
+TEST(FileScannerV2Test,
OwnerCloseDoesNotWaitForTheFetchLockOrSurfaceALateFetchError) {
+ RuntimeState state {TQueryOptions(), TQueryGlobals()};
+ ScriptedRemoteSplitSource source(&state);
+ std::promise<void> entered;
+ auto started = entered.get_future();
+ std::promise<void> release;
+ auto released = release.get_future().share();
+ std::atomic_int closed {0};
+ source.fetch = [&](TFetchSplitBatchResult* result, const
TFetchSplitBatchRequest& request) {
+ if (request.__isset.close_source && request.close_source) {
+ ++closed;
+ TStatus thrift_status;
+ thrift_status.__set_status_code(TStatusCode::OK);
+ result->__set_status(thrift_status);
+ return Status::OK();
+ }
+ entered.set_value();
+ released.wait();
+ return Status::IOError("the source closed while this fetch was in
flight");
+ };
+ bool has_next = true;
+ TFileRangeDesc range;
+ auto fetching =
+ std::async(std::launch::async, [&] { return
source.get_next(&has_next, &range); });
+ EXPECT_EQ(started.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ auto closing = std::async(std::launch::async, [&] { return source.close();
});
+ EXPECT_EQ(closing.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ release.set_value();
+ EXPECT_TRUE(closing.get().ok());
+ EXPECT_EQ(fetching.wait_for(std::chrono::seconds(5)),
std::future_status::ready);
+ EXPECT_TRUE(fetching.get().ok());
+ EXPECT_FALSE(has_next);
+ EXPECT_EQ(closed.load(), 1);
+}
+
+TEST(FileScannerV2Test, OwnerCloseRetriesAnUnacknowledgedRequest) {
+ for (bool lost_reply : {false, true}) {
+ RuntimeState state {TQueryOptions(), TQueryGlobals()};
+ ScriptedRemoteSplitSource source(&state);
+ int calls = 0;
+ int effects = 0;
+ bool released = false;
+ source.fetch = [&](TFetchSplitBatchResult* result, const
TFetchSplitBatchRequest& request) {
+ EXPECT_TRUE(request.close_source);
+ ++calls;
+ if (!lost_reply && calls == 1) {
+ return Status::IOError("connection failed before FE received
the close");
+ }
+ if (!released) {
+ released = true;
+ ++effects;
+ }
+ if (lost_reply && calls == 1) {
+ return Status::IOError("FE closed the source but the reply was
lost");
+ }
+ TStatus thrift_status;
+ thrift_status.__set_status_code(TStatusCode::OK);
+ result->__set_status(thrift_status);
+ return Status::OK();
+ };
+ EXPECT_TRUE(source.close().ok());
+ EXPECT_EQ(calls, 2);
+ EXPECT_EQ(effects, 1);
+ EXPECT_TRUE(source.close().ok());
+ EXPECT_EQ(calls, 2);
+ }
+}
+
+namespace {
+
+void verify_failed_close_can_recover(bool transport_failure) {
+ RuntimeState state {TQueryOptions(), TQueryGlobals()};
+ ScriptedRemoteSplitSource source(&state);
+ bool recovered = false;
+ int calls = 0;
+ source.fetch = [&](TFetchSplitBatchResult* result, const
TFetchSplitBatchRequest& request) {
+ EXPECT_TRUE(request.close_source);
+ ++calls;
+ if (!recovered && transport_failure) {
+ return Status::IOError("FE unavailable");
+ }
+ TStatus thrift_status;
+ thrift_status.__set_status_code(recovered ? TStatusCode::OK :
TStatusCode::INTERNAL_ERROR);
+ result->__set_status(thrift_status);
+ return Status::OK();
+ };
+ EXPECT_FALSE(source.close().ok());
+ EXPECT_EQ(calls, 2);
+ bool has_next = true;
+ TFileRangeDesc range;
+ EXPECT_TRUE(source.get_next(&has_next, &range).ok());
+ EXPECT_FALSE(has_next);
+ recovered = true;
+ EXPECT_TRUE(source.close().ok());
+ EXPECT_EQ(calls, 3);
+ EXPECT_TRUE(source.close().ok());
+ EXPECT_EQ(calls, 3);
+}
+
+} // namespace
+
+TEST(FileScannerV2Test, FailedOwnerCloseIsBoundedAndCanBeRetriedAfterRecovery)
{
+ verify_failed_close_can_recover(false);
+ verify_failed_close_can_recover(true);
+}
+
} // namespace doris
diff --git a/be/test/exec/scan/vfile_scanner_exception_test.cpp
b/be/test/exec/scan/vfile_scanner_exception_test.cpp
index a68002c24c6..605b54a8d7c 100644
--- a/be/test/exec/scan/vfile_scanner_exception_test.cpp
+++ b/be/test/exec/scan/vfile_scanner_exception_test.cpp
@@ -54,7 +54,8 @@ private:
public:
TestSplitSourceConnectorStub(const TFileScanRange& scan_range) :
_scan_range(scan_range) {}
- Status get_next(bool* has_next, TFileRangeDesc* range) override {
+ Status get_next(bool* has_next, TFileRangeDesc* range,
+ const std::atomic_bool* /*scanner_stop*/) override {
std::lock_guard<std::mutex> l(_range_lock);
if (_range_index < _scan_range.ranges.size()) {
*has_next = true;
diff --git a/be/test/format/wal/wal_manager_test.cpp
b/be/test/format/wal/wal_manager_test.cpp
index a6dabd57bde..6f48e0b3936 100644
--- a/be/test/format/wal/wal_manager_test.cpp
+++ b/be/test/format/wal/wal_manager_test.cpp
@@ -47,7 +47,8 @@ private:
public:
TestSplitSourceConnector(const TFileScanRange& scan_range) :
_scan_range(scan_range) {}
- Status get_next(bool* has_next, TFileRangeDesc* range) override {
+ Status get_next(bool* has_next, TFileRangeDesc* range,
+ const std::atomic_bool* /*scanner_stop*/) override {
std::lock_guard<std::mutex> l(_range_lock);
if (_range_index < _scan_range.ranges.size()) {
*has_next = true;
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/blockrule/SqlBlockRuleMgr.java
b/fe/fe-core/src/main/java/org/apache/doris/blockrule/SqlBlockRuleMgr.java
index f68307a3d2a..547e75ce45a 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/blockrule/SqlBlockRuleMgr.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/blockrule/SqlBlockRuleMgr.java
@@ -267,6 +267,30 @@ public class SqlBlockRuleMgr implements Writable {
}
}
+ /** Whether this query needs exact partition/split counts before execution
can be admitted. */
+ public boolean hasScanCountLimitations(String user) {
+ if (ConnectContext.get().getState().isInternal()) {
+ return false;
+ }
+ for (SqlBlockRule rule : nameToSqlBlockRuleMap.values()) {
+ if (rule.getGlobal() && requiresExactScanCounts(rule)) {
+ return true;
+ }
+ }
+ for (String ruleName :
Env.getCurrentEnv().getAuth().getSqlBlockRules(user)) {
+ SqlBlockRule rule = nameToSqlBlockRuleMap.get(ruleName);
+ if (rule != null && requiresExactScanCounts(rule)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ private static boolean requiresExactScanCounts(SqlBlockRule rule) {
+ // Cardinality and partition-filter checks do not depend on
enumerating file splits.
+ return rule.getEnable() && (rule.getPartitionNum() != 0 ||
rule.getTabletNum() != 0);
+ }
+
/**
* Check number whether legal by user.
**/
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java
index 978ea0bd303..e64aebe1c51 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java
@@ -18,6 +18,7 @@
package org.apache.doris.datasource;
import org.apache.doris.analysis.SlotDescriptor;
+import org.apache.doris.analysis.StatementBase;
import org.apache.doris.analysis.TableSample;
import org.apache.doris.analysis.TableScanParams;
import org.apache.doris.analysis.TableSnapshot;
@@ -118,6 +119,7 @@ public abstract class FileQueryScanNode extends
FileScanNode {
protected TableScanParams scanParams;
private Optional<MvccSnapshot> relationSnapshot = Optional.empty();
private boolean relationSnapshotInitialized = false;
+ private Boolean batchModeEnabled;
protected FileSplitter fileSplitter;
protected SummaryProfile summaryProfile;
@@ -149,6 +151,26 @@ public abstract class FileQueryScanNode extends
FileScanNode {
? null : statementContext.getExternalScanTaskCache();
}
+ /** Freeze the planning mode before asynchronous split conversion loses
the query thread context. */
+ protected boolean isBatchModeEnabled() {
+ if (batchModeEnabled == null) {
+ // TABLESAMPLE selects files once across the table: independent
partition batches
+ // would reset the ROWS budget, percentage rounding and random
selection.
+ batchModeEnabled =
sessionVariable.getEnableExternalTableBatchMode() && tableSample == null;
+ ConnectContext context = ConnectContext.get();
+ if (batchModeEnabled && context != null) {
+ StatementContext statementContext =
context.getStatementContext();
+ StatementBase statement = statementContext == null ? null :
statementContext.getParsedStatement();
+ // EXPLAIN VERBOSE needs the full file/delete details.
Scan-count rules run before execution
+ // and must not compare their limits with the batch path's
sample/approximate counts.
+ batchModeEnabled = (statement == null ||
!statement.isVerbose())
+ && !Env.getCurrentEnv().getSqlBlockRuleMgr()
+
.hasScanCountLimitations(context.getQualifiedUser());
+ }
+ }
+ return batchModeEnabled;
+ }
+
/** Whether statement-scoped scan-task reuse is active for this node. */
protected boolean canReuseExternalScanTasks() {
return sessionVariable.enableExternalScanTaskReuse &&
externalScanTaskCache != null;
@@ -476,11 +498,14 @@ public abstract class FileQueryScanNode extends
FileScanNode {
admissionResult = fileCacheAdmissionCheck();
}
- if (isBatchMode()) {
+ if (sessionVariable.getEnableExternalTableBatchMode() &&
isBatchMode()) {
+ long maxWaitTime = sessionVariable.getFetchSplitsMaxWaitTime();
+ Preconditions.checkArgument(maxWaitTime > 0,
"fetch_splits_max_wait_time_ms must be positive");
// File splits are generated lazily, and fetched by backends while
scanning.
// Only provide the unique ID of split source to backend.
splitAssignment = new SplitAssignment(backendPolicy, this,
this::splitToScanRange,
- locationProperties, pathPartitionKeys, admissionResult);
+ locationProperties, pathPartitionKeys, admissionResult,
+ sessionVariable.splitBatchSize,
sessionVariable.splitQueueCapacity);
splitAssignment.init();
if (executor != null) {
executor.getSummaryProfile().setGetSplitsFinishTime();
@@ -495,7 +520,6 @@ public abstract class FileQueryScanNode extends
FileScanNode {
throw new UserException("Approximate split number should not
be negative");
}
totalFileSize = fileSplit.getLength() * selectedSplitNum;
- long maxWaitTime = sessionVariable.getFetchSplitsMaxWaitTime();
// Not accurate, only used to estimate concurrency.
// Here, we must take the max of 1, because
// in the case of multiple BEs, `numApproximateSplits() /
backendPolicy.numBackends()` may be 0,
@@ -509,6 +533,7 @@ public abstract class FileQueryScanNode extends
FileScanNode {
TSplitSource tSource = new TSplitSource();
tSource.setSplitSourceId(splitSource.getUniqueId());
tSource.setNumSplits(numSplitsPerBE);
+ tSource.setSupportsClose(true);
curLocations.getScanRange().getExtScanRange().getFileScanRange().setSplitSource(tSource);
TScanRangeLocation location = new TScanRangeLocation();
location.setBackendId(backend.getId());
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java
index 3d739d1f66a..4f8a62a888a 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java
@@ -22,6 +22,7 @@ import org.apache.doris.spi.Split;
import org.apache.doris.system.Backend;
import org.apache.doris.thrift.TScanRangeLocations;
+import com.google.common.base.Preconditions;
import com.google.common.collect.Multimap;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
@@ -46,20 +47,23 @@ import java.util.concurrent.atomic.AtomicBoolean;
public class SplitAssignment {
private static final Logger LOG =
LogManager.getLogger(SplitAssignment.class);
private final Set<Long> sources = new HashSet<>();
+ private final Set<Long> stoppedBackends = ConcurrentHashMap.newKeySet();
private final FederationBackendPolicy backendPolicy;
private final SplitGenerator splitGenerator;
- private final ConcurrentHashMap<Backend,
BlockingQueue<Collection<TScanRangeLocations>>> assignment
+ private final ConcurrentHashMap<Long, BlockingQueue<TScanRangeLocations>>
assignment
= new ConcurrentHashMap<>();
private final SplitToScanRange splitToScanRange;
private final Map<String, String> locationProperties;
private final List<String> pathPartitionKeys;
private final boolean fileCacheAdmission;
+ private final int splitBatchSize;
+ private final int queueCapacity;
private final Object assignLock = new Object();
private Split sampleSplit = null;
private final AtomicBoolean isStopped = new AtomicBoolean(false);
private final AtomicBoolean scheduleFinished = new AtomicBoolean(false);
- private UserException exception = null;
+ private volatile UserException exception = null;
private final List<Closeable> closeableResources = new ArrayList<>();
public SplitAssignment(
@@ -69,6 +73,16 @@ public class SplitAssignment {
Map<String, String> locationProperties,
List<String> pathPartitionKeys,
boolean fileCacheAdmission) {
+ this(backendPolicy, splitGenerator, splitToScanRange,
locationProperties,
+ pathPartitionKeys, fileCacheAdmission, 8, 1024);
+ }
+
+ public SplitAssignment(FederationBackendPolicy backendPolicy,
SplitGenerator splitGenerator,
+ SplitToScanRange splitToScanRange, Map<String, String>
locationProperties,
+ List<String> pathPartitionKeys, boolean fileCacheAdmission, int
splitBatchSize, int queueCapacity) {
+ Preconditions.checkArgument(splitBatchSize > 0 && queueCapacity > 0);
+ this.splitBatchSize = splitBatchSize;
+ this.queueCapacity = queueCapacity;
this.backendPolicy = backendPolicy;
this.splitGenerator = splitGenerator;
this.splitToScanRange = splitToScanRange;
@@ -107,26 +121,54 @@ public class SplitAssignment {
private void appendBatch(Multimap<Backend, Split> batch) throws
UserException {
for (Backend backend : batch.keySet()) {
+ if (stoppedBackends.contains(backend.getId())) {
+ continue;
+ }
Collection<Split> splits = batch.get(backend);
List<TScanRangeLocations> locations = new
ArrayList<>(splits.size());
for (Split split : splits) {
locations.add(splitToScanRange.getScanRange(backend,
locationProperties, split, pathPartitionKeys,
fileCacheAdmission));
}
- while (needMoreSplit()) {
- BlockingQueue<Collection<TScanRangeLocations>> queue =
- assignment.computeIfAbsent(backend, be -> new
LinkedBlockingQueue<>(10000));
- try {
- if (queue.offer(locations, 100, TimeUnit.MILLISECONDS)) {
- break;
+ BlockingQueue<TScanRangeLocations> queue =
assignment.computeIfAbsent(
+ backend.getId(), be -> new
LinkedBlockingQueue<>(queueCapacity));
+ for (TScanRangeLocations location : locations) {
+ while (needMoreSplit() &&
!stoppedBackends.contains(backend.getId())) {
+ try {
+ if (queue.offer(location, 100, TimeUnit.MILLISECONDS))
{
+ if (stoppedBackends.contains(backend.getId())) {
+ queue.clear();
+ break;
+ }
+ synchronized (assignLock) {
+ if (sampleSplit == null) {
+ sampleSplit = splits.iterator().next();
+ assignLock.notifyAll();
+ }
+ }
+ break;
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ setException(new UserException("Interrupted while
publishing scan splits", e));
}
- } catch (InterruptedException e) {
- addUserException(new UserException("Failed to offer batch
split by interrupted", e));
}
}
}
}
+ public void stopBackend(long backendId) {
+ if (stoppedBackends.add(backendId)) {
+ BlockingQueue<TScanRangeLocations> queue =
assignment.get(backendId);
+ if (queue != null) {
+ queue.clear();
+ }
+ if (stoppedBackends.size() == backendPolicy.numBackends()) {
+ stop();
+ }
+ }
+ }
+
public void registerSource(long uniqueId) {
sources.add(uniqueId);
}
@@ -143,15 +185,15 @@ public class SplitAssignment {
if (splits.isEmpty()) {
return;
}
- Multimap<Backend, Split> batch = null;
- synchronized (assignLock) {
- if (sampleSplit == null) {
- sampleSplit = splits.get(0);
- assignLock.notify();
+ for (int begin = 0; begin < splits.size() && needMoreSplit();) {
+ int end = begin + Math.min(splitBatchSize, splits.size() - begin);
+ Multimap<Backend, Split> batch;
+ synchronized (assignLock) {
+ batch = backendPolicy.computeScanRangeAssignment(new
ArrayList<>(splits.subList(begin, end)));
}
- batch = backendPolicy.computeScanRangeAssignment(splits);
+ appendBatch(batch);
+ begin = end;
}
- appendBatch(batch);
}
private void notifyAssignment() {
@@ -160,13 +202,14 @@ public class SplitAssignment {
}
}
- public BlockingQueue<Collection<TScanRangeLocations>>
getAssignedSplits(Backend backend) throws UserException {
+ public BlockingQueue<TScanRangeLocations> getAssignedSplits(Backend
backend) throws UserException {
if (exception != null) {
throw exception;
}
- BlockingQueue<Collection<TScanRangeLocations>> splits =
assignment.computeIfAbsent(backend,
- be -> new LinkedBlockingQueue<>());
- if (scheduleFinished.get() && splits.isEmpty() || isStopped.get()) {
+ BlockingQueue<TScanRangeLocations> splits =
assignment.computeIfAbsent(backend.getId(),
+ be -> new LinkedBlockingQueue<>(queueCapacity));
+ if (scheduleFinished.get() && splits.isEmpty() || isStopped.get()
+ || stoppedBackends.contains(backend.getId())) {
return null;
}
return splits;
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java
index e24af768781..2fcddcec8b6 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java
@@ -21,9 +21,9 @@ import org.apache.doris.common.UserException;
import org.apache.doris.system.Backend;
import org.apache.doris.thrift.TScanRangeLocations;
+import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
-import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.BlockingQueue;
@@ -52,6 +52,7 @@ public class SplitSource {
private final long maxWaitTime;
public SplitSource(Backend backend, SplitAssignment splitAssignment, long
maxWaitTime) {
+ Preconditions.checkArgument(maxWaitTime > 0,
"fetch_splits_max_wait_time_ms must be positive");
this.uniqueId = UNIQUE_ID_GENERATOR.getAndIncrement();
this.backend = backend;
this.splitAssignment = splitAssignment;
@@ -60,6 +61,12 @@ public class SplitSource {
splitAssignment.registerSource(uniqueId);
}
+ /** Close this backend's shared consumer without stopping other backends
of the scan. */
+ public void close() {
+ isLastBatch.set(true);
+ splitAssignment.stopBackend(backend.getId());
+ }
+
public long getUniqueId() {
return uniqueId;
}
@@ -67,40 +74,47 @@ public class SplitSource {
/**
* Get the next batch of file splits. If there's no more split, return
empty list.
*/
- public List<TScanRangeLocations> getNextBatch(int maxBatchSize) throws
UserException {
+ public boolean isLastBatch() {
+ return isLastBatch.get();
+ }
+
+ public synchronized List<TScanRangeLocations> getNextBatch(int
maxBatchSize) throws UserException {
+ Preconditions.checkArgument(maxBatchSize > 0);
if (isLastBatch.get()) {
return Collections.emptyList();
}
List<TScanRangeLocations> scanRanges =
Lists.newArrayListWithExpectedSize(maxBatchSize);
- long startTime = System.currentTimeMillis();
- while (scanRanges.size() < maxBatchSize && System.currentTimeMillis()
- startTime < maxWaitTime) {
- BlockingQueue<Collection<TScanRangeLocations>> splits =
splitAssignment.getAssignedSplits(backend);
+ long started = System.nanoTime();
+ long waitNanos = TimeUnit.MILLISECONDS.toNanos(maxWaitTime);
+ while (scanRanges.size() < maxBatchSize) {
+ BlockingQueue<TScanRangeLocations> splits =
splitAssignment.getAssignedSplits(backend);
if (splits == null) {
isLastBatch.set(true);
break;
}
- while (scanRanges.size() < maxBatchSize) {
- try {
- Collection<TScanRangeLocations> splitCollection =
splits.poll(WAIT_TIME_OUT, TimeUnit.MILLISECONDS);
- if (splitCollection != null) {
- scanRanges.addAll(splitCollection);
- }
- if (!scanRanges.isEmpty() && System.currentTimeMillis() -
startTime > maxWaitTime) {
- return scanRanges;
- }
- if (splitCollection == null) {
- break;
- }
- } catch (InterruptedException e) {
- throw new UserException(e.getMessage(), e);
+ TScanRangeLocations split = splits.poll();
+ if (split != null) {
+ scanRanges.add(split);
+ continue;
+ }
+ if (!scanRanges.isEmpty()) {
+ break;
+ }
+ long wait = waitNanos - (System.nanoTime() - started);
+ if (wait <= 0) {
+ break;
+ }
+ try {
+ split = splits.poll(Math.min(wait,
TimeUnit.MILLISECONDS.toNanos(WAIT_TIME_OUT)),
+ TimeUnit.NANOSECONDS);
+ if (split != null) {
+ scanRanges.add(split);
}
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new UserException("Interrupted while waiting for scan
splits", e);
}
}
-
- if (scanRanges.isEmpty() && !isLastBatch.get()) {
- // This is timeout
- throw new UserException("Timeout. Max wait time(ms): " +
maxWaitTime);
- }
return scanRanges;
}
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java
index 1e45cbb3f7c..b1fa1addd33 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java
@@ -227,12 +227,15 @@ public class HiveScanNode extends FileQueryScanNode {
return allFiles;
} catch (Throwable t) {
LOG.warn("get file split failed for table: {}",
hmsTable.getName(), t);
- throw new UserException(
- "get file split failed for table: " + hmsTable.getName() + ",
err: " + Util.getRootCauseMessage(t),
- t);
+ throw createSplitPlanningException(t);
}
}
+ private UserException createSplitPlanningException(Throwable failure) {
+ return new UserException("get file split failed for table: " +
hmsTable.getName()
+ + ", err: " + Util.getRootCauseMessage(failure), failure);
+ }
+
@Override
public void startSplit(int numBackends) {
if (prunedPartitions.isEmpty()) {
@@ -245,8 +248,10 @@ public class HiveScanNode extends FileQueryScanNode {
String bindBrokerName = hmsTable.getCatalog().bindBrokerName();
AtomicInteger numFinishedPartitions = new AtomicInteger(0);
long startTime = System.currentTimeMillis();
+ List<List<HivePartition>> partitionBatches = Lists.partition(
+ prunedPartitions, sessionVariable.splitPartitionBatchSize);
CompletableFuture.runAsync(() -> {
- for (HivePartition partition : prunedPartitions) {
+ for (List<HivePartition> partitions : partitionBatches) {
if (batchException.get() != null || splitAssignment.isStop()) {
break;
}
@@ -256,22 +261,20 @@ public class HiveScanNode extends FileQueryScanNode {
try {
List<Split> allFiles = Lists.newArrayList();
getFileSplitByPartitions(
- cache,
Collections.singletonList(partition), allFiles, bindBrokerName,
+ cache, partitions, allFiles,
bindBrokerName,
numBackends, true);
- if (allFiles.size() > numSplitsPerPartition.get())
{
- numSplitsPerPartition.set(allFiles.size());
- }
+ updateNumSplitsPerPartition(allFiles);
if (splitAssignment.needMoreSplit()) {
splitAssignment.addToQueue(allFiles);
}
} catch (Exception e) {
- batchException.set(new
UserException(e.getMessage(), e));
+
batchException.set(createSplitPlanningException(e));
} finally {
splittersOnFlight.release();
if (batchException.get() != null) {
splitAssignment.setException(batchException.get());
}
- if (numFinishedPartitions.incrementAndGet() ==
prunedPartitions.size()) {
+ if (numFinishedPartitions.incrementAndGet() ==
partitionBatches.size()) {
if (getSummaryProfile() != null) {
getSummaryProfile().addExternalTableGetPartitionFilesTime(
System.currentTimeMillis() -
startTime);
@@ -282,7 +285,7 @@ public class HiveScanNode extends FileQueryScanNode {
}, scheduleExecutor);
} catch (Exception e) {
// When submitting a task, an exception will be thrown if
the task pool(scheduleExecutor) is full
- batchException.set(new UserException(e.getMessage(), e));
+ batchException.set(createSplitPlanningException(e));
break;
}
}
@@ -294,6 +297,9 @@ public class HiveScanNode extends FileQueryScanNode {
@Override
public boolean isBatchMode() {
+ if (!isBatchModeEnabled()) {
+ return false;
+ }
if (!partitionInit) {
try {
prunedPartitions = getPartitions();
@@ -311,6 +317,17 @@ public class HiveScanNode extends FileQueryScanNode {
return numSplitsPerPartition.get() * prunedPartitions.size();
}
+ private void updateNumSplitsPerPartition(List<Split> splits) {
+ // Scheduling batches contain several partitions; extrapolate the
largest single partition.
+ Map<List<String>, Integer> partitionSplitCounts = Maps.newHashMap();
+ int maxSplitCount = 0;
+ for (Split split : splits) {
+ int count = partitionSplitCounts.merge(((FileSplit)
split).getPartitionValues(), 1, Integer::sum);
+ maxSplitCount = Math.max(maxSplitCount, count);
+ }
+ numSplitsPerPartition.accumulateAndGet(maxSplitCount, Math::max);
+ }
+
private void getFileSplitByPartitions(HiveExternalMetaCache cache,
List<HivePartition> partitions,
List<Split> allFiles, String bindBrokerName, int numBackends,
boolean isBatchMode) throws IOException, UserException {
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java
index 4b7739604fb..feba79913f5 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java
@@ -1112,6 +1112,9 @@ public class HudiScanNode extends HiveScanNode {
@Override
public boolean isBatchMode() {
+ if (!isBatchModeEnabled()) {
+ return false;
+ }
if (incrementalRead && !incrementalRelation.fallbackFullTableScan()) {
return false;
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java
index 998e5152df8..a9cf412aa18 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java
@@ -3038,7 +3038,7 @@ public class IcebergScanNode extends FileQueryScanNode {
if (prepareTableLevelSnapshotCount()) {
return false;
}
- if (!sessionVariable.getEnableExternalTableBatchMode()) {
+ if (!isBatchModeEnabled()) {
isBatchMode = false;
return false;
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/maxcompute/source/MaxComputeScanNode.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/maxcompute/source/MaxComputeScanNode.java
index dd055ed6e37..54d33267a38 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/maxcompute/source/MaxComputeScanNode.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/maxcompute/source/MaxComputeScanNode.java
@@ -36,6 +36,7 @@ import org.apache.doris.common.UserException;
import org.apache.doris.common.maxcompute.MCProperties;
import org.apache.doris.common.util.LocationPath;
import org.apache.doris.datasource.FileQueryScanNode;
+import org.apache.doris.datasource.SplitAssignment;
import org.apache.doris.datasource.TableFormatType;
import org.apache.doris.datasource.maxcompute.MaxComputeExternalCatalog;
import org.apache.doris.datasource.maxcompute.MaxComputeExternalTable;
@@ -81,10 +82,7 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
-import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
public class MaxComputeScanNode extends FileQueryScanNode {
@@ -207,6 +205,9 @@ public class MaxComputeScanNode extends FileQueryScanNode {
@Override
public boolean isBatchMode() {
+ if (!isBatchModeEnabled()) {
+ return false;
+ }
if (table.getPartitionColumns().isEmpty()) {
return false;
}
@@ -243,52 +244,71 @@ public class MaxComputeScanNode extends FileQueryScanNode
{
(key, value) -> requiredPartitionSpecs.add(new
PartitionSpec(key))
);
- int batchNumPartitions = sessionVariable.getNumPartitionsInBatchMode();
-
Executor scheduleExecutor =
Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor();
- AtomicReference<UserException> batchException = new
AtomicReference<>(null);
- AtomicInteger numFinishedPartitions = new AtomicInteger(0);
+ new PartitionBatchTask(requiredPartitionSpecs.size(),
sessionVariable.splitPartitionBatchSize,
+ scheduleExecutor, splitAssignment, (begin, end) -> {
+ TableBatchReadSession session =
createTableBatchReadSession(
+ requiredPartitionSpecs.subList(begin, end));
+ List<Split> splits = getSplitByTableSession(session);
+ if (splitAssignment.needMoreSplit()) {
+ splitAssignment.addToQueue(splits);
+ }
+ }).submit();
+ }
- CompletableFuture.runAsync(() -> {
- for (int beginIndex = 0; beginIndex <
requiredPartitionSpecs.size(); beginIndex += batchNumPartitions) {
- int endIndex = Math.min(beginIndex + batchNumPartitions,
requiredPartitionSpecs.size());
- if (batchException.get() != null || splitAssignment.isStop()) {
- break;
- }
- List<PartitionSpec> requiredBatchPartitionSpecs =
requiredPartitionSpecs.subList(beginIndex, endIndex);
- int curBatchSize = endIndex - beginIndex;
-
- try {
- CompletableFuture.runAsync(() -> {
- try {
- TableBatchReadSession tableBatchReadSession =
-
createTableBatchReadSession(requiredBatchPartitionSpecs);
- List<Split> batchSplit =
getSplitByTableSession(tableBatchReadSession);
-
- if (splitAssignment.needMoreSplit()) {
- splitAssignment.addToQueue(batchSplit);
- }
- } catch (Exception e) {
- batchException.set(new
UserException(e.getMessage(), e));
- } finally {
- if (batchException.get() != null) {
-
splitAssignment.setException(batchException.get());
- }
-
- if (numFinishedPartitions.addAndGet(curBatchSize)
== requiredPartitionSpecs.size()) {
- splitAssignment.finishSchedule();
- }
- }
- }, scheduleExecutor);
- } catch (Exception e) {
- batchException.set(new UserException(e.getMessage(), e));
- }
+ @FunctionalInterface
+ interface PartitionBatchConsumer {
+ void accept(int begin, int end) throws Exception;
+ }
- if (batchException.get() != null) {
- splitAssignment.setException(batchException.get());
+ /** One queued continuation per scan, independent of the threshold that
enables batch mode. */
+ static final class PartitionBatchTask implements Runnable {
+ private final int partitionCount;
+ private final int batchSize;
+ private final Executor executor;
+ private final SplitAssignment assignment;
+ private final PartitionBatchConsumer consumer;
+ private int nextPartition;
+
+ PartitionBatchTask(int partitionCount, int batchSize, Executor
executor,
+ SplitAssignment assignment, PartitionBatchConsumer consumer) {
+ com.google.common.base.Preconditions.checkArgument(batchSize > 0);
+ this.partitionCount = partitionCount;
+ this.batchSize = batchSize;
+ this.executor = executor;
+ this.assignment = assignment;
+ this.consumer = consumer;
+ }
+
+ void submit() {
+ try {
+ executor.execute(this);
+ } catch (java.util.concurrent.RejectedExecutionException e) {
+ assignment.setException(new UserException(e.getMessage(), e));
+ assignment.finishSchedule();
+ }
+ }
+
+ @Override
+ public void run() {
+ if (!assignment.needMoreSplit() || nextPartition ==
partitionCount) {
+ assignment.finishSchedule();
+ return;
+ }
+ int end = nextPartition + Math.min(batchSize, partitionCount -
nextPartition);
+ try {
+ consumer.accept(nextPartition, end);
+ nextPartition = end;
+ if (nextPartition == partitionCount ||
!assignment.needMoreSplit()) {
+ assignment.finishSchedule();
+ } else {
+ submit();
}
+ } catch (Exception e) {
+ assignment.setException(new UserException(e.getMessage(), e));
+ assignment.finishSchedule();
}
- }, scheduleExecutor);
+ }
}
@Override
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
index 68ecf81c998..de60dd1e10c 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
@@ -84,7 +84,6 @@ import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
-
/**
* System variable.
**/
@@ -603,6 +602,10 @@ public class SessionVariable implements Serializable,
Writable {
public static final String ICEBERG_WRITE_TARGET_FILE_SIZE_BYTES =
"iceberg_write_target_file_size_bytes";
public static final String ENABLE_ICEBERG_MERGE_PARTITIONING =
"enable_iceberg_merge_partitioning";
+ public static final String SPLIT_BATCH_SIZE = "split_batch_size";
+ public static final String SPLIT_QUEUE_CAPACITY = "split_queue_capacity";
+ public static final String SPLIT_PARTITION_BATCH_SIZE =
"split_partition_batch_size";
+
public static final String NUM_PARTITIONS_IN_BATCH_MODE =
"num_partitions_in_batch_mode";
public static final String NUM_FILES_IN_BATCH_MODE =
"num_files_in_batch_mode";
@@ -2723,13 +2726,39 @@ public class SessionVariable implements Serializable,
Writable {
@VariableMgr.VarAttr(name = ICEBERG_WRITE_TARGET_FILE_SIZE_BYTES,
needForward = true)
public long icebergWriteTargetFileSizeBytes = 0L;
+ @VariableMgr.VarAttr(name = SPLIT_BATCH_SIZE, checker =
"checkSplitBatchSize", needForward = true)
+ public int splitBatchSize = 8;
+
+ @VariableMgr.VarAttr(name = SPLIT_QUEUE_CAPACITY, checker =
"checkSplitQueueCapacity", needForward = true)
+ public int splitQueueCapacity = 1024;
+
+ @VariableMgr.VarAttr(name = SPLIT_PARTITION_BATCH_SIZE, checker =
"checkSplitPartitionBatchSize",
+ needForward = true)
+ public int splitPartitionBatchSize = 16;
+
+ public void checkFetchSplitsMaxWaitTime(String value) throws Exception {
+ checkFieldLongValue(FETCH_SPLITS_MAX_WAIT_TIME, 1, value);
+ }
+
+ public void checkSplitBatchSize(String value) throws Exception {
+ checkFieldValue(SPLIT_BATCH_SIZE, 1, value);
+ }
+
+ public void checkSplitQueueCapacity(String value) throws Exception {
+ checkFieldValue(SPLIT_QUEUE_CAPACITY, 1, value);
+ }
+
+ public void checkSplitPartitionBatchSize(String value) throws Exception {
+ checkFieldValue(SPLIT_PARTITION_BATCH_SIZE, 1, value);
+ }
+
@VariableMgr.VarAttr(
name = NUM_PARTITIONS_IN_BATCH_MODE,
fuzzy = true,
description = {"如果分区数量超过阈值，BE 将通过 batch 方式获取 scan ranges。作用于
Hive、Hudi、MaxCompute 表。",
"If the number of partitions exceeds the threshold, scan
ranges will be got through batch mode."},
needForward = true)
- public int numPartitionsInBatchMode = 1024;
+ public int numPartitionsInBatchMode = 1;
@VariableMgr.VarAttr(
name = NUM_FILES_IN_BATCH_MODE,
@@ -2737,10 +2766,11 @@ public class SessionVariable implements Serializable,
Writable {
description = {"如果文件数量超过阈值，BE 将通过 batch 方式获取 scan ranges",
"If the number of files exceeds the threshold, scan ranges
will be got through batch mode."},
needForward = true)
- public int numFilesInBatchMode = 1024;
+ public int numFilesInBatchMode = 1;
@VariableMgr.VarAttr(
name = FETCH_SPLITS_MAX_WAIT_TIME,
+ checker = "checkFetchSplitsMaxWaitTime",
description = {"batch 方式中 BE 获取 splits 的最大等待时间",
"The max wait time of getting splits in batch mode."},
needForward = true)
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/service/FrontendServiceImpl.java
b/fe/fe-core/src/main/java/org/apache/doris/service/FrontendServiceImpl.java
index 0c751482760..23c3a74f1d4 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/service/FrontendServiceImpl.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/service/FrontendServiceImpl.java
@@ -1103,6 +1103,15 @@ public class FrontendServiceImpl implements
FrontendService.Iface {
TFetchSplitBatchResult result = new TFetchSplitBatchResult();
SplitSource splitSource =
Env.getCurrentEnv().getSplitSourceManager().getSplitSource(request.getSplitSourceId());
+ if (request.isSetCloseSource() && request.isCloseSource()) {
+ if (splitSource != null) {
+ splitSource.close();
+
Env.getCurrentEnv().getSplitSourceManager().removeSplitSource(request.getSplitSourceId());
+ }
+ result.setSplits(Collections.emptyList());
+ result.status = new TStatus(TStatusCode.OK);
+ return result;
+ }
if (splitSource == null) {
// Return a structured error status instead of throwing a bare
TException, so the BE
// receives a well-formed error (with a non-empty message) through
the normal result
@@ -1117,7 +1126,8 @@ public class FrontendServiceImpl implements
FrontendService.Iface {
try {
List<TScanRangeLocations> locations =
splitSource.getNextBatch(request.getMaxNumSplits());
result.setSplits(locations);
- result.status = new TStatus(TStatusCode.OK);
+ result.status = new TStatus(locations.isEmpty() &&
!splitSource.isLastBatch()
+ ? TStatusCode.INCOMPLETE : TStatusCode.OK);
return result;
} catch (Exception e) {
LOG.warn("failed to fetch split batch with source id {}",
request.getSplitSourceId(), e);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/blockrule/SqlBlockRuleMgrTest.java
b/fe/fe-core/src/test/java/org/apache/doris/blockrule/SqlBlockRuleMgrTest.java
index a87fcf1a84b..822c4853064 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/blockrule/SqlBlockRuleMgrTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/blockrule/SqlBlockRuleMgrTest.java
@@ -17,13 +17,18 @@
package org.apache.doris.blockrule;
+import org.apache.doris.catalog.Env;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.metric.MetricRepo;
+import org.apache.doris.mysql.privilege.Auth;
import org.apache.doris.persist.gson.GsonUtils;
+import org.apache.doris.qe.ConnectContext;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
import java.util.List;
import java.util.Map;
@@ -35,6 +40,40 @@ public class SqlBlockRuleMgrTest {
MetricRepo.init();
}
+ @Test
+ public void testExactScanCountsOnlyForApplicableEnabledCountRules() {
+ SqlBlockRuleMgr rules = new SqlBlockRuleMgr();
+ ConnectContext context = new ConnectContext();
+ context.setThreadLocalInfo();
+ Env env = Mockito.mock(Env.class);
+ Auth auth = Mockito.mock(Auth.class);
+ Mockito.when(env.getAuth()).thenReturn(auth);
+ Mockito.when(auth.getSqlBlockRules("user")).thenReturn(new String[]
{"missing"});
+ try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+ mockedEnv.when(Env::getCurrentEnv).thenReturn(env);
+ rules.getNameToSqlBlockRuleMap().put("user_limit", new
SqlBlockRule(
+ "user_limit", "NULL", "NULL", 1L, 0L, 0L, false, false,
true));
+ Assert.assertFalse(rules.hasScanCountLimitations("user"));
+ Mockito.when(auth.getSqlBlockRules("user")).thenReturn(new
String[] {"missing", "user_limit"});
+ Assert.assertTrue(rules.hasScanCountLimitations("user"));
+ context.getState().setInternal(true);
+ Assert.assertFalse(rules.hasScanCountLimitations("user"));
+ context.getState().setInternal(false);
+
+ Mockito.when(auth.getSqlBlockRules("user")).thenReturn(new
String[0]);
+ rules.getNameToSqlBlockRuleMap().put("disabled", new SqlBlockRule(
+ "disabled", "NULL", "NULL", 1L, 1L, 0L, false, true,
false));
+ rules.getNameToSqlBlockRuleMap().put("metadata_only", new
SqlBlockRule(
+ "metadata_only", "SELECT .*", "NULL", 0L, 0L, 1L, true,
true, true));
+ Assert.assertFalse(rules.hasScanCountLimitations("user"));
+ rules.getNameToSqlBlockRuleMap().put("global_limit", new
SqlBlockRule(
+ "global_limit", "NULL", "NULL", 0L, 1L, 0L, false, true,
true));
+ Assert.assertTrue(rules.hasScanCountLimitations("user"));
+ } finally {
+ ConnectContext.remove();
+ }
+ }
+
@Test
public void testToInfoString() {
SqlBlockRuleMgr mgr = new SqlBlockRuleMgr();
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/BoundedSplitDeliveryTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/BoundedSplitDeliveryTest.java
new file mode 100644
index 00000000000..5182864e502
--- /dev/null
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/BoundedSplitDeliveryTest.java
@@ -0,0 +1,148 @@
+// 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.
+
+package org.apache.doris.datasource;
+
+import org.apache.doris.common.UserException;
+import org.apache.doris.qe.SessionVariable;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TScanRangeLocations;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.mockito.Mockito;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+@Timeout(30)
+class BoundedSplitDeliveryTest {
+ @Test
+ void nonPositiveFetchWaitIsRejected() {
+ Backend backend = new Backend(1, "localhost", 0);
+ SplitAssignment assignment = Mockito.mock(SplitAssignment.class);
+ Assertions.assertThrows(IllegalArgumentException.class, () -> new
SplitSource(backend, assignment, 0));
+ Assertions.assertThrows(IllegalArgumentException.class, () -> new
SplitSource(backend, assignment, -1));
+ }
+
+ @Test
+ void defaultsAndValidation() throws Exception {
+ SessionVariable variables = new SessionVariable();
+ Assertions.assertEquals(1, variables.getNumPartitionsInBatchMode());
+ Assertions.assertEquals(1, variables.getNumFilesInBatchMode());
+ Assertions.assertEquals(8, variables.splitBatchSize);
+ Assertions.assertEquals(1024, variables.splitQueueCapacity);
+ Assertions.assertEquals(16, variables.splitPartitionBatchSize);
+ Assertions.assertThrows(Exception.class, () ->
variables.checkSplitQueueCapacity("0"));
+ Assertions.assertThrows(Exception.class, () ->
variables.checkSplitBatchSize("-1"));
+ Assertions.assertThrows(Exception.class, () ->
variables.checkSplitPartitionBatchSize("0"));
+ variables.checkSplitQueueCapacity("1");
+ Assertions.assertThrows(Exception.class, () ->
variables.checkFetchSplitsMaxWaitTime("0"));
+ Assertions.assertThrows(Exception.class, () ->
variables.checkFetchSplitsMaxWaitTime("-1"));
+ variables.checkFetchSplitsMaxWaitTime("1");
+ }
+
+ @Test
+ void producerBatchLargerThanQueueStartsAndDrainsWithoutLoss() throws
Exception {
+ Backend backend = new Backend(1, "localhost", 0);
+ FederationBackendPolicy policy =
Mockito.mock(FederationBackendPolicy.class);
+
Mockito.when(policy.computeScanRangeAssignment(Mockito.anyList())).thenAnswer(call
-> {
+ Multimap<Backend, Split> assigned = ArrayListMultimap.create();
+ for (Split split : call.<List<Split>>getArgument(0)) {
+ assigned.put(backend, split);
+ }
+ return assigned;
+ });
+ List<Split> splits = new ArrayList<>();
+ for (int i = 0; i < 9; i++) {
+ splits.add(Mockito.mock(Split.class));
+ }
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ AtomicReference<SplitAssignment> holder = new AtomicReference<>();
+ CompletableFuture<Void> published = new CompletableFuture<>();
+ SplitGenerator generator = new SplitGenerator() {
+ @Override
+ public void startSplit(int ignored) {
+ executor.execute(() -> {
+ try {
+ holder.get().addToQueue(splits);
+ holder.get().finishSchedule();
+ published.complete(null);
+ } catch (Throwable e) {
+ published.completeExceptionally(e);
+ }
+ });
+ }
+ };
+ SplitAssignment assignment = new SplitAssignment(policy, generator,
+ (be, properties, split, keys, admission) -> new
TScanRangeLocations(),
+ Collections.emptyMap(), Collections.emptyList(), true, 8, 2);
+ holder.set(assignment);
+ try {
+ // init must return after the first published split, before the
full batch fits.
+ assignment.init();
+ Assertions.assertNotNull(assignment.getSampleSplit());
+ Assertions.assertEquals(2,
assignment.getAssignedSplits(backend).size()
+ +
assignment.getAssignedSplits(backend).remainingCapacity());
+ SplitSource source = new SplitSource(backend, assignment, 1000);
+ int count = 0;
+ while (!source.isLastBatch()) {
+ List<TScanRangeLocations> batch = source.getNextBatch(3);
+ Assertions.assertTrue(batch.size() <= 3);
+ count += batch.size();
+ }
+ published.get(10, TimeUnit.SECONDS);
+ Assertions.assertEquals(9, count);
+ } finally {
+ assignment.stop();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void liveEmptySourceIsIncompleteAndFinishedSourceIsEof() throws Exception {
+ SplitAssignment assignment = new
SplitAssignment(Mockito.mock(FederationBackendPolicy.class),
+ new SplitGenerator() { }, (be, properties, split, keys,
admission) -> new TScanRangeLocations(),
+ Collections.emptyMap(), Collections.emptyList(), true);
+ SplitSource source = new SplitSource(new Backend(1, "localhost", 0),
assignment, 1);
+ Assertions.assertTrue(source.getNextBatch(1).isEmpty());
+ Assertions.assertFalse(source.isLastBatch());
+ assignment.finishSchedule();
+ Assertions.assertTrue(source.getNextBatch(1).isEmpty());
+ Assertions.assertTrue(source.isLastBatch());
+ }
+
+ @Test
+ void producerFailureIsAnErrorRatherThanIncomplete() {
+ SplitAssignment assignment = new
SplitAssignment(Mockito.mock(FederationBackendPolicy.class),
+ new SplitGenerator() { }, (be, properties, split, keys,
admission) -> new TScanRangeLocations(),
+ Collections.emptyMap(), Collections.emptyList(), true);
+ SplitSource source = new SplitSource(new Backend(1, "localhost", 0),
assignment, 1);
+ assignment.setException(new UserException("listing failed"));
+ Assertions.assertThrows(UserException.class, () ->
source.getNextBatch(1));
+ }
+}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/FileQueryScanNodeTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/FileQueryScanNodeTest.java
index 9b57f781305..167b1b8cb6b 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/datasource/FileQueryScanNodeTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/FileQueryScanNodeTest.java
@@ -119,6 +119,22 @@ public class FileQueryScanNodeTest {
}
}
+ @Test
+ public void testHudiHonorsCommonBatchEligibility() {
+ FileQueryScanNode node =
Mockito.mock(org.apache.doris.datasource.hudi.source.HudiScanNode.class,
+ Mockito.CALLS_REAL_METHODS);
+ Mockito.doReturn(false).when(node).isBatchModeEnabled();
+ Assert.assertFalse(node.isBatchMode());
+ }
+
+ @Test
+ public void testMaxComputeHonorsCommonBatchEligibility() {
+ FileQueryScanNode node =
Mockito.mock(org.apache.doris.datasource.maxcompute.source.MaxComputeScanNode.class,
+ Mockito.CALLS_REAL_METHODS);
+ Mockito.doReturn(false).when(node).isBatchModeEnabled();
+ Assert.assertFalse(node.isBatchMode());
+ }
+
@Before
public void setUp() {
table = Mockito.mock(TableIf.class);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitAssignmentTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitAssignmentTest.java
index 050749a6d11..c9b16eb3592 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitAssignmentTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitAssignmentTest.java
@@ -33,7 +33,6 @@ import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import java.util.ArrayList;
-import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
@@ -241,7 +240,7 @@ public class SplitAssignmentTest {
Assertions.assertEquals(mockSplit, splitAssignment.getSampleSplit());
// Verify assignment queue is created and contains data
- BlockingQueue<Collection<TScanRangeLocations>> queue =
splitAssignment.getAssignedSplits(mockBackend);
+ BlockingQueue<TScanRangeLocations> queue =
splitAssignment.getAssignedSplits(mockBackend);
Assertions.assertNotNull(queue);
}
@@ -305,7 +304,7 @@ public class SplitAssignmentTest {
}
// Verify the queue has data
- BlockingQueue<Collection<TScanRangeLocations>> queue =
splitAssignment.getAssignedSplits(mockBackend);
+ BlockingQueue<TScanRangeLocations> queue =
splitAssignment.getAssignedSplits(mockBackend);
Assertions.assertNotNull(queue);
}
@@ -399,7 +398,7 @@ public class SplitAssignmentTest {
Assertions.assertNotNull(splitAssignment.getSampleSplit());
Assertions.assertEquals(mockSplit, splitAssignment.getSampleSplit());
- BlockingQueue<Collection<TScanRangeLocations>> queue =
splitAssignment.getAssignedSplits(mockBackend);
+ BlockingQueue<TScanRangeLocations> queue =
splitAssignment.getAssignedSplits(mockBackend);
Assertions.assertNotNull(queue);
splitsProvider.join(1000);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitConsumerCloseTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitConsumerCloseTest.java
new file mode 100644
index 00000000000..aa6c3cf9de8
--- /dev/null
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitConsumerCloseTest.java
@@ -0,0 +1,136 @@
+// 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.
+
+package org.apache.doris.datasource;
+
+import org.apache.doris.catalog.Env;
+import org.apache.doris.service.FrontendServiceImpl;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TFetchSplitBatchRequest;
+import org.apache.doris.thrift.TScanRangeLocations;
+import org.apache.doris.thrift.TStatusCode;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+@Timeout(30)
+class SplitConsumerCloseTest {
+ private static Env env;
+ private static FrontendServiceImpl service;
+
+ @BeforeAll
+ static void initializeRpcFixture() {
+ // Keep Env/service cold initialization outside the bounded
queue-close race.
+ env = Mockito.mock(Env.class);
+ service = Mockito.mock(FrontendServiceImpl.class,
Mockito.CALLS_REAL_METHODS);
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ void
ownerCloseReleasesAFullBackendQueueWithoutDroppingOtherBackendSplits() throws
Exception {
+ Backend stopped = new Backend(1, "localhost", 0);
+ Backend live = new Backend(2, "localhost", 0);
+ FederationBackendPolicy policy =
Mockito.mock(FederationBackendPolicy.class);
+ Mockito.when(policy.numBackends()).thenReturn(2);
+ Split first = Mockito.mock(Split.class);
+ Split second = Mockito.mock(Split.class);
+
Mockito.when(policy.computeScanRangeAssignment(Mockito.anyList())).thenAnswer(call
-> {
+ Multimap<Backend, Split> result = ArrayListMultimap.create();
+ for (Split split : call.<java.util.List<Split>>getArgument(0)) {
+ result.put(split == first ? stopped : live, split);
+ }
+ return result;
+ });
+ SplitAssignment assignment = new SplitAssignment(policy, new
SplitGenerator() {},
+ (be, props, split, keys, admission) -> new
TScanRangeLocations(),
+ Collections.emptyMap(), Collections.emptyList(), true, 2, 1);
+ CountDownLatch offering = new CountDownLatch(1);
+ CountDownLatch allowOffer = new CountDownLatch(1);
+ AtomicBoolean delayNext = new AtomicBoolean(false);
+ BlockingQueue<TScanRangeLocations> stoppedQueue = new
LinkedBlockingQueue<TScanRangeLocations>(1) {
+ @Override
+ public boolean offer(TScanRangeLocations value, long timeout,
TimeUnit unit) throws InterruptedException {
+ if (delayNext.compareAndSet(true, false)) {
+ offering.countDown();
+ Assertions.assertTrue(allowOffer.await(5,
TimeUnit.SECONDS));
+ }
+ return super.offer(value, timeout, unit);
+ }
+ };
+ Field queues = SplitAssignment.class.getDeclaredField("assignment");
+ queues.setAccessible(true);
+ ((Map<Long, BlockingQueue<TScanRangeLocations>>)
queues.get(assignment)).put(stopped.getId(), stoppedQueue);
+ assignment.addToQueue(Collections.singletonList(first));
+ SplitSource stoppedSource = new SplitSource(stopped, assignment, 1000);
+ SplitSource liveSource = new SplitSource(live, assignment, 1000);
+ SplitSourceManager manager = new SplitSourceManager();
+ manager.registerSplitSource(stoppedSource);
+ manager.registerSplitSource(liveSource);
+ Mockito.when(env.getSplitSourceManager()).thenReturn(manager);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ delayNext.set(true);
+ try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+ mockedEnv.when(Env::getCurrentEnv).thenReturn(env);
+ Future<?> produced = executor.submit(() -> {
+ assignment.addToQueue(Arrays.asList(first, second));
+ assignment.finishSchedule();
+ return null;
+ });
+ Assertions.assertTrue(offering.await(5, TimeUnit.SECONDS));
+ TFetchSplitBatchRequest request = new TFetchSplitBatchRequest()
+
.setSplitSourceId(stoppedSource.getUniqueId()).setMaxNumSplits(8).setCloseSource(true);
+ Assertions.assertEquals(TStatusCode.OK,
service.fetchSplitBatch(request).getStatus().getStatusCode());
+
Assertions.assertNull(manager.getSplitSource(stoppedSource.getUniqueId()));
+ Assertions.assertFalse(assignment.isStop());
+ Assertions.assertEquals(TStatusCode.OK,
service.fetchSplitBatch(request).getStatus().getStatusCode());
+ allowOffer.countDown();
+ produced.get(5, TimeUnit.SECONDS);
+ Assertions.assertTrue(stoppedQueue.isEmpty(), "an offer racing
close must not retain descriptors");
+ Assertions.assertNull(assignment.getAssignedSplits(stopped));
+ Assertions.assertEquals(1, liveSource.getNextBatch(2).size());
+ Assertions.assertTrue(liveSource.isLastBatch());
+ request.setSplitSourceId(liveSource.getUniqueId());
+ Assertions.assertEquals(TStatusCode.OK,
service.fetchSplitBatch(request).getStatus().getStatusCode());
+ Assertions.assertTrue(assignment.isStop());
+ } finally {
+ allowOffer.countDown();
+ assignment.stop();
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5,
TimeUnit.SECONDS));
+ }
+ }
+}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java
index a194e8f67c6..c3e05409d3e 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java
@@ -17,6 +17,7 @@
package org.apache.doris.datasource.hive.source;
+import org.apache.doris.analysis.StatementBase;
import org.apache.doris.analysis.TableSample;
import org.apache.doris.analysis.TupleDescriptor;
import org.apache.doris.analysis.TupleId;
@@ -60,6 +61,82 @@ import java.util.concurrent.ThreadPoolExecutor;
public class HiveScanNodeTest {
private static final long MB = 1024L * 1024L;
+ @Test
+ public void testNativeBatchAndSynchronousFailuresKeepRootCause() throws
Exception {
+ HiveScanNode node = createHiveScanNode();
+ Field tableField = HiveScanNode.class.getDeclaredField("hmsTable");
+ tableField.setAccessible(true);
+ HMSExternalTable table = (HMSExternalTable) tableField.get(node);
+ Mockito.when(table.getName()).thenReturn("hive_test");
+ Field initialized =
HiveScanNode.class.getDeclaredField("partitionInit");
+ initialized.setAccessible(true);
+ initialized.set(node, true);
+ Field partitions =
HiveScanNode.class.getDeclaredField("prunedPartitions");
+ partitions.setAccessible(true);
+ partitions.set(node,
Collections.singletonList(Mockito.mock(HivePartition.class)));
+ org.apache.doris.datasource.SplitAssignment assignment =
+
Mockito.mock(org.apache.doris.datasource.SplitAssignment.class);
+ Field splitAssignment =
org.apache.doris.planner.ScanNode.class.getDeclaredField("splitAssignment");
+ splitAssignment.setAccessible(true);
+ splitAssignment.set(node, assignment);
+ org.apache.doris.catalog.Env env =
Mockito.mock(org.apache.doris.catalog.Env.class);
+ org.apache.doris.datasource.ExternalMetaCacheMgr manager =
+
Mockito.mock(org.apache.doris.datasource.ExternalMetaCacheMgr.class);
+ HiveExternalMetaCache cache =
Mockito.mock(HiveExternalMetaCache.class);
+ java.util.concurrent.ExecutorService executor =
java.util.concurrent.Executors.newSingleThreadExecutor();
+ Mockito.when(env.getExtMetaCacheMgr()).thenReturn(manager);
+ Mockito.when(manager.hive(Mockito.anyLong())).thenReturn(cache);
+ Mockito.when(manager.getScheduleExecutor()).thenReturn(executor);
+ RuntimeException wrapper = new RuntimeException("file cache wrapper",
+ new org.apache.doris.common.UserException("Partition location
does not exist: /missing"));
+ Mockito.when(cache.getFilesByPartitions(Mockito.anyList(),
Mockito.anyBoolean(), Mockito.anyBoolean(),
+ Mockito.isNull(), Mockito.same(table))).thenThrow(wrapper);
+
java.util.concurrent.CompletableFuture<org.apache.doris.common.UserException>
failure =
+ new java.util.concurrent.CompletableFuture<>();
+ Mockito.doAnswer(call -> {
+ failure.complete(call.getArgument(0));
+ return null;
+ }).when(assignment).setException(Mockito.any());
+ try (org.mockito.MockedStatic<org.apache.doris.catalog.Env> mockedEnv =
+ Mockito.mockStatic(org.apache.doris.catalog.Env.class)) {
+
mockedEnv.when(org.apache.doris.catalog.Env::getCurrentEnv).thenReturn(env);
+ node.startSplit(1);
+ org.apache.doris.common.UserException async = failure.get(5,
java.util.concurrent.TimeUnit.SECONDS);
+ Assert.assertTrue(async.getMessage().contains("Partition location
does not exist: /missing"));
+ Assert.assertTrue(async.getMessage().contains("hive_test"));
+ org.apache.doris.common.UserException sync = Assert.assertThrows(
+ org.apache.doris.common.UserException.class, () ->
node.getSplits(1));
+ Assert.assertEquals(sync.getMessage(), async.getMessage());
+ Assert.assertSame(wrapper, async.getCause());
+ } finally {
+ executor.shutdownNow();
+ Assert.assertTrue(executor.awaitTermination(5,
java.util.concurrent.TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ public void testVerboseExplainUsesExactFilePlanning() throws Exception {
+ HiveScanNode node = createHiveScanNode();
+ Field initialized =
HiveScanNode.class.getDeclaredField("partitionInit");
+ initialized.setAccessible(true);
+ initialized.set(node, true);
+ Field partitions =
HiveScanNode.class.getDeclaredField("prunedPartitions");
+ partitions.setAccessible(true);
+ partitions.set(node,
Collections.singletonList(Mockito.mock(HivePartition.class)));
+ ConnectContext context = new ConnectContext();
+ StatementContext statementContext = new StatementContext();
+ StatementBase statement = Mockito.mock(StatementBase.class);
+ Mockito.when(statement.isVerbose()).thenReturn(true);
+ statementContext.setParsedStatement(statement);
+ context.setStatementContext(statementContext);
+ context.setThreadLocalInfo();
+ try {
+ Assert.assertFalse(node.isBatchMode());
+ } finally {
+ ConnectContext.remove();
+ }
+ }
+
@Test
public void testStatementCacheReusesListingOnlyForSamePartitionIdentity()
throws Exception {
ConnectContext previousContext = ConnectContext.get();
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveSplitEstimateTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveSplitEstimateTest.java
new file mode 100644
index 00000000000..01a9c94df6f
--- /dev/null
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveSplitEstimateTest.java
@@ -0,0 +1,216 @@
+// 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.
+
+package org.apache.doris.datasource.hive.source;
+
+import org.apache.doris.analysis.TableSample;
+import org.apache.doris.analysis.TupleDescriptor;
+import org.apache.doris.analysis.TupleId;
+import org.apache.doris.blockrule.SqlBlockRuleMgr;
+import org.apache.doris.catalog.Column;
+import org.apache.doris.catalog.Env;
+import org.apache.doris.catalog.Type;
+import org.apache.doris.common.util.LocationPath;
+import org.apache.doris.datasource.ExternalMetaCacheMgr;
+import org.apache.doris.datasource.FileQueryScanNode;
+import org.apache.doris.datasource.FileSplitter;
+import org.apache.doris.datasource.SplitAssignment;
+import org.apache.doris.datasource.hive.HMSExternalCatalog;
+import org.apache.doris.datasource.hive.HMSExternalTable;
+import org.apache.doris.datasource.hive.HiveExternalMetaCache;
+import org.apache.doris.datasource.hive.HiveExternalMetaCache.FileCacheValue;
+import org.apache.doris.datasource.hive.HiveExternalMetaCache.HiveFileStatus;
+import org.apache.doris.datasource.hive.HivePartition;
+import org.apache.doris.planner.PlanNodeId;
+import org.apache.doris.planner.ScanContext;
+import org.apache.doris.planner.ScanNode;
+import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.SessionVariable;
+import org.apache.doris.thrift.TFileFormatType;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+public class HiveSplitEstimateTest {
+ @Test
+ public void testUniformPartitionsAreIndependentOfBatchSize() throws
Exception {
+ for (int batchSize : new int[] {16, 2, 1}) {
+ Assert.assertEquals(60, runBatchEstimate(batchSize, false, 10, 10,
10, 10, 10, 10));
+ }
+ }
+
+ @Test
+ public void testUnevenAndEmptyPartitionsKeepTheLargestPartitionEstimate()
throws Exception {
+ for (int batchSize : new int[] {16, 2, 1}) {
+ Assert.assertEquals(100, runBatchEstimate(batchSize, false, 0, 1,
25, 3));
+ }
+ }
+
+ @Test
+ public void testSampledPartitionsKeepTheSameEstimate() throws Exception {
+ for (int batchSize : new int[] {16, 2}) {
+ Assert.assertEquals(60, runBatchEstimate(batchSize, true, 10, 10,
10, 10, 10, 10));
+ }
+ }
+
+ @Test
+ public void testUnpartitionedScanKeepsItsSplitCount() throws Exception {
+ Assert.assertEquals(25, runBatchEstimate(16, false, 25));
+ }
+
+ @Test
+ public void testRowsSampleIsSelectedOnceAcrossAllPartitions() throws
Exception {
+ assertTableWideSample(new TableSample(false, 256L, 0L));
+ }
+
+ @Test
+ public void testPercentSampleIsSelectedOnceAcrossAllPartitions() throws
Exception {
+ assertTableWideSample(new TableSample(true, 1L, 0L));
+ }
+
+ private void assertTableWideSample(TableSample sample) throws Exception {
+ int[] fileCounts = new int[32];
+ Arrays.fill(fileCounts, 1);
+ for (int batchSize : new int[] {16, 2, 1}) {
+ Assert.assertEquals(1, runSplitPlanning(batchSize, sample, true,
fileCounts));
+ }
+ }
+
+ private int runBatchEstimate(int batchSize, boolean sample, int...
fileCounts) throws Exception {
+ return runSplitPlanning(batchSize, sample ? new TableSample(true,
100L, 0L) : null,
+ false, fileCounts);
+ }
+
+ private int runSplitPlanning(int batchSize, TableSample sample, boolean
choosePlanningMode,
+ int... fileCounts) throws Exception {
+ SessionVariable variables = new SessionVariable();
+ variables.splitPartitionBatchSize = batchSize;
+ HMSExternalTable table = Mockito.mock(HMSExternalTable.class);
+ HMSExternalCatalog catalog = Mockito.mock(HMSExternalCatalog.class);
+ Mockito.when(table.getCatalog()).thenReturn(catalog);
+
Mockito.when(table.getFileFormatType(variables)).thenReturn(TFileFormatType.FORMAT_PARQUET);
+ Mockito.when(catalog.bindBrokerName()).thenReturn("");
+ TupleDescriptor descriptor = new TupleDescriptor(new TupleId(0));
+ descriptor.setTable(table);
+ HiveScanNode node = new HiveScanNode(
+ new PlanNodeId(0), descriptor, false, variables, null,
ScanContext.EMPTY);
+ Field splitter =
FileQueryScanNode.class.getDeclaredField("fileSplitter");
+ splitter.setAccessible(true);
+ splitter.set(node, new FileSplitter(variables.maxInitialSplitSize,
+ variables.maxSplitSize, variables.maxInitialSplitNum));
+ node.setTableSample(sample);
+
Mockito.when(table.getFullSchema()).thenReturn(Collections.singletonList(new
Column("v", Type.INT)));
+ List<HivePartition> partitions = new ArrayList<>();
+ Map<HivePartition, FileCacheValue> files = new HashMap<>();
+ int totalFiles = 0;
+ for (int partition = 0; partition < fileCounts.length; partition++) {
+ List<String> values = fileCounts.length == 1
+ ? null :
Collections.singletonList(Integer.toString(partition));
+ String path = "hdfs://warehouse/t/p=" + partition;
+ HivePartition hivePartition = new HivePartition(null, false,
"parquet", path,
+ values, Collections.emptyMap());
+ partitions.add(hivePartition);
+ FileCacheValue cacheValue = new FileCacheValue();
+ cacheValue.setPartitionValues(values);
+ for (int file = 0; file < fileCounts[partition]; file++) {
+ HiveFileStatus status = new HiveFileStatus();
+ status.setPath(LocationPath.of(path + "/file" + file +
".parquet"));
+ status.setLength(1024);
+ cacheValue.getFiles().add(status);
+ totalFiles++;
+ }
+ files.put(hivePartition, cacheValue);
+ }
+ Field pruned = HiveScanNode.class.getDeclaredField("prunedPartitions");
+ pruned.setAccessible(true);
+ pruned.set(node, partitions);
+ Field initialized =
HiveScanNode.class.getDeclaredField("partitionInit");
+ initialized.setAccessible(true);
+ initialized.set(node, true);
+ SplitAssignment assignment = Mockito.mock(SplitAssignment.class);
+ Mockito.when(assignment.needMoreSplit()).thenReturn(true);
+ Field assignmentField =
ScanNode.class.getDeclaredField("splitAssignment");
+ assignmentField.setAccessible(true);
+ assignmentField.set(node, assignment);
+ CompletableFuture<Void> finished = new CompletableFuture<>();
+ Mockito.doAnswer(call -> {
+ finished.complete(null);
+ return null;
+ }).when(assignment).finishSchedule();
+ Mockito.doAnswer(call -> {
+ finished.completeExceptionally(call.getArgument(0));
+ return null;
+ }).when(assignment).setException(Mockito.any());
+ AtomicInteger delivered = new AtomicInteger();
+ Mockito.doAnswer(call -> {
+ delivered.addAndGet(call.<List<?>>getArgument(0).size());
+ return null;
+ }).when(assignment).addToQueue(Mockito.anyList());
+ HiveExternalMetaCache cache =
Mockito.mock(HiveExternalMetaCache.class);
+ Mockito.when(cache.getFilesByPartitions(Mockito.anyList(),
Mockito.anyBoolean(), Mockito.anyBoolean(),
+ Mockito.isNull(), Mockito.same(table))).thenAnswer(call ->
call.<List<HivePartition>>getArgument(0)
+
.stream().map(files::get).collect(Collectors.toList()));
+ Env env = Mockito.mock(Env.class);
+ ExternalMetaCacheMgr manager =
Mockito.mock(ExternalMetaCacheMgr.class);
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ Mockito.when(env.getExtMetaCacheMgr()).thenReturn(manager);
+
Mockito.when(env.getSqlBlockRuleMgr()).thenReturn(Mockito.mock(SqlBlockRuleMgr.class));
+ Mockito.when(manager.hive(Mockito.anyLong())).thenReturn(cache);
+ Mockito.when(manager.getScheduleExecutor()).thenReturn(executor);
+ try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+ mockedEnv.when(Env::getCurrentEnv).thenReturn(env);
+ if (choosePlanningMode) {
+ ConnectContext context = new ConnectContext();
+ context.setThreadLocalInfo();
+ try {
+ if (!node.isBatchMode()) {
+ return node.getSplits(1).size();
+ }
+ } finally {
+ ConnectContext.remove();
+ }
+ }
+ node.startSplit(1);
+ finished.get(5, TimeUnit.SECONDS);
+ Mockito.verify(assignment,
Mockito.never()).setException(Mockito.any());
+ if (choosePlanningMode) {
+ return delivered.get();
+ }
+ Assert.assertEquals(totalFiles, delivered.get());
+ return node.numApproximateSplits();
+ } finally {
+ executor.shutdownNow();
+ Assert.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java
index 12bfe0356a2..5408d4869d2 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java
@@ -23,11 +23,15 @@ import org.apache.doris.analysis.IntLiteral;
import org.apache.doris.analysis.SlotDescriptor;
import org.apache.doris.analysis.SlotId;
import org.apache.doris.analysis.SlotRef;
+import org.apache.doris.analysis.StatementBase;
import org.apache.doris.analysis.TableName;
import org.apache.doris.analysis.TableScanParams;
import org.apache.doris.analysis.TableSnapshot;
import org.apache.doris.analysis.TupleDescriptor;
import org.apache.doris.analysis.TupleId;
+import org.apache.doris.analysis.UserIdentity;
+import org.apache.doris.blockrule.SqlBlockRule;
+import org.apache.doris.blockrule.SqlBlockRuleMgr;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.DatabaseIf;
import org.apache.doris.catalog.Env;
@@ -65,6 +69,7 @@ import
org.apache.doris.datasource.iceberg.cache.IcebergManifestCacheLoader;
import org.apache.doris.datasource.iceberg.cache.ManifestCacheValue;
import org.apache.doris.datasource.iceberg.helper.IcebergWriterHelper;
import org.apache.doris.datasource.mvcc.MvccTableInfo;
+import org.apache.doris.mysql.privilege.Auth;
import org.apache.doris.nereids.StatementContext;
import org.apache.doris.planner.PlanNodeId;
import org.apache.doris.planner.ScanContext;
@@ -103,6 +108,7 @@ import org.apache.iceberg.FileContent;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.FileMetadata;
import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.ManifestFile;
import org.apache.iceberg.ManifestFiles;
import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.MetadataTableType;
@@ -123,6 +129,7 @@ import org.apache.iceberg.TableScan;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.hadoop.HadoopTables;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.FileIO;
import org.apache.iceberg.types.Types;
import org.apache.iceberg.util.ScanTaskUtil;
import org.apache.thrift.TDeserializer;
@@ -169,6 +176,87 @@ public class IcebergScanNodeTest {
@Rule
public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+ @Test
+ public void testVerboseExplainRequiresExactFilePlanning() throws Exception
{
+ assertBatchPlanningAllowed(true, false, 0, 0, true, false);
+ }
+
+ @Test
+ public void testSqlScanLimitsRequireExactFilePlanning() throws Exception {
+ assertBatchPlanningAllowed(false, false, 1, 0, true, false);
+ assertBatchPlanningAllowed(false, false, 0, 1, true, false);
+ assertBatchPlanningAllowed(false, true, 1, 0, true, false);
+ assertBatchPlanningAllowed(false, true, 0, 1, true, false);
+ }
+
+ @Test
+ public void testUnrestrictedQueriesKeepBatchPlanning() throws Exception {
+ assertBatchPlanningAllowed(false, false, 0, 0, true, true);
+ assertBatchPlanningAllowed(false, false, 1, 1, false, true);
+ }
+
+ @Test
+ public void testOrdinaryExplainKeepsTheRealBatchPlan() throws Exception {
+ assertBatchPlanningAllowed(true, false, false, 0, 0, true, true);
+ }
+
+ private void assertBatchPlanningAllowed(boolean explain, boolean global,
long partitions, long splits,
+ boolean enabled, boolean expected) throws Exception {
+ assertBatchPlanningAllowed(explain, explain, global, partitions,
splits, enabled, expected);
+ }
+
+ private void assertBatchPlanningAllowed(boolean explain, boolean verbose,
boolean global,
+ long partitions, long splits, boolean enabled, boolean expected)
throws Exception {
+ SessionVariable variables = new SessionVariable();
+ variables.enableExternalTableBatchMode = true;
+ variables.numFilesInBatchMode = 1;
+ TestIcebergScanNode node = new TestIcebergScanNode(variables);
+ Table table = Mockito.mock(Table.class);
+ FileIO io = Mockito.mock(FileIO.class);
+ TableScan scan = Mockito.mock(TableScan.class);
+ Snapshot snapshot = Mockito.mock(Snapshot.class);
+ ManifestFile manifest = Mockito.mock(ManifestFile.class);
+ Mockito.when(table.io()).thenReturn(io);
+ Mockito.when(table.specs()).thenReturn(Collections.singletonMap(0,
PartitionSpec.unpartitioned()));
+ Mockito.when(scan.table()).thenReturn(table);
+
Mockito.when(scan.schema()).thenReturn(PartitionSpec.unpartitioned().schema());
+ Mockito.when(scan.snapshot()).thenReturn(snapshot);
+ Mockito.when(scan.filter()).thenReturn(Expressions.alwaysTrue());
+ Mockito.when(snapshot.schemaId()).thenReturn(0);
+
Mockito.when(snapshot.dataManifests(io)).thenReturn(Collections.singletonList(manifest));
+ Mockito.when(manifest.partitionSpecId()).thenReturn(0);
+ Mockito.when(manifest.hasAddedFiles()).thenReturn(true);
+ Mockito.when(manifest.addedFilesCount()).thenReturn(2);
+ Mockito.when(manifest.existingFilesCount()).thenReturn(0);
+ node.setTableScan(scan);
+ setIcebergTable(node, table);
+ setPreExecutionAuthenticator(node, new ExecutionAuthenticator() {});
+
+ ConnectContext context = new ConnectContext();
+
context.setCurrentUserIdentity(UserIdentity.createAnalyzedUserIdentWithIp("batch_rule_user",
"%"));
+ StatementContext statementContext = new StatementContext();
+ StatementBase statement = Mockito.mock(StatementBase.class);
+ Mockito.when(statement.isExplain()).thenReturn(explain);
+ Mockito.when(statement.isVerbose()).thenReturn(verbose);
+ statementContext.setParsedStatement(statement);
+ context.setStatementContext(statementContext);
+ context.setThreadLocalInfo();
+ Env env = Mockito.mock(Env.class);
+ Auth auth = Mockito.mock(Auth.class);
+ SqlBlockRuleMgr rules = new SqlBlockRuleMgr();
+ rules.getNameToSqlBlockRuleMap().put("scan_limit", new SqlBlockRule(
+ "scan_limit", "NULL", "NULL", partitions, splits, 0L, false,
global, enabled));
+ Mockito.when(env.getAuth()).thenReturn(auth);
+ Mockito.when(env.getSqlBlockRuleMgr()).thenReturn(rules);
+ Mockito.when(auth.getSqlBlockRules("batch_rule_user")).thenReturn(new
String[] {"scan_limit"});
+ try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+ mockedEnv.when(Env::getCurrentEnv).thenReturn(env);
+ Assert.assertEquals(expected, node.isRealBatchMode());
+ } finally {
+ ConnectContext.remove();
+ }
+ }
+
@Test
public void testDoesNotUseHiveParquetInt96TimeZone() {
IcebergScanNode node = Mockito.mock(IcebergScanNode.class,
Mockito.CALLS_REAL_METHODS);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/maxcompute/source/MaxComputePartitionBatchTaskTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/maxcompute/source/MaxComputePartitionBatchTaskTest.java
new file mode 100644
index 00000000000..0ea63361ffa
--- /dev/null
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/maxcompute/source/MaxComputePartitionBatchTaskTest.java
@@ -0,0 +1,94 @@
+// 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.
+
+package org.apache.doris.datasource.maxcompute.source;
+
+import org.apache.doris.datasource.SplitAssignment;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
+
+public class MaxComputePartitionBatchTaskTest {
+ @Test
+ public void testBoundedContinuationUsesDedicatedBatchSize() {
+ SplitAssignment assignment = Mockito.mock(SplitAssignment.class);
+ Mockito.when(assignment.needMoreSplit()).thenReturn(true);
+ Queue<Runnable> queue = new ArrayDeque<>();
+ List<Integer> sizes = new ArrayList<>();
+ Executor executor = task -> {
+ Assert.assertTrue("one queued continuation per scan",
queue.isEmpty());
+ queue.add(task);
+ };
+ new MaxComputeScanNode.PartitionBatchTask(33, 16, executor, assignment,
+ (begin, end) -> sizes.add(end - begin)).submit();
+ while (!queue.isEmpty()) {
+ queue.remove().run();
+ }
+ Assert.assertEquals(Arrays.asList(16, 16, 1), sizes);
+ Mockito.verify(assignment).finishSchedule();
+ }
+
+ @Test
+ public void testRejectedInitialAndContinuationSubmissionsAreTerminal() {
+ for (boolean rejectInitial : Arrays.asList(true, false)) {
+ SplitAssignment assignment = Mockito.mock(SplitAssignment.class);
+ Mockito.when(assignment.needMoreSplit()).thenReturn(true);
+ Queue<Runnable> queue = new ArrayDeque<>();
+ Executor executor = task -> {
+ if (rejectInitial || queue.isEmpty()) {
+ throw new RejectedExecutionException("full");
+ }
+ queue.add(task);
+ };
+ if (!rejectInitial) {
+ queue.add(() -> {});
+ }
+ new MaxComputeScanNode.PartitionBatchTask(2, 1, executor,
assignment, (begin, end) -> {}).submit();
+ if (!rejectInitial) {
+ queue.remove();
+ queue.remove().run();
+ }
+ Mockito.verify(assignment).setException(Mockito.any());
+ Mockito.verify(assignment).finishSchedule();
+ }
+ }
+
+ @Test
+ public void testFailureAndCancellationDoNotScheduleMoreWork() {
+ SplitAssignment failed = Mockito.mock(SplitAssignment.class);
+ Mockito.when(failed.needMoreSplit()).thenReturn(true);
+ new MaxComputeScanNode.PartitionBatchTask(2, 1, Runnable::run, failed,
(begin, end) -> {
+ throw new java.io.IOException("session failed");
+ }).submit();
+ Mockito.verify(failed).setException(Mockito.any());
+ Mockito.verify(failed).finishSchedule();
+ SplitAssignment stopped = Mockito.mock(SplitAssignment.class);
+ new MaxComputeScanNode.PartitionBatchTask(2, 1, Runnable::run,
stopped, (begin, end) -> {
+ Assert.fail("a cancelled scan must not open another session");
+ }).submit();
+ Mockito.verify(stopped).finishSchedule();
+ }
+}
diff --git a/gensrc/thrift/FrontendService.thrift
b/gensrc/thrift/FrontendService.thrift
index 87027b02a4c..45e63623fa8 100644
--- a/gensrc/thrift/FrontendService.thrift
+++ b/gensrc/thrift/FrontendService.thrift
@@ -1681,6 +1681,8 @@ struct TSyncQueryColumns {
struct TFetchSplitBatchRequest {
1: optional i64 split_source_id
2: optional i32 max_num_splits
+ // The backend's shared scan consumer has stopped and will not fetch again.
+ 3: optional bool close_source
}
struct TFetchSplitBatchResult {
diff --git a/gensrc/thrift/PlanNodes.thrift b/gensrc/thrift/PlanNodes.thrift
index fb336e14dc1..4a31a5e9cda 100644
--- a/gensrc/thrift/PlanNodes.thrift
+++ b/gensrc/thrift/PlanNodes.thrift
@@ -779,6 +779,8 @@ struct TFileRangeDesc {
struct TSplitSource {
1: optional i64 split_source_id
2: optional i32 num_splits
+ // This FE accepts TFetchSplitBatchRequest.close_source. Absent on older
FEs.
+ 3: optional bool supports_close
}
// TFileScanRange represents a set of descriptions of a file and the rules for
reading and converting it.
diff --git
a/regression-test/suites/external_table_p0/hive/test_hive_partitions.groovy
b/regression-test/suites/external_table_p0/hive/test_hive_partitions.groovy
index b8b3a39a2f9..8e64896e988 100644
--- a/regression-test/suites/external_table_p0/hive/test_hive_partitions.groovy
+++ b/regression-test/suites/external_table_p0/hive/test_hive_partitions.groovy
@@ -195,10 +195,15 @@ suite("test_hive_partitions",
"p0,external,hive,external_docker,external_docker_
sql """set num_partitions_in_batch_mode=1"""
explain {
sql ("select * from partition_table")
- verbose (true)
-
contains "(approximate)inputSplitNum=60"
}
+ explain {
+ sql ("select * from partition_table")
+ verbose (true)
+ contains "inputSplitNum=6"
+ contains "partition=6/6"
+ notContains "approximate"
+ }
sql """unset variable num_partitions_in_batch_mode"""
} finally {
}
diff --git
a/regression-test/suites/external_table_p0/iceberg/test_iceberg_variant_read.groovy
b/regression-test/suites/external_table_p0/iceberg/test_iceberg_variant_read.groovy
index b0190b28813..17f15874854 100644
---
a/regression-test/suites/external_table_p0/iceberg/test_iceberg_variant_read.groovy
+++
b/regression-test/suites/external_table_p0/iceberg/test_iceberg_variant_read.groovy
@@ -1412,6 +1412,8 @@ public class AppendVariantEqualityDelete {
// seeking its unshredded value. More than four rows qualify, forcing
local TopN overshoot to be
// truncated after the merge exchange while the mapper-eligible projected
path crosses the wire.
explain {
+ // The exact two-file fixture check needs complete planning rather
than a batch estimate.
+ verbose true
sql """
SELECT id, CAST(projected['n'] AS INT)
FROM (
diff --git
a/regression-test/suites/external_table_p0/iceberg/write/test_iceberg_write_uuid_inputs.groovy
b/regression-test/suites/external_table_p0/iceberg/write/test_iceberg_write_uuid_inputs.groovy
index 86acef3627d..f7a32d564de 100644
---
a/regression-test/suites/external_table_p0/iceberg/write/test_iceberg_write_uuid_inputs.groovy
+++
b/regression-test/suites/external_table_p0/iceberg/write/test_iceberg_write_uuid_inputs.groovy
@@ -98,6 +98,7 @@ suite("test_iceberg_write_uuid_inputs",
"p0,external,iceberg,external_docker,ext
// Both the UUID and NULL identity partitions must survive
FE partition parsing.
explain {
sql "SELECT id, u FROM ${table}"
+ verbose true
contains "partition=2/"
}
sql "INSERT OVERWRITE TABLE ${table}
PARTITION(u='${compact}') VALUES (31)"
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
