This is an automated email from the ASF dual-hosted git repository.
924060929 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 439906a36a4 [opt](hive) Decouple partition metadata batches from
directory tasks (#70247)
439906a36a4 is described below
commit 439906a36a4933f15fdfc785e00f51ea16639510
Author: 924060929 <[email protected]>
AuthorDate: Fri Nov 6 09:11:21 2026 +0800
[opt](hive) Decouple partition metadata batches from directory tasks
(#70247)
### Background
Partitioned Hive Batch Mode fetches partition objects from Hive
Metastore (HMS), lists directories and streams splits into bounded
queues. Small directory work units allow scans to start early and
producers to yield. HMS has a separate property limiting partition names
per RPC, but the streaming path previously also used the small
directory-task size for those RPCs.
### What problem does this PR solve?
Issue Number: None
Related PR: None
PR Type: Performance
Problem Summary:
A cold scan of 1000 partitions could make 63 partition-object requests
with a directory batch size of 16. Separate bounded metadata acquisition
from directory work: fetch a small first HMS batch, then larger bounded
batches, while readers continue claiming small directory slices. Deliver
the first slice without walking the next large cached batch, and resume
sources through callbacks when metadata becomes available. Profiling
also exposed redundant per-split destination selection under the
assignment lock when only one backend is eligible; that case now assigns
remote batches directly to their only possible owner. Session defaults
and BE code are unchanged.
### Details
- **Change map:** 636 changed production lines, 1159 test lines, 10
documentation lines and 4 build/API-baseline lines; no generated files.
Hive loader/source integration owns shared batch cursors and
cancellation; an HMS cache-only probe avoids remote reads when cached;
an optional scan SPI and generic readiness handling release workers
while metadata is pending. The common backend policy adds a 19-line
single-destination fast path.
- **Why it was slow:** small directory units caused excessive HMS round
trips. Preparing a subsequent large cached batch synchronously inside
the last-slice `take()` also delayed delivery of the first small slice.
Separately, single-backend streaming assignment still hashed every
split, built candidate/alternative-host lists and updated weight maps
per split under `SplitAssignment.assignBatch`'s monitor.
- **Optimization:** fetch `min(16,
hive.hms_partitions_batch_size_per_rpc)` names first, then use the
existing positive RPC cap (default 5000). Share one ready result list;
readers claim at most the existing directory batch size. Admit the next
large cache probe without waiting; an idle reader may help its own
probe, with `FutureTask` electing one executor. Only the existing
METADATA scheduler performs remote loads, and readiness callbacks resume
readers on the existing SPLIT scheduler. For one eligible backend and
wholly remote batches, preserve split order and exact accumulated
weights while skipping candidate selection. Local-only and
multiple-backend batches retain the general path.
- **Results:** the cold single-query full scan fetches 1000 HMS
partitions in two physical RPCs (16 and 984), compared with 63 in the
parent streaming path. The controlled SQL comparison below isolates the
subsequent first-slice and assignment corrections; all three columns
already contain bounded HMS batching. A separate same-JVM experiment
confirms lower FE CPU/allocation and a LIMIT improvement; full-scan
latency still has substantial variance, so no fixed 22 ms penalty or
stable full-scan speedup is claimed.
- **Applicability and cost:** targets partitioned Hive scans with HMS
partition-cache misses; a directory's file Listing retains its existing
behavior. Buffered metadata is bounded by one configured HMS result
batch plus slices held by bounded readers. A remote call already running
at cancellation may finish; buffered/late results are discarded and
subsequent requests are prevented. Cache helpers do no remote I/O.
Single-backend assignment adds a linear remote-accessibility check and
replaces per-split hash/candidate/map work with one exact weight
accumulation. No additional pool or cache is introduced.
- **Behavior change:** metadata RPC grouping and source readiness
change. Split batch size, directory batch size, queue capacity, pruning,
invalidation, query-pinned handles, early stop, scheduler quanta and
backpressure retain their meanings. Connectors that do not opt in retain
their path. Connector API changes from 29 to 30; deploy the matching
connector bundle with FE. No FE/BE wire-format change.
#### Controlled complete-SQL comparison
All packages use parent `315c2e003317ef44e8a8fbb445a1ec8452d5647e`.
“Initial batching” is `5e7b1c2ca33d3c20119e95eebaab9bd35d90e6c2`;
“First-slice correction” is `c1e67745edb537a96109f6a2f43d2bd75a249417`;
“Final” also includes the single-backend fast path. Initial/corrected
packages differ only in the Hive loader classes. Corrected/final
measurement packages differ in exactly `FederationBackendPolicy.class`,
compiled by the successful full FE build, to isolate the optimization.
| Scenario | Initial batching (ms) | First-slice correction (ms) | Final
(ms) |
|---|---:|---:|---:|
| Single query, cold, LIMIT 10 | 20.75 | 20.20 | 20.80 |
| Single query, cold, full scan | 155.88 | 160.56 | 166.42 |
| Single query, warm, LIMIT 10 | 8.14 | 7.26 | 7.32 |
| Single query, warm, full scan | 86.38 | 81.79 | 84.69 |
| Four-query wave, cold, LIMIT 10 | 46.43 | 47.03 | 45.83 |
| Four-query wave, cold, full scan | 370.16 | 375.65 | 368.44 |
| Four-query wave, warm, LIMIT 10 | 11.21 | 11.27 | 9.61 |
| Four-query wave, warm, full scan | 235.75 | 232.23 | 236.50 |
ARM64 Docker, identical Release BE SHA256
`f1ac268df25c8f9a6f8095e3dda20f6b0dbc591cb4609b03612dd5a913bfd0b1`, one
engine at a time, 6 CPUs and 14 GiB, FE heap and BE memory limit 4 GiB
each. MinIO retains its CPU set and memory limit with CFS quota removed;
every timed wave verifies zero throttling. User VMs remain stopped. The
fixture has 1000 partitions, ten small Parquet files per partition and
five rows per file: 10000 files, 50000 rows.
Settings: `enable_external_table_batch_mode=true`,
`num_partitions_in_batch_mode=1`, `split_batch_size=8`,
`split_queue_capacity=1024`, `split_partition_batch_size=16`,
`split_producer_parallelism=0`, `enable_sql_cache=false`,
`enable_stats=false`. Cold/warm refer to FE metadata and file-list
caches; JVM and BE footer/data caches are warmed. Cache gates drain
previous tasks and retain connector identity. Each phase runs the same
3000-SQL warmup. Order: initial/corrected/final/final/corrected/initial;
each cell has 32 samples per variant. Four-query latency ends when all
four synchronized persistent clients complete. All 1920 formal SQL
executions validate results. Profile, bytecode tracing and sampling are
disabled during formal timing. Across 240 independent query Profiles,
all 120 full scans verify 10000 footer hits, zero footer reads, 20000
page hits and zero page misses/writes.
Phase effects are material: unchanged corrected code has warm four-query
full-scan medians of 221.36 and 240.81 ms across its two phases.
Separate randomized, paired diagnostics therefore keep the same JVM and
BE process. Choosing first-return, inline-cache-probe or wait-for-probe
ordering yields full-scan medians of 209.61 / 210.33 / 208.47 ms, with
45 waves per mode and twelve split-fetch RPCs per query in all modes;
this does not establish a fixed first-return penalty.
With only the new assignment fast path enabled/disabled in that
same-process style of diagnostic (60 randomized waves per mode),
full-scan medians are 226.74 / 237.40 ms. The median paired difference
is -5.31 ms, bootstrap 95% interval [-9.18, +1.28] ms: stable full-scan
latency improvement is not established. LIMIT medians are 10.30 / 11.37
ms, paired difference -1.15 ms, interval [-1.60, -0.52] ms. Full-wave FE
process CPU falls from 215 to 170 ms. This temporary diagnostic guard is
absent from production code and the formal comparison above.
Independent 240-query flamegraph windows in the same JVM reduce
assignment CPU samples from 5.03 to 0.76 ms/query and sampled
allocations from 16.77 to 1.23 MiB/query. Hash-ring samples disappear.
The assignment monitor's cumulative waiting falls from 676.63 to 39.77
ms/query; these are sums across overlapping producer waits, not SQL
elapsed time. The first-slice helper remains around 0.36–0.43 ms
CPU/query. The full scan continues to spend most elapsed time waiting
for scan results, so reduced producer work need not translate directly
into the same wall-time reduction.
### Testing
- Current head, local ARM64 Docker limited to 2 CPUs: `./build.sh --fe
-j2`, including Checkstyle, and 130 focused tests all pass:
`FederationBackendPolicyTest`, `FederationBackendPolicySpreadTest`,
`SplitProducerTest`, `SplitSourceTest`,
`PluginDrivenScanNodeHiveStreamingTest`.
- The new single-backend legacy-hash test fails before the optimization
with 140 unnecessary split-identity reads and passes after it. Six
selection settings, exact cumulative weights across empty/small/large
batches, local-only validation and overflow are covered; existing
multiple-backend tests pass.
- The Hive/HMS/SPI implementation previously passed 578 Hive, 184 HMS,
85 SPI and 71 focused core tests. Those broader connector suites precede
the isolated backend-policy addition; current-head core and complete-SQL
validation are listed separately above.
- `takingFirstDirectoryBatchDoesNotWalkTheNextLargeCachedBatch`
reproduces the first-slice barrier. Deterministic coverage also includes
helper/worker execution exactly once, full scheduler queues, one-slot
handoff, cancellation during misses, queued cancellation, remote
failures/late results, bounded batches, dropped partitions and readiness
completion races.
- `git diff --check` passes. No BE change or rebuild. Runtime test
services and temporary quota changes are restored after testing.
### Notes for Reviewers
Start with `HivePartitionBatchLoader` and
`HiveScanPlanProvider.streamPartitionSources`, then the generic adapter
in `PluginDrivenScanNode` and readiness handling in `SplitProducer`. The
loader monitor protects cursors/results/failure only; cache work, RPCs,
scheduling and future completion run outside it. `FutureTask.get()` is
used by `done()` after completion and never waits there. A reader
retains only its slice. Backend-policy state keeps its existing
assignment lock; the single-destination fast path reduces work inside
that lock.
### Release note
Reduce cold Hive partition planning round trips while retaining early
split delivery, cancellation and bounded queues. Reduce FE CPU and
temporary allocations for remote split assignment to one eligible
backend.
### AI Disclosure
- Tool: OpenAI Codex.
- Generated: Implementation, tests, self-review and this description.
- Human review: The author directed the design and investigation. Agent
self-review is complete; human line-by-line review is pending.
---
fe/fe-connector/README.md | 10 +
.../connector/hive/HiveCatalogProperties.java | 4 +
.../connector/hive/HivePartitionBatchLoader.java | 317 ++++++++++++
.../doris/connector/hive/HiveScanPlanProvider.java | 135 ++++-
.../hive/HivePartitionBatchLoaderTest.java | 543 +++++++++++++++++++++
.../connector/hive/HiveScanBatchModeTest.java | 269 ++++++++++
.../doris/connector/hms/CachingHmsClient.java | 19 +
.../org/apache/doris/connector/hms/HmsClient.java | 11 +
.../doris/connector/hms/CachingHmsClientTest.java | 66 +++
.../spi/scan/ConnectorScanPlanProvider.java | 13 +
.../connector/spi/scan/ConnectorSplitSource.java | 14 +
.../connector/spi/ConnectorPluginSurfaceTest.java | 2 +-
.../test/resources/connector-plugin-surface.txt | 2 +
fe/fe-connector/pom.xml | 2 +-
.../datasource/scan/FederationBackendPolicy.java | 19 +
.../datasource/scan/PluginDrivenScanNode.java | 71 +++
.../doris/datasource/split/SplitProducer.java | 33 ++
.../PluginDrivenScanNodePartitionBatchTest.java | 63 +++
.../doris/datasource/split/SplitProducerTest.java | 114 +++++
.../doris/planner/FederationBackendPolicyTest.java | 102 ++++
20 files changed, 1803 insertions(+), 6 deletions(-)
diff --git a/fe/fe-connector/README.md b/fe/fe-connector/README.md
index 1b61b82c0ab..cfcfd5d496c 100644
--- a/fe/fe-connector/README.md
+++ b/fe/fe-connector/README.md
@@ -118,6 +118,16 @@ before releasing read transactions and statement metadata.
An SDK call that
cannot be interrupted may finish later; it must not issue another request
after observing cancellation.
+Partitioned connectors may implement `streamPartitionSources` to batch metadata
+independently of directory work. The hook opens a bounded set of cooperating
+sources without remote I/O. A source reports a metadata dependency through
+`isBlocked`; the engine releases its worker until completion. If `hasNext`
+returns false while waiting, retain that dependency, including a completed one,
+until the next read attempt so readiness cannot be mistaken for end of input.
+Use the existing METADATA scheduler for metadata loading and let SPLIT sources
+claim small directory batches; neither stage may wait for child tasks. The
+default hook keeps ordinary partition-batch sources unchanged.
+
`SET first_split_wait_timeout_ms = 2500` bounds the first split wait, including
enumeration executor queueing. The default is 30000 ms; values must be
positive.
`SET_VAR` and FE forwarding are supported. Preparation before enumeration is
diff --git
a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveCatalogProperties.java
b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveCatalogProperties.java
index c3e243f5124..ebbeb07e8c8 100644
---
a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveCatalogProperties.java
+++
b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveCatalogProperties.java
@@ -246,6 +246,10 @@ public final class HiveCatalogProperties {
return hmsClientPoolSize;
}
+ public int getHmsPartitionsBatchSizePerRpc() {
+ return hmsPartitionsBatchSizePerRpc;
+ }
+
public boolean isHmsEventsIncrementalSyncEnabled() {
return enableHmsEventsIncrementalSync;
}
diff --git
a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HivePartitionBatchLoader.java
b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HivePartitionBatchLoader.java
new file mode 100644
index 00000000000..3cbb4b8845c
--- /dev/null
+++
b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HivePartitionBatchLoader.java
@@ -0,0 +1,317 @@
+// 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.connector.hive;
+
+import org.apache.doris.connector.hms.HmsPartitionInfo;
+import org.apache.doris.connector.spi.ConnectorTaskScheduler;
+import org.apache.doris.connector.spi.scan.ConnectorScanContext;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+
+/**
+ * One scan's serial HMS producer. Readers claim small directory batches from
one bounded result list.
+ * Claiming its last slice permits one next metadata batch, overlapping HMS
with directory work without
+ * waiting for children or accumulating an unbounded queue. Large cache probes
are queued separately
+ * from remote reads, so an idle reader can help its own cache probe without
doing remote I/O.
+ */
+final class HivePartitionBatchLoader {
+ private static final int INITIAL_BATCH_SIZE = 16;
+ private final ConnectorTaskScheduler scheduler;
+ private final ConnectorScanContext scanContext;
+ private final List<String> names;
+ private final int maxBatchSize;
+ private final Function<List<String>, Optional<List<HmsPartitionInfo>>>
cached;
+ private final Function<List<String>, List<HmsPartitionInfo>> loader;
+ // Serializes synchronous cache hits and inline/offline executors without
recursive submission.
+ private final AtomicInteger pumping = new AtomicInteger();
+ // All fields below are guarded by this. I/O, scheduler calls and future
completion run outside it.
+ private int nextName;
+ private List<HmsPartitionInfo> ready = Collections.emptyList();
+ private int nextReady;
+ private CompletableFuture<Void> available = new CompletableFuture<>();
+ private LoadTask current;
+ private Throwable failure;
+ private int readers;
+
+ HivePartitionBatchLoader(ConnectorTaskScheduler scheduler,
ConnectorScanContext scanContext,
+ List<String> names, int maxBatchSize, int readers,
+ Function<List<String>, Optional<List<HmsPartitionInfo>>> cached,
+ Function<List<String>, List<HmsPartitionInfo>> loader) {
+ this.scheduler = scheduler;
+ this.scanContext = scanContext;
+ this.names = names;
+ this.maxBatchSize = maxBatchSize;
+ this.readers = readers;
+ this.cached = cached;
+ this.loader = loader;
+ scanContext.whenCancelled(() -> fail(new CancellationException("Split
generation stopped")));
+ }
+
+ CompletableFuture<Void> isBlocked() {
+ pump();
+ LoadTask task;
+ synchronized (this) {
+ task = current;
+ }
+ if (task != null) {
+ task.helpCacheProbe();
+ }
+ synchronized (this) {
+ return failure != null || (ready.isEmpty() && (current != null ||
nextName < names.size()))
+ ? available : null;
+ }
+ }
+
+ synchronized boolean isFinished() {
+ return failure == null && ready.isEmpty() && current == null &&
nextName == names.size();
+ }
+
+ List<HmsPartitionInfo> take(int size) {
+ List<HmsPartitionInfo> result;
+ synchronized (this) {
+ if (failure != null) {
+ throw new CompletionException(failure);
+ }
+ if (ready.isEmpty()) {
+ return Collections.emptyList();
+ }
+ int end = Math.min(nextReady + size, ready.size());
+ // Own a small reference array: a slow reader must not retain an
entire completed HMS batch.
+ result = new ArrayList<>(ready.subList(nextReady, end));
+ nextReady = end;
+ if (nextReady == ready.size()) {
+ ready = Collections.emptyList();
+ nextReady = 0;
+ available = new CompletableFuture<>();
+ }
+ }
+ pump();
+ return result;
+ }
+
+ void readerClosed() {
+ boolean last;
+ synchronized (this) {
+ last = --readers == 0;
+ }
+ if (last) {
+ fail(new CancellationException("All partition readers closed"));
+ }
+ }
+
+ private void pump() {
+ if (pumping.getAndIncrement() != 0) {
+ return;
+ }
+ int missed = 1;
+ do {
+ LoadTask task = null;
+ synchronized (this) {
+ if (failure == null && current == null && ready.isEmpty() &&
nextName < names.size()) {
+ int size = nextName == 0 ? Math.min(INITIAL_BATCH_SIZE,
maxBatchSize) : maxBatchSize;
+ int end = nextName + Math.min(size, names.size() -
nextName);
+ task = new LoadTask(names.subList(nextName, end));
+ nextName = end;
+ current = task;
+ }
+ }
+ if (task != null) {
+ try {
+ scanContext.checkCancelled();
+ if (task.deferredCacheProbe) {
+ // Return the directory slice before walking the next
large cached batch.
+ task.cacheProbe.admit();
+ } else {
+ Optional<List<HmsPartitionInfo>> hit =
cached.apply(task.batch);
+ if (hit.isPresent()) {
+ publish(task, hit.get());
+ } else {
+ task.remoteLoad.admit();
+ }
+ }
+ } catch (Throwable e) {
+ fail(e);
+ }
+ }
+ missed = pumping.addAndGet(-missed);
+ } while (missed != 0);
+ }
+
+ private void publish(LoadTask task, List<HmsPartitionInfo> partitions) {
+ CompletableFuture<Void> signal = null;
+ boolean empty;
+ synchronized (this) {
+ if (failure != null) {
+ return;
+ }
+ current = null;
+ ready = partitions;
+ empty = ready.isEmpty();
+ if (!empty || nextName == names.size()) {
+ signal = available;
+ }
+ }
+ task.complete(null);
+ if (signal != null) {
+ signal.complete(null);
+ }
+ if (empty) {
+ pump();
+ }
+ }
+
+ private void fail(Throwable cause) {
+ CompletableFuture<Void> signal;
+ LoadTask task;
+ synchronized (this) {
+ if (failure != null) {
+ return;
+ }
+ failure = cause;
+ ready = Collections.emptyList();
+ if (available.isDone()) {
+ available = new CompletableFuture<>();
+ }
+ signal = available;
+ task = current;
+ current = null;
+ }
+ signal.completeExceptionally(cause);
+ if (task != null) {
+ task.completeExceptionally(cause);
+ }
+ }
+
+ private final class LoadTask extends CompletableFuture<Void> {
+ private final List<String> batch;
+ private final boolean deferredCacheProbe;
+ private final ScheduledStage cacheProbe;
+ private final ScheduledStage remoteLoad;
+
+ private LoadTask(List<String> batch) {
+ this.batch = batch;
+ this.deferredCacheProbe = batch.size() > INITIAL_BATCH_SIZE;
+ this.remoteLoad = new ScheduledStage(() -> {
+ List<HmsPartitionInfo> partitions = loader.apply(batch);
+ scanContext.checkCancelled();
+ publish(this, partitions);
+ });
+ this.cacheProbe = new ScheduledStage(() -> {
+ Optional<List<HmsPartitionInfo>> hit = cached.apply(batch);
+ // A cancellation during a cache miss must not start another
HMS request.
+ scanContext.checkCancelled();
+ if (hit.isPresent()) {
+ publish(this, hit.get());
+ } else {
+ remoteLoad.admit();
+ }
+ });
+ whenComplete((unused, error) -> {
+ cacheProbe.cancel(false);
+ remoteLoad.cancel(false);
+ });
+ }
+
+ private void helpCacheProbe() {
+ if (deferredCacheProbe) {
+ // FutureTask elects one executor, even if a metadata worker
dequeues it concurrently.
+ // Helpers only probe the cache; a miss submits remoteLoad and
returns without waiting.
+ cacheProbe.run();
+ }
+ }
+
+ private final class ScheduledStage extends FutureTask<Void> {
+ private final AtomicInteger admitting = new AtomicInteger();
+ private boolean admitted;
+ private final Runnable onSpace = this::admit;
+
+ private ScheduledStage(Runnable action) {
+ super(() -> {
+ ClassLoader previous =
Thread.currentThread().getContextClassLoader();
+ try {
+
Thread.currentThread().setContextClassLoader(HivePartitionBatchLoader.class.getClassLoader());
+ scanContext.checkCancelled();
+ if (!LoadTask.this.isDone()) {
+ action.run();
+ }
+ return null;
+ } finally {
+ Thread.currentThread().setContextClassLoader(previous);
+ }
+ });
+ }
+
+ private void admit() {
+ // Capacity callbacks can race admission, including inline
notification.
+ if (admitting.getAndIncrement() != 0) {
+ return;
+ }
+ int missed = 1;
+ do {
+ // Cancelling one stage can synchronously notify another
stage's capacity callback.
+ if (!isDone() && !LoadTask.this.isDone() && !admitted) {
+ try {
+ admitted = scheduler.tryExecute(this,
this::setException, onSpace);
+ if (admitted) {
+ scheduler.removeSpaceListener(onSpace);
+ }
+ } catch (RuntimeException e) {
+ setException(e);
+ }
+ }
+ if (isDone()) {
+ scheduler.removeSpaceListener(onSpace);
+ scheduler.remove(this);
+ }
+ missed = admitting.addAndGet(-missed);
+ } while (missed != 0);
+ }
+
+ @Override
+ protected void done() {
+ scheduler.removeSpaceListener(onSpace);
+ // A helper may have executed this task while its queue entry
was still pending.
+ scheduler.remove(this);
+ if (LoadTask.this.isDone()) {
+ return;
+ }
+ try {
+ // FutureTask invokes done only after completion; get
never waits here.
+ get();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ fail(e);
+ } catch (ExecutionException e) {
+ fail(e.getCause());
+ } catch (CancellationException e) {
+ fail(e);
+ }
+ }
+ }
+ }
+}
diff --git
a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveScanPlanProvider.java
b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveScanPlanProvider.java
index f9e8e4364ed..c2312582227 100644
---
a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveScanPlanProvider.java
+++
b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveScanPlanProvider.java
@@ -56,6 +56,7 @@ import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
@@ -277,6 +278,83 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
return new HiveSplitSource(session, hiveHandle,
request.getRequiredPartitions(), request.getScanContext());
}
+ @Override
+ public Optional<List<ConnectorSplitSource>>
streamPartitionSources(ConnectorSession session,
+ ConnectorScanRequest request, int partitionsPerTask, int
maxSources) {
+ HiveTableHandle handle = (HiveTableHandle) request.getTableHandle();
+ rejectTransactionalRead(handle);
+ List<String> names = request.getRequiredPartitions();
+ int sourceCount = Math.min(maxSources, (names.size() - 1) /
partitionsPerTask + 1);
+ HivePartitionBatchLoader loader = new HivePartitionBatchLoader(
+
session.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA),
request.getScanContext(), names,
+ catalogProperties.getHmsPartitionsBatchSizePerRpc(),
sourceCount,
+ batch -> cachedBatchPartitions(handle, batch),
+ batch -> resolveBatchPartitions(handle, batch,
request.getScanContext()));
+ List<ConnectorSplitSource> sources = new ArrayList<>(sourceCount);
+ for (int i = 0; i < sourceCount; i++) {
+ sources.add(new ConnectorSplitSource() {
+ private HiveSplitSource source;
+ private boolean closed;
+
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ return closed ? null : source == null ? loader.isBlocked()
: source.isBlocked();
+ }
+
+ @Override
+ public boolean hasNext() {
+ if (closed) {
+ return false;
+ }
+ if (source == null) {
+ source = new HiveSplitSource(session, handle, names,
request.getScanContext(),
+ partitionsPerTask);
+ source.partitionLoader = loader;
+ }
+ return source.hasNext();
+ }
+
+ @Override
+ public ConnectorScanRange next() {
+ return source.next();
+ }
+
+ @Override
+ public void close() {
+ if (!closed) {
+ closed = true;
+ if (source == null) {
+ loader.readerClosed();
+ } else {
+ source.close();
+ }
+ }
+ }
+ });
+ }
+ return Optional.of(sources);
+ }
+
+ private Optional<List<HmsPartitionInfo>>
cachedBatchPartitions(HiveTableHandle handle, List<String> names) {
+ Map<String, HmsPartitionInfo> retained =
handle.getPrunedPartitionsByName();
+ if (!retained.isEmpty()) {
+ List<HmsPartitionInfo> result = new ArrayList<>(names.size());
+ for (String name : names) {
+ HmsPartitionInfo partition = retained.get(name);
+ if (partition == null) {
+ return Optional.empty();
+ }
+ result.add(partition);
+ }
+ return Optional.of(result);
+ }
+ Optional<HmsPartitionBatchResult> result =
hmsClient.getCachedPartitionsWithStats(
+ handle.getDbName(), handle.getTableName(), names);
+ result.ifPresent(hit -> partitionBatchProfile.record(
+ handle.getDbName(), handle.getTableName(), hit.getStats()));
+ return result.map(HmsPartitionBatchResult::getPartitions);
+ }
+
private static void rejectTransactionalRead(HiveTableHandle handle) {
// Both full ACID and insert-only layouts need write-id visibility;
plain file scans
// (including metadata-only counts) cannot return a transactionally
correct result.
@@ -300,6 +378,8 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
private final boolean splittable;
private final long targetSplitSize;
private final int partitionBatchSize;
+ private HivePartitionBatchLoader partitionLoader;
+ private CompletableFuture<Void> blockedOnPartitions;
private Iterator<PartitionScanInfo> partitions =
Collections.emptyIterator();
private Iterator<HiveFileStatus> files = Collections.emptyIterator();
private Iterator<ConnectorScanRange> ranges =
Collections.emptyIterator();
@@ -311,6 +391,12 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
HiveSplitSource(ConnectorSession session, HiveTableHandle handle,
List<String> requiredPartitions,
ConnectorScanContext scanContext) {
+ this(session, handle, requiredPartitions, scanContext, Math.max(1,
+ sessionInt(session, "split_partition_batch_size",
DEFAULT_PARTITION_BATCH_SIZE)));
+ }
+
+ HiveSplitSource(ConnectorSession session, HiveTableHandle handle,
List<String> requiredPartitions,
+ ConnectorScanContext scanContext, int partitionsPerTask) {
rejectTransactionalRead(handle);
this.handle = handle;
this.scanContext = scanContext;
@@ -323,13 +409,21 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
this.isLzo = isLzoInputFormat(handle.getInputFormat());
this.splittable = tableFormat.isSplittable() && !isLzo;
this.targetSplitSize = getTargetSplitSize(session);
- this.partitionBatchSize = Math.max(1,
- sessionInt(session, "split_partition_batch_size",
DEFAULT_PARTITION_BATCH_SIZE));
+ this.partitionBatchSize = partitionsPerTask;
recordPruningProfile(handle);
}
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ if (closed || partitionLoader == null || ranges.hasNext() ||
files.hasNext() || partitions.hasNext()) {
+ return null;
+ }
+ return blockedOnPartitions != null ? blockedOnPartitions :
partitionLoader.isBlocked();
+ }
+
@Override
public boolean hasNext() {
+ blockedOnPartitions = null;
// Empty partitions can keep hasNext() inside this source across
many directory requests.
while (!closed && !scheduler.isClosed()) {
scanContext.checkCancelled();
@@ -345,6 +439,15 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
continue;
}
if (!loadNextPartitions()) {
+ if (partitionLoader != null) {
+ // Keep this dependency even if it completes before
the caller checks it:
+ // an empty read while waiting must not be mistaken
for end of input.
+ blockedOnPartitions = partitionLoader.isBlocked();
+ if (blockedOnPartitions == null &&
!partitionLoader.isFinished()) {
+ // Another reader may have taken the list after
our readiness check.
+ continue;
+ }
+ }
return false;
}
}
@@ -369,6 +472,11 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
}
private boolean loadNextPartitions() {
+ if (partitionLoader != null) {
+ partitions =
convertPartitions(partitionLoader.take(partitionBatchSize),
+ handle.getPartitionKeyNames()).iterator();
+ return partitions.hasNext();
+ }
List<String> keys = handle.getPartitionKeyNames();
if (keys == null || keys.isEmpty()) {
if (unpartitionedInitialized) {
@@ -399,8 +507,18 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
}
int end = Math.min(partitionNames.size(), nextPartition +
partitionBatchSize);
scanContext.checkCancelled();
- partitions = convertPartitions(resolveBatchPartitions(handle,
- partitionNames.subList(nextPartition, end), scanContext),
keys).iterator();
+ List<String> batch = partitionNames.subList(nextPartition, end);
+ // A retained pruning view belongs to this handle; never replace
it with a newer cache view.
+ Optional<HmsPartitionBatchResult> cached =
handle.getPrunedPartitionsByName().isEmpty()
+ ?
hmsClient.getCachedPartitionsWithStats(handle.getDbName(),
handle.getTableName(), batch)
+ : Optional.empty();
+ if (cached.isPresent()) {
+ partitionBatchProfile.record(handle.getDbName(),
handle.getTableName(), cached.get().getStats());
+ partitions = convertPartitions(cached.get().getPartitions(),
keys).iterator();
+ nextPartition = end;
+ return true;
+ }
+ partitions = convertPartitions(resolveBatchPartitions(handle,
batch, scanContext), keys).iterator();
nextPartition = end;
return true;
}
@@ -415,10 +533,18 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
@Override
public void close() {
+ if (closed) {
+ return;
+ }
+ if (partitionLoader != null) {
+ partitionLoader.readerClosed();
+ partitionLoader = null;
+ }
closed = true;
partitions = Collections.emptyIterator();
files = Collections.emptyIterator();
ranges = Collections.emptyIterator();
+ blockedOnPartitions = null;
partitionNames = null;
partition = null;
}
@@ -464,6 +590,7 @@ public class HiveScanPlanProvider implements
ConnectorScanPlanProvider {
ConnectorSession session,
ConnectorScanRequest request,
List<String> partitionBatch) {
+ rejectTransactionalRead((HiveTableHandle) request.getTableHandle());
// Batch mode deliberately does not retain completed ranges in the
statement scope: its
// purpose is to bound FE memory while splits are streamed to the
coordinator. Caching every
// batch until statement close would materialize the full scan again
and defeat that bound.
diff --git
a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HivePartitionBatchLoaderTest.java
b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HivePartitionBatchLoaderTest.java
new file mode 100644
index 00000000000..c06603dbbb9
--- /dev/null
+++
b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HivePartitionBatchLoaderTest.java
@@ -0,0 +1,543 @@
+// 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.connector.hive;
+
+import org.apache.doris.connector.hms.HmsPartitionInfo;
+import org.apache.doris.connector.spi.ConnectorTaskScheduler;
+import org.apache.doris.connector.spi.scan.ConnectorScanContext;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+class HivePartitionBatchLoaderTest {
+ @Test
+ void firstBatchStartsSmallAndRemainingBatchIsSharedWithoutDuplicates() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ List<Integer> requests = new ArrayList<>();
+ HivePartitionBatchLoader loader = loader(scheduler, new
ConnectorScanContext(), 1000, 1024, batch -> {
+ requests.add(batch.size());
+ return partitions(batch);
+ });
+ Assertions.assertTrue(requests.isEmpty());
+ CompletableFuture<Void> initial = loader.isBlocked();
+ Assertions.assertFalse(initial.isDone());
+ scheduler.runNext();
+ Assertions.assertTrue(initial.isDone());
+ List<HmsPartitionInfo> all = new ArrayList<>(loader.take(16));
+ Assertions.assertEquals(1, scheduler.tasks.size(), "next HMS batch
overlaps the first directory batch");
+ scheduler.runAll();
+ Assertions.assertEquals(List.of(16, 984), requests);
+ while (!loader.isFinished()) {
+ List<HmsPartitionInfo> batch = loader.take(16);
+ Assertions.assertTrue(batch.size() <= 16);
+ all.addAll(batch);
+ }
+ Assertions.assertEquals(1000, all.size());
+ Assertions.assertEquals(1000, new
HashSet<>(all.stream().map(HmsPartitionInfo::getValues)
+ .collect(java.util.stream.Collectors.toList())).size());
+ }
+
+ @Test
+ void takingFirstDirectoryBatchDoesNotWalkTheNextLargeCachedBatch() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ List<Integer> probes = new ArrayList<>();
+ HivePartitionBatchLoader loader = new
HivePartitionBatchLoader(scheduler, new ConnectorScanContext(),
+ names(1000), 1024, 3, batch -> {
+ probes.add(batch.size());
+ return Optional.of(partitions(batch));
+ }, batch -> {
+ throw new AssertionError("cached metadata must not issue
an HMS request");
+ });
+ Assertions.assertNull(loader.isBlocked());
+ Assertions.assertEquals(16, loader.take(16).size());
+ Assertions.assertEquals(List.of(16), probes,
+ "claiming the first slice must not prepare the following 984
cached partitions inline");
+ Assertions.assertEquals(1, scheduler.tasks.size());
+ Runnable queued = scheduler.tasks.element();
+ Assertions.assertNull(loader.isBlocked(), "an idle reader can help the
cache probe");
+ queued.run();
+ Assertions.assertTrue(scheduler.tasks.isEmpty(), "the helped task must
be removed from the queue");
+ Assertions.assertEquals(List.of(16, 984), probes);
+ int count = 16;
+ while (!loader.isFinished()) {
+ count += loader.take(16).size();
+ }
+ Assertions.assertEquals(1000, count);
+ }
+
+ @Test
+ void helperAndMetadataWorkerCannotRunTheSameCacheProbeTwice() throws
Exception {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicInteger probes = new AtomicInteger();
+ HivePartitionBatchLoader loader = new
HivePartitionBatchLoader(scheduler, new ConnectorScanContext(),
+ names(1000), 1024, 3, batch -> {
+ probes.incrementAndGet();
+ if (batch.size() > 16) {
+ entered.countDown();
+ try {
+ Assertions.assertTrue(release.await(5,
TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(e);
+ }
+ }
+ return Optional.of(partitions(batch));
+ }, batch -> {
+ throw new AssertionError("cache hit must not issue an HMS
request");
+ });
+ loader.isBlocked();
+ loader.take(16);
+ ExecutorService reader = Executors.newSingleThreadExecutor();
+ try {
+ Future<CompletableFuture<Void>> helper =
reader.submit(loader::isBlocked);
+ Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS));
+ scheduler.runNext();
+ Assertions.assertFalse(helper.isDone(), "the metadata worker must
not duplicate the helper's probe");
+ release.countDown();
+ Assertions.assertNull(helper.get(5, TimeUnit.SECONDS));
+ Assertions.assertEquals(2, probes.get());
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ } finally {
+ release.countDown();
+ reader.shutdownNow();
+ }
+ }
+
+ @Test
+ void helpedCacheMissHandsOneQueueSlotToRemoteLoadWithoutWaiting() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ scheduler.capacity = 1;
+ AtomicInteger requests = new AtomicInteger();
+ HivePartitionBatchLoader loader = new
HivePartitionBatchLoader(scheduler, new ConnectorScanContext(),
+ names(1000), 1024, 3,
+ batch -> batch.size() == 16 ? Optional.of(partitions(batch)) :
Optional.empty(), batch -> {
+ requests.incrementAndGet();
+ return partitions(batch);
+ });
+ loader.isBlocked();
+ Assertions.assertEquals(16, loader.take(16).size());
+ CompletableFuture<Void> pending = loader.isBlocked();
+ Assertions.assertFalse(pending.isDone());
+ Assertions.assertEquals(0, requests.get(), "helping a cache miss must
not perform the RPC");
+ Assertions.assertEquals(1, scheduler.tasks.size());
+ Assertions.assertNull(scheduler.onSpace, "removing the helped cache
task must admit the remote task");
+ scheduler.runNext();
+ Assertions.assertTrue(pending.isDone());
+ Assertions.assertEquals(1, requests.get());
+ }
+
+ @Test
+ void failureInLargeCachedProbeWakesReaders() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ IllegalStateException failure = new IllegalStateException("cache probe
failed");
+ HivePartitionBatchLoader loader = new
HivePartitionBatchLoader(scheduler, new ConnectorScanContext(),
+ names(1000), 1024, 3, batch -> {
+ if (batch.size() > 16) {
+ throw failure;
+ }
+ return Optional.of(partitions(batch));
+ }, batch -> {
+ throw new AssertionError("failed cache probe must not
issue an HMS request");
+ });
+ Assertions.assertNull(loader.isBlocked());
+ Assertions.assertEquals(16, loader.take(16).size());
+ scheduler.runNext();
+ CompletableFuture<Void> pending = loader.isBlocked();
+ Assertions.assertSame(failure,
Assertions.assertThrows(CompletionException.class, pending::join).getCause());
+ Assertions.assertSame(failure,
Assertions.assertThrows(CompletionException.class,
+ () -> loader.take(16)).getCause());
+ }
+
+ @Test
+ void cancellingQueuedLargeCachedProbePreventsItsExecution() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ ConnectorScanContext context = new ConnectorScanContext();
+ List<Integer> probes = new ArrayList<>();
+ HivePartitionBatchLoader loader = new
HivePartitionBatchLoader(scheduler, context,
+ names(1000), 1024, 3, batch -> {
+ probes.add(batch.size());
+ return Optional.of(partitions(batch));
+ }, batch -> {
+ throw new AssertionError("stopped cached scan must not
issue an HMS request");
+ });
+ loader.isBlocked();
+ loader.take(16);
+ Runnable stale = scheduler.tasks.element();
+ context.cancel();
+ Assertions.assertTrue(loader.isBlocked().isCompletedExceptionally());
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ stale.run();
+ Assertions.assertEquals(List.of(16), probes);
+ }
+
+ @Test
+ void cancellationDuringLargeCacheMissDoesNotStartAnHmsRequest() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ ConnectorScanContext context = new ConnectorScanContext();
+ AtomicInteger requests = new AtomicInteger();
+ HivePartitionBatchLoader loader = new
HivePartitionBatchLoader(scheduler, context,
+ names(1000), 1024, 3, batch -> {
+ if (batch.size() > 16) {
+ context.cancel();
+ return Optional.empty();
+ }
+ return Optional.of(partitions(batch));
+ }, batch -> {
+ requests.incrementAndGet();
+ return partitions(batch);
+ });
+ loader.isBlocked();
+ loader.take(16);
+ scheduler.runNext();
+ Assertions.assertTrue(loader.isBlocked().isCompletedExceptionally());
+ Assertions.assertEquals(0, requests.get());
+ }
+
+ @Test
+ void maximumRpcLimitDoesNotOverflowAfterTheFirstBatch() {
+ List<Integer> requests = new ArrayList<>();
+ HivePartitionBatchLoader loader =
loader(ConnectorTaskScheduler.DIRECT, new ConnectorScanContext(),
+ 17, Integer.MAX_VALUE, batch -> {
+ requests.add(batch.size());
+ return partitions(batch);
+ });
+ Assertions.assertNull(loader.isBlocked());
+ List<HmsPartitionInfo> all = new ArrayList<>(loader.take(16));
+ all.addAll(loader.take(16));
+ Assertions.assertTrue(loader.isFinished());
+ Assertions.assertEquals(17, all.size());
+ Assertions.assertEquals(List.of(16, 1), requests);
+ }
+
+ @Test
+ void configuredRpcLimitAlsoBoundsFirstBatch() {
+ List<Integer> requests = new ArrayList<>();
+ HivePartitionBatchLoader loader =
loader(ConnectorTaskScheduler.DIRECT, new ConnectorScanContext(),
+ 10, 4, batch -> {
+ requests.add(batch.size());
+ return partitions(batch);
+ });
+ int count = 0;
+ while (!loader.isFinished()) {
+ Assertions.assertNull(loader.isBlocked());
+ count += loader.take(16).size();
+ }
+ Assertions.assertEquals(10, count);
+ Assertions.assertEquals(List.of(4, 4, 2), requests);
+ }
+
+ @Test
+ void oneReadyListBoundsMetadataAheadOfDirectoryConsumers() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ HivePartitionBatchLoader loader = loader(scheduler, new
ConnectorScanContext(), 200, 32,
+ HivePartitionBatchLoaderTest::partitions);
+ loader.isBlocked();
+ scheduler.runNext();
+ loader.take(16);
+ scheduler.runAll();
+ Assertions.assertNull(loader.isBlocked());
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ loader.take(16);
+ Assertions.assertTrue(scheduler.tasks.isEmpty(), "half-consumed batch
must stop the parent");
+ loader.take(16);
+ Assertions.assertEquals(1, scheduler.tasks.size());
+ }
+
+ @Test
+ void fullyCachedMetadataDoesNotNeedMetadataCapacity() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ scheduler.full = true;
+ HivePartitionBatchLoader loader = new
HivePartitionBatchLoader(scheduler, new ConnectorScanContext(),
+ names(100), 32, 3, batch -> Optional.of(partitions(batch)),
batch -> {
+ throw new AssertionError("cache hit must not load
remotely");
+ });
+ int count = 0;
+ while (!loader.isFinished()) {
+ Assertions.assertNull(loader.isBlocked());
+ count += loader.take(16).size();
+ }
+ Assertions.assertEquals(100, count);
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ Assertions.assertNull(scheduler.onSpace);
+ }
+
+ @Test
+ void fullSchedulerReturnsImmediatelyAndCancellationRemovesSpaceWaiter() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ scheduler.full = true;
+ ConnectorScanContext context = new ConnectorScanContext();
+ AtomicInteger requests = new AtomicInteger();
+ HivePartitionBatchLoader loader = loader(scheduler, context, 100, 32,
batch -> {
+ requests.incrementAndGet();
+ return partitions(batch);
+ });
+ CompletableFuture<Void> pending = loader.isBlocked();
+ Assertions.assertFalse(pending.isDone());
+ Runnable stale = scheduler.onSpace;
+ context.cancel();
+ Assertions.assertTrue(pending.isCompletedExceptionally());
+ Assertions.assertNull(scheduler.onSpace);
+ scheduler.full = false;
+ stale.run();
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ Assertions.assertEquals(0, requests.get());
+ }
+
+ @Test
+ void admissionResumesAfterSpaceWithoutDuplicateRpc() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ scheduler.full = true;
+ AtomicInteger requests = new AtomicInteger();
+ HivePartitionBatchLoader loader = loader(scheduler, new
ConnectorScanContext(), 10, 32, batch -> {
+ requests.incrementAndGet();
+ return partitions(batch);
+ });
+ CompletableFuture<Void> pending = loader.isBlocked();
+ scheduler.full = false;
+ scheduler.onSpace.run();
+ scheduler.runNext();
+ Assertions.assertTrue(pending.isDone());
+ Assertions.assertEquals(1, requests.get());
+ }
+
+ @Test
+ void spaceCallbackBeforeAdmissionReturnsCannotEnqueueTwice() {
+ Queue<Runnable> tasks = new ArrayDeque<>();
+ AtomicInteger admissions = new AtomicInteger();
+ List<Runnable> callbacks = new ArrayList<>();
+ ConnectorTaskScheduler scheduler = new ConnectorTaskScheduler() {
+ @Override
+ public boolean tryExecute(Runnable task, Consumer<Throwable>
onFailure, Runnable onSpace) {
+ callbacks.add(onSpace);
+ if (admissions.incrementAndGet() == 1) {
+ onSpace.run();
+ return false;
+ }
+ tasks.add(task);
+ return true;
+ }
+
+ @Override
+ public void removeSpaceListener(Runnable callback) {
+ }
+
+ @Override
+ public void execute(Runnable task) {
+ throw new UnsupportedOperationException();
+ }
+ };
+ HivePartitionBatchLoader loader = loader(scheduler, new
ConnectorScanContext(), 10, 32,
+ HivePartitionBatchLoaderTest::partitions);
+ CompletableFuture<Void> pending = loader.isBlocked();
+ callbacks.get(0).run();
+ Assertions.assertEquals(2, admissions.get());
+ Assertions.assertEquals(1, tasks.size());
+ tasks.remove().run();
+ Assertions.assertTrue(pending.isDone());
+ Assertions.assertEquals(10, loader.take(16).size());
+ }
+
+ @Test
+ void closingLastReaderCancelsQueuedParentAndWakesWaiters() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ AtomicInteger requests = new AtomicInteger();
+ HivePartitionBatchLoader loader = loader(scheduler, new
ConnectorScanContext(), 100, 32, batch -> {
+ requests.incrementAndGet();
+ return partitions(batch);
+ });
+ CompletableFuture<Void> pending = loader.isBlocked();
+ Runnable stale = scheduler.tasks.element();
+ loader.readerClosed();
+ loader.readerClosed();
+ Assertions.assertFalse(pending.isDone());
+ loader.readerClosed();
+ Assertions.assertTrue(pending.isCompletedExceptionally());
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ stale.run();
+ Assertions.assertEquals(0, requests.get());
+ }
+
+ @Test
+ void failedAdmissionWakesReadersAndDoesNotRunAStaleTask() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ AtomicInteger requests = new AtomicInteger();
+ HivePartitionBatchLoader loader = loader(scheduler, new
ConnectorScanContext(), 100, 32, batch -> {
+ requests.incrementAndGet();
+ return partitions(batch);
+ });
+ CompletableFuture<Void> pending = loader.isBlocked();
+ Runnable stale = scheduler.tasks.element();
+ IllegalStateException failure = new IllegalStateException("metadata
scheduler closed");
+ scheduler.onFailure.accept(failure);
+ Assertions.assertSame(failure,
Assertions.assertThrows(CompletionException.class, pending::join).getCause());
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ stale.run();
+ Assertions.assertEquals(0, requests.get());
+ }
+
+ @Test
+ void loadingFailureIsVisibleToEveryReader() {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ IllegalStateException failure = new IllegalStateException("HMS
failed");
+ HivePartitionBatchLoader loader = loader(scheduler, new
ConnectorScanContext(), 100, 32, batch -> {
+ throw failure;
+ });
+ CompletableFuture<Void> a = loader.isBlocked();
+ CompletableFuture<Void> b = loader.isBlocked();
+ scheduler.runNext();
+ Assertions.assertSame(failure,
Assertions.assertThrows(CompletionException.class, a::join).getCause());
+ Assertions.assertSame(failure,
Assertions.assertThrows(CompletionException.class, b::join).getCause());
+ Assertions.assertSame(failure,
Assertions.assertThrows(CompletionException.class,
+ () -> loader.take(16)).getCause());
+ }
+
+ @Test
+ void cancellationDiscardsUninterruptibleRpcResultAndStopsFurtherCalls()
throws Exception {
+ QueuedScheduler scheduler = new QueuedScheduler();
+ ConnectorScanContext context = new ConnectorScanContext();
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicInteger requests = new AtomicInteger();
+ HivePartitionBatchLoader loader = loader(scheduler, context, 100, 32,
batch -> {
+ requests.incrementAndGet();
+ entered.countDown();
+ try {
+ Assertions.assertTrue(release.await(5, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ throw new RuntimeException(e);
+ }
+ return partitions(batch);
+ });
+ CompletableFuture<Void> pending = loader.isBlocked();
+ ExecutorService worker = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> running = worker.submit(scheduler::runNext);
+ Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS));
+ context.cancel();
+ Assertions.assertTrue(pending.isCompletedExceptionally());
+ release.countDown();
+ running.get(5, TimeUnit.SECONDS);
+ Assertions.assertEquals(1, requests.get());
+ Assertions.assertTrue(scheduler.tasks.isEmpty());
+ } finally {
+ release.countDown();
+ worker.shutdownNow();
+ }
+ }
+
+ @Test
+ void droppedPartitionsDoNotRecurseOrPrematurelyFinish() {
+ HivePartitionBatchLoader loader =
loader(ConnectorTaskScheduler.DIRECT, new ConnectorScanContext(),
+ 20000, 1, batch -> Collections.emptyList());
+ Assertions.assertNull(loader.isBlocked());
+ Assertions.assertTrue(loader.isFinished());
+ }
+
+ private static HivePartitionBatchLoader loader(ConnectorTaskScheduler
scheduler, ConnectorScanContext context,
+ int count, int batchSize, Function<List<String>,
List<HmsPartitionInfo>> fetch) {
+ return new HivePartitionBatchLoader(scheduler, context, names(count),
batchSize, 3,
+ batch -> Optional.empty(), fetch);
+ }
+
+ private static List<String> names(int count) {
+ List<String> result = new ArrayList<>();
+ for (int i = 0; i < count; i++) {
+ result.add("p=" + i);
+ }
+ return result;
+ }
+
+ private static List<HmsPartitionInfo> partitions(List<String> names) {
+ List<HmsPartitionInfo> result = new ArrayList<>();
+ for (String name : names) {
+ result.add(new HmsPartitionInfo(List.of(name.substring(2)), name,
+ null, null, null, Collections.emptyMap()));
+ }
+ return result;
+ }
+
+ private static final class QueuedScheduler implements
ConnectorTaskScheduler {
+ private final Queue<Runnable> tasks = new ArrayDeque<>();
+ private boolean full;
+ private int capacity = Integer.MAX_VALUE;
+ private Runnable onSpace;
+ private Consumer<Throwable> onFailure;
+
+ @Override
+ public boolean tryExecute(Runnable task, Consumer<Throwable>
onFailure, Runnable callback) {
+ if (full || tasks.size() >= capacity) {
+ onSpace = callback;
+ return false;
+ }
+ this.onFailure = onFailure;
+ tasks.add(task);
+ return true;
+ }
+
+ @Override
+ public void removeSpaceListener(Runnable callback) {
+ if (onSpace == callback) {
+ onSpace = null;
+ }
+ }
+
+ @Override
+ public void remove(Runnable task) {
+ if (tasks.remove(task) && !full && tasks.size() < capacity &&
onSpace != null) {
+ Runnable callback = onSpace;
+ onSpace = null;
+ callback.run();
+ }
+ }
+
+ @Override
+ public void execute(Runnable task) {
+ throw new UnsupportedOperationException();
+ }
+
+ private void runAll() {
+ while (!tasks.isEmpty()) {
+ runNext();
+ }
+ }
+
+ private void runNext() {
+ tasks.remove().run();
+ }
+ }
+}
diff --git
a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveScanBatchModeTest.java
b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveScanBatchModeTest.java
index 61841f38e40..863671d19f2 100644
---
a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveScanBatchModeTest.java
+++
b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveScanBatchModeTest.java
@@ -44,6 +44,7 @@ import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import java.io.IOException;
+import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -53,8 +54,10 @@ import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
+import java.util.Queue;
import java.util.Set;
import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
@@ -99,6 +102,229 @@ public class HiveScanBatchModeTest {
private static final String LAZY_SIMPLE_SERDE =
"org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe";
private static final String OPEN_CSV_SERDE =
"org.apache.hadoop.hive.serde2.OpenCSVSerde";
+ @Test
+ public void completedMetadataDependencyMustNotLookLikeEndOfInput() throws
Exception {
+ Queue<Runnable> metadata = new ArrayDeque<>();
+ ConnectorTaskScheduler scheduler = new ConnectorTaskScheduler() {
+ @Override
+ public boolean tryExecute(Runnable task, Consumer<Throwable>
onFailure, Runnable onSpace) {
+ metadata.add(task);
+ return true;
+ }
+
+ @Override
+ public void removeSpaceListener(Runnable callback) {
+ }
+
+ @Override
+ public void execute(Runnable task) {
+ metadata.add(task);
+ }
+ };
+ FakeHmsClient hms = new FakeHmsClient();
+ CountingLister lister = new CountingLister();
+ List<String> names = new ArrayList<>();
+ for (int i = 0; i < 17; i++) {
+ names.add("year=2024/month=" + i);
+ }
+ ConnectorSplitSource source = provider(hms,
lister).streamPartitionSources(
+ new FakeSession(Collections.emptyMap(), scheduler),
+ ConnectorScanRequest.builder(parquetHandle(PART_KEYS),
Collections.emptyList())
+ .requiredPartitions(names).build(), 16,
1).orElseThrow().get(0);
+ try {
+ Assertions.assertFalse(source.isBlocked().isDone());
+ metadata.remove().run();
+ Set<String> paths = new HashSet<>();
+ for (int i = 0; i < 16; i++) {
+ Assertions.assertTrue(source.hasNext());
+
Assertions.assertTrue(paths.add(source.next().getPath().orElseThrow()));
+ }
+ CompletableFuture<Void> pending = source.isBlocked();
+ Assertions.assertNotNull(pending);
+ Assertions.assertFalse(pending.isDone());
+ Assertions.assertFalse(source.hasNext());
+
+ // Complete after the empty read, before the caller checks whether
it means EOF.
+ metadata.remove().run();
+ Assertions.assertTrue(pending.isDone());
+ Assertions.assertSame(pending, source.isBlocked(), "retain
readiness until the next read attempt");
+ Assertions.assertTrue(source.hasNext());
+
Assertions.assertTrue(paths.add(source.next().getPath().orElseThrow()));
+ Assertions.assertFalse(source.hasNext());
+ Assertions.assertNull(source.isBlocked());
+ Assertions.assertEquals(17, paths.size());
+ Assertions.assertEquals(17, lister.totalCalls);
+ } finally {
+ source.close();
+ }
+ }
+
+ @Test
+ public void
parallelSourcesKeepCapturedPartitionViewAfterSharedCacheChanges() throws
Exception {
+ FakeHmsClient hms = new FakeHmsClient() {
+ @Override
+ public Optional<HmsPartitionBatchResult>
getCachedPartitionsWithStats(
+ String db, String table, List<String> names) {
+ throw new AssertionError("captured partition view must take
precedence over newer cache data");
+ }
+ };
+ List<String> names = Arrays.asList("year=2024/month=01",
"year=2024/month=02");
+ List<HmsPartitionInfo> captured = Arrays.asList(
+ new HmsPartitionInfo(Arrays.asList("2024", "01"), names.get(0),
+ null, null, null, Collections.emptyMap()),
+ new HmsPartitionInfo(Arrays.asList("2024", "02"), names.get(1),
+ null, null, null, Collections.emptyMap()));
+ HiveTableHandle handle =
parquetHandle(PART_KEYS).toBuilder().prunedPartitions(captured)
+ .prunedPartitionsByName(Map.of(names.get(0), captured.get(0),
names.get(1), captured.get(1))).build();
+ CountingLister lister = new CountingLister();
+ HiveScanPlanProvider provider = provider(hms, lister);
+ List<ConnectorSplitSource> sources =
provider.streamPartitionSources(new FakeSession(),
+ ConnectorScanRequest.builder(handle,
Collections.emptyList()).requiredPartitions(names).build(),
+ 1, 2).get();
+ int count = 0;
+ for (ConnectorSplitSource source : sources) {
+ try {
+ while (source.hasNext()) {
+ source.next();
+ count++;
+ }
+ } finally {
+ source.close();
+ }
+ }
+ Assertions.assertEquals(2, count);
+ Assertions.assertEquals(new HashSet<>(names),
lister.callsPerLocation.keySet());
+ Assertions.assertTrue(hms.requestedPartitionNames.isEmpty());
+ }
+
+ @Test
+ public void metadataProducerContinuesWhileThreeDirectoryReadersRun()
throws Exception {
+ ExecutorService metadata = Executors.newSingleThreadExecutor();
+ ExecutorService directories = Executors.newFixedThreadPool(3);
+ CountDownLatch secondBatch = new CountDownLatch(1);
+ AtomicInteger rpc = new AtomicInteger();
+ FakeHmsClient hms = new FakeHmsClient() {
+ @Override
+ public List<HmsPartitionInfo> getPartitions(String db, String
table, List<String> names) {
+ if (rpc.incrementAndGet() == 2) {
+ secondBatch.countDown();
+ }
+ return super.getPartitions(db, table, names);
+ }
+ };
+ ConnectorTaskScheduler scheduler = new ConnectorTaskScheduler() {
+ @Override
+ public boolean tryExecute(Runnable task, Consumer<Throwable>
onFailure, Runnable onSpace) {
+ metadata.execute(task);
+ return true;
+ }
+
+ @Override
+ public void removeSpaceListener(Runnable callback) {
+ }
+
+ @Override
+ public void execute(Runnable task) {
+ metadata.execute(task);
+ }
+ };
+ Set<String> paths = java.util.concurrent.ConcurrentHashMap.newKeySet();
+ HiveScanPlanProvider provider = provider(hms, (location, fs) -> {
+ try {
+ Assertions.assertTrue(secondBatch.await(5, TimeUnit.SECONDS),
+ "parent must fetch the next batch while the first
child is listing");
+ } catch (InterruptedException e) {
+ throw new RuntimeException(e);
+ }
+ return Collections.singletonList(new HiveFileStatus(location +
"/data", 10, 1));
+ });
+ List<String> names = new ArrayList<>();
+ for (int i = 0; i < 1000; i++) {
+ names.add("year=2024/month=" + i);
+ }
+ List<ConnectorSplitSource> sources = provider.streamPartitionSources(
+ new FakeSession(Collections.emptyMap(), scheduler),
+ ConnectorScanRequest.builder(parquetHandle(PART_KEYS),
Collections.emptyList())
+ .requiredPartitions(names).build(), 16, 3).get();
+ try {
+ List<Future<?>> readers = new ArrayList<>();
+ for (ConnectorSplitSource source : sources) {
+ readers.add(directories.submit(() -> {
+ try {
+ while (true) {
+ java.util.concurrent.CompletableFuture<Void>
blocked = source.isBlocked();
+ if (blocked != null) {
+ blocked.get(5, TimeUnit.SECONDS);
+ }
+ if (!source.hasNext()) {
+ if (source.isBlocked() != null) {
+ continue;
+ }
+ break;
+ }
+
Assertions.assertTrue(paths.add(source.next().getPath().get()));
+ }
+ source.close();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }));
+ }
+ for (Future<?> reader : readers) {
+ reader.get(10, TimeUnit.SECONDS);
+ }
+ Assertions.assertEquals(1000, paths.size());
+ Assertions.assertEquals(2, rpc.get());
+ } finally {
+ for (ConnectorSplitSource source : sources) {
+ source.close();
+ }
+ directories.shutdownNow();
+ metadata.shutdownNow();
+ }
+ }
+
+ @Test
+ public void
parallelDirectorySourcesShareLargeHmsBatchesAndPreserveAllSplits() throws
Exception {
+ FakeHmsClient hms = new FakeHmsClient();
+ CountingLister lister = new CountingLister();
+ HiveScanPlanProvider provider = provider(hms, lister);
+ List<String> names = new ArrayList<>();
+ for (int i = 0; i < 1000; i++) {
+ names.add("year=2024/month=" + i);
+ }
+ List<ConnectorSplitSource> sources =
provider.streamPartitionSources(new FakeSession(),
+ ConnectorScanRequest.builder(parquetHandle(PART_KEYS),
Collections.emptyList())
+ .requiredPartitions(names).build(), 16, 3).get();
+ Assertions.assertTrue(hms.requestedPartitionNames.isEmpty(), "planning
must not perform HMS I/O");
+ Set<String> paths = new HashSet<>();
+ try {
+ boolean more;
+ do {
+ more = false;
+ for (ConnectorSplitSource source : sources) {
+ java.util.concurrent.CompletableFuture<Void> blocked =
source.isBlocked();
+ if (blocked != null) {
+ blocked.get(5, TimeUnit.SECONDS);
+ }
+ if (source.hasNext()) {
+ more = true;
+
Assertions.assertTrue(paths.add(source.next().getPath().get()));
+ }
+ }
+ } while (more);
+ } finally {
+ for (ConnectorSplitSource source : sources) {
+ source.close();
+ }
+ }
+ Assertions.assertEquals(1000, paths.size());
+ Assertions.assertEquals(1000, lister.totalCalls);
+ Assertions.assertEquals(2, hms.requestedPartitionNames.size());
+ Assertions.assertEquals(16, hms.requestedPartitionNames.get(0).size());
+ Assertions.assertEquals(984,
hms.requestedPartitionNames.get(1).size());
+ }
+
@Test
public void millionFilesOnlyCreateTheConsumedSplitPrefix() throws
Exception {
AtomicInteger normalizedFiles = new AtomicInteger();
@@ -291,6 +517,49 @@ public class HiveScanBatchModeTest {
Assertions.assertEquals(Collections.singleton("year=2024/month=02"),
lister.callsPerLocation.keySet());
}
+ @Test
+ public void streamingCachedPartitionsDoNotSubmitMetadataTasks() throws
Exception {
+ ConnectorTaskScheduler scheduler = new ConnectorTaskScheduler() {
+ @Override
+ public boolean tryExecute(Runnable task, Consumer<Throwable>
onFailure, Runnable onSpace) {
+ Assertions.fail("A fully cached batch must not enter the
metadata executor");
+ return false;
+ }
+
+ @Override
+ public void removeSpaceListener(Runnable onSpace) {
+ }
+
+ @Override
+ public void execute(Runnable task) {
+ Assertions.fail("A fully cached batch must not enter the
metadata executor");
+ }
+ };
+ FakeHmsClient hms = new FakeHmsClient() {
+ @Override
+ public Optional<HmsPartitionBatchResult>
getCachedPartitionsWithStats(
+ String dbName, String tableName, List<String> names) {
+ return Optional.of(new
HmsPartitionBatchResult(Collections.singletonList(part("cached")),
+
HmsPartitionBatchStats.builder().requestedItems(1).build()));
+ }
+
+ @Override
+ public List<HmsPartitionInfo> getPartitions(String dbName, String
tableName, List<String> names) {
+ Assertions.fail("A fully cached batch must not fetch partition
objects");
+ return Collections.emptyList();
+ }
+ };
+ HiveScanPlanProvider provider = provider(hms, new CountingLister());
+ ConnectorScanRequest request =
ConnectorScanRequest.builder(parquetHandle(PART_KEYS),
+
Collections.emptyList()).requiredPartitions(Collections.singletonList("year=2024/month=01")).build();
+ try (ConnectorSplitSource source = provider.streamSplits(new
FakeSession(Collections.emptyMap(), scheduler),
+ request)) {
+ Assertions.assertTrue(source.hasNext());
+ Assertions.assertEquals("cached/000000_0",
source.next().getPath().get());
+ Assertions.assertFalse(source.hasNext());
+ }
+ }
+
@Test
public void streamingReloadsWholeBatchWhenCapturedMapIsIncomplete() throws
Exception {
List<List<String>> requested = new ArrayList<>();
diff --git
a/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/CachingHmsClient.java
b/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/CachingHmsClient.java
index 9ba86488eb2..19bc8f2b258 100644
---
a/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/CachingHmsClient.java
+++
b/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/CachingHmsClient.java
@@ -34,6 +34,7 @@ import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
+import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
@@ -271,6 +272,24 @@ public class CachingHmsClient implements HmsClient {
return getPartitionsWithStats(dbName, tableName, partNames, true);
}
+ @Override
+ public Optional<HmsPartitionBatchResult> getCachedPartitionsWithStats(
+ String dbName, String tableName, List<String> partNames) {
+ long start = System.nanoTime();
+ HmsPartitionRequest request = new HmsPartitionRequest(dbName,
tableName, partNames);
+ List<HmsPartitionInfo> cached = new ArrayList<>(partNames.size());
+ for (HmsPartitionIdentity.ParsedPartitionName partition :
request.getPartitions()) {
+ HmsPartitionInfo value = partitionsCache.getIfPresent(
+ new PartitionKey(dbName, tableName,
partition.getValues()));
+ if (value == null) {
+ return Optional.empty();
+ }
+ cached.add(value);
+ }
+ return Optional.of(new HmsPartitionBatchResult(cached,
HmsPartitionBatchStats.builder()
+
.requestedItems(partNames.size()).logicalElapsedNanos(System.nanoTime() -
start).build()));
+ }
+
@Override
public HmsPartitionBatchResult getExistingPartitionsWithStats(String
dbName, String tableName,
List<String> partNames, ConnectorScanContext scanContext) {
diff --git
a/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/HmsClient.java
b/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/HmsClient.java
index 7bd76166771..493cf5ab67e 100644
---
a/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/HmsClient.java
+++
b/fe/fe-connector/fe-connector-hms/src/main/java/org/apache/doris/connector/hms/HmsClient.java
@@ -23,6 +23,7 @@ import java.io.Closeable;
import java.util.Collections;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.function.Function;
/**
@@ -217,6 +218,16 @@ public interface HmsClient extends Closeable {
return new HmsPartitionBatchResult(partitions, stats);
}
+ /**
+ * Returns all requested objects, in request order, only if they are
already cached. Never performs remote
+ * I/O or waits for an in-flight load. An empty Optional means at least
one object is unavailable locally,
+ * including for clients without a cache. Used to keep warm lookups off
the metadata executor.
+ */
+ default Optional<HmsPartitionBatchResult> getCachedPartitionsWithStats(
+ String dbName, String tableName, List<String> partNames) {
+ return Optional.empty();
+ }
+
/**
* Returns the requested partitions that still exist together with
transport-invocation batching statistics.
* Implementations that do not expose transport details retain the
omission-tolerant result and report a
diff --git
a/fe/fe-connector/fe-connector-hms/src/test/java/org/apache/doris/connector/hms/CachingHmsClientTest.java
b/fe/fe-connector/fe-connector-hms/src/test/java/org/apache/doris/connector/hms/CachingHmsClientTest.java
index 1816b19ca0f..1ead83538a0 100644
---
a/fe/fe-connector/fe-connector-hms/src/test/java/org/apache/doris/connector/hms/CachingHmsClientTest.java
+++
b/fe/fe-connector/fe-connector-hms/src/test/java/org/apache/doris/connector/hms/CachingHmsClientTest.java
@@ -342,6 +342,72 @@ public class CachingHmsClientTest {
// ---- getPartitions ----
+ @Test
+ public void cachedPartitionProbeDoesNotWaitForAnInFlightLoad() throws
Exception {
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ RecordingHmsClient delegate = new RecordingHmsClient() {
+ @Override
+ public List<HmsPartitionInfo> getPartitions(String db, String
table, List<String> names) {
+ entered.countDown();
+ awaitLatch(release);
+ return super.getPartitions(db, table, names);
+ }
+ };
+ CachingHmsClient cache = cache(delegate, Collections.emptyMap());
+ ExecutorService workers = Executors.newFixedThreadPool(2);
+ try {
+ List<String> names = Collections.singletonList("p=1");
+ Future<List<HmsPartitionInfo>> loading = workers.submit(() ->
cache.getPartitions("db", "t", names));
+ Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS));
+ Future<Boolean> probe = workers.submit(() ->
cache.getCachedPartitionsWithStats("db", "t", names)
+ .isPresent());
+ Assertions.assertFalse(probe.get(5, TimeUnit.SECONDS));
+ Assertions.assertFalse(loading.isDone());
+ release.countDown();
+ HmsPartitionInfo loaded = loading.get(5, TimeUnit.SECONDS).get(0);
+ Assertions.assertSame(loaded,
cache.getCachedPartitionsWithStats("db", "t", names)
+ .orElseThrow().getPartitions().get(0));
+ Assertions.assertEquals(1, delegate.getPartitionsCalls);
+ } finally {
+ release.countDown();
+ workers.shutdownNow();
+ cache.close();
+ }
+ }
+
+ @Test
+ public void cachedPartitionProbeNeverLoadsAndRespectsInvalidation() {
+ RecordingHmsClient delegate = new RecordingHmsClient();
+ CatalogMetaCache owner = CatalogMetaCache.unmanaged();
+ CachingHmsClient cache = new CachingHmsClient(owner, delegate,
Collections.emptyMap());
+ Assertions.assertFalse(cache.getCachedPartitionsWithStats("db", "t",
Arrays.asList("p=1", "p=2"))
+ .isPresent());
+ Assertions.assertEquals(0, delegate.getPartitionsCalls);
+ List<HmsPartitionInfo> loaded = cache.getPartitions("db", "t",
Arrays.asList("p=1", "p=2"));
+ HmsPartitionBatchResult hit = cache.getCachedPartitionsWithStats("db",
"t", Arrays.asList("p=2", "p=1"))
+ .orElseThrow(AssertionError::new);
+ Assertions.assertSame(loaded.get(1), hit.getPartitions().get(0));
+ Assertions.assertSame(loaded.get(0), hit.getPartitions().get(1));
+ Assertions.assertEquals(2, hit.getStats().getRequestedItems());
+ Assertions.assertEquals(0, hit.getStats().getTransportInvocations());
+ Assertions.assertFalse(cache.getCachedPartitionsWithStats("db", "t",
Arrays.asList("p=1", "p=3"))
+ .isPresent());
+ Assertions.assertEquals(1, delegate.getPartitionsCalls);
+ owner.invalidateTable("db", "t");
+ Assertions.assertFalse(cache.getCachedPartitionsWithStats("db", "t",
Collections.singletonList("p=1"))
+ .isPresent());
+ Assertions.assertEquals(1, delegate.getPartitionsCalls);
+ }
+
+ @Test
+ public void nonCachingClientProbeDoesNotFetchPartitions() {
+ RecordingHmsClient raw = new RecordingHmsClient();
+ Assertions.assertFalse(raw.getCachedPartitionsWithStats("db", "t",
Collections.singletonList("p=1"))
+ .isPresent());
+ Assertions.assertEquals(0, raw.getPartitionsCalls);
+ }
+
@Test
public void getPartitionsSharesPerPartitionEntriesAcrossRequests() {
RecordingHmsClient delegate = new RecordingHmsClient();
diff --git
a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProvider.java
b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProvider.java
index a4965b3db2c..e64d7089c76 100644
---
a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProvider.java
+++
b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProvider.java
@@ -232,6 +232,19 @@ public interface ConnectorScanPlanProvider {
return 1024;
}
+ /**
+ * Optionally opens cooperating sources for the complete pruned partition
set in {@code request}.
+ * A connector can batch metadata independently of directory work and
share the results between
+ * sources. No remote I/O may run here. Each source has one serial reader,
may report an asynchronous
+ * dependency through {@link ConnectorSplitSource#isBlocked()}, and must
release its shared resources
+ * when all sources close or the scan is cancelled. The engine owns
admission and backpressure.
+ * Empty means to use the engine's ordinary partition-batch sources.
+ */
+ default Optional<List<ConnectorSplitSource>>
streamPartitionSources(ConnectorSession session,
+ ConnectorScanRequest request, int partitionsPerTask, int
maxSources) {
+ return Optional.empty();
+ }
+
/**
* Whether this connector's scan ranges carry meaningful byte lengths
* ({@link ConnectorScanRange#getLength()}) so the engine can apply {@code
TABLESAMPLE} by
diff --git
a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorSplitSource.java
b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorSplitSource.java
index 0cb35f10a9d..4ee96a848f7 100644
---
a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorSplitSource.java
+++
b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/scan/ConnectorSplitSource.java
@@ -18,6 +18,7 @@
package org.apache.doris.connector.spi.scan;
import java.io.Closeable;
+import java.util.concurrent.CompletableFuture;
/**
* A lazy, closeable source of {@link ConnectorScanRange}s for streaming
(batched) split generation,
@@ -37,6 +38,19 @@ import java.io.Closeable;
*/
public interface ConnectorSplitSource extends Closeable {
+ /**
+ * An asynchronous metadata dependency for cooperating sources opened by
+ * {@link ConnectorScanPlanProvider#streamPartitionSources}, or null when
immediately runnable
+ * (including end of input). Ordinary {@code streamSplits} sources remain
synchronous.
+ * Readers must observe exceptional completion and retry after normal
completion. While blocked,
+ * {@code hasNext()} may return false without indicating end of input; in
that case this method must
+ * retain the dependency, even after completion, until the next {@code
hasNext()} attempt. Completion
+ * callbacks must only schedule work; they must not advance or close this
source. The default incurs no allocation.
+ */
+ default CompletableFuture<Void> isBlocked() {
+ return null;
+ }
+
/**
* Nonblocking, idempotent hint that a backend has requested a batch
beyond its first one.
* May run concurrently on an engine RPC thread. A connector may use this
to expand speculative
diff --git
a/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java
b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java
index a37bb7aa064..73b1fd28d14 100644
---
a/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java
+++
b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java
@@ -89,7 +89,7 @@ public class ConnectorPluginSurfaceTest {
version.load(in);
}
// Major 29 adds typed table-absence errors after the upstream API 28
changes.
- Assertions.assertEquals("29.0", version.getProperty("api.version"));
+ Assertions.assertEquals("30.0", version.getProperty("api.version"));
}
/** Root entry points plus provider/handle types returned to connector
plugins. */
diff --git
a/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt
b/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt
index 440e7ddafdf..fdd2b94e68b 100644
---
a/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt
+++
b/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt
@@ -219,6 +219,7 @@
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#planScanForPartiti
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#populateScanLevelParams(org.apache.doris.thrift.TFileScanRangeParams,java.util.Map):void
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#releaseReadTransaction(java.lang.String):void
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#scannedPartitionCount(java.util.List):java.util.OptionalLong
+org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#streamPartitionSources(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.scan.ConnectorScanRequest,int,int):java.util.Optional
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#streamSplits(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.scan.ConnectorScanRequest):org.apache.doris.connector.spi.scan.ConnectorSplitSource
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#streamingSplitEstimate(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle,java.util.Optional,boolean):long
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#supportsBatchScan(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle):boolean
@@ -272,6 +273,7 @@
org.apache.doris.connector.spi.scan.ConnectorScanRequest$Builder#requiredPartiti
org.apache.doris.connector.spi.scan.ConnectorScanRequest$Builder#scanContext(org.apache.doris.connector.spi.scan.ConnectorScanContext):org.apache.doris.connector.spi.scan.ConnectorScanRequest$Builder
org.apache.doris.connector.spi.scan.ConnectorSplitSource#close():void
org.apache.doris.connector.spi.scan.ConnectorSplitSource#hasNext():boolean
+org.apache.doris.connector.spi.scan.ConnectorSplitSource#isBlocked():java.util.concurrent.CompletableFuture
org.apache.doris.connector.spi.scan.ConnectorSplitSource#next():org.apache.doris.connector.spi.scan.ConnectorScanRange
org.apache.doris.connector.spi.scan.ConnectorSplitSource#onMoreSplitsRequested():void
org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:FILE_FORMAT_TYPE:java.lang.String=file_format_type
diff --git a/fe/fe-connector/pom.xml b/fe/fe-connector/pom.xml
index a5555398e7e..636bc50720a 100644
--- a/fe/fe-connector/pom.xml
+++ b/fe/fe-connector/pom.xml
@@ -55,7 +55,7 @@ under the License.
of the latter two means bumping this property as well (and
fe-extension-spi means bumping
all five families).
-->
- <connector.plugin.api.version>29.0</connector.plugin.api.version>
+ <connector.plugin.api.version>30.0</connector.plugin.api.version>
</properties>
<modules>
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java
index 8f9a6161815..84fa22d37bd 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java
@@ -267,6 +267,25 @@ public class FederationBackendPolicy {
ListMultimap<Backend, Split> assignment = ArrayListMultimap.create();
Collections.shuffle(splits, new Random(FIXED_SHUFFLE_SEED));
+ if (backends.size() == 1 &&
splits.stream().allMatch(Split::isRemotelyAccessible)) {
+ // Every remotely readable split has the same owner. Avoid
per-split hashing,
+ // candidate allocation and weight-map updates under the batch
assignment lock.
+ // Local-only splits still need the locality checks in the general
path below.
+ Backend backend = backends.get(0);
+ boolean alternativesUsed = !isSpreadEnabled();
+ long weight = 0;
+ for (Split split : splits) {
+ weight = Math.addExact(weight,
split.getSplitWeight().getRawValue());
+ if (alternativesUsed) {
+ split.setAlternativeHosts(Collections.emptyList());
+ }
+ }
+ assignment.putAll(backend, splits);
+ if (!splits.isEmpty()) {
+ addAssignedWeight(backend, weight);
+ }
+ return assignment;
+ }
List<Split> remainingSplits;
ResettableRandomizedIterator<Backend> randomCandidates = new
ResettableRandomizedIterator<>(backends);
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/PluginDrivenScanNode.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/PluginDrivenScanNode.java
index 4900b57a8b5..2ee7c5a4e8e 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/PluginDrivenScanNode.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/PluginDrivenScanNode.java
@@ -2377,6 +2377,7 @@ public class PluginDrivenScanNode extends
FileQueryScanNode {
});
// Admission does no connector I/O and remains owned by the planning
thread.
+ int startedSources = 0;
try {
splitAssignment.startEnumeration();
Executor scheduleExecutor =
connectorSession.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT);
@@ -2384,6 +2385,7 @@ public class PluginDrivenScanNode extends
FileQueryScanNode {
if (batchException.get() != null || splitAssignment.isStop()) {
break;
}
+ startedSources++;
profileFinalizer.taskSubmitted();
new SplitProducer(splitAssignment, scheduleExecutor, source,
failure -> {
try {
@@ -2399,6 +2401,13 @@ public class PluginDrivenScanNode extends
FileQueryScanNode {
} catch (Exception e) {
publishBatchFailure(batchException, splitAssignment, new
UserException(e.getMessage(), e));
} finally {
+ for (int i = startedSources; i < sources.size(); i++) {
+ try {
+ sources.get(i).close();
+ } catch (Exception e) {
+ publishBatchFailure(batchException, splitAssignment, new
UserException(e.getMessage(), e));
+ }
+ }
profileFinalizer.closeDispatch();
}
}
@@ -2412,6 +2421,14 @@ public class PluginDrivenScanNode extends
FileQueryScanNode {
if (producerLimit > 0) {
sourceCount = Math.min(sourceCount, producerLimit);
}
+ int maxSources = sourceCount;
+ Optional<List<ConnectorSplitSource>> cooperating =
onPluginClassLoader(scanProvider,
+ () -> scanProvider.streamPartitionSources(connectorSession,
+ request.withRequiredPartitions(partitions),
partitionsPerBatch, maxSources));
+ if (cooperating.isPresent()) {
+ return cooperating.get().stream().map(source ->
cooperatingBatchSource(scanProvider, source))
+ .collect(Collectors.toList());
+ }
List<SplitProducer.BatchSource> sources = new ArrayList<>(sourceCount);
// Bound active partition batches even when full queues make producers
yield their workers.
// Hive streams within each batch; the SPI adapter retains one list
for older implementations.
@@ -2468,6 +2485,60 @@ public class PluginDrivenScanNode extends
FileQueryScanNode {
return sources;
}
+ private SplitProducer.BatchSource
cooperatingBatchSource(ConnectorScanPlanProvider scanProvider,
+ ConnectorSplitSource source) {
+ return new SplitProducer.BatchSource() {
+ private final PluginDrivenSplit.Factory splitFactory = new
PluginDrivenSplit.Factory();
+ private CompletableFuture<Void> blocked;
+
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ return blocked;
+ }
+
+ @Override
+ public List<Split> getNextBatch() {
+ return readMetadata(() -> onPluginClassLoader(scanProvider, ()
-> {
+ blocked = null;
+ List<Split> splits = new
ArrayList<>(splitAssignment.getSplitBatchSize());
+ while (splits.size() < splitAssignment.getSplitBatchSize()
+ && splitAssignment.needMoreSplit()) {
+ CompletableFuture<Void> dependency =
source.isBlocked();
+ if (dependency != null) {
+ if (!dependency.isDone()) {
+ blocked = dependency;
+ break;
+ }
+ dependency.join();
+ }
+ if (source.hasNext()) {
+ splits.add(toSplit(source.next(), splitFactory));
+ continue;
+ }
+ blocked = source.isBlocked();
+ if (blocked == null || !blocked.isDone()) {
+ break;
+ }
+ blocked.join();
+ }
+ return splits;
+ }));
+ }
+
+ @Override
+ public void close() {
+ onPluginClassLoader(scanProvider, () -> {
+ try {
+ source.close();
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ return null;
+ });
+ }
+ };
+ }
+
private SplitProducer.BatchSource
createStreamingBatchSource(ConnectorScanPlanProvider scanProvider,
ConnectorScanRequest request, SummaryProfile summaryProfile) {
RuntimeProfile generationProfile = new RuntimeProfile("Split
Generation " + getId());
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java
b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java
index a85a9f63557..8bc279ea801 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java
@@ -33,6 +33,7 @@ import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
@@ -55,6 +56,10 @@ public class SplitProducer implements Runnable, Closeable {
public interface BatchSource extends Closeable {
List<Split> getNextBatch() throws Exception;
+ default CompletableFuture<Void> isBlocked() {
+ return null;
+ }
+
@Override
default void close() throws IOException {
}
@@ -81,6 +86,7 @@ public class SplitProducer implements Runnable, Closeable {
private Backend pendingBackend;
private List<TScanRangeLocations> pendingRanges;
private int pendingRangeOffset;
+ private CompletableFuture<Void> waitingForSource;
public SplitProducer(SplitAssignment assignment, Executor executor,
BatchSource source,
Consumer<Exception> completion) {
@@ -197,8 +203,18 @@ public class SplitProducer implements Runnable, Closeable {
break;
}
if (pendingRanges == null && !pending.hasNext()) {
+ if (awaitSource()) {
+ more = false;
+ break;
+ }
List<Split> batch = source.getNextBatch();
if (batch.isEmpty()) {
+ // Completion may race the empty read. A dependency,
even one that just
+ // completed, means retry rather than EOF. Never
occupy a worker waiting.
+ if (source.isBlocked() != null) {
+ more = !awaitSource();
+ break;
+ }
exhausted = true;
break;
}
@@ -252,6 +268,22 @@ public class SplitProducer implements Runnable, Closeable {
return false;
}
+ private boolean awaitSource() {
+ CompletableFuture<Void> dependency = source.isBlocked();
+ if (dependency == null) {
+ return false;
+ }
+ if (dependency.isDone()) {
+ dependency.join();
+ return false;
+ }
+ if (waitingForSource != dependency) {
+ waitingForSource = dependency;
+ dependency.whenComplete((unused, error) -> schedule());
+ }
+ return true;
+ }
+
private UserException publishFailure(Throwable failure) {
if (failure == null || cancelled || assignment.isStop()) {
return null;
@@ -269,6 +301,7 @@ public class SplitProducer implements Runnable, Closeable {
pending = Collections.emptyIterator();
pendingRanges = null;
pendingBackend = null;
+ waitingForSource = null;
try {
source.close();
} catch (Throwable e) {
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodePartitionBatchTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodePartitionBatchTest.java
index a14a3fe76e6..4f4fa763659 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodePartitionBatchTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodePartitionBatchTest.java
@@ -56,12 +56,75 @@ import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Queue;
import java.util.Set;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
class PluginDrivenScanNodePartitionBatchTest {
+ @Test
+ void cooperatingSourcesYieldWhileMetadataIsUnavailable() throws Exception {
+ Fixture fixture = new Fixture(1, Collections.emptySet());
+ ConnectorScanRequest request = ConnectorScanRequest.builder(new
ConnectorTableHandle() {
+ }, Collections.emptyList()).build();
+ Iterator<ConnectorScanRange> ranges =
fixture.provider.planScanForPartitionBatch(
+ null, request, Collections.singletonList("0")).iterator();
+ AtomicBoolean loading = new AtomicBoolean(true);
+ CompletableFuture<Void> ready = new CompletableFuture<>();
+ AtomicInteger closed = new AtomicInteger();
+ ConnectorSplitSource connectorSource = new ConnectorSplitSource() {
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ return loading.get() ? ready : null;
+ }
+
+ @Override
+ public boolean hasNext() {
+ return !loading.get() && ranges.hasNext();
+ }
+
+ @Override
+ public ConnectorScanRange next() {
+ return ranges.next();
+ }
+
+ @Override
+ public void close() {
+ closed.incrementAndGet();
+ }
+ };
+ AtomicReference<ConnectorScanRequest> captured = new
AtomicReference<>();
+ Mockito.doAnswer(call -> {
+ captured.set(call.getArgument(1));
+ Assertions.assertEquals(16, call.<Integer>getArgument(2));
+ Assertions.assertEquals(3, call.<Integer>getArgument(3));
+ return Optional.of(Collections.singletonList(connectorSource));
+ }).when(fixture.provider).streamPartitionSources(Mockito.any(),
Mockito.any(), Mockito.anyInt(),
+ Mockito.anyInt());
+ List<SplitProducer.BatchSource> sources = fixture.sources(1000, 16, 3);
+ Assertions.assertEquals(1, sources.size());
+ Assertions.assertEquals(1000,
captured.get().getRequiredPartitions().size());
+ try (SplitProducer.BatchSource source = sources.get(0)) {
+ Assertions.assertTrue(source.getNextBatch().isEmpty());
+ Assertions.assertSame(ready, source.isBlocked());
+ Assertions.assertFalse(ready.isDone());
+ loading.set(false);
+ ready.complete(null);
+ for (long identity = 0; identity < 3; identity++) {
+ Assertions.assertEquals(identity,
source.getNextBatch().get(0).getStart());
+ Assertions.assertNull(source.isBlocked());
+ }
+ Assertions.assertTrue(source.getNextBatch().isEmpty());
+ Assertions.assertNull(source.isBlocked());
+ }
+ Assertions.assertEquals(1, closed.get());
+ Mockito.verify(fixture.provider,
Mockito.never()).streamSplits(Mockito.any(), Mockito.any());
+ }
+
@ParameterizedTest
@ValueSource(ints = {1, 1024})
void batchesRetainConfiguredSizeAcrossPartitions(int partitionsPerBatch)
throws Exception {
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
index 2e847e8a05a..87a18e4605d 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
@@ -49,6 +49,7 @@ import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
@@ -94,6 +95,119 @@ class SplitProducerTest {
};
}
+ @Test
+ void producingThenWaitingDoesNotKeepResubmitting() throws Exception {
+ Fixture fixture = new Fixture(1, 3);
+ ManualExecutor executor = new ManualExecutor();
+ CompletableFuture<Void> dependency = new CompletableFuture<>();
+ RecordingSource source = new RecordingSource(1, 0, 2) {
+ private boolean first = true;
+ private boolean paused;
+
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ return paused ? dependency : null;
+ }
+
+ @Override
+ public List<Split> getNextBatch() throws IOException {
+ paused = first;
+ first = false;
+ return super.getNextBatch();
+ }
+ };
+ fixture.start(executor, source, () -> 0L);
+ executor.runUntilIdle();
+ Assertions.assertEquals(1,
fixture.assignment.getAssignedSplits(fixture.backend).size());
+ Assertions.assertEquals(0, fixture.completions.get());
+ dependency.complete(null);
+ executor.runUntilIdle();
+ Assertions.assertEquals(Arrays.asList(0L, 2L),
identities(fixture.source().getNextBatch(2)));
+ Assertions.assertEquals(1, fixture.completions.get());
+ }
+
+ @Test
+ void asynchronousMetadataWaitReleasesWorkerAndResumesOnCompletion() throws
Exception {
+ Fixture fixture = new Fixture(1, 3);
+ ManualExecutor executor = new ManualExecutor();
+ CompletableFuture<Void> dependency = new CompletableFuture<>();
+ RecordingSource source = new RecordingSource(1, 0) {
+ private boolean awaiting = true;
+
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ return awaiting ? dependency : null;
+ }
+
+ @Override
+ public List<Split> getNextBatch() throws IOException {
+ awaiting = false;
+ return super.getNextBatch();
+ }
+ };
+ fixture.start(executor, source);
+ executor.runUntilIdle();
+ Assertions.assertEquals(0, fixture.completions.get());
+
Assertions.assertTrue(fixture.assignment.getAssignedSplits(fixture.backend).isEmpty());
+ dependency.complete(null);
+ executor.runUntilIdle();
+ Assertions.assertEquals(Collections.singletonList(0L),
identities(fixture.source().getNextBatch(1)));
+ Assertions.assertEquals(1, fixture.completions.get());
+ }
+
+ @Test
+ void metadataCompletionRacingEmptyReadDoesNotBecomeEndOfInput() throws
Exception {
+ Fixture fixture = new Fixture(1, 3);
+ ManualExecutor executor = new ManualExecutor();
+ RecordingSource source = new RecordingSource(1, 0) {
+ private final CompletableFuture<Void> dependency = new
CompletableFuture<>();
+ private CompletableFuture<Void> blocked;
+ private boolean first = true;
+
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ return blocked;
+ }
+
+ @Override
+ public List<Split> getNextBatch() throws IOException {
+ if (first) {
+ first = false;
+ blocked = dependency;
+ dependency.complete(null);
+ return Collections.emptyList();
+ }
+ blocked = null;
+ return super.getNextBatch();
+ }
+ };
+ fixture.start(executor, source);
+ executor.runUntilIdle();
+ Assertions.assertEquals(Collections.singletonList(0L),
identities(fixture.source().getNextBatch(1)));
+ Assertions.assertEquals(1, fixture.completions.get());
+ }
+
+ @Test
+ void metadataFailureWakesProducerAndReachesConsumer() throws Exception {
+ Fixture fixture = new Fixture(1, 3);
+ ManualExecutor executor = new ManualExecutor();
+ CompletableFuture<Void> dependency = new CompletableFuture<>();
+ RecordingSource source = new RecordingSource(1, 0) {
+ @Override
+ public CompletableFuture<Void> isBlocked() {
+ return dependency;
+ }
+ };
+ fixture.start(executor, source);
+ executor.runUntilIdle();
+ dependency.completeExceptionally(new IllegalStateException("metadata
failed"));
+ executor.runUntilIdle();
+ Assertions.assertTrue(Assertions.assertThrows(UserException.class,
+ () ->
fixture.source().getNextBatch(1)).getMessage().contains("metadata failed"));
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ }
+
@Test
void producersFillConfiguredBufferBeforeAnyBackendFetch() throws Exception
{
Fixture fixture = new Fixture(1, 3);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java
b/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java
index 6532d28bde4..ba19b23aa65 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java
@@ -38,6 +38,7 @@ import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
@@ -54,6 +55,7 @@ import java.util.Objects;
import java.util.Random;
import java.util.Set;
import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
public class FederationBackendPolicyTest {
private Env env = Mockito.mock(Env.class);
@@ -71,6 +73,106 @@ public class FederationBackendPolicyTest {
mockedEnvStatic.close();
}
+ @ParameterizedTest
+ @CsvSource({"ROUND_ROBIN, 1, false", "RANDOM, 1, false",
"CONSISTENT_HASHING, 1, false",
+ "CONSISTENT_HASHING, 0, true", "CONSISTENT_HASHING, 3, true",
"RANDOM, 0, true"})
+ public void
testSingleBackendRemoteBatchesSkipHashing(NodeSelectionStrategy strategy,
+ int candidates, boolean spread) throws UserException {
+ Backend backend = new Backend(13001L, "192.0.2.1", 9050);
+ SystemInfoService service = new SystemInfoService();
+ backend.setAlive(true);
+ service.addBackend(backend);
+ mockedEnvStatic.when(Env::getCurrentSystemInfo).thenReturn(service);
+ Mockito.when(env.getComputeGroupMgr()).thenReturn(new
ComputeGroupMgr(service));
+ FederationBackendPolicy policy = new FederationBackendPolicy(strategy,
candidates);
+ policy.init();
+ AtomicInteger identityReads = new AtomicInteger();
+ long expectedWeight = 0;
+ for (int size : List.of(0, 1, 8, 128, 3)) {
+ List<Split> splits = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ long weight = i + 1;
+ FileSplit split = new
FileSplit(LocationPath.of("s3://bucket/part-" + i),
+ 0, 1, 1, 0, new String[]{"another-host"},
Collections.emptyList()) {
+ @Override
+ public String getSplitIdentity() {
+ identityReads.incrementAndGet();
+ return getPathString();
+ }
+
+ @Override
+ public SplitWeight getSplitWeight() {
+ return SplitWeight.fromRawValue(weight);
+ }
+ };
+ split.setAlternativeHosts(List.of("previous-alternative"));
+ splits.add(split);
+ expectedWeight += weight;
+ }
+ Multimap<Backend, Split> assigned =
policy.computeScanRangeAssignment(splits);
+ Assertions.assertEquals(splits, new
ArrayList<>(assigned.get(backend)));
+ Assertions.assertEquals(size, assigned.size());
+ Assertions.assertEquals(expectedWeight,
policy.getAssignedWeightPerBackend().get(backend));
+ for (Split split : splits) {
+ Assertions.assertEquals(spread ?
List.of("previous-alternative") : Collections.emptyList(),
+ split.getAlternativeHosts());
+ }
+ }
+ Assertions.assertEquals(0, identityReads.get(), "One eligible backend
does not need split hashing");
+ }
+
+ @Test
+ public void testSingleBackendPreservesLocalOnlyRestrictions() throws
UserException {
+ Backend backend = new Backend(13002L, "192.0.2.2", 9050);
+ SystemInfoService service = new SystemInfoService();
+ backend.setAlive(true);
+ service.addBackend(backend);
+ mockedEnvStatic.when(Env::getCurrentSystemInfo).thenReturn(service);
+ Mockito.when(env.getComputeGroupMgr()).thenReturn(new
ComputeGroupMgr(service));
+ FederationBackendPolicy policy = new FederationBackendPolicy();
+ policy.init();
+ FileSplit local = new FileSplit(LocationPath.of("file:///tmp/local"),
+ 0, 1, 1, 0, new String[]{backend.getHost()},
Collections.emptyList()) {
+ @Override
+ public boolean isRemotelyAccessible() {
+ return false;
+ }
+ };
+ FileSplit remote = new FileSplit(LocationPath.of("s3://bucket/remote"),
+ 0, 1, 1, 0, new String[]{"another-host"},
Collections.emptyList());
+ Multimap<Backend, Split> assigned = policy.computeScanRangeAssignment(
+ new ArrayList<>(List.of(local, remote)));
+ Assertions.assertEquals(2, assigned.get(backend).size());
+ Assertions.assertEquals(2 * SplitWeight.standard().getRawValue(),
+ policy.getAssignedWeightPerBackend().get(backend));
+ local.setHosts(new String[]{"another-host"});
+ Assertions.assertThrows(UserException.class,
+ () -> policy.computeScanRangeAssignment(new
ArrayList<>(List.of(local))));
+ }
+
+ @Test
+ public void testSingleBackendRejectsWeightOverflow() throws UserException {
+ Backend backend = new Backend(13003L, "192.0.2.3", 9050);
+ SystemInfoService service = new SystemInfoService();
+ backend.setAlive(true);
+ service.addBackend(backend);
+ mockedEnvStatic.when(Env::getCurrentSystemInfo).thenReturn(service);
+ Mockito.when(env.getComputeGroupMgr()).thenReturn(new
ComputeGroupMgr(service));
+ FederationBackendPolicy policy = new FederationBackendPolicy();
+ policy.init();
+ FileSplit heavy = new FileSplit(LocationPath.of("s3://bucket/heavy"),
+ 0, 1, 1, 0, null, Collections.emptyList()) {
+ @Override
+ public SplitWeight getSplitWeight() {
+ return SplitWeight.fromRawValue(Long.MAX_VALUE);
+ }
+ };
+ FileSplit extra = new FileSplit(LocationPath.of("s3://bucket/extra"),
+ 0, 1, 1, 0, null, Collections.emptyList());
+ Assertions.assertThrows(ArithmeticException.class,
+ () -> policy.computeScanRangeAssignment(new
ArrayList<>(List.of(heavy, extra))));
+ }
+
@Test
public void testRoundRobinUsesCandidateOrderAcrossBatches() throws
UserException {
SystemInfoService service = new SystemInfoService();
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
