You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@doris.apache.org by hu...@apache.org on 2026/10/16 02:26:53 UTC
(doris) branch master updated: [improvement](fe) Stream Hive splits and reduce planning overhead (#69133)
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 6c8627ab977 [improvement](fe) Stream Hive splits and reduce planning overhead (#69133)
6c8627ab977 is described below
commit 6c8627ab97743194b03caf14c20faa68778d8b9b
Author: Calvin Kirs <gu...@selectdb.com>
AuthorDate: Fri Oct 16 10:26:43 2026 +0800
[improvement](fe) Stream Hive splits and reduce planning overhead (#69133)
### What problem does this PR solve?
Problem Summary:
Hive batch scans currently materialize all splits within a partition
before publishing them. A `LIMIT` query can therefore wait for work it
will never consume. The consumer can also create an unbounded queue
before the producer applies its capacity, and a fetch can return more
splits than the BE requested.
This change makes split generation incremental, bounds queued work, and
preserves each split across producer and consumer batch boundaries. It
also reduces repeated FE variable export and Hive metadata work on the
query planning path.
#### Implementation
- Generate Hive splits lazily across partitions, files, and file
offsets. Use the request-based connector streaming entry point and reuse
the existing split boundary and format rules. ACID and TABLESAMPLE
special paths retain their existing synchronous behavior.
- Add positive session settings `split_batch_size=128` and
`split_queue_capacity=32`. Both queue creation paths use the same
bounded factory. Capacity is measured in batches per BE: at most 4,096
queued splits with the defaults, excluding producer batches, consumer
residuals, and BE in-flight work.
- Honor the existing BE fetch batch size, retaining the unused portion
of a producer batch for the next request. Preserve residuals at EOF and
propagate errors/interruption.
- Yield the shared scheduling worker when a BE queue is full, then
resume after consumption releases capacity. Preserve assignment and
partial submission progress so successful backend submissions are not
replayed. Serialize source advancement and cleanup; handle cancellation
and submission rejection.
- Prefer eligible partition batching over whole-scan streaming,
preserving concurrent directory listing across Hive partition batches.
Bound active sources per scan by the existing scheduler parallelism and
enumerate each disjoint request scope lazily. Non-streaming connectors
retain their existing eager batch hook through a deferred adapter. A
source closes its exhausted batch before claiming the next one; full
queues yield their workers.
- Share async scan preparation and producer completion across whole-scan
and partition-batch streaming. Synchronous Hive partition batches
collect the same scoped source; whole-scan completeness tracking is
retained. Remove the obsolete queue-default constructor and reuse
variable-specific standard validation.
- Precompute visible variable descriptors while reading live session
values on each export. Reuse successful Hive table metadata within a
statement and the existing positive HMS table cache for existence
checks. Misses/errors are not memoized, fresh reads still reach HMS, and
existing cache invalidation/TTL behavior is retained.
The connector plugin API advances from 15.0 on upstream base
`091925a5629` to 17.0 with the recorded surface updated. External
connector plugins need rebuilding against the new API. There is no BE
protocol or storage-format change.
The request-based streaming API preserves partition scope without an
overload that drops request fields. Hive overrides it with lazy
enumeration; partition-batching connectors can inherit the adapter over
`planScanForPartitionBatch`. Iceberg remains on its whole-scan streaming
path.
#### Local performance measurements
The measurements below identify their exact snapshots. Sections 1–3 were
measured before the final scheduling correctness fixes and upstream
synchronization. Section 4 remeasures the partition-concurrency fix at
`9835d58d88d`. These measurements do not include the subsequent internal
refactor.
Fixture: a Hive nonpartitioned table with 10,000 Parquet files, 100 rows
per file, 1,000,000 rows total. A partitioned fixture has the same file
count in one selected partition. All timings are milliseconds, measured
by client `execute` plus full result fetch on a persistent connection.
Queries execute serially on a shared host without CPU isolation.
```sql
EXPLAIN SELECT * FROM split_hms.split_repro.files_10000 LIMIT 10;
SELECT * FROM split_hms.split_repro.files_10000 LIMIT 10;
```
**1. Controlled split-streaming comparison against the starting master**
Starting master: `84e5e6cd52b`; streaming-stage snapshot: `7be63ee0417`.
Both use the same Doris `4.1.3-rc02-7126cf65d96` BE and the same
settings within each pair. Each SQL has 10 warmups and 100 measured
samples; this stage used its original default cache settings. The later
variable/HMS optimizations are not included in this comparison.
| Fixture | SQL | Starting master median | Streaming-stage median |
Median latency reduction | Starting master P95 | Streaming-stage P95 |
|---|---|---:|---:|---:|---:|---:|
| Nonpartitioned, 10k files | EXPLAIN | 57.50 | 25.11 | 56.33% | 72.71 |
60.99 |
| Nonpartitioned, 10k files | SELECT LIMIT 10 | 340.47 | 69.91 | 79.47%
| 392.70 | 123.88 |
| One partition, 10k files | EXPLAIN | 51.00 | 25.22 | 50.55% | 66.06 |
58.56 |
| One partition, 10k files | SELECT LIMIT 10 | 153.55 | 76.32 | 50.29% |
186.48 | 119.82 |
Separate A/B runs for the planning/cache changes used 500 warmups and
100 measured samples per SQL. Variable export plus statement metadata
reuse reduced EXPLAIN median from 14.21 to 9.60 ms and SELECT median
from 38.33 to 16.26 ms. Reusing the positive Hive table cache for
existence checks reduced EXPLAIN from 7.62 to 3.81 ms and SELECT from
13.20 to 9.54 ms. These are separate comparisons; their percentages must
not be added. Packet captures for 30 warm EXPLAIN statements showed HMS
calls falling from 60 to 30, then to 0.
**2. Full optimization snapshot versus original Doris 4.1.3**
Full optimization snapshot: `80cbc3e54d9` FE with the same 4.1.3 BE.
Original version: `apache/doris:all-in-one-4.1.3`, actual FE/BE build
`4.1.3-rc02-7126cf65d96`.
Same nonpartitioned fixture, FE heaps 2 GiB, 500 warmup query pairs and
100 measured pairs per round, two rounds with reversed engine order. The
pooled result contains 200 raw samples per SQL/version. SQL result
cache, BE file cache, and Profile are disabled; FE metadata/file-list
caches are warm. The current FE also disables bucketed hash aggregation
and the local shuffle planner for compatibility with the older BE.
The original BE exceeded its 2 GB memory limit during warmup, before
formal sampling. Only the original BE was increased to 3 GB for the
completed comparison; the optimized FE's BE remained at 2 GB. This is a
comparison of the recorded configurations, not an equal-memory product
ranking.
| Version | BE memory limit | EXPLAIN median | EXPLAIN P95 | SELECT
median | SELECT P95 |
|---|---:|---:|---:|---:|---:|
| Original Doris 4.1.3 | 3 GB | 50.65 | 152.04 | 372.11 | 1549.85 |
| Optimized FE `80cbc3e54d9` + 4.1.3 BE | 2 GB | 4.42 | 34.81 | 11.91 |
84.63 |
The observed median reductions are 91.27% for EXPLAIN and 96.80% for
SELECT in this comparison. For the optimized FE, round 1 recorded
EXPLAIN 12.65 ms and SELECT 37.21 ms; round 2 recorded EXPLAIN 3.31 ms
and SELECT 7.76 ms. Both rounds followed 500 warmups. These rounds ran
after an FE restart on a shared host; the causes of the variation were
not isolated, and both rounds are retained. Master also contains changes
beyond this PR relative to 4.1.3, so the entire difference cannot be
attributed to these commits.
**3. Additional observations**
- A separate two-round run of `80cbc3e54d9`, with 500 warmups and 100
measurements per round, recorded pooled EXPLAIN median/P95 of 3.25/4.96
ms and SELECT median/P95 of 7.75/15.82 ms. It did not simultaneously
remeasure the starting master or original 4.1.3. It must not replace the
optimized values in the paired tables above.
- With the same FE, file listing warm, and default P128/Q32, artificial
fine ranges over 10,000 files produced 1,000,000 potential splits.
Synchronous versus streaming EXPLAIN medians were 3329.12 versus 3.49
ms; streaming LIMIT median was 59.24 ms. Each configuration had 45
measured samples. A separate profile generated only 4,352 splits for the
LIMIT. This is a million-split stress case, not a million-real-file
full-scan measurement.
- Full COUNT/SUM still consumes all splits. Small same-FE
synchronous/streaming comparisons recorded nonpartitioned medians of
1598.00/1592.62 ms and single-partition medians of 1699.13/1548.57 ms (2
warmups, 12 samples each). Large LIMIT savings should not be generalized
to all full scans.
### Release note
Hive batch scans generate splits incrementally with bounded producer
batches and per-BE queues. BE split fetches respect their requested
batch size. Full queues no longer occupy shared scheduling workers, and
partition-batch pre-generation is bounded per scan. Warm Hive planning
avoids redundant metadata requests. Connector plugin API is updated to
17.0; rebuild external connector plugins against the new API.
### Check List (For Author)
- Test
- [x] Regression test
- `test_hive_streaming_splits`: 49 result groups in one suite, standard
generated expected output, then comparison against the final
pre-synchronization feature snapshot `f33859be44a` with current FE +
original 4.1.3 BE and one scheduler worker. Covers
synchronous/P3Q1/P7Q2/P128Q32, COUNT/SUM/distinct identities, partition
pruning, empty inputs, JOIN/UNION, LIMIT, CTAS, hints, and invalid
settings.
- [x] Unit Test
- 103 targeted tests in ten classes passed at `f33859be44a`.
- After synchronization with master `091925a5629`, 311 tests in 49
classes passed using `run-fe-ut.sh`, including the complete connector
SPI suite at API 16.0, Hive/HMS metadata and variable export tests,
split scheduling/consumption tests, and the upstream
backend-capability/projection tests.
- Review fix `9835d58d88d`: 643 distinct tests in 65 classes passed
across the broad run and reruns after fixing the old test fixture.
Includes complete SPI tests at API 17.0, 191 Iceberg scan cases, 35 Hive
scan cases, and the actual Hive mode-selection/dispatch concurrency
tests. A temporary restoration of the old routing fails exactly at the
concurrent-progress assertion; restoring the fix passes all 17 targeted
routing cases.
- Follow-up refactor `16c9999c698`: 500 tests in 42 classes passed with
`run-fe-ut.sh`; includes real Hive single-source/partition routing,
Parquet/LZO eager/lazy equivalence, scheduling/consumption, all
PluginDrivenScanNode tests, 191 Iceberg cases, and actual VariableMgr
positive/negative setting validation.
- [x] Manual test (add detailed scripts or steps below)
- Performance measurements and conditions are listed above. Million-row
COUNT/SUM, per-file identities, and actual million-file LIMIT were also
checked; complete aggregates over the actual million-file fixture remain
outstanding.
- FE `build.sh --fe` and Checkstyle passed at `f33859be44a`.
- The earlier upstream-synchronized head `9e2879a5b70` passed `build.sh
--fe` (87 modules), including Checkstyle and connector packaging.
- Review fix `9835d58d88d` also passed `build.sh --fe` and Checkstyle;
the packaged SPI advertises API 17.0. Its performance was subsequently
remeasured in isolated services; conditions and tradeoffs are in section
4.
- Follow-up refactor `16c9999c698`: `DISABLE_BUILD_UI=ON build.sh --fe`
and Checkstyle passed; UI rebuilding was excluded for this Java-only
refactor. No post-refactor performance rerun is claimed.
- Behavior changed:
- [x] Yes. Incremental/bounded Hive split generation, strict BE
consumption batches, cooperative scheduling, positive metadata reuse,
and connector API 17.0 as described above.
- Does this need documentation?
- [x] Yes. The two session settings and connector API version change
need user/plugin documentation; the documentation PR is pending.
---
.../connector/hive/HiveConnectorMetadata.java | 16 +-
.../doris/connector/hive/HiveFileListingCache.java | 6 +
.../doris/connector/hive/HiveScanPlanProvider.java | 303 +++++++++++----
.../HiveConnectorMetadataTableInfoMemoTest.java | 202 ++++++++++
.../connector/hive/HiveScanBatchModeTest.java | 245 +++++++++++-
.../doris/connector/hms/CachingHmsClient.java | 17 +-
.../doris/connector/hms/CachingHmsClientTest.java | 97 ++++-
.../connector/iceberg/IcebergScanPlanProvider.java | 6 +-
.../iceberg/IcebergScanPlanProviderTest.java | 78 ++--
.../spi/scan/ConnectorScanPlanProvider.java | 88 +++--
.../connector/spi/scan/ConnectorSplitSource.java | 4 +-
.../connector/spi/ConnectorPluginSurfaceTest.java | 4 +-
.../ConnectorScanPlanProviderBatchScanTest.java | 48 ++-
.../test/resources/connector-plugin-surface.txt | 2 +-
fe/fe-connector/pom.xml | 2 +-
.../doris/datasource/scan/FileQueryScanNode.java | 14 +-
.../datasource/scan/PluginDrivenScanNode.java | 315 +++++++++-------
.../doris/datasource/split/SplitAssignment.java | 162 ++++++--
.../doris/datasource/split/SplitProducer.java | 183 +++++++++
.../apache/doris/datasource/split/SplitSource.java | 71 ++--
.../java/org/apache/doris/qe/SessionVariable.java | 28 ++
.../main/java/org/apache/doris/qe/VariableMgr.java | 16 +-
.../PluginDrivenScanNodeHiveStreamingTest.java | 325 ++++++++++++++++
.../PluginDrivenScanNodePartitionBatchTest.java | 211 +++++++++++
.../datasource/split/SplitAssignmentTest.java | 204 +++++-----
.../doris/datasource/split/SplitProducerTest.java | 411 +++++++++++++++++++++
.../doris/datasource/split/SplitSourceTest.java | 207 +++++++++++
.../split/SplitStreamingConcurrencyTest.java | 164 ++++++++
.../org/apache/doris/qe/SessionVariablesTest.java | 18 +
.../org/apache/doris/qe/VariableMgrExportTest.java | 86 +++++
.../hive/test_hive_streaming_splits.out | 236 ++++++++++++
.../hive/test_hive_streaming_splits.groovy | 135 +++++++
32 files changed, 3442 insertions(+), 462 deletions(-)
diff --git a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorMetadata.java b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorMetadata.java
index 402a69bcd85..64e503c219f 100644
--- a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorMetadata.java
+++ b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorMetadata.java
@@ -38,6 +38,7 @@ import org.apache.doris.connector.spi.ConnectorDatabaseMetadata;
import org.apache.doris.connector.spi.ConnectorMetadata;
import org.apache.doris.connector.spi.ConnectorPartitionInfo;
import org.apache.doris.connector.spi.ConnectorSession;
+import org.apache.doris.connector.spi.ConnectorStatementScopes;
import org.apache.doris.connector.spi.ConnectorStorageContext;
import org.apache.doris.connector.spi.ConnectorTableSchema;
import org.apache.doris.connector.spi.ConnectorTableStatistics;
@@ -94,6 +95,7 @@ import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
import java.util.OptionalLong;
import java.util.Set;
@@ -119,6 +121,7 @@ import java.util.stream.Collectors;
public class HiveConnectorMetadata implements ConnectorMetadata {
private static final Logger LOG = LogManager.getLogger(HiveConnectorMetadata.class);
+ static final String TABLE_INFO_NAMESPACE = "hms.table-info";
/**
* The HMS table parameter iceberg writes its table comment into (mirrored from the iceberg table property
@@ -436,10 +439,19 @@ public class HiveConnectorMetadata implements ConnectorMetadata {
@Override
public Optional<ConnectorTableHandle> getTableHandle(
ConnectorSession session, String dbName, String tableName) {
- if (!hmsClient.tableExists(dbName, tableName)) {
+ // Cache raw metadata, deriving handles outside the loader so sibling resolution can use
+ // its own statement memos without nesting inside this computeIfAbsent.
+ HmsTableInfo tableInfo = ConnectorStatementScopes.resolveInStatement(
+ session, TABLE_INFO_NAMESPACE, dbName, tableName, () -> {
+ if (!hmsClient.tableExists(dbName, tableName)) {
+ // ConcurrentHashMap does not retain null: CTAS must resolve a target after creating it.
+ return null;
+ }
+ return Objects.requireNonNull(hmsClient.getTable(dbName, tableName));
+ });
+ if (tableInfo == null) {
return Optional.empty();
}
- HmsTableInfo tableInfo = hmsClient.getTable(dbName, tableName);
HiveTableType tableType = HiveTableFormatDetector.detect(tableInfo);
// Foreign-handle divert: an iceberg-on-HMS or hudi-on-HMS table registered in this HMS catalog is served by
// the embedded iceberg / hudi SIBLING connector, not by hive. Return the sibling's OWN table handle (the
diff --git a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveFileListingCache.java b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveFileListingCache.java
index 5d9116abff4..09f28d3b879 100644
--- a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveFileListingCache.java
+++ b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveFileListingCache.java
@@ -187,6 +187,12 @@ public class HiveFileListingCache {
cache.isEnabled() && cache.isWeightBounded())).files;
}
+ /** Returns a cheap batch-planning hint, or -1 when listing would require I/O. */
+ long cachedFileCount(String dbName, String tableName, String location, List<String> partitionValues) {
+ FileListingValue value = cache.getIfPresent(new FileListingKey(dbName, tableName, location, partitionValues));
+ return value == null ? -1 : value.files.size();
+ }
+
/** Drops every cached listing for one table. Backs {@code REFRESH TABLE}. */
public void invalidateTable(String dbName, String tableName) {
owner.invalidateTable(dbName, tableName);
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 28bcaef410d..f3e4a916b38 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
@@ -33,6 +33,7 @@ import org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider;
import org.apache.doris.connector.spi.scan.ConnectorScanProfile;
import org.apache.doris.connector.spi.scan.ConnectorScanRange;
import org.apache.doris.connector.spi.scan.ConnectorScanRequest;
+import org.apache.doris.connector.spi.scan.ConnectorSplitSource;
import org.apache.doris.connector.spi.scan.ScanNodePropertyKeys;
import org.apache.doris.filesystem.FileEntry;
import org.apache.doris.filesystem.FileSystem;
@@ -48,10 +49,12 @@ import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
+import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
+import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
@@ -70,13 +73,8 @@ import java.util.function.UnaryOperator;
* <li>Return {@link HiveScanRange} for each file split</li>
* </ol>
*
- * <p>Current limitations (Phase 2 initial):</p>
- * <ul>
- * <li>No file listing cache (lists files directly each time)</li>
- * <li>No ACID transaction support (non-transactional tables only)</li>
- * <li>No file-count streaming split mode; partition-batch mode is limited to partitioned
- * non-transactional tables (see {@link #supportsBatchScan})</li>
- * </ul>
+ * <p>Non-transactional tables can stream file splits without materializing a partition's ranges.
+ * Transactional tables retain snapshot-aware synchronous planning.</p>
*/
public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
@@ -88,6 +86,7 @@ public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
/** Maximum number of partitions to list from HMS. */
private static final int MAX_PARTITIONS = 100000;
+ private static final int DEFAULT_BATCH_THRESHOLD = 1024;
/**
* Connector-internal signal (consumed by {@link #populateScanLevelParams}) marking a full-ACID
@@ -248,6 +247,177 @@ public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
return partKeyNames != null && !partKeyNames.isEmpty() && !hiveHandle.isTransactional();
}
+ @Override
+ public long streamingSplitEstimate(ConnectorSession session, ConnectorTableHandle handle,
+ Optional<ConnectorExpression> filter, boolean countPushdown) {
+ HiveTableHandle hiveHandle = (HiveTableHandle) handle;
+ if (hiveHandle.isTransactional()
+ || Boolean.FALSE.equals(session.getProperty("enable_external_table_batch_mode", Boolean.class))) {
+ return -1;
+ }
+ int fileThreshold = sessionInt(session, "num_files_in_batch_mode", DEFAULT_BATCH_THRESHOLD);
+ long fileCount = cachedFileCount(hiveHandle);
+ // Unknown file counts must not trigger a full listing on the planning thread. Use a
+ // nonzero concurrency hint and let the lazy source perform the listing asynchronously.
+ if (fileCount < 0) {
+ return Math.max(DEFAULT_BATCH_THRESHOLD, fileThreshold);
+ }
+ return fileCount >= fileThreshold ? fileCount : -1;
+ }
+
+ private long cachedFileCount(HiveTableHandle handle) {
+ List<String> partitionKeys = handle.getPartitionKeyNames();
+ if (partitionKeys == null || partitionKeys.isEmpty()) {
+ return fileListingCache.cachedFileCount(handle.getDbName(), handle.getTableName(),
+ handle.getLocation(), Collections.emptyList());
+ }
+ List<HmsPartitionInfo> partitions = handle.getPrunedPartitions();
+ if (partitions == null) {
+ return -1;
+ }
+ long fileCount = 0;
+ for (HmsPartitionInfo partition : partitions) {
+ long count = fileListingCache.cachedFileCount(handle.getDbName(), handle.getTableName(),
+ partition.getLocation(), partition.getValues());
+ if (count < 0) {
+ return -1;
+ }
+ fileCount += count;
+ }
+ return fileCount;
+ }
+
+ @Override
+ public ConnectorSplitSource streamSplits(ConnectorSession session, ConnectorScanRequest request) {
+ HiveTableHandle hiveHandle = (HiveTableHandle) request.getTableHandle();
+ if (hiveHandle.isTransactional()) {
+ throw new DorisConnectorException("Transactional Hive scans require snapshot-aware planning");
+ }
+ return new HiveSplitSource(session, hiveHandle, request.getRequiredPartitions());
+ }
+
+ private static int sessionInt(ConnectorSession session, String name, int defaultValue) {
+ Integer value = session.getProperty(name, Integer.class);
+ return value == null ? defaultValue : value;
+ }
+
+ private final class HiveSplitSource implements ConnectorSplitSource {
+ private final HiveTableHandle handle;
+ private final FileSystem fs;
+ private final HiveFileFormat tableFormat;
+ private final boolean isLzo;
+ private final boolean splittable;
+ private final long targetSplitSize;
+ private final int partitionBatchSize;
+ private Iterator<PartitionScanInfo> partitions = Collections.emptyIterator();
+ private Iterator<HiveFileStatus> files = Collections.emptyIterator();
+ private Iterator<ConnectorScanRange> ranges = Collections.emptyIterator();
+ private List<String> partitionNames;
+ private PartitionScanInfo partition;
+ private int nextPartition;
+ private boolean unpartitionedInitialized;
+ private volatile boolean closed;
+
+ HiveSplitSource(ConnectorSession session, HiveTableHandle handle, List<String> requiredPartitions) {
+ this.handle = handle;
+ this.partitionNames = requiredPartitions.isEmpty() ? null : requiredPartitions;
+ this.fs = storage().getFileSystem(session);
+ this.tableFormat = HiveFileFormat.detect(handle.getInputFormat(), handle.getSerializationLib(),
+ readHiveJsonInOneColumn(session), handle.isFirstColumnString());
+ this.isLzo = isLzoInputFormat(handle.getInputFormat());
+ this.splittable = tableFormat.isSplittable() && !isLzo;
+ this.targetSplitSize = getTargetSplitSize(session);
+ this.partitionBatchSize = Math.max(1,
+ sessionInt(session, "num_partitions_in_batch_mode", DEFAULT_BATCH_THRESHOLD));
+ recordPruningProfile(handle);
+ }
+
+ @Override
+ public boolean hasNext() {
+ while (!closed) {
+ if (ranges.hasNext()) {
+ return true;
+ }
+ if (files.hasNext()) {
+ loadNextFileRanges();
+ continue;
+ }
+ if (partitions.hasNext()) {
+ loadNextPartitionFiles();
+ continue;
+ }
+ if (!loadNextPartitions()) {
+ return false;
+ }
+ }
+ return false;
+ }
+
+ private void loadNextFileRanges() {
+ HiveFileStatus file = files.next();
+ if (isLzo && !isLzoDataFile(file.getPath())) {
+ return;
+ }
+ HiveFileFormat format = partition.fileFormat == null ? tableFormat : partition.fileFormat;
+ ranges = splitFileRanges(file.getPath(), file.getLength(), file.getModificationTime(),
+ partition, format, splittable, targetSplitSize, null, null);
+ }
+
+ private void loadNextPartitionFiles() {
+ partition = partitions.next();
+ files = listPartitionFiles(handle.getDbName(), handle.getTableName(),
+ partition, fs, null).iterator();
+ }
+
+ private boolean loadNextPartitions() {
+ List<String> keys = handle.getPartitionKeyNames();
+ if (keys == null || keys.isEmpty()) {
+ if (unpartitionedInitialized) {
+ return false;
+ }
+ unpartitionedInitialized = true;
+ partitions = Collections.singletonList(new PartitionScanInfo(
+ handle.getLocation(), Collections.emptyMap(), null)).iterator();
+ return true;
+ }
+ List<HmsPartitionInfo> pruned = handle.getPrunedPartitions();
+ if (partitionNames == null && pruned != null) {
+ if (nextPartition == pruned.size()) {
+ return false;
+ }
+ int end = Math.min(pruned.size(), nextPartition + partitionBatchSize);
+ partitions = convertPartitions(pruned.subList(nextPartition, end), keys).iterator();
+ nextPartition = end;
+ return true;
+ }
+ if (partitionNames == null) {
+ partitionNames = hmsClient.listPartitionNames(
+ handle.getDbName(), handle.getTableName(), MAX_PARTITIONS);
+ }
+ if (nextPartition == partitionNames.size()) {
+ return false;
+ }
+ int end = Math.min(partitionNames.size(), nextPartition + partitionBatchSize);
+ partitions = convertPartitions(loadPartitionsWithProfile(handle.getDbName(), handle.getTableName(),
+ partitionNames.subList(nextPartition, end)), keys).iterator();
+ nextPartition = end;
+ return true;
+ }
+
+ @Override
+ public ConnectorScanRange next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return ranges.next();
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ }
+ }
+
/**
* Hive scan ranges (base/insert files) carry positive byte lengths, so the engine can apply
* {@code TABLESAMPLE} by size-weighted split selection — restoring the legacy
@@ -282,8 +452,8 @@ public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
* strings are the metastore-rendered {@code key=value/...} partition names (the keys of the Nereids
* selected-partition map), exactly the form {@code hmsClient.getPartitions} accepts, so batch-scoped
* resolution is a clean {@code getPartitions(batch)} round-trip. Only the non-ACID path is reachable here
- * ({@link #supportsBatchScan} excludes transactional tables). Reuses the SAME helpers as {@link #planScan}:
- * format detect, split size, hadoop conf, {@link #convertPartitions}, {@link #listAndSplitFiles}.
+ * ({@link #supportsBatchScan} excludes transactional tables). Collects the same partition-scoped source
+ * that {@link #streamSplits} consumes lazily, preserving its file filtering and split boundaries.
*/
@Override
public List<ConnectorScanRange> planScanForPartitionBatch(
@@ -293,43 +463,16 @@ public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
// 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.
- return doPlanScanForPartitionBatch(session, request, partitionBatch);
- }
-
- private List<ConnectorScanRange> doPlanScanForPartitionBatch(
- ConnectorSession session,
- ConnectorScanRequest request,
- List<String> partitionBatch) {
- HiveTableHandle hiveHandle = (HiveTableHandle) request.getTableHandle();
- String dbName = hiveHandle.getDbName();
- String tableName = hiveHandle.getTableName();
- recordPruningProfile(hiveHandle);
-
- // Resolve ONLY this batch's partitions (scoped to partitionBatch), NOT handle.getPrunedPartitions().
- List<HmsPartitionInfo> hmsPartitions = loadPartitionsWithProfile(dbName, tableName, partitionBatch);
- List<PartitionScanInfo> partitions = convertPartitions(
- hmsPartitions, hiveHandle.getPartitionKeyNames());
- if (partitions.isEmpty()) {
+ // An empty batch contains no partitions; an empty streaming request instead means scan all.
+ if (partitionBatch.isEmpty()) {
return Collections.emptyList();
}
-
- HiveFileFormat fileFormat = HiveFileFormat.detect(
- hiveHandle.getInputFormat(), hiveHandle.getSerializationLib(),
- readHiveJsonInOneColumn(session), hiveHandle.isFirstColumnString());
- long targetSplitSize = getTargetSplitSize(session);
- boolean isLzo = isLzoInputFormat(hiveHandle.getInputFormat());
- // LZO text is not splittable (see planScan); mask it out of the TEXT-derived splittable flag.
- boolean splittable = fileFormat.isSplittable() && !isLzo;
- // Only the non-ACID path is reachable here (supportsBatchScan excludes transactional tables), so this
- // borrows the engine's per-catalog Doris FileSystem to list — no Hadoop Configuration is needed.
- FileSystem fs = storage().getFileSystem(session);
-
List<ConnectorScanRange> ranges = new ArrayList<>();
- for (PartitionScanInfo partition : partitions) {
- HiveFileFormat partFormat = partition.fileFormat != null
- ? partition.fileFormat : fileFormat;
- listAndSplitFiles(dbName, tableName, partition, partFormat,
- splittable, isLzo, targetSplitSize, fs, ranges, null);
+ try (HiveSplitSource source = new HiveSplitSource(
+ session, (HiveTableHandle) request.getTableHandle(), partitionBatch)) {
+ while (source.hasNext()) {
+ ranges.add(source.next());
+ }
}
return ranges;
}
@@ -643,6 +786,17 @@ public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
PartitionScanInfo partition, HiveFileFormat fileFormat,
boolean splittable, boolean isLzo, long targetSplitSize, FileSystem fs,
List<ConnectorScanRange> ranges, PlanCompleteness completeness) {
+ for (HiveFileStatus file : listPartitionFiles(dbName, tableName, partition, fs, completeness)) {
+ if (isLzo && !isLzoDataFile(file.getPath())) {
+ continue;
+ }
+ splitFile(file.getPath(), file.getLength(), file.getModificationTime(),
+ partition, fileFormat, splittable, targetSplitSize, null, null, ranges);
+ }
+ }
+
+ private List<HiveFileStatus> listPartitionFiles(String dbName, String tableName,
+ PartitionScanInfo partition, FileSystem fs, PlanCompleteness completeness) {
List<HiveFileStatus> files;
try {
files = fileListingCache.listDataFiles(dbName, tableName, partition.location,
@@ -660,19 +814,9 @@ public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
completeness.markIncomplete();
}
LOG.warn("Cannot list files in partition: {}", partition.location, e);
- return;
- }
- for (HiveFileStatus file : files) {
- // LZO text tables: only *.lzo files are data. Exclude *.lzo.index sidecars (and any other
- // non-*.lzo entry), mirroring Hive's LzoTextInputFormat.listStatus() and legacy
- // HiveExternalMetaCache's HiveUtil.isLzoDataFile filter. HiveFileListingCache strips only
- // _/.-prefixed hidden files, so without this a *.lzo.index sidecar is read as an extra text row.
- if (isLzo && !isLzoDataFile(file.getPath())) {
- continue;
- }
- splitFile(file.getPath(), file.getLength(), file.getModificationTime(),
- partition, fileFormat, splittable, targetSplitSize, null, null, ranges);
+ return Collections.emptyList();
}
+ return files;
}
private static final class PlanCompleteness {
@@ -754,29 +898,38 @@ public class HiveScanPlanProvider implements ConnectorScanPlanProvider {
HiveFileFormat fileFormat, boolean splittable, long targetSplitSize,
String acidPartitionLocation, List<String> acidDeltas,
List<ConnectorScanRange> ranges) {
- if (fileSize == 0) {
- return;
- }
- // Normalize the BE-facing native data-file path scheme (s3a://->s3://): the connector lists files via the
- // engine FileSystem with the raw scheme (Hadoop wants s3a), but BE's native S3 reader rejects s3a. ACID
- // delete-delta / partition paths are normalized separately at their emit sites.
- filePath = normalizeNativeUri(filePath);
+ splitFileRanges(filePath, fileSize, modTime, partition, fileFormat, splittable,
+ targetSplitSize, acidPartitionLocation, acidDeltas).forEachRemaining(ranges::add);
+ }
- if (!splittable || targetSplitSize <= 0 || fileSize <= targetSplitSize) {
- // Single range for the whole file
- ranges.add(newRangeBuilder(filePath, 0, fileSize, fileSize, modTime, fileFormat,
- partition, acidPartitionLocation, acidDeltas).build());
- return;
+ private Iterator<ConnectorScanRange> splitFileRanges(String filePath, long fileSize, long modTime,
+ PartitionScanInfo partition, HiveFileFormat fileFormat, boolean splittable, long targetSplitSize,
+ String acidPartitionLocation, List<String> acidDeltas) {
+ if (fileSize == 0) {
+ return Collections.emptyIterator();
}
+ String normalizedPath = normalizeNativeUri(filePath);
+ long splitSize = splittable && targetSplitSize > 0 ? targetSplitSize : fileSize;
+ return new Iterator<ConnectorScanRange>() {
+ private long offset;
+
+ @Override
+ public boolean hasNext() {
+ return offset < fileSize;
+ }
- // Split file into ranges
- long offset = 0;
- while (offset < fileSize) {
- long splitSize = Math.min(targetSplitSize, fileSize - offset);
- ranges.add(newRangeBuilder(filePath, offset, splitSize, fileSize, modTime, fileFormat,
- partition, acidPartitionLocation, acidDeltas).build());
- offset += splitSize;
- }
+ @Override
+ public ConnectorScanRange next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ long length = Math.min(splitSize, fileSize - offset);
+ ConnectorScanRange range = newRangeBuilder(normalizedPath, offset, length, fileSize,
+ modTime, fileFormat, partition, acidPartitionLocation, acidDeltas).build();
+ offset += length;
+ return range;
+ }
+ };
}
private static HiveScanRange.Builder newRangeBuilder(String filePath, long start, long length,
diff --git a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveConnectorMetadataTableInfoMemoTest.java b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveConnectorMetadataTableInfoMemoTest.java
new file mode 100644
index 00000000000..105df566ca5
--- /dev/null
+++ b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveConnectorMetadataTableInfoMemoTest.java
@@ -0,0 +1,202 @@
+// 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.HmsClient;
+import org.apache.doris.connector.hms.HmsDatabaseInfo;
+import org.apache.doris.connector.hms.HmsPartitionInfo;
+import org.apache.doris.connector.hms.HmsTableInfo;
+import org.apache.doris.connector.spi.ConnectorColumn;
+import org.apache.doris.connector.spi.ConnectorSession;
+import org.apache.doris.connector.spi.ConnectorStatementScope;
+import org.apache.doris.connector.spi.ConnectorType;
+import org.apache.doris.connector.spi.handle.ConnectorTableHandle;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+class HiveConnectorMetadataTableInfoMemoTest {
+ private final RecordingHmsClient client = new RecordingHmsClient();
+ private final HiveConnectorMetadata metadata = new HiveConnectorMetadata(
+ client, HiveTestProperties.minimal(), new FakeConnectorContext());
+
+ @Test
+ void reusesPositiveMetadataWithinOneStatement() {
+ TestStatementScope scope = new TestStatementScope();
+ ConnectorSession session = new ScopeSession(1, "q1", scope);
+ HiveTableHandle first = handle(session, "db", "t");
+ client.generation = 2;
+ HiveTableHandle second = handle(session, "db", "t");
+ Assertions.assertEquals(first.getLocation(), second.getLocation());
+ Assertions.assertEquals(1, client.existsCalls);
+ Assertions.assertEquals(1, client.tableCalls);
+ Assertions.assertTrue(HiveConnectorMetadata.TABLE_INFO_NAMESPACE.startsWith(
+ new HiveConnectorProvider().getType() + "."));
+ }
+
+ @Test
+ void isolatesQueryCatalogAndTableIdentities() {
+ TestStatementScope scope = new TestStatementScope();
+ handle(new ScopeSession(1, "q1", scope), "db", "t");
+ client.generation = 2;
+ Assertions.assertTrue(handle(new ScopeSession(1, "q2", scope), "db", "t").getLocation().endsWith("/v2"));
+ handle(new ScopeSession(2, "q2", scope), "db", "t");
+ handle(new ScopeSession(1, "q2", scope), "other_db", "t");
+ handle(new ScopeSession(1, "q2", scope), "db", "other_t");
+ Assertions.assertEquals(5, client.existsCalls);
+ Assertions.assertEquals(5, client.tableCalls);
+ }
+
+ @Test
+ void doesNotRememberAbsenceBeforeCreateOrConcurrentCreation() {
+ ConnectorSession session = new ScopeSession(1, "q1", new TestStatementScope());
+ client.exists = false;
+ Assertions.assertFalse(metadata.getTableHandle(session, "db", "t").isPresent());
+ client.exists = true;
+ handle(session, "db", "t");
+ handle(session, "db", "t");
+ Assertions.assertEquals(2, client.existsCalls);
+ Assertions.assertEquals(1, client.tableCalls);
+ }
+
+ @Test
+ void propagatesLoadFailureWithoutRememberingIt() {
+ ConnectorSession session = new ScopeSession(1, "q1", new TestStatementScope());
+ client.failure = new IllegalStateException("HMS unavailable");
+ Assertions.assertSame(client.failure, Assertions.assertThrows(IllegalStateException.class,
+ () -> metadata.getTableHandle(session, "db", "t")));
+ client.failure = null;
+ handle(session, "db", "t");
+ Assertions.assertEquals(2, client.existsCalls);
+ Assertions.assertEquals(1, client.tableCalls);
+ }
+
+ @Test
+ void noStatementScopeKeepsRevalidating() {
+ ConnectorSession session = new ScopeSession(1, "q1", ConnectorStatementScope.NONE);
+ handle(session, "db", "t");
+ handle(session, "db", "t");
+ handle(null, "db", "t");
+ handle(null, "db", "t");
+ Assertions.assertEquals(4, client.existsCalls);
+ Assertions.assertEquals(4, client.tableCalls);
+ }
+
+ @Test
+ void explicitFreshReadBypassesTheMemo() {
+ ConnectorSession session = new ScopeSession(1, "q1", new TestStatementScope());
+ HiveTableHandle first = handle(session, "db", "t");
+ client.generation = 2;
+ String ddl = metadata.renderShowCreateTableDdl(session, first).get();
+ Assertions.assertTrue(ddl.contains("/v2"), ddl);
+ Assertions.assertTrue(handle(session, "db", "t").getLocation().endsWith("/v1"));
+ Assertions.assertEquals(1, client.freshCalls);
+ Assertions.assertEquals(1, client.existsCalls);
+ Assertions.assertEquals(1, client.tableCalls);
+ }
+
+ private HiveTableHandle handle(ConnectorSession session, String database, String table) {
+ ConnectorTableHandle handle = metadata.getTableHandle(session, database, table).get();
+ return (HiveTableHandle) handle;
+ }
+
+ private static final class RecordingHmsClient implements HmsClient {
+ private int existsCalls;
+ private int tableCalls;
+ private int freshCalls;
+ private int generation = 1;
+ private boolean exists = true;
+ private RuntimeException failure;
+
+ @Override
+ public boolean tableExists(String database, String table) {
+ existsCalls++;
+ if (failure != null) {
+ throw failure;
+ }
+ return exists;
+ }
+
+ @Override
+ public HmsTableInfo getTable(String database, String table) {
+ tableCalls++;
+ return table(database, table);
+ }
+
+ @Override
+ public HmsTableInfo getTableFresh(String database, String table) {
+ freshCalls++;
+ return table(database, table);
+ }
+
+ private HmsTableInfo table(String database, String table) {
+ return HmsTableInfo.builder().dbName(database).tableName(table)
+ .tableType("MANAGED_TABLE")
+ .inputFormat("org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat")
+ .outputFormat("org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat")
+ .serializationLib("org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe")
+ .location("s3://bucket/" + database + "/" + table + "/v" + generation)
+ .columns(Collections.singletonList(new ConnectorColumn(
+ "id", ConnectorType.of("BIGINT"), null, true, null)))
+ .partitionKeys(Collections.emptyList()).parameters(Collections.emptyMap()).build();
+ }
+
+ @Override
+ public List<String> listDatabases() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public HmsDatabaseInfo getDatabase(String database) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public List<String> listTables(String database) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public Map<String, String> getDefaultColumnValues(String database, String table) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public List<String> listPartitionNames(String database, String table, int maxParts) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public List<HmsPartitionInfo> getPartitions(String database, String table, List<String> names) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public HmsPartitionInfo getPartition(String database, String table, List<String> values) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void close() {
+ }
+ }
+}
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 aeb7bdf9beb..7af485014df 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
@@ -32,6 +32,7 @@ import org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider;
import org.apache.doris.connector.spi.scan.ConnectorScanProfile;
import org.apache.doris.connector.spi.scan.ConnectorScanRange;
import org.apache.doris.connector.spi.scan.ConnectorScanRequest;
+import org.apache.doris.connector.spi.scan.ConnectorSplitSource;
import org.apache.doris.connector.spi.scan.ScanNodePropertyKeys;
import org.apache.doris.filesystem.FileSystem;
import org.apache.doris.thrift.TFileCompressType;
@@ -45,11 +46,13 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
/**
@@ -86,6 +89,226 @@ 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 millionFilesOnlyCreateTheConsumedSplitPrefix() throws Exception {
+ AtomicInteger normalizedFiles = new AtomicInteger();
+ AtomicInteger listings = new AtomicInteger();
+ HiveFileListingCache cache = new HiveFileListingCache(HiveTestProperties.minimal(), (location, fs) -> {
+ listings.incrementAndGet();
+ return Collections.nCopies(1000000, new HiveFileStatus(location + "/data.parquet", 10, 1));
+ });
+ FakeConnectorContext context = new FakeConnectorContext() {
+ @Override
+ public String normalizeStorageUri(String uri) {
+ normalizedFiles.incrementAndGet();
+ return uri;
+ }
+ };
+ HiveScanPlanProvider provider = new HiveScanPlanProvider(null, HiveTestProperties.minimal(), context,
+ new HiveReadTransactionManager(), cache);
+ HiveTableHandle handle = parquetHandle(Collections.emptyList());
+ try (ConnectorSplitSource source = provider.streamSplits(new FakeSession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1).build())) {
+ Assertions.assertEquals(0, listings.get());
+ for (int i = 0; i < 128; i++) {
+ Assertions.assertTrue(source.hasNext());
+ Assertions.assertEquals(10, source.next().getLength());
+ }
+ Assertions.assertEquals(128, normalizedFiles.get());
+ Assertions.assertEquals(1, listings.get());
+ source.close();
+ Assertions.assertFalse(source.hasNext());
+ Assertions.assertEquals(128, normalizedFiles.get());
+ }
+ }
+
+ @Test
+ public void hugeSingleFileDoesNotMaterializeItsSplitList() throws Exception {
+ HiveScanPlanProvider provider = provider(null, (location, fs) -> Collections.singletonList(
+ new HiveFileStatus(location + "/large.parquet", Long.MAX_VALUE, 1)));
+ FakeSession session = new FakeSession(Collections.singletonMap("file_split_size", "1"));
+ try (ConnectorSplitSource source = provider.streamSplits(session,
+ ConnectorScanRequest.builder(parquetHandle(Collections.emptyList()), Collections.emptyList())
+ .filter(Optional.empty()).limit(-1).build())) {
+ for (int i = 0; i < 3; i++) {
+ Assertions.assertTrue(source.hasNext());
+ HiveScanRange range = (HiveScanRange) source.next();
+ Assertions.assertEquals(i, range.getStart());
+ Assertions.assertEquals(1, range.getLength());
+ }
+ }
+ }
+
+ @Test
+ public void streamingAndSynchronousSplitsMatchAcrossPartitionsAndFileBoundaries() throws Exception {
+ HiveScanPlanProvider provider = provider(new FakeHmsClient(), (location, fs) -> Arrays.asList(
+ new HiveFileStatus(location + "/empty.parquet", 0, 1),
+ new HiveFileStatus(location + "/small.parquet", 3, 2),
+ new HiveFileStatus(location + "/large.parquet", 35, 3)));
+ HiveTableHandle handle = parquetHandle(PART_KEYS).toBuilder().prunedPartitions(Arrays.asList(
+ new HmsPartitionInfo(Arrays.asList("2024", "01"), "year=2024/month=01",
+ null, null, null, Collections.emptyMap()),
+ new HmsPartitionInfo(Arrays.asList("2024", "02"), "year=2024/month=02",
+ null, null, null, Collections.emptyMap()))).build();
+ FakeSession session = new FakeSession(Collections.singletonMap("file_split_size", "8"));
+ List<ConnectorScanRange> synchronous = provider.planScan(session,
+ ConnectorScanRequest.builder(handle, Collections.emptyList()).build());
+ List<ConnectorScanRange> streaming = new ArrayList<>();
+ try (ConnectorSplitSource source = provider.streamSplits(session,
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1).build())) {
+ while (source.hasNext()) {
+ streaming.add(source.next());
+ }
+ }
+ long[] starts = {0, 0, 8, 16, 24, 32};
+ long[] lengths = {3, 8, 8, 8, 8, 3};
+ for (List<ConnectorScanRange> ranges : Arrays.asList(synchronous, streaming)) {
+ Assertions.assertEquals(12, ranges.size());
+ for (int partition = 0; partition < 2; partition++) {
+ String month = partition == 0 ? "01" : "02";
+ Map<String, String> partitionValues = new HashMap<>();
+ partitionValues.put("year", "2024");
+ partitionValues.put("month", month);
+ for (int i = 0; i < starts.length; i++) {
+ HiveScanRange range = (HiveScanRange) ranges.get(partition * starts.length + i);
+ String fileName = i == 0 ? "small.parquet" : "large.parquet";
+ Assertions.assertEquals("year=2024/month=" + month + "/" + fileName, range.getPath().get());
+ Assertions.assertEquals(starts[i], range.getStart());
+ Assertions.assertEquals(lengths[i], range.getLength());
+ Assertions.assertEquals(partitionValues, range.getPartitionValues());
+ }
+ }
+ }
+ }
+
+ @Test
+ public void closingFirstPartitionDoesNotListRemainingPartitions() throws Exception {
+ CountingLister lister = new CountingLister();
+ HiveScanPlanProvider provider = provider(new FakeHmsClient(), lister);
+ HiveTableHandle handle = parquetHandle(PART_KEYS).toBuilder().prunedPartitions(Arrays.asList(
+ part("year=2024/month=01"), part("year=2024/month=02"))).build();
+ try (ConnectorSplitSource source = provider.streamSplits(new FakeSession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1).build())) {
+ Assertions.assertEquals(0, lister.totalCalls);
+ Assertions.assertTrue(source.hasNext());
+ source.next();
+ Assertions.assertEquals(1, lister.totalCalls);
+ source.close();
+ Assertions.assertFalse(source.hasNext());
+ Assertions.assertNull(lister.callsPerLocation.get("year=2024/month=02"));
+ }
+ }
+
+ @Test
+ public void streamingUsesEnginePartitionSelectionWithoutListingOtherPartitions() throws Exception {
+ CountingLister lister = new CountingLister();
+ // listPartitionNames fails if called: the engine already supplies the selected names.
+ HiveScanPlanProvider provider = provider(new FakeHmsClient(), lister);
+ HiveTableHandle handle = parquetHandle(PART_KEYS);
+ ConnectorScanRequest request = ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .requiredPartitions(Collections.singletonList("year=2024/month=02")).build();
+ try (ConnectorSplitSource source = provider.streamSplits(new FakeSession(), request)) {
+ Assertions.assertTrue(source.hasNext());
+ HiveScanRange range = (HiveScanRange) source.next();
+ Assertions.assertEquals("year=2024/month=02/000000_0", range.getPath().get());
+ Assertions.assertFalse(source.hasNext());
+ }
+ Assertions.assertEquals(1, lister.totalCalls);
+ Assertions.assertEquals(Collections.singleton("year=2024/month=02"), lister.callsPerLocation.keySet());
+ }
+
+ @Test
+ public void partitionBatchAndStreamingMatchAcrossFormatsAndListingFailures() throws Exception {
+ List<String> batch = Arrays.asList("year=2024/month=01", "year=2024/month=02", "year=2024/month=03");
+ for (boolean lzo : Arrays.asList(false, true)) {
+ Map<String, Integer> listings = new HashMap<>();
+ HiveScanPlanProvider provider = provider(new FakeHmsClient(), (location, fs) -> {
+ listings.merge(location, 1, Integer::sum);
+ if (location.equals(batch.get(0))) {
+ return Collections.emptyList();
+ }
+ if (location.equals(batch.get(1))) {
+ throw new HiveDirectoryListingException("unreadable partition", new IOException("retry"));
+ }
+ String extension = lzo ? ".lzo" : ".parquet";
+ List<HiveFileStatus> files = new ArrayList<>(Arrays.asList(
+ new HiveFileStatus(location + "/empty" + extension, 0, 1),
+ new HiveFileStatus(location + "/small" + extension, 3, 2),
+ new HiveFileStatus(location + "/large" + extension, 35, 3)));
+ if (lzo) {
+ files.add(new HiveFileStatus(location + "/large.lzo.index", 20, 4));
+ files.add(new HiveFileStatus(location + "/other.txt", 20, 4));
+ }
+ return files;
+ });
+ HiveTableHandle handle = parquetHandle(PART_KEYS).toBuilder()
+ .inputFormat(lzo ? LZO_TEXT_INPUT_FORMAT : PARQUET_INPUT_FORMAT)
+ .serializationLib(lzo ? LAZY_SIMPLE_SERDE : PARQUET_SERDE)
+ .prunedPartitions(Collections.singletonList(part("year=2023/month=12"))).build();
+ Map<String, String> settings = new HashMap<>();
+ settings.put("file_split_size", "8");
+ settings.put("num_partitions_in_batch_mode", "1");
+ FakeSession session = new FakeSession(settings);
+ ConnectorScanRequest request = ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .requiredPartitions(Collections.singletonList("year=2023/month=12")).build();
+
+ Assertions.assertTrue(provider.planScanForPartitionBatch(session, request,
+ Collections.emptyList()).isEmpty());
+ Assertions.assertTrue(listings.isEmpty(), "An empty batch must not enumerate the handle's partitions");
+ List<ConnectorScanRange> synchronous = provider.planScanForPartitionBatch(session, request, batch);
+ List<ConnectorScanRange> streaming = new ArrayList<>();
+ try (ConnectorSplitSource source = provider.streamSplits(session, request.withRequiredPartitions(batch))) {
+ while (source.hasNext()) {
+ streaming.add(source.next());
+ }
+ }
+ long[] starts = lzo ? new long[] {0, 0} : new long[] {0, 0, 8, 16, 24, 32};
+ long[] lengths = lzo ? new long[] {3, 35} : new long[] {3, 8, 8, 8, 8, 3};
+ for (List<ConnectorScanRange> ranges : Arrays.asList(synchronous, streaming)) {
+ Assertions.assertEquals(starts.length, ranges.size());
+ for (int i = 0; i < starts.length; i++) {
+ ConnectorScanRange range = ranges.get(i);
+ String file = i == 0 ? "/small" : "/large";
+ Assertions.assertEquals(batch.get(2) + file + (lzo ? ".lzo" : ".parquet"),
+ range.getPath().get());
+ Assertions.assertEquals(starts[i], range.getStart());
+ Assertions.assertEquals(lengths[i], range.getLength());
+ }
+ }
+ Assertions.assertEquals(new HashSet<>(batch), listings.keySet());
+ Assertions.assertEquals(2, (int) listings.get(batch.get(1)),
+ "A tolerated directory failure must not be cached by either enumeration path");
+ }
+ }
+
+ @Test
+ public void streamingDecisionUsesCachedCountWithoutListing() {
+ AtomicInteger listings = new AtomicInteger();
+ HiveFileListingCache cache = new HiveFileListingCache(HiveTestProperties.minimal(), (location, fs) -> {
+ listings.incrementAndGet();
+ return Collections.nCopies(3, new HiveFileStatus(location + "/data.parquet", 10, 1));
+ });
+ HiveScanPlanProvider provider = new HiveScanPlanProvider(null, HiveTestProperties.minimal(),
+ new FakeConnectorContext(), new HiveReadTransactionManager(), cache);
+ HiveTableHandle handle = parquetHandle(Collections.emptyList());
+ FakeSession session = new FakeSession(Collections.singletonMap("num_files_in_batch_mode", "2"));
+ Assertions.assertTrue(provider.streamingSplitEstimate(session, handle, Optional.empty(), false) >= 0);
+ Assertions.assertEquals(0, listings.get());
+ cache.listDataFiles("db", "t", handle.getLocation(), null);
+ Assertions.assertEquals(3, provider.streamingSplitEstimate(session, handle, Optional.empty(), false));
+ Assertions.assertEquals(-1, provider.streamingSplitEstimate(new FakeSession(), handle, Optional.empty(), false));
+ Assertions.assertEquals(1, listings.get());
+ }
+
+ private static HiveTableHandle parquetHandle(List<String> partitionKeys) {
+ return new HiveTableHandle.Builder("db", "t", HiveTableType.HIVE)
+ .inputFormat(PARQUET_INPUT_FORMAT).serializationLib(PARQUET_SERDE)
+ .location("s3://bucket/table").partitionKeyNames(partitionKeys).build();
+ }
+
@Test
public void scanReuseNamespaceUsesConnectorType() {
String prefix = new HiveConnectorProvider().getType() + ".";
@@ -970,6 +1193,16 @@ public class HiveScanBatchModeTest {
/** Minimal {@link ConnectorSession} (no split-size override, empty session properties). */
private static final class FakeSession implements ConnectorSession {
+ private final Map<String, String> properties;
+
+ FakeSession() {
+ this(Collections.emptyMap());
+ }
+
+ FakeSession(Map<String, String> properties) {
+ this.properties = properties;
+ }
+
private final Map<String, Object> statementMemos = new HashMap<>();
private final ConnectorStatementScope statementScope = new ConnectorStatementScope() {
@Override
@@ -1011,7 +1244,17 @@ public class HiveScanBatchModeTest {
@Override
public <T> T getProperty(String name, Class<T> type) {
- return null;
+ String value = properties.get(name);
+ if (value == null) {
+ return null;
+ }
+ if (type == Integer.class) {
+ return type.cast(Integer.valueOf(value));
+ }
+ if (type == Boolean.class) {
+ return type.cast(Boolean.valueOf(value));
+ }
+ return type.cast(value);
}
@Override
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 c6b6dc93898..f9125dfe66b 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
@@ -43,7 +43,7 @@ import java.util.function.Function;
/**
* A caching {@link HmsClient} decorator: it wraps another {@code HmsClient} (in production the pooled
* {@link ThriftHmsClient}) and serves the three scan-hot-path read methods from a bounded, TTL-expiring
- * cache, delegating every other method verbatim.
+ * cache. Existence checks reuse cached table metadata; other methods delegate verbatim.
*
* <p><b>Why this exists.</b> Without this decorator the hive connector would cache nothing — {@code getTable},
* {@code listPartitionNames} and {@code getPartitions} would be fresh Thrift RPCs on every scan. Legacy fe-core
@@ -74,6 +74,8 @@ import java.util.function.Function;
* column-stats fast path, off the scan hot path, so it caches at low priority but on the same
* machinery as the rest.</li>
* </ul>
+ * {@code tableExists} also checks the table cache: a cached table proves existence, while a miss
+ * delegates without caching the boolean result. It shares the table entry's TTL and invalidation.
*
* <p><b>Pass-through.</b> Every other read, plus every write / DDL / ACID method, is passed straight
* through to the delegate. Cache invalidation belongs to the connector's shared {@link CatalogMetaCache}
@@ -190,6 +192,14 @@ public class CachingHmsClient implements HmsClient {
key -> delegate.getTable(key.dbName, key.tableName));
}
+ @Override
+ public boolean tableExists(String dbName, String tableName) {
+ // Only successfully loaded metadata proves existence. Do not cache absence: CTAS resolves
+ // its target again after creation. A disabled or invalidated table cache falls through to HMS.
+ return tableCache.getIfPresent(new TableKey(dbName, tableName)) != null
+ || delegate.tableExists(dbName, tableName);
+ }
+
@Override
public HmsTableInfo getTableFresh(String dbName, String tableName) {
// Fresh (cache-bypassing) table read for SHOW CREATE TABLE, which must reflect the latest remote schema
@@ -531,11 +541,6 @@ public class CachingHmsClient implements HmsClient {
return delegate.listTables(dbName);
}
- @Override
- public boolean tableExists(String dbName, String tableName) {
- return delegate.tableExists(dbName, tableName);
- }
-
@Override
public Map<String, String> getDefaultColumnValues(String dbName, String tableName) {
return delegate.getDefaultColumnValues(dbName, tableName);
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 cae43cfe58f..1816b19ca0f 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
@@ -162,6 +162,94 @@ public class CachingHmsClientTest {
Assertions.assertEquals(2, delegate.getColumnStatsCalls);
}
+ // ---- tableExists shares positive table metadata ----
+
+ @Test
+ public void tableExistsUsesLoadedMetadataAcrossLookups() {
+ RecordingHmsClient delegate = new RecordingHmsClient();
+ CachingHmsClient cache = cache(delegate, Collections.emptyMap());
+
+ Assertions.assertTrue(cache.tableExists("db", "t"));
+ HmsTableInfo table = cache.getTable("db", "t");
+ // A later statement's existence probe and getTable share the same catalog-owned entry.
+ Assertions.assertTrue(cache.tableExists("db", "t"));
+ Assertions.assertTrue(cache.tableExists("db", "t"));
+ Assertions.assertSame(table, cache.getTable("db", "t"));
+ Assertions.assertEquals(1, delegate.tableExistsCalls);
+ Assertions.assertEquals(1, delegate.getTableCalls);
+ }
+
+ @Test
+ public void tableExistsCacheIsScopedByDatabaseAndTable() {
+ RecordingHmsClient delegate = new RecordingHmsClient();
+ CachingHmsClient cache = cache(delegate, Collections.emptyMap());
+ cache.getTable("db", "t");
+ delegate.tableExistsResult = false;
+
+ Assertions.assertTrue(cache.tableExists("db", "t"));
+ Assertions.assertFalse(cache.tableExists("other_db", "t"));
+ Assertions.assertFalse(cache.tableExists("db", "other_t"));
+ Assertions.assertEquals(2, delegate.tableExistsCalls);
+ }
+
+ @Test
+ public void tableExistsDoesNotRetainBooleanResults() {
+ RecordingHmsClient delegate = new RecordingHmsClient();
+ CachingHmsClient cache = cache(delegate, Collections.emptyMap());
+ delegate.tableExistsResult = false;
+ Assertions.assertFalse(cache.tableExists("db", "t"));
+
+ // A missing CTAS target can be created and resolved again without a negative cache entry.
+ delegate.tableExistsResult = true;
+ Assertions.assertTrue(cache.tableExists("db", "t"));
+ delegate.tableExistsResult = false;
+ Assertions.assertFalse(cache.tableExists("db", "t"));
+ Assertions.assertEquals(3, delegate.tableExistsCalls);
+ Assertions.assertEquals(0, delegate.getTableCalls);
+ }
+
+ @Test
+ public void tableExistsHonorsDisabledTableCache() {
+ for (Map<String, String> properties : Arrays.asList(
+ props("meta.cache.hive.table.enable", "false"),
+ props("meta.cache.hive.table.ttl-second", "0"),
+ props("schema.cache.ttl-second", "0"))) {
+ RecordingHmsClient delegate = new RecordingHmsClient();
+ CachingHmsClient cache = cache(delegate, properties);
+ cache.getTable("db", "t");
+ delegate.tableExistsResult = false;
+ Assertions.assertFalse(cache.tableExists("db", "t"));
+ Assertions.assertEquals(1, delegate.tableExistsCalls);
+ }
+ }
+
+ @Test
+ public void tableExistsHonorsTableInvalidation() {
+ RecordingHmsClient delegate = new RecordingHmsClient();
+ CatalogMetaCache owner = CatalogMetaCache.unmanaged();
+ CachingHmsClient cache = new CachingHmsClient(owner, delegate, Collections.emptyMap());
+ cache.getTable("db", "t");
+ delegate.tableExistsResult = false;
+ Assertions.assertTrue(cache.tableExists("db", "t"));
+
+ owner.invalidateTable("db", "t");
+ Assertions.assertFalse(cache.tableExists("db", "t"));
+ Assertions.assertEquals(1, delegate.tableExistsCalls);
+ }
+
+ @Test
+ public void tableExistsPropagatesFailuresWithoutCachingThem() {
+ RecordingHmsClient delegate = new RecordingHmsClient();
+ CachingHmsClient cache = cache(delegate, Collections.emptyMap());
+ delegate.tableExistsError = new HmsClientException("HMS unavailable");
+ Assertions.assertSame(delegate.tableExistsError, Assertions.assertThrows(HmsClientException.class,
+ () -> cache.tableExists("db", "t")));
+
+ delegate.tableExistsError = null;
+ Assertions.assertTrue(cache.tableExists("db", "t"));
+ Assertions.assertEquals(2, delegate.tableExistsCalls);
+ }
+
// ---- listPartitionNames ----
@Test
@@ -1365,6 +1453,9 @@ public class CachingHmsClientTest {
*/
private static class RecordingHmsClient implements HmsClient {
int getTableCalls;
+ int tableExistsCalls;
+ boolean tableExistsResult = true;
+ HmsClientException tableExistsError;
int listPartitionNamesCalls;
int getPartitionsCalls;
int getColumnStatsCalls;
@@ -1497,7 +1588,11 @@ public class CachingHmsClientTest {
@Override
public boolean tableExists(String dbName, String tableName) {
- return false;
+ tableExistsCalls++;
+ if (tableExistsError != null) {
+ throw tableExistsError;
+ }
+ return tableExistsResult;
}
@Override
diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java
index 73beceb47ed..515d67b5f56 100644
--- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java
+++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java
@@ -571,8 +571,10 @@ public class IcebergScanPlanProvider implements ConnectorScanPlanProvider {
* method must independently preserve an explicitly empty pinned snapshot.
*/
@Override
- public ConnectorSplitSource streamSplits(ConnectorSession session, ConnectorTableHandle handle,
- List<ConnectorColumnHandle> columns, Optional<ConnectorExpression> filter, long limit) {
+ public ConnectorSplitSource streamSplits(ConnectorSession session, ConnectorScanRequest request) {
+ ConnectorTableHandle handle = request.getTableHandle();
+ List<ConnectorColumnHandle> columns = request.getColumns();
+ Optional<ConnectorExpression> filter = request.getFilter();
IcebergTableHandle iceHandle = (IcebergTableHandle) handle;
validateMetadataColumnReader(session, columns);
if (iceHandle.isResolvedEmptySnapshot()) {
diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScanPlanProviderTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScanPlanProviderTest.java
index 305fbf81347..37135377acc 100644
--- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScanPlanProviderTest.java
+++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScanPlanProviderTest.java
@@ -452,8 +452,9 @@ public class IcebergScanPlanProviderTest {
estimate.getMessage());
DorisConnectorException lazy = Assertions.assertThrows(DorisConnectorException.class, () -> {
- try (ConnectorSplitSource source = provider.streamSplits(
- emptySession(), handle, Collections.emptyList(), Optional.empty(), -1L)) {
+ try (ConnectorSplitSource source = provider.streamSplits(emptySession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build())) {
source.hasNext();
}
});
@@ -1121,9 +1122,9 @@ public class IcebergScanPlanProviderTest {
IcebergScanPlanProvider provider = new IcebergScanPlanProvider(
IcebergCatalogProperties.of(Collections.emptyMap()), opsReturning(table));
List<ConnectorScanRange> ranges = new ArrayList<>();
- try (ConnectorSplitSource source = provider.streamSplits(
- new FakeScanSession("UTC", Collections.emptyMap()),
- new IcebergTableHandle("db1", "pt"), Collections.emptyList(), Optional.empty(), -1L)) {
+ try (ConnectorSplitSource source = provider.streamSplits(new FakeScanSession("UTC", Collections.emptyMap()),
+ ConnectorScanRequest.builder(new IcebergTableHandle("db1", "pt"), Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build())) {
while (source.hasNext()) {
ranges.add(source.next());
}
@@ -1456,8 +1457,8 @@ public class IcebergScanPlanProviderTest {
DorisConnectorException ex = Assertions.assertThrows(DorisConnectorException.class,
() -> provider.streamSplits(session,
- new IcebergTableHandle("db1", "stream_metadata"), columns,
- Optional.empty(), -1L));
+ ConnectorScanRequest.builder(new IcebergTableHandle("db1", "stream_metadata"), columns)
+ .filter(Optional.empty()).limit(-1L).build()));
Assertions.assertEquals(
"Iceberg metadata columns are only supported by FileScannerV2 native Parquet/ORC reader; "
+ "actual reader is JNI",
@@ -1474,8 +1475,8 @@ public class IcebergScanPlanProviderTest {
DorisConnectorException ex = Assertions.assertThrows(DorisConnectorException.class,
() -> provider.streamSplits(session,
- new IcebergTableHandle("db1", "stream_metadata_v2"), columns,
- Optional.empty(), -1L));
+ ConnectorScanRequest.builder(new IcebergTableHandle("db1", "stream_metadata_v2"), columns)
+ .filter(Optional.empty()).limit(-1L).build()));
Assertions.assertEquals(
"Iceberg metadata columns require FileScannerV2 native Parquet/ORC reader",
ex.getMessage());
@@ -1845,8 +1846,9 @@ public class IcebergScanPlanProviderTest {
Assertions.assertEquals(1L,
provider.streamingSplitEstimate(session, handle, Optional.empty(), false),
"schema-carrier planning must not retain tasks or disable streaming batch dispatch");
- try (ConnectorSplitSource source = provider.streamSplits(
- session, handle, Collections.emptyList(), Optional.empty(), -1L)) {
+ try (ConnectorSplitSource source = provider.streamSplits(session,
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build())) {
Assertions.assertTrue(source.hasNext(),
"the first split must be dispatchable without draining all remaining byte splits");
Assertions.assertEquals(0L, source.next().getStart());
@@ -2469,8 +2471,9 @@ public class IcebergScanPlanProviderTest {
provider.streamingSplitEstimate(batchSession(2, true), handle, filter, false),
"the nonmatching manifest must not enable streaming at a two-file threshold");
}
- List<ConnectorScanRange> streamed = drain(provider.streamSplits(
- batchSession(1, true), handle, Collections.emptyList(), filter, -1L));
+ List<ConnectorScanRange> streamed = drain(provider.streamSplits(batchSession(1, true),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(filter).limit(-1L).build()));
Assertions.assertEquals(sortedPaths(ranges), sortedPaths(streamed));
}
}
@@ -2543,8 +2546,8 @@ public class IcebergScanPlanProviderTest {
for (int i = 0; i < 2; i++) {
boolean scanHistorical = (i == 0) == historicalFirst;
List<ConnectorScanRange> ranges = drain(provider.streamSplits(batchSession(1, true),
- scanHistorical ? historical : current, Collections.emptyList(),
- Optional.of(scanHistorical ? intFilter : bigintFilter), -1L));
+ ConnectorScanRequest.builder(scanHistorical ? historical : current, Collections.emptyList())
+ .filter(Optional.of(scanHistorical ? intFilter : bigintFilter)).limit(-1L).build()));
Assertions.assertEquals(Collections.singletonList(location + "old.parquet"), sortedPaths(ranges),
(scanHistorical ? "historical" : "current") + " scan, historical first=" + historicalFirst);
}
@@ -2560,8 +2563,9 @@ public class IcebergScanPlanProviderTest {
Optional<ConnectorExpression> filter = Optional.of(eqInt("x", 1));
Assertions.assertEquals(1L, provider.streamingSplitEstimate(session, handle, filter, false));
- List<ConnectorScanRange> ranges = drain(provider.streamSplits(
- session, handle, Collections.emptyList(), filter, -1L));
+ List<ConnectorScanRange> ranges = drain(provider.streamSplits(session,
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(filter).limit(-1L).build()));
Assertions.assertEquals(1, ranges.size());
Assertions.assertTrue(ranges.get(0).getPath().get().endsWith("old.parquet"));
}
@@ -2605,7 +2609,8 @@ public class IcebergScanPlanProviderTest {
new IcebergTableHandle("db1", "t1"), Optional.empty(), false),
"the pre-pin batch decision must observe the concurrent first append");
List<ConnectorScanRange> ranges = drain(provider.streamSplits(emptySession(),
- emptyRead, Collections.emptyList(), Optional.empty(), -1L));
+ ConnectorScanRequest.builder(emptyRead, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
Assertions.assertTrue(ranges.isEmpty(),
"the post-pin streaming path must preserve the explicitly empty read boundary");
}
@@ -3126,7 +3131,8 @@ public class IcebergScanPlanProviderTest {
// streamed counterpart of planScan's eager enumeration.
IcebergScanPlanProvider provider = providerOver(threeFileTable());
List<ConnectorScanRange> ranges = drain(provider.streamSplits(emptySession(),
- new IcebergTableHandle("db1", "t1"), Collections.emptyList(), Optional.empty(), -1L));
+ ConnectorScanRequest.builder(new IcebergTableHandle("db1", "t1"), Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
Set<String> paths = ranges.stream()
.map(r -> ((IcebergScanRange) r).getOriginalPath()).collect(ImmutableSet.toImmutableSet());
Assertions.assertEquals(
@@ -3144,7 +3150,8 @@ public class IcebergScanPlanProviderTest {
IcebergTableHandle scoped = new IcebergTableHandle("db1", "t1").withRewriteFileScope(
ImmutableSet.of("s3://b/db/t1/f1.parquet", "s3://b/db/t1/f3.parquet"));
List<ConnectorScanRange> ranges = drain(provider.streamSplits(emptySession(),
- scoped, Collections.emptyList(), Optional.empty(), -1L));
+ ConnectorScanRequest.builder(scoped, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
Set<String> paths = ranges.stream()
.map(r -> ((IcebergScanRange) r).getOriginalPath()).collect(ImmutableSet.toImmutableSet());
Assertions.assertEquals(
@@ -3158,7 +3165,8 @@ public class IcebergScanPlanProviderTest {
// guard in next() -> NPE/wrong instead of NoSuchElementException -> red.
IcebergScanPlanProvider provider = providerOver(threeFileTable());
ConnectorSplitSource source = provider.streamSplits(new FakeScanSession("UTC", Collections.emptyMap()),
- new IcebergTableHandle("db1", "t1"), Collections.emptyList(), Optional.empty(), -1L);
+ ConnectorScanRequest.builder(new IcebergTableHandle("db1", "t1"), Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build());
while (source.hasNext()) {
source.next();
}
@@ -3173,7 +3181,8 @@ public class IcebergScanPlanProviderTest {
// underlying planFiles() iterable. MUTATION: dropping the `iterator != null` guard in close() -> NPE -> red.
IcebergScanPlanProvider provider = providerOver(threeFileTable());
ConnectorSplitSource source = provider.streamSplits(emptySession(),
- new IcebergTableHandle("db1", "t1"), Collections.emptyList(), Optional.empty(), -1L);
+ ConnectorScanRequest.builder(new IcebergTableHandle("db1", "t1"), Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build());
Assertions.assertDoesNotThrow(source::close);
}
@@ -3822,12 +3831,16 @@ public class IcebergScanPlanProviderTest {
// Cache OFF: the streaming SDK planFiles() path (the pre-PERF-04 behavior).
List<ConnectorScanRange> sdk = drain(new IcebergScanPlanProvider(IcebergCatalogProperties.of(Collections.emptyMap()), opsReturning(table))
- .streamSplits(emptySession(), handle, Collections.emptyList(), Optional.empty(), -1L));
+ .streamSplits(emptySession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
// Cache ON: streaming must now read manifests THROUGH the cache and yield the SAME files (PERF-04 C17).
// MUTATION: streaming still bypasses the cache (scan.planFiles()) -> cache stays empty -> red.
IcebergManifestCache cache = new IcebergManifestCache();
List<ConnectorScanRange> cached = drain(manifestProvider(manifestCacheProps(), table, cache)
- .streamSplits(emptySession(), handle, Collections.emptyList(), Optional.empty(), -1L));
+ .streamSplits(emptySession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
Assertions.assertEquals(sortedPaths(sdk), sortedPaths(cached));
Assertions.assertEquals(3, cached.size());
@@ -3846,10 +3859,14 @@ public class IcebergScanPlanProviderTest {
Optional<ConnectorExpression> wherePeq1 = Optional.of(eqInt("p", 1));
List<ConnectorScanRange> sdk = drain(new IcebergScanPlanProvider(IcebergCatalogProperties.of(Collections.emptyMap()), opsReturning(table))
- .streamSplits(emptySession(), handle, Collections.emptyList(), wherePeq1, -1L));
+ .streamSplits(emptySession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(wherePeq1).limit(-1L).build()));
IcebergManifestCache cache = new IcebergManifestCache();
List<ConnectorScanRange> cached = drain(manifestProvider(manifestCacheProps(), table, cache)
- .streamSplits(emptySession(), handle, Collections.emptyList(), wherePeq1, -1L));
+ .streamSplits(emptySession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(wherePeq1).limit(-1L).build()));
// Partition prune (ManifestEvaluator + residual) keeps only p=1 in BOTH paths. MUTATION: the lazy iterator
// dropping the residual/metrics prune -> p=2 leaks in -> sizes differ -> red.
@@ -3875,7 +3892,9 @@ public class IcebergScanPlanProviderTest {
IcebergManifestCache cache = new IcebergManifestCache();
List<ConnectorScanRange> cached = drain(manifestProvider(manifestCacheProps(), table, cache)
- .streamSplits(emptySession(), handle, Collections.emptyList(), Optional.empty(), -1L));
+ .streamSplits(emptySession(),
+ ConnectorScanRequest.builder(handle, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
Assertions.assertEquals(3, cached.size(), "the lazy iterator must flat-map across all data manifests");
}
@@ -4054,8 +4073,9 @@ public class IcebergScanPlanProviderTest {
IcebergManifestCache cache = new IcebergManifestCache();
List<ConnectorScanRange> ranges = drain(manifestProvider(manifestCacheProps(), reloaded, cache)
- .streamSplits(emptySession(), new IcebergTableHandle("db1", "dropped_cache_key"),
- Collections.emptyList(), Optional.empty(), -1L));
+ .streamSplits(emptySession(),
+ ConnectorScanRequest.builder(new IcebergTableHandle("db1", "dropped_cache_key"), Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
Assertions.assertEquals(1, ranges.size());
Assertions.assertEquals(1, deleteCount(ranges.get(0)));
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 8a52b83f74c..0d2203d44a3 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
@@ -26,8 +26,10 @@ import org.apache.doris.thrift.TFileScanRangeParams;
import org.apache.doris.thrift.TTableFormatFileDesc;
import java.util.Collections;
+import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.OptionalLong;
import java.util.Set;
@@ -206,9 +208,11 @@ public interface ConnectorScanPlanProvider {
* <p>When {@code true}, a partition-aware ScanNode (e.g. {@code PluginDrivenScanNode}) may
* enter batch mode: instead of enumerating all splits synchronously via {@link #planScan},
* it slices the pruned partitions into batches and calls {@link #planScanForPartitionBatch}
- * per batch on a background executor, streaming splits as they are produced (mirrors legacy
- * {@code MaxComputeScanNode.startSplit}). The default is {@code false}, so connectors stay on
- * the synchronous {@code planScan} path unless they opt in.</p>
+ * per batch on a background executor (mirrors legacy {@code MaxComputeScanNode.startSplit}).
+ * Both {@link #planScanForPartitionBatch} and {@link #streamSplits} MUST honor each request's
+ * partition subset. The default streaming adapter defers the existing eager batch hook; a
+ * connector can override it to enumerate each batch lazily. The default {@code false} disables
+ * partition batching; file-count streaming is decided independently by {@link #streamingSplitEstimate}.</p>
*
* @param session the current session
* @param handle the table handle
@@ -320,6 +324,55 @@ public interface ConnectorScanPlanProvider {
return planScan(session, request.withRequiredPartitions(partitionBatch));
}
+ /**
+ * Opens a lazy source for a scan request. The engine advances and closes each source serially,
+ * but may drive sources for disjoint partition batches concurrently when {@link #supportsBatchScan}
+ * is true. Such sources MUST restrict enumeration to the request's partition subset.
+ *
+ * <p>For partition-batching connectors, the default defers {@link #planScanForPartitionBatch}
+ * until consumption, preserving existing batch implementations. Connectors that can stream
+ * within a partition override this to avoid materializing the batch's ranges. Streaming-only
+ * connectors must override this method and keep their whole-scan semantics.</p>
+ *
+ * @param session the current session
+ * @param request the scan scope, columns, filter and limit; snapshot/time-travel is already pinned
+ * @return a lazy, closeable source of scan ranges
+ */
+ default ConnectorSplitSource streamSplits(ConnectorSession session, ConnectorScanRequest request) {
+ if (!supportsBatchScan(session, request.getTableHandle())) {
+ throw new UnsupportedOperationException("Connector did not implement streaming split generation");
+ }
+ return new ConnectorSplitSource() {
+ private Iterator<ConnectorScanRange> ranges;
+ private boolean closed;
+
+ @Override
+ public boolean hasNext() {
+ if (closed) {
+ return false;
+ }
+ if (ranges == null) {
+ ranges = planScanForPartitionBatch(session, request, request.getRequiredPartitions()).iterator();
+ }
+ return ranges.hasNext();
+ }
+
+ @Override
+ public ConnectorScanRange next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return ranges.next();
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ ranges = Collections.emptyIterator();
+ }
+ };
+ }
+
/**
* Decides whether this scan should use streaming (lazy) split generation and, if so, estimates
* the number of splits it will produce. This is the file-count counterpart of the partition-count
@@ -332,7 +385,7 @@ public interface ConnectorScanPlanProvider {
* <p>When this returns a non-negative value, the engine drives streaming split generation via
* {@link #streamSplits} (pulling ranges with backpressure) instead of the synchronous
* {@link #planScan}; the returned estimate doubles as the BE concurrency hint. A negative value
- * keeps the scan on the synchronous {@code planScan} path (the default).</p>
+ * disables file-count streaming; partition batching may still use {@link #streamSplits}.</p>
*
* <p>The decision MUST be cheap (e.g. manifest metadata counts), NOT a full split enumeration —
* the heavy enumeration is deferred to {@link #streamSplits}.</p>
@@ -363,33 +416,6 @@ public interface ConnectorScanPlanProvider {
return false;
}
- /**
- * Builds a lazy {@link ConnectorSplitSource} for streaming split generation. Called once, on a
- * background task, only when {@link #streamingSplitEstimate} returned a non-negative value. The
- * engine pulls ranges from the source with backpressure and pumps them into the split queue
- * (mirrors legacy {@code IcebergScanNode.doStartSplit}).
- *
- * <p>The source MUST defer the heavy planning (e.g. {@code TableScan.planFiles()}) until ranges
- * are consumed, so FE heap usage stays bounded for very large scans. The default throws, so a
- * connector that returns a non-negative {@link #streamingSplitEstimate} must override this.</p>
- *
- * @param session the current session
- * @param handle the table handle (snapshot/time-travel already pinned by the engine)
- * @param columns the columns to read
- * @param filter an optional remaining filter expression
- * @param limit the maximum number of rows to return, or -1 for no limit
- * @return a lazy, closeable source of scan ranges
- */
- default ConnectorSplitSource streamSplits(
- ConnectorSession session,
- ConnectorTableHandle handle,
- List<ConnectorColumnHandle> columns,
- Optional<ConnectorExpression> filter,
- long limit) {
- throw new UnsupportedOperationException(
- "streamSplits requires streamingSplitEstimate(...) >= 0; connector did not implement it");
- }
-
/**
* Returns scan-node-level properties shared across all scan ranges.
*
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 ddb3227936c..1bf80d8d582 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
@@ -32,8 +32,8 @@ import java.io.Closeable;
*
* <p>Implementations MUST defer the heavy planning (e.g. iceberg {@code TableScan.planFiles()})
* until ranges are actually consumed, and MUST release the underlying resources in {@link #close()}.
- * Instances are single-pass and not thread-safe; the engine drives one source from a single
- * background task.</p>
+ * Instances are single-pass and not thread-safe; the engine advances and closes each source
+ * serially. Successive batches may run on different background workers.</p>
*/
public interface ConnectorSplitSource extends Closeable {
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 696ed4ecf7e..f0fa7ddbc41 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
@@ -88,8 +88,8 @@ public class ConnectorPluginSurfaceTest {
Assertions.assertNotNull(in, "missing connector plugin API version resource");
version.load(in);
}
- // Major 15 carries backend capability and aggregate-pushdown accessors on ConnectorScanRequest.
- Assertions.assertEquals("15.0", version.getProperty("api.version"));
+ // Major 17 makes the scoped request the canonical streaming entry point.
+ Assertions.assertEquals("17.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/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProviderBatchScanTest.java b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProviderBatchScanTest.java
index d2e8015b772..2f7f521c83c 100644
--- a/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProviderBatchScanTest.java
+++ b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/scan/ConnectorScanPlanProviderBatchScanTest.java
@@ -30,7 +30,9 @@ import org.junit.jupiter.api.Test;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
+import java.util.NoSuchElementException;
import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
/**
* FIX-BATCH-MODE-SPLIT (P4-T06e / NG-7) — guards the two additive SPI defaults on
@@ -60,6 +62,48 @@ public class ConnectorScanPlanProviderBatchScanTest {
private static final ConnectorTableHandle HANDLE = new ConnectorTableHandle() {
};
+ @Test
+ public void partitionBatchSourceDefersScopedPlanningAndReleasesItsRanges() throws Exception {
+ AtomicInteger calls = new AtomicInteger();
+ ConnectorScanRange range = () -> Collections.emptyMap();
+ ConnectorScanRequest request = ConnectorScanRequest.builder(HANDLE, Collections.emptyList())
+ .requiredPartitions(Arrays.asList("pt=1", "pt=2"))
+ .backendCapabilities(Collections.singleton("capability"))
+ .aggregatePushdown(true).build();
+ ConnectorScanPlanProvider provider = new ConnectorScanPlanProvider() {
+ @Override
+ public List<ConnectorScanRange> planScan(ConnectorSession session, ConnectorScanRequest scoped) {
+ calls.incrementAndGet();
+ Assertions.assertEquals(Collections.singletonList("pt=2"), scoped.getRequiredPartitions());
+ Assertions.assertTrue(scoped.isAggregatePushdown());
+ Assertions.assertEquals(request.getBackendCapabilities(), scoped.getBackendCapabilities());
+ return Collections.singletonList(range);
+ }
+
+ @Override
+ public boolean supportsBatchScan(ConnectorSession session, ConnectorTableHandle handle) {
+ return true;
+ }
+ };
+ ConnectorScanRequest scoped = request.withRequiredPartitions(Collections.singletonList("pt=2"));
+ try (ConnectorSplitSource source = provider.streamSplits(null, scoped)) {
+ Assertions.assertEquals(0, calls.get());
+ Assertions.assertTrue(source.hasNext());
+ Assertions.assertTrue(source.hasNext());
+ Assertions.assertSame(range, source.next());
+ Assertions.assertFalse(source.hasNext());
+ Assertions.assertThrows(NoSuchElementException.class, source::next);
+ Assertions.assertEquals(1, calls.get());
+ source.close();
+ Assertions.assertFalse(source.hasNext());
+ }
+ try (ConnectorSplitSource source = provider.streamSplits(null, scoped)) {
+ source.close();
+ Assertions.assertFalse(source.hasNext());
+ Assertions.assertEquals(1, calls.get(), "Closing an unused source must not plan its batch");
+ }
+ }
+
@Test
public void testSupportsBatchScanDefaultsFalse() {
// Default MUST be false: any connector that does not opt in stays on the synchronous path.
@@ -81,7 +125,9 @@ public class ConnectorScanPlanProviderBatchScanTest {
// non-negative streamingSplitEstimate without implementing streamSplits, which is a connector bug.
ConnectorScanPlanProvider provider = new RecordingProvider();
Assertions.assertThrows(UnsupportedOperationException.class,
- () -> provider.streamSplits(null, null, Collections.emptyList(), Optional.empty(), -1L));
+ () -> provider.streamSplits(null,
+ ConnectorScanRequest.builder(HANDLE, Collections.emptyList())
+ .filter(Optional.empty()).limit(-1L).build()));
}
@Test
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 5e3b47b67b9..d0eaa168f45 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
@@ -161,7 +161,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#streamSplits(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle,java.util.List,java.util.Optional,long):org.apache.doris.connector.spi.scan.ConnectorSplitSource
+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
org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider#supportsFileCache():boolean
diff --git a/fe/fe-connector/pom.xml b/fe/fe-connector/pom.xml
index c207b98365f..daa012ed574 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>15.0</connector.plugin.api.version>
+ <connector.plugin.api.version>17.0</connector.plugin.api.version>
</properties>
<modules>
diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FileQueryScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FileQueryScanNode.java
index 1edeb1b7c76..b1d0d4f480e 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FileQueryScanNode.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FileQueryScanNode.java
@@ -416,7 +416,8 @@ public abstract class FileQueryScanNode extends FileScanNode {
// File splits are generated lazily, and fetched by backends while scanning.
// Only provide the unique ID of split source to backend.
splitAssignment = new SplitAssignment(backendPolicy, this, this::splitToScanRange,
- locationProperties, pathPartitionKeys, admissionResult);
+ locationProperties, pathPartitionKeys, admissionResult,
+ sessionVariable.getSplitBatchSize(), sessionVariable.getSplitQueueCapacity());
splitAssignment.init();
if (executor != null) {
executor.getSummaryProfile().setGetSplitsFinishTime();
@@ -748,10 +749,13 @@ public abstract class FileQueryScanNode extends FileScanNode {
@Override
public void stop() {
if (splitAssignment != null) {
- splitAssignment.stop();
- SplitSourceManager manager = Env.getCurrentEnv().getSplitSourceManager();
- for (Long sourceId : splitAssignment.getSources()) {
- manager.removeSplitSource(sourceId);
+ try {
+ splitAssignment.stop();
+ } finally {
+ SplitSourceManager manager = Env.getCurrentEnv().getSplitSourceManager();
+ for (Long sourceId : splitAssignment.getSources()) {
+ manager.removeSplitSource(sourceId);
+ }
}
}
}
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 42ae2d116e3..e569392b6b5 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
@@ -35,6 +35,7 @@ import org.apache.doris.catalog.TableIf;
import org.apache.doris.catalog.Type;
import org.apache.doris.common.Config;
import org.apache.doris.common.UserException;
+import org.apache.doris.common.profile.Counter;
import org.apache.doris.common.profile.RuntimeProfile;
import org.apache.doris.common.profile.SummaryProfile;
import org.apache.doris.connector.spi.Connector;
@@ -72,6 +73,7 @@ import org.apache.doris.datasource.plugin.PluginDrivenSysExternalTable;
import org.apache.doris.datasource.split.FileSplit;
import org.apache.doris.datasource.split.PluginDrivenSplit;
import org.apache.doris.datasource.split.SplitAssignment;
+import org.apache.doris.datasource.split.SplitProducer;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan.SelectedPartitions;
import org.apache.doris.planner.PlanNodeId;
@@ -90,10 +92,13 @@ import org.apache.doris.thrift.TFileRangeDesc;
import org.apache.doris.thrift.TFileTextScanRangeParams;
import org.apache.doris.thrift.TPushAggOp;
import org.apache.doris.thrift.TTableFormatFileDesc;
+import org.apache.doris.thrift.TUnit;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
+import java.io.IOException;
+import java.io.UncheckedIOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -1969,10 +1974,18 @@ public class PluginDrivenScanNode extends FileQueryScanNode {
return false;
}
boolean hasSlots = !desc.getSlots().isEmpty();
+ // Keep partition-listing concurrency for large partitioned scans. Each bounded producer
+ // streams within its claimed partition batch through the scoped request.
+ boolean supportsBatchScan = onPluginClassLoader(scanProvider,
+ () -> scanProvider.supportsBatchScan(connectorSession, currentHandle));
+ if (shouldUseBatchMode(selectedPartitions, hasSlots,
+ supportsBatchScan, sessionVariable.getNumPartitionsInBatchMode())) {
+ return true;
+ }
// Streaming (file-count) batch flavor (FIX-M3): the connector owns the whole decision (e.g.
// Iceberg's matched-file count vs num_files_in_batch_mode); the engine only requires output slots
- // (a scan with no slots needs no file ranges). Checked before the partition-count flavor because a
- // connector implements at most one — Iceberg streams files, MaxCompute slices partitions.
+ // (a scan with no slots needs no file ranges). Covers non-partitioned scans and partitioned
+ // scans below the partition-batch threshold.
if (hasSlots) {
boolean countPushdown = isTableLevelCountStarPushdown();
long estimate = onPluginClassLoader(scanProvider, () -> scanProvider.streamingSplitEstimate(
@@ -1983,12 +1996,7 @@ public class PluginDrivenScanNode extends FileQueryScanNode {
return true;
}
}
- // Partition-count batch flavor (legacy MaxCompute): its connector odpsTable.getFileNum()>0 check is
- // folded into supportsBatchScan; the partition threshold is evaluated generically.
- boolean supportsBatchScan = onPluginClassLoader(scanProvider,
- () -> scanProvider.supportsBatchScan(connectorSession, currentHandle));
- return shouldUseBatchMode(selectedPartitions, hasSlots,
- supportsBatchScan, sessionVariable.getNumPartitionsInBatchMode());
+ return false;
}
/**
@@ -2059,13 +2067,9 @@ public class PluginDrivenScanNode extends FileQueryScanNode {
*/
@Override
public void startSplit(int numBackends) {
+ // EXPLAIN also initializes batch scans, but has no BE to drain their bounded queues.
+ QeProcessorImpl.INSTANCE.registerQueryFinishCallback(connectorSession.getQueryId(), this::stop);
QueryFinishCleanup queryFinishCleanup = registerQueryFinishCleanup(connectorSession);
- if (streamingBatch) {
- // File-count streaming flavor (FIX-M3): pump a connector-driven lazy source instead of
- // slicing partitions. Mutually exclusive with the partition-slicing path below.
- startStreamingSplit(queryFinishCleanup);
- return;
- }
try {
checkSysTableScanConstraints();
} catch (UserException e) {
@@ -2080,55 +2084,67 @@ public class PluginDrivenScanNode extends FileQueryScanNode {
this.selectedPartitionNum = partitionCounts[0];
this.totalPartitionNum = partitionCounts[1];
}
- if (selectedPartitions.selectedPartitions.isEmpty()) {
+ if (!streamingBatch && selectedPartitions.selectedPartitions.isEmpty()) {
// Unreachable under the isBatchMode gate (size >= num_partitions_in_batch_mode >= 1);
// kept for fidelity with legacy MaxComputeScanNode.startSplit's empty short-circuit.
return;
}
- // Mirror getSplits()'s projection + filter pushdown (but NOT the limit) before going async.
- // tryPushDownProjection mutates currentHandle, so capture the resolved handle afterwards.
- // buildColumnHandles (a time-travel schema/handle mismatch) and pinMvccSnapshot both throw;
- // startSplit cannot throw checked exceptions, so surface either through the SplitAssignment
- // error channel (same protocol as checkSysTableScanConstraints above) rather than dropping it.
- final List<ConnectorColumnHandle> columns;
- final Optional<ConnectorExpression> remainingFilter;
+ final ConnectorScanRequest request;
try {
- columns = buildColumnHandles();
- tryPushDownProjection(columns);
- remainingFilter = buildRemainingFilter();
- // Pin the MVCC snapshot onto currentHandle before the resolved handle is captured below,
- // so the async batch path (planScanForPartitionBatch) reads at the pinned snapshot.
- pinMvccSnapshot();
+ request = prepareBatchScan(queryFinishCleanup);
} catch (UserException e) {
+ // startSplit cannot throw checked exceptions; use the same error channel for both flavors.
splitAssignment.setException(e);
return;
}
- // Scope the scan to a distributed rewrite group's files (no-op for every non-rewrite read).
+ ConnectorScanPlanProvider scanProvider = resolveScanProvider();
+ SummaryProfile summaryProfile = SummaryProfile.getSummaryProfile(ConnectContext.get());
+ List<SplitProducer.BatchSource> sources;
+ if (streamingBatch) {
+ List<String> requiredPartitions = resolveRequiredPartitions(selectedPartitions,
+ onPluginClassLoader(scanProvider, scanProvider::ignorePartitionPruneShortCircuit));
+ if (requiredPartitions != null && requiredPartitions.isEmpty()) {
+ splitAssignment.finishSchedule();
+ return;
+ }
+ sources = Collections.singletonList(createStreamingBatchSource(scanProvider,
+ request.withRequiredPartitions(requiredPartitions), summaryProfile));
+ } else {
+ sources = createPartitionBatchSources(scanProvider, request,
+ new ArrayList<>(selectedPartitions.selectedPartitions.keySet()),
+ sessionVariable.getNumPartitionsInBatchMode(), Config.max_external_cache_loader_thread_pool_size);
+ }
+ startSplitProducers(scanProvider, sources,
+ summaryProfile == null ? null : summaryProfile.getExecutionSummary());
+ }
+
+ private ConnectorScanRequest prepareBatchScan(QueryFinishCleanup queryFinishCleanup) throws UserException {
+ // Projection mutates currentHandle. Pin snapshot and rewrite scope before capturing it in the
+ // request, just as getSplits does. Batch scans deliberately omit LIMIT and COUNT(*) pushdown.
+ List<ConnectorColumnHandle> columns = buildColumnHandles();
+ tryPushDownProjection(columns);
+ Optional<ConnectorExpression> remainingFilter = buildRemainingFilter();
+ // The pre-pin streaming estimate is only a scheduling hint; the source reads the pinned handle.
+ pinMvccSnapshot();
pinRewriteFileScope();
- final ConnectorTableHandle handle = currentHandle;
- final ConnectorScanPlanProvider scanProvider = resolveScanProvider();
- queryFinishCleanup.bind(scanProvider);
- // One request for the whole batched scan; each batch re-scopes it to its own partitions. No row
- // limit and no COUNT(*) pushdown on this path (batch mode is entered before either applies),
- // matching what the batched call passed before the request object existed.
- final ConnectorScanRequest batchRequest = ConnectorScanRequest.builder(handle, columns)
- .filter(remainingFilter)
- .build();
- final List<String> allPartitions =
- new ArrayList<>(selectedPartitions.selectedPartitions.keySet());
- final int batchSize = sessionVariable.getNumPartitionsInBatchMode();
- SummaryProfile batchSummaryProfile = SummaryProfile.getSummaryProfile(ConnectContext.get());
- final RuntimeProfile batchExecutionSummary = batchSummaryProfile == null
- ? null : batchSummaryProfile.getExecutionSummary();
+ queryFinishCleanup.bind(resolveScanProvider());
+ return ConnectorScanRequest.builder(currentHandle, columns).filter(remainingFilter).build();
+ }
+ private void startSplitProducers(ConnectorScanPlanProvider scanProvider, List<SplitProducer.BatchSource> sources,
+ RuntimeProfile executionSummary) {
Executor scheduleExecutor = Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor();
AtomicReference<UserException> batchException = new AtomicReference<>(null);
SubmittedTaskFinalizer profileFinalizer = new SubmittedTaskFinalizer(() -> {
try {
- List<ConnectorScanProfile> profiles = onPluginClassLoader(scanProvider,
- () -> scanProvider.collectScanProfiles(connectorSession));
- writeScanProfilesInto(batchExecutionSummary, profiles);
+ // Whole-scan streaming previously collected profiles only on successful completion;
+ // partition batching also retains partial profiles when a producer fails or is stopped.
+ if (!streamingBatch || splitAssignment.needMoreSplit()) {
+ List<ConnectorScanProfile> profiles = onPluginClassLoader(scanProvider,
+ () -> scanProvider.collectScanProfiles(connectorSession));
+ writeScanProfilesInto(executionSummary, profiles);
+ }
} catch (Exception e) {
UserException profileFailure = new UserException(e.getMessage(), e);
UserException primaryFailure = batchException.get();
@@ -2144,48 +2160,34 @@ public class PluginDrivenScanNode extends FileQueryScanNode {
Runnable dispatch = () -> {
try {
- for (int begin = 0; begin < allPartitions.size(); begin += batchSize) {
- int end = Math.min(begin + batchSize, allPartitions.size());
+ for (SplitProducer.BatchSource source : sources) {
if (batchException.get() != null || splitAssignment.isStop()) {
break;
}
- List<String> batch = allPartitions.subList(begin, end);
profileFinalizer.taskSubmitted();
- try {
- CompletableFuture.runAsync(() -> {
- try {
- List<ConnectorScanRange> ranges = onPluginClassLoader(scanProvider,
- () -> scanProvider.planScanForPartitionBatch(
- connectorSession, batchRequest, batch));
- List<Split> batchSplits = new ArrayList<>(ranges.size());
- for (ConnectorScanRange range : ranges) {
- batchSplits.add(toSplit(range));
- }
- if (splitAssignment.needMoreSplit()) {
- splitAssignment.addToQueue(batchSplits);
- }
- } catch (Exception e) {
+ new SplitProducer(splitAssignment, scheduleExecutor, source, failure -> {
+ try {
+ if (failure != null) {
publishBatchFailure(batchException, splitAssignment,
- new UserException(e.getMessage(), e));
- } finally {
- profileFinalizer.taskFinished();
+ new UserException(failure.getMessage(), failure));
}
- }, scheduleExecutor);
- } catch (Exception e) {
- try {
- publishBatchFailure(batchException, splitAssignment,
- new UserException(e.getMessage(), e));
} finally {
profileFinalizer.taskFinished();
}
- }
+ }).start();
}
} finally {
profileFinalizer.closeDispatch();
}
};
try {
- CompletableFuture.runAsync(dispatch, scheduleExecutor);
+ // Keep the single-source path's direct submission; partition dispatch still runs on the
+ // scheduler so admitting its lanes does not hold the planning thread.
+ if (streamingBatch) {
+ dispatch.run();
+ } else {
+ CompletableFuture.runAsync(dispatch, scheduleExecutor);
+ }
} catch (Exception e) {
try {
publishBatchFailure(batchException, splitAssignment,
@@ -2196,79 +2198,108 @@ public class PluginDrivenScanNode extends FileQueryScanNode {
}
}
- /**
- * Streams splits from a connector-driven lazy {@link ConnectorSplitSource} into
- * {@link #splitAssignment} with backpressure, mirroring legacy {@code IcebergScanNode.doStartSplit}.
- * Used by the file-count streaming batch flavor (see {@link #computeBatchMode}); the partition-count
- * flavor stays on the {@link #startSplit} partition-slicing path. Deliberately does NOT push the limit
- * (passes {@code -1}): the LIMIT-split optimization stays on the non-batch {@link #getSplits} path only.
- */
- private void startStreamingSplit(QueryFinishCleanup queryFinishCleanup) {
- try {
- checkSysTableScanConstraints();
- } catch (UserException e) {
- // startSplit cannot throw checked exceptions; surface the fail-loud guard through the
- // SplitAssignment error channel so the query fails rather than silently ignoring the guard.
- splitAssignment.setException(e);
- return;
+ List<SplitProducer.BatchSource> createPartitionBatchSources(ConnectorScanPlanProvider scanProvider,
+ ConnectorScanRequest request, List<String> partitions, int partitionsPerBatch, int parallelism) {
+ AtomicInteger nextPartition = new AtomicInteger();
+ int sourceCount = Math.min(parallelism, (partitions.size() - 1) / partitionsPerBatch + 1);
+ 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.
+ for (int sourceIndex = 0; sourceIndex < sourceCount; sourceIndex++) {
+ sources.add(new SplitProducer.BatchSource() {
+ private ConnectorSplitSource source;
+
+ @Override
+ public List<Split> getNextBatch() {
+ return onPluginClassLoader(scanProvider, () -> {
+ List<Split> splits = new ArrayList<>(splitAssignment.getSplitBatchSize());
+ while (splits.size() < splitAssignment.getSplitBatchSize()
+ && splitAssignment.needMoreSplit()) {
+ if (source != null && source.hasNext()) {
+ splits.add(toSplit(source.next()));
+ continue;
+ }
+ closeSource();
+ int begin = nextPartition.getAndAdd(partitionsPerBatch);
+ if (begin >= partitions.size()) {
+ break;
+ }
+ int end = Math.min(begin + partitionsPerBatch, partitions.size());
+ source = scanProvider.streamSplits(connectorSession,
+ request.withRequiredPartitions(partitions.subList(begin, end)));
+ }
+ return splits;
+ });
+ }
+
+ private void closeSource() {
+ if (source != null) {
+ try {
+ source.close();
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ } finally {
+ source = null;
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ onPluginClassLoader(scanProvider, () -> {
+ closeSource();
+ return null;
+ });
+ }
+ });
}
- // Mirror getSplits()'s projection + filter pushdown (but NOT the limit) before going async.
- // tryPushDownProjection mutates currentHandle, so capture the resolved handle afterwards.
- // buildColumnHandles (a time-travel schema/handle mismatch) and pinMvccSnapshot both throw;
- // startSplit cannot throw checked exceptions, so surface either through the SplitAssignment
- // error channel so the query fails rather than silently ignoring it.
- final List<ConnectorColumnHandle> columns;
- final Optional<ConnectorExpression> remainingFilter;
- try {
- columns = buildColumnHandles();
- tryPushDownProjection(columns);
- remainingFilter = buildRemainingFilter();
- // Pin the MVCC snapshot + rewrite scope onto currentHandle before capturing the resolved
- // handle, so the lazy source plans at the PINNED snapshot (rows are always correct). The
- // streamingSplitEstimate gate ran pre-pin (current snapshot), so it only affects the
- // approximate batch decision + concurrency hint, never the rows. Narrow caveat: a
- // time-travel query to a past snapshot much LARGER than current (table since
- // compacted/expired) may under-count -> pick the eager path -> lose streaming's OOM
- // protection for that scan. Acceptable (perf-only, rare); documented in the FIX-M3 design.
- pinMvccSnapshot();
- } catch (UserException e) {
- splitAssignment.setException(e);
- return;
+ return sources;
+ }
+
+ private SplitProducer.BatchSource createStreamingBatchSource(ConnectorScanPlanProvider scanProvider,
+ ConnectorScanRequest request, SummaryProfile summaryProfile) {
+ RuntimeProfile generationProfile = new RuntimeProfile("Split Generation " + getId());
+ generationProfile.addInfoString("BatchSize", String.valueOf(sessionVariable.getSplitBatchSize()));
+ generationProfile.addInfoString("QueueCapacity", String.valueOf(sessionVariable.getSplitQueueCapacity()));
+ Counter generatedSplits = generationProfile.addCounter("GeneratedSplits", TUnit.UNIT,
+ RuntimeProfile.ROOT_COUNTER);
+ generatedSplits.setLevel(1);
+ if (summaryProfile != null) {
+ summaryProfile.getExecutionSummary().addChild(generationProfile, true);
}
- pinRewriteFileScope();
- final ConnectorTableHandle handle = currentHandle;
- final ConnectorScanPlanProvider scanProvider = resolveScanProvider();
- queryFinishCleanup.bind(scanProvider);
- Executor scheduleExecutor = Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor();
- CompletableFuture.runAsync(() -> {
- ConnectorSplitSource source = null;
- try {
- source = onPluginClassLoader(scanProvider,
- () -> scanProvider.streamSplits(connectorSession, handle, columns, remainingFilter, -1L));
- // Pull ranges with backpressure (needMoreSplit) and pump them one at a time, exactly like
- // legacy doStartSplit. The bounded SplitAssignment queue throttles the lazy source so FE
- // heap stays bounded for million-file scans.
- while (splitAssignment.needMoreSplit() && source.hasNext()) {
- List<Split> one = new ArrayList<>(1);
- one.add(toSplit(source.next()));
- splitAssignment.addToQueue(one);
- }
- splitAssignment.finishSchedule();
- } catch (Exception e) {
- splitAssignment.setException(new UserException(e.getMessage(), e));
- } finally {
- // Close in a finally that SWALLOWS close errors (NOT try-with-resources, whose close runs
- // before the catch): a close() failure must not fail a scan whose splits were already
- // enumerated + finishSchedule()-d (legacy doStartSplit swallowed close errors identically).
- if (source != null) {
- try {
- source.close();
- } catch (Exception ce) {
- LOG.warn("Failed to close streaming split source for {}", handle, ce);
+ return new SplitProducer.BatchSource() {
+ private ConnectorSplitSource source;
+
+ @Override
+ public List<Split> getNextBatch() {
+ return onPluginClassLoader(scanProvider, () -> {
+ if (source == null) {
+ source = scanProvider.streamSplits(connectorSession, request);
+ }
+ List<Split> splits = new ArrayList<>(splitAssignment.getSplitBatchSize());
+ while (splits.size() < splitAssignment.getSplitBatchSize()
+ && splitAssignment.needMoreSplit() && source.hasNext()) {
+ splits.add(toSplit(source.next()));
}
+ generatedSplits.setValue(generatedSplits.getValue() + splits.size());
+ return splits;
+ });
+ }
+
+ @Override
+ public void close() {
+ if (source != null) {
+ onPluginClassLoader(scanProvider, () -> {
+ try {
+ source.close();
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ return null;
+ });
}
}
- }, scheduleExecutor);
+ };
}
@Override
diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitAssignment.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitAssignment.java
index 61d7ab9ec0d..ca904044815 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitAssignment.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitAssignment.java
@@ -23,6 +23,7 @@ import org.apache.doris.spi.Split;
import org.apache.doris.system.Backend;
import org.apache.doris.thrift.TScanRangeLocations;
+import com.google.common.base.Preconditions;
import com.google.common.collect.Multimap;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
@@ -30,6 +31,7 @@ import org.apache.logging.log4j.Logger;
import java.io.Closeable;
import java.util.ArrayList;
import java.util.Collection;
+import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
@@ -55,13 +57,16 @@ public class SplitAssignment {
private final Map<String, String> locationProperties;
private final List<String> pathPartitionKeys;
private final boolean fileCacheAdmission;
+ private final int splitBatchSize;
+ private final int queueCapacity;
private final Object assignLock = new Object();
private Split sampleSplit = null;
private final AtomicBoolean isStopped = new AtomicBoolean(false);
private final AtomicBoolean scheduleFinished = new AtomicBoolean(false);
- private UserException exception = null;
+ private volatile UserException exception = null;
private final List<Closeable> closeableResources = new ArrayList<>();
+ private final Map<Backend, Set<Runnable>> queueSpaceListeners = new HashMap<>();
public SplitAssignment(
FederationBackendPolicy backendPolicy,
@@ -69,13 +74,19 @@ public class SplitAssignment {
SplitToScanRange splitToScanRange,
Map<String, String> locationProperties,
List<String> pathPartitionKeys,
- boolean fileCacheAdmission) {
+ boolean fileCacheAdmission,
+ int splitBatchSize,
+ int queueCapacity) {
+ Preconditions.checkArgument(splitBatchSize > 0, "Split batch size must be positive");
+ Preconditions.checkArgument(queueCapacity > 0, "Split queue capacity must be positive");
this.backendPolicy = backendPolicy;
this.splitGenerator = splitGenerator;
this.splitToScanRange = splitToScanRange;
this.locationProperties = locationProperties;
this.pathPartitionKeys = pathPartitionKeys;
this.fileCacheAdmission = fileCacheAdmission;
+ this.splitBatchSize = splitBatchSize;
+ this.queueCapacity = queueCapacity;
}
public void init() throws UserException {
@@ -108,26 +119,85 @@ public class SplitAssignment {
private void appendBatch(Multimap<Backend, Split> batch) throws UserException {
for (Backend backend : batch.keySet()) {
- Collection<Split> splits = batch.get(backend);
- List<TScanRangeLocations> locations = new ArrayList<>(splits.size());
- for (Split split : splits) {
- locations.add(splitToScanRange.getScanRange(backend, locationProperties, split, pathPartitionKeys,
- fileCacheAdmission));
- }
+ Collection<TScanRangeLocations> locations = createScanRanges(backend, batch.get(backend));
while (needMoreSplit()) {
BlockingQueue<Collection<TScanRangeLocations>> queue =
- assignment.computeIfAbsent(backend, be -> new LinkedBlockingQueue<>(10000));
+ getOrCreateAssignmentQueue(backend);
try {
if (queue.offer(locations, 100, TimeUnit.MILLISECONDS)) {
break;
}
} catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
addUserException(new UserException("Failed to offer batch split by interrupted", e));
}
}
}
}
+ Multimap<Backend, Split> assignBatch(List<Split> splits) throws UserException {
+ Preconditions.checkArgument(splits.size() <= splitBatchSize, "Producer batch exceeds split batch size");
+ synchronized (assignLock) {
+ return backendPolicy.computeScanRangeAssignment(splits);
+ }
+ }
+
+ Collection<TScanRangeLocations> createScanRanges(Backend backend, Collection<Split> splits)
+ throws UserException {
+ List<TScanRangeLocations> locations = new ArrayList<>(splits.size());
+ for (Split split : splits) {
+ locations.add(splitToScanRange.getScanRange(backend, locationProperties, split, pathPartitionKeys,
+ fileCacheAdmission));
+ }
+ // Publish readiness after range conversion initializes the scan-level parameters.
+ synchronized (assignLock) {
+ if (sampleSplit == null) {
+ sampleSplit = splits.iterator().next();
+ assignLock.notifyAll();
+ }
+ }
+ return locations;
+ }
+
+ boolean offerBatch(Backend backend, Collection<TScanRangeLocations> locations, Runnable resume) {
+ // Coordinate the failed offer and listener registration with consumer notifications so
+ // capacity released just before suspension cannot be missed.
+ synchronized (queueSpaceListeners) {
+ // An error/stop may race range conversion. Do not register a waiter after its wakeup.
+ if (!needMoreSplit()) {
+ return false;
+ }
+ if (getOrCreateAssignmentQueue(backend).offer(locations)) {
+ removeQueueSpaceListener(backend, resume);
+ return true;
+ }
+ queueSpaceListeners.computeIfAbsent(backend, key -> new HashSet<>()).add(resume);
+ return false;
+ }
+ }
+
+ void removeQueueSpaceListener(Backend backend, Runnable listener) {
+ synchronized (queueSpaceListeners) {
+ Set<Runnable> listeners = queueSpaceListeners.get(backend);
+ if (listeners != null) {
+ listeners.remove(listener);
+ if (listeners.isEmpty()) {
+ queueSpaceListeners.remove(backend);
+ }
+ }
+ }
+ }
+
+ void notifyQueueSpace(Backend backend) {
+ Set<Runnable> listeners;
+ synchronized (queueSpaceListeners) {
+ listeners = queueSpaceListeners.remove(backend);
+ }
+ if (listeners != null) {
+ listeners.forEach(Runnable::run);
+ }
+ }
+
public void registerSource(long uniqueId) {
sources.add(uniqueId);
}
@@ -141,18 +211,16 @@ public class SplitAssignment {
}
public void addToQueue(List<Split> splits) throws UserException {
- if (splits.isEmpty()) {
- return;
- }
- Multimap<Backend, Split> batch = null;
- synchronized (assignLock) {
- if (sampleSplit == null) {
- sampleSplit = splits.get(0);
- assignLock.notify();
- }
- batch = backendPolicy.computeScanRangeAssignment(splits);
+ for (int begin = 0; begin < splits.size() && needMoreSplit();) {
+ int end = begin + Math.min(splitBatchSize, splits.size() - begin);
+ List<Split> batchSplits = new ArrayList<>(splits.subList(begin, end));
+ appendBatch(assignBatch(batchSplits));
+ begin = end;
}
- appendBatch(batch);
+ }
+
+ public int getSplitBatchSize() {
+ return splitBatchSize;
}
private void notifyAssignment() {
@@ -161,12 +229,15 @@ public class SplitAssignment {
}
}
+ private BlockingQueue<Collection<TScanRangeLocations>> getOrCreateAssignmentQueue(Backend backend) {
+ return assignment.computeIfAbsent(backend, be -> new LinkedBlockingQueue<>(queueCapacity));
+ }
+
public BlockingQueue<Collection<TScanRangeLocations>> getAssignedSplits(Backend backend) throws UserException {
if (exception != null) {
throw exception;
}
- BlockingQueue<Collection<TScanRangeLocations>> splits = assignment.computeIfAbsent(backend,
- be -> new LinkedBlockingQueue<>());
+ BlockingQueue<Collection<TScanRangeLocations>> splits = getOrCreateAssignmentQueue(backend);
if (scheduleFinished.get() && splits.isEmpty() || isStopped.get()) {
return null;
}
@@ -176,9 +247,15 @@ public class SplitAssignment {
public void setException(UserException e) {
addUserException(e);
notifyAssignment();
+ Set<Runnable> listeners = new HashSet<>();
+ synchronized (queueSpaceListeners) {
+ queueSpaceListeners.values().forEach(listeners::addAll);
+ queueSpaceListeners.clear();
+ }
+ listeners.forEach(Runnable::run);
}
- private void addUserException(UserException e) {
+ private synchronized void addUserException(UserException e) {
if (exception != null) {
exception.addSuppressed(e);
} else {
@@ -192,18 +269,15 @@ public class SplitAssignment {
}
public void stop() {
- if (isStop()) {
+ if (!isStopped.compareAndSet(false, true)) {
return;
}
- isStopped.set(true);
- closeableResources.forEach((closeable) -> {
- try {
- closeable.close();
- } catch (Exception e) {
- LOG.warn("close resource error:{}", e.getMessage(), e);
- // ignore
- }
- });
+ List<Closeable> resources;
+ synchronized (closeableResources) {
+ resources = new ArrayList<>(closeableResources);
+ closeableResources.clear();
+ }
+ resources.forEach(this::closeResource);
notifyAssignment();
if (exception != null) {
throw new RuntimeException(exception);
@@ -215,6 +289,26 @@ public class SplitAssignment {
}
public void addCloseable(Closeable resource) {
- closeableResources.add(resource);
+ synchronized (closeableResources) {
+ if (!isStop()) {
+ closeableResources.add(resource);
+ return;
+ }
+ }
+ closeResource(resource);
+ }
+
+ void removeCloseable(Closeable resource) {
+ synchronized (closeableResources) {
+ closeableResources.remove(resource);
+ }
+ }
+
+ private void closeResource(Closeable resource) {
+ try {
+ resource.close();
+ } catch (Exception e) {
+ LOG.warn("close resource error:{}", e.getMessage(), e);
+ }
}
}
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
new file mode 100644
index 00000000000..6d2f0e26207
--- /dev/null
+++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java
@@ -0,0 +1,183 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.datasource.split;
+
+import org.apache.doris.common.UserException;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TScanRangeLocations;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+
+/**
+ * Serially advances one producer batch per task, yielding the shared executor when a BE queue is
+ * full. Each producer retains at most one assigned batch and resumes when a consumer frees space.
+ */
+public class SplitProducer implements Runnable, Closeable {
+ private static final Logger LOG = LogManager.getLogger(SplitProducer.class);
+
+ public interface BatchSource extends Closeable {
+ List<Split> getNextBatch() throws Exception;
+
+ @Override
+ default void close() throws IOException {
+ }
+ }
+
+ private final SplitAssignment assignment;
+ private final Executor executor;
+ private final BatchSource source;
+ private final Consumer<Exception> completion;
+ private final Runnable resume = this::schedule;
+ private final Object scheduleLock = new Object();
+ private boolean scheduled;
+ private boolean wakeupRequested;
+ private boolean completed;
+ private volatile boolean cancelled;
+ private Iterator<Map.Entry<Backend, Collection<Split>>> pending = Collections.emptyIterator();
+ private Backend pendingBackend;
+ private Collection<TScanRangeLocations> pendingRanges;
+
+ public SplitProducer(SplitAssignment assignment, Executor executor, BatchSource source,
+ Consumer<Exception> completion) {
+ this.assignment = assignment;
+ this.executor = executor;
+ this.source = source;
+ this.completion = completion;
+ }
+
+ public void start() {
+ assignment.addCloseable(this);
+ schedule();
+ }
+
+ private void schedule() {
+ synchronized (scheduleLock) {
+ if (completed) {
+ return;
+ }
+ wakeupRequested = true;
+ if (scheduled) {
+ return;
+ }
+ scheduled = true;
+ }
+ submit();
+ }
+
+ private void submit() {
+ try {
+ executor.execute(this);
+ } catch (RuntimeException e) {
+ complete(cancelled || !assignment.needMoreSplit() ? null : e);
+ }
+ }
+
+ @Override
+ public void run() {
+ synchronized (scheduleLock) {
+ wakeupRequested = false;
+ }
+ boolean more = false;
+ try {
+ more = produceBatch();
+ } catch (Exception e) {
+ complete(e);
+ }
+ synchronized (scheduleLock) {
+ if (completed || (!more && !wakeupRequested)) {
+ scheduled = false;
+ return;
+ }
+ }
+ submit();
+ }
+
+ private boolean produceBatch() throws Exception {
+ if (cancelled || !assignment.needMoreSplit()) {
+ complete(null);
+ return false;
+ }
+ if (pendingRanges == null && !pending.hasNext()) {
+ List<Split> batch = source.getNextBatch();
+ if (batch.isEmpty()) {
+ complete(null);
+ return false;
+ }
+ pending = assignment.assignBatch(batch).asMap().entrySet().iterator();
+ }
+ while (!cancelled && assignment.needMoreSplit()) {
+ if (pendingRanges == null) {
+ if (!pending.hasNext()) {
+ return true;
+ }
+ Map.Entry<Backend, Collection<Split>> entry = pending.next();
+ pendingBackend = entry.getKey();
+ pendingRanges = assignment.createScanRanges(pendingBackend, entry.getValue());
+ }
+ if (!assignment.offerBatch(pendingBackend, pendingRanges, resume)) {
+ if (cancelled || !assignment.needMoreSplit()) {
+ complete(null);
+ }
+ return false;
+ }
+ pendingRanges = null;
+ }
+ complete(null);
+ return false;
+ }
+
+ private void complete(Exception failure) {
+ synchronized (scheduleLock) {
+ if (completed) {
+ return;
+ }
+ completed = true;
+ }
+ assignment.removeQueueSpaceListener(pendingBackend, resume);
+ assignment.removeCloseable(this);
+ try {
+ completion.accept(failure);
+ } catch (Exception e) {
+ assignment.setException(new UserException(e.getMessage(), e));
+ } finally {
+ try {
+ source.close();
+ } catch (Exception e) {
+ LOG.warn("Failed to close split producer source", e);
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ cancelled = true;
+ schedule();
+ }
+}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitSource.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitSource.java
index 0163223b243..b1be85133b2 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitSource.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitSource.java
@@ -21,10 +21,12 @@ import org.apache.doris.common.UserException;
import org.apache.doris.system.Backend;
import org.apache.doris.thrift.TScanRangeLocations;
-import com.google.common.collect.Lists;
+import com.google.common.base.Preconditions;
+import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
+import java.util.Iterator;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
@@ -50,6 +52,7 @@ public class SplitSource {
private final SplitAssignment splitAssignment;
private final AtomicBoolean isLastBatch;
private final long maxWaitTime;
+ private Iterator<TScanRangeLocations> pendingSplits = Collections.emptyIterator();
public SplitSource(Backend backend, SplitAssignment splitAssignment, long maxWaitTime) {
this.uniqueId = UNIQUE_ID_GENERATOR.getAndIncrement();
@@ -67,40 +70,60 @@ public class SplitSource {
/**
* Get the next batch of file splits. If there's no more split, return empty list.
*/
- public List<TScanRangeLocations> getNextBatch(int maxBatchSize) throws UserException {
+ public synchronized List<TScanRangeLocations> getNextBatch(int maxBatchSize) throws UserException {
+ Preconditions.checkArgument(maxBatchSize > 0, "Requested split batch size must be positive");
if (isLastBatch.get()) {
return Collections.emptyList();
}
- List<TScanRangeLocations> scanRanges = Lists.newArrayListWithExpectedSize(maxBatchSize);
- long startTime = System.currentTimeMillis();
- while (scanRanges.size() < maxBatchSize && System.currentTimeMillis() - startTime < maxWaitTime) {
+ List<TScanRangeLocations> scanRanges = new ArrayList<>(
+ Math.min(maxBatchSize, splitAssignment.getSplitBatchSize()));
+ long startTime = System.nanoTime();
+ long waitNanos = TimeUnit.MILLISECONDS.toNanos(maxWaitTime);
+ while (scanRanges.size() < maxBatchSize) {
+ // Check producer errors even while draining a previously fetched batch.
BlockingQueue<Collection<TScanRangeLocations>> splits = splitAssignment.getAssignedSplits(backend);
+ if (splitAssignment.isStop()) {
+ pendingSplits = Collections.emptyIterator();
+ isLastBatch.set(true);
+ break;
+ }
+ while (pendingSplits.hasNext() && scanRanges.size() < maxBatchSize) {
+ scanRanges.add(pendingSplits.next());
+ }
+ if (scanRanges.size() == maxBatchSize) {
+ return scanRanges;
+ }
+ pendingSplits = Collections.emptyIterator();
+ // A finished, empty queue is EOF only after the pending batch has been drained.
if (splits == null) {
isLastBatch.set(true);
break;
}
- while (scanRanges.size() < maxBatchSize) {
- try {
- Collection<TScanRangeLocations> splitCollection = splits.poll(WAIT_TIME_OUT, TimeUnit.MILLISECONDS);
- if (splitCollection != null) {
- scanRanges.addAll(splitCollection);
- }
- if (!scanRanges.isEmpty() && System.currentTimeMillis() - startTime > maxWaitTime) {
- return scanRanges;
- }
- if (splitCollection == null) {
- break;
- }
- } catch (InterruptedException e) {
- throw new UserException(e.getMessage(), e);
+ Collection<TScanRangeLocations> batch = splits.poll();
+ if (batch != null) {
+ splitAssignment.notifyQueueSpace(backend);
+ pendingSplits = batch.iterator();
+ continue;
+ }
+ if (!scanRanges.isEmpty()) {
+ break;
+ }
+ long remainingNanos = waitNanos - (System.nanoTime() - startTime);
+ if (remainingNanos <= 0) {
+ throw new UserException("Timeout. Max wait time(ms): " + maxWaitTime);
+ }
+ try {
+ batch = splits.poll(Math.min(remainingNanos, TimeUnit.MILLISECONDS.toNanos(WAIT_TIME_OUT)),
+ TimeUnit.NANOSECONDS);
+ if (batch != null) {
+ splitAssignment.notifyQueueSpace(backend);
+ pendingSplits = batch.iterator();
}
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new UserException(e.getMessage(), e);
}
}
-
- if (scanRanges.isEmpty() && !isLastBatch.get()) {
- // This is timeout
- throw new UserException("Timeout. Max wait time(ms): " + maxWaitTime);
- }
return scanRanges;
}
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
index 30dd95266f7..425a5ec735e 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
@@ -590,6 +590,10 @@ public class SessionVariable implements Serializable, Writable {
public static final String NUM_FILES_IN_BATCH_MODE = "num_files_in_batch_mode";
+ public static final String SPLIT_BATCH_SIZE = "split_batch_size";
+
+ public static final String SPLIT_QUEUE_CAPACITY = "split_queue_capacity";
+
public static final String FETCH_SPLITS_MAX_WAIT_TIME = "fetch_splits_max_wait_time_ms";
/**
@@ -2597,6 +2601,22 @@ public class SessionVariable implements Serializable, Writable {
needForward = true)
public int numFilesInBatchMode = 1024;
+ @VarAttrDef.VarAttr(name = SPLIT_BATCH_SIZE, checker = "checkSplitBatchSize", needForward = true,
+ description = "Maximum number of splits in a producer batch.")
+ public int splitBatchSize = 128;
+
+ @VarAttrDef.VarAttr(name = SPLIT_QUEUE_CAPACITY, checker = "checkSplitQueueCapacity", needForward = true,
+ description = "Maximum number of queued split batches per backend.")
+ public int splitQueueCapacity = 32;
+
+ public void checkSplitBatchSize(String value) throws Exception {
+ checkFieldValue(SPLIT_BATCH_SIZE, 1, value);
+ }
+
+ public void checkSplitQueueCapacity(String value) throws Exception {
+ checkFieldValue(SPLIT_QUEUE_CAPACITY, 1, value);
+ }
+
@VarAttrDef.VarAttr(
name = FETCH_SPLITS_MAX_WAIT_TIME,
description = "The max wait time of getting splits in batch mode.",
@@ -4931,6 +4951,14 @@ public class SessionVariable implements Serializable, Writable {
return numFilesInBatchMode;
}
+ public int getSplitBatchSize() {
+ return splitBatchSize;
+ }
+
+ public int getSplitQueueCapacity() {
+ return splitQueueCapacity;
+ }
+
public long getFetchSplitsMaxWaitTime() {
return fetchSplitsMaxWaitTime;
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java b/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java
index ec4bc17148e..629378a7353 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/VariableMgr.java
@@ -111,6 +111,8 @@ public class VariableMgr {
// If a session variable "foo" is an experimental variable,
// its display name is "experimental_foo"
private static ImmutableMap<String, VarContext> ctxByDisplayVarName;
+ // Visibility and aliases come from immutable annotations; values must still be read on every export.
+ private static final ImmutableMap<String, VarContext> SHOWN_CTX_BY_VAR_NAME;
// This variable is equivalent to the default value of session variables.
// Whenever a new session is established, the value in this object is copied to the session-level variable.
@@ -172,6 +174,13 @@ public class VariableMgr {
ImmutableSortedMap.Builder<String, VarContext> builder = getStringVarContextBuilder(defaultSessionVariable);
ctxByVarName = builder.build();
ctxByDisplayVarName = getDisplaySessionVars();
+ ImmutableMap.Builder<String, VarContext> shownBuilder = ImmutableMap.builder();
+ for (Map.Entry<String, VarContext> entry : ctxByVarName.entrySet()) {
+ if (isShown(entry.getValue())) {
+ shownBuilder.put(entry);
+ }
+ }
+ SHOWN_CTX_BY_VAR_NAME = shownBuilder.build();
}
public static SessionVariable getDefaultSessionVariable() {
@@ -976,14 +985,11 @@ public class VariableMgr {
* without a hard-coded whitelist.</p>
*/
public static Map<String, String> toMap(SessionVariable sessionVar) {
- Map<String, String> result = Maps.newHashMap();
+ Map<String, String> result = Maps.newHashMapWithExpectedSize(SHOWN_CTX_BY_VAR_NAME.size());
rlock.lock();
try {
- for (Map.Entry<String, VarContext> entry : ctxByVarName.entrySet()) {
+ for (Map.Entry<String, VarContext> entry : SHOWN_CTX_BY_VAR_NAME.entrySet()) {
VarContext ctx = entry.getValue();
- if (!isShown(ctx)) {
- continue;
- }
result.put(entry.getKey(), getValue(sessionVar, ctx.getField()));
}
} finally {
diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeHiveStreamingTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeHiveStreamingTest.java
new file mode 100644
index 00000000000..79c603ee3ba
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeHiveStreamingTest.java
@@ -0,0 +1,325 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.datasource.scan;
+
+import org.apache.doris.analysis.SlotDescriptor;
+import org.apache.doris.analysis.TupleDescriptor;
+import org.apache.doris.catalog.Env;
+import org.apache.doris.catalog.PartitionItem;
+import org.apache.doris.catalog.TableIf;
+import org.apache.doris.common.Config;
+import org.apache.doris.common.jmockit.Deencapsulation;
+import org.apache.doris.common.util.DebugUtil;
+import org.apache.doris.connector.hive.HiveCatalogProperties;
+import org.apache.doris.connector.hive.HiveFileListingCache;
+import org.apache.doris.connector.hive.HiveReadTransactionManager;
+import org.apache.doris.connector.hive.HiveScanPlanProvider;
+import org.apache.doris.connector.hive.HiveScanRange;
+import org.apache.doris.connector.hive.HiveTableHandle;
+import org.apache.doris.connector.hive.HiveTableType;
+import org.apache.doris.connector.hms.HmsClient;
+import org.apache.doris.connector.hms.HmsPartitionInfo;
+import org.apache.doris.connector.spi.Connector;
+import org.apache.doris.connector.spi.ConnectorContext;
+import org.apache.doris.connector.spi.ConnectorMetadata;
+import org.apache.doris.connector.spi.ConnectorSession;
+import org.apache.doris.connector.spi.ConnectorStatementScope;
+import org.apache.doris.connector.spi.ConnectorStorageContext;
+import org.apache.doris.datasource.ExternalMetaCacheMgr;
+import org.apache.doris.datasource.split.PluginDrivenSplit;
+import org.apache.doris.datasource.split.SplitAssignment;
+import org.apache.doris.datasource.split.SplitGenerator;
+import org.apache.doris.datasource.split.SplitSource;
+import org.apache.doris.datasource.split.SplitSourceManager;
+import org.apache.doris.filesystem.FileEntry;
+import org.apache.doris.filesystem.FileIterator;
+import org.apache.doris.filesystem.FileSystem;
+import org.apache.doris.filesystem.Location;
+import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan.SelectedPartitions;
+import org.apache.doris.qe.QeProcessorImpl;
+import org.apache.doris.qe.SessionVariable;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TPushAggOp;
+import org.apache.doris.thrift.TScanRangeLocation;
+import org.apache.doris.thrift.TScanRangeLocations;
+import org.apache.doris.thrift.TUniqueId;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+class PluginDrivenScanNodeHiveStreamingTest {
+ @Test
+ @Timeout(30)
+ void hiveModeSelectionPreservesBoundedPartitionConcurrencyAndUniqueRanges() throws Exception {
+ Set<Integer> emptyPartitions = new HashSet<>(Arrays.asList(0, 4, 7));
+ try (Fixture fixture = new Fixture(emptyPartitions)) {
+ fixture.start();
+ Assertions.assertTrue(fixture.entered.await(10, TimeUnit.SECONDS),
+ "Two partition batches must reach cold listing concurrently through startSplit");
+ Assertions.assertEquals(2, fixture.listings.get());
+ fixture.release.countDown();
+
+ List<Long> actual = fixture.drain();
+ Set<Long> expected = expectedRanges(emptyPartitions);
+ Assertions.assertEquals(expected.size(), actual.size(), "Every range is emitted once");
+ Assertions.assertEquals(expected, new HashSet<>(actual));
+ Assertions.assertEquals(9, fixture.listings.get(), "No partition batch is omitted or loaded twice");
+ Assertions.assertEquals(2, fixture.peakListings.get(), "The producer count bounds remote listing");
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void fileCountStreamingDrainsThroughTheSharedProducerFinalizer() throws Exception {
+ Set<Integer> emptyPartitions = new HashSet<>(Arrays.asList(0, 4, 7));
+ try (Fixture fixture = new Fixture(emptyPartitions)) {
+ SessionVariable variables = Deencapsulation.getField(fixture.node, "sessionVariable");
+ Mockito.when(variables.getNumPartitionsInBatchMode()).thenReturn(1024);
+ fixture.release.countDown();
+ fixture.start();
+ Assertions.assertTrue((Boolean) Deencapsulation.getField(fixture.node, "streamingBatch"));
+
+ List<Long> actual = fixture.drain();
+ Set<Long> expected = expectedRanges(emptyPartitions);
+ Assertions.assertEquals(expected.size(), actual.size());
+ Assertions.assertEquals(expected, new HashSet<>(actual));
+ Assertions.assertEquals(9, fixture.listings.get());
+ Assertions.assertEquals(1, fixture.peakListings.get(), "Whole-scan streaming uses one producer");
+ }
+ }
+
+ private static Set<Long> expectedRanges(Set<Integer> emptyPartitions) {
+ Set<Long> expected = new HashSet<>();
+ for (int partition = 0; partition < 9; partition++) {
+ if (!emptyPartitions.contains(partition)) {
+ for (int offset = 0; offset < 17; offset++) {
+ expected.add(partition * 100L + offset);
+ }
+ }
+ }
+ return expected;
+ }
+
+ @Test
+ @Timeout(30)
+ void cancellationStopsHivePartitionBatchesBeforeFurtherListing() throws Exception {
+ try (Fixture fixture = new Fixture(Collections.emptySet())) {
+ fixture.start();
+ Assertions.assertTrue(fixture.entered.await(10, TimeUnit.SECONDS));
+ fixture.assignment.stop();
+ fixture.release.countDown();
+ fixture.executor.shutdown();
+ Assertions.assertTrue(fixture.executor.awaitTermination(10, TimeUnit.SECONDS));
+ Assertions.assertEquals(2, fixture.listings.get());
+ }
+ }
+
+ private static class Fixture implements AutoCloseable {
+ private static final AtomicInteger NEXT_QUERY = new AtomicInteger();
+ final CountDownLatch entered = new CountDownLatch(2);
+ final CountDownLatch release = new CountDownLatch(1);
+ final AtomicInteger listings = new AtomicInteger();
+ final AtomicInteger activeListings = new AtomicInteger();
+ final AtomicInteger peakListings = new AtomicInteger();
+ final ExecutorService executor = Executors.newFixedThreadPool(4);
+ final Backend backend = new Backend(0, "127.0.0.1", 9050);
+ final TUniqueId queryId = new TUniqueId(69133, NEXT_QUERY.incrementAndGet());
+ final int previousParallelism = Config.max_external_cache_loader_thread_pool_size;
+ final MockedStatic<Env> envStatic;
+ final PluginDrivenScanNode node;
+ final SplitAssignment assignment;
+
+ Fixture(Set<Integer> emptyPartitions) throws Exception {
+ FileSystem fs = Mockito.mock(FileSystem.class);
+ Mockito.when(fs.forLocation(Mockito.any())).thenReturn(fs);
+ Mockito.when(fs.list(Mockito.any())).thenAnswer(call -> {
+ Location location = call.getArgument(0);
+ listings.incrementAndGet();
+ peakListings.accumulateAndGet(activeListings.incrementAndGet(), Math::max);
+ entered.countDown();
+ try {
+ if (!release.await(20, TimeUnit.SECONDS)) {
+ throw new IOException("Timed out waiting for concurrent directory listing");
+ }
+ } finally {
+ activeListings.decrementAndGet();
+ }
+ int partition = Integer.parseInt(location.uri().substring(location.uri().lastIndexOf('=') + 1));
+ List<FileEntry> files = emptyPartitions.contains(partition) ? Collections.emptyList()
+ : Collections.singletonList(new FileEntry(Location.of(location.uri() + "/data.parquet"),
+ 17, false, 1, Collections.emptyList()));
+ Iterator<FileEntry> entries = files.iterator();
+ return new FileIterator() {
+ @Override
+ public boolean hasNext() {
+ return entries.hasNext();
+ }
+
+ @Override
+ public FileEntry next() {
+ return entries.next();
+ }
+
+ @Override
+ public void close() {
+ }
+ };
+ });
+ ConnectorStorageContext storage = Mockito.mock(ConnectorStorageContext.class, Mockito.CALLS_REAL_METHODS);
+ Mockito.when(storage.getFileSystem(Mockito.any())).thenReturn(fs);
+ ConnectorContext context = Mockito.mock(ConnectorContext.class);
+ Mockito.when(context.getStorageContext()).thenReturn(storage);
+ HmsClient hms = Mockito.mock(HmsClient.class, Mockito.CALLS_REAL_METHODS);
+ List<String> partitionNames = new ArrayList<>();
+ for (int partition = 0; partition < 9; partition++) {
+ partitionNames.add("pt=" + partition);
+ }
+ Mockito.when(hms.listPartitionNames(Mockito.anyString(), Mockito.anyString(), Mockito.anyInt()))
+ .thenReturn(partitionNames);
+ Mockito.when(hms.getPartitions(Mockito.anyString(), Mockito.anyString(), Mockito.anyList()))
+ .thenAnswer(call -> {
+ List<HmsPartitionInfo> partitions = new ArrayList<>();
+ for (String name : call.<List<String>>getArgument(2)) {
+ partitions.add(new HmsPartitionInfo(Collections.singletonList(name.substring(3)),
+ "s3://bucket/" + name, null, null, null, Collections.emptyMap()));
+ }
+ return partitions;
+ });
+ HiveCatalogProperties properties = HiveCatalogProperties.of(Collections.singletonMap(
+ HiveCatalogProperties.HIVE_METASTORE_URIS, "thrift://localhost:9083"));
+ HiveScanPlanProvider provider = new HiveScanPlanProvider(hms, properties, context,
+ new HiveReadTransactionManager(), new HiveFileListingCache(properties));
+ HiveTableHandle handle = new HiveTableHandle.Builder("db", "t", HiveTableType.HIVE)
+ .partitionKeyNames(Collections.singletonList("pt")).location("s3://bucket/")
+ .inputFormat("org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat")
+ .serializationLib("org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe").build();
+ ConnectorSession session = Mockito.mock(ConnectorSession.class, Mockito.CALLS_REAL_METHODS);
+ Mockito.when(session.getQueryId()).thenReturn(DebugUtil.printId(queryId));
+ Mockito.when(session.getStatementScope()).thenReturn(ConnectorStatementScope.NONE);
+ Mockito.when(session.getProperty("file_split_size", String.class)).thenReturn("1");
+ Mockito.when(session.getProperty("num_partitions_in_batch_mode", Integer.class)).thenReturn(2);
+ // The old implementation chooses its single-source streaming path for this same Hive scan.
+ Assertions.assertTrue(provider.streamingSplitEstimate(session, handle, Optional.empty(), false) >= 0);
+
+ node = Mockito.mock(PluginDrivenScanNode.class, Mockito.CALLS_REAL_METHODS);
+ Connector connector = Mockito.mock(Connector.class);
+ Mockito.when(connector.getScanPlanProvider(Mockito.any())).thenReturn(provider);
+ TupleDescriptor desc = Mockito.mock(TupleDescriptor.class);
+ Mockito.when(desc.getSlots()).thenReturn(new ArrayList<>(Collections.singletonList(
+ Mockito.mock(SlotDescriptor.class))));
+ SessionVariable variables = Mockito.mock(SessionVariable.class);
+ Mockito.when(variables.getNumPartitionsInBatchMode()).thenReturn(2);
+ Mockito.when(variables.getSplitBatchSize()).thenReturn(2);
+ Mockito.when(variables.getSplitQueueCapacity()).thenReturn(1);
+ Deencapsulation.setField(node, "connector", connector);
+ Deencapsulation.setField(node, "connectorSession", session);
+ Deencapsulation.setField(node, "currentHandle", handle);
+ Deencapsulation.setField(node, "cachedMetadata", Mockito.mock(ConnectorMetadata.class));
+ Deencapsulation.setField(node, "sessionVariable", variables);
+ Deencapsulation.setField(node, "desc", desc);
+ Mockito.doReturn(Mockito.mock(TableIf.class)).when(node).getTargetTable();
+ Mockito.doReturn(TPushAggOp.NONE).when(node).getPushDownAggNoGroupingOp();
+ Map<String, PartitionItem> selected = new LinkedHashMap<>();
+ for (int partition = 0; partition < 9; partition++) {
+ selected.put("pt=" + partition, Mockito.mock(PartitionItem.class));
+ }
+ node.setSelectedPartitions(new SelectedPartitions(9, selected, false));
+ FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class);
+ Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(call -> {
+ Multimap<Backend, Split> assigned = ArrayListMultimap.create();
+ assigned.putAll(backend, call.<List<Split>>getArgument(0));
+ return assigned;
+ });
+ assignment = new SplitAssignment(policy, Mockito.mock(SplitGenerator.class),
+ (be, props, split, keys, admission) -> {
+ HiveScanRange range = (HiveScanRange) ((PluginDrivenSplit) split).getConnectorScanRange();
+ long identity = Long.parseLong(range.getPartitionValues().get("pt")) * 100 + range.getStart();
+ return new TScanRangeLocations().setLocations(Collections.singletonList(
+ new TScanRangeLocation().setBackendId(identity)));
+ }, Collections.emptyMap(), Collections.emptyList(), true, 2, 1);
+ Deencapsulation.setField(node, "splitAssignment", assignment);
+ ExternalMetaCacheMgr manager = Mockito.mock(ExternalMetaCacheMgr.class);
+ Mockito.when(manager.getScheduleExecutor()).thenReturn(executor);
+ Env env = Mockito.mock(Env.class);
+ Mockito.when(env.getExtMetaCacheMgr()).thenReturn(manager);
+ Mockito.when(env.getSplitSourceManager()).thenReturn(Mockito.mock(SplitSourceManager.class));
+ envStatic = Mockito.mockStatic(Env.class);
+ envStatic.when(Env::getCurrentEnv).thenReturn(env);
+ Config.max_external_cache_loader_thread_pool_size = 2;
+ }
+
+ void start() {
+ Assertions.assertTrue(node.isBatchMode());
+ node.startSplit(1);
+ }
+
+ List<Long> drain() throws Exception {
+ SplitSource consumer = new SplitSource(backend, assignment, 1000);
+ List<Long> actual = new ArrayList<>();
+ for (int request = 0; request < 1000; request++) {
+ List<TScanRangeLocations> batch = consumer.getNextBatch(3);
+ Assertions.assertTrue(batch.size() <= 3);
+ if (batch.isEmpty()) {
+ return actual;
+ }
+ for (TScanRangeLocations range : batch) {
+ actual.add(range.getLocations().get(0).getBackendId());
+ }
+ }
+ throw new AssertionError("Hive streaming scan did not finish");
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ assignment.stop();
+ release.countDown();
+ QeProcessorImpl.INSTANCE.unregisterQuery(queryId);
+ executor.shutdown();
+ Assertions.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+ } finally {
+ executor.shutdownNow();
+ envStatic.close();
+ Config.max_external_cache_loader_thread_pool_size = previousParallelism;
+ }
+ }
+ }
+}
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
new file mode 100644
index 00000000000..36323a963f5
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodePartitionBatchTest.java
@@ -0,0 +1,211 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.datasource.scan;
+
+import org.apache.doris.common.UserException;
+import org.apache.doris.common.jmockit.Deencapsulation;
+import org.apache.doris.connector.spi.handle.ConnectorTableHandle;
+import org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider;
+import org.apache.doris.connector.spi.scan.ConnectorScanRange;
+import org.apache.doris.connector.spi.scan.ConnectorScanRequest;
+import org.apache.doris.datasource.split.SplitAssignment;
+import org.apache.doris.datasource.split.SplitGenerator;
+import org.apache.doris.datasource.split.SplitProducer;
+import org.apache.doris.datasource.split.SplitSource;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TScanRangeLocation;
+import org.apache.doris.thrift.TScanRangeLocations;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicInteger;
+
+class PluginDrivenScanNodePartitionBatchTest {
+ @Test
+ void fullQueueDoesNotLoadEveryPartitionBatch() throws Exception {
+ Fixture fixture = new Fixture(1, Collections.emptySet());
+ fixture.start(100, 1, 2);
+ fixture.executor.runUntilIdle();
+ Assertions.assertEquals(2, fixture.loaded.size(), "Paused producers must retain only two partition lists");
+ Assertions.assertEquals(1, fixture.assignment.getAssignedSplits(fixture.backend).size());
+ Assertions.assertEquals(0, fixture.finished.get());
+
+ List<Long> actual = fixture.drain();
+ List<Long> expected = new ArrayList<>();
+ for (long partition = 0; partition < 100; partition++) {
+ for (long range = 0; range < 3; range++) {
+ expected.add(partition * 10 + range);
+ }
+ }
+ Collections.sort(actual);
+ Assertions.assertEquals(expected, actual);
+ Assertions.assertEquals(100, fixture.loaded.size());
+ Assertions.assertEquals(100, new HashSet<>(fixture.loaded).size());
+ Assertions.assertEquals(2, fixture.finished.get());
+ }
+
+ @Test
+ void emptyPartitionBatchesAndPartialLastBatchDoNotEndTheScanEarly() throws Exception {
+ for (int parallelism : Arrays.asList(1, 8)) {
+ Fixture fixture = new Fixture(2, new HashSet<>(Arrays.asList("0", "1", "4")));
+ fixture.start(5, 2, parallelism);
+ List<Long> actual = fixture.drain();
+ Collections.sort(actual);
+ Assertions.assertEquals(Arrays.asList(20L, 21L, 22L, 30L, 31L, 32L), actual);
+ Assertions.assertEquals(Arrays.asList(Arrays.asList("0", "1"), Arrays.asList("2", "3"),
+ Collections.singletonList("4")), fixture.loaded);
+ Assertions.assertEquals(Math.min(parallelism, 3), fixture.finished.get(),
+ "Do not create more producers than partition batches");
+ }
+ }
+
+ @Test
+ void cancellationDuringEmptyPartitionLoadingStopsFurtherCalls() throws Exception {
+ Fixture fixture = new Fixture(1, Collections.emptySet());
+ Mockito.doAnswer(call -> {
+ fixture.loaded.add(new ArrayList<>(call.<List<String>>getArgument(2)));
+ fixture.assignment.stop();
+ return Collections.emptyList();
+ }).when(fixture.provider).planScanForPartitionBatch(Mockito.any(), Mockito.any(), Mockito.any());
+ fixture.start(100, 1, 2);
+ fixture.executor.runUntilIdle();
+ Assertions.assertEquals(Collections.singletonList(Collections.singletonList("0")), fixture.loaded);
+ Assertions.assertEquals(2, fixture.finished.get());
+ }
+
+ private static class Fixture {
+ final Backend backend = new Backend(0, "127.0.0.1", 9050);
+ final ManualExecutor executor = new ManualExecutor();
+ final List<List<String>> loaded = new ArrayList<>();
+ final AtomicInteger finished = new AtomicInteger();
+ final SplitAssignment assignment;
+ final PluginDrivenScanNode node;
+ final ConnectorScanPlanProvider provider;
+
+ Fixture(int batchSize, Set<String> emptyPartitions) throws Exception {
+ FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class);
+ Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(call -> {
+ Multimap<Backend, Split> result = ArrayListMultimap.create();
+ result.putAll(backend, call.<List<Split>>getArgument(0));
+ return result;
+ });
+ assignment = new SplitAssignment(policy, Mockito.mock(SplitGenerator.class),
+ (be, properties, split, keys, admission) -> new TScanRangeLocations().setLocations(
+ Collections.singletonList(new TScanRangeLocation().setBackendId(split.getStart()))),
+ Collections.emptyMap(), Collections.emptyList(), true, batchSize, 1);
+ node = Mockito.mock(PluginDrivenScanNode.class, Mockito.CALLS_REAL_METHODS);
+ Deencapsulation.setField(node, "splitAssignment", assignment);
+ provider = Mockito.mock(ConnectorScanPlanProvider.class, Mockito.CALLS_REAL_METHODS);
+ Mockito.when(provider.supportsBatchScan(Mockito.any(), Mockito.any())).thenReturn(true);
+ Mockito.doAnswer(call -> {
+ List<String> partitions = call.getArgument(2);
+ loaded.add(new ArrayList<>(partitions));
+ List<ConnectorScanRange> ranges = new ArrayList<>();
+ for (String partition : partitions) {
+ if (emptyPartitions.contains(partition)) {
+ continue;
+ }
+ for (int offset = 0; offset < 3; offset++) {
+ long identity = Long.parseLong(partition) * 10 + offset;
+ ranges.add(new ConnectorScanRange() {
+ @Override
+ public long getStart() {
+ return identity;
+ }
+
+ @Override
+ public Map<String, String> getProperties() {
+ return Collections.emptyMap();
+ }
+ });
+ }
+ }
+ return ranges;
+ }).when(provider).planScanForPartitionBatch(Mockito.any(), Mockito.any(), Mockito.any());
+ }
+
+ void start(int partitionCount, int partitionsPerBatch, int parallelism) {
+ List<String> partitions = new ArrayList<>();
+ for (int partition = 0; partition < partitionCount; partition++) {
+ partitions.add(String.valueOf(partition));
+ }
+ PluginDrivenScanNode.SubmittedTaskFinalizer finalizer =
+ new PluginDrivenScanNode.SubmittedTaskFinalizer(assignment::finishSchedule);
+ ConnectorScanRequest request = ConnectorScanRequest.builder(new ConnectorTableHandle() {
+ }, Collections.emptyList()).build();
+ for (SplitProducer.BatchSource source : node.createPartitionBatchSources(
+ provider, request, partitions, partitionsPerBatch, parallelism)) {
+ finalizer.taskSubmitted();
+ new SplitProducer(assignment, executor, source, failure -> {
+ if (failure != null) {
+ assignment.setException(new UserException(failure.getMessage(), failure));
+ }
+ finished.incrementAndGet();
+ finalizer.taskFinished();
+ }).start();
+ }
+ finalizer.closeDispatch();
+ }
+
+ List<Long> drain() throws Exception {
+ SplitSource consumer = new SplitSource(backend, assignment, 1000);
+ List<Long> identities = new ArrayList<>();
+ for (int request = 0; request < 1000; request++) {
+ executor.runUntilIdle();
+ List<TScanRangeLocations> ranges = consumer.getNextBatch(1);
+ if (ranges.isEmpty()) {
+ return identities;
+ }
+ Assertions.assertEquals(1, ranges.size());
+ identities.add(ranges.get(0).getLocations().get(0).getBackendId());
+ }
+ throw new AssertionError("Partition sources did not finish");
+ }
+ }
+
+ private static class ManualExecutor implements Executor {
+ final Queue<Runnable> tasks = new ArrayDeque<>();
+
+ @Override
+ public void execute(Runnable task) {
+ tasks.add(task);
+ }
+
+ void runUntilIdle() {
+ for (int task = 0; task < 1000 && !tasks.isEmpty(); task++) {
+ tasks.remove().run();
+ }
+ Assertions.assertTrue(tasks.isEmpty(), "Producers must yield when the queue is full");
+ }
+ }
+}
diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitAssignmentTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitAssignmentTest.java
index b3eb77158b3..2a40fce91dc 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitAssignmentTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitAssignmentTest.java
@@ -38,7 +38,11 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
public class SplitAssignmentTest {
@@ -76,7 +80,9 @@ public class SplitAssignmentTest {
mockSplitToScanRange,
locationProperties,
pathPartitionKeys,
- true
+ true,
+ 128,
+ 32
);
}
@@ -113,49 +119,20 @@ public class SplitAssignmentTest {
splitGeneratorThread.join(1000); // Wait for thread to complete
}
- @Test
- void testInitTimeout() throws Exception {
- // Use spy to simulate timeout behavior quickly instead of waiting 30 seconds
- SplitAssignment testAssignment = Mockito.spy(new SplitAssignment(
- mockBackendPolicy,
- mockSplitGenerator,
- mockSplitToScanRange,
- locationProperties,
- pathPartitionKeys,
- true
- ));
-
- Mockito.doThrow(new UserException("Failed to get first split after waiting for 0 seconds."))
- .when(testAssignment).init();
-
- // Test & Verify - should timeout immediately now
- UserException exception = Assertions.assertThrows(UserException.class, () -> testAssignment.init());
- Assertions.assertTrue(exception.getMessage().contains("Failed to get first split after waiting for"));
- }
-
@Test
void testInitInterrupted() throws Exception {
- CountDownLatch initStarted = new CountDownLatch(1);
- CountDownLatch shouldInterrupt = new CountDownLatch(1);
-
- Thread initThread = new Thread(() -> {
- try {
- initStarted.countDown();
- shouldInterrupt.await();
- splitAssignment.init();
- } catch (Exception e) {
- // Expected interruption
- }
- });
-
- initThread.start();
- initStarted.await();
-
- // Interrupt the init thread
- initThread.interrupt();
- shouldInterrupt.countDown();
-
- initThread.join(1000);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<UserException> result = executor.submit(() -> {
+ Thread.currentThread().interrupt();
+ return Assertions.assertThrows(UserException.class, splitAssignment::init);
+ });
+ Assertions.assertInstanceOf(InterruptedException.class, result.get(5, TimeUnit.SECONDS).getCause());
+ Mockito.verify(mockSplitGenerator).startSplit(mockBackendPolicy.numBackends());
+ } finally {
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
}
@Test
@@ -200,51 +177,112 @@ public class SplitAssignmentTest {
}
@Test
- void testAddToQueueSampleSplitAlreadySet() throws Exception {
+ void testConsumerCreatesBoundedQueueFirst() throws Exception {
+ BlockingQueue<Collection<TScanRangeLocations>> queue = splitAssignment.getAssignedSplits(mockBackend);
+ Assertions.assertEquals(32, queue.remainingCapacity());
+ Assertions.assertSame(queue, splitAssignment.getAssignedSplits(mockBackend));
+ }
+
+ @Test
+ void testSamplePublishedAfterScanParametersAreInitialized() throws Exception {
Multimap<Backend, Split> batch = ArrayListMultimap.create();
batch.put(mockBackend, mockSplit);
+ Mockito.when(mockBackendPolicy.computeScanRangeAssignment(Mockito.any())).thenReturn(batch);
+ CountDownLatch converting = new CountDownLatch(1);
+ CountDownLatch converted = new CountDownLatch(1);
+ Mockito.when(mockSplitToScanRange.getScanRange(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(),
+ Mockito.anyBoolean())).thenAnswer(invocation -> {
+ converting.countDown();
+ Assertions.assertTrue(converted.await(5, TimeUnit.SECONDS));
+ return mockScanRangeLocations;
+ });
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> producer = executor.submit(() -> {
+ splitAssignment.addToQueue(Collections.singletonList(mockSplit));
+ return null;
+ });
+ Assertions.assertTrue(converting.await(5, TimeUnit.SECONDS));
+ Assertions.assertNull(splitAssignment.getSampleSplit());
+ converted.countDown();
+ producer.get(5, TimeUnit.SECONDS);
+ Assertions.assertSame(mockSplit, splitAssignment.getSampleSplit());
+ } finally {
+ converted.countDown();
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+ @Test
+ void testProducerCreatesBoundedQueueFirst() throws Exception {
+ Multimap<Backend, Split> batch = ArrayListMultimap.create();
+ batch.put(mockBackend, mockSplit);
Mockito.when(mockBackendPolicy.computeScanRangeAssignment(Mockito.any())).thenReturn(batch);
Mockito.when(mockSplitToScanRange.getScanRange(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(),
Mockito.anyBoolean())).thenReturn(mockScanRangeLocations);
- // Setup: First call to set sample split
- List<Split> firstSplits = Collections.singletonList(mockSplit);
-
- splitAssignment.addToQueue(firstSplits);
- Split firstSampleSplit = splitAssignment.getSampleSplit();
+ splitAssignment.addToQueue(Collections.singletonList(mockSplit));
- // Test: Second call should not change sample split
- List<Split> secondSplits = Collections.singletonList(mockSplit);
+ BlockingQueue<Collection<TScanRangeLocations>> queue = splitAssignment.getAssignedSplits(mockBackend);
+ Assertions.assertEquals(1, queue.size());
+ Assertions.assertEquals(31, queue.remainingCapacity());
+ Assertions.assertSame(mockScanRangeLocations, queue.element().iterator().next());
+ }
- splitAssignment.addToQueue(secondSplits);
+ @Test
+ void testConsumerCreatedQueueBlocksProducerUntilSpaceAvailable() throws Exception {
+ BlockingQueue<Collection<TScanRangeLocations>> queue = splitAssignment.getAssignedSplits(mockBackend);
+ Collection<TScanRangeLocations> queuedBatch = Collections.singletonList(mockScanRangeLocations);
+ for (int i = 0; i < 32; i++) {
+ Assertions.assertTrue(queue.offer(queuedBatch));
+ }
+ Assertions.assertFalse(queue.offer(queuedBatch));
- // Verify sample split unchanged
- Assertions.assertEquals(firstSampleSplit, splitAssignment.getSampleSplit());
+ Multimap<Backend, Split> batch = ArrayListMultimap.create();
+ batch.put(mockBackend, mockSplit);
+ Mockito.when(mockBackendPolicy.computeScanRangeAssignment(Mockito.any())).thenReturn(batch);
+ CountDownLatch scanRangeCreated = new CountDownLatch(1);
+ Mockito.when(mockSplitToScanRange.getScanRange(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(),
+ Mockito.anyBoolean())).thenAnswer(invocation -> {
+ scanRangeCreated.countDown();
+ return mockScanRangeLocations;
+ });
+
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> producer = executor.submit(() -> {
+ splitAssignment.addToQueue(Collections.singletonList(mockSplit));
+ return null;
+ });
+ Assertions.assertTrue(scanRangeCreated.await(5, TimeUnit.SECONDS));
+ Assertions.assertThrows(TimeoutException.class, () -> producer.get(200, TimeUnit.MILLISECONDS));
+ Assertions.assertEquals(32, queue.size());
+
+ Assertions.assertSame(queuedBatch, queue.poll());
+ producer.get(5, TimeUnit.SECONDS);
+ Assertions.assertEquals(32, queue.size());
+ } finally {
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
}
@Test
- void testAddToQueueWithQueueBlockingScenario() throws Exception {
+ void testAddToQueueSampleSplitAlreadySet() throws Exception {
Multimap<Backend, Split> batch = ArrayListMultimap.create();
batch.put(mockBackend, mockSplit);
+ Split secondSplit = Mockito.mock(Split.class);
+ Multimap<Backend, Split> secondBatch = ArrayListMultimap.create();
+ secondBatch.put(mockBackend, secondSplit);
- Mockito.when(mockBackendPolicy.computeScanRangeAssignment(Mockito.any())).thenReturn(batch);
+ Mockito.when(mockBackendPolicy.computeScanRangeAssignment(Mockito.any())).thenReturn(batch, secondBatch);
Mockito.when(mockSplitToScanRange.getScanRange(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(),
Mockito.anyBoolean())).thenReturn(mockScanRangeLocations);
- // This test simulates a scenario where appendBatch might experience queue blocking
- // by adding multiple batches rapidly
-
- List<Split> splits = Collections.singletonList(mockSplit);
-
- // First, fill up the queue by adding many batches
- for (int i = 0; i < 10; i++) {
- splitAssignment.addToQueue(splits);
- }
-
- // Verify the queue has data
- BlockingQueue<Collection<TScanRangeLocations>> queue = splitAssignment.getAssignedSplits(mockBackend);
- Assertions.assertNotNull(queue);
+ splitAssignment.addToQueue(Collections.singletonList(mockSplit));
+ splitAssignment.addToQueue(Collections.singletonList(secondSplit));
+ Assertions.assertSame(mockSplit, splitAssignment.getSampleSplit());
}
@Test
@@ -329,36 +367,6 @@ public class SplitAssignmentTest {
splitsProvider.join(1000);
}
- // ==================== appendBatch() behavior tests ====================
-
- @Test
- void testAppendBatchTimeoutBehavior() throws Exception {
- Multimap<Backend, Split> batch = ArrayListMultimap.create();
- batch.put(mockBackend, mockSplit);
-
- Mockito.when(mockBackendPolicy.computeScanRangeAssignment(Mockito.any())).thenReturn(batch);
- Mockito.when(mockSplitToScanRange.getScanRange(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(),
- Mockito.anyBoolean())).thenReturn(mockScanRangeLocations);
-
- // This test verifies that appendBatch properly handles queue offer timeouts
- // We'll simulate this by first filling the assignment and then trying to add more
-
- List<Split> splits = Collections.singletonList(mockSplit);
-
- // Add multiple splits to potentially cause queue pressure
- for (int i = 0; i < 50; i++) {
- try {
- splitAssignment.addToQueue(splits);
- } catch (Exception e) {
- // Expected if queue gets full and times out
- break;
- }
- }
-
- // Verify that splits were processed
- Assertions.assertNotNull(splitAssignment.getSampleSplit());
- }
-
@Test
void testInitWhenNeedMoreSplitReturnsFalse() throws Exception {
// Test init behavior when needMoreSplit() returns false
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
new file mode 100644
index 00000000000..1dfe1b89c91
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
@@ -0,0 +1,411 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.datasource.split;
+
+import org.apache.doris.common.UserException;
+import org.apache.doris.common.util.LocationPath;
+import org.apache.doris.datasource.scan.FederationBackendPolicy;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TScanRangeLocation;
+import org.apache.doris.thrift.TScanRangeLocations;
+
+import com.google.common.collect.LinkedListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+class SplitProducerTest {
+ @Test
+ void twoScansInitializeOnOneWorkerBeforeAnyConsumerStarts() throws Exception {
+ Fixture first = new Fixture(1, 1);
+ Fixture second = new Fixture(1, 1);
+ ExecutorService worker = Executors.newSingleThreadExecutor();
+ ExecutorService planner = Executors.newSingleThreadExecutor();
+ RecordingSource firstSource = new RecordingSource(1, 0, 2, 4, 6, 8);
+ RecordingSource secondSource = new RecordingSource(1, 10, 12, 14, 16, 18);
+ Mockito.doAnswer(call -> {
+ first.start(worker, firstSource);
+ return null;
+ }).when(first.generator).startSplit(Mockito.anyInt());
+ Mockito.doAnswer(call -> {
+ second.start(worker, secondSource);
+ return null;
+ }).when(second.generator).startSplit(Mockito.anyInt());
+ try {
+ Future<?> initialized = planner.submit(() -> {
+ first.assignment.init();
+ second.assignment.init();
+ return null;
+ });
+ initialized.get(5, TimeUnit.SECONDS);
+ worker.submit(() -> { }).get(5, TimeUnit.SECONDS);
+ Assertions.assertEquals(1, first.assignment.getAssignedSplits(first.backend).size());
+ Assertions.assertEquals(1, second.assignment.getAssignedSplits(second.backend).size());
+ Assertions.assertEquals(Arrays.asList(0L, 2L, 4L, 6L, 8L), drain(first.source(), 2));
+ Assertions.assertEquals(Arrays.asList(10L, 12L, 14L, 16L, 18L), drain(second.source(), 1));
+ Assertions.assertTrue(firstSource.closed.await(5, TimeUnit.SECONDS));
+ Assertions.assertTrue(secondSource.closed.await(5, TimeUnit.SECONDS));
+ Assertions.assertEquals(1, firstSource.closeCalls.get());
+ Assertions.assertEquals(1, secondSource.closeCalls.get());
+ } finally {
+ first.assignment.stop();
+ second.assignment.stop();
+ planner.shutdownNow();
+ worker.shutdownNow();
+ Assertions.assertTrue(planner.awaitTermination(5, TimeUnit.SECONDS));
+ Assertions.assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void partialMultiBackendSubmissionDoesNotReplayOrReassignTheBatch() throws Exception {
+ Fixture fixture = new Fixture(2, 1);
+ ManualExecutor executor = new ManualExecutor();
+ fixture.assignment.addToQueue(Collections.singletonList(split(1)));
+ RecordingSource source = new RecordingSource(2, 0, 3);
+ fixture.start(executor, source);
+ executor.runNext();
+ Assertions.assertTrue(executor.tasks.isEmpty());
+ Assertions.assertEquals(0, fixture.completions.get());
+ Assertions.assertEquals(Collections.singletonList(0L), identities(fixture.source().getNextBatch(1)));
+ Assertions.assertTrue(executor.tasks.isEmpty());
+ SplitSource other = new SplitSource(fixture.otherBackend, fixture.assignment, 1000);
+ Assertions.assertEquals(Collections.singletonList(1L), identities(other.getNextBatch(1)));
+ executor.runUntilIdle();
+ Assertions.assertEquals(Collections.singletonList(3L), identities(other.getNextBatch(1)));
+ Assertions.assertTrue(other.getNextBatch(1).isEmpty());
+ Assertions.assertTrue(fixture.source().getNextBatch(1).isEmpty());
+ Mockito.verify(fixture.policy, Mockito.times(2)).computeScanRangeAssignment(Mockito.any());
+ Assertions.assertEquals(1, fixture.conversions.get(0L));
+ Assertions.assertEquals(1, fixture.conversions.get(3L));
+ Assertions.assertEquals(1, fixture.completions.get());
+ Assertions.assertEquals(1, source.closeCalls.get());
+ }
+
+ @Test
+ void capacityReleasedBeforeTheProducerSuspendsIsNotLost() throws Exception {
+ Fixture fixture = new Fixture(1, 1);
+ fixture.assignment = Mockito.spy(fixture.assignment);
+ fixture.assignment.addToQueue(Collections.singletonList(split(0)));
+ SplitSource consumer = fixture.source();
+ Mockito.doAnswer(call -> {
+ boolean offered = (boolean) call.callRealMethod();
+ if (!offered) {
+ Assertions.assertEquals(Collections.singletonList(0L), identities(consumer.getNextBatch(1)));
+ }
+ return offered;
+ }).when(fixture.assignment).offerBatch(Mockito.any(), Mockito.any(), Mockito.any());
+ ManualExecutor executor = new ManualExecutor();
+ RecordingSource source = new RecordingSource(1, 2);
+ fixture.start(executor, source);
+ executor.runNext();
+ Assertions.assertEquals(1, executor.tasks.size());
+ executor.runUntilIdle();
+ Assertions.assertEquals(Collections.singletonList(2L), identities(consumer.getNextBatch(1)));
+ Assertions.assertTrue(consumer.getNextBatch(1).isEmpty());
+ Assertions.assertEquals(1, source.closeCalls.get());
+ }
+
+ @Test
+ void cancellationClosesASuspendedProducerWithoutAConsumer() throws Exception {
+ Fixture fixture = new Fixture(1, 1);
+ fixture.assignment.addToQueue(Collections.singletonList(split(0)));
+ ManualExecutor executor = new ManualExecutor();
+ RecordingSource source = new RecordingSource(1, 2);
+ fixture.start(executor, source);
+ executor.runNext();
+ Assertions.assertTrue(executor.tasks.isEmpty());
+ Assertions.assertEquals(0, source.closeCalls.get());
+ fixture.assignment.stop();
+ executor.runUntilIdle();
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ Assertions.assertTrue(fixture.source().getNextBatch(1).isEmpty());
+ fixture.assignment.stop();
+ Assertions.assertTrue(executor.tasks.isEmpty());
+ }
+
+ @Test
+ void cancellationDoesNotCloseTheSourceDuringAnActiveRead() throws Exception {
+ Fixture fixture = new Fixture(1, 1);
+ CountDownLatch reading = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ RecordingSource source = new RecordingSource(1, 0) {
+ @Override
+ public List<Split> getNextBatch() throws IOException {
+ reading.countDown();
+ try {
+ Assertions.assertTrue(release.await(5, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException(e);
+ }
+ return super.getNextBatch();
+ }
+ };
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ try {
+ fixture.start(executor, source);
+ Assertions.assertTrue(reading.await(5, TimeUnit.SECONDS));
+ fixture.assignment.stop();
+ Assertions.assertEquals(0, source.closeCalls.get());
+ release.countDown();
+ Assertions.assertTrue(source.closed.await(5, TimeUnit.SECONDS));
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ } finally {
+ release.countDown();
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void rejectedCancellationCleanupStillClosesWithoutAddingAnError() throws Exception {
+ Fixture fixture = new Fixture(1, 1);
+ fixture.assignment.addToQueue(Collections.singletonList(split(0)));
+ ManualExecutor executor = new ManualExecutor();
+ RecordingSource source = new RecordingSource(1, 2);
+ fixture.start(executor, source);
+ executor.runNext();
+ executor.reject = true;
+ fixture.assignment.stop();
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ Assertions.assertTrue(fixture.source().getNextBatch(1).isEmpty());
+ }
+
+ @Test
+ void rejectedInitialSubmissionPublishesFailureAndClosesTheSource() throws Exception {
+ Fixture fixture = new Fixture(1, 1);
+ ManualExecutor executor = new ManualExecutor();
+ executor.reject = true;
+ RecordingSource source = new RecordingSource(1, 0);
+ fixture.start(executor, source);
+ UserException failure = Assertions.assertThrows(UserException.class,
+ () -> fixture.assignment.getAssignedSplits(fixture.backend));
+ Assertions.assertInstanceOf(RejectedExecutionException.class, failure.getCause());
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ }
+
+ @Test
+ void rejectedResumptionClosesThePendingSource() throws Exception {
+ Fixture fixture = new Fixture(1, 1);
+ fixture.assignment.addToQueue(Collections.singletonList(split(0)));
+ ManualExecutor executor = new ManualExecutor();
+ RecordingSource source = new RecordingSource(1, 2);
+ fixture.start(executor, source);
+ executor.runNext();
+ executor.reject = true;
+ UserException failure = Assertions.assertThrows(UserException.class, () -> fixture.source().getNextBatch(1));
+ Assertions.assertInstanceOf(RejectedExecutionException.class, failure.getCause());
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ }
+
+ @Test
+ void sourceFailureWakesOtherSuspendedProducers() throws Exception {
+ Fixture fixture = new Fixture(1, 1);
+ fixture.assignment.addToQueue(Collections.singletonList(split(0)));
+ ManualExecutor executor = new ManualExecutor();
+ RecordingSource pending = new RecordingSource(1, 2);
+ fixture.start(executor, pending);
+ executor.runNext();
+ RecordingSource failed = new RecordingSource(1) {
+ @Override
+ public List<Split> getNextBatch() throws IOException {
+ throw new IOException("listing failed");
+ }
+ };
+ fixture.start(executor, failed);
+ executor.runUntilIdle();
+ UserException failure = Assertions.assertThrows(UserException.class,
+ () -> fixture.assignment.getAssignedSplits(fixture.backend));
+ Assertions.assertInstanceOf(IOException.class, failure.getCause());
+ Assertions.assertEquals(1, pending.closeCalls.get());
+ Assertions.assertEquals(1, failed.closeCalls.get());
+ Assertions.assertEquals(2, fixture.completions.get());
+ }
+
+ @Test
+ void failureRacingWaiterRegistrationDoesNotLeaveTheSourceSuspended() throws Exception {
+ for (boolean beforeOffer : Arrays.asList(true, false)) {
+ Fixture fixture = new Fixture(1, 1);
+ fixture.assignment = Mockito.spy(fixture.assignment);
+ fixture.assignment.addToQueue(Collections.singletonList(split(0)));
+ UserException expected = new UserException("another producer failed");
+ Mockito.doAnswer(call -> {
+ if (beforeOffer) {
+ fixture.assignment.setException(expected);
+ }
+ boolean offered = (boolean) call.callRealMethod();
+ if (!beforeOffer) {
+ fixture.assignment.setException(expected);
+ }
+ return offered;
+ }).when(fixture.assignment).offerBatch(Mockito.any(), Mockito.any(), Mockito.any());
+ ManualExecutor executor = new ManualExecutor();
+ RecordingSource source = new RecordingSource(1, 2);
+ fixture.start(executor, source);
+ executor.runUntilIdle();
+ Assertions.assertSame(expected, Assertions.assertThrows(UserException.class,
+ () -> fixture.assignment.getAssignedSplits(fixture.backend)));
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ }
+ }
+
+ private static List<Long> drain(SplitSource source, int batchSize) throws Exception {
+ List<Long> result = new ArrayList<>();
+ while (true) {
+ List<TScanRangeLocations> batch = source.getNextBatch(batchSize);
+ Assertions.assertTrue(batch.size() <= batchSize);
+ if (batch.isEmpty()) {
+ return result;
+ }
+ result.addAll(identities(batch));
+ }
+ }
+
+ private static List<Long> identities(List<TScanRangeLocations> ranges) {
+ List<Long> result = new ArrayList<>();
+ for (TScanRangeLocations range : ranges) {
+ result.add(range.getLocations().get(0).getBackendId());
+ }
+ return result;
+ }
+
+ private static Split split(long identity) {
+ return new FileSplit(LocationPath.of("s3://bucket/file.parquet"), identity, 1, 100, 0,
+ null, Collections.emptyList());
+ }
+
+ private static class RecordingSource implements SplitProducer.BatchSource {
+ private final List<Split> splits = new ArrayList<>();
+ private final int batchSize;
+ private int next;
+ final AtomicInteger closeCalls = new AtomicInteger();
+ final CountDownLatch closed = new CountDownLatch(1);
+
+ RecordingSource(int batchSize, long... identities) {
+ this.batchSize = batchSize;
+ for (long identity : identities) {
+ splits.add(split(identity));
+ }
+ }
+
+ @Override
+ public List<Split> getNextBatch() throws IOException {
+ int end = Math.min(splits.size(), next + batchSize);
+ List<Split> batch = splits.subList(next, end);
+ next = end;
+ return batch;
+ }
+
+ @Override
+ public void close() {
+ closeCalls.incrementAndGet();
+ closed.countDown();
+ }
+ }
+
+ private static class ManualExecutor implements Executor {
+ final Queue<Runnable> tasks = new ArrayDeque<>();
+ boolean reject;
+
+ @Override
+ public void execute(Runnable task) {
+ if (reject) {
+ throw new RejectedExecutionException("executor stopped");
+ }
+ tasks.add(task);
+ }
+
+ void runNext() {
+ Assertions.assertFalse(tasks.isEmpty());
+ tasks.remove().run();
+ }
+
+ void runUntilIdle() {
+ for (int i = 0; i < 100 && !tasks.isEmpty(); i++) {
+ runNext();
+ }
+ Assertions.assertTrue(tasks.isEmpty(), "Producer must yield on a full queue or terminate at EOF");
+ }
+ }
+
+ private static class Fixture {
+ final Backend backend = new Backend(0, "127.0.0.1", 9050);
+ final Backend otherBackend = new Backend(1, "127.0.0.1", 9051);
+ final SplitGenerator generator = Mockito.mock(SplitGenerator.class);
+ final FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class);
+ final AtomicInteger completions = new AtomicInteger();
+ final Map<Long, Integer> conversions = new HashMap<>();
+ SplitAssignment assignment;
+
+ Fixture(int batchSize, int queueCapacity) throws Exception {
+ Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(call -> {
+ Multimap<Backend, Split> result = LinkedListMultimap.create();
+ for (Split split : call.<List<Split>>getArgument(0)) {
+ result.put(split.getStart() % 2 == 0 ? backend : otherBackend, split);
+ }
+ return result;
+ });
+ assignment = new SplitAssignment(policy, generator, (be, properties, split, keys, admission) -> {
+ conversions.merge(split.getStart(), 1, Integer::sum);
+ return new TScanRangeLocations().setLocations(Collections.singletonList(
+ new TScanRangeLocation().setBackendId(split.getStart())));
+ }, Collections.emptyMap(), Collections.emptyList(), true, batchSize, queueCapacity);
+ }
+
+ void start(Executor executor, RecordingSource source) {
+ new SplitProducer(assignment, executor, source, failure -> {
+ if (failure != null) {
+ assignment.setException(new UserException(failure.getMessage(), failure));
+ }
+ completions.incrementAndGet();
+ assignment.finishSchedule();
+ }).start();
+ }
+
+ SplitSource source() {
+ return new SplitSource(backend, assignment, 1000);
+ }
+ }
+}
diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitSourceTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitSourceTest.java
new file mode 100644
index 00000000000..5af23f18954
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitSourceTest.java
@@ -0,0 +1,207 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.datasource.split;
+
+import org.apache.doris.common.UserException;
+import org.apache.doris.datasource.scan.FederationBackendPolicy;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TScanRangeLocations;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class SplitSourceTest {
+ private Backend backend;
+ private SplitAssignment assignment;
+ private List<Split> splits;
+ private List<TScanRangeLocations> ranges;
+
+ @BeforeEach
+ void setUp() throws Exception {
+ backend = Mockito.mock(Backend.class);
+ FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class);
+ Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(invocation -> {
+ Multimap<Backend, Split> batch = ArrayListMultimap.create();
+ for (Split split : invocation.<List<Split>>getArgument(0)) {
+ batch.put(backend, split);
+ }
+ return batch;
+ });
+ splits = new ArrayList<>();
+ ranges = new ArrayList<>();
+ IdentityHashMap<Split, TScanRangeLocations> locations = new IdentityHashMap<>();
+ for (int i = 0; i < 8; i++) {
+ Split split = Mockito.mock(Split.class);
+ TScanRangeLocations range = Mockito.mock(TScanRangeLocations.class);
+ splits.add(split);
+ ranges.add(range);
+ locations.put(split, range);
+ }
+ assignment = new SplitAssignment(policy, Mockito.mock(SplitGenerator.class),
+ (be, properties, split, keys, admission) -> locations.get(split),
+ Collections.emptyMap(), Collections.emptyList(), true, 3, 3);
+ }
+
+ @Test
+ void drainPendingBatchBeforeFinishedQueueBecomesEof() throws Exception {
+ assignment.addToQueue(splits);
+ assignment.finishSchedule();
+ SplitSource source = new SplitSource(backend, assignment, 1000);
+
+ List<TScanRangeLocations> actual = new ArrayList<>();
+ for (int requested : Arrays.asList(2, 1, 4, 1)) {
+ List<TScanRangeLocations> batch = source.getNextBatch(requested);
+ Assertions.assertEquals(requested, batch.size());
+ actual.addAll(batch);
+ }
+ Assertions.assertEquals(ranges, actual);
+ Assertions.assertTrue(source.getNextBatch(2).isEmpty());
+ Assertions.assertTrue(source.getNextBatch(2).isEmpty());
+ }
+
+ @Test
+ void producerBatchesAndConsumerCreatedQueueRespectConfiguredLimits() throws Exception {
+ BlockingQueue<Collection<TScanRangeLocations>> queue = assignment.getAssignedSplits(backend);
+ Assertions.assertEquals(3, queue.remainingCapacity());
+ assignment.addToQueue(splits);
+ Assertions.assertEquals(3, queue.size());
+ Assertions.assertEquals(Arrays.asList(3, 3, 2),
+ queue.stream().map(Collection::size).collect(java.util.stream.Collectors.toList()));
+ }
+
+ @Test
+ void fullQueueResumesAfterConsumerFetch() throws Exception {
+ assignment.addToQueue(splits);
+ CountDownLatch started = new CountDownLatch(1);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> producer = executor.submit(() -> {
+ started.countDown();
+ assignment.addToQueue(Collections.singletonList(splits.get(0)));
+ return null;
+ });
+ Assertions.assertTrue(started.await(5, TimeUnit.SECONDS));
+ Assertions.assertThrows(TimeoutException.class, () -> producer.get(200, TimeUnit.MILLISECONDS));
+ SplitSource source = new SplitSource(backend, assignment, 1000);
+ Assertions.assertEquals(ranges.subList(0, 1), source.getNextBatch(1));
+ producer.get(5, TimeUnit.SECONDS);
+ Assertions.assertEquals(3, assignment.getAssignedSplits(backend).size());
+ } finally {
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void availablePartialBatchDoesNotWaitForProducer() throws Exception {
+ assignment.addToQueue(splits.subList(0, 1));
+ SplitSource source = new SplitSource(backend, assignment, 5000);
+ Assertions.assertTimeoutPreemptively(Duration.ofSeconds(1), () ->
+ Assertions.assertEquals(ranges.subList(0, 1), source.getNextBatch(100)));
+ }
+
+ @Test
+ void timeoutIsAnErrorRatherThanEof() {
+ SplitSource source = new SplitSource(backend, assignment, 20);
+ UserException failure = Assertions.assertThrows(UserException.class, () -> source.getNextBatch(1));
+ Assertions.assertTrue(failure.getMessage().contains("Timeout"));
+ }
+
+ @Test
+ void errorIsCheckedWhilePendingSplitsRemain() throws Exception {
+ assignment.addToQueue(splits.subList(0, 3));
+ SplitSource source = new SplitSource(backend, assignment, 1000);
+ source.getNextBatch(1);
+ UserException failure = new UserException("failed to generate splits");
+ assignment.setException(failure);
+ Assertions.assertSame(failure, Assertions.assertThrows(UserException.class, () -> source.getNextBatch(1)));
+ }
+
+ @Test
+ void concurrentRequestsDoNotDuplicateOrLosePendingSplits() throws Exception {
+ assignment.addToQueue(splits);
+ assignment.finishSchedule();
+ SplitSource source = new SplitSource(backend, assignment, 1000);
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ try {
+ Future<List<TScanRangeLocations>> first = executor.submit(() -> source.getNextBatch(2));
+ Future<List<TScanRangeLocations>> second = executor.submit(() -> source.getNextBatch(2));
+ List<TScanRangeLocations> actual = new ArrayList<>(first.get(5, TimeUnit.SECONDS));
+ actual.addAll(second.get(5, TimeUnit.SECONDS));
+ actual.addAll(source.getNextBatch(8));
+ Set<TScanRangeLocations> unique = Collections.newSetFromMap(new IdentityHashMap<>());
+ unique.addAll(actual);
+ Assertions.assertEquals(8, actual.size());
+ Assertions.assertEquals(8, unique.size());
+ Assertions.assertTrue(unique.containsAll(ranges));
+ } finally {
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void cancellationUnblocksProducerAndClosesLateResources() throws Exception {
+ assignment.addToQueue(splits);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> producer = executor.submit(() -> {
+ assignment.addToQueue(Collections.singletonList(splits.get(0)));
+ return null;
+ });
+ Assertions.assertThrows(TimeoutException.class, () -> producer.get(200, TimeUnit.MILLISECONDS));
+ assignment.stop();
+ producer.get(5, TimeUnit.SECONDS);
+ CountDownLatch closed = new CountDownLatch(1);
+ assignment.addCloseable(closed::countDown);
+ Assertions.assertEquals(0, closed.getCount());
+ Assertions.assertTrue(new SplitSource(backend, assignment, 1000).getNextBatch(1).isEmpty());
+ } finally {
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void invalidRequestSizeDoesNotReturnFalseEof() {
+ SplitSource source = new SplitSource(backend, assignment, 1000);
+ Assertions.assertThrows(IllegalArgumentException.class, () -> source.getNextBatch(0));
+ Assertions.assertThrows(IllegalArgumentException.class, () -> source.getNextBatch(-1));
+ }
+}
diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitStreamingConcurrencyTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitStreamingConcurrencyTest.java
new file mode 100644
index 00000000000..40c09b31d96
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitStreamingConcurrencyTest.java
@@ -0,0 +1,164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.datasource.split;
+
+import org.apache.doris.common.util.LocationPath;
+import org.apache.doris.datasource.scan.FederationBackendPolicy;
+import org.apache.doris.spi.Split;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TScanRangeLocation;
+import org.apache.doris.thrift.TScanRangeLocations;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+class SplitStreamingConcurrencyTest {
+ @Test
+ void consumerCreatedQueuesStayBoundedAcrossBackends() throws Exception {
+ verifyConcurrentStreaming(true);
+ }
+
+ @Test
+ void producerCreatedQueuesStayBoundedAcrossBackends() throws Exception {
+ verifyConcurrentStreaming(false);
+ }
+
+ private void verifyConcurrentStreaming(boolean consumerFirst) throws Exception {
+ int producerBatchSize = 7;
+ int queueCapacity = 2;
+ int splitCount = 10003;
+ int[] requestedSizes = {1, 10, 16};
+ List<Backend> backends = new ArrayList<>();
+ for (int i = 0; i < requestedSizes.length; i++) {
+ backends.add(new Backend(i, "127.0.0.1", 9050 + i));
+ }
+ FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class);
+ Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(invocation -> {
+ List<Split> batch = invocation.getArgument(0);
+ Assertions.assertTrue(batch.size() <= producerBatchSize);
+ Multimap<Backend, Split> assigned = ArrayListMultimap.create();
+ for (Split split : batch) {
+ assigned.put(backends.get((int) (split.getStart() % backends.size())), split);
+ }
+ return assigned;
+ });
+ CountDownLatch queuesFilled = new CountDownLatch(1);
+ SplitAssignment assignment = new SplitAssignment(policy, Mockito.mock(SplitGenerator.class),
+ (backend, properties, split, keys, admission) -> {
+ // Reaching batch three proves all queues received the first two batches.
+ if (split.getStart() >= producerBatchSize * queueCapacity) {
+ queuesFilled.countDown();
+ }
+ // Preserve a unique split identity through assignment and residual batches.
+ return new TScanRangeLocations().setLocations(Collections.singletonList(
+ new TScanRangeLocation().setBackendId(split.getStart())));
+ }, Collections.emptyMap(), Collections.emptyList(), true, producerBatchSize, queueCapacity);
+ List<Split> splits = new ArrayList<>();
+ LocationPath path = LocationPath.of("s3://bucket/file.parquet");
+ for (int i = 0; i < splitCount; i++) {
+ splits.add(new FileSplit(path, i, 1, splitCount, 0, null, Collections.emptyList()));
+ }
+ List<BlockingQueue<Collection<TScanRangeLocations>>> queues = new ArrayList<>();
+ if (consumerFirst) {
+ for (Backend backend : backends) {
+ queues.add(assignment.getAssignedSplits(backend));
+ }
+ }
+ ExecutorService executor = Executors.newFixedThreadPool(backends.size() + 1);
+ try {
+ Future<?> producer = executor.submit(() -> {
+ assignment.addToQueue(splits);
+ assignment.finishSchedule();
+ return null;
+ });
+ Assertions.assertTrue(queuesFilled.await(5, TimeUnit.SECONDS));
+ if (!consumerFirst) {
+ for (Backend backend : backends) {
+ queues.add(assignment.getAssignedSplits(backend));
+ }
+ }
+ for (BlockingQueue<Collection<TScanRangeLocations>> queue : queues) {
+ Assertions.assertEquals(queueCapacity, queue.size());
+ Assertions.assertEquals(0, queue.remainingCapacity());
+ verifyQueue(queue, producerBatchSize, queueCapacity);
+ }
+ List<Future<Set<Long>>> consumers = new ArrayList<>();
+ for (int i = 0; i < backends.size(); i++) {
+ int backendIndex = i;
+ SplitSource source = new SplitSource(backends.get(i), assignment, 10000);
+ consumers.add(executor.submit(() -> {
+ Set<Long> identities = new HashSet<>();
+ while (true) {
+ verifyQueue(queues.get(backendIndex), producerBatchSize, queueCapacity);
+ List<TScanRangeLocations> batch = source.getNextBatch(requestedSizes[backendIndex]);
+ Assertions.assertTrue(batch.size() <= requestedSizes[backendIndex]);
+ if (batch.isEmpty()) {
+ Assertions.assertTrue(source.getNextBatch(1).isEmpty());
+ return identities;
+ }
+ for (TScanRangeLocations range : batch) {
+ long identity = range.getLocations().get(0).getBackendId();
+ Assertions.assertEquals(backendIndex, identity % backends.size());
+ Assertions.assertTrue(identities.add(identity), "Duplicate split " + identity);
+ }
+ }
+ }));
+ }
+ producer.get(30, TimeUnit.SECONDS);
+ Set<Long> actual = new HashSet<>();
+ for (Future<Set<Long>> consumer : consumers) {
+ actual.addAll(consumer.get(30, TimeUnit.SECONDS));
+ }
+ Assertions.assertEquals(splitCount, actual.size());
+ for (long i = 0; i < splitCount; i++) {
+ Assertions.assertTrue(actual.contains(i), "Missing split " + i);
+ }
+ System.out.printf("Split streaming: consumerFirst=%s, backends=%d, splits=%d, "
+ + "producerBatchLimit=%d, observedQueuePeak=%d, requests=[1,10,16]%n",
+ consumerFirst, backends.size(), splitCount, producerBatchSize, queueCapacity);
+ } finally {
+ assignment.stop();
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ private void verifyQueue(BlockingQueue<Collection<TScanRangeLocations>> queue,
+ int producerBatchSize, int queueCapacity) {
+ Assertions.assertTrue(queue.size() <= queueCapacity);
+ for (Collection<TScanRangeLocations> batch : queue) {
+ Assertions.assertTrue(batch.size() <= producerBatchSize);
+ }
+ }
+}
diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java
index b978122e2f8..9c4f041db66 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java
@@ -443,6 +443,24 @@ public class SessionVariablesTest extends TestWithFeService {
Assertions.assertEquals(2, sv.ivfNprobe);
}
+ @Test
+ public void testSplitBatchAndQueueLimitsUseTheirOwnVariableCheckers() throws Exception {
+ SessionVariable variable = new SessionVariable();
+ for (String name : new String[] {SessionVariable.SPLIT_BATCH_SIZE, SessionVariable.SPLIT_QUEUE_CAPACITY}) {
+ Field field = VariableMgr.getVarContext(name).getField();
+ for (int value : new int[] {1, 9}) {
+ VariableMgr.setVar(variable, new SetVar(SetType.SESSION, name, new IntLiteral(value)));
+ Assertions.assertEquals(value, field.getInt(variable));
+ }
+ for (int value : new int[] {0, -1}) {
+ DdlException failure = Assertions.assertThrows(DdlException.class,
+ () -> VariableMgr.setVar(variable, new SetVar(SetType.SESSION, name, new IntLiteral(value))));
+ Assertions.assertTrue(failure.getMessage().contains(name + " value should greater than or equal 1"));
+ Assertions.assertEquals(9, field.getInt(variable));
+ }
+ }
+ }
+
@Test
public void testFileCacheQueryLimitBytesToThrift() throws Exception {
SessionVariable variable = new SessionVariable();
diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/VariableMgrExportTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/VariableMgrExportTest.java
new file mode 100644
index 00000000000..87dd3c178a8
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/qe/VariableMgrExportTest.java
@@ -0,0 +1,86 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.qe;
+
+import org.apache.doris.analysis.SetType;
+import org.apache.doris.analysis.SetVar;
+import org.apache.doris.analysis.StringLiteral;
+import org.apache.doris.connector.ConnectorSessionBuilder;
+import org.apache.doris.connector.spi.ConnectorSession;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+class VariableMgrExportTest {
+ @Test
+ void exportsCurrentValuesWithoutChangingPreviousSnapshots() throws Exception {
+ SessionVariable variables = new SessionVariable();
+ variables.batchSize = 100;
+ variables.enableProfile = false;
+ Map<String, String> first = VariableMgr.toMap(variables);
+
+ VariableMgr.setVar(variables, new SetVar(SetType.SESSION, SessionVariable.BATCH_SIZE,
+ new StringLiteral("200")));
+ variables.enableProfile = true;
+ Map<String, String> second = VariableMgr.toMap(variables);
+ Assertions.assertEquals("100", first.get(SessionVariable.BATCH_SIZE));
+ Assertions.assertEquals("false", first.get(SessionVariable.ENABLE_PROFILE));
+ Assertions.assertEquals("200", second.get(SessionVariable.BATCH_SIZE));
+ Assertions.assertEquals("true", second.get(SessionVariable.ENABLE_PROFILE));
+ }
+
+ @Test
+ void exportsAliasesAndExcludesRemovedAndInvisibleVariables() throws Exception {
+ SessionVariable variables = new SessionVariable();
+ VariableMgr.setVar(variables, new SetVar(SetType.SESSION, "enable_pushdown_string_minmax",
+ new StringLiteral("false")));
+ Map<String, String> values = VariableMgr.toMap(variables);
+ Assertions.assertEquals("false", values.get("enable_pushdown_string_minmax"));
+ Assertions.assertEquals("false", values.get(SessionVariable.FORCE_PUSHDOWN_ZONEMAP_MINMAX));
+ Assertions.assertFalse(values.containsKey(SessionVariable.ENABLE_LOCAL_EXCHANGE));
+ Assertions.assertFalse(values.containsKey(SessionVariable.ENABLE_NEREIDS_PLANNER));
+ }
+
+ @Test
+ void rereadsGlobalValuesOnEachExport() {
+ int original = GlobalVariable.lowerCaseTableNames;
+ try {
+ GlobalVariable.lowerCaseTableNames = 1;
+ Map<String, String> first = VariableMgr.toMap(new SessionVariable());
+ GlobalVariable.lowerCaseTableNames = 2;
+ Map<String, String> second = VariableMgr.toMap(new SessionVariable());
+ Assertions.assertEquals("1", first.get(GlobalVariable.LOWER_CASE_TABLE_NAMES));
+ Assertions.assertEquals("2", second.get(GlobalVariable.LOWER_CASE_TABLE_NAMES));
+ } finally {
+ GlobalVariable.lowerCaseTableNames = original;
+ }
+ }
+
+ @Test
+ void connectorSessionsObserveTheirOwnCurrentSettings() {
+ ConnectContext context = new ConnectContext();
+ context.getSessionVariable().batchSize = 100;
+ ConnectorSession first = ConnectorSessionBuilder.from(context).build();
+ context.getSessionVariable().batchSize = 200;
+ ConnectorSession second = ConnectorSessionBuilder.from(context).build();
+ Assertions.assertEquals(100, first.getProperty(SessionVariable.BATCH_SIZE, Integer.class));
+ Assertions.assertEquals(200, second.getProperty(SessionVariable.BATCH_SIZE, Integer.class));
+ }
+}
diff --git a/regression-test/data/external_table_p0/hive/test_hive_streaming_splits.out b/regression-test/data/external_table_p0/hive/test_hive_streaming_splits.out
new file mode 100644
index 00000000000..cd6e775b10b
--- /dev/null
+++ b/regression-test/data/external_table_p0/hive/test_hive_streaming_splits.out
@@ -0,0 +1,236 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !empty_unpartitioned --
+0 \N
+
+-- !empty_partitioned --
+0 \N
+
+-- !synchronous_unpartitioned_aggregate --
+120 7140 120 660
+
+-- !synchronous_unpartitioned_files --
+0 10 10 45 0 9
+1 10 10 145 10 19
+10 10 10 1045 100 109
+11 10 10 1145 110 119
+2 10 10 245 20 29
+3 10 10 345 30 39
+4 10 10 445 40 49
+5 10 10 545 50 59
+6 10 10 645 60 69
+7 10 10 745 70 79
+8 10 10 845 80 89
+9 10 10 945 90 99
+
+-- !synchronous_unpartitioned_limit --
+10
+
+-- !synchronous_partitioned_aggregate --
+120 7140 120 660 240
+
+-- !synchronous_partitioned_files --
+1 0 10 10 45 0 9
+1 3 10 10 345 30 39
+1 6 10 10 645 60 69
+1 9 10 10 945 90 99
+2 1 10 10 145 10 19
+2 10 10 10 1045 100 109
+2 4 10 10 445 40 49
+2 7 10 10 745 70 79
+3 11 10 10 1145 110 119
+3 2 10 10 245 20 29
+3 5 10 10 545 50 59
+3 8 10 10 845 80 89
+
+-- !synchronous_partition_range --
+80 5160 80 200
+
+-- !synchronous_single_partition_limit --
+10
+
+-- !synchronous_empty_partitions --
+0 \N
+
+-- !synchronous_join_aggregate --
+120 7140 240
+
+-- !synchronous_union_aggregate --
+240 14280
+
+-- !synchronous_union_limit --
+10
+
+-- !streaming_3_1_unpartitioned_aggregate --
+120 7140 120 660
+
+-- !streaming_3_1_unpartitioned_files --
+0 10 10 45 0 9
+1 10 10 145 10 19
+10 10 10 1045 100 109
+11 10 10 1145 110 119
+2 10 10 245 20 29
+3 10 10 345 30 39
+4 10 10 445 40 49
+5 10 10 545 50 59
+6 10 10 645 60 69
+7 10 10 745 70 79
+8 10 10 845 80 89
+9 10 10 945 90 99
+
+-- !streaming_3_1_unpartitioned_limit --
+10
+
+-- !streaming_3_1_partitioned_aggregate --
+120 7140 120 660 240
+
+-- !streaming_3_1_partitioned_files --
+1 0 10 10 45 0 9
+1 3 10 10 345 30 39
+1 6 10 10 645 60 69
+1 9 10 10 945 90 99
+2 1 10 10 145 10 19
+2 10 10 10 1045 100 109
+2 4 10 10 445 40 49
+2 7 10 10 745 70 79
+3 11 10 10 1145 110 119
+3 2 10 10 245 20 29
+3 5 10 10 545 50 59
+3 8 10 10 845 80 89
+
+-- !streaming_3_1_partition_range --
+80 5160 80 200
+
+-- !streaming_3_1_single_partition_limit --
+10
+
+-- !streaming_3_1_empty_partitions --
+0 \N
+
+-- !streaming_3_1_join_aggregate --
+120 7140 240
+
+-- !streaming_3_1_union_aggregate --
+240 14280
+
+-- !streaming_3_1_union_limit --
+10
+
+-- !streaming_7_2_unpartitioned_aggregate --
+120 7140 120 660
+
+-- !streaming_7_2_unpartitioned_files --
+0 10 10 45 0 9
+1 10 10 145 10 19
+10 10 10 1045 100 109
+11 10 10 1145 110 119
+2 10 10 245 20 29
+3 10 10 345 30 39
+4 10 10 445 40 49
+5 10 10 545 50 59
+6 10 10 645 60 69
+7 10 10 745 70 79
+8 10 10 845 80 89
+9 10 10 945 90 99
+
+-- !streaming_7_2_unpartitioned_limit --
+10
+
+-- !streaming_7_2_partitioned_aggregate --
+120 7140 120 660 240
+
+-- !streaming_7_2_partitioned_files --
+1 0 10 10 45 0 9
+1 3 10 10 345 30 39
+1 6 10 10 645 60 69
+1 9 10 10 945 90 99
+2 1 10 10 145 10 19
+2 10 10 10 1045 100 109
+2 4 10 10 445 40 49
+2 7 10 10 745 70 79
+3 11 10 10 1145 110 119
+3 2 10 10 245 20 29
+3 5 10 10 545 50 59
+3 8 10 10 845 80 89
+
+-- !streaming_7_2_partition_range --
+80 5160 80 200
+
+-- !streaming_7_2_single_partition_limit --
+10
+
+-- !streaming_7_2_empty_partitions --
+0 \N
+
+-- !streaming_7_2_join_aggregate --
+120 7140 240
+
+-- !streaming_7_2_union_aggregate --
+240 14280
+
+-- !streaming_7_2_union_limit --
+10
+
+-- !streaming_128_32_unpartitioned_aggregate --
+120 7140 120 660
+
+-- !streaming_128_32_unpartitioned_files --
+0 10 10 45 0 9
+1 10 10 145 10 19
+10 10 10 1045 100 109
+11 10 10 1145 110 119
+2 10 10 245 20 29
+3 10 10 345 30 39
+4 10 10 445 40 49
+5 10 10 545 50 59
+6 10 10 645 60 69
+7 10 10 745 70 79
+8 10 10 845 80 89
+9 10 10 945 90 99
+
+-- !streaming_128_32_unpartitioned_limit --
+10
+
+-- !streaming_128_32_partitioned_aggregate --
+120 7140 120 660 240
+
+-- !streaming_128_32_partitioned_files --
+1 0 10 10 45 0 9
+1 3 10 10 345 30 39
+1 6 10 10 645 60 69
+1 9 10 10 945 90 99
+2 1 10 10 145 10 19
+2 10 10 10 1045 100 109
+2 4 10 10 445 40 49
+2 7 10 10 745 70 79
+3 11 10 10 1145 110 119
+3 2 10 10 245 20 29
+3 5 10 10 545 50 59
+3 8 10 10 845 80 89
+
+-- !streaming_128_32_partition_range --
+80 5160 80 200
+
+-- !streaming_128_32_single_partition_limit --
+10
+
+-- !streaming_128_32_empty_partitions --
+0 \N
+
+-- !streaming_128_32_join_aggregate --
+120 7140 240
+
+-- !streaming_128_32_union_aggregate --
+240 14280
+
+-- !streaming_128_32_union_limit --
+10
+
+-- !ctas --
+120 7140 120 660
+
+-- !statement_hint --
+120 7140 1 1
+
+-- !settings_after_hint --
+128 32
+
diff --git a/regression-test/suites/external_table_p0/hive/test_hive_streaming_splits.groovy b/regression-test/suites/external_table_p0/hive/test_hive_streaming_splits.groovy
new file mode 100644
index 00000000000..7a0fca11193
--- /dev/null
+++ b/regression-test/suites/external_table_p0/hive/test_hive_streaming_splits.groovy
@@ -0,0 +1,135 @@
+// 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.
+
+suite("test_hive_streaming_splits", "p0,external") {
+ if (!"true".equalsIgnoreCase(context.config.otherConfigs.get("enableHiveTest"))) {
+ return
+ }
+
+ // An existing catalog also allows this suite to run against an isolated S3-backed HMS fixture.
+ String catalog = context.config.otherConfigs.get("hiveSplitCatalog")
+ if (!catalog) {
+ catalog = "hive_streaming_splits"
+ String host = context.config.otherConfigs.get("externalEnvIp")
+ String port = context.config.otherConfigs.get("hive3HmsPort")
+ sql "drop catalog if exists ${catalog}"
+ sql """create catalog ${catalog} properties (
+ 'type' = 'hms', 'hive.metastore.uris' = 'thrift://${host}:${port}')"""
+ }
+ def settings = ["enable_external_table_batch_mode", "num_files_in_batch_mode",
+ "num_partitions_in_batch_mode", "split_batch_size", "split_queue_capacity"]
+ def saved = settings.collectEntries { name -> [(name): sql("show variables like '${name}'")[0][1]] }
+ try {
+ sql "create database if not exists ${catalog}.hive_streaming_splits"
+ sql "use ${catalog}.hive_streaming_splits"
+ // Truncate while the previous partition metadata still exists, so repeated runs clear its files.
+ if (!sql("show tables like 'streaming_unpartitioned'").isEmpty()) {
+ sql "truncate table streaming_unpartitioned"
+ }
+ if (!sql("show tables like 'streaming_partitioned'").isEmpty()) {
+ sql "truncate table streaming_partitioned"
+ }
+ if (!sql("show tables like 'streaming_ctas'").isEmpty()) {
+ sql "truncate table streaming_ctas"
+ }
+ sql "drop table if exists streaming_unpartitioned"
+ sql "drop table if exists streaming_partitioned"
+ sql "drop table if exists streaming_ctas"
+ sql """create table streaming_unpartitioned (id bigint, file_id int)
+ engine=hive properties ('file_format'='parquet')"""
+ sql """create table streaming_partitioned (id bigint, file_id int, p int)
+ engine=hive partition by list (p) () properties ('file_format'='parquet')"""
+ sql "set enable_external_table_batch_mode=true"
+ sql "set num_files_in_batch_mode=1"
+ sql "set num_partitions_in_batch_mode=1"
+ qt_empty_unpartitioned "select count(*), sum(id) from streaming_unpartitioned"
+ qt_empty_partitioned "select count(*), sum(id) from streaming_partitioned"
+ // Separate inserts create enough files to cross several small producer batches.
+ for (int file = 0; file < 12; file++) {
+ sql """insert into streaming_unpartitioned
+ select number + ${file * 10}, ${file} from numbers('number'='10')"""
+ sql """insert into streaming_partitioned
+ select number + ${file * 10}, ${file}, ${file % 3 + 1} from numbers('number'='10')"""
+ }
+ def checkResults = { String label ->
+ quickTest("${label}_unpartitioned_aggregate", """select count(*), sum(id), count(distinct id), sum(file_id)
+ from streaming_unpartitioned""", true)
+ quickTest("${label}_unpartitioned_files", """select file_id, count(*), count(distinct id),
+ sum(id), min(id), max(id) from streaming_unpartitioned group by file_id order by file_id""", true)
+ quickTest("${label}_unpartitioned_limit", """select count(*) from
+ (select id from streaming_unpartitioned limit 10) limited""", true)
+ quickTest("${label}_partitioned_aggregate", """select count(*), sum(id), count(distinct id), sum(file_id), sum(p)
+ from streaming_partitioned""", true)
+ quickTest("${label}_partitioned_files", """select p, file_id, count(*), count(distinct id),
+ sum(id), min(id), max(id) from streaming_partitioned group by p, file_id order by p, file_id""", true)
+ quickTest("${label}_partition_range", """select count(*), sum(id), count(distinct id), sum(p)
+ from streaming_partitioned where p between 2 and 3""", true)
+ quickTest("${label}_single_partition_limit", """select count(*) from
+ (select id from streaming_partitioned where p=1 limit 10) limited""", true)
+ quickTest("${label}_empty_partitions", """select count(*), sum(id)
+ from streaming_partitioned where p>3""", true)
+ quickTest("${label}_join_aggregate", """select count(*), sum(u.id), sum(p.p)
+ from streaming_unpartitioned u join streaming_partitioned p on u.id=p.id""", true)
+ quickTest("${label}_union_aggregate", """select count(*), sum(id) from
+ (select id from streaming_unpartitioned union all
+ select id from streaming_partitioned) combined""", true)
+ quickTest("${label}_union_limit", """select count(*) from
+ (select id from streaming_unpartitioned union all
+ select id from streaming_partitioned limit 10) limited""", true)
+ }
+ sql "set enable_external_table_batch_mode=false"
+ sql "set num_partitions_in_batch_mode=0"
+ explain {
+ sql "select * from streaming_partitioned limit 10"
+ notContains "(approximate)inputSplitNum"
+ }
+ checkResults("synchronous")
+ sql "set enable_external_table_batch_mode=true"
+ sql "set num_partitions_in_batch_mode=1"
+ for (def limits : [[3, 1], [7, 2], [128, 32]]) {
+ sql "set split_batch_size=${limits[0]}"
+ sql "set split_queue_capacity=${limits[1]}"
+ explain {
+ sql "select * from streaming_unpartitioned limit 10"
+ contains "(approximate)inputSplitNum"
+ }
+ explain {
+ sql "select * from streaming_partitioned where p=1 limit 10"
+ contains "(approximate)inputSplitNum"
+ }
+ checkResults("streaming_${limits[0]}_${limits[1]}")
+ }
+ // CTAS probes an absent target before creating and resolving it in the same statement.
+ sql """create table streaming_ctas properties ('file_format'='parquet')
+ as select id, file_id from streaming_unpartitioned"""
+ qt_ctas "select count(*), sum(id), count(distinct id), sum(file_id) from streaming_ctas"
+ qt_statement_hint """select /*+ SET_VAR(split_batch_size=1, split_queue_capacity=1) */
+ count(*), sum(id), max(@@split_batch_size), max(@@split_queue_capacity)
+ from streaming_unpartitioned"""
+ qt_settings_after_hint "select @@split_batch_size, @@split_queue_capacity"
+ test {
+ sql "set split_batch_size=0"
+ exception "split_batch_size value should greater than or equal 1"
+ }
+ test {
+ sql "set split_queue_capacity=-1"
+ exception "split_queue_capacity value should greater than or equal 1"
+ }
+ } finally {
+ saved.each { name, value -> sql "set ${name}=${value}" }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@doris.apache.org
For additional commands, e-mail: commits-help@doris.apache.org