diff --git a/fe/fe-common/src/main/java/org/apache/doris/common/Config.java b/fe/fe-common/src/main/java/org/apache/doris/common/Config.java index d032e271112078..3d1cecbf2d3929 100644 --- a/fe/fe-common/src/main/java/org/apache/doris/common/Config.java +++ b/fe/fe-common/src/main/java/org/apache/doris/common/Config.java @@ -2209,6 +2209,15 @@ public class Config extends ConfigBase { @ConfField(mutable = false, masterOnly = false) public static int max_external_cache_loader_thread_pool_size = 64; + @ConfField(description = "Workers for fair split production") + public static int external_scan_split_threads = 64; + + @ConfField(description = "Workers for independent metadata tasks") + public static int external_scan_metadata_threads = 64; + + @ConfField(description = "Maximum queued tasks per query in each external scan scheduler") + public static int external_scan_task_queue_capacity = 128; + /** * Max cache num of external catalog's file * Decrease this value if FE's memory is small diff --git a/fe/fe-connector/README.md b/fe/fe-connector/README.md index 89d1a0fb422070..06c95738a01b18 100644 --- a/fe/fe-connector/README.md +++ b/fe/fe-connector/README.md @@ -120,6 +120,9 @@ tests and embedded use. ## Reading the API +For the scan task scheduler SPI, admission, cancellation and configuration, see +[Cooperative external scan scheduling](docs/scan-task-scheduling.md). + This document never lists SPI methods. The truth lives in code, behind four mechanisms: diff --git a/fe/fe-connector/docs/scan-task-scheduling.md b/fe/fe-connector/docs/scan-task-scheduling.md new file mode 100644 index 00000000000000..ebfb05f5a39f67 --- /dev/null +++ b/fe/fe-connector/docs/scan-task-scheduling.md @@ -0,0 +1,125 @@ +# Cooperative external scan scheduling + +A statement obtains its stage schedulers through +`ConnectorSession.getTaskScheduler`. The engine owns the task groups and closes +them with the statement. Connectors submit neutral `Runnable` tasks through this +SPI; they do not import FE executor classes. + +The two stages have independent workers: + +- `SPLIT` advances scan producers and assigns ranges to the BE split queues. +- `METADATA` performs independent operations such as cold Manifest reads. + +A producer may need metadata tasks to finish. Keeping those tasks out of the +producer pool avoids a pool full of parents waiting for children queued behind +them. A task must not block waiting for children submitted to its own executor. + +## Admission and continuation + +Each executor has a concurrent FIFO ready queue containing task groups. A +worker claims a task from the next group, republishes that group if it still +has pending work, and then executes the task outside queue synchronization. +This allows several workers to serve one query and gives other ready queries +an opportunity at subsequent task boundaries. It does not preempt a running +task or guarantee equal CPU time. + +`tryExecute` never waits for space. A full group queue registers the supplied +`onSpace` callback and returns `false`; the caller retains the task until it +can retry. Reuse the same callback identity and remove it when the task is +cancelled or no longer needs admission. Failures are delivered through the +failure callback rather than silently dropping work. + +A serial producer can use `tryExecuteContinuation` for bounded further work. +The group still returns behind other ready groups, but a worker need not wake +another worker just to execute its own serial continuation. Initial tasks, +external admission and recovery from backpressure retain the normal wakeup +protocol. Split production yields after at most 16 source batches or after +observing 1 ms elapsed between batches. One blocking batch can exceed that +interval; the bound is cooperative. + +## Iceberg Manifest reads + +Cold Manifest prefetch uses a bounded window and a statement-wide read budget. +Within that window, completed reads can be consumed before a slower earlier +Manifest. First-split production does not wait for speculative prefetch when +one Manifest already supplies output; subsequent BE demand enables further +prefetch. Filtered or rewrite-scoped scans must still advance while no output +has been found. + +Closing a scan stops new submissions and wakes completion waiters. Load failures +propagate through the corresponding future. A running shared cache load is not +interrupted just because one waiting scan stops. Authentication and plugin +thread-context class loading remain scoped around each actual metadata read. +A scheduled reader that finds another query already loading the same Manifest +observes that shared completion and returns its metadata worker immediately. +The completion callback only publishes the bounded window entry; it does not +read another Manifest or run split production on the loading query's thread. +Returned asynchronous tasks remain owned by their task group until terminal +completion, so closing the group still wakes the waiting producer. Closing a +scan detaches its observer without cancelling the shared cache load. Cache +publication continues to use the original scope/key generation fence. +Cached partition values stay immutable: adapt a file task to the pinned scan +schema at the scan boundary. Select the format from each file range rather +than planning files merely to infer a table-wide format. + +The captured scheduler exposes its closed state even after query cleanup removes +the statement registry. A Manifest reader waiting for another scan's read permit +checks this state at most 100 ms apart and rechecks after admission. Cancellation +therefore stops that waiter without releasing another scan's permit or starting +another remote read. This does not interrupt an already running shared cache load. +Hive checks the same captured state between directory requests, including when +empty partitions are skipped within one iterator call. + +## Configuration and limits + +| Setting | Default | Meaning | +|---|---:|---| +| FE `external_scan_split_threads` | 64 | Worker threads in the producer scheduler | +| FE `external_scan_metadata_threads` | 64 | Worker threads in the metadata scheduler | +| FE `external_scan_task_queue_capacity` | 128 | Pending tasks per query in each stage | +| Session `iceberg_manifest_load_concurrency` | 8 | Manifest read budget per query and window bound per scan; 1 uses sequential loading | + +These task queues are distinct from `split_queue_capacity`, which bounds the +produced Split buffer. Queue capacity does not reserve that many running +threads. Worker counts limit the respective pools, and the query read budget +also covers caller-assisted Manifest loading. + +Workers are ordinary threads. Synchronous remote calls occupy workers until +they return; FIFO admission cannot recover a worker blocked inside an SDK. +Effective remote timeouts and dependency fault isolation are separate concerns. + +Planning-only attempts also run their query-finish callbacks when statement +execution returns. This releases batch Split producers and metadata leases for +EXPLAIN and planning failures even when no coordinator was registered. Normal +coordinator unregistration remains idempotent. Deferred Arrow Flight queries keep +those callbacks until their retained executor finalizes the query after consumption. + +Backend assignment has its own policy-state lock. First-split readiness and stop +notifications use a separate short-held lock, so an already converted range can +wake the planner while another lane is still assigning its batch. The policy's +mutable load-balancing state remains serialized. + +The engine captures task-group cancellation on the statement executor before +asynchronous work starts. KILL and timeout close those groups even before a +coordinator exists; a group captured after cancellation is immediately closed. +Cancellation listeners run outside their registry operation and are claimed once. +This stops queued work and wakes future waiters without interrupting a running +shared cache loader. Normal completion retains the separate read-lease cleanup. + +### Consumer completion notification + +Each backend has a bounded split queue and its own data-ready wait channel. A consumer registers before checking queue contents and terminal state, then waits only when neither is available. Successful enqueue, producer completion, failure and cancellation signal that channel. The empty queue is not periodically polled to discover EOF: otherwise finishing just after a poll begins adds 100 ms despite completed production. Producers avoid the channel monitor when no consumer waits; callbacks that resume producers run after releasing the consumer channel monitor. Cancellation signals consumers before closing SDK resources. + +### Initial delivery and producer lookahead + +`split_initial_batch_size` (default 64) bounds the first Split fetch per backend, also capped by the BE request size. Subsequent fetches use the BE request size. The bound is applied before compatible file ranges are compacted into RPC envelopes. It is independent of `split_batch_size` (producer batch size) and `split_queue_capacity` (buffered batches per backend). Producers continue listing files and preparing splits before BE consumption, yielding only for task fairness and configured queue backpressure. A small initial RPC therefore does not suspend useful directory prefetch while BE starts. + +Actual Manifest readers retain a lease on the exact read-table generation until the callable exits. Query cleanup releases the statement owner, while an already running shared reader keeps its FileIO alive. Future cancellation is not evidence that the callable has finished. + +## Batch mode switch + +`enable_external_table_batch_mode` governs both partition-batch and file-streaming +planning. The engine checks it before calling provider capability or file-count +estimation methods. Setting it to `false` selects synchronous split planning for +connectors that otherwise support batch mode. Connectors without a batch or +streaming implementation keep their synchronous path regardless of this switch. diff --git a/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/MetaCache.java b/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/MetaCache.java index f3b425e0463876..387095f14e83fe 100644 --- a/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/MetaCache.java +++ b/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/MetaCache.java @@ -19,6 +19,7 @@ import java.util.Objects; import java.util.Optional; +import java.util.concurrent.CompletionStage; import java.util.function.BiConsumer; import java.util.function.BooleanSupplier; import java.util.function.Consumer; @@ -51,6 +52,16 @@ public V get(K key, Function loader) { return delegate.get(nonNullKey, definition.scope(nonNullKey), loader); } + /** + * Loads on the caller when elected, or observes an existing load without blocking this caller. + * Cancelling or completing a future obtained from the stage does not affect the shared load. + * Callbacks run on the completing thread and must not perform blocking work. + */ + public CompletionStage getOrLoadFuture(K key, Function loader) { + K nonNullKey = Objects.requireNonNull(key, "key can not be null"); + return delegate.getOrLoadFuture(nonNullKey, definition.scope(nonNullKey), loader); + } + /** * Loads a missing value while allowing an adapter to publish an auxiliary side effect immediately before the * cache value at the same guarded commit point. The coordinator must invoke the supplied commit callback once; diff --git a/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/ScopedMetaCache.java b/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/ScopedMetaCache.java index b342dfa5023964..faf4b83588af17 100644 --- a/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/ScopedMetaCache.java +++ b/fe/fe-connector/fe-connector-cache/src/main/java/org/apache/doris/connector/cache/ScopedMetaCache.java @@ -43,6 +43,7 @@ import java.util.TreeMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.CompletionStage; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.Executor; @@ -181,6 +182,23 @@ public V get(K key, ScopePath path, Function loader) { public V getWithPublicationAction(K key, ScopePath path, Function loader, BiConsumer> publicationCoordinator) { + return getWithPublicationAction(key, path, loader, publicationCoordinator, + Function.identity(), this::awaitLoad); + } + + /** + * Runs an elected loader on the caller, but returns without waiting when another caller owns the load. + * The returned stage cannot complete or cancel the shared load. Completion callbacks must stay short: + * they may execute on the elected loader before its publication lease is released. + */ + public CompletionStage getOrLoadFuture(K key, ScopePath path, Function loader) { + return getWithPublicationAction(key, path, loader, (loaded, commit) -> commit.accept(NO_OP), + CompletableFuture::completedFuture, CompletableFuture::minimalCompletionStage); + } + + private R getWithPublicationAction(K key, ScopePath path, Function loader, + BiConsumer> publicationCoordinator, Function loadedResult, + Function, R> sharedResult) { Objects.requireNonNull(key, "key can not be null"); Objects.requireNonNull(path, "path can not be null"); Function loadFunction = Objects.requireNonNull(loader, "loader can not be null"); @@ -206,14 +224,14 @@ public V getWithPublicationAction(K key, ScopePath path, Function loader, notifyDiscarded(key, loaded); } } - return loaded; + return loadedResult.apply(loaded); } VersionedValue presentVersioned = currentVersionedValue(key, path); if (presentVersioned != null) { recordAccess(true); scheduleRefresh(key, path, loader, presentVersioned); - return presentVersioned.value; + return loadedResult.apply(presentVersioned.value); } recordAccess(false); try (PublicationLease lease = acquirePublicationLease(key, path, true)) { @@ -221,7 +239,7 @@ public V getWithPublicationAction(K key, ScopePath path, Function loader, CompletableFuture ownLoad = new CompletableFuture<>(); CompletableFuture existingLoad = inFlightLoads.putIfAbsent(loadAddress, ownLoad); if (existingLoad != null) { - return awaitLoad(existingLoad); + return sharedResult.apply(existingLoad); } try { afterLoadElection.run(); @@ -237,7 +255,7 @@ public V getWithPublicationAction(K key, ScopePath path, Function loader, if (electedValuePresent) { VersionedValue present = electedValue.get(); ownLoad.complete(present.value); - return present.value; + return loadedResult.apply(present.value); } V loaded = loadAndRecord(key, loadFunction); if (loaded != null) { @@ -261,7 +279,7 @@ public V getWithPublicationAction(K key, ScopePath path, Function loader, } } ownLoad.complete(loaded); - return loaded; + return loadedResult.apply(loaded); } catch (RuntimeException | Error throwable) { ownLoad.completeExceptionally(throwable); throw throwable; diff --git a/fe/fe-connector/fe-connector-cache/src/test/java/org/apache/doris/connector/cache/MetaCacheSharedLoadTest.java b/fe/fe-connector/fe-connector-cache/src/test/java/org/apache/doris/connector/cache/MetaCacheSharedLoadTest.java new file mode 100644 index 00000000000000..ca865de8565ff7 --- /dev/null +++ b/fe/fe-connector/fe-connector-cache/src/test/java/org/apache/doris/connector/cache/MetaCacheSharedLoadTest.java @@ -0,0 +1,203 @@ +// 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.cache; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; + +class MetaCacheSharedLoadTest { + private static MetaCache cache(CatalogMetaCache owner, boolean enabled) { + return owner.create(MetaCacheDefinition.builder("shared-stage", + CacheSpec.of(enabled, -1, 100), ignored -> ScopePath.table("db", "tbl")).build()); + } + + private static void await(CountDownLatch latch) { + try { + Assertions.assertTrue(latch.await(10, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + + @Test + void sharedObserverReturnsBeforeTheElectedLoadAndCannotCancelIt() throws Exception { + ExecutorService workers = Executors.newFixedThreadPool(2); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger loads = new AtomicInteger(); + try (CatalogMetaCache owner = CatalogMetaCache.unmanaged()) { + MetaCache cache = cache(owner, true); + Function loader = key -> { + loads.incrementAndGet(); + entered.countDown(); + await(release); + return "value"; + }; + Future> elected = workers.submit(() -> cache.getOrLoadFuture("key", loader)); + await(entered); + CompletionStage shared = workers.submit(() -> cache.getOrLoadFuture("key", loader)) + .get(2, TimeUnit.SECONDS); + Assertions.assertFalse(shared.toCompletableFuture().isDone()); + Assertions.assertEquals("other query", workers.submit(() -> "other query").get(2, TimeUnit.SECONDS)); + Assertions.assertTrue(shared.toCompletableFuture().cancel(false)); + CompletableFuture survivor = shared.toCompletableFuture(); + CompletableFuture brokenCallback = shared.thenApply(value -> { + throw new IllegalArgumentException("observer failed"); + }).toCompletableFuture(); + release.countDown(); + Assertions.assertEquals("value", elected.get(5, TimeUnit.SECONDS).toCompletableFuture().get()); + Assertions.assertEquals("value", survivor.get(5, TimeUnit.SECONDS)); + Assertions.assertThrows(ExecutionException.class, brokenCallback::get); + Assertions.assertEquals(1, loads.get()); + Assertions.assertEquals("value", cache.getIfPresent("key")); + Assertions.assertEquals("value", cache.getOrLoadFuture("key", key -> { + throw new AssertionError("cached value must not load"); + }).toCompletableFuture().get()); + } finally { + release.countDown(); + workers.shutdownNow(); + Assertions.assertTrue(workers.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + void sharedFailureReachesObserversAndAllowsRetry() throws Exception { + for (Throwable expected : new Throwable[] {new IllegalStateException("remote failure"), + new AssertionError("remote error")}) { + ExecutorService worker = Executors.newSingleThreadExecutor(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + try (CatalogMetaCache owner = CatalogMetaCache.unmanaged()) { + MetaCache cache = cache(owner, true); + Future elected = worker.submit(() -> cache.getOrLoadFuture("key", key -> { + entered.countDown(); + await(release); + if (expected instanceof Error) { + throw (Error) expected; + } + throw (RuntimeException) expected; + })); + await(entered); + CompletionStage shared = cache.getOrLoadFuture("key", key -> { + throw new AssertionError("duplicate load"); + }); + release.countDown(); + Assertions.assertSame(expected, + Assertions.assertThrows(ExecutionException.class, () -> elected.get(5, TimeUnit.SECONDS)) + .getCause()); + Assertions.assertSame(expected, Assertions.assertThrows(ExecutionException.class, + () -> shared.toCompletableFuture().get(5, TimeUnit.SECONDS)).getCause()); + Assertions.assertEquals("retried", cache.getOrLoadFuture("key", key -> "retried") + .toCompletableFuture().get()); + Assertions.assertEquals(0, cache.metrics().getInFlightLoadCount()); + } finally { + release.countDown(); + worker.shutdownNow(); + } + } + } + + @Test + void invalidationSeparatesNewLoadWithoutChangingTheOldObservers() throws Exception { + for (String operation : new String[] {"key", "table", "catalog", "close"}) { + ExecutorService worker = Executors.newSingleThreadExecutor(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + try (CatalogMetaCache owner = CatalogMetaCache.unmanaged()) { + MetaCache cache = cache(owner, true); + Future elected = worker.submit(() -> cache.getOrLoadFuture("key", key -> { + entered.countDown(); + await(release); + return "old snapshot"; + })); + await(entered); + CompletionStage observer = cache.getOrLoadFuture("key", key -> { + throw new AssertionError("duplicate old load"); + }); + switch (operation) { + case "key": + cache.invalidateKey("key"); + break; + case "table": + owner.invalidateTable("db", "tbl"); + break; + case "catalog": + owner.invalidateCatalog(); + break; + case "close": + owner.close(); + break; + default: + throw new AssertionError(operation); + } + if (!operation.equals("close")) { + Assertions.assertEquals("new snapshot", cache.getOrLoadFuture("key", key -> "new snapshot") + .toCompletableFuture().get(2, TimeUnit.SECONDS)); + } + release.countDown(); + elected.get(5, TimeUnit.SECONDS); + Assertions.assertEquals("old snapshot", observer.toCompletableFuture().get(5, TimeUnit.SECONDS)); + if (operation.equals("close")) { + Assertions.assertThrows(IllegalStateException.class, + () -> cache.getOrLoadFuture("key", key -> "closed")); + } else { + Assertions.assertEquals("new snapshot", cache.getIfPresent("key"), + "old completion must not replace a newer generation"); + } + } finally { + release.countDown(); + worker.shutdownNow(); + } + } + } + + @Test + void disabledAndNullLoadsRetainNoValue() throws Exception { + for (boolean enabled : new boolean[] {false, true}) { + AtomicInteger loads = new AtomicInteger(); + try (CatalogMetaCache owner = CatalogMetaCache.unmanaged()) { + MetaCache cache = cache(owner, enabled); + for (int i = 0; i < 2; i++) { + Assertions.assertNull(cache.getOrLoadFuture("null", key -> { + loads.incrementAndGet(); + return null; + }).toCompletableFuture().get()); + } + Assertions.assertEquals(2, loads.get()); + Assertions.assertEquals(0, cache.size()); + if (!enabled) { + Assertions.assertEquals("uncached", cache.getOrLoadFuture("key", key -> "uncached") + .toCompletableFuture().get()); + Assertions.assertNull(cache.getIfPresent("key")); + } + } + } + } +} 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 18c4b142adc4ef..83b22cac5d68fb 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 @@ -25,6 +25,7 @@ import org.apache.doris.connector.spi.ConnectorContext; import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.ConnectorStorageContext; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.connector.spi.DorisConnectorException; import org.apache.doris.connector.spi.handle.ConnectorColumnHandle; import org.apache.doris.connector.spi.handle.ConnectorTableHandle; @@ -304,6 +305,7 @@ private static int sessionInt(ConnectorSession session, String name, int default private final class HiveSplitSource implements ConnectorSplitSource { private final HiveTableHandle handle; + private final ConnectorTaskScheduler scheduler; private final FileSystem fs; private final HiveFileFormat tableFormat; private final boolean isLzo; @@ -321,6 +323,7 @@ private final class HiveSplitSource implements ConnectorSplitSource { HiveSplitSource(ConnectorSession session, HiveTableHandle handle, List requiredPartitions) { this.handle = handle; + this.scheduler = session.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT); this.partitionNames = requiredPartitions.isEmpty() ? null : requiredPartitions; this.fs = storage().getFileSystem(session); this.tableFormat = HiveFileFormat.detect(handle.getInputFormat(), handle.getSerializationLib(), @@ -335,7 +338,10 @@ private final class HiveSplitSource implements ConnectorSplitSource { @Override public boolean hasNext() { - while (!closed) { + // Empty partitions can keep hasNext() inside the source for many directory requests. + // Observe query cleanup here, before advancing to another listing, rather than waiting + // for the producer to regain control and close this source after its active read. + while (!closed && !scheduler.isClosed()) { if (ranges.hasNext()) { return true; } 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 e4b45ceaa12401..8950b912c5bebf 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 @@ -26,6 +26,7 @@ import org.apache.doris.connector.hms.HmsTableInfo; import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.ConnectorStatementScope; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.connector.spi.DorisConnectorException; import org.apache.doris.connector.spi.handle.ConnectorColumnHandle; import org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider; @@ -53,7 +54,9 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; import java.util.function.Supplier; /** @@ -184,6 +187,46 @@ public void streamingAndSynchronousSplitsMatchAcrossPartitionsAndFileBoundaries( } } + @Test + public void cancellationBetweenEmptyPartitionsDoesNotListRemainingPartitions() throws Exception { + AtomicBoolean cancelled = new AtomicBoolean(); + AtomicInteger listings = new AtomicInteger(); + ConnectorTaskScheduler scheduler = new ConnectorTaskScheduler() { + @Override + public boolean isClosed() { + return cancelled.get(); + } + + @Override + public boolean tryExecute(Runnable task, Consumer failure, Runnable onSpace) { + throw new UnsupportedOperationException("This source only observes the query lifetime"); + } + + @Override + public void removeSpaceListener(Runnable onSpace) { + } + + @Override + public void execute(Runnable task) { + throw new UnsupportedOperationException("This source only observes the query lifetime"); + } + }; + HiveScanPlanProvider provider = provider(new FakeHmsClient(), (location, fs) -> { + listings.incrementAndGet(); + cancelled.set(true); + return Collections.emptyList(); + }); + HiveTableHandle handle = parquetHandle(PART_KEYS).toBuilder().prunedPartitions(Arrays.asList( + part("year=2024/month=01"), part("year=2024/month=02"))).build(); + FakeSession session = new FakeSession(Collections.singletonMap("num_partitions_in_batch_mode", "1"), + scheduler); + try (ConnectorSplitSource source = provider.streamSplits(session, + ConnectorScanRequest.builder(handle, Collections.emptyList()).build())) { + Assertions.assertFalse(source.hasNext()); + Assertions.assertEquals(1, listings.get(), "cancellation must stop listing the next empty partition"); + } + } + @Test public void closingFirstPartitionDoesNotListRemainingPartitions() throws Exception { CountingLister lister = new CountingLister(); @@ -1263,13 +1306,24 @@ public void close() { /** Minimal {@link ConnectorSession} (no split-size override, empty session properties). */ private static final class FakeSession implements ConnectorSession { private final Map properties; + private final ConnectorTaskScheduler scheduler; FakeSession() { this(Collections.emptyMap()); } FakeSession(Map properties) { + this(properties, ConnectorTaskScheduler.DIRECT); + } + + FakeSession(Map properties, ConnectorTaskScheduler scheduler) { this.properties = properties; + this.scheduler = scheduler; + } + + @Override + public ConnectorTaskScheduler getTaskScheduler(ConnectorTaskScheduler.Stage stage) { + return scheduler; } private final Map statementMemos = new HashMap<>(); diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCacheSizeEstimator.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCacheSizeEstimator.java index bbd05df29822c8..fdffb7dc184fc1 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCacheSizeEstimator.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCacheSizeEstimator.java @@ -246,10 +246,7 @@ private static long estimateWrappedReferenceList(List values) { } static long estimateManifestKey(IcebergManifestEntryKey key) { - long bytes = add(MANIFEST_KEY_SHALLOW_BYTES, JvmSizeUtils.stringSize(key.getManifestPath())); - // The key retains its bound partition type. A non-empty value's files share it, but an empty value (a delete - // manifest whose entries are all DELETED) holds nothing that would account for it. - return add(bytes, estimateIcebergType(key.getPartitionType())); + return add(MANIFEST_KEY_SHALLOW_BYTES, JvmSizeUtils.stringSize(key.getManifestPath())); } static long estimateManifestValue(ManifestCacheValue value) { diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCommentCache.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCommentCache.java index fb6cd7677320a1..35a1e96068543c 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCommentCache.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergCommentCache.java @@ -43,8 +43,8 @@ * disclosure). Vended-credentials uses a single static catalog identity, so every user loads the same table and * sharing a comment is safe. * - *

No credential gate (unlike {@link IcebergTableCache}, like {@link IcebergPartitionCache} / - * {@link IcebergFormatCache}): the cached value is a bare comment {@link String} with no {@code FileIO} / + *

No credential gate (unlike {@link IcebergTableCache}, like {@link IcebergPartitionCache}): the cached value + * is a bare comment {@link String} with no {@code FileIO} / * credential. A comment changes only via external DDL, picked up through the REFRESH invalidate hooks. The * connector builds it with the table entry's {@code meta.cache.iceberg.table.(enable|ttl-second|capacity)}; * {@code enable=false}, {@code ttl-second <= 0} or {@code capacity=0} disables it (read live), so a no-cache diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnector.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnector.java index 2d372f21eadb0d..12c88ced2f834e 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnector.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnector.java @@ -221,9 +221,6 @@ public class IcebergConnector implements Connector { // The value is pure metadata (no FileIO/credential), but under session=user it is an authorization-sensitive // projection (a shared hit would disclose one user's partitions), so it is disabled there (see constructor). private final IcebergPartitionCache partitionCache; // null under session=user - // PERF-03: cross-query inferred-file-format cache (the whole-table planFiles() fallback result, keyed by - // (table, snapshotId)). Same authorization-sensitive treatment as partitionCache: disabled under session=user. - private final IcebergFormatCache formatCache; // null under session=user // PERF-05: cross-query table-comment cache (value = the 'comment' property string). Built ONLY for a REST // vended-credentials catalog that is NOT session=user -- plain catalogs already reuse tableCache (PERF-01) for // the comment path, and session=user must stay live because the loadTable itself carries per-user @@ -236,7 +233,7 @@ public class IcebergConnector implements Connector { // rebuild, not just the remote scan. Two typed fields because the two SPI hooks return structurally different // derived types and neither derives from the other; the shared raw PARTITIONS scan underneath is already // deduped by partitionCache, so two derived caches never double-scan remotely. Authorization-sensitive - // projection like partitionCache/formatCache: disabled (null) under iceberg.rest.session=user so a shared + // projection like partitionCache: disabled (null) under iceberg.rest.session=user so a shared // (no user dimension) hit cannot disclose one user's partition view; kept for every other flavor. private final ConnectorMetadataCache // null under session=user mvccPartitionViewCache; @@ -313,11 +310,6 @@ public IcebergConnector(Map properties, ConnectorContext context ? null : new IcebergPartitionCache( metaCache, cacheSpec("partition")); - // PERF-03: inferred-file-format cache. Same authorization-sensitive treatment as partitionCache (disabled - // under session=user, kept otherwise); readers already tolerate a null cache (resolveFileFormatName). - this.formatCache = isUserSessionEnabled() - ? null - : new IcebergFormatCache(metaCache, tableFactsCacheSpec); // PERF-05: table-comment cache, built ONLY for a REST vended-credentials catalog that is NOT session=user. // Plain catalogs (tableCache on) already serve the comment path from tableCache; session=user is excluded // because a shared comment cache would bypass the per-user loadTable authorization (a metadata disclosure). @@ -811,11 +803,6 @@ IcebergPartitionCache partitionCacheForTest() { return partitionCache; } - /** Test-only: the cross-query inferred-file-format cache (PERF-03), or {@code null} for a session=user catalog. */ - IcebergFormatCache formatCacheForTest() { - return formatCache; - } - /** Test-only: the table-comment cache (PERF-05), or {@code null} unless vended-credentials and non-session. */ IcebergCommentCache commentCacheForTest() { return commentCache; @@ -840,7 +827,7 @@ public ConnectorScanPlanProvider getScanPlanProvider() { // threaded for parity with the legacy single per-catalog IcebergMetadataOps. return new IcebergScanPlanProvider(catalogProps, this::newCatalogBackedOps, context, manifestCache, - tableCache, formatCache, catalogResourceTracker); + tableCache, catalogResourceTracker); } @Override diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergFormatCache.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergFormatCache.java deleted file mode 100644 index 49f3794b346703..00000000000000 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergFormatCache.java +++ /dev/null @@ -1,158 +0,0 @@ -// 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.iceberg; - -import org.apache.doris.connector.cache.CacheSpec; -import org.apache.doris.connector.cache.CatalogMetaCache; -import org.apache.doris.connector.cache.MetaCache; -import org.apache.doris.connector.cache.MetaCacheDefinition; -import org.apache.doris.connector.cache.MetaCacheSizeEstimators; -import org.apache.doris.connector.cache.ScopePath; - -import org.apache.iceberg.catalog.TableIdentifier; - -import java.util.Objects; -import java.util.function.Supplier; - -/** - * Per-catalog cache of an iceberg table's inferred write file-format name (PERF-03), keyed by - * {@code (TableIdentifier, snapshotId)}. Kills the per-query revival of the #64134 heavy fallback: when a table - * carries neither the {@code write-format} nickname nor {@code write.format.default} (migrated tables, and any - * table whose writer never set the property — parquet is an implicit default, not stored), - * {@link IcebergScanPlanProvider#getScanNodeProperties} resolves the scan-level {@code file_format_type} through - * {@link IcebergWriterHelper#getFileFormat}, which falls back to an unfiltered {@code table.newScan().planFiles()} - * — a whole-table manifest scan (remote FileIO) — just to read the FIRST data file's format. That fallback was - * re-run on every query/EXPLAIN with no cross-query reuse; this cache collapses it to one scan per - * {@code (table, snapshot)}. - * - *

Snapshot-keyed, so always correct. A snapshot is immutable, so its first data file's format is a - * pure function of the key; a new commit yields a new snapshot id (a new key -> a live scan). Within the TTL - * the snapshot id itself is held stable by {@link IcebergLatestSnapshotCache}, which is what makes the key stable - * across queries. The key snapshot is the table's {@code currentSnapshot().snapshotId()} — exactly what the - * inference reads (it scans {@code table.newScan()}, never the handle's time-travel pin), matching legacy - * {@code IcebergUtils.getFileFormat} and PERF-02's partition cache verbatim. - * - *

Only the inference fallback is cached. The two cheap property probes ({@code write-format}, - * {@code write.format.default}) stay outside this cache — the caller consults it only when both miss. - * - *

No credential gate (unlike {@link IcebergTableCache}, like {@link IcebergPartitionCache}): the cached - * value is a bare format-name {@link String} with no {@code FileIO} / credential, so it is safe to share across - * users and is built unconditionally (only the table entry's cache keys disable it). - * - *

Backed identically to {@link IcebergPartitionCache}: a contextual, access-TTL {@link MetaCache} with - * manual miss-load, so the inference runs OUTSIDE Caffeine's compute lock and a failed scan's exception - * propagates verbatim and is NOT cached (the next query retries — legacy parity). The connector builds it with - * the table entry's {@code meta.cache.iceberg.table.(enable|ttl-second|capacity)}; {@code enable=false}, - * {@code ttl-second <= 0} or {@code capacity=0} disables it (read live). Lives on the long-lived - * per-catalog {@link IcebergConnector}; a REFRESH CATALOG rebuilds the connector and thus the cache. - */ -final class IcebergFormatCache { - - /** Immutable composite key: a table's inferred format is distinct per pinned snapshot id. */ - static final class Key { - final TableIdentifier id; - final long snapshotId; - - Key(TableIdentifier id, long snapshotId) { - this.id = id; - this.snapshotId = snapshotId; - } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (!(o instanceof Key)) { - return false; - } - Key that = (Key) o; - return snapshotId == that.snapshotId && Objects.equals(id, that.id); - } - - @Override - public int hashCode() { - return Objects.hash(id, snapshotId); - } - } - - private final CatalogMetaCache owner; - private final MetaCache entry; - - IcebergFormatCache(long ttlSeconds, int maxSize) { - // "<= 0 disables" connector TTL contract, folded to CacheSpec's disable sentinel (CacheSpec.ofConnectorTtl). - this(CatalogMetaCache.unmanaged(), CacheSpec.ofConnectorTtl(ttlSeconds, maxSize)); - } - - IcebergFormatCache(CatalogMetaCache owner, CacheSpec spec) { - this.owner = owner; - this.entry = owner.create(MetaCacheDefinition - .builder("iceberg-format", spec, IcebergFormatCache::scope) - .sizeEstimator(MetaCacheSizeEstimators.reflective()) - .build()); - } - - /** Caching is on unless the spec disables it; a disabled cache means "always infer live". */ - boolean isEnabled() { - return entry.isEnabled(); - } - - /** - * Returns the cached format name for {@code key} if present and unexpired, else runs {@code loader} (the live - * whole-table format inference), caches and returns it. Disabled cache -> {@code loader} every call. The - * loader runs OUTSIDE Caffeine's compute lock (single-flight per key) and its exception propagates unwrapped - * and is NOT cached. - */ - String getOrLoad(Key key, Supplier loader) { - return entry.get(key, ignored -> loader.get()); - } - - /** Drops every cached snapshot entry for one table so the next read infers live (REFRESH TABLE). */ - void invalidate(TableIdentifier id) { - owner.invalidateTable(id.namespace().toString(), id.name()); - } - - /** Drops every cached entry for one database (REFRESH DATABASE / DROP DATABASE); db match = namespace equality. */ - void invalidateDb(String dbName) { - owner.invalidateDatabase(dbName); - } - - /** Drops all cached entries (REFRESH CATALOG). */ - void invalidateAll() { - owner.invalidateCatalog(); - } - - /** Test-only: current number of cached entries (accurate map membership, not Caffeine's estimate). */ - int size() { - return Math.toIntExact(entry.size()); - } - - /** Test-only: how many times the live loader (the whole-table format inference) actually ran — the metric gate. */ - long loadCountForTest() { - return entry.loadSuccessCount(); - } - - /** Test-only: the enable / TTL / capacity / weight spec this cache was built with. */ - CacheSpec cacheSpecForTest() { - return entry.cacheSpec(); - } - - private static ScopePath scope(Key key) { - return ScopePath.table(key.id.namespace().toString(), key.id.name()); - } -} diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestCache.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestCache.java index ca54be70761e88..0ca91242ebcb3d 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestCache.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestCache.java @@ -24,26 +24,32 @@ import org.apache.doris.connector.cache.MetaCacheSizeEstimators; import org.apache.doris.connector.cache.ScopePath; -import com.google.common.base.Preconditions; import org.apache.iceberg.DataFile; import org.apache.iceberg.DeleteFile; import org.apache.iceberg.ManifestContent; import org.apache.iceberg.ManifestFile; import org.apache.iceberg.ManifestFiles; import org.apache.iceberg.ManifestReader; +import org.apache.iceberg.PartitionField; import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Snapshot; import org.apache.iceberg.Table; +import org.apache.iceberg.types.Types; import java.io.IOException; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; +import java.util.concurrent.CompletionStage; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.LongAdder; import java.util.function.LongSupplier; import java.util.function.Supplier; @@ -123,9 +129,9 @@ public int hashCode() { /** One in-flight statement's manifest-cache access tally, plus a touch stamp for the leak sweep. */ private static final class ScanStats { - long hits; - long misses; - long failures; + final LongAdder hits = new LongAdder(); + final LongAdder misses = new LongAdder(); + final LongAdder failures = new LongAdder(); volatile long lastTouchNanos; ScanStats(long nowNanos) { @@ -202,33 +208,53 @@ Set getOrLoadEqualityDeleteFieldIds( * loader runs OUTSIDE Caffeine's compute lock (manual miss-load; single-flight per key), so a same-key * miss parses at most once. */ - ManifestCacheValue getManifestCacheValue(ManifestFile manifest, Table table, - Map specsById) { - IcebergManifestEntryKey key = keyOf(manifest, specsById); + ManifestCacheValue getManifestCacheValue(ManifestFile manifest, Table table) { + IcebergManifestEntryKey key = IcebergManifestEntryKey.of(manifest); return entry.get(key, k -> loadManifestCacheValue( - manifest, table, specsById, k.getContent(), entry.isWeightBounded())); + manifest, table, k.getContent(), entry.isWeightBounded())); + } + + /** Runs the elected read here; a shared in-flight read only registers a completion observer. */ + CompletionStage getManifestCacheValueFuture( + ManifestFile manifest, Table table, String queryId) { + IcebergManifestEntryKey key = IcebergManifestEntryKey.of(manifest); + recordAccess(key, queryId); + return entry.getOrLoadFuture(key, k -> loadManifestCacheValue( + manifest, table, k.getContent(), entry.isWeightBounded())); + } + + /** Cached-only probe: avoids futures and executor traffic on the warm planning path. */ + ManifestCacheValue getCachedManifestValue(ManifestFile manifest, String queryId) { + ManifestCacheValue value = entry.getIfPresent(IcebergManifestEntryKey.of(manifest)); + if (value != null && queryId != null && !queryId.isEmpty()) { + touch(queryId).hits.increment(); + } + return value; } /** - * Same as {@link #getManifestCacheValue(ManifestFile, Table, Map)} but tallies this access as a hit or miss under + * Same as {@link #getManifestCacheValue(ManifestFile, Table)} but tallies this access as a hit or miss under * {@code queryId} (a hit == the entry was already cached BEFORE this load, matching the legacy - * {@code getIfPresent(key) != null} probe). Within one statement the manifest-level plan processes manifests - * sequentially, so the per-query counters need no synchronization. A blank queryId is not recorded. + * {@code getIfPresent(key) != null} probe). Concurrent prefetch shares the per-query atomic counters. + * A blank queryId is not recorded. */ - ManifestCacheValue getManifestCacheValue(ManifestFile manifest, Table table, - Map specsById, String queryId) { - IcebergManifestEntryKey key = keyOf(manifest, specsById); - boolean hit = entry.getIfPresent(key) != null; + ManifestCacheValue getManifestCacheValue(ManifestFile manifest, Table table, String queryId) { + IcebergManifestEntryKey key = IcebergManifestEntryKey.of(manifest); + recordAccess(key, queryId); + return entry.get(key, k -> loadManifestCacheValue( + manifest, table, k.getContent(), entry.isWeightBounded())); + } + + private void recordAccess(IcebergManifestEntryKey key, String queryId) { if (queryId != null && !queryId.isEmpty()) { + boolean hit = entry.getIfPresent(key) != null; ScanStats stats = touch(queryId); if (hit) { - stats.hits++; + stats.hits.increment(); } else { - stats.misses++; + stats.misses.increment(); } } - return entry.get(key, k -> loadManifestCacheValue( - manifest, table, specsById, k.getContent(), entry.isWeightBounded())); } /** @@ -237,7 +263,7 @@ ManifestCacheValue getManifestCacheValue(ManifestFile manifest, Table table, */ void recordFailure(String queryId) { if (queryId != null && !queryId.isEmpty()) { - touch(queryId).failures++; + touch(queryId).failures.increment(); } } @@ -252,7 +278,7 @@ long[] takeStats(String queryId) { ScanStats stats = statsByQuery.remove(queryId); return stats == null ? new long[] {0L, 0L, 0L} - : new long[] {stats.hits, stats.misses, stats.failures}; + : new long[] {stats.hits.sum(), stats.misses.sum(), stats.failures.sum()}; } /** Fetch (or lazily create) this query's tally, opportunistically aging out leaked entries first. */ @@ -274,34 +300,23 @@ private void sweepExpiredStats(long nowNanos) { statsByQuery.entrySet().removeIf(e -> nowNanos - e.getValue().lastTouchNanos >= statsTtlNanos); } - /** - * {@code specsById} is the scan's spec map (bound to the scan's schema), the same one its residuals and delete - * index use: the reader converts partition values to these types, so they must match what the scan evaluates. - */ - private static IcebergManifestEntryKey keyOf(ManifestFile manifest, Map specsById) { - PartitionSpec boundSpec = specsById.get(manifest.partitionSpecId()); - Preconditions.checkArgument(boundSpec != null, "No bound partition spec %s for iceberg manifest %s", - manifest.partitionSpecId(), manifest.path()); - return IcebergManifestEntryKey.of(manifest, boundSpec); - } - private static ManifestCacheValue loadManifestCacheValue(ManifestFile manifest, Table table, - Map specsById, ManifestContent content, boolean estimateWeight) { + ManifestContent content, boolean estimateWeight) { try { if (content == ManifestContent.DELETES) { - return ManifestCacheValue.forDeleteFiles(loadDeleteFiles(manifest, table, specsById), estimateWeight); + return ManifestCacheValue.forDeleteFiles(loadDeleteFiles(manifest, table), estimateWeight); } - return ManifestCacheValue.forDataFiles(loadDataFiles(manifest, table, specsById), estimateWeight); + return ManifestCacheValue.forDataFiles(loadDataFiles(manifest, table), estimateWeight); } catch (IOException e) { throw new RuntimeException("Failed to read iceberg manifest " + manifest.path(), e); } } - private static List loadDataFiles(ManifestFile manifest, Table table, - Map specsById) throws IOException { + private static List loadDataFiles(ManifestFile manifest, Table table) throws IOException { List dataFiles = new ArrayList<>(); // .copy() is mandatory: the ManifestReader iterator reuses the same object across iterations. - try (ManifestReader reader = ManifestFiles.read(manifest, table.io(), specsById)) { + try (ManifestReader reader = ManifestFiles.read(manifest, table.io(), + manifestReadSpecs(manifest, table))) { for (DataFile dataFile : reader) { dataFiles.add(dataFile.copy()); } @@ -309,11 +324,46 @@ private static List loadDataFiles(ManifestFile manifest, Table table, return dataFiles; } - private static List loadDeleteFiles(ManifestFile manifest, Table table, - Map specsById) throws IOException { + /** Recover the manifest writer's partition source types, or read its header when that schema expired. */ + private static Map manifestReadSpecs(ManifestFile manifest, Table table) { + PartitionSpec spec = table.specs().get(manifest.partitionSpecId()); + // Fixed-result transforms (bucket/year/month/day/hour) remain readable without their source + // schema. Iceberg resolves their result from UNKNOWN; do not require expired source schemas. + boolean missingSource = spec.rawPartitionType().fields().stream() + .anyMatch(field -> field.type().typeId() == org.apache.iceberg.types.Type.TypeID.UNKNOWN); + if (!missingSource) { + return Collections.singletonMap(spec.specId(), spec); + } + Long writerSnapshotId = manifest.snapshotId(); + Snapshot writerSnapshot = writerSnapshotId == null ? null : table.snapshot(writerSnapshotId); + Schema writerSchema = writerSnapshot == null ? null : table.schemas().get(writerSnapshot.schemaId()); + if (writerSchema == null) { + // The manifest header retains its physical partition type after snapshot/schema expiration. + return null; + } + Map sources = new LinkedHashMap<>(); + Set reservedNames = new HashSet<>(); + spec.fields().forEach(field -> reservedNames.add(field.name())); + for (PartitionField field : spec.fields()) { + Types.NestedField source = writerSchema.findField(field.sourceId()); + if (source == null) { + return null; + } + String name = "source_" + field.sourceId(); + while (!reservedNames.add(name)) { + name += "_"; + } + sources.put(field.sourceId(), Types.NestedField.optional(field.sourceId(), name, source.type())); + } + PartitionSpec readable = spec.toUnbound().bind(new Schema( + new ArrayList<>(sources.values())), true); + return Collections.singletonMap(readable.specId(), readable); + } + + private static List loadDeleteFiles(ManifestFile manifest, Table table) throws IOException { List deleteFiles = new ArrayList<>(); try (ManifestReader reader = - ManifestFiles.readDeleteManifest(manifest, table.io(), specsById)) { + ManifestFiles.readDeleteManifest(manifest, table.io(), manifestReadSpecs(manifest, table))) { for (DeleteFile deleteFile : reader) { deleteFiles.add(deleteFile.copy()); } diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKey.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKey.java index 9e40896d52b473..90552bef6d5d3a 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKey.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKey.java @@ -19,8 +19,6 @@ import org.apache.iceberg.ManifestContent; import org.apache.iceberg.ManifestFile; -import org.apache.iceberg.PartitionSpec; -import org.apache.iceberg.types.Types; import java.util.Objects; @@ -33,25 +31,18 @@ * cached payload, and a table-level invalidation (REFRESH TABLE) intentionally does NOT drop it (legacy * {@code testInvalidateTableKeepsManifestCache} parity). Runtime loader context (the manifest/table instances) * must not be stored here. - * - *

The key also carries the partition type the manifest is decoded with: the reader converts partition values to - * the bound spec's types, and after a type promotion (e.g. INT to BIGINT) a scan pinned to the old schema and a scan - * on the current one bind the same manifest to different types. Each binding needs its own decoded copy. */ public class IcebergManifestEntryKey { private final String manifestPath; private final ManifestContent content; - private final Types.StructType partitionType; - public IcebergManifestEntryKey(String manifestPath, ManifestContent content, Types.StructType partitionType) { + public IcebergManifestEntryKey(String manifestPath, ManifestContent content) { this.manifestPath = Objects.requireNonNull(manifestPath, "manifestPath can not be null"); this.content = Objects.requireNonNull(content, "content can not be null"); - this.partitionType = Objects.requireNonNull(partitionType, "partitionType can not be null"); } - /** The key for {@code manifest} decoded with {@code boundSpec}, the scan's spec for its partition spec id. */ - public static IcebergManifestEntryKey of(ManifestFile manifest, PartitionSpec boundSpec) { - return new IcebergManifestEntryKey(manifest.path(), manifest.content(), boundSpec.partitionType()); + public static IcebergManifestEntryKey of(ManifestFile manifest) { + return new IcebergManifestEntryKey(manifest.path(), manifest.content()); } public String getManifestPath() { @@ -62,10 +53,6 @@ public ManifestContent getContent() { return content; } - public Types.StructType getPartitionType() { - return partitionType; - } - @Override public boolean equals(Object o) { if (this == o) { @@ -76,12 +63,11 @@ public boolean equals(Object o) { } IcebergManifestEntryKey that = (IcebergManifestEntryKey) o; return Objects.equals(manifestPath, that.manifestPath) - && content == that.content - && partitionType.equals(that.partitionType); + && content == that.content; } @Override public int hashCode() { - return Objects.hash(manifestPath, content, partitionType); + return Objects.hash(manifestPath, content); } } diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestPrefetcher.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestPrefetcher.java new file mode 100644 index 00000000000000..4cfa315580ec06 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestPrefetcher.java @@ -0,0 +1,385 @@ +// 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.iceberg; + +import org.apache.doris.connector.spi.ConnectorContext; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; + +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.io.CloseableIterator; + +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; +import java.util.function.Function; +import java.util.function.Supplier; + +/** + * Completion-ordered, single-consumer manifest prefetch. close() may run on the query-cancellation thread. + * Only the immutable manifest-list iterator and cached-only probes run under the short registry lock; + * remote loading, future waits and cancellation all run outside it. At most window results are retained. + */ +final class IcebergManifestPrefetcher implements CloseableIterator { + static final class LoadedManifest { + final ManifestFile manifest; + final ManifestCacheValue files; + + LoadedManifest(ManifestFile manifest, ManifestCacheValue files) { + this.manifest = manifest; + this.files = files; + } + } + + private static final class Pending { + final ManifestFile manifest; + final ManifestCacheValue cached; + RunnableFuture future; + + Pending(ManifestFile manifest, ManifestCacheValue cached, RunnableFuture future) { + this.manifest = manifest; + this.cached = cached; + this.future = future; + } + } + + private static final Pending CLOSED = new Pending(null, null, null); + + private final CloseableIterator manifests; + private final int window; + private final Function cached; + private final Function loader; + private final Function> sharedLoader; + private final ConnectorTaskScheduler scheduler; + private final BooleanSupplier allowPrefetch; + private final Deque pending = new ArrayDeque<>(); + // Entries stay in pending until consumed, so completed results also count against the window. + // Completion callbacks never acquire the registry lock. Reserve one slot for the close signal. + private final BlockingQueue ready; + private volatile boolean closed; + private boolean coldManifestConsumed; + private final ConnectorContext context; + private final IcebergManifestReadContext readContext; + + IcebergManifestPrefetcher(CloseableIterator manifests, int window, + Function cached, + Function loader, ConnectorContext context, Semaphore readPermits) { + this(manifests, window, cached, loader, context, readPermits, () -> true); + } + + IcebergManifestPrefetcher(CloseableIterator manifests, int window, + Function cached, + Function loader, ConnectorContext context, Semaphore readPermits, + BooleanSupplier allowPrefetch) { + this(manifests, window, cached, loader, context, readPermits, allowPrefetch, null); + } + + IcebergManifestPrefetcher(CloseableIterator manifests, int window, + Function cached, + Function loader, ConnectorContext context, Semaphore readPermits, + BooleanSupplier allowPrefetch, ConnectorTaskScheduler scheduler) { + this(manifests, window, cached, loader, context, readPermits, allowPrefetch, scheduler, + manifest -> CompletableFuture.completedFuture(loader.apply(manifest))); + } + + IcebergManifestPrefetcher(CloseableIterator manifests, int window, + Function cached, + Function loader, ConnectorContext context, Semaphore readPermits, + BooleanSupplier allowPrefetch, ConnectorTaskScheduler scheduler, + Function> sharedLoader) { + this(manifests, window, cached, loader, context, readPermits, allowPrefetch, scheduler, + sharedLoader, () -> () -> { }); + } + + IcebergManifestPrefetcher(CloseableIterator manifests, int window, + Function cached, + Function loader, ConnectorContext context, Semaphore readPermits, + BooleanSupplier allowPrefetch, ConnectorTaskScheduler scheduler, + Function> sharedLoader, + Supplier retainLoadOwner) { + this.context = context; + this.readContext = new IcebergManifestReadContext(context, readPermits, + () -> closed || (scheduler != null && scheduler.isClosed()), retainLoadOwner); + this.sharedLoader = sharedLoader; + this.scheduler = scheduler; + this.manifests = manifests; + this.window = window; + this.ready = new ArrayBlockingQueue<>(window + 1); + this.cached = cached; + this.loader = loader; + this.allowPrefetch = allowPrefetch; + } + + @Override + public boolean hasNext() { + synchronized (pending) { + return !closed && (!pending.isEmpty() || manifests.hasNext()); + } + } + + @Override + public LoadedManifest next() { + Pending head = null; + List submissions = null; + try { + synchronized (pending) { + if (closed) { + throw new CancellationException("Iceberg manifest source is closed"); + } + if (pending.isEmpty()) { + if (!manifests.hasNext()) { + throw new NoSuchElementException(); + } + ManifestFile manifest = manifests.next(); + ManifestCacheValue value = cached.apply(manifest); + if (value != null) { + // Warm cache: no futures, submissions, or speculative iterator traversal. + return new LoadedManifest(manifest, value); + } + if (window == 1 || !coldManifestConsumed || !allowPrefetch.getAsBoolean()) { + // A required serial read has no executor handoff. Expand only when useful work + // needs more metadata, without making progress depend on backend requests. + head = new Pending(manifest, null, null); + } else { + submissions = new ArrayList<>(window); + pending.addLast(submit(manifest, null, submissions)); + } + } + if (head == null) { + // If a matching file was found, pause further speculative submissions until the + // backend asks for another batch. Already-started shared reads may still complete. + while (allowPrefetch.getAsBoolean() && pending.size() < window && manifests.hasNext()) { + ManifestFile manifest = manifests.next(); + if (submissions == null) { + submissions = new ArrayList<>(window); + } + pending.addLast(submit(manifest, cached.apply(manifest), submissions)); + } + } + } + // The offline scheduler may execute inline. Even that path must never do IO under + // the pending lock: close() must be able to cancel and wake the consumer at any time. + if (submissions != null) { + submissions.forEach(ScheduledLoad::admit); + } + if (head == null) { + head = awaitReady(); + } + ManifestCacheValue value; + if (head.future == null && head.cached == null) { + value = load(head.manifest); + } else if (head.cached != null) { + value = head.cached; + } else { + // Readiness is published only after the load completes. get() supplies the + // completion barrier and propagates failures, including Errors. + value = head.future.get(); + } + synchronized (pending) { + if (closed) { + throw new CancellationException("Iceberg manifest source is closed"); + } + if (head.future != null || head.cached != null) { + pending.remove(head); + } + coldManifestConsumed = true; + } + return new LoadedManifest(head.manifest, value); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw fail(new RuntimeException("Interrupted while loading Iceberg manifest", e)); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + throw fail(new RuntimeException("Failed to load Iceberg manifest " + head.manifest.path(), cause)); + } catch (RuntimeException e) { + throw fail(e); + } + } + + private Pending awaitReady() throws InterruptedException { + Pending help; + synchronized (pending) { + if (closed) { + throw new CancellationException("Iceberg manifest source is closed"); + } + Pending completed = ready.poll(); + if (completed != null) { + return completed; + } + help = pending.getFirst(); + } + // If every worker is occupied by a parent, run our own queued work to guarantee progress. + // If a worker already owns it, run() is a no-op: wait for ANY completion, not this future. + if (scheduler == null) { + help.future.run(); + } + Pending completed = ready.take(); + if (closed) { + throw new CancellationException("Iceberg manifest source is closed"); + } + return completed; + } + + private Pending submit(ManifestFile manifest, ManifestCacheValue value, List submissions) { + Pending task = new Pending(manifest, value, null); + if (value != null) { + ready.add(task); + return task; + } + if (scheduler != null) { + ScheduledLoad load = new ScheduledLoad(task); + task.future = load; + submissions.add(load); + return task; + } + Callable load = () -> { + try { + return load(manifest); + } finally { + if (!closed) { + // A close racing this publication can leave at most window completions plus + // its wakeup signal. add() must never block a metadata worker on consumption. + ready.add(task); + } + } + }; + task.future = context == null ? new FutureTask<>(load) : context.submitMetadataTask(load); + return task; + } + + /** A rejected admission keeps only this window entry; it never blocks a scheduler worker. */ + private final class ScheduledLoad extends CompletableFuture + implements RunnableFuture { + private final Pending pendingLoad; + private final Runnable onSpace = this::admit; + private final AtomicBoolean started = new AtomicBoolean(); + // A shared load must not retain a cancelled scan until remote IO eventually completes. + private final AtomicReference completionTarget = new AtomicReference<>(this); + + private ScheduledLoad(Pending pendingLoad) { + this.pendingLoad = pendingLoad; + whenComplete((value, failure) -> { + completionTarget.set(null); + scheduler.removeSpaceListener(onSpace); + if (!closed) { + ready.add(pendingLoad); + } + }); + } + + private void admit() { + if (isDone()) { + return; + } + try { + scheduler.tryExecute(this, this::completeExceptionally, onSpace); + } catch (RuntimeException e) { + completeExceptionally(e); + } + // Cancellation can race callback registration in tryExecute. + if (isDone()) { + scheduler.removeSpaceListener(onSpace); + } + } + + @Override + public void run() { + if (isDone() || !started.compareAndSet(false, true)) { + return; + } + try { + CompletionStage shared = executeLoad(pendingLoad.manifest, sharedLoader); + AtomicReference target = completionTarget; + shared.whenComplete((value, failure) -> completeSharedLoad(target, value, failure)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + completeExceptionally(e); + } catch (Throwable failure) { + completeExceptionally(failure); + } + } + } + + private static void completeSharedLoad(AtomicReference target, + ManifestCacheValue value, Throwable failure) { + ScheduledLoad load = target.getAndSet(null); + if (load != null) { + if (failure == null) { + load.complete(value); + } else { + load.completeExceptionally(failure); + } + } + } + + private ManifestCacheValue load(ManifestFile manifest) throws InterruptedException { + return executeLoad(manifest, loader); + } + + private T executeLoad(ManifestFile manifest, Function action) throws InterruptedException { + return readContext.read(manifest, action); + } + + private RuntimeException fail(RuntimeException failure) { + try { + close(); + } catch (IOException closeFailure) { + failure.addSuppressed(closeFailure); + } + return failure; + } + + @Override + public void close() throws IOException { + List unused; + synchronized (pending) { + if (closed) { + return; + } + closed = true; + unused = new ArrayList<>(pending); + pending.clear(); + ready.clear(); + ready.add(CLOSED); + } + for (Pending task : unused) { + if (task.future != null) { + // MetaCache single-flights this immutable load across queries. Interrupting its elected + // loader would also fail unrelated waiters. Cancel queued work, but let an ongoing shared + // read finish; closed is checked before any subsequent manifest read starts. + task.future.cancel(false); + } + } + manifests.close(); + } +} diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestReadContext.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestReadContext.java new file mode 100644 index 00000000000000..f832e17c9a9ef2 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergManifestReadContext.java @@ -0,0 +1,81 @@ +// 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.iceberg; + +import org.apache.doris.connector.spi.ConnectorContext; + +import org.apache.iceberg.ManifestFile; + +import java.util.concurrent.CancellationException; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.function.BooleanSupplier; +import java.util.function.Function; +import java.util.function.Supplier; + +/** Query-wide admission and lifetime for both data and delete manifest reads. */ +final class IcebergManifestReadContext { + private final ConnectorContext context; + private final Semaphore permits; + private final BooleanSupplier closed; + private final Supplier retainOwner; + + IcebergManifestReadContext(ConnectorContext context, Semaphore permits, BooleanSupplier closed, + Supplier retainOwner) { + this.context = context; + this.permits = permits; + this.closed = closed; + this.retainOwner = retainOwner; + } + + T read(ManifestFile manifest, Function action) throws InterruptedException { + return read(manifest.path(), () -> action.apply(manifest)); + } + + T read(String description, Supplier action) throws InterruptedException { + checkOpen(); + // A second scan may be waiting while another scan owns all read permits. Check cancellation + // periodically, without interrupting an elected cache load that other queries still need. + while (!permits.tryAcquire(100, TimeUnit.MILLISECONDS)) { + checkOpen(); + } + ClassLoader previous = Thread.currentThread().getContextClassLoader(); + try { + checkOpen(); + Thread.currentThread().setContextClassLoader(IcebergManifestReadContext.class.getClassLoader()); + // Cancellation of the task Future can precede the end of remote I/O. + try (AutoCloseable owner = retainOwner.get()) { + return context == null ? action.get() + : context.executeAuthenticated(action::get); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException("Failed to load Iceberg manifest " + description, e); + } + } finally { + Thread.currentThread().setContextClassLoader(previous); + permits.release(); + } + } + + private void checkOpen() { + if (closed.getAsBoolean()) { + throw new CancellationException("Iceberg manifest source is closed"); + } + } +} diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergPartitionValueAdapter.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergPartitionValueAdapter.java new file mode 100644 index 00000000000000..7a957622b35763 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergPartitionValueAdapter.java @@ -0,0 +1,312 @@ +// 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.iceberg; + +import org.apache.iceberg.ContentFile; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileContent; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionData; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** Query-local partition adaptation; all other content-file metadata is delegated unchanged. */ +final class IcebergPartitionValueAdapter { + private final int[] sourcePositions; + private final Type[] sourceTypes; + private final Type[] targetTypes; + private final PartitionData template; + + IcebergPartitionValueAdapter(Types.StructType source, Types.StructType target) { + int size = target.fields().size(); + sourcePositions = new int[size]; + sourceTypes = new Type[size]; + targetTypes = new Type[size]; + template = new PartitionData(target); + for (int i = 0; i < size; i++) { + Types.NestedField field = target.fields().get(i); + targetTypes[i] = field.type(); + sourcePositions[i] = -1; + for (int j = 0; j < source.fields().size(); j++) { + if (source.fields().get(j).fieldId() == field.fieldId()) { + sourcePositions[i] = j; + sourceTypes[i] = source.fields().get(j).type(); + break; + } + } + if (sourcePositions[i] < 0) { + throw new IllegalArgumentException("Missing cached partition field " + field.fieldId()); + } + } + } + + DataFile adapt(DataFile file) { + return new DataView(file, adapt((PartitionData) file.partition())); + } + + DeleteFile adapt(DeleteFile file) { + return new DeleteView(file, adapt((PartitionData) file.partition())); + } + + PartitionData adapt(PartitionData source) { + // Copies reuse the Avro schema; do not reconstruct/serialize it for every file. + PartitionData original = source.copy(); + PartitionData result = template.copy(); + for (int i = 0; i < sourcePositions.length; i++) { + if (sourcePositions[i] >= 0) { + result.set(i, convert(original.get(sourcePositions[i], Object.class), sourceTypes[i], targetTypes[i])); + } + } + return result; + } + + private static Object convert(Object value, Type source, Type target) { + if (value == null) { + return null; + } + if (target.typeId() == Type.TypeID.UNKNOWN) { + // A dropped source has no target type. Different cache loaders may have decoded the same + // physical value before or after promotion, so preserve its value in a common representation. + // Iceberg's UNKNOWN comparator uses natural ordering: Integer/Long and Float/Double cannot + // compare across Java classes and would silently fail data/delete partition-key matching. + if (source.typeId() == Type.TypeID.INTEGER) { + return ((Integer) value).longValue(); + } + if (source.typeId() == Type.TypeID.FLOAT) { + return ((Float) value).doubleValue(); + } + // Decimal promotion changes precision only; its Java type and scale already agree. + return value; + } + if (source.equals(target)) { + return value; + } + if (source.typeId() == Type.TypeID.INTEGER && target.typeId() == Type.TypeID.LONG) { + return ((Integer) value).longValue(); + } + if (source.typeId() == Type.TypeID.LONG && target.typeId() == Type.TypeID.INTEGER) { + return Math.toIntExact((Long) value); + } + if (source.typeId() == Type.TypeID.FLOAT && target.typeId() == Type.TypeID.DOUBLE) { + return ((Float) value).doubleValue(); + } + if (source.typeId() == Type.TypeID.DOUBLE && target.typeId() == Type.TypeID.FLOAT) { + double original = (Double) value; + float narrowed = (float) original; + if (!Double.isNaN(original) && Double.compare(original, (double) narrowed) != 0) { + throw new IllegalArgumentException("Partition value cannot be represented as float: " + value); + } + return narrowed; + } + if (source.typeId() == Type.TypeID.DECIMAL && target.typeId() == Type.TypeID.DECIMAL) { + Types.DecimalType decimalType = (Types.DecimalType) target; + BigDecimal decimal = ((BigDecimal) value).setScale(decimalType.scale()); + if (decimal.precision() > decimalType.precision()) { + throw new IllegalArgumentException("Partition value exceeds " + target + ": " + value); + } + return decimal; + } + throw new IllegalArgumentException("Cannot adapt partition value from " + source + " to " + target); + } + + private abstract static class FileView> implements ContentFile { + protected final F file; + private final PartitionData partition; + + FileView(F file, PartitionData partition) { + this.file = file; + this.partition = partition; + } + + protected abstract F wrap(F copy, PartitionData copiedPartition); + + @Override + public PartitionData partition() { + return partition; + } + + @Override + public F copy() { + return wrap(file.copy(), partition.copy()); + } + + @Override + public F copyWithoutStats() { + return wrap(file.copyWithoutStats(), partition.copy()); + } + + @Override + public F copyWithStats(Set requestedColumnIds) { + return wrap(file.copyWithStats(requestedColumnIds), partition.copy()); + } + + @Override + public String manifestLocation() { + return file.manifestLocation(); + } + + @Override + public Long pos() { + return file.pos(); + } + + @Override + public int specId() { + return file.specId(); + } + + @Override + public FileContent content() { + return file.content(); + } + + @Override + public CharSequence path() { + return file.path(); + } + + @Override + public String location() { + return file.location(); + } + + @Override + public FileFormat format() { + return file.format(); + } + + @Override + public long recordCount() { + return file.recordCount(); + } + + @Override + public long fileSizeInBytes() { + return file.fileSizeInBytes(); + } + + @Override + public Map columnSizes() { + return file.columnSizes(); + } + + @Override + public Map valueCounts() { + return file.valueCounts(); + } + + @Override + public Map nullValueCounts() { + return file.nullValueCounts(); + } + + @Override + public Map nanValueCounts() { + return file.nanValueCounts(); + } + + @Override + public Map lowerBounds() { + return file.lowerBounds(); + } + + @Override + public Map upperBounds() { + return file.upperBounds(); + } + + @Override + public ByteBuffer keyMetadata() { + return file.keyMetadata(); + } + + @Override + public List splitOffsets() { + return file.splitOffsets(); + } + + @Override + public List equalityFieldIds() { + return file.equalityFieldIds(); + } + + @Override + public Integer sortOrderId() { + return file.sortOrderId(); + } + + @Override + public Long dataSequenceNumber() { + return file.dataSequenceNumber(); + } + + @Override + public Long fileSequenceNumber() { + return file.fileSequenceNumber(); + } + + @Override + public Long firstRowId() { + return file.firstRowId(); + } + } + + private static final class DataView extends FileView implements DataFile { + DataView(DataFile file, PartitionData partition) { + super(file, partition); + } + + @Override + protected DataFile wrap(DataFile copy, PartitionData copiedPartition) { + return new DataView(copy, copiedPartition); + } + } + + private static final class DeleteView extends FileView implements DeleteFile { + DeleteView(DeleteFile file, PartitionData partition) { + super(file, partition); + } + + @Override + protected DeleteFile wrap(DeleteFile copy, PartitionData copiedPartition) { + return new DeleteView(copy, copiedPartition); + } + + @Override + public String referencedDataFile() { + return file.referencedDataFile(); + } + + @Override + public Long contentOffset() { + return file.contentOffset(); + } + + @Override + public Long contentSizeInBytes() { + return file.contentSizeInBytes(); + } + } +} 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 fe1f1fd1da02a5..bc69b7c8e0514a 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 @@ -22,6 +22,7 @@ import org.apache.doris.connector.spi.ConnectorContext; import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.ConnectorStorageContext; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.connector.spi.DorisConnectorException; import org.apache.doris.connector.spi.handle.ConnectorColumnHandle; import org.apache.doris.connector.spi.handle.ConnectorTableHandle; @@ -117,7 +118,10 @@ import java.util.Optional; import java.util.OptionalLong; import java.util.Set; +import java.util.concurrent.CancellationException; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import java.util.function.Supplier; import java.util.function.UnaryOperator; @@ -244,10 +248,6 @@ public class IcebergScanPlanProvider implements ConnectorScanPlanProvider { // when the connector's credential gate disables the cross-query layer; when null resolveTable still uses the // per-statement scope and falls back to a direct remote load. private final IcebergTableCache tableCache; - // PERF-03: cross-query inferred-file-format cache shared with the connector, owned by the long-lived - // IcebergConnector and injected via getScanPlanProvider. Nullable — null via the offline-test ctors; when null - // getScanNodeProperties resolves file_format_type live (matching pre-PERF-03 behaviour, node-memoized per query). - private final IcebergFormatCache formatCache; private final IcebergCatalogResourceTracker resourceTracker; // FIX-SCAN-METRICS: per-query stash of the iceberg SDK scan diagnostics captured by the attached @@ -302,29 +302,16 @@ public IcebergScanPlanProvider(IcebergCatalogProperties catalogProps, this(catalogProps, catalogOpsResolver, context, manifestCache, tableCache, null); } - /** - * Full ctor used by {@link IcebergConnector#getScanPlanProvider()}, adding the PERF-03 cross-query - * inferred-file-format cache ({@code formatCache}). The 5-arg ctor delegates here with a null format cache - * (offline tests + pre-cache paths resolve {@code file_format_type} live). - */ - public IcebergScanPlanProvider(IcebergCatalogProperties catalogProps, - Function catalogOpsResolver, - ConnectorContext context, IcebergManifestCache manifestCache, IcebergTableCache tableCache, - IcebergFormatCache formatCache) { - this(catalogProps, catalogOpsResolver, context, manifestCache, tableCache, formatCache, null); - } - IcebergScanPlanProvider(IcebergCatalogProperties catalogProps, Function catalogOpsResolver, ConnectorContext context, IcebergManifestCache manifestCache, IcebergTableCache tableCache, - IcebergFormatCache formatCache, IcebergCatalogResourceTracker resourceTracker) { + IcebergCatalogResourceTracker resourceTracker) { this.catalogProps = catalogProps; this.properties = catalogProps.getRaw(); this.catalogOpsResolver = catalogOpsResolver; this.context = context; this.manifestCache = manifestCache; this.tableCache = tableCache; - this.formatCache = formatCache; this.resourceTracker = resourceTracker; } @@ -599,11 +586,12 @@ public ConnectorSplitSource streamSplits(ConnectorSession session, ConnectorScan long sliceSize = fileSplitSize > 0 ? fileSplitSize : sessionLong(session, MAX_FILE_SPLIT_SIZE, DEFAULT_MAX_FILE_SPLIT_SIZE); try { + ManifestReadControl readControl = new ManifestReadControl(false); CloseableIterable tasks = streamingFileScanTasks( - scan, session, table, filter, sliceSize); + scan, session, table, filter, sliceSize, readControl, iceHandle); return new IcebergStreamingSplitSource(tasks, table, scanSchema, formatVersion, partitioned, orderedPartitionKeys, zone, uriNormalizer, sliceSize, - iceHandle.getRewriteFileScope(), iceHandle, requiresRowId(columns)); + iceHandle.getRewriteFileScope(), iceHandle, readControl, requiresRowId(columns)); } catch (RuntimeException e) { throw IcebergExceptionUtils.wrapMetadataReadFailure(iceHandle, e); } @@ -636,12 +624,18 @@ public void close() { * streaming source's {@code hasNext}. Cache disabled -> the SDK path, byte-unchanged. */ private CloseableIterable streamingFileScanTasks(TableScan scan, ConnectorSession session, - Table table, Optional filter, long sliceSize) { + Table table, Optional filter, long sliceSize, + ManifestReadControl readControl, IcebergTableHandle readHandle) { if (isManifestCacheEnabled()) { try { return TableScanUtil.splitFiles( - cacheBackedFileScanTasks(scan, session, table, filter, null), sliceSize); + cacheBackedFileScanTasks(scan, session, table, filter, null, readControl, readHandle), + sliceSize); + } catch (CancellationException e) { + // A stopped query must not restart file planning through the SDK fallback. + throw e; } catch (Exception e) { + requireActiveManifestScheduler(session); LOG.warn("Iceberg streaming plan with manifest cache failed, falling back to SDK scan: {}", e.getMessage(), e); manifestCache.recordFailure(session.getQueryId()); @@ -650,6 +644,21 @@ private CloseableIterable streamingFileScanTasks(TableScan scan, C return TableScanUtil.splitFiles(scan.planFiles(), sliceSize); } + private static final class ManifestReadControl { + private volatile boolean producedSplit; + private volatile boolean moreSplitsRequested; + + ManifestReadControl(boolean moreSplitsRequested) { + this.moreSplitsRequested = moreSplitsRequested; + } + + boolean allowPrefetch() { + // With no output split yet, including rewrite-scope filtering, metadata enumeration + // must make progress in parallel. After producing output, wait for further backend demand. + return !producedSplit || moreSplitsRequested; + } + } + /** * Lazy {@link ConnectorSplitSource} over an iceberg scan's byte-offset-split {@link FileScanTask}s: maps each * task to an {@link IcebergScanRange} on demand (via {@link #buildRangeForTask}) so the engine can pump them @@ -675,7 +684,9 @@ private final class IcebergStreamingSplitSource implements ConnectorSplitSource // streamSplits() BEFORE the source is returned, leaking the planFiles() iterable (the engine pump's // close() never receives it). Lazy open instead routes any failure through hasNext()->the engine's // setException + finally-close, with tasks still closed. null until first use. - private CloseableIterator iterator; + private volatile CloseableIterator iterator; + private final AtomicBoolean closed = new AtomicBoolean(); + private final ManifestReadControl readControl; // Look-ahead buffer so hasNext() can skip data files filtered out by the rewrite scope. private IcebergScanRange buffered; // Per-file invariant cache (PERF-11): per split-source = per scan; the pump is single-threaded, so the @@ -686,8 +697,9 @@ private final class IcebergStreamingSplitSource implements ConnectorSplitSource int formatVersion, boolean partitioned, List orderedPartitionKeys, ZoneId zone, UnaryOperator uriNormalizer, long sliceSize, Set rewriteScope, - IcebergTableHandle handle, boolean requiresRowId) { + IcebergTableHandle handle, ManifestReadControl readControl, boolean requiresRowId) { this.tasks = tasks; + this.readControl = readControl; this.table = table; this.scanSchema = scanSchema; this.formatVersion = formatVersion; @@ -701,24 +713,32 @@ private final class IcebergStreamingSplitSource implements ConnectorSplitSource this.requiresRowId = requiresRowId; } + @Override + public void onMoreSplitsRequested() { + readControl.moreSplitsRequested = true; + } + @Override public boolean hasNext() { try { + if (closed.get()) { + return false; + } if (buffered != null) { return true; } if (iterator == null) { iterator = tasks.iterator(); } - while (iterator.hasNext()) { + while (!closed.get() && iterator.hasNext()) { IcebergScanRange range = buildRangeForTask( iterator.next(), table, scanSchema, formatVersion, partitioned, - orderedPartitionKeys, zone, uriNormalizer, sliceSize, rewriteScope, null, - scratch); + orderedPartitionKeys, zone, uriNormalizer, sliceSize, rewriteScope, null, scratch); if (range != null) { if (requiresRowId) { rejectRowIdRangeWithoutPartitionData(range); } + readControl.producedSplit = true; buffered = range; return true; } @@ -742,6 +762,9 @@ public ConnectorScanRange next() { @Override public void close() throws IOException { + if (!closed.compareAndSet(false, true)) { + return; + } try { if (iterator != null) { iterator.close(); @@ -813,7 +836,7 @@ private List planScanInternal( if (countPushdown && filter.isEmpty()) { Optional> countRanges = planCountPushdown( table, scan, scanSchema, formatVersion, partitioned, orderedPartitionKeys, zone, - uriNormalizer, session, filter); + uriNormalizer, session, filter, iceHandle); if (countRanges.isPresent()) { return countRanges.get(); } @@ -848,7 +871,7 @@ private List planScanInternal( // once per file, not per slice. The streaming source below holds its own. PerFileScratch scratch = new PerFileScratch(); List ranges = new ArrayList<>(); - try (SplitPlan plan = planFileScanTask(scan, session, table, filter)) { + try (SplitPlan plan = planFileScanTask(scan, session, table, filter, iceHandle)) { for (FileScanTask task : plan.tasks) { // Shared per-task mapping (rewrite-scope skip + M-2 weight denominator + v3 stash side-effect), // identical to the streaming path's IcebergStreamingSplitSource so both produce the same ranges. @@ -1368,9 +1391,9 @@ private static Schema pinnedSchema(Table table, IcebergTableHandle handle) { * perform a normal scan. */ private Optional> planCountPushdown(Table table, TableScan scan, - Schema scanSchema, int formatVersion, boolean partitioned, List orderedPartitionKeys, - ZoneId zone, - UnaryOperator uriNormalizer, ConnectorSession session, Optional filter) { + Schema scanSchema, int formatVersion, boolean partitioned, List orderedPartitionKeys, ZoneId zone, + UnaryOperator uriNormalizer, ConnectorSession session, Optional filter, + IcebergTableHandle readHandle) { Snapshot snapshot = scan.snapshot(); if (snapshot == null) { return Optional.of(Collections.emptyList()); @@ -1398,17 +1421,16 @@ private Optional> planCountPushdown(Table table, TableS if (!visibleRows.isPresent()) { return Optional.empty(); } - return planManifestCountRange( - table, scan, scanSchema, visibleRows.getAsLong(), formatVersion, partitioned, - orderedPartitionKeys, zone, uriNormalizer, session, filter, netPositionDeletes); + return planManifestCountRange(table, scan, scanSchema, visibleRows.getAsLong(), formatVersion, + partitioned, orderedPartitionKeys, zone, uriNormalizer, session, filter, + netPositionDeletes, readHandle); } // Older manifest lists may omit aggregate counters. Preserve correctness by falling back to the // bounded per-file enumeration instead of trusting snapshot summary metadata. - return planCountPushdownFromFileTasks( - table, scan, scanSchema, formatVersion, partitioned, orderedPartitionKeys, zone, - uriNormalizer, session, filter, netPositionDeletes, - netPositionDeletes ? positionDeleteRows.getAsLong() : 0); + return planCountPushdownFromFileTasks(table, scan, scanSchema, formatVersion, partitioned, + orderedPartitionKeys, zone, uriNormalizer, session, filter, netPositionDeletes, + netPositionDeletes ? positionDeleteRows.getAsLong() : 0, readHandle); } /** @@ -1443,11 +1465,11 @@ && subtractPositionDeleteRows(manifestCount.getAsLong(), netPositionDeletes ? positionDeleteRows.getAsLong() : 0).isPresent(); } - private Optional> planManifestCountRange( - Table table, TableScan scan, Schema scanSchema, long exactCount, int formatVersion, - boolean partitioned, List orderedPartitionKeys, ZoneId zone, - UnaryOperator uriNormalizer, ConnectorSession session, - Optional filter, boolean netPositionDeletes) { + private Optional> planManifestCountRange(Table table, TableScan scan, + Schema scanSchema, long exactCount, + int formatVersion, boolean partitioned, List orderedPartitionKeys, ZoneId zone, + UnaryOperator uriNormalizer, ConnectorSession session, Optional filter, + boolean netPositionDeletes, IcebergTableHandle readHandle) { if (!isManifestCacheEnabled()) { return buildManifestCountRange( table, scanSchema, scan.planFiles(), exactCount, formatVersion, partitioned, @@ -1456,10 +1478,15 @@ private Optional> planManifestCountRange( String statsQueryId = session != null ? session.getQueryId() : null; try { return buildManifestCountRange( - table, scanSchema, cacheBackedFileScanTasks(scan, session, table, filter, statsQueryId), + table, scanSchema, + cacheBackedFileScanTasks(scan, session, table, filter, statsQueryId, readHandle), exactCount, formatVersion, partitioned, orderedPartitionKeys, zone, uriNormalizer, netPositionDeletes); + } catch (CancellationException e) { + // A stopped query must not restart file planning through the SDK fallback. + throw e; } catch (Exception e) { + requireActiveManifestScheduler(session); LOG.warn("Iceberg count-pushdown representative plan with manifest cache failed, " + "falling back to SDK scan: {}", e.getMessage(), e); manifestCache.recordFailure(statsQueryId); @@ -1496,7 +1523,7 @@ private Optional> planCountPushdownFromFileTasks( Table table, TableScan scan, Schema scanSchema, int formatVersion, boolean partitioned, List orderedPartitionKeys, ZoneId zone, UnaryOperator uriNormalizer, ConnectorSession session, Optional filter, boolean netPositionDeletes, - long positionDeleteRows) { + long positionDeleteRows, IcebergTableHandle readHandle) { if (!isManifestCacheEnabled()) { return accumulateCountPushdownFileTasks( table, scanSchema, scan.planFiles(), formatVersion, partitioned, @@ -1506,10 +1533,14 @@ private Optional> planCountPushdownFromFileTasks( try { return accumulateCountPushdownFileTasks( table, scanSchema, - cacheBackedFileScanTasks(scan, session, table, filter, statsQueryId), formatVersion, + cacheBackedFileScanTasks(scan, session, table, filter, statsQueryId, readHandle), formatVersion, partitioned, orderedPartitionKeys, zone, uriNormalizer, netPositionDeletes, positionDeleteRows); + } catch (CancellationException e) { + // A stopped query must not restart file planning through the SDK fallback. + throw e; } catch (Exception e) { + requireActiveManifestScheduler(session); LOG.warn("Iceberg count-pushdown plan with manifest cache failed, falling back to SDK scan: {}", e.getMessage(), e); manifestCache.recordFailure(statsQueryId); @@ -1770,10 +1801,10 @@ private PerFileScratch computePerFileInvariants(Table table, Schema scanSchema, } partitionValues = ordered; } - // Fail loud on a non-orc/parquet data file, mirroring legacy IcebergScanNode.getFileFormatType() (which - // throws DdlException at plan start). Without this guard the per-range format would silently stay the - // node default FORMAT_JNI and BE would route the file to its system-table JNI reader. (System tables, - // which legitimately use JNI, are the separate P6.5 path, not this normal-read path.) + // Validate the physical format before publishing a native range. An unsupported data file must + // fail here, rather than inherit the scan-level Parquet default and reach the wrong reader. + // Metadata tables legitimately use JNI and take a separate path; normal data ranges always + // carry their own explicit Parquet/ORC format, including mixed-format snapshots. String fileFormat = dataFile.format().name().toLowerCase(Locale.ROOT); if (!"parquet".equals(fileFormat) && !"orc".equals(fileFormat)) { throw new IllegalStateException( @@ -2009,7 +2040,8 @@ public Map getScanNodeProperties( scanSchema = pinnedSchema(table, iceHandle); exactScan = buildScan(table, iceHandle, filter, session); try { - applicableEqualityDeleteFieldIds = cachedApplicableEqualityDeleteFieldIds(table, exactScan); + applicableEqualityDeleteFieldIds = cachedApplicableEqualityDeleteFieldIds( + table, exactScan, session, iceHandle); } catch (RuntimeException e) { throw IcebergExceptionUtils.wrapMetadataReadFailure(iceHandle, e); } @@ -2021,27 +2053,12 @@ public Map getScanNodeProperties( "Current Iceberg scan semantics"); } } - // Scan-level file_format_type is NOT merely a per-range default: BE selects FileScannerV2 vs the V1 - // FileScanner from it (FileScanLocalState::_should_use_file_scanner_v2), and that selector runs before - // any split is fetched -> FORMAT_JNI here pins the ENTIRE scan to V1 whatever IcebergScanRange sets per - // range. V1 is the only tree carrying TableSchemaChangeHelper, so pinning there also re-exposes every - // V1-only reader bug. Emit the table's real data format (legacy IcebergScanNode.getFileFormatType - // parity, same IcebergUtils.getFileFormat resolution, which throws on a non-parquet/orc table); the - // per-file format still travels per range, since one table may mix parquet and orc data files. - // PERF-03: the non-system format resolution falls back to an unfiltered whole-table planFiles() when the - // table sets neither write-format nor write.format.default; memoize that inference per (table, snapshot) - // across queries via formatCache (pure metadata, no credential gate). Null cache (offline) resolves live. - String fileFormatType = "jni"; - if (!systemTable) { - try { - fileFormatType = IcebergWriterHelper.getFileFormat(table, - TableIdentifier.of(iceHandle.getDbName(), iceHandle.getTableName()), formatCache) - .name().toLowerCase(Locale.ROOT); - } catch (RuntimeException e) { - throw IcebergExceptionUtils.wrapMetadataReadFailure(iceHandle, e); - } - } - props.put(ScanNodePropertyKeys.FILE_FORMAT_TYPE, fileFormatType); + // This is only the scan-level native-reader default, not the format of the table's files. + // Every normal Iceberg range carries its own validated Parquet/ORC format, which BE prefers + // over this default (including mixed-format snapshots and row-id fetches). Do not enumerate + // data manifests here just to infer a format: scanning starts in planScan/streamSplits. + // Keep JNI for metadata tables, whose serialized splits use the Java system-table reader. + props.put(ScanNodePropertyKeys.FILE_FORMAT_TYPE, systemTable ? "jni" : "parquet"); // [D-065] System (metadata) tables ($snapshots/$files/...) read via the JNI serialized-split path // (planSystemTableScan): the metadata-table schema travels INSIDE the serialized FileScanTask, so BE // needs neither the base-table path_partition_keys (a metadata table is not base-spec partitioned -> @@ -2270,18 +2287,35 @@ static Set applicableEqualityDeleteFieldIds(Table table, TableScan scan return fieldIds; } - private Set cachedApplicableEqualityDeleteFieldIds(Table table, TableScan scan) { + private Set cachedApplicableEqualityDeleteFieldIds(Table table, TableScan scan, + ConnectorSession session, IcebergTableHandle readHandle) { Snapshot snapshot = scan.snapshot(); if (snapshot == null || "0".equals(snapshot.summary().get(TOTAL_EQUALITY_DELETES))) { return Collections.emptySet(); } + Supplier> loader = () -> { + Semaphore permits = IcebergStatementScope.manifestReadBudget(session, manifestReadConcurrency(session)); + ConnectorTaskScheduler scheduler = manifestReadScheduler(session); + IcebergManifestReadContext reads = new IcebergManifestReadContext(context, permits, + () -> scheduler != null && scheduler.isClosed(), IcebergStatementScope.readTableLeaseFactory( + session, readHandle.getDbName(), readHandle.getTableName())); + try { + return reads.read("equality-delete fields for " + table.name(), + () -> applicableEqualityDeleteFieldIds(table, scan)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + CancellationException cancelled = new CancellationException("Interrupted while reading delete fields"); + cancelled.initCause(e); + throw cancelled; + } + }; if (manifestCache == null) { - return applicableEqualityDeleteFieldIds(table, scan); + return loader.get(); } // Snapshot contents are immutable, so this compact projection can be shared even when the optional // full manifest cache is disabled; a loader failure must still escape and remain retryable. return manifestCache.getOrLoadEqualityDeleteFieldIds( - table.location(), snapshot.snapshotId(), () -> applicableEqualityDeleteFieldIds(table, scan)); + table.location(), snapshot.snapshotId(), loader); } /** @@ -2909,13 +2943,18 @@ private SplitPlan splitFiles(TableScan scan, ConnectorSession session) { * any failure falls back to {@code splitFiles} (legacy parity), so a cache bug can never break a query. */ private SplitPlan planFileScanTask(TableScan scan, ConnectorSession session, Table table, - Optional filter) { + Optional filter, + IcebergTableHandle readHandle) { if (!isManifestCacheEnabled()) { return splitFiles(scan, session); } try { - return planFileScanTaskWithManifestCache(scan, session, table, filter); + return planFileScanTaskWithManifestCache(scan, session, table, filter, readHandle); + } catch (CancellationException e) { + // A stopped query must not restart file planning through the SDK fallback. + throw e; } catch (Exception e) { + requireActiveManifestScheduler(session); LOG.warn("Iceberg plan with manifest cache failed, falling back to SDK scan: {}", e.getMessage(), e); // Mirror the legacy manifestCacheFailures bump so VERBOSE EXPLAIN can report the fallback. manifestCache.recordFailure(session.getQueryId()); @@ -2933,14 +2972,15 @@ private SplitPlan planFileScanTask(TableScan scan, ConnectorSession session, Tab * schema use the table's CURRENT schema (legacy parity). */ private SplitPlan planFileScanTaskWithManifestCache(TableScan scan, - ConnectorSession session, Table table, Optional filter) throws IOException { + ConnectorSession session, Table table, Optional filter, + IcebergTableHandle readHandle) throws IOException { // Null-safe queryId (offline tests pass a null session): a null id selects the no-stats cache overload, // matching the pre-PERF-04 body, which read scan.snapshot() and returned early for an empty table BEFORE // ever calling session.getQueryId(). String statsQueryId = session != null ? session.getQueryId() : null; List tasks = new ArrayList<>(); try (CloseableIterable whole = - cacheBackedFileScanTasks(scan, session, table, filter, statsQueryId)) { + cacheBackedFileScanTasks(scan, session, table, filter, statsQueryId, readHandle)) { for (FileScanTask task : whole) { tasks.add(task); } @@ -2966,16 +3006,23 @@ private SplitPlan planFileScanTaskWithManifestCache(TableScan scan, * {@code planFiles()} (which also reads all delete manifests up front); delete manifests are far fewer than * data manifests. Phase 2 is lazy: the returned iterable's iterator flat-maps the matching data * manifests, yielding one surviving task at a time WITHOUT materializing the list, so a million-file streaming - * scan keeps FE heap bounded (peak = the delete index + one manifest's files + the split queue). + * scan keeps FE heap bounded (peak = the delete index + a bounded manifest window + the split queue). * *

{@code statsQueryId} is nullable: a non-null id tallies cache hits/misses under that query (the - * single-threaded synchronous + COUNT paths); {@code null} selects the no-stats overload for the STREAMING - * path, whose Phase 2 runs on the engine pump thread while Phase 1 ran on the calling thread — tallying the - * per-query {@code ScanStats} counters from two threads would race (the cache documents them as - * single-thread-per-query), and streaming reports no manifest-cache stats today anyway. + * synchronous + COUNT paths); {@code null} preserves the existing no-stats streaming path. + * Cache counters support concurrent prefetch workers. */ private CloseableIterable cacheBackedFileScanTasks(TableScan scan, - ConnectorSession session, Table table, Optional filter, String statsQueryId) { + ConnectorSession session, Table table, Optional filter, String statsQueryId, + IcebergTableHandle readHandle) { + return cacheBackedFileScanTasks(scan, session, table, filter, statsQueryId, + new ManifestReadControl(true), readHandle); + } + + private CloseableIterable cacheBackedFileScanTasks(TableScan scan, + ConnectorSession session, Table table, Optional filter, String statsQueryId, + ManifestReadControl readControl, IcebergTableHandle readHandle) { + ConnectorTaskScheduler scheduler = requireActiveManifestScheduler(session); Snapshot snapshot = scan.snapshot(); if (snapshot == null) { return CloseableIterable.withNoopClose(Collections.emptyList()); @@ -2992,6 +3039,14 @@ private CloseableIterable cacheBackedFileScanTasks(TableScan scan, new InclusiveMetricsEvaluator(scanSchema, filterExpr, caseSensitive); String schemaJson = SchemaParser.toJson(scanSchema); + // Every scan in this statement shares admission, including eager delete-manifest reads. + int concurrency = manifestReadConcurrency(session); + Semaphore readPermits = IcebergStatementScope.manifestReadBudget(session, concurrency); + Supplier retainLoadOwner = IcebergStatementScope.readTableLeaseFactory(session, + readHandle.getDbName(), readHandle.getTableName()); + IcebergManifestReadContext deleteReadContext = new IcebergManifestReadContext(context, readPermits, + () -> scheduler != null && scheduler.isClosed(), retainLoadOwner); + // Phase 1 (eager): partition-prune + cache-load delete manifests into the delete-file index. List deleteFiles = new ArrayList<>(); for (ManifestFile manifest : snapshot.deleteManifests(table.io())) { @@ -3005,7 +3060,19 @@ private CloseableIterable cacheBackedFileScanTasks(TableScan scan, if (!ManifestEvaluator.forPartitionFilter(filterExpr, spec, caseSensitive).eval(manifest)) { continue; } - deleteFiles.addAll(manifestCacheGet(manifest, table, specsById, statsQueryId).getDeleteFiles()); + ManifestCacheValue value = manifestCache.getCachedManifestValue(manifest, statsQueryId); + if (value == null) { + try { + value = deleteReadContext.read(manifest, m -> manifestCacheGet(m, table, statsQueryId)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + CancellationException cancelled = new CancellationException( + "Interrupted while reading Iceberg delete manifest " + manifest.path()); + cancelled.initCause(e); + throw cancelled; + } + } + deleteFiles.addAll(value.forScan(spec).getDeleteFiles()); } DeleteFileIndex deleteIndex = DeleteFileIndex.builderFor(deleteFiles) .schemasById(table.schemas()) @@ -3013,29 +3080,87 @@ private CloseableIterable cacheBackedFileScanTasks(TableScan scan, .caseSensitive(caseSensitive) .build(); - // Phase 2 (lazy): flat-map the matching data manifests, read through the cache on demand. - CloseableIterable dataManifests = + CloseableIterable matching = getMatchingManifest(snapshot.dataManifests(table.io()), specsById, filterExpr); + CloseableIterable dataManifests = CloseableIterable.filter(matching, + manifest -> manifest.content() == ManifestContent.DATA + && specsById.containsKey(manifest.partitionSpecId()) + && residualEvaluators.containsKey(manifest.partitionSpecId())); return new CloseableIterable() { + private ManifestCacheFileScanTaskIterator iterator; + private boolean closed; + @Override - public CloseableIterator iterator() { - return new ManifestCacheFileScanTaskIterator(dataManifests.iterator(), table, specsById, - residualEvaluators, metricsEvaluator, deleteIndex, schemaJson, statsQueryId); + public synchronized CloseableIterator iterator() { + if (closed) { + throw new CancellationException("Iceberg manifest source is closed"); + } + if (iterator != null) { + throw new IllegalStateException("Iceberg manifest source is single-pass"); + } + IcebergManifestPrefetcher prefetcher = new IcebergManifestPrefetcher( + dataManifests.iterator(), concurrency, + manifest -> manifestCache.getCachedManifestValue(manifest, statsQueryId), + manifest -> manifestCacheGet(manifest, table, statsQueryId), context, readPermits, + readControl::allowPrefetch, scheduler, + manifest -> manifestCache.getManifestCacheValueFuture(manifest, table, statsQueryId), + retainLoadOwner); + iterator = new ManifestCacheFileScanTaskIterator(prefetcher, specsById, + residualEvaluators, metricsEvaluator, deleteIndex, schemaJson); + return iterator; } @Override public void close() throws IOException { - dataManifests.close(); + ManifestCacheFileScanTaskIterator active; + synchronized (this) { + if (closed) { + return; + } + closed = true; + active = iterator; + } + try { + if (active != null) { + active.close(); + } + } finally { + dataManifests.close(); + } } }; } + private static int manifestReadConcurrency(ConnectorSession session) { + int concurrency = (int) sessionLong(session, "iceberg_manifest_load_concurrency", 8); + if (concurrency < 1 || concurrency > 64) { + throw new IllegalArgumentException("iceberg_manifest_load_concurrency must be between 1 and 64"); + } + return concurrency; + } + + private ConnectorTaskScheduler manifestReadScheduler(ConnectorSession session) { + // Offline callers have no query lifetime or task groups, including empty-table planning. + if (session == null) { + return null; + } + return context == null ? session.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA) + : context.getMetadataTaskScheduler(session); + } + + private ConnectorTaskScheduler requireActiveManifestScheduler(ConnectorSession session) { + ConnectorTaskScheduler scheduler = manifestReadScheduler(session); + if (scheduler != null && scheduler.isClosed()) { + throw new CancellationException("Iceberg query task group is closed"); + } + return scheduler; + } + /** Dispatch a manifest read to the stats-tallying or no-stats cache overload (see cacheBackedFileScanTasks). */ - private ManifestCacheValue manifestCacheGet(ManifestFile manifest, Table table, - Map specsById, String statsQueryId) { + private ManifestCacheValue manifestCacheGet(ManifestFile manifest, Table table, String statsQueryId) { return statsQueryId != null - ? manifestCache.getManifestCacheValue(manifest, table, specsById, statsQueryId) - : manifestCache.getManifestCacheValue(manifest, table, specsById); + ? manifestCache.getManifestCacheValue(manifest, table, statsQueryId) + : manifestCache.getManifestCacheValue(manifest, table); } /** @@ -3047,32 +3172,28 @@ private ManifestCacheValue manifestCacheGet(ManifestFile manifest, Table table, * table / per manifest spec), byte-identical to the pre-PERF-04 per-file computation. */ private final class ManifestCacheFileScanTaskIterator implements CloseableIterator { - private final CloseableIterator manifestIt; - private final Table table; + private final IcebergManifestPrefetcher manifestIt; + private volatile boolean closed; private final Map specsById; private final Map residualEvaluators; private final InclusiveMetricsEvaluator metricsEvaluator; private final DeleteFileIndex deleteIndex; private final String schemaJson; - private final String statsQueryId; private Iterator dataFileIt; private ResidualEvaluator currentResidual; private String currentSpecJson; private FileScanTask next; - ManifestCacheFileScanTaskIterator(CloseableIterator manifestIt, Table table, + ManifestCacheFileScanTaskIterator(IcebergManifestPrefetcher manifestIt, Map specsById, Map residualEvaluators, - InclusiveMetricsEvaluator metricsEvaluator, DeleteFileIndex deleteIndex, String schemaJson, - String statsQueryId) { + InclusiveMetricsEvaluator metricsEvaluator, DeleteFileIndex deleteIndex, String schemaJson) { this.manifestIt = manifestIt; - this.table = table; this.specsById = specsById; this.residualEvaluators = residualEvaluators; this.metricsEvaluator = metricsEvaluator; this.deleteIndex = deleteIndex; this.schemaJson = schemaJson; - this.statsQueryId = statsQueryId; } @Override @@ -3095,7 +3216,7 @@ public FileScanTask next() { // Fill `next` with the next surviving task, draining the current data manifest's files and advancing to // further data manifests as needed. Per-file logic mirrors the pre-PERF-04 materialized Phase 2 exactly. private void advance() { - while (next == null) { + while (!closed && next == null) { if (dataFileIt != null && dataFileIt.hasNext()) { DataFile dataFile = dataFileIt.next(); if (!metricsEvaluator.eval(dataFile)) { @@ -3111,7 +3232,8 @@ private void advance() { if (!manifestIt.hasNext()) { return; } - ManifestFile manifest = manifestIt.next(); + IcebergManifestPrefetcher.LoadedManifest loaded = manifestIt.next(); + ManifestFile manifest = loaded.manifest; if (manifest.content() != ManifestContent.DATA) { dataFileIt = null; continue; @@ -3124,12 +3246,13 @@ private void advance() { } currentResidual = residual; currentSpecJson = PartitionSpecParser.toJson(spec); - dataFileIt = manifestCacheGet(manifest, table, specsById, statsQueryId).getDataFiles().iterator(); + dataFileIt = loaded.files.forScan(spec).getDataFiles().iterator(); } } @Override public void close() throws IOException { + closed = true; manifestIt.close(); } } diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergStatementScope.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergStatementScope.java index 6bb44ce3612ed2..4ef66c19e11d3d 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergStatementScope.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergStatementScope.java @@ -31,6 +31,7 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; @@ -72,6 +73,8 @@ final class IcebergStatementScope { static final String WRITE_SCHEMA_NAMESPACE = "iceberg.write-schema"; static final String ACTIVE_WRITE_SCHEMA_NAMESPACE = "iceberg.active-write-schema"; + static final String MANIFEST_READ_BUDGET_NAMESPACE = "iceberg.manifest-read-budget"; + private IcebergStatementScope() {} /** A frozen read or its live lease proves existence even if REFRESH retires the catalog cache. */ @@ -109,6 +112,34 @@ static Table sharedBorrowedTable(ConnectorSession session, String dbName, String return borrowed.table; } + /** Shared across all Iceberg scans/catalogs in one execution, isolated across reused query scopes. */ + static Semaphore manifestReadBudget(ConnectorSession session, int concurrency) { + if (session == null) { + return new Semaphore(concurrency); + } + return session.getStatementScope().computeIfAbsent( + MANIFEST_READ_BUDGET_NAMESPACE + ":" + session.getQueryId(), () -> new Semaphore(concurrency)); + } + + /** Capture the existing statement owner; each actual callable retains it until its body exits. */ + static Supplier readTableLeaseFactory(ConnectorSession session, String dbName, String tableName) { + if (session == null || session.getStatementScope() == ConnectorStatementScope.NONE) { + return () -> () -> { }; + } + Object owner = ConnectorStatementScopes.resolveInStatement(session, TABLE_NAMESPACE, dbName, tableName, + () -> { + throw new IllegalStateException("Read table was not resolved: " + dbName + "." + tableName); + }); + if (owner instanceof TrackedTable) { + return () -> ((TrackedTable) owner).retainLease()::close; + } + if (owner instanceof ScopedBorrow) { + return () -> ((ScopedBorrow) owner).lease.retainLease()::close; + } + // A plain table is catalog-owned, with no statement-local FileIO cleanup. + return () -> () -> { }; + } + /** Runs a metadata-only load under a bounded lease when no statement scope exists. */ static T withBorrowedTable(ConnectorSession session, String dbName, String tableName, Supplier loader, Supplier unscopedLoader, diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergTableCache.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergTableCache.java index b727259d491fcd..dbfc851fe73d53 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergTableCache.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergTableCache.java @@ -249,6 +249,14 @@ Table snapshotReadTable() { return owner.snapshotReadTable(); } + TableLease retainLease() { + TableLease retained = owner.tryBorrow(); + if (retained == null) { + throw new IllegalStateException("Iceberg table owner is already closed"); + } + return retained; + } + @Override public void close() { if (closed.compareAndSet(false, true)) { diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriterHelper.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriterHelper.java index 06f5ceb4df0a6a..ecd1f957c0512a 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriterHelper.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriterHelper.java @@ -38,11 +38,9 @@ import org.apache.iceberg.PartitionData; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; -import org.apache.iceberg.Snapshot; import org.apache.iceberg.SortOrder; import org.apache.iceberg.Table; import org.apache.iceberg.TableProperties; -import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.WriteResult; import org.apache.iceberg.types.Conversions; @@ -55,7 +53,6 @@ import org.apache.logging.log4j.Logger; import java.io.IOException; -import java.io.UncheckedIOException; import java.nio.ByteBuffer; import java.time.ZoneId; import java.util.ArrayList; @@ -452,17 +449,6 @@ static FileFormat getFileFormat(Table table) { return toFileFormat(resolveFileFormatName(table, table.properties())); } - /** - * PERF-03 cache-aware overload used by the scan path ({@code IcebergScanPlanProvider.getScanNodeProperties}). - * Identical resolution to {@link #getFileFormat(Table)}, except the whole-table {@code planFiles()} inference - * fallback (fired only when the table sets neither {@code write-format} nor {@code write.format.default}) is - * memoized per {@code (id, currentSnapshotId)} through {@code cache}. The two cheap property probes stay - * uncached. A {@code null} cache/id (offline tests) or an empty table (no current snapshot) resolves live. - */ - static FileFormat getFileFormat(Table table, TableIdentifier id, IcebergFormatCache cache) { - return toFileFormat(resolveFileFormatName(table, table.properties(), id, cache)); - } - static FileFormat toFileFormat(String fileFormatName) { if (fileFormatName.toLowerCase().contains(ORC_NAME)) { return FileFormat.ORC; @@ -483,69 +469,18 @@ private static String resolveFileFormatName(Table table, Map pro return inferFileFormatFromDataFiles(table); } - /** - * PERF-03: like {@link #resolveFileFormatName(Table, Map)} but routes the inference fallback (the heavy - * unfiltered {@code planFiles()}) through the cross-query {@code (id, snapshotId)} cache. Only the inference is - * cached; the property probes are cheap. A failed inference (the loader throws) is NOT cached, so the next - * query retries — for this query it degrades to the parquet default, mirroring the swallow of the live path. - */ - private static String resolveFileFormatName(Table table, Map properties, - TableIdentifier id, IcebergFormatCache cache) { - if (properties.containsKey(WRITE_FORMAT)) { - return properties.get(WRITE_FORMAT); - } - if (properties.containsKey(TableProperties.DEFAULT_FILE_FORMAT)) { - return properties.get(TableProperties.DEFAULT_FILE_FORMAT); - } - Snapshot snapshot = table.currentSnapshot(); - if (cache == null || id == null || snapshot == null) { - return inferFileFormatFromDataFiles(table); - } - try { - return cache.getOrLoad(new IcebergFormatCache.Key(id, snapshot.snapshotId()), - () -> inferFirstDataFileFormat(table)); - } catch (RuntimeException e) { - LOG.warn("Failed to infer file format from data files for table {}, defaulting to {}", - table.name(), PARQUET_NAME, e); - return PARQUET_NAME; - } - } - - /** - * The whole-table format inference (the #64134 heavy op): reads the FIRST data file's format from the current - * snapshot via an unfiltered {@code planFiles()}. The swallowing wrapper used by the write path and the direct - * live path (a failure degrades to parquet). The cache loader instead calls {@link #inferFirstDataFileFormat} - * so a transient failure is not memoized. - */ + /** Write-side inference only; data scans obtain the format directly from each data file. */ private static String inferFileFormatFromDataFiles(Table table) { if (table.currentSnapshot() == null) { return PARQUET_NAME; } - try { - return inferFirstDataFileFormat(table); - } catch (RuntimeException e) { - LOG.warn("Failed to infer file format from data files for table {}, defaulting to {}", - table.name(), PARQUET_NAME, e); - return PARQUET_NAME; - } - } - - /** - * Raw first-data-file format inference that PROPAGATES failures (unchecked) so the cache loader does not - * memoize a transient remote-IO error. The caller guarantees a non-null current snapshot; an empty snapshot - * (no data files) yields the deterministic parquet default (cacheable). The try-with-resources - * {@code CloseableIterable.close()} checked {@link IOException} is wrapped as {@link UncheckedIOException} so - * the method stays unchecked (a {@code Supplier} loader cannot declare checked throws). - */ - private static String inferFirstDataFileFormat(Table table) { try (CloseableIterable files = table.newScan().planFiles()) { Iterator it = files.iterator(); - if (it.hasNext()) { - return it.next().file().format().name().toLowerCase(); - } + return it.hasNext() ? it.next().file().format().name().toLowerCase() : PARQUET_NAME; + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to infer file format from data files for table {}, defaulting to {}", + table.name(), PARQUET_NAME, e); return PARQUET_NAME; - } catch (IOException e) { - throw new UncheckedIOException(e); } } diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/ManifestCacheValue.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/ManifestCacheValue.java index 2c2b2d60925216..d6ab52da77b84e 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/ManifestCacheValue.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/ManifestCacheValue.java @@ -20,8 +20,12 @@ import org.apache.doris.connector.cache.MetaCacheSizeEstimate; import org.apache.doris.connector.cache.MetaCacheSizeEstimator; +import com.google.common.collect.Lists; import org.apache.iceberg.DataFile; import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.PartitionData; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.types.Types; import java.util.Collections; import java.util.List; @@ -30,7 +34,9 @@ * Cached manifest payload containing the parsed files of one iceberg manifest. * *

Ported verbatim from the legacy fe-core {@code org.apache.doris.datasource.iceberg.cache.ManifestCacheValue} - * (references only iceberg-SDK types, so no fe-core dependency). A DATA manifest yields data files; a DELETES + * (references only iceberg-SDK types, so no fe-core dependency). Parsed partition types reflect the schema used + * by the first loader; scans must use {@link #forScan} before interpreting partition values. + * A DATA manifest yields data files; a DELETES * manifest yields delete files (T08, mirrors {@link IcebergManifestCache}). */ public class ManifestCacheValue { @@ -64,6 +70,22 @@ static ManifestCacheValue forDeleteFiles(List deleteFiles, boolean e return new ManifestCacheValue(Collections.emptyList(), deleteFiles, estimateWeight); } + /** Adapt lazily at the cache/scan boundary, before predicates or delete matching inspect partitions. */ + ManifestCacheValue forScan(PartitionSpec spec) { + if (dataFiles.isEmpty() && deleteFiles.isEmpty()) { + return this; + } + PartitionData partition = (PartitionData) (dataFiles.isEmpty() + ? deleteFiles.get(0).partition() : dataFiles.get(0).partition()); + Types.StructType target = spec.rawPartitionType(); + if (partition.getPartitionType().equals(target)) { + return this; + } + IcebergPartitionValueAdapter adapter = new IcebergPartitionValueAdapter(partition.getPartitionType(), target); + return new ManifestCacheValue(Lists.transform(dataFiles, adapter::adapt), + Lists.transform(deleteFiles, adapter::adapt), false); + } + public List getDataFiles() { return dataFiles; } diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorCacheTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorCacheTest.java index be4736618e3d1b..778e6e1e71bdb4 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorCacheTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorCacheTest.java @@ -198,7 +198,6 @@ public void invalidateHooksAreNoThrowForSessionUserWithNulledCaches() { IcebergConnector connector = new IcebergConnector(session, new RecordingConnectorContext()); Assertions.assertNull(connector.latestSnapshotCacheForTest()); Assertions.assertNull(connector.partitionCacheForTest()); - Assertions.assertNull(connector.formatCacheForTest()); Assertions.assertDoesNotThrow(() -> connector.invalidateTable("db1", "t1")); Assertions.assertDoesNotThrow(() -> connector.invalidateDb("db1")); Assertions.assertDoesNotThrow(connector::invalidateAll); @@ -214,7 +213,7 @@ public void refreshCatalogInvalidateAllDropsManifestCache() { Table table = tableWithOneManifest(); ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); IcebergManifestCache manifestCache = connector.manifestCacheForTest(); - manifestCache.getManifestCacheValue(manifest, table, table.specs()); + manifestCache.getManifestCacheValue(manifest, table); Assertions.assertEquals(1, manifestCache.size(), "the manifest is cached after a load"); // REFRESH TABLE must NOT drop the manifest cache (path-keyed immutable content; legacy parity). @@ -255,7 +254,7 @@ public void weightLimitsReachEveryConnectorOwnedCache() throws Exception { long oneMb = 1024L * 1024L; Map catalogLimited = maxWeightsOf(props("meta.cache.max-weight", "1MB")); for (String cache : new String[] {"iceberg-table", "iceberg-latest-snapshot", "iceberg-partition", - "iceberg-format", "iceberg-manifest", "iceberg-equality-delete-field-ids", + "iceberg-manifest", "iceberg-equality-delete-field-ids", "iceberg.mvcc-partition-view", "iceberg.list-partitions-view"}) { Assertions.assertEquals(Long.valueOf(oneMb), catalogLimited.get(cache), "a catalog limit must bound " + cache); @@ -351,8 +350,8 @@ public void independentPartitionTtlKeepsItsNoExpirationSemantics() throws Except } @Test - public void disabledTableEntryReadsSnapshotFormatAndCommentLive() throws Exception { - // table.enable=false and capacity=0 turn the table entry off, like ttl-second<=0. The latest-snapshot, format + public void disabledTableEntryReadsSnapshotAndCommentLive() throws Exception { + // table.enable=false and capacity=0 turn the table entry off, like ttl-second<=0. The latest-snapshot // and comment caches take the table entry's keys, so they read live too, and the latest schema bypasses the // name-keyed schema cache. MUTATION: building them from ttl-second alone -> the first pin stays cached -> red. for (String[] off : new String[][] { @@ -368,7 +367,6 @@ public void disabledTableEntryReadsSnapshotFormatAndCommentLive() throws Excepti Assertions.assertEquals(2L, snapshots.getOrLoad(id, loader).snapshotId, off[0]); Assertions.assertEquals(0, snapshots.size()); Assertions.assertFalse(connector.tableCacheForTest().isEnabled(), off[0]); - Assertions.assertFalse(connector.formatCacheForTest().isEnabled(), off[0]); Assertions.assertEquals(OptionalLong.of(0L), connector.schemaCacheTtlSecondOverride(), off[0]); } Map vended = restProps(true, false); @@ -381,7 +379,7 @@ public void disabledTableEntryReadsSnapshotFormatAndCommentLive() throws Excepti } @Test - public void tableEntryKeysExceptMaxWeightSizeSnapshotFormatAndCommentCaches() throws Exception { + public void tableEntryKeysExceptMaxWeightSizeSnapshotAndCommentCaches() throws Exception { // The table entry's ttl-second and capacity also size these caches. Its max-weight bounds the table cache // only; the catalog and FE-global limits still bound them. MUTATION: the fixed default capacity -> red; // handing them the whole table spec, max-weight included -> red. @@ -391,7 +389,6 @@ public void tableEntryKeysExceptMaxWeightSizeSnapshotFormatAndCommentCaches() th properties.put("meta.cache.iceberg.table.max-weight", "1MB"); try (IcebergConnector connector = new IcebergConnector(properties, new RecordingConnectorContext())) { List.of(connector.latestSnapshotCacheForTest().cacheSpecForTest(), - connector.formatCacheForTest().cacheSpecForTest(), connector.commentCacheForTest().cacheSpecForTest()).forEach(spec -> { Assertions.assertTrue(spec.isCacheEnabled()); Assertions.assertEquals(600L, spec.getTtlSecond()); @@ -578,60 +575,6 @@ public void refreshHooksInvalidatePartitionCache() { Assertions.assertEquals(0, cache.size(), "REFRESH CATALOG drops everything"); } - private static IcebergFormatCache.Key fmtKey(String db, String tbl, long snapshotId) { - return new IcebergFormatCache.Key(TableIdentifier.of(db, tbl), snapshotId); - } - - @Test - public void formatCacheBuiltUnlessSessionUser() { - // The inferred-format cache stores a pure metadata format-name string (no FileIO/credential), so like the - // partition cache it stays built for a REST vended-credentials catalog. But under iceberg.rest.session=user - // it is an AUTHORIZATION-sensitive projection, so it is disabled (null) there (same treatment as the - // partition cache). MUTATION: dropping the session=user gate -> non-null for session -> red; gating on the - // vended flag -> null for vended -> red. - Assertions.assertNotNull( - new IcebergConnector(Collections.emptyMap(), new RecordingConnectorContext()).formatCacheForTest(), - "a plain catalog builds the format cache"); - Map vended = new HashMap<>(); - vended.put(IcebergCatalogProperties.ICEBERG_CATALOG_TYPE, IcebergCatalogProperties.TYPE_REST); - vended.put("iceberg.rest.vended-credentials-enabled", "true"); - Assertions.assertNotNull( - new IcebergConnector(vended, new RecordingConnectorContext()).formatCacheForTest(), - "a vended-credentials catalog still builds the format cache (a format name carries no credentials)"); - Map session = new HashMap<>(); - session.put(IcebergCatalogProperties.ICEBERG_CATALOG_TYPE, IcebergCatalogProperties.TYPE_REST); - session.put("iceberg.rest.session", "user"); - Assertions.assertNull( - new IcebergConnector(session, new RecordingConnectorContext()).formatCacheForTest(), - "a session=user catalog must NOT build the format cache (per-user authz must not be bypassed)"); - } - - @Test - public void refreshHooksInvalidateFormatCache() { - // The REFRESH hooks must clear the inferred-format cache too (else a rewrite that changed the write format - // would stay invisible beyond the pin): REFRESH TABLE drops that table's snapshot entries, REFRESH DATABASE - // that db's, REFRESH CATALOG everything. MUTATION: an invalidate* hook not touching formatCache -> a stale - // entry survives -> a size assert below red. - IcebergConnector connector = - new IcebergConnector(Collections.emptyMap(), new RecordingConnectorContext()); - IcebergFormatCache cache = connector.formatCacheForTest(); - Assertions.assertNotNull(cache); - cache.getOrLoad(fmtKey("db1", "t1", 1L), () -> "parquet"); - cache.getOrLoad(fmtKey("db1", "t1", 2L), () -> "orc"); - cache.getOrLoad(fmtKey("db1", "t2", 1L), () -> "parquet"); - cache.getOrLoad(fmtKey("db2", "t1", 1L), () -> "parquet"); - Assertions.assertEquals(4, cache.size()); - - connector.invalidateTable("db1", "t1"); - Assertions.assertEquals(2, cache.size(), "REFRESH TABLE drops both snapshot entries of db1.t1"); - - connector.invalidateDb("db1"); - Assertions.assertEquals(1, cache.size(), "REFRESH DATABASE drops db1's remaining table"); - - connector.invalidateAll(); - Assertions.assertEquals(0, cache.size(), "REFRESH CATALOG drops everything"); - } - private static Map restProps(boolean vended, boolean sessionUser) { Map m = new HashMap<>(); m.put(IcebergCatalogProperties.ICEBERG_CATALOG_TYPE, IcebergCatalogProperties.TYPE_REST); diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergFormatCacheTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergFormatCacheTest.java deleted file mode 100644 index d4eb3e719a24c6..00000000000000 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergFormatCacheTest.java +++ /dev/null @@ -1,181 +0,0 @@ -// 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.iceberg; - -import org.apache.iceberg.catalog.TableIdentifier; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; - -import java.util.concurrent.atomic.AtomicInteger; - -/** - * Unit tests for {@link IcebergFormatCache} (PERF-03). Mirrors {@link IcebergPartitionCacheTest} but keys by - * {@code (TableIdentifier, snapshotId)} and stores the inferred format-name String. Covers within-TTL stability, - * the {@code ttl <= 0} disable, invalidation, and the not-cached-on-failure guarantee that makes a transient - * remote-IO failure retry on the next query (legacy parity). - */ -public class IcebergFormatCacheTest { - - private static IcebergFormatCache.Key key(String db, String tbl, long snapshotId) { - return new IcebergFormatCache.Key(TableIdentifier.of(db, tbl), snapshotId); - } - - @Test - public void cachesWithinTtlAndServesTheSameValue() { - AtomicInteger loads = new AtomicInteger(); - IcebergFormatCache c = new IcebergFormatCache(100, 1000); - - String first = c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - return "orc"; - }); - // Second read of the SAME (table, snapshot) within TTL returns the cached format, NOT a fresh whole-table - // planFiles() inference -> this is what collapses the per-query #64134 fallback. MUTATION: inferring live - // every call -> returns "parquet" / loads==2 -> red. - String second = c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - return "parquet"; - }); - Assertions.assertEquals("orc", first); - Assertions.assertEquals("orc", second, "within TTL the cached format must be served"); - Assertions.assertEquals(1, loads.get(), "the live inference must run exactly once within TTL"); - Assertions.assertEquals(1, c.loadCountForTest(), "the metric-gate load count must be 1"); - Assertions.assertTrue(c.isEnabled()); - } - - @Test - public void differentSnapshotIdIsADifferentKey() { - AtomicInteger loads = new AtomicInteger(); - IcebergFormatCache c = new IcebergFormatCache(100, 1000); - c.getOrLoad(key("db", "t", 1L), () -> { - loads.incrementAndGet(); - return "parquet"; - }); - // A new commit yields a new snapshot id -> a distinct key -> a live inference (freshness across snapshots, - // e.g. a rewrite that changed the write format). MUTATION: keying by table only -> loads==1 -> red. - String s2 = c.getOrLoad(key("db", "t", 2L), () -> { - loads.incrementAndGet(); - return "orc"; - }); - Assertions.assertEquals("orc", s2); - Assertions.assertEquals(2, loads.get()); - Assertions.assertEquals(2, c.size()); - } - - @Test - public void ttlZeroDisablesCachingAlwaysLive() { - AtomicInteger loads = new AtomicInteger(); - IcebergFormatCache c = new IcebergFormatCache(0, 1000); - c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - return "orc"; - }); - String second = c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - return "parquet"; - }); - // ttl-second=0 (the no-cache catalog) infers live every time. MUTATION: caching despite ttl<=0 -> - // second=="orc" / loads==1 -> red. - Assertions.assertEquals("parquet", second, "ttl-second=0 must always infer live"); - Assertions.assertEquals(2, loads.get()); - Assertions.assertFalse(c.isEnabled()); - Assertions.assertEquals(0, c.size(), "ttl-second=0 must not store anything"); - } - - @Test - public void negativeTtlDisablesCachingAlwaysLive() { - AtomicInteger loads = new AtomicInteger(); - IcebergFormatCache c = new IcebergFormatCache(-1, 1000); - c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - return "orc"; - }); - String second = c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - return "parquet"; - }); - Assertions.assertEquals("parquet", second, "ttl-second=-1 must always infer live"); - Assertions.assertEquals(2, loads.get()); - Assertions.assertFalse(c.isEnabled()); - } - - @Test - public void invalidateDropsAllSnapshotsOfOneTable() { - AtomicInteger loads = new AtomicInteger(); - IcebergFormatCache c = new IcebergFormatCache(100, 1000); - c.getOrLoad(key("db", "t", 1L), () -> "parquet"); - c.getOrLoad(key("db", "t", 2L), () -> "orc"); - c.getOrLoad(key("db", "other", 1L), () -> "parquet"); - Assertions.assertEquals(3, c.size()); - - // REFRESH TABLE db.t must drop BOTH snapshot entries of db.t and leave db.other intact. MUTATION: - // invalidating a single (table, snapshot) key -> the other snapshot survives -> size 2 here -> red. - c.invalidate(TableIdentifier.of("db", "t")); - Assertions.assertEquals(1, c.size(), "only db.other's entry must survive"); - String reload = c.getOrLoad(key("db", "t", 1L), () -> { - loads.incrementAndGet(); - return "orc"; - }); - Assertions.assertEquals("orc", reload); - Assertions.assertEquals(1, loads.get()); - } - - @Test - public void invalidateDbClearsOnlyThatDbsTables() { - IcebergFormatCache c = new IcebergFormatCache(100, 1000); - c.getOrLoad(key("db1", "t1", 1L), () -> "parquet"); - c.getOrLoad(key("db1", "t2", 1L), () -> "orc"); - c.getOrLoad(key("db2", "t1", 1L), () -> "parquet"); - Assertions.assertEquals(3, c.size()); - c.invalidateDb("db1"); - Assertions.assertEquals(1, c.size(), "only db2's entry must survive REFRESH DATABASE db1"); - } - - @Test - public void invalidateAllClearsEverything() { - IcebergFormatCache c = new IcebergFormatCache(100, 1000); - c.getOrLoad(key("db", "t1", 1L), () -> "parquet"); - c.getOrLoad(key("db", "t2", 1L), () -> "orc"); - Assertions.assertEquals(2, c.size()); - c.invalidateAll(); - Assertions.assertEquals(0, c.size()); - } - - @Test - public void loaderFailureIsNotCachedSoNextQueryRetries() { - // A transient remote-IO failure during inference (planFiles/close) must NOT be memoized, or the parquet - // fallback would stick for the whole TTL. The MetaCache manual-miss-load path re-throws the loader's - // RuntimeException verbatim and does not store it. MUTATION: caching on failure -> the second call would - // not re-run the loader / size==1 -> red. - IcebergFormatCache c = new IcebergFormatCache(100, 1000); - AtomicInteger loads = new AtomicInteger(); - Assertions.assertThrows(RuntimeException.class, () -> c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - throw new RuntimeException("transient manifest read failure"); - })); - Assertions.assertEquals(0, c.size(), "a failed inference must not be cached"); - // The next query at the same key re-runs the loader (retry), and a success is then cached. - String ok = c.getOrLoad(key("db", "t", 5L), () -> { - loads.incrementAndGet(); - return "orc"; - }); - Assertions.assertEquals("orc", ok); - Assertions.assertEquals(2, loads.get(), "the loader must re-run after a failure (no sticky failure)"); - Assertions.assertEquals(1, c.size()); - } -} diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestCacheTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestCacheTest.java index d4e8edc6d425f7..91288ac1e6af58 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestCacheTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestCacheTest.java @@ -18,29 +18,54 @@ package org.apache.doris.connector.iceberg; import org.apache.doris.connector.cache.CatalogMetaCache; -import org.apache.doris.connector.cache.JvmSizeUtils; -import org.apache.doris.connector.cache.MetaCacheSizeEstimate; import org.apache.iceberg.DataFile; import org.apache.iceberg.DataFiles; -import org.apache.iceberg.ManifestContent; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.DeleteFileIndex; +import org.apache.iceberg.FileMetadata; import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.ManifestFiles; +import org.apache.iceberg.ManifestReader; import org.apache.iceberg.Metrics; +import org.apache.iceberg.PartitionData; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.Table; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.expressions.ResidualEvaluator; import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.io.SeekableInputStream; +import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import java.io.IOException; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.math.BigDecimal; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicInteger; /** @@ -49,6 +74,590 @@ */ public class IcebergManifestCacheTest { + @Test + public void stablePartitionTypesOpenOneStreamAndWarmCacheOpensNone() throws Exception { + assertManifestStreamCount(Types.IntegerType.get(), 7, 1); + assertManifestStreamCount(Types.StringType.get(), "north", 1); + assertManifestStreamCount(Types.DateType.get(), 2, 1); + assertManifestStreamCount(Types.TimestampType.withoutZone(), 123456L, 1); + assertManifestStreamCount(Types.FloatType.get(), 1.25F, 1); + assertManifestStreamCount(Types.BooleanType.get(), true, 1); + assertManifestStreamCount(Types.BinaryType.get(), ByteBuffer.wrap(new byte[] {1, 2}), 1); + } + + @Test + public void promotablePartitionTypesAlsoOpenOneStream() throws Exception { + assertManifestStreamCount(Types.LongType.get(), 7L, 1); + assertManifestStreamCount(Types.DoubleType.get(), 1.25D, 1); + assertManifestStreamCount(Types.DecimalType.of(9, 2), new BigDecimal("1.25"), 1); + } + + @Test + public void unpartitionedManifestOpensOneStream() throws Exception { + Table table = tableWithTwoDataFiles(); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + CountingFileIO io = new CountingFileIO(table.io()); + try (ManifestReader reader = ManifestFiles.read(manifest, io, null)) { + Assertions.assertEquals(2, countFiles(reader)); + } + Assertions.assertEquals(2, io.opens.get(), "the SDK header and entry readers open separate streams"); + Table wrapped = tableWithIo(table, io); + IcebergManifestCache cache = new IcebergManifestCache(); + ManifestCacheValue value = cache.getManifestCacheValue(manifest, wrapped); + Assertions.assertEquals(2, value.getDataFiles().size()); + Assertions.assertEquals(3, io.opens.get(), "the cache loader must add only one stream"); + Assertions.assertSame(value, cache.getManifestCacheValue(manifest, wrapped)); + Assertions.assertEquals(3, io.opens.get()); + } + + @Test + public void bucketPartitionRemainsSingleReadAfterSourceTypePromotion() throws Exception { + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).bucket("id", 16).build(); + Table table = tableWithPartitionValue(schema, spec, 7); + long snapshotId = table.currentSnapshot().snapshotId(); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + IcebergManifestCache cache = new IcebergManifestCache(); + ManifestCacheValue before = cache.getManifestCacheValue(manifest, table); + table.updateSchema().updateColumn("id", Types.LongType.get()).commit(); + Assertions.assertEquals(snapshotId, table.currentSnapshot().snapshotId()); + Assertions.assertSame(before, cache.getManifestCacheValue(manifest, table)); + CountingFileIO io = new CountingFileIO(table.io()); + DataFile cold = new IcebergManifestCache().getManifestCacheValue(manifest, tableWithIo(table, io)) + .getDataFiles().get(0); + Assertions.assertEquals(1, io.opens.get()); + Assertions.assertEquals(before.getDataFiles().get(0).partition().get(0, Integer.class), + cold.partition().get(0, Integer.class)); + } + + @Test + public void partitionSpecEvolutionUsesEachManifestsSpecId() throws Exception { + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.StringType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("id").build(); + Table table = tableWithPartitionValue(schema, spec, "north"); + ManifestFile oldManifest = table.currentSnapshot().dataManifests(table.io()).get(0); + table.updateSpec().renameField("id", "region").commit(); + Assertions.assertNotEquals(oldManifest.partitionSpecId(), table.spec().specId()); + CountingFileIO io = new CountingFileIO(table.io()); + DataFile file = new IcebergManifestCache().getManifestCacheValue(oldManifest, tableWithIo(table, io)) + .getDataFiles().get(0); + Assertions.assertEquals(1, io.opens.get()); + Assertions.assertEquals(oldManifest.partitionSpecId(), file.specId()); + Assertions.assertEquals("north", file.partition().get(0, String.class)); + } + + @Test + public void droppingPartitionSourceRetainsValuesWithOneRead() throws Exception { + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.StringType.get()), + Types.NestedField.optional(2, "payload", Types.StringType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("id").build(); + Table table = tableWithPartitionValue(schema, spec, "north"); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + table.updateSpec().removeField("id").commit(); + table.updateSchema().deleteColumn("id").commit(); + CountingFileIO io = new CountingFileIO(table.io()); + DataFile file = new IcebergManifestCache().getManifestCacheValue(manifest, tableWithIo(table, io)) + .getDataFiles().get(0); + Assertions.assertEquals(1, io.opens.get(), "recover the dropped source type from retained table schemas"); + Assertions.assertEquals("north", file.partition().get(0, String.class)); + try (ManifestReader nativeReader = ManifestFiles.read(manifest, table.io(), table.specs())) { + DataFile nativeFile = nativeReader.iterator().next(); + ManifestCacheValue current = ManifestCacheValue.forDataFiles(Collections.singletonList(file)) + .forScan(table.specs().get(manifest.partitionSpecId())); + Assertions.assertEquals(nativeFile.partition().get(0, Object.class), + current.getDataFiles().get(0).partition().get(0, Object.class)); + Assertions.assertEquals("north", file.partition().get(0, String.class), "cached value remains intact"); + } + } + + @ParameterizedTest + @CsvSource({"true", "false"}) + public void droppedPromotedSourceUsesWriterSchemaOrManifestHeader(boolean writerSnapshotRetained) + throws Exception { + Schema initial = new Schema(Types.NestedField.optional(1, "p", Types.IntegerType.get())); + PartitionSpec spec = PartitionSpec.builderFor(initial).identity("p").build(); + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("schema-id-order", Collections.emptyMap()); + catalog.createNamespace(Namespace.of("db1")); + Table table = catalog.createTable(TableIdentifier.of("db1", "t1"), initial, spec); + table.updateSchema().updateColumn("p", Types.LongType.get()).commit(); + PartitionData partition = new PartitionData(table.spec().rawPartitionType()); + partition.set(0, 7L); + table.newFastAppend().appendFile(DataFiles.builder(table.spec()).withPath("/data/long.parquet") + .withFileSizeInBytes(100).withRecordCount(1).withPartition(partition).build()).commit(); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + Assertions.assertNotNull(table.snapshot(manifest.snapshotId())); + table.updateSpec().removeField("p").commit(); + table.updateSchema().deleteColumn("p").commit(); + + // Schema IDs are identifiers, not a chronology: an older INT schema can have a larger ID. + Map schemas = new LinkedHashMap<>(table.schemas()); + schemas.put(100, new Schema(100, Collections.singletonList( + Types.NestedField.optional(1, "p", Types.IntegerType.get())))); + CountingFileIO io = new CountingFileIO(table.io()); + Table source = (Table) Proxy.newProxyInstance(Table.class.getClassLoader(), new Class[] {Table.class}, + (proxy, method, args) -> { + if (method.getName().equals("schemas")) { + return schemas; + } + if (!writerSnapshotRetained && method.getName().equals("snapshot") + && manifest.snapshotId().equals(args[0])) { + return null; + } + if (method.getName().equals("io")) { + return io; + } + try { + return method.invoke(table, args); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + }); + ManifestCacheValue cached = new IcebergManifestCache().getManifestCacheValue(manifest, source); + DataFile projected = cached.forScan(table.specs().get(manifest.partitionSpecId())).getDataFiles().get(0); + Assertions.assertEquals(7L, projected.partition().get(0, Long.class)); + Assertions.assertEquals(writerSnapshotRetained ? 1 : 2, io.opens.get()); + } + + @Test + public void droppedDeletePartitionSourceMatchesNativeProjectionWithoutChangingCache() throws Exception { + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.StringType.get()), + Types.NestedField.optional(2, "payload", Types.StringType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("id").build(); + Table table = tableWithPartitionValue(schema, spec, "north"); + PartitionData partition = new PartitionData(spec.partitionType()); + partition.set(0, "north"); + DeleteFile deleted = FileMetadata.deleteFileBuilder(spec).ofEqualityDeletes(2) + .withPath("/data/delete.parquet").withFileSizeInBytes(100).withRecordCount(1) + .withPartition(partition).build(); + table.newRowDelta().addDeletes(deleted).commit(); + ManifestFile manifest = table.currentSnapshot().deleteManifests(table.io()).get(0); + table.updateSpec().removeField("id").commit(); + table.updateSchema().deleteColumn("id").commit(); + ManifestCacheValue cached = new IcebergManifestCache().getManifestCacheValue(manifest, table); + try (ManifestReader reader = ManifestFiles.readDeleteManifest(manifest, table.io(), table.specs())) { + Object nativeValue = reader.iterator().next().partition().get(0, Object.class); + ManifestCacheValue projected = cached.forScan(table.specs().get(manifest.partitionSpecId())); + Assertions.assertEquals("north", nativeValue); + Assertions.assertEquals(nativeValue, projected.getDeleteFiles().get(0).partition().get(0, Object.class)); + Assertions.assertEquals("north", cached.getDeleteFiles().get(0).partition().get(0, String.class)); + Assertions.assertSame(cached, cached.forScan(spec), "historical schema keeps the cached physical values"); + } + } + + @Test + public void retainedBucketManifestSurvivesSourceSchemaExpiration() throws Exception { + Schema schema = new Schema(Types.NestedField.optional(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "payload", Types.StringType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).bucket("id", 16).build(); + Table table = tableWithPartitionValue(schema, spec, 7); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + table.updateSpec().removeField("id_bucket").commit(); + table.updateSchema().deleteColumn("id").commit(); + table.newFastAppend().appendFile(DataFiles.builder(table.spec()).withPath("/data/new.parquet") + .withFileSizeInBytes(100).withRecordCount(1).build()).commit(); + table.expireSnapshots().expireOlderThan(System.currentTimeMillis() + 1000) + .retainLast(1).cleanExpiredMetadata(true).commit(); + Assertions.assertTrue(table.schemas().values().stream().noneMatch(value -> value.findField(1) != null)); + CountingFileIO io = new CountingFileIO(table.io()); + DataFile file = new IcebergManifestCache().getManifestCacheValue(manifest, tableWithIo(table, io)) + .getDataFiles().get(0); + Assertions.assertEquals(7, file.partition().get(0, Integer.class)); + Assertions.assertEquals(1, io.opens.get()); + } + + @ParameterizedTest + @CsvSource({"false, false", "true, false", "false, true", "true, true"}) + void retainedVariableTypeManifestSurvivesSourceSchemaExpiration(boolean truncate, boolean deletes) + throws Exception { + Schema schema = new Schema(Types.NestedField.optional(1, "id", Types.StringType.get()), + Types.NestedField.optional(2, "payload", Types.StringType.get())); + PartitionSpec spec = truncate ? PartitionSpec.builderFor(schema).truncate("id", 5).build() + : PartitionSpec.builderFor(schema).identity("id").build(); + Table table = tableWithPartitionValue(schema, spec, "north"); + ManifestFile manifest; + if (deletes) { + PartitionData partition = new PartitionData(spec.partitionType()); + partition.set(0, "north"); + table.newRowDelta().addDeletes(FileMetadata.deleteFileBuilder(spec).ofPositionDeletes() + .withPath("/data/delete.parquet").withPartition(partition) + .withFileSizeInBytes(100).withRecordCount(1).build()).commit(); + manifest = table.currentSnapshot().deleteManifests(table.io()).get(0); + } else { + manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + } + table.updateSpec().removeField(spec.fields().get(0).name()).commit(); + table.updateSchema().deleteColumn("id").commit(); + table.newFastAppend().appendFile(DataFiles.builder(table.spec()).withPath("/data/new.parquet") + .withFileSizeInBytes(100).withRecordCount(1).build()).commit(); + table.expireSnapshots().expireOlderThan(System.currentTimeMillis() + 1000) + .retainLast(1).cleanExpiredMetadata(true).commit(); + Assertions.assertTrue(table.schemas().values().stream().noneMatch(value -> value.findField(1) != null)); + CountingFileIO io = new CountingFileIO(table.io()); + IcebergManifestCache cache = new IcebergManifestCache(); + Table wrapped = tableWithIo(table, io); + ManifestCacheValue value = cache.getManifestCacheValue(manifest, wrapped); + Object physical = deletes ? value.getDeleteFiles().get(0).partition().get(0, Object.class) + : value.getDataFiles().get(0).partition().get(0, Object.class); + Assertions.assertEquals("north", physical.toString()); + Assertions.assertEquals(2, io.opens.get(), "only expired source schemas require the SDK header recovery"); + Assertions.assertSame(value, cache.getManifestCacheValue(manifest, wrapped)); + Assertions.assertEquals(2, io.opens.get(), "warm reads do not reopen the header or entries"); + ManifestCacheValue projected = value.forScan(table.specs().get(spec.specId())); + Object preserved = deletes ? projected.getDeleteFiles().get(0).partition().get(0, Object.class) + : projected.getDataFiles().get(0).partition().get(0, Object.class); + Assertions.assertEquals("north", preserved.toString()); + if (deletes) { + ManifestFile dataManifest = table.currentSnapshot().dataManifests(table.io()).stream() + .filter(file -> file.partitionSpecId() == spec.specId()).findFirst().orElseThrow(); + DataFile data = cache.getManifestCacheValue(dataManifest, wrapped) + .forScan(table.specs().get(spec.specId())).getDataFiles().get(0); + DeleteFileIndex index = DeleteFileIndex.builderFor(projected.getDeleteFiles()) + .schemasById(table.schemas()).specsById(table.specs()).build(); + Assertions.assertEquals(1, index.forDataFile(data.dataSequenceNumber(), data).length, + "physical partition recovery must preserve position-delete matching"); + } + } + + @Test + public void recoveredSourceNamesCannotCollideWithPartitionNames() throws Exception { + Schema schema = new Schema(Types.NestedField.optional(1, "p", Types.StringType.get()), + Types.NestedField.optional(2, "q", Types.IntegerType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("p").bucket("q", 16, "source_1").build(); + Table table = tableWithPartitionValue(schema, spec, "north"); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + table.updateSpec().removeField("p").commit(); + table.updateSchema().deleteColumn("p").commit(); + CountingFileIO io = new CountingFileIO(table.io()); + DataFile file = new IcebergManifestCache().getManifestCacheValue(manifest, tableWithIo(table, io)) + .getDataFiles().get(0); + Assertions.assertEquals("north", file.partition().get(0, String.class)); + Assertions.assertEquals(1, io.opens.get()); + } + + @Test + public void floatingPointAndDecimalPromotionsAdaptBothCacheOrders() throws Exception { + assertPromotionAdaptsPartitionType(Types.FloatType.get(), Types.DoubleType.get(), 1.25F); + assertPromotionAdaptsPartitionType(Types.DecimalType.of(9, 2), Types.DecimalType.of(18, 2), + new BigDecimal("1.25")); + } + + private static void assertPromotionAdaptsPartitionType(Type.PrimitiveType from, Type.PrimitiveType to, + Object partitionValue) { + Schema schema = new Schema(Types.NestedField.required(1, "id", from)); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("id").build(); + Table table = tableWithPartitionValue(schema, spec, partitionValue); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + IcebergManifestCache cache = new IcebergManifestCache(); + ManifestCacheValue warm = cache.getManifestCacheValue(manifest, table); + table.updateSchema().updateColumn("id", to).commit(); + Assertions.assertSame(warm, cache.getManifestCacheValue(manifest, table)); + CountingFileIO io = new CountingFileIO(table.io()); + DataFile cold = new IcebergManifestCache().getManifestCacheValue(manifest, tableWithIo(table, io)) + .getDataFiles().get(0); + Assertions.assertEquals(1, io.opens.get()); + for (ManifestCacheValue payload : Arrays.asList(warm, + ManifestCacheValue.forDataFiles(Collections.singletonList(cold)))) { + for (PartitionSpec target : Arrays.asList(spec, table.spec())) { + DataFile adapted = payload.forScan(target).getDataFiles().get(0); + Assertions.assertEquals(target.rawPartitionType(), + ((PartitionData) adapted.partition()).getPartitionType()); + Assertions.assertEquals(Expressions.alwaysTrue(), + ResidualEvaluator.of(target, + Expressions.equal("id", partitionValue), true) + .residualFor(adapted.partition())); + } + } + Assertions.assertEquals(partitionValue, warm.getDataFiles().get(0).partition().get(0, Object.class)); + } + + private static void assertManifestStreamCount(Type type, Object partitionValue, int expectedOpens) + throws IOException { + Schema schema = new Schema(Types.NestedField.required(1, "id", type)); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("id").build(); + Table table = tableWithPartitionValue(schema, spec, partitionValue); + CountingFileIO io = new CountingFileIO(table.io()); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + IcebergManifestCache cache = new IcebergManifestCache(); + Table wrapped = tableWithIo(table, io); + ManifestCacheValue value = cache.getManifestCacheValue(manifest, wrapped); + Assertions.assertEquals(partitionValue, value.getDataFiles().get(0).partition().get(0, Object.class)); + try (ManifestReader reader = ManifestFiles.read(manifest, table.io(), null)) { + Assertions.assertEquals(((PartitionData) reader.iterator().next().partition()).getPartitionType(), + ((PartitionData) value.getDataFiles().get(0).partition()).getPartitionType()); + } + Assertions.assertEquals(expectedOpens, io.opens.get(), type.toString()); + Assertions.assertSame(value, cache.getManifestCacheValue(manifest, wrapped)); + Assertions.assertEquals(expectedOpens, io.opens.get(), "a parsed-manifest cache hit must not open a stream"); + } + + private static Table tableWithPartitionValue(Schema schema, PartitionSpec spec, Object partitionValue) { + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("test", Collections.emptyMap()); + catalog.createNamespace(Namespace.of("db1")); + Table table = catalog.createTable(TableIdentifier.of("db1", "physical"), schema, spec); + PartitionData partition = new PartitionData(spec.partitionType()); + partition.set(0, partitionValue); + table.newFastAppend().appendFile(DataFiles.builder(spec).withPath("/data/physical.parquet") + .withFileSizeInBytes(100).withRecordCount(1).withPartition(partition).build()).commit(); + return table; + } + + private static int countFiles(ManifestReader reader) { + int count = 0; + Iterator files = reader.iterator(); + while (files.hasNext()) { + files.next(); + count++; + } + return count; + } + + private static Table tableWithIo(Table table, FileIO io) { + return (Table) Proxy.newProxyInstance(Table.class.getClassLoader(), new Class[] {Table.class}, + (proxy, method, args) -> { + if (method.getName().equals("io")) { + return io; + } + try { + return method.invoke(table, args); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + }); + } + + private static final class CountingFileIO implements FileIO { + private final FileIO delegate; + private final AtomicInteger opens = new AtomicInteger(); + + private CountingFileIO(FileIO delegate) { + this.delegate = delegate; + } + + @Override + public Map properties() { + Map properties = new HashMap<>(delegate.properties()); + properties.put("io.manifest.cache-enabled", "false"); + return properties; + } + + @Override + public InputFile newInputFile(String path) { + InputFile input = delegate.newInputFile(path); + return new InputFile() { + @Override + public long getLength() { + return input.getLength(); + } + + @Override + public SeekableInputStream newStream() { + opens.incrementAndGet(); + return input.newStream(); + } + + @Override + public String location() { + return input.location(); + } + + @Override + public boolean exists() { + return input.exists(); + } + }; + } + + @Override + public OutputFile newOutputFile(String path) { + return delegate.newOutputFile(path); + } + + @Override + public void deleteFile(String path) { + delegate.deleteFile(path); + } + } + + @Test + public void schemaPromotionAdaptsCachedPartitionsWithoutMutation() throws Exception { + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("test", Collections.emptyMap()); + catalog.createNamespace(Namespace.of("db1")); + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("id").build(); + Table table = catalog.createTable(TableIdentifier.of("db1", "promoted"), schema, spec); + table.newFastAppend().appendFile(DataFiles.builder(spec).withPath("/data/old.parquet") + .withFileSizeInBytes(100).withRecordCount(1).withPartitionPath("id=7").build()).commit(); + long snapshotId = table.currentSnapshot().snapshotId(); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + IcebergManifestCache cache = new IcebergManifestCache(); + ManifestCacheValue before = cache.getManifestCacheValue(manifest, table); + + table.updateSchema().updateColumn("id", Types.LongType.get()).commit(); + Assertions.assertEquals(snapshotId, table.currentSnapshot().snapshotId(), + "a metadata-only schema change keeps the snapshot and manifest identities"); + Assertions.assertSame(before, cache.getManifestCacheValue(manifest, table)); + DataFile warm = before.getDataFiles().get(0); + DataFile cold = new IcebergManifestCache().getManifestCacheValue(manifest, table).getDataFiles().get(0); + Assertions.assertEquals(7, warm.partition().get(0, Integer.class)); + Assertions.assertEquals(7L, cold.partition().get(0, Long.class)); + ExecutorService pool = Executors.newFixedThreadPool(4); + try { + List> reads = new ArrayList<>(); + for (ManifestCacheValue payload : Arrays.asList(before, + ManifestCacheValue.forDataFiles(Collections.singletonList(cold)))) { + for (PartitionSpec target : Arrays.asList(spec, table.spec())) { + reads.add(() -> { + for (int i = 0; i < 100; i++) { + DataFile adapted = payload.forScan(target).getDataFiles().get(0); + Assertions.assertEquals(Expressions.alwaysTrue(), + ResidualEvaluator.of(target, + Expressions.equal("id", 7), true) + .residualFor(adapted.partition())); + Assertions.assertEquals(warm.dataSequenceNumber(), adapted.dataSequenceNumber()); + Assertions.assertEquals(warm.fileSequenceNumber(), adapted.fileSequenceNumber()); + Assertions.assertEquals(target.rawPartitionType(), + ((PartitionData) adapted.copy().partition()).getPartitionType()); + } + return null; + }); + } + } + for (Future future : pool.invokeAll(reads)) { + future.get(); + } + } finally { + pool.shutdownNow(); + } + Assertions.assertEquals(7, warm.partition().get(0, Integer.class)); + Assertions.assertEquals(7L, cold.partition().get(0, Long.class)); + } + + @Test + public void promotedPartitionsStillMatchEqualityDeletes() { + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); + PartitionSpec oldSpec = PartitionSpec.builderFor(schema).identity("id").build(); + Table table = tableWithPartitionValue(schema, oldSpec, 7); + ManifestFile dataManifest = table.currentSnapshot().dataManifests(table.io()).get(0); + DeleteFile delete = FileMetadata.deleteFileBuilder(oldSpec) + .ofEqualityDeletes(1).withPath("/delete.parquet").withRecordCount(1).withFileSizeInBytes(100) + .withPartitionPath("id=7").build(); + table.newRowDelta().addDeletes(delete).commit(); + ManifestFile deleteManifest = table.currentSnapshot().deleteManifests(table.io()).get(0); + IcebergManifestCache cache = new IcebergManifestCache(); + ManifestCacheValue data = cache.getManifestCacheValue(dataManifest, table); + ManifestCacheValue deletes = cache.getManifestCacheValue(deleteManifest, table); + table.updateSchema().updateColumn("id", Types.LongType.get()).commit(); + for (PartitionSpec target : Arrays.asList(oldSpec, table.spec())) { + DeleteFileIndex index = DeleteFileIndex + .builderFor(deletes.forScan(target).getDeleteFiles()) + .schemasById(table.schemas()).specsById(Collections.singletonMap(target.specId(), target)).build(); + DataFile adapted = data.forScan(target).getDataFiles().get(0); + DeleteFile[] matches = index.forDataFile(adapted.dataSequenceNumber(), adapted); + Assertions.assertEquals(1, matches.length); + Assertions.assertEquals(delete.location(), matches[0].location()); + Assertions.assertEquals(deletes.getDeleteFiles().get(0).dataSequenceNumber(), + matches[0].dataSequenceNumber()); + } + Assertions.assertEquals(7, deletes.getDeleteFiles().get(0).partition().get(0, Integer.class)); + } + + @ParameterizedTest + @CsvSource({"false, false", "true, false", "false, true", "true, true"}) + void promotedPartitionsMatchDeletesAfterSourceCleanup(boolean warmDeletes, boolean truncate) throws Exception { + Schema schema = new Schema(Types.NestedField.optional(1, "p", Types.IntegerType.get()), + Types.NestedField.required(2, "id", Types.IntegerType.get())); + PartitionSpec original = truncate ? PartitionSpec.builderFor(schema).truncate("p", 4).build() + : PartitionSpec.builderFor(schema).identity("p").build(); + int value = truncate ? 4 : 7; + Table table = tableWithPartitionValue(schema, original, value); + ManifestFile dataManifest = table.currentSnapshot().dataManifests(table.io()).get(0); + PartitionData partition = new PartitionData(original.rawPartitionType()); + partition.set(0, value); + DeleteFile delete = FileMetadata.deleteFileBuilder(original).ofEqualityDeletes(2) + .withPath("/data/delete.parquet").withFileSizeInBytes(100).withRecordCount(1) + .withPartition(partition).build(); + table.newRowDelta().addDeletes(delete).commit(); + ManifestFile deleteManifest = table.currentSnapshot().deleteManifests(table.io()).get(0); + table.updateSchema().updateColumn("p", Types.LongType.get()).commit(); + IcebergManifestCache cache = new IcebergManifestCache(); + ManifestCacheValue warm = cache.getManifestCacheValue(warmDeletes ? deleteManifest : dataManifest, table); + table.updateSpec().removeField(original.fields().get(0).name()).commit(); + table.updateSchema().deleteColumn("p").commit(); + table.newFastAppend().appendFile(DataFiles.builder(table.spec()).withPath("/data/new.parquet") + .withFileSizeInBytes(100).withRecordCount(1).build()).commit(); + table.expireSnapshots().expireOlderThan(System.currentTimeMillis() + 1000) + .retainLast(1).cleanExpiredMetadata(true).commit(); + Assertions.assertTrue(table.schemas().values().stream().noneMatch(retained -> retained.findField(1) != null)); + ManifestCacheValue data = cache.getManifestCacheValue(dataManifest, table); + ManifestCacheValue deletes = cache.getManifestCacheValue(deleteManifest, table); + Assertions.assertSame(warm, warmDeletes ? deletes : data); + Assertions.assertEquals((long) value, (warmDeletes ? deletes.getDeleteFiles().get(0) + : data.getDataFiles().get(0)).partition().get(0, Object.class)); + Assertions.assertEquals(value, (warmDeletes ? data.getDataFiles().get(0) + : deletes.getDeleteFiles().get(0)).partition().get(0, Object.class)); + PartitionSpec current = table.specs().get(original.specId()); + ExecutorService pool = Executors.newFixedThreadPool(2); + try { + List> reads = new ArrayList<>(); + for (PartitionSpec target : Arrays.asList(original, current)) { + reads.add(() -> { + for (int i = 0; i < 20; i++) { + DataFile adapted = data.forScan(target).getDataFiles().get(0); + DeleteFileIndex index = DeleteFileIndex.builderFor(deletes.forScan(target).getDeleteFiles()) + .schemasById(table.schemas()) + .specsById(Collections.singletonMap(target.specId(), target)).build(); + Assertions.assertEquals(1, index.forDataFile(adapted.dataSequenceNumber(), adapted).length, + "mixed warm/cold numeric values must match after source-schema cleanup"); + PartitionData otherPartition = new PartitionData(original.rawPartitionType()); + otherPartition.set(0, value + 4); + DataFile other = DataFiles.builder(original).withPath("/data/other.parquet") + .withPartition(otherPartition).withRecordCount(1).withFileSizeInBytes(100).build(); + DataFile projected = ManifestCacheValue.forDataFiles(Collections.singletonList(other)) + .forScan(target).getDataFiles().get(0); + Assertions.assertEquals(0, index.forDataFile(adapted.dataSequenceNumber(), projected).length, + "normalization must not attach deletes from another partition"); + } + return null; + }); + } + for (Future result : pool.invokeAll(reads)) { + result.get(); + } + } finally { + pool.shutdownNow(); + } + Assertions.assertEquals((long) value, (warmDeletes ? deletes.getDeleteFiles().get(0) + : data.getDataFiles().get(0)).partition().get(0, Object.class), "shared cache must remain unchanged"); + } + + @Test + public void concurrentPrefetchCountsEveryCacheAccess() throws Exception { + Table table = tableWithTwoDataFiles(); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + IcebergManifestCache cache = new IcebergManifestCache(); + cache.getManifestCacheValue(manifest, table); + ExecutorService pool = Executors.newFixedThreadPool(8); + try { + List> jobs = new ArrayList<>(); + for (int i = 0; i < 8; i++) { + jobs.add(pool.submit(() -> { + for (int j = 0; j < 100; j++) { + Assertions.assertNotNull(cache.getCachedManifestValue(manifest, "parallel")); + } + })); + } + for (Future job : jobs) { + job.get(5, java.util.concurrent.TimeUnit.SECONDS); + } + Assertions.assertArrayEquals(new long[] {800, 0, 0}, cache.takeStats("parallel")); + } finally { + pool.shutdownNow(); + } + } + @Test public void manifestEstimationCountsBackingArraysAndRejectsHiddenHeapBuffers() { byte[] backing = new byte[1024 * 1024]; @@ -96,7 +705,7 @@ public void loadsDataFilesAndCachesByManifestPath() { ManifestFile manifest = manifests.get(0); IcebergManifestCache cache = new IcebergManifestCache(); - ManifestCacheValue first = cache.getManifestCacheValue(manifest, table, table.specs()); + ManifestCacheValue first = cache.getManifestCacheValue(manifest, table); // WHY: a DATA manifest yields the two appended data files. MUTATION: returning the delete-file branch // (empty) -> 0 -> red. Assertions.assertEquals(2, first.getDataFiles().size()); @@ -105,7 +714,7 @@ public void loadsDataFilesAndCachesByManifestPath() { // A second read of the SAME manifest path returns the SAME cached value (no re-parse). MUTATION: not // caching (re-loading) -> a different instance -> red. - ManifestCacheValue second = cache.getManifestCacheValue(manifest, table, table.specs()); + ManifestCacheValue second = cache.getManifestCacheValue(manifest, table); Assertions.assertSame(first, second, "the manifest payload must be served from the cache on a repeat"); Assertions.assertEquals(1, cache.size()); } @@ -118,36 +727,14 @@ public void weightBoundedManifestIsEstimatedAndCached() { IcebergManifestCache cache = new IcebergManifestCache(owner, Collections.singletonMap("meta.cache.iceberg.manifest.max-weight", "1MB")); - ManifestCacheValue first = cache.getManifestCacheValue(manifest, table, table.specs()); - ManifestCacheValue second = cache.getManifestCacheValue(manifest, table, table.specs()); + ManifestCacheValue first = cache.getManifestCacheValue(manifest, table); + ManifestCacheValue second = cache.getManifestCacheValue(manifest, table); Assertions.assertSame(first, second); Assertions.assertEquals(2, first.getDataFiles().size()); } } - @Test - public void manifestKeyWeightCountsTheBoundPartitionType() { - // WHY: the key retains the bound partition type, and an empty value (a delete manifest whose entries are all - // DELETED) holds no file that references it, so only the key estimate can account for it. MUTATION: - // estimateManifestKey ignoring the partition type -> both entries weigh the same -> red. - ManifestCacheValue empty = ManifestCacheValue.forDeleteFiles(Collections.emptyList(), true); - MetaCacheSizeEstimate narrow = IcebergCacheSizeEstimator.estimateManifestEntry(deleteManifestKey(1), empty); - MetaCacheSizeEstimate wide = IcebergCacheSizeEstimator.estimateManifestEntry(deleteManifestKey(10), empty); - Assertions.assertTrue(narrow.isComplete() && wide.isComplete()); - long extraFieldBytes = wide.getBytes() - narrow.getBytes(); - Assertions.assertTrue(extraFieldBytes >= 9 * JvmSizeUtils.instanceSize(Types.NestedField.class), - "narrow=" + narrow.getBytes() + ", wide=" + wide.getBytes()); - } - - private static IcebergManifestEntryKey deleteManifestKey(int partitionFieldCount) { - List fields = new ArrayList<>(); - for (int i = 0; i < partitionFieldCount; i++) { - fields.add(Types.NestedField.optional(1000 + i, "p" + i, Types.IntegerType.get())); - } - return new IcebergManifestEntryKey("/m/deletes.avro", ManifestContent.DELETES, Types.StructType.of(fields)); - } - @Test public void equalityDeleteFieldIdsLoadOncePerSnapshot() { IcebergManifestCache cache = new IcebergManifestCache(); @@ -198,9 +785,9 @@ public void capacityOverflowFlushesWholesale() { // maxSize 1: the first load fills it; a re-read still hits (same key). The wholesale-flush valve is the // legacy behavior — re-reads are harmless since the value is immutable. Smoke that size stays bounded. IcebergManifestCache cache = new IcebergManifestCache(1); - cache.getManifestCacheValue(manifest, table, table.specs()); + cache.getManifestCacheValue(manifest, table); Assertions.assertEquals(1, cache.size()); - cache.getManifestCacheValue(manifest, table, table.specs()); + cache.getManifestCacheValue(manifest, table); Assertions.assertTrue(cache.size() <= 1, "a bounded cache must not grow past its max size"); } @@ -209,7 +796,7 @@ public void invalidateAllClearsEveryEntry() { Table table = tableWithTwoDataFiles(); ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); IcebergManifestCache cache = new IcebergManifestCache(); - cache.getManifestCacheValue(manifest, table, table.specs()); + cache.getManifestCacheValue(manifest, table); AtomicInteger equalityLoads = new AtomicInteger(); cache.getOrLoadEqualityDeleteFieldIds(table.location(), table.currentSnapshot().snapshotId(), () -> { equalityLoads.incrementAndGet(); @@ -231,8 +818,7 @@ public void invalidateAllClearsEveryEntry() { public void copiesDataFilesSoReaderReuseDoesNotAlias() { Table table = tableWithTwoDataFiles(); ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); - List files = new IcebergManifestCache() - .getManifestCacheValue(manifest, table, table.specs()).getDataFiles(); + List files = new IcebergManifestCache().getManifestCacheValue(manifest, table).getDataFiles(); // WHY: ManifestReader reuses one object across iterations, so the cache must .copy() each file. Two // distinct paths prove the entries were copied out (not the same reused instance). MUTATION: dropping // .copy() -> both entries alias the last file -> both paths equal -> red. diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKeyTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKeyTest.java index ecd4c08d213357..96371cfa606fac 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKeyTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestEntryKeyTest.java @@ -18,25 +18,19 @@ package org.apache.doris.connector.iceberg; import org.apache.iceberg.ManifestContent; -import org.apache.iceberg.types.Types; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; /** - * Unit tests for {@link IcebergManifestEntryKey} (T08). The key is path + content + bound partition type; two - * manifests that share a path (across tables) hit the same cache entry, which is what lets a REFRESH TABLE keep the - * manifest cache. + * Unit tests for {@link IcebergManifestEntryKey} (T08). The key is path + content; two manifests that share a + * path (across tables) hit the same cache entry, which is what lets a REFRESH TABLE keep the manifest cache. */ public class IcebergManifestEntryKeyTest { - private static final Types.StructType INT_PARTITION = Types.StructType.of( - Types.NestedField.optional(1000, "p", Types.IntegerType.get())); - private static final Types.StructType LONG_PARTITION = Types.StructType.of( - Types.NestedField.optional(1000, "p", Types.LongType.get())); @Test public void samePathSameContentAreEqual() { - IcebergManifestEntryKey a = new IcebergManifestEntryKey("/m/shared.avro", ManifestContent.DATA, INT_PARTITION); - IcebergManifestEntryKey b = new IcebergManifestEntryKey("/m/shared.avro", ManifestContent.DATA, INT_PARTITION); + IcebergManifestEntryKey a = new IcebergManifestEntryKey("/m/shared.avro", ManifestContent.DATA); + IcebergManifestEntryKey b = new IcebergManifestEntryKey("/m/shared.avro", ManifestContent.DATA); // WHY: the path-keyed cache must treat two references to the same manifest path as one entry (cross-table // sharing). MUTATION: keying on identity instead of (path, content) -> not equal -> red. Assertions.assertEquals(a, b); @@ -45,9 +39,8 @@ public void samePathSameContentAreEqual() { @Test public void differentContentSamePathAreNotEqual() { - IcebergManifestEntryKey data = new IcebergManifestEntryKey("/m/x.avro", ManifestContent.DATA, INT_PARTITION); - IcebergManifestEntryKey deletes = - new IcebergManifestEntryKey("/m/x.avro", ManifestContent.DELETES, INT_PARTITION); + IcebergManifestEntryKey data = new IcebergManifestEntryKey("/m/x.avro", ManifestContent.DATA); + IcebergManifestEntryKey deletes = new IcebergManifestEntryKey("/m/x.avro", ManifestContent.DELETES); // A data manifest and a delete manifest are distinct cache entries even at the same path. MUTATION: // ignoring content -> equal -> red. Assertions.assertNotEquals(data, deletes); @@ -56,17 +49,7 @@ public void differentContentSamePathAreNotEqual() { @Test public void differentPathAreNotEqual() { Assertions.assertNotEquals( - new IcebergManifestEntryKey("/m/a.avro", ManifestContent.DATA, INT_PARTITION), - new IcebergManifestEntryKey("/m/b.avro", ManifestContent.DATA, INT_PARTITION)); - } - - @Test - public void differentBoundPartitionTypeAreNotEqual() { - // The reader converts partition values to the bound spec's types, so a manifest decoded for a scan before an - // INT -> BIGINT promotion and one decoded for a current scan hold different values. MUTATION: ignoring the - // partition type -> equal -> one scan reads the other's Integer/Long values -> red. - Assertions.assertNotEquals( - new IcebergManifestEntryKey("/m/x.avro", ManifestContent.DATA, INT_PARTITION), - new IcebergManifestEntryKey("/m/x.avro", ManifestContent.DATA, LONG_PARTITION)); + new IcebergManifestEntryKey("/m/a.avro", ManifestContent.DATA), + new IcebergManifestEntryKey("/m/b.avro", ManifestContent.DATA)); } } diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestPrefetcherTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestPrefetcherTest.java new file mode 100644 index 00000000000000..9810b95c75c948 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestPrefetcherTest.java @@ -0,0 +1,448 @@ +// 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.iceberg; + +import org.apache.doris.connector.cache.CacheSpec; +import org.apache.doris.connector.cache.CatalogMetaCache; +import org.apache.doris.connector.cache.MetaCache; +import org.apache.doris.connector.cache.MetaCacheDefinition; +import org.apache.doris.connector.cache.ScopePath; +import org.apache.doris.connector.spi.ConnectorContext; + +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.CloseableIterator; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Proxy; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; + +public class IcebergManifestPrefetcherTest { + private static final ManifestCacheValue EMPTY = ManifestCacheValue.forDataFiles(Collections.emptyList()); + + private static CloseableIterator manifests(int count) { + List files = new ArrayList<>(); + for (int i = 0; i < count; i++) { + String path = "manifest-" + i; + files.add((ManifestFile) Proxy.newProxyInstance(ManifestFile.class.getClassLoader(), + new Class[] {ManifestFile.class}, (proxy, method, args) -> { + if (method.getName().equals("path")) { + return path; + } + throw new UnsupportedOperationException(method.getName()); + })); + } + return CloseableIterable.withNoopClose(files).iterator(); + } + + private static class Context implements ConnectorContext { + final List> submitted = new ArrayList<>(); + final ExecutorService executor; + final AtomicInteger authenticated = new AtomicInteger(); + + Context(ExecutorService executor) { + this.executor = executor; + } + + @Override + public String getCatalogName() { + return "test"; + } + + @Override + public long getCatalogId() { + return 1; + } + + @Override + public synchronized RunnableFuture submitMetadataTask(Callable callable) { + FutureTask future = new FutureTask<>(callable); + submitted.add(future); + if (executor != null) { + executor.execute(future); + } + return future; + } + + @Override + public T executeAuthenticated(Callable callable) throws Exception { + Assertions.assertSame(IcebergManifestPrefetcher.class.getClassLoader(), + Thread.currentThread().getContextClassLoader()); + authenticated.incrementAndGet(); + return callable.call(); + } + } + + @Test + public void warmCacheAvoidsSubmissionAndSpeculativeTraversal() throws Exception { + Context context = new Context(null); + AtomicInteger probes = new AtomicInteger(); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(20), 8, + manifest -> { + probes.incrementAndGet(); + return EMPTY; + }, + manifest -> { + throw new AssertionError("warm cache loaded remotely"); + }, context, new Semaphore(8))) { + Assertions.assertEquals("manifest-0", source.next().manifest.path()); + Assertions.assertEquals(1, probes.get()); + Assertions.assertTrue(context.submitted.isEmpty()); + } + } + + @Test + public void saturatedExecutorDegradesToOwnTaskAndCloseCancelsUnusedWindow() throws Exception { + Context context = new Context(null); // accepts tasks but has no worker available + AtomicInteger loads = new AtomicInteger(); + IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(100), 4, + manifest -> null, manifest -> { + loads.incrementAndGet(); + return EMPTY; + }, context, new Semaphore(4)); + Assertions.assertEquals("manifest-0", source.next().manifest.path()); + Assertions.assertEquals(0, context.submitted.size(), "first manifest must not submit executor work"); + Assertions.assertEquals(1, loads.get()); + Assertions.assertEquals("manifest-1", source.next().manifest.path()); + Assertions.assertEquals(4, context.submitted.size(), "subsequent demand expands the bounded window"); + Assertions.assertEquals(2, loads.get()); + source.close(); + source.close(); + Assertions.assertFalse(source.hasNext()); + for (RunnableFuture task : context.submitted) { + task.run(); + } + Assertions.assertEquals(2, loads.get(), "cancelled queued tasks must never start a read"); + Assertions.assertEquals(2, context.authenticated.get()); + } + + @Test + public void demandControlsPrefetchWithoutBlockingSerialProgress() throws Exception { + Context context = new Context(null); + java.util.concurrent.atomic.AtomicBoolean demand = new java.util.concurrent.atomic.AtomicBoolean(); + AtomicInteger loads = new AtomicInteger(); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(20), 4, + manifest -> null, manifest -> { + loads.incrementAndGet(); + return EMPTY; + }, context, new Semaphore(4), demand::get)) { + Assertions.assertEquals("manifest-0", source.next().manifest.path()); + Assertions.assertEquals("manifest-1", source.next().manifest.path()); + Assertions.assertTrue(context.submitted.isEmpty()); + demand.set(true); + Assertions.assertEquals("manifest-2", source.next().manifest.path()); + Assertions.assertEquals(4, context.submitted.size()); + demand.set(false); + Assertions.assertEquals("manifest-3", source.next().manifest.path()); + Assertions.assertEquals(4, context.submitted.size(), "paused demand must not refill the window"); + Assertions.assertEquals(4, loads.get()); + } + Assertions.assertTrue(context.submitted.stream().allMatch(Future::isDone)); + } + + @Test + public void sequentialSettingDoesNotSubmitAndKeepsOrder() throws Exception { + Context context = new Context(null); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(3), 1, + manifest -> null, manifest -> EMPTY, context, new Semaphore(1))) { + for (int i = 0; i < 3; i++) { + Assertions.assertEquals("manifest-" + i, source.next().manifest.path()); + } + Assertions.assertFalse(source.hasNext()); + Assertions.assertTrue(context.submitted.isEmpty()); + } + } + + @Test + public void loadFailurePropagatesAndCancelsTheRemainingWindow() throws Exception { + Context context = new Context(null); + IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(10), 4, + manifest -> null, manifest -> { + if (manifest.path().equals("manifest-0")) { + return EMPTY; + } + throw new IllegalStateException("manifest read failed"); + }, + context, new Semaphore(4)); + Assertions.assertEquals("manifest-0", source.next().manifest.path()); + RuntimeException failure = Assertions.assertThrows(RuntimeException.class, source::next); + Assertions.assertEquals("manifest read failed", failure.getCause().getMessage()); + Assertions.assertFalse(source.hasNext()); + Assertions.assertTrue(context.submitted.stream().allMatch(Future::isDone)); + } + + @Test + public void multipleScansShareReadLimitAndReturnEveryManifestOnce() throws Exception { + ExecutorService workers = Executors.newFixedThreadPool(6); + ExecutorService parents = Executors.newFixedThreadPool(2); + Context context = new Context(workers); + Semaphore permits = new Semaphore(2); + CountDownLatch twoReading = new CountDownLatch(2); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger active = new AtomicInteger(); + AtomicInteger peak = new AtomicInteger(); + Function load = manifest -> { + int current = active.incrementAndGet(); + peak.accumulateAndGet(current, Math::max); + twoReading.countDown(); + try { + Assertions.assertTrue(release.await(5, TimeUnit.SECONDS)); + return EMPTY; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } finally { + active.decrementAndGet(); + } + }; + try (IcebergManifestPrefetcher first = new IcebergManifestPrefetcher(manifests(12), 4, + manifest -> null, load, context, permits); + IcebergManifestPrefetcher second = new IcebergManifestPrefetcher(manifests(12), 4, + manifest -> null, load, context, permits)) { + List>> results = new ArrayList<>(); + for (IcebergManifestPrefetcher source : List.of(first, second)) { + results.add(parents.submit(() -> { + List paths = new ArrayList<>(); + while (source.hasNext()) { + paths.add(source.next().manifest.path()); + } + return paths; + })); + } + Assertions.assertTrue(twoReading.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(2, peak.get()); + release.countDown(); + for (Future> result : results) { + List paths = result.get(5, TimeUnit.SECONDS); + Set expected = new HashSet<>(); + for (int i = 0; i < 12; i++) { + expected.add("manifest-" + i); + } + Assertions.assertEquals(expected, new HashSet<>(paths)); + Assertions.assertEquals(12, paths.size()); + } + Assertions.assertEquals(2, peak.get()); + Assertions.assertEquals(2, permits.availablePermits()); + Assertions.assertEquals(24, context.authenticated.get()); + } finally { + release.countDown(); + workers.shutdownNow(); + parents.shutdownNow(); + } + } + + @Test + public void cancellationPreservesSharedLoadAndPreventsFurtherRequests() throws Exception { + ExecutorService parents = Executors.newFixedThreadPool(2); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch secondReading = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger loads = new AtomicInteger(); + AtomicInteger interrupted = new AtomicInteger(); + try (CatalogMetaCache owner = CatalogMetaCache.unmanaged()) { + MetaCache cache = owner.create(MetaCacheDefinition + .builder("shared-manifest", CacheSpec.of(true, -1, 10), + ignored -> ScopePath.catalog()).build()); + Function loader = manifest -> cache.get(manifest.path(), key -> { + loads.incrementAndGet(); + entered.countDown(); + try { + Assertions.assertTrue(release.await(5, TimeUnit.SECONDS)); + return EMPTY; + } catch (InterruptedException e) { + interrupted.incrementAndGet(); + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }); + try (IcebergManifestPrefetcher first = new IcebergManifestPrefetcher(manifests(30), 4, + manifest -> null, loader, null, new Semaphore(4)); + IcebergManifestPrefetcher second = new IcebergManifestPrefetcher(manifests(1), 4, + manifest -> null, manifest -> { + secondReading.countDown(); + return loader.apply(manifest); + }, null, new Semaphore(4))) { + Future cancelled = parents.submit(first::next); + Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS)); + Future survivor = parents.submit(() -> + second.next().manifest.path()); + Assertions.assertTrue(secondReading.await(5, TimeUnit.SECONDS)); + first.close(); + release.countDown(); + Assertions.assertEquals("manifest-0", survivor.get(5, TimeUnit.SECONDS)); + Assertions.assertThrows(java.util.concurrent.ExecutionException.class, + () -> cancelled.get(5, TimeUnit.SECONDS)); + Assertions.assertFalse(first.hasNext()); + Assertions.assertEquals(0, interrupted.get(), "one query must not interrupt a shared cache loader"); + Assertions.assertEquals(1, loads.get(), "cancelled queued manifests must not start new requests"); + } + } finally { + release.countDown(); + parents.shutdownNow(); + } + } + + private static void await(CountDownLatch latch) { + try { + Assertions.assertTrue(latch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + private static Context startedHeadContext(ExecutorService workers, CountDownLatch started) { + return new Context(workers) { + @Override + public synchronized RunnableFuture submitMetadataTask(Callable callable) { + RunnableFuture task = super.submitMetadataTask(callable); + if (submitted.size() == 1) { + await(started); + } + return task; + } + }; + } + + @Test + public void completedManifestBypassesSlowHeadWithoutGrowingPausedWindow() throws Exception { + ExecutorService workers = Executors.newFixedThreadPool(3); + ExecutorService consumer = Executors.newSingleThreadExecutor(); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + Context context = startedHeadContext(workers, started); + AtomicInteger loads = new AtomicInteger(); + java.util.concurrent.atomic.AtomicBoolean demand = new java.util.concurrent.atomic.AtomicBoolean(true); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(20), 3, + manifest -> null, manifest -> { + loads.incrementAndGet(); + if (manifest.path().equals("manifest-1")) { + started.countDown(); + await(release); + } + return EMPTY; + }, context, new Semaphore(3), demand::get)) { + Assertions.assertEquals("manifest-0", source.next().manifest.path()); + String first = consumer.submit(() -> source.next().manifest.path()).get(5, TimeUnit.SECONDS); + Assertions.assertNotEquals("manifest-1", first, "a completed manifest must bypass the slow head"); + Assertions.assertEquals(3, context.submitted.size()); + demand.set(false); + String second = consumer.submit(() -> source.next().manifest.path()).get(5, TimeUnit.SECONDS); + Assertions.assertEquals(Set.of("manifest-2", "manifest-3"), Set.of(first, second)); + Assertions.assertEquals(3, context.submitted.size(), "paused consumption must not refill the window"); + Assertions.assertEquals(4, loads.get(), "completed and running results share one bounded window"); + release.countDown(); + Assertions.assertEquals("manifest-1", source.next().manifest.path()); + demand.set(true); + Set seen = new HashSet<>(Set.of("manifest-0", first, second, "manifest-1")); + while (source.hasNext()) { + Assertions.assertTrue(seen.add(source.next().manifest.path()), "duplicate manifest"); + } + Assertions.assertEquals(20, seen.size()); + } finally { + release.countDown(); + workers.shutdownNow(); + consumer.shutdownNow(); + } + } + + @Test + public void completedFailureBypassesSlowHeadAndClosesWindow() throws Exception { + ExecutorService workers = Executors.newFixedThreadPool(2); + ExecutorService consumer = Executors.newSingleThreadExecutor(); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + Context context = startedHeadContext(workers, started); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(3), 2, + manifest -> null, manifest -> { + if (manifest.path().equals("manifest-1")) { + started.countDown(); + await(release); + } else if (manifest.path().equals("manifest-2")) { + throw new IllegalStateException("later manifest failed"); + } + return EMPTY; + }, context, new Semaphore(2))) { + source.next(); + Future next = consumer.submit(source::next); + java.util.concurrent.ExecutionException failure = Assertions.assertThrows( + java.util.concurrent.ExecutionException.class, () -> next.get(5, TimeUnit.SECONDS)); + Assertions.assertEquals("later manifest failed", failure.getCause().getCause().getMessage()); + Assertions.assertFalse(source.hasNext()); + Assertions.assertTrue(context.submitted.stream().allMatch(Future::isDone)); + } finally { + release.countDown(); + workers.shutdownNow(); + consumer.shutdownNow(); + } + } + + @Test + public void closeWakesCompletionWaitWithoutInterruptingSharedRead() throws Exception { + ExecutorService workers = Executors.newSingleThreadExecutor(); + ExecutorService consumer = Executors.newSingleThreadExecutor(); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + Context context = startedHeadContext(workers, started); + AtomicInteger loads = new AtomicInteger(); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher(manifests(10), 3, + manifest -> null, manifest -> { + loads.incrementAndGet(); + if (manifest.path().equals("manifest-1")) { + started.countDown(); + await(release); + } + return EMPTY; + }, context, new Semaphore(3))) { + source.next(); + Future next = consumer.submit(source::next); + await(started); + source.close(); + java.util.concurrent.ExecutionException failure = Assertions.assertThrows( + java.util.concurrent.ExecutionException.class, () -> next.get(5, TimeUnit.SECONDS)); + Assertions.assertInstanceOf(java.util.concurrent.CancellationException.class, failure.getCause()); + release.countDown(); + for (RunnableFuture task : context.submitted) { + task.run(); + } + Assertions.assertEquals(2, loads.get(), "closing must cancel queued reads"); + Assertions.assertFalse(source.hasNext()); + } finally { + release.countDown(); + workers.shutdownNow(); + consumer.shutdownNow(); + } + } + +} diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestReadContextTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestReadContextTest.java new file mode 100644 index 00000000000000..97ca9b1e4b3840 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergManifestReadContextTest.java @@ -0,0 +1,310 @@ +// 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.iceberg; + +import org.apache.doris.connector.spi.ConnectorSession; +import org.apache.doris.connector.spi.ConnectorStatementScopes; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; + +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableScan; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.time.ZoneId; +import java.util.Arrays; +import java.util.Collections; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.UnaryOperator; + +@Timeout(30) +class IcebergManifestReadContextTest { + @ParameterizedTest + @CsvSource({"false, false", "true, false", "false, true", "true, true"}) + void deleteManifestReadsShareTheDataReadBudget(boolean cancel, boolean equality) throws Exception { + CountDownLatch waiting = new CountDownLatch(1); + Semaphore budget = new Semaphore(1) { + @Override + public boolean tryAcquire(long timeout, TimeUnit unit) throws InterruptedException { + waiting.countDown(); + return super.tryAcquire(timeout, unit); + } + }; + TestStatementScope scope = new TestStatementScope(); + AtomicBoolean closed = new AtomicBoolean(); + ConnectorTaskScheduler scheduler = (ConnectorTaskScheduler) Proxy.newProxyInstance( + ConnectorTaskScheduler.class.getClassLoader(), new Class[] {ConnectorTaskScheduler.class}, + (proxy, method, args) -> { + if (method.getName().equals("isClosed")) { + return closed.get(); + } + throw new AssertionError("unexpected scheduler operation: " + method.getName()); + }); + ConnectorSession delegate = new IcebergStatementScopeTest.ScopeSession(7, "union", scope); + ConnectorSession session = (ConnectorSession) Proxy.newProxyInstance( + ConnectorSession.class.getClassLoader(), new Class[] {ConnectorSession.class}, + (proxy, method, args) -> { + if (method.getName().equals("getSessionProperties")) { + return Map.of("iceberg_manifest_load_concurrency", "1"); + } + if (method.getName().equals("getTaskScheduler")) { + return scheduler; + } + return method.invoke(delegate, args); + }); + scope.computeIfAbsent(IcebergStatementScope.MANIFEST_READ_BUDGET_NAMESPACE + ":union", () -> budget); + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("delete-budget", Collections.emptyMap()); + catalog.createNamespace(Namespace.of("db1")); + Table original = catalog.createTable(TableIdentifier.of("db1", "t1"), schema, + PartitionSpec.unpartitioned(), Map.of("format-version", "2")); + original.newAppend().appendFile(DataFiles.builder(original.spec()) + .withPath("s3://bucket/data.parquet").withFileSizeInBytes(128).withRecordCount(1).build()).commit(); + org.apache.iceberg.DeleteFile delete = equality + ? FileMetadata.deleteFileBuilder(original.spec()).ofEqualityDeletes(1) + .withPath("s3://bucket/delete.parquet").withFileSizeInBytes(128).withRecordCount(1).build() + : FileMetadata.deleteFileBuilder(original.spec()).ofPositionDeletes() + .withPath("s3://bucket/delete.parquet").withFileSizeInBytes(128).withRecordCount(1).build(); + original.newRowDelta().addDeletes(delete).commit(); + String deletePath = original.currentSnapshot().deleteManifests(original.io()).get(0).path(); + AtomicInteger deleteReads = new AtomicInteger(); + FileIO io = (FileIO) Proxy.newProxyInstance(FileIO.class.getClassLoader(), new Class[] {FileIO.class}, + (proxy, method, args) -> { + if (method.getName().equals("newInputFile")) { + Object file = args[0]; + String path = file instanceof ManifestFile ? ((ManifestFile) file).path() : file.toString(); + if (deletePath.equals(path)) { + deleteReads.incrementAndGet(); + } + } + return method.invoke(original.io(), args); + }); + Table table = (Table) Proxy.newProxyInstance(Table.class.getClassLoader(), new Class[] {Table.class}, + (proxy, method, args) -> { + if (method.getName().equals("io")) { + return io; + } + return method.invoke(original, args); + }); + ConnectorStatementScopes.resolveInStatement(session, IcebergStatementScope.TABLE_NAMESPACE, + "db1", "t1", () -> table); + TableScan scan = original.newScan(); + IcebergManifestCache cache = new IcebergManifestCache(); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + IcebergCatalogProperties.of(Map.of(IcebergConnector.MANIFEST_CACHE_ENABLE, "true")), + new RecordingIcebergCatalogOps(), null, cache); + Method plan = IcebergScanPlanProvider.class.getDeclaredMethod("planFileScanTask", + TableScan.class, ConnectorSession.class, Table.class, Optional.class, IcebergTableHandle.class); + plan.setAccessible(true); + ExecutorService worker = Executors.newSingleThreadExecutor(); + budget.acquire(); // The first scan is currently loading a cold data manifest. + try { + Future secondScan = worker.submit(() -> equality + ? provider.getScanNodeProperties(session, new IcebergTableHandle("db1", "t1"), + Collections.emptyList(), Optional.empty()) + : plan.invoke(provider, scan, session, table, Optional.empty(), + new IcebergTableHandle("db1", "t1"))); + Assertions.assertTrue(waiting.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(0, deleteReads.get(), "delete reads cannot bypass the shared query budget"); + if (cancel) { + closed.set(true); + ExecutionException failure = Assertions.assertThrows(ExecutionException.class, + () -> secondScan.get(5, TimeUnit.SECONDS)); + Throwable cause = failure; + while (cause.getCause() != null) { + cause = cause.getCause(); + } + Assertions.assertInstanceOf(CancellationException.class, cause); + Assertions.assertEquals(0, deleteReads.get(), "cancellation must not restart SDK planning"); + Assertions.assertEquals(0, budget.availablePermits(), "the other scan still owns its permit"); + budget.release(); + } else { + budget.release(); + secondScan.get(5, TimeUnit.SECONDS); + Assertions.assertTrue(deleteReads.get() > 0); + Assertions.assertEquals(1, budget.availablePermits()); + } + } finally { + worker.shutdownNow(); + scope.closeAll(); + } + } + + @ParameterizedTest + @CsvSource({"streaming, acquisition", "materialized, acquisition", "representative, acquisition", + "count, acquisition", "streaming, failure", "materialized, failure", + "representative, failure", "count, failure"}) + void cancellationBeforeSchedulerAcquisitionCannotRestartSdkPlanning(String path, String trigger) throws Exception { + AtomicBoolean closed = new AtomicBoolean(); + ConnectorTaskScheduler scheduler = (ConnectorTaskScheduler) Proxy.newProxyInstance( + ConnectorTaskScheduler.class.getClassLoader(), new Class[] {ConnectorTaskScheduler.class}, + (proxy, method, args) -> { + if (method.getName().equals("isClosed")) { + return closed.get(); + } + throw new AssertionError("unexpected scheduler operation: " + method.getName()); + }); + ConnectorSession delegate = new IcebergStatementScopeTest.ScopeSession(7, "cancelled", new TestStatementScope()); + ConnectorSession session = (ConnectorSession) Proxy.newProxyInstance( + ConnectorSession.class.getClassLoader(), new Class[] {ConnectorSession.class}, + (proxy, method, args) -> { + if (method.getName().equals("getTaskScheduler")) { + if (trigger.equals("acquisition")) { + // QueryTaskGroups.get reports cancellation after the captured query was closed. + throw new CancellationException("Query task groups are closed"); + } + return scheduler; + } + return method.invoke(delegate, args); + }); + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("cancel-before-acquisition", Collections.emptyMap()); + catalog.createNamespace(Namespace.of("db1")); + Table table = catalog.createTable(TableIdentifier.of("db1", "t1"), schema, + PartitionSpec.unpartitioned(), Map.of("format-version", "2")); + table.newAppend().appendFile(DataFiles.builder(table.spec()).withPath("s3://bucket/data.parquet") + .withFileSizeInBytes(128).withRecordCount(1).build()).commit(); + table.newRowDelta().addDeletes(FileMetadata.deleteFileBuilder(table.spec()).ofPositionDeletes() + .withPath("s3://bucket/delete.parquet").withFileSizeInBytes(128).withRecordCount(1).build()).commit(); + AtomicInteger sdkPlans = new AtomicInteger(); + TableScan original = table.newScan(); + TableScan scan = (TableScan) Proxy.newProxyInstance(TableScan.class.getClassLoader(), + new Class[] {TableScan.class}, (proxy, method, args) -> { + if (method.getName().equals("planFiles")) { + sdkPlans.incrementAndGet(); + } + if (method.getName().equals("snapshot") && trigger.equals("failure")) { + closed.set(true); + throw new IllegalStateException("cache setup failed during cancellation"); + } + return method.invoke(original, args); + }); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + IcebergCatalogProperties.of(Map.of(IcebergConnector.MANIFEST_CACHE_ENABLE, "true")), + new RecordingIcebergCatalogOps(), null, new IcebergManifestCache()); + IcebergTableHandle handle = new IcebergTableHandle("db1", "t1"); + String methodName; + Object[] arguments; + switch (path) { + case "streaming": + methodName = "streamingFileScanTasks"; + arguments = new Object[] {scan, session, table, Optional.empty(), 128L, null, handle}; + break; + case "materialized": + methodName = "planFileScanTask"; + arguments = new Object[] {scan, session, table, Optional.empty(), handle}; + break; + case "representative": + methodName = "planManifestCountRange"; + arguments = new Object[] {table, scan, scan.schema(), 1L, 2, false, Collections.emptyList(), + ZoneId.of("UTC"), UnaryOperator.identity(), + session, Optional.empty(), true, handle}; + break; + case "count": + methodName = "planCountPushdownFromFileTasks"; + arguments = new Object[] {table, scan, scan.schema(), 2, false, Collections.emptyList(), + ZoneId.of("UTC"), UnaryOperator.identity(), + session, Optional.empty(), true, 0L, handle}; + break; + default: + throw new AssertionError(path); + } + Method entry = Arrays.stream(IcebergScanPlanProvider.class.getDeclaredMethods()) + .filter(method -> method.getName().equals(methodName)).findFirst().orElseThrow(); + entry.setAccessible(true); + InvocationTargetException error = Assertions.assertThrows( + InvocationTargetException.class, () -> entry.invoke(provider, arguments)); + Assertions.assertInstanceOf(CancellationException.class, error.getCause()); + Assertions.assertEquals(0, sdkPlans.get(), "cancelled cache planning must not open a fresh SDK delete index"); + } + + @Test + void offlineEmptyScanDoesNotAcquireAQueryScheduler() throws Exception { + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("offline-empty", Collections.emptyMap()); + catalog.createNamespace(Namespace.of("db1")); + Table table = catalog.createTable(TableIdentifier.of("db1", "t1"), + new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get()))); + RecordingConnectorContext context = new RecordingConnectorContext(); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + IcebergCatalogProperties.of(Map.of(IcebergConnector.MANIFEST_CACHE_ENABLE, "true")), + new RecordingIcebergCatalogOps(), context, new IcebergManifestCache()); + Method plan = IcebergScanPlanProvider.class.getDeclaredMethod("cacheBackedFileScanTasks", + TableScan.class, ConnectorSession.class, Table.class, Optional.class, + String.class, IcebergTableHandle.class); + plan.setAccessible(true); + try (CloseableIterable tasks = (CloseableIterable) plan.invoke(provider, table.newScan(), null, + table, Optional.empty(), null, new IcebergTableHandle("db1", "t1"))) { + Assertions.assertFalse(tasks.iterator().hasNext()); + } + Assertions.assertEquals(0, context.metadataSubmissions.get()); + } + + private static ManifestFile sourceManifest() { + return (ManifestFile) Proxy.newProxyInstance(ManifestFile.class.getClassLoader(), + new Class[] {ManifestFile.class}, (proxy, method, args) -> "delete.avro"); + } + + @Test + void failedReadReleasesPermitAndOwner() { + Semaphore budget = new Semaphore(1); + AtomicInteger owners = new AtomicInteger(); + IcebergManifestReadContext reads = new IcebergManifestReadContext(null, budget, () -> false, + () -> { + owners.incrementAndGet(); + return owners::decrementAndGet; + }); + IllegalStateException failure = new IllegalStateException("broken delete manifest"); + Assertions.assertSame(failure, Assertions.assertThrows(IllegalStateException.class, + () -> reads.read(sourceManifest(), m -> { + throw failure; + }))); + Assertions.assertEquals(0, owners.get()); + Assertions.assertEquals(1, budget.availablePermits()); + } +} diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergPartitionValueAdapterTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergPartitionValueAdapterTest.java new file mode 100644 index 00000000000000..271d2bc2883329 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergPartitionValueAdapterTest.java @@ -0,0 +1,137 @@ +// 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.iceberg; + +import org.apache.iceberg.ContentFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.PartitionData; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeWrapper; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Method; +import java.math.BigDecimal; +import java.util.Collections; + +public class IcebergPartitionValueAdapterTest { + @Test + public void deleteFileViewPreservesAllMetadataAndCopies() throws Exception { + Schema schema = new Schema(Types.NestedField.optional(1, "p", Types.LongType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("p").build(); + PartitionData partition = new PartitionData(spec.rawPartitionType()); + partition.set(0, 7L); + DeleteFile file = FileMetadata.deleteFileBuilder(spec).ofPositionDeletes() + .withPath("/delete.puffin").withFormat("PUFFIN").withPartition(partition) + .withRecordCount(2).withFileSizeInBytes(100) + .withReferencedDataFile("/data.parquet").withContentOffset(10).withContentSizeInBytes(20).build(); + Types.StructType target = Types.StructType.of( + Types.NestedField.optional(1000, "p", Types.IntegerType.get())); + DeleteFile view = new IcebergPartitionValueAdapter(spec.rawPartitionType(), target).adapt(file); + for (Method method : DeleteFile.class.getMethods()) { + if (method.getParameterCount() == 0 && !method.getName().equals("partition") + && !method.getName().startsWith("copy")) { + Assertions.assertEquals(method.invoke(file), method.invoke(view), method.getName()); + } + } + for (ContentFile copy : new ContentFile[] {view, view.copy(), view.copyWithoutStats(), + view.copyWithStats(Collections.emptySet())}) { + Assertions.assertEquals(7, copy.partition().get(0, Integer.class)); + } + ((PartitionData) view.partition()).set(0, 8); + Assertions.assertEquals(7L, file.partition().get(0, Long.class)); + } + + @Test + public void adaptationUsesFieldIdsAndPreservesNulls() { + Types.StructType sourceType = Types.StructType.of( + Types.NestedField.optional(1000, "old_name", Types.IntegerType.get()), + Types.NestedField.optional(1001, "nullable", Types.StringType.get())); + PartitionData source = new PartitionData(sourceType); + source.set(0, 7); + Types.StructType target = Types.StructType.of(sourceType.fields().get(1), + Types.NestedField.optional(1000, "new_name", Types.LongType.get())); + PartitionData result = new IcebergPartitionValueAdapter(sourceType, target).adapt(source); + Assertions.assertNull(result.get(0, Object.class)); + Assertions.assertEquals(7L, result.get(1, Long.class)); + Assertions.assertEquals(7, source.get(0, Integer.class)); + } + + @Test + public void lossyHistoricalConversionsFailInsteadOfTruncating() { + Assertions.assertThrows(ArithmeticException.class, + () -> convert(Types.LongType.get(), 3000000000L, Types.IntegerType.get())); + Assertions.assertThrows(IllegalArgumentException.class, + () -> convert(Types.DoubleType.get(), 0.1D, Types.FloatType.get())); + Assertions.assertThrows(IllegalArgumentException.class, + () -> convert(Types.DecimalType.of(18, 2), new BigDecimal("1234567890.12"), + Types.DecimalType.of(9, 2))); + Assertions.assertThrows(IllegalArgumentException.class, + () -> convert(Types.StringType.get(), "7", Types.IntegerType.get())); + } + + @Test + public void floatingSpecialValuesSurviveHistoricalConversion() { + Assertions.assertEquals(Float.NaN, convert(Types.DoubleType.get(), Double.NaN, Types.FloatType.get())); + Assertions.assertEquals(Float.NEGATIVE_INFINITY, + convert(Types.DoubleType.get(), Double.NEGATIVE_INFINITY, Types.FloatType.get())); + Assertions.assertEquals(-0.0F, convert(Types.DoubleType.get(), -0.0D, Types.FloatType.get())); + } + + @Test + public void droppedSourcePromotionsHaveEqualPartitionKeys() { + assertUnknownPartitionKeysEqual(Types.IntegerType.get(), 7, Types.LongType.get(), 7L); + assertUnknownPartitionKeysEqual(Types.IntegerType.get(), Integer.MIN_VALUE, + Types.LongType.get(), (long) Integer.MIN_VALUE); + assertUnknownPartitionKeysEqual(Types.FloatType.get(), 1.25F, Types.DoubleType.get(), 1.25D); + assertUnknownPartitionKeysEqual(Types.FloatType.get(), 0.1F, Types.DoubleType.get(), (double) 0.1F); + assertUnknownPartitionKeysEqual(Types.FloatType.get(), -0.0F, Types.DoubleType.get(), -0.0D); + assertUnknownPartitionKeysEqual(Types.FloatType.get(), Float.NaN, Types.DoubleType.get(), Double.NaN); + assertUnknownPartitionKeysEqual(Types.FloatType.get(), Float.POSITIVE_INFINITY, + Types.DoubleType.get(), Double.POSITIVE_INFINITY); + assertUnknownPartitionKeysEqual(Types.DecimalType.of(9, 2), new BigDecimal("1.25"), + Types.DecimalType.of(18, 2), new BigDecimal("1.25")); + Assertions.assertEquals(9007199254740993L, + convert(Types.LongType.get(), 9007199254740993L, Types.UnknownType.get())); + Assertions.assertNull(convert(Types.IntegerType.get(), null, Types.UnknownType.get())); + } + + private void assertUnknownPartitionKeysEqual(org.apache.iceberg.types.Type from, Object original, + org.apache.iceberg.types.Type promotedType, Object promoted) { + Types.StructType target = Types.StructType.of( + Types.NestedField.optional(1000, "p", Types.UnknownType.get())); + PartitionData first = new PartitionData(target); + first.set(0, convert(from, original, Types.UnknownType.get())); + PartitionData second = new PartitionData(target); + second.set(0, convert(promotedType, promoted, Types.UnknownType.get())); + StructLikeWrapper left = StructLikeWrapper.forType(target).set(first); + StructLikeWrapper right = StructLikeWrapper.forType(target).set(second); + Assertions.assertEquals(left, right); + Assertions.assertEquals(left.hashCode(), right.hashCode()); + } + + private Object convert(org.apache.iceberg.types.Type from, Object value, org.apache.iceberg.types.Type to) { + PartitionData source = new PartitionData(Types.StructType.of(Types.NestedField.optional(1000, "p", from))); + source.set(0, value); + return new IcebergPartitionValueAdapter(source.getPartitionType(), + Types.StructType.of(Types.NestedField.optional(1000, "p", to))).adapt(source).get(0, Object.class); + } +} 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 bb46a63b3552b0..5e3e82bad787df 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 @@ -1560,7 +1560,7 @@ public void planScanPopulatesPerFilePartitionAndFormatCarriers() { } @Test - public void getScanNodePropertiesEmitsPathPartitionKeysAndRealDataFormat() { + public void getScanNodePropertiesEmitsPathPartitionKeysAndNativeDefault() { Schema schema = new Schema( Types.NestedField.required(1, "id", Types.IntegerType.get()), Types.NestedField.required(2, "P", Types.IntegerType.get()), @@ -1576,17 +1576,12 @@ public void getScanNodePropertiesEmitsPathPartitionKeysAndRealDataFormat() { // MUTATION: omitting path_partition_keys -> BE double-fills partition columns -> DCHECK -> this red. // MUTATION: re-lowercasing "P" -> "p,region" != "P,region" -> red. Assertions.assertEquals("P,region", props.get("path_partition_keys")); - // A base table MUST NOT report jni here: BE reads the scan-level format to pick FileScannerV2 vs V1 - // (_should_use_file_scanner_v2) before it can see any split, so "jni" would pin every iceberg scan to - // the V1 tree -- the only one carrying TableSchemaChangeHelper, whose StructNode DCHECK then SIGABRTs - // the BE on a Top-N lazy-materialization rowid slot (TeamCity 995122). No write.format.default and no - // snapshot on this fixture -> IcebergUtils.getFileFormat parity default = parquet. - // MUTATION: reverting to a hardcoded "jni" -> red here, and V2 is silently lost in production. + // A native default avoids a format-discovery scan. The physical format belongs to each range. Assertions.assertEquals("parquet", props.get("file_format_type")); } @Test - public void getScanNodePropertiesResolvesOrcTableFormat() { + public void getScanNodePropertiesDoesNotUseWriteFormatForRead() { Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); Table table = createTable("orct", schema, PartitionSpec.unpartitioned(), Collections.singletonMap(TableProperties.DEFAULT_FILE_FORMAT, "orc")); @@ -1595,10 +1590,10 @@ public void getScanNodePropertiesResolvesOrcTableFormat() { Map props = provider.getScanNodeProperties( null, new IcebergTableHandle("db1", "orct"), Collections.emptyList(), Optional.empty()); - // Pins that the scan-level format is RESOLVED from the table, not a constant: an orc table must say orc - // so BE's genSlotToSchemaIdMapForOrc / V2 selection see the truth (legacy getFileFormatType parity). - // MUTATION: hardcoding "parquet" instead of IcebergWriterHelper.getFileFormat(table) -> red. - Assertions.assertEquals("orc", props.get("file_format_type")); + Assertions.assertEquals("parquet", props.get("file_format_type")); + // An ORC write default does not describe old retained files; no data file is needed here. + Assertions.assertTrue(provider.planScan(null, ConnectorScanRequest.builder( + new IcebergTableHandle("db1", "orct"), Collections.emptyList()).build()).isEmpty()); } @Test @@ -4982,6 +4977,149 @@ private static List sortedPaths(List ranges) { return paths; } + @Test + public void scanPropertiesDoNotReadManifestsAndSplitsRetainOrc() throws IOException { + Table table = createTable("format_cache", SCHEMA, PartitionSpec.unpartitioned()); + for (int i = 0; i < 8; i++) { + table.newFastAppend().appendFile(dataFile(table.spec(), + "s3://b/db/format_cache/" + i + ".orc", 1024, null, null, FileFormat.ORC)).commit(); + } + RecordingConnectorContext context = new RecordingConnectorContext(); + IcebergManifestCache cache = new IcebergManifestCache(); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + IcebergCatalogProperties.of(manifestCacheProps()), opsReturning(table), context, cache); + Map props = provider.getScanNodeProperties(emptySession(), + new IcebergTableHandle("db1", "format_cache"), Collections.emptyList(), Optional.empty()); + Assertions.assertEquals("parquet", props.get("file_format_type")); + Assertions.assertEquals(0, cache.size(), "scan properties must not load data manifests"); + Assertions.assertEquals(0, context.metadataSubmissions.get(), "scan properties must not prefetch"); + try (ConnectorSplitSource source = provider.streamSplits(emptySession(), + scanOf(new IcebergTableHandle("db1", "format_cache")).build())) { + Assertions.assertTrue(source.hasNext()); + Assertions.assertEquals(TFileFormatType.FORMAT_ORC, populate(source.next()).getFormatType()); + Assertions.assertEquals(1, cache.size()); + source.onMoreSplitsRequested(); + Assertions.assertEquals(7, drain(source).size()); + } + Assertions.assertTrue(context.metadataSubmissions.get() > 0, + "full scans must retain bounded manifest prefetch after backend demand"); + Assertions.assertEquals(8, cache.size()); + } + + @Test + public void scanPropertiesLeaveManifestReadErrorsForSplitGeneration() { + Table table = createTable("format_failure", SCHEMA, PartitionSpec.unpartitioned()); + table.newFastAppend().appendFile(dataFile(table.spec(), + "s3://b/db/format_failure/one.orc", 1024, null, null, FileFormat.ORC)).commit(); + ManifestFile first = table.currentSnapshot().dataManifests(table.io()).get(0); + FailOnceFileIO io = new FailOnceFileIO(table.io(), first.path()); + Table wrapped = tableWithIo(table, io); + IcebergScanPlanProvider provider = manifestProvider(manifestCacheProps(), wrapped, + new IcebergManifestCache()); + Map props = provider.getScanNodeProperties(emptySession(), + new IcebergTableHandle("db1", "format_failure"), Collections.emptyList(), Optional.empty()); + Assertions.assertEquals("parquet", props.get("file_format_type")); + Assertions.assertFalse(io.failed.get(), "format initialization must not open the data manifest"); + Assertions.assertThrows(RuntimeException.class, () -> { + try (ConnectorSplitSource source = provider.streamSplits(emptySession(), + scanOf(new IcebergTableHandle("db1", "format_failure")).build())) { + source.hasNext(); + } + }); + Assertions.assertTrue(io.failed.get(), "the actual split scan must surface the manifest failure"); + } + + @Test + public void nativeDefaultDoesNotOverrideMixedFileFormats() { + Table table = createTable("mixed_formats", SCHEMA, PartitionSpec.unpartitioned(), + Collections.singletonMap(TableProperties.DEFAULT_FILE_FORMAT, "orc")); + table.newFastAppend() + .appendFile(dataFile(table.spec(), "s3://b/db/mixed/a.parquet", 1024, null, null)) + .appendFile(dataFile(table.spec(), "s3://b/db/mixed/b.orc", 1024, null, null, FileFormat.ORC)) + .commit(); + for (Map properties : Arrays.asList(manifestCacheProps(), Collections.emptyMap())) { + IcebergScanPlanProvider provider = manifestProvider(properties, table, new IcebergManifestCache()); + Map props = provider.getScanNodeProperties(emptySession(), + new IcebergTableHandle("db1", "mixed_formats"), Collections.emptyList(), Optional.empty()); + Assertions.assertEquals("parquet", props.get("file_format_type")); + List ranges = provider.planScan(emptySession(), ConnectorScanRequest.builder( + new IcebergTableHandle("db1", "mixed_formats"), Collections.emptyList()).build()); + Assertions.assertEquals(TFileFormatType.FORMAT_PARQUET, + populate(byPath(ranges, "a.parquet")).getFormatType()); + Assertions.assertEquals(TFileFormatType.FORMAT_ORC, + populate(byPath(ranges, "b.orc")).getFormatType()); + } + } + + @Test + public void streamingManifestPrefetchWaitsForFurtherBackendDemand() throws IOException { + Table table = createTable("demand", SCHEMA, PartitionSpec.unpartitioned()); + for (int i = 0; i < 8; i++) { + table.newFastAppend().appendFile(dataFile(table.spec(), + "s3://b/db/demand/" + i + ".parquet", 1024, null, null)).commit(); + } + Assertions.assertEquals(8, table.currentSnapshot().dataManifests(table.io()).size()); + RecordingConnectorContext context = new RecordingConnectorContext(); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + IcebergCatalogProperties.of(manifestCacheProps()), opsReturning(table), context, + new IcebergManifestCache()); + try (ConnectorSplitSource source = provider.streamSplits(emptySession(), + scanOf(new IcebergTableHandle("db1", "demand")).build())) { + List paths = new ArrayList<>(); + paths.add(source.next().getPath().get()); + paths.add(source.next().getPath().get()); + Assertions.assertEquals(0, context.metadataSubmissions.get()); + source.onMoreSplitsRequested(); + while (source.hasNext()) { + paths.add(source.next().getPath().get()); + } + Assertions.assertEquals(8, new java.util.HashSet<>(paths).size()); + Assertions.assertTrue(context.metadataSubmissions.get() > 0); + } + } + + @Test + public void rewriteScopeFilteringKeepsPrefetchUntilAnOutputSplitExists() throws IOException { + Table table = createTable("rewrite_demand", SCHEMA, PartitionSpec.unpartitioned()); + for (int i = 0; i < 8; i++) { + table.newFastAppend().appendFile(dataFile(table.spec(), + "s3://b/db/rewrite/" + i + ".parquet", 1024, null, null)).commit(); + } + RecordingConnectorContext context = new RecordingConnectorContext(); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + IcebergCatalogProperties.of(manifestCacheProps()), opsReturning(table), context, + new IcebergManifestCache()); + IcebergTableHandle handle = new IcebergTableHandle("db1", "rewrite_demand") + .withRewriteFileScope(Collections.singleton("s3://b/db/rewrite/not-present.parquet")); + try (ConnectorSplitSource source = provider.streamSplits(emptySession(), scanOf(handle).build())) { + Assertions.assertFalse(source.hasNext()); + Assertions.assertTrue(context.metadataSubmissions.get() > 0); + } + } + + @Test + public void fullyFilteredStreamingScanPrefetchesWithoutBackendDemand() throws IOException { + Table table = createTable("filtered_demand", SCHEMA, PartitionSpec.unpartitioned()); + Map bounds = Collections.singletonMap(1, + Conversions.toByteBuffer(Types.IntegerType.get(), 10)); + for (int i = 0; i < 8; i++) { + DataFile file = DataFiles.builder(table.spec()).withPath("s3://b/db/filtered/" + i + ".parquet") + .withFormat(FileFormat.PARQUET).withFileSizeInBytes(1024) + .withMetrics(new Metrics(1L, null, null, null, null, bounds, bounds)).build(); + table.newFastAppend().appendFile(file).commit(); + } + RecordingConnectorContext context = new RecordingConnectorContext(); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + IcebergCatalogProperties.of(manifestCacheProps()), opsReturning(table), context, + new IcebergManifestCache()); + try (ConnectorSplitSource source = provider.streamSplits(emptySession(), + scanOf(new IcebergTableHandle("db1", "filtered_demand")) + .filter(Optional.of(eqInt("id", -1))).build())) { + Assertions.assertFalse(source.hasNext()); + Assertions.assertTrue(context.metadataSubmissions.get() > 0); + } + } + @Test public void planScanManifestCacheEnabledMatchesSdkPathAndConsumesCache() { Table table = createTable("t1", SCHEMA, PartitionSpec.unpartitioned()); @@ -5883,10 +6021,8 @@ public void planScanCombinesPartitionPruneDeleteAndPathNormalizeOnOneRange() { // --- T09 helpers --- private static FakeIcebergTable fakeTable(String name) { - // write.format.default: getScanNodeProperties now resolves the scan-level format off the table (it - // drives BE's FileScannerV2-vs-V1 pick -- see getScanNodePropertiesEmitsPathPartitionKeysAndRealDataFormat). - // A declared format keeps these credential tests off IcebergUtils' last-resort planFiles() inference, - // which FakeIcebergTable deliberately refuses; a real table normally declares it too. + // Keep a normal write default on this credential fixture. Read-side scan properties use a native + // default and must not call FakeIcebergTable.newScan() to infer the physical file format. Map properties = Collections.singletonMap(TableProperties.DEFAULT_FILE_FORMAT, "parquet"); FakeIcebergTable table = new FakeIcebergTable( diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScheduledManifestTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScheduledManifestTest.java new file mode 100644 index 00000000000000..6638c9df306506 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScheduledManifestTest.java @@ -0,0 +1,406 @@ +// 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.iceberg; + +import org.apache.doris.connector.spi.ConnectorTaskScheduler; + +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.io.CloseableIterable; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.lang.reflect.Proxy; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; +import java.util.function.Function; + +@Timeout(30) +class IcebergScheduledManifestTest { + private static final ManifestCacheValue EMPTY = ManifestCacheValue.forDataFiles(Collections.emptyList()); + + private static final class Scheduler implements ConnectorTaskScheduler { + final Queue tasks = new ArrayDeque<>(); + final Set listeners = new LinkedHashSet<>(); + final CountDownLatch full = new CountDownLatch(1); + final CountDownLatch admitted = new CountDownLatch(1); + volatile boolean closed; + + @Override + public boolean isClosed() { + return closed; + } + + @Override + public synchronized boolean tryExecute(Runnable task, Consumer onFailure, Runnable onSpace) { + if (closed) { + throw new RejectedExecutionException("query closed"); + } + if (!tasks.isEmpty()) { + listeners.add(onSpace); + full.countDown(); + return false; + } + tasks.add(task); + admitted.countDown(); + return true; + } + + @Override + public synchronized void removeSpaceListener(Runnable onSpace) { + listeners.remove(onSpace); + } + + @Override + public void execute(Runnable task) { + Assertions.fail("must use bounded admission"); + } + + void runOne() { + Runnable task; + List notifications; + synchronized (this) { + task = tasks.remove(); + notifications = new ArrayList<>(listeners); + listeners.clear(); + } + notifications.forEach(Runnable::run); + task.run(); + } + } + + private IcebergManifestPrefetcher source(Scheduler scheduler, AtomicInteger loads) { + List manifests = new ArrayList<>(); + for (int i = 0; i < 20; i++) { + manifests.add(sourceManifest("m-" + i)); + } + return new IcebergManifestPrefetcher(CloseableIterable.withNoopClose(manifests).iterator(), 3, + manifest -> null, manifest -> { + loads.incrementAndGet(); + return EMPTY; + }, + null, new Semaphore(3), () -> true, scheduler); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void queryCancellationStopsAReadWaitingForItsSharedPermit(boolean releasePermit) throws Exception { + Scheduler scheduler = new Scheduler(); + AtomicInteger loads = new AtomicInteger(); + CountDownLatch waiting = new CountDownLatch(1); + Semaphore permits = new Semaphore(0) { + @Override + public void acquire() throws InterruptedException { + waiting.countDown(); + super.acquire(); + } + + @Override + public boolean tryAcquire(long timeout, TimeUnit unit) throws InterruptedException { + waiting.countDown(); + return super.tryAcquire(timeout, unit); + } + }; + ExecutorService parent = Executors.newSingleThreadExecutor(); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher( + CloseableIterable.withNoopClose(Collections.singletonList( + sourceManifest("waiting"))).iterator(), 1, manifest -> null, + manifest -> { + loads.incrementAndGet(); + return EMPTY; + }, null, permits, () -> true, scheduler)) { + Future read = parent.submit(source::next); + Assertions.assertTrue(waiting.await(5, TimeUnit.SECONDS)); + scheduler.closed = true; + if (releasePermit) { + permits.release(); + } + ExecutionException failure = Assertions.assertThrows( + ExecutionException.class, () -> read.get(2, TimeUnit.SECONDS)); + Assertions.assertInstanceOf(CancellationException.class, failure.getCause()); + Assertions.assertEquals(0, loads.get(), "cancellation must not start another remote read"); + Assertions.assertEquals(releasePermit ? 1 : 0, permits.availablePermits(), + "a waiter must not release another scan's permit"); + } finally { + permits.release(); + parent.shutdownNow(); + Assertions.assertTrue(parent.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static ManifestFile sourceManifest(String path) { + return (ManifestFile) Proxy.newProxyInstance(ManifestFile.class.getClassLoader(), + new Class[] {ManifestFile.class}, (proxy, method, args) -> { + if (method.getName().equals("path")) { + return path; + } + throw new UnsupportedOperationException(method.getName()); + }); + } + + @Test + void fullMetadataQueueDoesNotBlockSubmissionAndResumesWithinTheWindow() throws Exception { + Scheduler scheduler = new Scheduler(); + AtomicInteger loads = new AtomicInteger(); + ExecutorService parent = Executors.newSingleThreadExecutor(); + try (IcebergManifestPrefetcher source = source(scheduler, loads)) { + Assertions.assertEquals("m-0", source.next().manifest.path()); + Future next = parent.submit(source::next); + Assertions.assertTrue(scheduler.full.await(10, TimeUnit.SECONDS)); + Assertions.assertEquals(1, loads.get(), "a waiting parent must not bypass query scheduling"); + scheduler.runOne(); + Assertions.assertEquals("m-1", next.get(10, TimeUnit.SECONDS).manifest.path()); + Assertions.assertEquals(2, loads.get()); + source.close(); + Assertions.assertTrue(scheduler.listeners.isEmpty()); + while (!scheduler.tasks.isEmpty()) { + scheduler.runOne(); + } + Assertions.assertEquals(2, loads.get(), "closed window must not issue another IO"); + } finally { + parent.shutdownNow(); + } + } + + @Test + void rejectionWakesParentAndPreservesTheCause() throws Exception { + Scheduler scheduler = new Scheduler(); + AtomicInteger loads = new AtomicInteger(); + try (IcebergManifestPrefetcher source = source(scheduler, loads)) { + source.next(); + scheduler.closed = true; + RuntimeException failure = Assertions.assertThrows(RuntimeException.class, source::next); + Assertions.assertInstanceOf(RejectedExecutionException.class, failure.getCause()); + Assertions.assertEquals(1, loads.get()); + Assertions.assertTrue(scheduler.listeners.isEmpty()); + } + } + + @Test + void cancelWakesAParentEvenWhenTheMetadataWorkerIsUnavailable() throws Exception { + Scheduler scheduler = new Scheduler(); + AtomicInteger loads = new AtomicInteger(); + ExecutorService parent = Executors.newSingleThreadExecutor(); + try (IcebergManifestPrefetcher source = source(scheduler, loads)) { + source.next(); + Future next = parent.submit(source::next); + Assertions.assertTrue(scheduler.full.await(10, TimeUnit.SECONDS)); + source.close(); + Assertions.assertThrows(java.util.concurrent.ExecutionException.class, + () -> next.get(10, TimeUnit.SECONDS)); + Assertions.assertEquals(1, loads.get()); + Assertions.assertTrue(scheduler.listeners.isEmpty()); + } finally { + parent.shutdownNow(); + } + } + + private IcebergManifestPrefetcher sharedSource(Scheduler scheduler, int count, + Function> loader) { + List manifests = new ArrayList<>(); + for (int i = 0; i < count; i++) { + manifests.add(sourceManifest("m-" + i)); + } + return new IcebergManifestPrefetcher(CloseableIterable.withNoopClose(manifests).iterator(), 3, + manifest -> null, manifest -> EMPTY, null, new Semaphore(3), () -> true, scheduler, loader); + } + + @Test + void cancelledFutureRetainsLoadOwnerUntilCallableActuallyExits() throws Exception { + Scheduler scheduler = new Scheduler(); + ExecutorService parent = Executors.newSingleThreadExecutor(); + ExecutorService worker = Executors.newSingleThreadExecutor(); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch finishRead = new CountDownLatch(1); + AtomicInteger retained = new AtomicInteger(); + try (IcebergManifestPrefetcher source = new IcebergManifestPrefetcher( + CloseableIterable.withNoopClose(List.of(sourceManifest("first"), sourceManifest("second"))).iterator(), + 2, manifest -> null, manifest -> EMPTY, null, new Semaphore(2), () -> true, scheduler, + manifest -> { + reading.countDown(); + try { + Assertions.assertTrue(finishRead.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + throw new AssertionError(e); + } + return CompletableFuture.completedFuture(EMPTY); + }, () -> { + retained.incrementAndGet(); + return retained::decrementAndGet; + })) { + source.next(); + Future next = parent.submit(source::next); + Assertions.assertTrue(scheduler.admitted.await(5, TimeUnit.SECONDS)); + Future running = worker.submit(scheduler::runOne); + Assertions.assertTrue(reading.await(5, TimeUnit.SECONDS)); + source.close(); + Assertions.assertEquals(1, retained.get(), "cancel(false) must not close a running load's FileIO"); + Assertions.assertThrows(ExecutionException.class, () -> next.get(5, TimeUnit.SECONDS)); + finishRead.countDown(); + running.get(5, TimeUnit.SECONDS); + Assertions.assertEquals(0, retained.get()); + } finally { + finishRead.countDown(); + parent.shutdownNow(); + worker.shutdownNow(); + } + } + + @Test + void sharedFutureReturnsWorkerAndCompletesInReadinessOrder() throws Exception { + Scheduler scheduler = new Scheduler(); + ExecutorService parent = Executors.newSingleThreadExecutor(); + CompletableFuture slow = new CompletableFuture<>(); + CompletableFuture fast = new CompletableFuture<>(); + AtomicInteger registered = new AtomicInteger(); + try (IcebergManifestPrefetcher source = sharedSource(scheduler, 3, manifest -> { + registered.incrementAndGet(); + return manifest.path().equals("m-1") ? slow : fast; + })) { + Assertions.assertEquals("m-0", source.next().manifest.path()); + Future next = parent.submit(() -> source.next().manifest.path()); + Assertions.assertTrue(scheduler.full.await(5, TimeUnit.SECONDS)); + scheduler.runOne(); + Assertions.assertEquals(1, registered.get(), "worker must return before a shared load completes"); + scheduler.runOne(); + Assertions.assertEquals(2, registered.get(), "other scheduled work must be able to run"); + Assertions.assertFalse(next.isDone()); + fast.complete(EMPTY); + Assertions.assertEquals("m-2", next.get(5, TimeUnit.SECONDS)); + slow.complete(EMPTY); + Assertions.assertEquals("m-1", source.next().manifest.path()); + Assertions.assertFalse(source.hasNext()); + } finally { + slow.complete(EMPTY); + fast.complete(EMPTY); + parent.shutdownNow(); + } + } + + @Test + void closeDetachesPendingObserverWithoutCancellingSharedRead() throws Exception { + Scheduler scheduler = new Scheduler(); + ExecutorService parent = Executors.newSingleThreadExecutor(); + CompletableFuture shared = new CompletableFuture<>(); + AtomicInteger registered = new AtomicInteger(); + try (IcebergManifestPrefetcher source = sharedSource(scheduler, 10, manifest -> { + registered.incrementAndGet(); + return shared; + })) { + source.next(); + Future next = parent.submit(source::next); + Assertions.assertTrue(scheduler.full.await(5, TimeUnit.SECONDS)); + scheduler.runOne(); + Assertions.assertEquals(1, registered.get()); + source.close(); + Assertions.assertThrows(ExecutionException.class, () -> next.get(5, TimeUnit.SECONDS)); + Assertions.assertFalse(shared.isDone()); + Assertions.assertTrue(scheduler.listeners.isEmpty()); + while (!scheduler.tasks.isEmpty()) { + scheduler.runOne(); + } + Assertions.assertEquals(1, registered.get(), "early stop must prevent queued remote requests"); + shared.complete(EMPTY); + Assertions.assertFalse(source.hasNext()); + } finally { + shared.complete(EMPTY); + parent.shutdownNow(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void sharedFailurePreservesCauseBeforeOrAfterObserverRegistration(boolean alreadyCompleted) throws Exception { + Scheduler scheduler = new Scheduler(); + ExecutorService parent = Executors.newSingleThreadExecutor(); + IllegalStateException expected = new IllegalStateException("shared manifest failed"); + CompletableFuture shared = new CompletableFuture<>(); + if (alreadyCompleted) { + shared.completeExceptionally(expected); + } + try (IcebergManifestPrefetcher source = sharedSource(scheduler, 10, manifest -> shared)) { + source.next(); + Future next = parent.submit(source::next); + Assertions.assertTrue(scheduler.full.await(5, TimeUnit.SECONDS)); + scheduler.runOne(); + shared.completeExceptionally(expected); + ExecutionException failure = Assertions.assertThrows(ExecutionException.class, + () -> next.get(5, TimeUnit.SECONDS)); + Assertions.assertSame(expected, failure.getCause().getCause()); + Assertions.assertFalse(source.hasNext()); + Assertions.assertTrue(scheduler.listeners.isEmpty()); + while (!scheduler.tasks.isEmpty()) { + scheduler.runOne(); + } + } finally { + shared.completeExceptionally(expected); + parent.shutdownNow(); + } + } + + @Test + void completionRegistrationAndCloseCannotLoseWakeup() throws Exception { + ExecutorService parent = Executors.newSingleThreadExecutor(); + try { + for (int i = 0; i < 100; i++) { + Scheduler scheduler = new Scheduler(); + CompletableFuture shared = new CompletableFuture<>(); + try (IcebergManifestPrefetcher source = sharedSource(scheduler, 2, manifest -> shared)) { + source.next(); + Future next = parent.submit(source::next); + Assertions.assertTrue(scheduler.admitted.await(5, TimeUnit.SECONDS)); + if ((i & 1) == 0) { + shared.complete(EMPTY); + } + scheduler.runOne(); + source.close(); + shared.complete(EMPTY); + try { + next.get(5, TimeUnit.SECONDS); + } catch (ExecutionException failure) { + Assertions.assertInstanceOf(CancellationException.class, failure.getCause()); + } + Assertions.assertTrue(scheduler.listeners.isEmpty()); + } + } + } finally { + parent.shutdownNow(); + } + } + +} diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergStatementScopeTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergStatementScopeTest.java index 2df7d37c00b982..bada56e0796db5 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergStatementScopeTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergStatementScopeTest.java @@ -60,6 +60,19 @@ private static Table mutableTable(String name) { return catalog.createTable(TableIdentifier.of("db1", name), SCHEMA); } + @Test + public void manifestReadBudgetIsSharedAcrossCatalogsButNotExecutions() { + TestStatementScope scope = new TestStatementScope(); + java.util.concurrent.Semaphore first = IcebergStatementScope.manifestReadBudget( + new ScopeSession(1L, "q1", scope), 8); + Assertions.assertSame(first, IcebergStatementScope.manifestReadBudget( + new ScopeSession(2L, "q1", scope), 8)); + Assertions.assertNotSame(first, IcebergStatementScope.manifestReadBudget( + new ScopeSession(1L, "q2", scope), 8)); + Assertions.assertNotSame(IcebergStatementScope.manifestReadBudget(null, 8), + IcebergStatementScope.manifestReadBudget(null, 8)); + } + @Test public void sameStatementSharesOneLoadedTable() { // Read metadata and scan planning resolve one frozen view through sharedTable; sharing one scope @@ -251,6 +264,28 @@ public void directTrackedTablePinsCatalogUntilStatementClose() { Assertions.assertEquals(1, catalogCloses.get()); } + @Test + public void runningReadLeaseDefersStatementOwnedFileIoCleanup() throws Exception { + IcebergCatalogResourceTracker tracker = new IcebergCatalogResourceTracker(); + TestStatementScope scope = new TestStatementScope(); + ScopeSession session = new ScopeSession(7L, "q1", scope); + AtomicInteger tableCloses = new AtomicInteger(); + AtomicInteger catalogCloses = new AtomicInteger(); + IcebergStatementScope.sharedTrackedTable(session, "db1", "t", tracker, () -> table("t"), + ignored -> tableCloses::incrementAndGet); + Supplier retain = IcebergStatementScope.readTableLeaseFactory(session, "db1", "t"); + AutoCloseable release = retain.get(); + scope.closeAll(); + tracker.close(catalogCloses::incrementAndGet); + Assertions.assertEquals(0, tableCloses.get()); + Assertions.assertEquals(0, catalogCloses.get()); + release.close(); + release.close(); + Assertions.assertEquals(1, tableCloses.get()); + Assertions.assertEquals(1, catalogCloses.get()); + Assertions.assertThrows(IllegalStateException.class, retain::get); + } + @Test public void unscopedMetadataOperationReleasesExactGenerationBeforeReturning() { IcebergCatalogResourceTracker tracker = new IcebergCatalogResourceTracker(); diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWriterHelperTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWriterHelperTest.java index c07f04896f54f3..9cf897e24c3464 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWriterHelperTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWriterHelperTest.java @@ -644,83 +644,12 @@ public void getFileFormatThrowsOnUnsupported() { () -> IcebergWriterHelper.getFileFormat(tableWith("write.format.default", "avro"))); } - // ─────────────────── getFileFormat: PERF-03 cross-query inference cache ─────────────────── - - /** Unpartitioned table carrying one data file of {@code fileFormat} and the given extra properties. */ - private Table tableWithDataFile(FileFormat fileFormat, String... props) { - Table table = tableWith(props); - table.newAppend().appendFile(DataFiles.builder(unpartitionedSpec) - .withPath("s3://b/db1/t/f0." + fileFormat.name().toLowerCase()) - .withFileSizeInBytes(100) - .withRecordCount(1) - .withFormat(fileFormat) - .build()).commit(); - return table; - } - - @Test - public void getFileFormatInferenceIsCachedAcrossQueriesAtSameSnapshot() { - // A table with NO write-format / write.format.default (migrated / engine-default table) resolves the scan - // level file_format_type via an unfiltered whole-table planFiles() inference (the #64134 heavy op). PERF-03 - // memoizes that inference per (table, currentSnapshotId): repeated getScanNodeProperties computes across - // queries collapse onto ONE remote inference. MUTATION: not threading the cache into the call -> each query - // re-scans -> loadCountForTest > 1 -> red. - Table table = tableWithDataFile(FileFormat.ORC); - TableIdentifier id = TableIdentifier.of("db1", "t"); - IcebergFormatCache cache = new IcebergFormatCache(100, 1000); - - FileFormat f1 = IcebergWriterHelper.getFileFormat(table, id, cache); - FileFormat f2 = IcebergWriterHelper.getFileFormat(table, id, cache); - FileFormat f3 = IcebergWriterHelper.getFileFormat(table, id, cache); - - Assertions.assertEquals(FileFormat.ORC, f1); - Assertions.assertEquals(FileFormat.ORC, f2); - Assertions.assertEquals(FileFormat.ORC, f3); - Assertions.assertEquals(1, cache.loadCountForTest(), - "the whole-table inference must run exactly once across repeated queries at one snapshot"); - Assertions.assertEquals(1, cache.size(), "one (table, snapshot) entry"); - // Parity: the cached resolution matches the uncached (live) resolution. - Assertions.assertEquals(IcebergWriterHelper.getFileFormat(table), f1, - "cached format must equal a live (uncached) inference"); - } - - @Test - public void getFileFormatPropertyPathIsNeverCached() { - // When the table carries a format property, resolution is a cheap property read that must NOT touch the - // cache (only the inference fallback is memoized). MUTATION: caching the property path -> size > 0 -> red. - Table table = tableWithDataFile(FileFormat.ORC, "write.format.default", "parquet"); - TableIdentifier id = TableIdentifier.of("db1", "t"); - IcebergFormatCache cache = new IcebergFormatCache(100, 1000); - - Assertions.assertEquals(FileFormat.PARQUET, IcebergWriterHelper.getFileFormat(table, id, cache)); - Assertions.assertEquals(0, cache.size(), "the property path must never populate the inference cache"); - Assertions.assertEquals(0, cache.loadCountForTest()); - } - - @Test - public void getFileFormatNullCacheResolvesLive() { - // Offline / pre-cache paths pass a null cache: resolution stays live and correct (no NPE). - Table table = tableWithDataFile(FileFormat.ORC); - Assertions.assertEquals(FileFormat.ORC, - IcebergWriterHelper.getFileFormat(table, TableIdentifier.of("db1", "t"), null)); - } - @Test - public void getFileFormatUnsupportedFirstFileCachesInferenceButRethrowsEachCall() { - // R3: an unsupported first data-file format (e.g. avro) is a SUCCESSFUL inference (the loader returns - // "avro"); the unsupported-format throw happens in the format mapping OUTSIDE getOrLoad. So the inference is - // cached once (loadCount == 1) yet every call re-throws (parity with legacy, which inferred + threw every - // query) without re-scanning. MUTATION: caching the mapped value / caching the throw -> loadCount != 1. - Table table = tableWithDataFile(FileFormat.AVRO); - TableIdentifier id = TableIdentifier.of("db1", "t"); - IcebergFormatCache cache = new IcebergFormatCache(100, 1000); - - Assertions.assertThrows(RuntimeException.class, - () -> IcebergWriterHelper.getFileFormat(table, id, cache)); - Assertions.assertThrows(RuntimeException.class, - () -> IcebergWriterHelper.getFileFormat(table, id, cache)); - Assertions.assertEquals(1, cache.loadCountForTest(), - "the inference runs once; the unsupported-format throw is in the mapping, not the loader"); - Assertions.assertEquals(1, cache.size()); + public void writeFormatStillInfersOrcFromExistingFiles() { + Table table = tableWith(); + table.newAppend().appendFile(DataFiles.builder(unpartitionedSpec) + .withPath("s3://b/db1/t/f0.orc").withFileSizeInBytes(100) + .withRecordCount(1).withFormat(FileFormat.ORC).build()).commit(); + Assertions.assertEquals(FileFormat.ORC, IcebergWriterHelper.getFileFormat(table)); } } diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/RecordingConnectorContext.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/RecordingConnectorContext.java index a2d2240af723c7..d7f6e74fb5b378 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/RecordingConnectorContext.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/RecordingConnectorContext.java @@ -22,6 +22,7 @@ import org.apache.doris.connector.spi.ConnectorContext; import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.ConnectorStorageContext; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.filesystem.FileSystem; import org.apache.doris.filesystem.properties.StorageProperties; import org.apache.doris.thrift.TFileType; @@ -31,6 +32,9 @@ import java.util.List; import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.UnaryOperator; /** @@ -54,6 +58,37 @@ public ConnectorStorageContext getStorageContext() { return this; } + final AtomicInteger metadataSubmissions = new AtomicInteger(); + + @Override + public RunnableFuture submitMetadataTask(Callable task) { + metadataSubmissions.incrementAndGet(); + return new FutureTask<>(task); + } + + @Override + public ConnectorTaskScheduler getMetadataTaskScheduler(ConnectorSession session) { + ConnectorTaskScheduler delegate = ConnectorContext.super.getMetadataTaskScheduler(session); + return new ConnectorTaskScheduler() { + @Override + public boolean tryExecute(Runnable task, java.util.function.Consumer onFailure, + Runnable onSpace) { + metadataSubmissions.incrementAndGet(); + return delegate.tryExecute(task, onFailure, onSpace); + } + + @Override + public void removeSpaceListener(Runnable onSpace) { + delegate.removeSpaceListener(onSpace); + } + + @Override + public void execute(Runnable task) { + delegate.execute(task); + } + }; + } + int authCount; boolean failAuth; /** Runs when a task has finished but {@link #executeAuthenticated} has not left its scope yet, so a test can diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorContext.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorContext.java index e0a6245733d1b9..58b177de8e6760 100644 --- a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorContext.java +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorContext.java @@ -20,6 +20,8 @@ import java.util.Collections; import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RunnableFuture; /** * Runtime context provided by fe-core to connector implementations. @@ -32,6 +34,23 @@ */ public interface ConnectorContext { + /** + * Offers independent metadata work to the engine's shared scheduling pool without waiting for queue space. + * The returned task may remain unstarted when the pool is full. Its owner MUST call {@code run()} before + * {@code get()}: FutureTask arbitrates execution exactly once, so the caller can help its own queued work + * even when every pool worker is occupied by a parent task. Always use get() to observe completion/errors. + * Cancel unused tasks; completion/cancellation removes accepted work from the engine queue. + * The callable owns its plugin classloader and authentication context. The default defers to caller help. + */ + default RunnableFuture submitMetadataTask(Callable task) { + return new FutureTask<>(task); + } + + /** Returns the captured query-owned scheduler for independent metadata children. */ + default ConnectorTaskScheduler getMetadataTaskScheduler(ConnectorSession session) { + return session.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA); + } + /** Returns the catalog name. */ String getCatalogName(); diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorSession.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorSession.java index e305338266ed77..421164745b21bd 100644 --- a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorSession.java +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorSession.java @@ -27,6 +27,15 @@ */ public interface ConnectorSession { + /** + * Query-owned scheduling captured with this session, shared by all scans/catalogs of the statement. + * SPLIT tasks may wait for METADATA tasks, never the reverse. Closing the query rejects outstanding + * work and prevents this captured handle from creating a new queue. The default is for offline use. + */ + default ConnectorTaskScheduler getTaskScheduler(ConnectorTaskScheduler.Stage stage) { + return ConnectorTaskScheduler.DIRECT; + } + /** Returns the unique query identifier. */ String getQueryId(); diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorTaskScheduler.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorTaskScheduler.java new file mode 100644 index 00000000000000..bec2160dbc2647 --- /dev/null +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorTaskScheduler.java @@ -0,0 +1,75 @@ +// 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.spi; + +import java.util.concurrent.Executor; +import java.util.function.Consumer; + +/** + * A query's bounded task queue. The engine owns its lifetime; connectors must not close it. + * Admission never waits. On a full queue the caller retains its task and retries after onSpace runs. + * The callback may race the return from tryExecute, so callers must coalesce retries. + * Accepted tasks run once or receive onFailure; failures include cancellation before execution. + * Tasks own their authentication and context classloader. Callbacks must be short and nonblocking. + */ +public interface ConnectorTaskScheduler extends Executor { + enum Stage { + SPLIT, + METADATA + } + + boolean tryExecute(Runnable task, Consumer onFailure, Runnable onSpace); + + /** Query completion closes the group; synchronous permit waiters can observe the same lifetime. */ + default boolean isClosed() { + return false; + } + + /** + * Requeue the currently running task as the caller's final action before returning. The scheduler + * may retain the current worker instead of waking another one, but must preserve normal queue order, + * admission, and failure semantics. Calls outside that task fall back to ordinary admission. + */ + default boolean tryExecuteContinuation(Runnable task, Consumer onFailure, Runnable onSpace) { + return tryExecute(task, onFailure, onSpace); + } + + void removeSpaceListener(Runnable onSpace); + + /** Offline sessions have no shared worker pool, and execute directly. */ + ConnectorTaskScheduler DIRECT = new ConnectorTaskScheduler() { + @Override + public boolean tryExecute(Runnable task, Consumer onFailure, Runnable onSpace) { + try { + task.run(); + } catch (Throwable failure) { + onFailure.accept(failure); + } + return true; + } + + @Override + public void removeSpaceListener(Runnable onSpace) { + } + + @Override + public void execute(Runnable task) { + task.run(); + } + }; +} diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ForwardingConnectorContext.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ForwardingConnectorContext.java index d65d9f4c4313b3..fb53c9c7e839b0 100644 --- a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ForwardingConnectorContext.java +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ForwardingConnectorContext.java @@ -20,6 +20,7 @@ import java.util.Map; import java.util.Objects; import java.util.concurrent.Callable; +import java.util.concurrent.RunnableFuture; /** * Base class for a {@link ConnectorContext} decorator: forwards every method to the wrapped context. @@ -100,6 +101,16 @@ public String sanitizeOutboundUrl(String url) { return delegate.sanitizeOutboundUrl(url); } + @Override + public ConnectorTaskScheduler getMetadataTaskScheduler(ConnectorSession session) { + return delegate.getMetadataTaskScheduler(session); + } + + @Override + public RunnableFuture submitMetadataTask(Callable task) { + return delegate.submitMetadataTask(task); + } + @Override public T executeAuthenticated(Callable task) throws Exception { return delegate.executeAuthenticated(task); 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 1bf80d8d582690..0cb35f10a9d064 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 @@ -37,6 +37,15 @@ */ public interface ConnectorSplitSource extends Closeable { + /** + * Nonblocking, idempotent hint that a backend has requested a batch beyond its first one. + * May run concurrently on an engine RPC thread. A connector may use this to expand speculative + * prefetch, but MUST keep producing without the hint: planning and fully filtered scans may + * need to enumerate metadata before any backend can request a batch. + */ + default void onMoreSplitsRequested() { + } + /** * Returns whether more ranges remain. May advance over internally-skipped tasks (e.g. files * filtered out by a rewrite scope), so it is the only safe way to test for completion. 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 466c865a38ae39..f0e0415c7f714c 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 @@ -91,12 +91,13 @@ public void connectorApiMajorTracksTheRecordedSurfaceChange() throws IOException // Major 20 adds deleteUsingMetadata to the connector write surface; major 21 adds // typed filesystem credential descriptions and statement-scoped metadata. // Major 22 adds the metadata-only DDL handle path; major 23 adds the provider's - // ConnectorSecuritySurface declaration. - Assertions.assertEquals("23.0", version.getProperty("api.version")); + // ConnectorSecuritySurface declaration; major 24 adds the scan task scheduler. + Assertions.assertEquals("24.0", version.getProperty("api.version")); } /** Root entry points plus provider/handle types returned to connector plugins. */ private static final List> FROZEN_TYPES = Arrays.asList( + ConnectorTaskScheduler.class, DiagnosticException.class, ConnectorProvider.class, ConnectorSecuritySurface.class, @@ -121,11 +122,12 @@ public void connectorApiMajorTracksTheRecordedSurfaceChange() throws IOException ConnectorWritePlanProvider.class, org.apache.doris.extension.spi.Plugin.class, org.apache.doris.extension.spi.PluginFactory.class, - org.apache.doris.extension.spi.PluginContext.class); + org.apache.doris.extension.spi.PluginContext.class, + org.apache.doris.connector.spi.scan.ConnectorSplitSource.class); /** Public enum constants linked directly by connector plugin bytecode. */ private static final List>> FROZEN_ENUM_TYPES = - Arrays.asList(ConnectorCapability.class, + Arrays.asList(ConnectorCapability.class, ConnectorTaskScheduler.Stage.class, org.apache.doris.connector.spi.write.ConnectorRowChangeStyle.class, org.apache.doris.connector.spi.write.ConnectorWriteDistribution.Mode.class, org.apache.doris.connector.spi.write.ConnectorWriteDistribution.WriterAssignment.class); 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 ac81ad74418c86..3cdbe7dc9f8fc9 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 @@ -71,8 +71,10 @@ org.apache.doris.connector.spi.ConnectorContext#getClusterId():int org.apache.doris.connector.spi.ConnectorContext#getConnectorConfig():java.util.Map org.apache.doris.connector.spi.ConnectorContext#getEnvironment():java.util.Map org.apache.doris.connector.spi.ConnectorContext#getHttpSecurityHook():org.apache.doris.connector.spi.ConnectorHttpSecurityHook +org.apache.doris.connector.spi.ConnectorContext#getMetadataTaskScheduler(org.apache.doris.connector.spi.ConnectorSession):org.apache.doris.connector.spi.ConnectorTaskScheduler org.apache.doris.connector.spi.ConnectorContext#getStorageContext():org.apache.doris.connector.spi.ConnectorStorageContext org.apache.doris.connector.spi.ConnectorContext#sanitizeOutboundUrl(java.lang.String):java.lang.String +org.apache.doris.connector.spi.ConnectorContext#submitMetadataTask(java.util.concurrent.Callable):java.util.concurrent.RunnableFuture org.apache.doris.connector.spi.ConnectorProvider#acceptedCreateTableEngineNames():java.util.Set org.apache.doris.connector.spi.ConnectorProvider#create():org.apache.doris.extension.spi.Plugin org.apache.doris.connector.spi.ConnectorProvider#create(java.util.Map,org.apache.doris.connector.spi.ConnectorContext):org.apache.doris.connector.spi.Connector @@ -107,6 +109,7 @@ org.apache.doris.connector.spi.ConnectorSession#getQueryId():java.lang.String org.apache.doris.connector.spi.ConnectorSession#getSessionId():java.lang.String org.apache.doris.connector.spi.ConnectorSession#getSessionProperties():java.util.Map org.apache.doris.connector.spi.ConnectorSession#getStatementScope():org.apache.doris.connector.spi.ConnectorStatementScope +org.apache.doris.connector.spi.ConnectorSession#getTaskScheduler(org.apache.doris.connector.spi.ConnectorTaskScheduler$Stage):org.apache.doris.connector.spi.ConnectorTaskScheduler org.apache.doris.connector.spi.ConnectorSession#getTimeZone():java.lang.String org.apache.doris.connector.spi.ConnectorSession#getUser():java.lang.String org.apache.doris.connector.spi.ConnectorSession#isExternalScanTaskReuseEnabled():boolean @@ -128,6 +131,13 @@ org.apache.doris.connector.spi.ConnectorTableSchema#getTableName():java.lang.Str org.apache.doris.connector.spi.ConnectorTableSchema#getWriteMetadataIdentity():java.lang.String org.apache.doris.connector.spi.ConnectorTableSchema#hashCode():int org.apache.doris.connector.spi.ConnectorTableSchema#toString():java.lang.String +org.apache.doris.connector.spi.ConnectorTaskScheduler#execute(java.lang.Runnable):void +org.apache.doris.connector.spi.ConnectorTaskScheduler#isClosed():boolean +org.apache.doris.connector.spi.ConnectorTaskScheduler#removeSpaceListener(java.lang.Runnable):void +org.apache.doris.connector.spi.ConnectorTaskScheduler#tryExecute(java.lang.Runnable,java.util.function.Consumer,java.lang.Runnable):boolean +org.apache.doris.connector.spi.ConnectorTaskScheduler#tryExecuteContinuation(java.lang.Runnable,java.util.function.Consumer,java.lang.Runnable):boolean +org.apache.doris.connector.spi.ConnectorTaskScheduler$Stage#enum:METADATA +org.apache.doris.connector.spi.ConnectorTaskScheduler$Stage#enum:SPLIT org.apache.doris.connector.spi.DiagnosticException#getDiagnosticMessage():java.lang.String org.apache.doris.connector.spi.DiagnosticException#getDiagnosticStackTrace(java.lang.Throwable):java.lang.String org.apache.doris.connector.spi.handle.ConnectorColumnHandle#equals(java.lang.Object):boolean @@ -221,6 +231,10 @@ org.apache.doris.connector.spi.scan.ConnectorScanRequest#isCountPushdown():boole org.apache.doris.connector.spi.scan.ConnectorScanRequest#isExplainOnly():boolean org.apache.doris.connector.spi.scan.ConnectorScanRequest#isPartitionsPrunedToEmpty():boolean org.apache.doris.connector.spi.scan.ConnectorScanRequest#withRequiredPartitions(java.util.List):org.apache.doris.connector.spi.scan.ConnectorScanRequest +org.apache.doris.connector.spi.scan.ConnectorSplitSource#close():void +org.apache.doris.connector.spi.scan.ConnectorSplitSource#hasNext():boolean +org.apache.doris.connector.spi.scan.ConnectorSplitSource#next():org.apache.doris.connector.spi.scan.ConnectorScanRange +org.apache.doris.connector.spi.scan.ConnectorSplitSource#onMoreSplitsRequested():void org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:FILE_FORMAT_TYPE:java.lang.String=file_format_type org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:LOCATION_PREFIX:java.lang.String=location. org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:PATH_PARTITION_KEYS:java.lang.String=path_partition_keys diff --git a/fe/fe-connector/pom.xml b/fe/fe-connector/pom.xml index 18455d127a4cb9..ad60cbf0d6984c 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). --> - 23.0 + 24.0 diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/ThreadPoolManager.java b/fe/fe-core/src/main/java/org/apache/doris/common/ThreadPoolManager.java index 6cbdab0a39993e..d6ab78bca14d58 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/ThreadPoolManager.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/ThreadPoolManager.java @@ -18,6 +18,7 @@ package org.apache.doris.common; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.kerberos.ExecutionAuthenticator; import org.apache.doris.metric.Metric; import org.apache.doris.metric.Metric.MetricUnit; @@ -104,6 +105,15 @@ public static void registerThreadPoolMetric(String poolName, ThreadPoolExecutor } } + public static void registerFairTaskExecutorMetric(String name, FairTaskExecutor executor) { + registerMetric(name, "pool_size", Metric.MetricType.GAUGE, MetricUnit.NOUNIT, executor::getPoolSize); + registerMetric(name, "active_thread_num", Metric.MetricType.GAUGE, MetricUnit.NOUNIT, executor::getActiveCount); + registerMetric(name, "task_in_queue", Metric.MetricType.GAUGE, MetricUnit.NOUNIT, executor::getQueuedTaskCount); + registerMetric(name, "query_groups", Metric.MetricType.GAUGE, MetricUnit.NOUNIT, executor::getGroupCount); + registerMetric(name, "completed_task_count", Metric.MetricType.COUNTER, MetricUnit.NOUNIT, + executor::getCompletedTaskCount); + } + private static void registerMetric(String poolName, String metricName, Metric.MetricType type, MetricUnit unit, Supplier supplier) { Metric gauge = new Metric("thread_pool", type, unit, "thread_pool statistics") { @@ -383,6 +393,32 @@ public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) { } } + /** Submit without queue-space waits; the owner can run the returned future before awaiting it. */ + public static RunnableFuture submitWithCallerHelp(ThreadPoolExecutor executor, Callable callable) { + HelpableTask task = new HelpableTask<>(callable, executor); + try { + executor.execute(task); + } catch (RejectedExecutionException rejected) { + // Keep it unstarted for its owner, rather than blocking inside BlockedPolicy. + } + return task; + } + + private static final class HelpableTask extends FutureTask { + private final ThreadPoolExecutor executor; + + private HelpableTask(Callable callable, ThreadPoolExecutor executor) { + super(callable); + this.executor = executor; + } + + @Override + protected void done() { + // Caller help or cancellation can finish a future before a worker dequeues it. + executor.remove(this); + } + } + /** * A handler for rejected task that try to be blocked until the pool enqueue task succeed or timeout, * used for fixed thread pool @@ -402,6 +438,9 @@ public BlockedPolicy(String threadPoolName, int timeoutSeconds) { @Override public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) { + if (r instanceof HelpableTask) { + throw new RejectedExecutionException("Metadata task deferred to caller help: " + threadPoolName); + } try { boolean ret = executor.getQueue().offer(r, timeoutSeconds, TimeUnit.SECONDS); if (!ret) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/util/FairTaskExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/common/util/FairTaskExecutor.java new file mode 100644 index 00000000000000..c525db2b19a39f --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/common/util/FairTaskExecutor.java @@ -0,0 +1,448 @@ +// 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.common.util; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.io.Closeable; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Executor; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.LongAdder; +import java.util.concurrent.locks.LockSupport; +import java.util.function.Consumer; + +/** + * Fixed workers fairly dequeue task groups, then execute outside scheduler locks. A group can use + * all workers; fairness applies to task admission, not preemption of blocking calls. Queue capacity + * belongs to each group, so one caller cannot fill another caller's queue. Tasks must yield by + * returning and resubmitting, and must not wait for children submitted to this same executor. + */ +public final class FairTaskExecutor implements Closeable { + private static final Logger LOG = LogManager.getLogger(FairTaskExecutor.class); + private final ConcurrentLinkedQueue ready = new ConcurrentLinkedQueue<>(); + private final ConcurrentLinkedQueue waiters = new ConcurrentLinkedQueue<>(); + private final List workers = new ArrayList<>(); + private final ThreadLocal currentWorker = new ThreadLocal<>(); + private final Set groups = java.util.concurrent.ConcurrentHashMap.newKeySet(); + private final AtomicBoolean started = new AtomicBoolean(); + private volatile boolean closed; + private final LongAdder active = new LongAdder(); + private final LongAdder completed = new LongAdder(); + + public FairTaskExecutor(int parallelism, ThreadFactory threadFactory) { + if (parallelism < 1) { + throw new IllegalArgumentException("Parallelism must be positive"); + } + for (int i = 0; i < parallelism; i++) { + workers.add(threadFactory.newThread(this::work)); + } + } + + public synchronized Group newGroup(int capacity) { + if (closed) { + throw new RejectedExecutionException("Executor is closed"); + } + if (capacity < 1) { + throw new IllegalArgumentException("Group queue capacity must be positive"); + } + Group group = new Group(capacity); + groups.add(group); + return group; + } + + private synchronized void start() { + if (closed) { + throw new RejectedExecutionException("Executor is closed"); + } + if (started.compareAndSet(false, true)) { + workers.forEach(Thread::start); + } + } + + private static final class WorkerState { + private Group group; + private Runnable task; + private boolean continuationQueued; + } + + private static final class Entry { + private final Runnable task; + private final Consumer rejected; + + private Entry(Runnable task, Consumer rejected) { + this.task = task; + this.rejected = rejected; + } + } + + public final class Group implements Executor, Closeable { + private final ConcurrentLinkedQueue pending = new ConcurrentLinkedQueue<>(); + private final Set runningFutures = java.util.concurrent.ConcurrentHashMap.newKeySet(); + private final AtomicInteger queued = new AtomicInteger(); + private final AtomicBoolean published = new AtomicBoolean(); + // Enqueue/dequeue use only CLQ and CAS. This lock protects capacity listeners and close; + // neither the normal admission path nor user callbacks run under it. + private final Object admissionLock = new Object(); + private final Map> spaceListeners = new HashMap<>(); + private volatile boolean hasSpaceListeners; + private volatile boolean stopped; + private final int capacity; + + private Group(int capacity) { + this.capacity = capacity; + } + + @Override + public void execute(Runnable task) { + if (!tryExecute(task, failure -> { + if (task instanceof Future) { + ((Future) task).cancel(false); + } else { + LOG.warn("Task was rejected or failed", failure); + } + }, null)) { + throw new RejectedExecutionException("Task group queue is full (capacity=" + capacity + ")"); + } + } + + /** False registers one capacity callback. The caller retains its task and retries on notification. */ + public boolean tryExecute(Runnable task, Consumer rejected, Runnable onSpace) { + return tryExecute(task, rejected, onSpace, null); + } + + /** The caller must return after admission; the current worker will poll the global queue again. */ + public boolean tryExecuteContinuation(Runnable task, Consumer rejected, Runnable onSpace) { + WorkerState worker = currentWorker.get(); + return tryExecute(task, rejected, onSpace, + worker != null && worker.group == this && worker.task == task ? worker : null); + } + + private boolean tryExecute(Runnable task, Consumer rejected, Runnable onSpace, WorkerState worker) { + if (stopped || closed) { + throw new RejectedExecutionException("Task group is closed"); + } + if (!started.get()) { + start(); + } + while (true) { + int count = queued.get(); + if (count >= capacity) { + synchronized (admissionLock) { + if (stopped || closed) { + throw new RejectedExecutionException("Task group is closed"); + } + if (onSpace != null) { + spaceListeners.put(onSpace, rejected); + hasSpaceListeners = true; + } + } + // Dequeue can precede registration. Either this recheck or the dequeuer + // notifies; removing the callback under admissionLock prevents double wakes. + if (queued.get() < capacity || stopped || closed) { + notifySpace(); + } + return false; + } + if (queued.compareAndSet(count, count + 1)) { + Entry entry = new Entry(task, rejected); + pending.add(entry); + // close may have drained before this reserved slot was published. Removal + // transfers ownership to exactly one of close, the worker, or this submitter. + if (stopped || closed) { + if (pending.remove(entry)) { + queued.decrementAndGet(); + reject(entry); + } + } else { + publish(this, worker == null); + if (worker != null) { + worker.continuationQueued = true; + } + } + return true; + } + } + } + + public void removeSpaceListener(Runnable listener) { + synchronized (admissionLock) { + spaceListeners.remove(listener); + hasSpaceListeners = !spaceListeners.isEmpty(); + } + } + + private void notifySpace() { + if (!hasSpaceListeners) { + return; + } + Map> callbacks; + synchronized (admissionLock) { + callbacks = new HashMap<>(spaceListeners); + spaceListeners.clear(); + hasSpaceListeners = false; + } + callbacks.forEach((callback, failureHandler) -> { + try { + callback.run(); + } catch (Throwable e) { + try { + failureHandler.accept(e); + } catch (Throwable failure) { + LOG.error("Task group capacity failure callback failed", failure); + } + } + }); + } + + public int queuedTasks() { + return queued.get(); + } + + public boolean isClosed() { + return stopped || closed; + } + + @Override + public void close() { + synchronized (admissionLock) { + stopped = true; + } + Entry entry; + while ((entry = pending.poll()) != null) { + queued.decrementAndGet(); + reject(entry); + } + // Wake parents waiting on futures without interrupting an elected shared cache loader. + // Plain resumable tasks cancel cooperatively; their source must not close during iteration. + for (Entry running : runningFutures) { + try { + ((Future) running.task).cancel(false); + } catch (Throwable failure) { + LOG.error("Task cancellation callback failed", failure); + } + } + notifySpace(); + ready.remove(this); + groups.remove(this); + } + } + + private void publish(Group group) { + publish(group, true); + } + + private void publish(Group group, boolean wakeWorker) { + if (!group.stopped && !group.pending.isEmpty() && group.published.compareAndSet(false, true)) { + ready.add(group); + if (group.stopped || closed) { + ready.remove(group); + } else if (wakeWorker) { + signalOne(); + } + } + } + + private void reject(Entry entry) { + try { + entry.rejected.accept(new RejectedExecutionException("Task group is closed")); + } catch (Throwable e) { + LOG.warn("Task rejection callback failed", e); + } + } + + private void work() { + WorkerState worker = new WorkerState(); + currentWorker.set(worker); + try { + work(worker); + } finally { + currentWorker.remove(); + } + } + + private void work(WorkerState worker) { + while (!closed) { + Group group = awaitGroup(); + if (group == null) { + return; + } + if (worker.continuationQueued) { + worker.continuationQueued = false; + // We retained a worker for its continuation, but may have dequeued another query. + // Wake a peer for remaining ready work before this task can block on external IO. + if (!ready.isEmpty()) { + signalOne(); + } + } + Entry entry = group.pending.poll(); + Future future = entry != null && entry.task instanceof Future ? (Future) entry.task : null; + if (entry != null) { + if (future != null) { + group.runningFutures.add(entry); + if (future instanceof CompletionStage) { + // A cooperative task may return while observing a shared asynchronous load. + // Keep its cancellation ownership until completion, without retaining a worker. + ((CompletionStage) future).whenComplete((value, failure) -> + group.runningFutures.remove(entry)); + } + } + group.queued.decrementAndGet(); + } + group.published.set(false); + publish(group); + group.notifySpace(); + if (entry == null) { + continue; + } + if (group.stopped || closed) { + reject(entry); + if (future != null) { + group.runningFutures.remove(entry); + } + continue; + } + active.increment(); + ClassLoader previousClassLoader = Thread.currentThread().getContextClassLoader(); + worker.group = group; + worker.task = entry.task; + try { + entry.task.run(); + } catch (Throwable e) { + // A failed caller must not permanently remove an executor worker. + try { + entry.rejected.accept(e); + } catch (Throwable callbackFailure) { + LOG.error("Task group failure callback failed", callbackFailure); + } + } finally { + worker.group = null; + worker.task = null; + if (future != null && !(future instanceof CompletionStage)) { + group.runningFutures.remove(entry); + } + active.decrement(); + completed.increment(); + Thread.currentThread().setContextClassLoader(previousClassLoader); + Thread.interrupted(); + } + } + } + + public int getPoolSize() { + return started.get() ? workers.size() : 0; + } + + public long getActiveCount() { + return active.sum(); + } + + public long getCompletedTaskCount() { + return completed.sum(); + } + + public int getGroupCount() { + return groups.size(); + } + + public long getQueuedTaskCount() { + return groups.stream().mapToLong(Group::queuedTasks).sum(); + } + + /** Closing cancels queued work but deliberately does not interrupt shared metadata IO. */ + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + long deadline = System.nanoTime() + unit.toNanos(timeout); + for (Thread worker : workers) { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0) { + return workers.stream().noneMatch(Thread::isAlive); + } + TimeUnit.NANOSECONDS.timedJoin(worker, remaining); + } + return workers.stream().noneMatch(Thread::isAlive); + } + + private static final class Waiter { + private final Thread thread = Thread.currentThread(); + private final AtomicBoolean waiting = new AtomicBoolean(true); + } + + private Group awaitGroup() { + while (!closed) { + Group group = ready.poll(); + if (group != null) { + return group; + } + Waiter waiter = new Waiter(); + waiters.add(waiter); + try { + group = ready.poll(); + if (group != null) { + return group; + } + while (!closed && waiter.waiting.get()) { + LockSupport.park(this); + Thread.interrupted(); + } + } finally { + boolean notified = !waiter.waiting.getAndSet(false); + waiters.remove(waiter); + if (notified && !ready.isEmpty()) { + signalOne(); + } + } + } + return null; + } + + private void signalOne() { + Waiter waiter; + while ((waiter = waiters.poll()) != null) { + if (waiter.waiting.compareAndSet(true, false)) { + LockSupport.unpark(waiter.thread); + return; + } + } + } + + @Override + public void close() { + synchronized (this) { + closed = true; + } + for (Group group : groups) { + group.close(); + } + for (Thread worker : workers) { + LockSupport.unpark(worker); + } + ready.clear(); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/util/LocationPath.java b/fe/fe-core/src/main/java/org/apache/doris/common/util/LocationPath.java index 32e35dcf1e44c4..123ba12a2ab283 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/util/LocationPath.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/util/LocationPath.java @@ -144,8 +144,7 @@ public static LocationPath ofAdapters(String location, throw new IllegalArgumentException("Invalid location: " + location + ", normalized location is null"); } } - String encodedLocation = encodedLocation(normalizedLocation); - URI uri = URI.create(encodedLocation); + URI uri = parseFileSystemUri(normalizedLocation); String fsIdentifier = Strings.nullToEmpty(uri.getScheme()) + "://" + Strings.nullToEmpty(uri.getAuthority()); if (StringUtils.isBlank(schema)) { schema = Strings.nullToEmpty(uri.getScheme()); @@ -163,8 +162,7 @@ public static StorageTypeId fromSchemaWithContext(String location, String schema public static LocationPath of(String location) { String schema = extractScheme(location); - String encodedLocation = encodedLocation(location); - URI uri = URI.create(encodedLocation); + URI uri = parseFileSystemUri(location); String fsIdentifier = Strings.nullToEmpty(uri.getScheme()) + "://" + Strings.nullToEmpty(uri.getAuthority()); return new LocationPath(schema, location, fsIdentifier, null); } @@ -186,8 +184,7 @@ public static LocationPath of(String location, StorageAdapter storageAdapter) { String schema = extractScheme(location); String normalizedLocation = storageAdapter.validateAndNormalizeUri(location); - String encodedLocation = encodedLocation(normalizedLocation); - URI uri = URI.create(encodedLocation); + URI uri = parseFileSystemUri(normalizedLocation); String fsIdentifier = Strings.nullToEmpty(uri.getScheme()) + "://" + Strings.nullToEmpty(uri.getAuthority()); if (StringUtils.isBlank(schema)) { @@ -248,8 +245,7 @@ public static LocationPath ofWithCache(String location, } } else { // Fallback to full URI parsing - String encodedLocation = encodedLocation(normalizedLocation); - URI uri = URI.create(encodedLocation); + URI uri = parseFileSystemUri(normalizedLocation); String authority = uri.getAuthority(); if (Strings.isNullOrEmpty(authority)) { throw new StoragePropertiesException("Invalid location, missing authority: " + normalizedLocation); @@ -339,7 +335,62 @@ private static StorageAdapter findStorageAdapter(StorageTypeId type, String sche } + private static String fileSystemUriPrefix(String location) { + // Only scheme and authority are consumed by the factories. For a hierarchical URI, the + // path is opaque here: encodedLocation escapes its special characters before URI parsing, + // and normalizedLocation retains the original bytes. Avoid encoding/parsing that path for + // every file and split. Keep the slash so an empty authority (hdfs:///path) stays valid. + int schemeEnd = location.indexOf(SCHEME_DELIM); + if (schemeEnd > 0 && location.indexOf('/') == schemeEnd + 1) { + int pathStart = location.indexOf('/', schemeEnd + SCHEME_DELIM.length()); + if (pathStart >= 0) { + location = location.substring(0, pathStart + 1); + } + } + return location; + } + + private static URI parseFileSystemUri(String location) { + return URI.create(encodedLocation(fileSystemUriPrefix(location))); + } + + /** A single producer's parser, retaining only its last filesystem prefix and file path. */ + public static class Factory { + private String lastPrefix; + private LocationPath lastPath; + + public LocationPath create(String location) { + if (lastPath != null && location.equals(lastPath.normalizedLocation)) { + return lastPath; + } + String prefix = fileSystemUriPrefix(location); + if (prefix.equals(lastPrefix)) { + // Only the opaque path changed; schema and authority have already been parsed. + lastPath = new LocationPath(lastPath.schema, location, lastPath.fsIdentifier, null); + } else { + LocationPath parsed = LocationPath.of(location); + lastPrefix = prefix; + lastPath = parsed; + } + return lastPath; + } + } + private static String encodedLocation(String location) { + // URLEncoder leaves these characters unchanged (the subsequent replacements restore ':' + // and '/'). Preserve URI.create validation while avoiding encoding an ordinary URI prefix. + boolean needsEncoding = false; + for (int i = 0; i < location.length(); i++) { + char c = location.charAt(i); + if (!(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' + || c == '-' || c == '_' || c == '.' || c == '*' || c == ':' || c == '/')) { + needsEncoding = true; + break; + } + } + if (!needsEncoding) { + return location; + } try { return URLEncoder.encode(location, StandardCharsets.UTF_8.name()) .replace("%2F", "/") diff --git a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java index d5989123614afc..aa866eba680e7d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java @@ -28,6 +28,7 @@ import org.apache.doris.nereids.StatementContext; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.GlobalVariable; +import org.apache.doris.qe.StmtExecutor; import org.apache.doris.qe.VariableMgr; import java.util.Collections; @@ -174,9 +175,20 @@ public ConnectorSession build() { } } } + ConnectorStatementScope scope = captureStatementScope(); + if (scope != ConnectorStatementScope.NONE && connectContext != null) { + StmtExecutor owner = connectContext.getExecutor(); + if (owner != null) { + scope.computeIfAbsent(QueryTaskGroups.SCOPE_KEY, () -> { + QueryTaskGroups groups = new QueryTaskGroups(); + owner.registerQueryCancellationCallback(groups::close); + return groups; + }); + } + } return new ConnectorSessionImpl(queryId, user, timeZone, locale, catalogId, catalogName, catalogProperties, sessionProperties, sid, cred, - captureStatementScope(), captureBeExecVersion()); + scope, captureBeExecVersion()); } /** diff --git a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionImpl.java b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionImpl.java index ffd66d2b0cdff5..00c9951e10064e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionImpl.java +++ b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionImpl.java @@ -22,6 +22,7 @@ import org.apache.doris.connector.spi.ConnectorDelegatedCredential; import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.ConnectorStatementScope; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.connector.spi.handle.ConnectorTransaction; import java.util.Collections; @@ -57,6 +58,7 @@ public class ConnectorSessionImpl implements ConnectorSession { // The backend execution version the planner pinned for the statement, captured at construction like the // statement scope; Config.be_exec_version when the session is built outside a planned statement. private final int beExecVersion; + private final QueryTaskGroups taskGroups; // Otherwise-immutable session; this is bound once by the insert executor at write time // for connectors using the SPI transaction model (e.g. maxcompute), and read back by the // connector's planWrite via getCurrentTransaction(). volatile for cross-thread visibility. @@ -88,6 +90,15 @@ public class ConnectorSessionImpl implements ConnectorSession { this.delegatedCredential = delegatedCredential; this.statementScope = statementScope != null ? statementScope : ConnectorStatementScope.NONE; this.beExecVersion = beExecVersion; + // Capture the handle now. Looking it up from a cleared scope on a worker could resurrect a + // finished query. Offline sessions use the SPI's direct scheduler and own no background work. + this.taskGroups = this.statementScope == ConnectorStatementScope.NONE ? null + : this.statementScope.computeIfAbsent(QueryTaskGroups.SCOPE_KEY, QueryTaskGroups::new); + } + + @Override + public ConnectorTaskScheduler getTaskScheduler(ConnectorTaskScheduler.Stage stage) { + return taskGroups == null ? ConnectorTaskScheduler.DIRECT : taskGroups.get(stage); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/connector/DefaultConnectorContext.java b/fe/fe-core/src/main/java/org/apache/doris/connector/DefaultConnectorContext.java index f519f6de28b368..801d92d2bfb5a8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/connector/DefaultConnectorContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/connector/DefaultConnectorContext.java @@ -23,6 +23,7 @@ import org.apache.doris.common.ClientPool; import org.apache.doris.common.Config; import org.apache.doris.common.EnvUtils; +import org.apache.doris.common.ThreadPoolManager; import org.apache.doris.common.Version; import org.apache.doris.common.util.LocationPath; import org.apache.doris.connector.spi.Connector; @@ -65,6 +66,8 @@ import java.util.Map; import java.util.Objects; import java.util.concurrent.Callable; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.ThreadPoolExecutor; import java.util.function.Function; import java.util.function.Supplier; import java.util.function.UnaryOperator; @@ -123,6 +126,12 @@ public void afterRequest() { } }; + @Override + public RunnableFuture submitMetadataTask(Callable task) { + return ThreadPoolManager.submitWithCallerHelp( + (ThreadPoolExecutor) Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor(), task); + } + public DefaultConnectorContext(String catalogName, long catalogId) { this(catalogName, catalogId, () -> NOOP_AUTH); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/connector/QueryTaskGroups.java b/fe/fe-core/src/main/java/org/apache/doris/connector/QueryTaskGroups.java new file mode 100644 index 00000000000000..e946beea2d4d44 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/connector/QueryTaskGroups.java @@ -0,0 +1,105 @@ +// 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; + +import org.apache.doris.catalog.Env; +import org.apache.doris.common.Config; +import org.apache.doris.common.util.FairTaskExecutor; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; + +import java.io.Closeable; +import java.util.EnumMap; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.function.Consumer; +import java.util.function.Function; + +/** Captured by every session of one statement before off-thread planning can begin. */ +final class QueryTaskGroups implements Closeable { + static final String SCOPE_KEY = "engine:query-task-groups"; + private final Map groups = + new EnumMap<>(ConnectorTaskScheduler.Stage.class); + private final Function executors; + private boolean closed; + + QueryTaskGroups() { + this(stage -> Env.getCurrentEnv().getExtMetaCacheMgr().getQueryTaskExecutor(stage)); + } + + QueryTaskGroups(Function executors) { + this.executors = executors; + } + + synchronized ConnectorTaskScheduler get(ConnectorTaskScheduler.Stage stage) { + if (closed) { + // This captured query has ended; callers must not treat cancellation as a retryable rejection. + throw new CancellationException("Query task groups are closed"); + } + return groups.computeIfAbsent(stage, key -> new Adapter( + executors.apply(key).newGroup(Config.external_scan_task_queue_capacity))); + } + + @Override + public void close() { + Adapter[] snapshot; + synchronized (this) { + if (closed) { + return; + } + closed = true; + snapshot = groups.values().toArray(new Adapter[0]); + } + // Never invoke task completion (which may re-enter query cleanup) under the registry lock. + for (Adapter group : snapshot) { + group.delegate.close(); + } + } + + private static final class Adapter implements ConnectorTaskScheduler { + private final FairTaskExecutor.Group delegate; + + private Adapter(FairTaskExecutor.Group delegate) { + this.delegate = delegate; + } + + @Override + public boolean tryExecute(Runnable task, Consumer onFailure, Runnable onSpace) { + return delegate.tryExecute(task, onFailure, onSpace); + } + + @Override + public boolean isClosed() { + return delegate.isClosed(); + } + + @Override + public boolean tryExecuteContinuation(Runnable task, Consumer onFailure, Runnable onSpace) { + return delegate.tryExecuteContinuation(task, onFailure, onSpace); + } + + @Override + public void removeSpaceListener(Runnable onSpace) { + delegate.removeSpaceListener(onSpace); + } + + @Override + public void execute(Runnable task) { + delegate.execute(task); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalMetaCacheMgr.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalMetaCacheMgr.java index de53810f544536..dfed5491552493 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalMetaCacheMgr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalMetaCacheMgr.java @@ -21,10 +21,12 @@ import org.apache.doris.common.Config; import org.apache.doris.common.ThreadPoolManager; import org.apache.doris.common.cache.NereidsSortedPartitionsCacheManager; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.connector.cache.CacheSpec; import org.apache.doris.connector.cache.CatalogMetaCache; import org.apache.doris.connector.cache.MetaCache; import org.apache.doris.connector.cache.MetaCacheGovernance; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.datasource.doris.DorisExternalMetaCache; import org.apache.doris.datasource.doris.RemoteDorisExternalCatalog; import org.apache.doris.datasource.metacache.ExternalCatalogMetaCache; @@ -34,6 +36,7 @@ import org.apache.doris.datasource.metacache.MetaCacheEntryStats; import org.apache.doris.datasource.mvcc.PluginDrivenPartitionView; import org.apache.doris.fs.FileSystemCache; +import org.apache.doris.nereids.util.MoreFieldsThread; import com.github.benmanes.caffeine.cache.stats.CacheStats; import com.google.common.collect.Maps; @@ -90,6 +93,8 @@ public class ExternalMetaCacheMgr { private ExecutorService fileListingExecutor; // This executor is used to schedule the getting split tasks private ExecutorService scheduleExecutor; + private FairTaskExecutor splitTaskExecutor; + private FairTaskExecutor metadataTaskExecutor; private final Map cacheTypes = Maps.newConcurrentMap(); // all catalogs could share the same fsCache. @@ -156,6 +161,32 @@ public ExecutorService getFileListingExecutor() { return fileListingExecutor; } + public synchronized FairTaskExecutor getQueryTaskExecutor(ConnectorTaskScheduler.Stage stage) { + if (stage == ConnectorTaskScheduler.Stage.SPLIT) { + if (splitTaskExecutor == null) { + splitTaskExecutor = newFairTaskExecutor("SplitTask", Config.external_scan_split_threads); + } + return splitTaskExecutor; + } + if (metadataTaskExecutor == null) { + metadataTaskExecutor = newFairTaskExecutor("MetadataTask", Config.external_scan_metadata_threads); + } + return metadataTaskExecutor; + } + + private FairTaskExecutor newFairTaskExecutor(String name, int threads) { + java.util.concurrent.atomic.AtomicInteger sequence = new java.util.concurrent.atomic.AtomicInteger(); + FairTaskExecutor executor = new FairTaskExecutor(threads, task -> { + Thread worker = new MoreFieldsThread(task); + worker.setName(name + "-" + sequence.incrementAndGet()); + worker.setDaemon(true); + worker.setContextClassLoader(ExternalMetaCacheMgr.class.getClassLoader()); + return worker; + }); + ThreadPoolManager.registerFairTaskExecutorMetric(name, executor); + return executor; + } + public ExecutorService getScheduleExecutor() { return scheduleExecutor; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java index b5080ef343385d..8f9a616181575b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/scan/FederationBackendPolicy.java @@ -329,8 +329,10 @@ public Multimap computeScanRangeAssignment(List splits) t candidateNodes = candidateCount >= backends.size() ? backends : consistentHash.getNode(split, candidateCount); } else { - candidateNodes = consistentHash.getNode( - split, Config.split_assigner_min_consistent_hash_candidate_num); + // The ring cannot find more candidates than eligible physical backends. + candidateNodes = consistentHash.getNode(split, + Math.min(Config.split_assigner_min_consistent_hash_candidate_num, + backends.size())); } } break; @@ -381,6 +383,13 @@ private void equateDistribution(ListMultimap assignment) { return; } + // Small streaming batches usually keep the accumulated weights within the allowed variance. + // Avoid rebuilding both priority queues when their first comparison would already stop redistribution. + if (assignedWeightCounts.lastKey() - assignedWeightCounts.firstKey() + <= SplitWeight.rawValueForStandardSplitCount(Config.split_assigner_max_split_num_variance)) { + return; + } + List allNodes = new ArrayList<>(); allNodes.addAll(backends); Collections.sort(allNodes, Comparator.comparing(Backend::getId)); 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 b1d0d4f480e832..e434dc523e735a 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 @@ -439,7 +439,8 @@ public void createScanRangeLocations() throws UserException { // and finally numSplitsPerBE is 0, resulting in no data being queried. int numSplitsPerBE = Math.max(selectedSplitNum / backendPolicy.numBackends(), 1); for (Backend backend : backendPolicy.getBackends()) { - SplitSource splitSource = new SplitSource(backend, splitAssignment, maxWaitTime); + SplitSource splitSource = new SplitSource(backend, splitAssignment, maxWaitTime, + sessionVariable.getSplitInitialBatchSize()); splitSources.add(splitSource); Env.getCurrentEnv().getSplitSourceManager().registerSplitSource(splitSource); TScanRangeLocations curLocations = newLocations(); 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 1808f6d0f4464f..89dccb471dc79e 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 @@ -28,7 +28,6 @@ import org.apache.doris.analysis.TupleDescriptor; import org.apache.doris.catalog.ArrayType; import org.apache.doris.catalog.Column; -import org.apache.doris.catalog.Env; import org.apache.doris.catalog.MapType; import org.apache.doris.catalog.StructType; import org.apache.doris.catalog.TableIf; @@ -42,6 +41,7 @@ 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.ConnectorTaskScheduler; import org.apache.doris.connector.spi.handle.ConnectorColumnHandle; import org.apache.doris.connector.spi.handle.ConnectorTableHandle; import org.apache.doris.connector.spi.handle.PassthroughQueryTableHandle; @@ -111,7 +111,6 @@ import java.util.OptionalLong; import java.util.Random; import java.util.Set; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -1013,6 +1012,13 @@ public static Runnable buildReadTransactionReleaseCallback( }); } + static Runnable buildSplitSourceDemandCallback(ConnectorScanPlanProvider provider, ConnectorSplitSource source) { + return () -> onPluginClassLoader(provider, () -> { + source.onMoreSplitsRequested(); + return null; + }); + } + static final class QueryFinishCleanup implements Runnable { private final String queryId; private final ConnectorStatementScope statementScope; @@ -1769,8 +1775,9 @@ public List getSplits(int numBackends) throws UserException { } List splits = new ArrayList<>(ranges.size()); + PluginDrivenSplit.Factory splitFactory = new PluginDrivenSplit.Factory(); for (ConnectorScanRange range : ranges) { - splits.add(toSplit(range)); + splits.add(toSplit(range, splitFactory)); } // FIX-E (explain gap): accumulate the native/total scan-range counts (for the connector // EXPLAIN line paimonNativeReadSplits) and, under COUNT(*) pushdown, the precomputed merged row @@ -1861,11 +1868,11 @@ private static boolean isExplainOnly() { // Every range the connector plans becomes a split here, on the planning thread or on the batch-mode // split generation threads, so that the plan's single-use ranges are known (cannotBeRedispatched). - private Split toSplit(ConnectorScanRange range) { + private Split toSplit(ConnectorScanRange range, PluginDrivenSplit.Factory splitFactory) { if (range.isSingleUse()) { plannedSingleUseRange = true; } - return new PluginDrivenSplit(range); + return splitFactory.create(range); } /** @@ -1969,6 +1976,11 @@ public boolean isBatchMode() { } private boolean computeBatchMode() { + // The query-level switch governs both partition batching and connector streaming. Check it + // before probing the provider, which may access remote metadata to estimate file counts. + if (!sessionVariable.getEnableExternalTableBatchMode()) { + return false; + } // getScanPlanProvider() may be null for connectors without scan capability; mirror the // null-guard in getSplits() so isBatchMode (run on the dispatch + explain paths) never NPEs. ConnectorScanPlanProvider scanProvider = resolveScanProvider(); @@ -2132,7 +2144,7 @@ public void startSplit(int numBackends) { } else { sources = createPartitionBatchSources(scanProvider, request, new ArrayList<>(selectedPartitions.selectedPartitions.keySet()), - sessionVariable.getNumPartitionsInBatchMode(), Config.max_external_cache_loader_thread_pool_size); + sessionVariable.getNumPartitionsInBatchMode(), Config.external_scan_split_threads); } startSplitProducers(scanProvider, sources, summaryProfile == null ? null : summaryProfile.getExecutionSummary()); @@ -2153,7 +2165,6 @@ private ConnectorScanRequest prepareBatchScan(QueryFinishCleanup queryFinishClea private void startSplitProducers(ConnectorScanPlanProvider scanProvider, List sources, RuntimeProfile executionSummary) { - Executor scheduleExecutor = Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor(); AtomicReference batchException = new AtomicReference<>(null); SubmittedTaskFinalizer profileFinalizer = new SubmittedTaskFinalizer(() -> { try { @@ -2177,43 +2188,30 @@ private void startSplitProducers(ConnectorScanPlanProvider scanProvider, List { - try { - for (SplitProducer.BatchSource source : sources) { - if (batchException.get() != null || splitAssignment.isStop()) { - break; - } - profileFinalizer.taskSubmitted(); - new SplitProducer(splitAssignment, scheduleExecutor, source, failure -> { - try { - if (failure != null) { - publishBatchFailure(batchException, splitAssignment, - new UserException(failure.getMessage(), failure)); - } - } finally { - profileFinalizer.taskFinished(); - } - }).start(); - } - } finally { - profileFinalizer.closeDispatch(); - } - }; + // Admitting the bounded number of lanes does no listing or planning. Each producer + // independently suspends on task-queue capacity without occupying a worker. try { - // 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); + Executor scheduleExecutor = connectorSession.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT); + for (SplitProducer.BatchSource source : sources) { + if (batchException.get() != null || splitAssignment.isStop()) { + break; + } + profileFinalizer.taskSubmitted(); + new SplitProducer(splitAssignment, scheduleExecutor, source, failure -> { + try { + if (failure != null) { + publishBatchFailure(batchException, splitAssignment, + new UserException(failure.getMessage(), failure)); + } + } finally { + profileFinalizer.taskFinished(); + } + }).start(); } } catch (Exception e) { - try { - publishBatchFailure(batchException, splitAssignment, - new UserException(e.getMessage(), e)); - } finally { - profileFinalizer.closeDispatch(); - } + publishBatchFailure(batchException, splitAssignment, new UserException(e.getMessage(), e)); + } finally { + profileFinalizer.closeDispatch(); } } @@ -2227,6 +2225,7 @@ List createPartitionBatchSources(ConnectorScanPlanPro for (int sourceIndex = 0; sourceIndex < sourceCount; sourceIndex++) { sources.add(new SplitProducer.BatchSource() { private ConnectorSplitSource source; + private final PluginDrivenSplit.Factory splitFactory = new PluginDrivenSplit.Factory(); @Override public List getNextBatch() { @@ -2235,7 +2234,7 @@ public List getNextBatch() { while (splits.size() < splitAssignment.getSplitBatchSize() && splitAssignment.needMoreSplit()) { if (source != null && source.hasNext()) { - splits.add(toSplit(source.next())); + splits.add(toSplit(source.next(), splitFactory)); continue; } closeSource(); @@ -2288,17 +2287,20 @@ private SplitProducer.BatchSource createStreamingBatchSource(ConnectorScanPlanPr } return new SplitProducer.BatchSource() { private ConnectorSplitSource source; + private final PluginDrivenSplit.Factory splitFactory = new PluginDrivenSplit.Factory(); @Override public List getNextBatch() { return onPluginClassLoader(scanProvider, () -> { if (source == null) { source = scanProvider.streamSplits(connectorSession, request); + splitAssignment.setMoreSplitsRequestedCallback( + buildSplitSourceDemandCallback(scanProvider, source)); } List splits = new ArrayList<>(splitAssignment.getSplitBatchSize()); while (splits.size() < splitAssignment.getSplitBatchSize() && splitAssignment.needMoreSplit() && source.hasNext()) { - splits.add(toSplit(source.next())); + splits.add(toSplit(source.next(), splitFactory)); } generatedSplits.setValue(generatedSplits.getValue() + splits.size()); return splits; diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/FileScanRangeCompactor.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/FileScanRangeCompactor.java new file mode 100644 index 00000000000000..d47a3eb446e599 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/FileScanRangeCompactor.java @@ -0,0 +1,136 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource.split; + +import org.apache.doris.thrift.TFileScanRange; +import org.apache.doris.thrift.TFileScanRangeParams; +import org.apache.doris.thrift.TQueryOptions; +import org.apache.doris.thrift.TScanRangeLocations; +import org.apache.doris.thrift.TScanRangeParams; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** Reduces repeated file-scan wrappers only at the transport boundary, after scan assignment. */ +public final class FileScanRangeCompactor { + private FileScanRangeCompactor() { + } + + /** + * Compact only after fetching: maxBatchSize counts the original split descriptions, not envelopes. + * Every range comes from FileQueryScanNode.splitToScanRange, so file-specific attributes remain in + * TFileRangeDesc. Keep adjacent envelopes separate when their scan parameters or locations differ. + */ + public static List compactSplitBatch(List scanRanges) { + if (scanRanges.size() < 2) { + return scanRanges; + } + List compacted = new ArrayList<>(); + TScanRangeLocations previous = null; + TFileScanRange previousFiles = null; + boolean copied = false; + for (TScanRangeLocations current : scanRanges) { + TFileScanRange files = current.getScanRange().getExtScanRange().getFileScanRange(); + if (previous != null + && Objects.equals(previous.getLocations(), current.getLocations()) + && canMerge(previousFiles, files)) { + if (!copied) { + // Keep compaction local to this response; leave the dequeued entries unchanged. + previous = new TScanRangeLocations(previous); + previousFiles = previous.getScanRange().getExtScanRange().getFileScanRange(); + compacted.set(compacted.size() - 1, previous); + copied = true; + } + previousFiles.getRanges().addAll(files.getRanges()); + } else { + compacted.add(current); + previous = current; + previousFiles = files; + copied = false; + } + } + return compacted; + } + + /** + * Compact synchronous file scans within each already assigned instance. BE limits the number of + * local scanners by the number of wrappers, so retain at least its scanner-concurrency upper bound: + * max_file_scanners_concurrency * parallel_instance (the BE may divide this by operator parallelism). + * Non-file scans and the assignment's original lists are left unchanged. + */ + public static Map> compactInstanceScanRanges( + Map> scanRanges, + Map fileScanParams, TQueryOptions options) { + if (fileScanParams.isEmpty()) { + return scanRanges; + } + int scannersPerInstance = options.getMaxFileScannersConcurrency() > 0 + ? options.getMaxFileScannersConcurrency() : 16; + long minimumWrappers = (long) scannersPerInstance * options.getParallelInstance(); + Map> compacted = new LinkedHashMap<>(); + for (Map.Entry> entry : scanRanges.entrySet()) { + List ranges = entry.getValue(); + compacted.put(entry.getKey(), fileScanParams.containsKey(entry.getKey()) + ? compactAssignedRanges(ranges, minimumWrappers) : ranges); + } + return compacted; + } + + private static List compactAssignedRanges( + List ranges, long minimumWrappers) { + // Floor division ensures compaction never leaves fewer wrappers than the BE can scan in parallel. + int maximumGroupSize = (int) Math.max(1, ranges.size() / minimumWrappers); + if (maximumGroupSize == 1) { + return ranges; + } + List compacted = new ArrayList<>(); + TScanRangeParams previous = null; + TFileScanRange previousFiles = null; + int groupSize = 0; + for (TScanRangeParams current : ranges) { + TFileScanRange files = current.getScanRange().getExtScanRange().getFileScanRange(); + if (previous != null && groupSize < maximumGroupSize + && previous.isSetVolumeId() == current.isSetVolumeId() + && previous.getVolumeId() == current.getVolumeId() + && canMerge(previousFiles, files)) { + if (groupSize == 1) { + previous = new TScanRangeParams(previous); + previousFiles = previous.getScanRange().getExtScanRange().getFileScanRange(); + compacted.set(compacted.size() - 1, previous); + } + previousFiles.getRanges().addAll(files.getRanges()); + groupSize++; + } else { + compacted.add(current); + previous = current; + previousFiles = files; + groupSize = 1; + } + } + return compacted; + } + + private static boolean canMerge(TFileScanRange left, TFileScanRange right) { + // A split-source reference is a distinct stream, not a list of file ranges to concatenate. + return !left.isSetSplitSource() && !right.isSetSplitSource() + && Objects.equals(left.getParams(), right.getParams()); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/PluginDrivenSplit.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/PluginDrivenSplit.java index 85a8c2b2c270a7..ac89c910c8ab1d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/PluginDrivenSplit.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/PluginDrivenSplit.java @@ -38,7 +38,11 @@ public class PluginDrivenSplit extends FileSplit { private final String splitIdentity; public PluginDrivenSplit(ConnectorScanRange scanRange) { - super(buildPath(scanRange), + this(scanRange, buildPath(scanRange)); + } + + private PluginDrivenSplit(ConnectorScanRange scanRange, LocationPath path) { + super(path, scanRange.getStart(), scanRange.getLength(), scanRange.getFileSize(), @@ -95,6 +99,16 @@ private static LocationPath buildPath(ConnectorScanRange scanRange) { return LocationPath.of(pathStr); } + /** Converts one producer stream while reusing filesystem identities across its files. */ + public static class Factory { + private final LocationPath.Factory paths = new LocationPath.Factory(); + + public PluginDrivenSplit create(ConnectorScanRange scanRange) { + String location = scanRange.getPath().orElse("connector://virtual"); + return new PluginDrivenSplit(scanRange, paths.create(location)); + } + } + private static List buildPartitionValues(ConnectorScanRange scanRange) { Map partValues = scanRange.getPartitionValues(); if (partValues == null) { 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 ca904044815bb1..2f0a44c960fde5 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 @@ -31,7 +31,6 @@ 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; @@ -51,7 +50,8 @@ public class SplitAssignment { private final Set sources = new HashSet<>(); private final FederationBackendPolicy backendPolicy; private final SplitGenerator splitGenerator; - private final ConcurrentHashMap>> assignment + // Queue identity is the stable backend ID; Backend.hashCode also hashes mutable endpoint/health fields. + private final ConcurrentHashMap assignment = new ConcurrentHashMap<>(); private final SplitToScanRange splitToScanRange; private final Map locationProperties; @@ -59,14 +59,36 @@ public class SplitAssignment { private final boolean fileCacheAdmission; private final int splitBatchSize; private final int queueCapacity; - private final Object assignLock = new Object(); + // Backend policy state is mutable; it must not delay publishing an already converted first split. + private final Object backendAssignmentLock = new Object(); + private final Object readinessLock = new Object(); private Split sampleSplit = null; private final AtomicBoolean isStopped = new AtomicBoolean(false); private final AtomicBoolean scheduleFinished = new AtomicBoolean(false); private volatile UserException exception = null; private final List closeableResources = new ArrayList<>(); - private final Map> queueSpaceListeners = new HashMap<>(); + private final ConcurrentHashMap> queueSpaceListeners = new ConcurrentHashMap<>(); + private volatile Runnable moreSplitsRequestedCallback = () -> {}; + + /** Data readiness and terminal state use the same per-backend wait channel. */ + private static final class AssignmentQueue { + private final BlockingQueue> batches; + // Written under this channel's monitor; producers avoid that monitor when nobody waits. + private volatile int waitingReaders; + + private AssignmentQueue(int capacity) { + batches = new LinkedBlockingQueue<>(capacity); + } + + private void signalReaders() { + if (waitingReaders > 0) { + synchronized (this) { + notifyAll(); + } + } + } + } public SplitAssignment( FederationBackendPolicy backendPolicy, @@ -91,13 +113,13 @@ public SplitAssignment( public void init() throws UserException { splitGenerator.startSplit(backendPolicy.numBackends()); - synchronized (assignLock) { + synchronized (readinessLock) { final int waitIntervalTimeMillis = 100; final int initTimeoutMillis = 30000; // 30s int waitTotalTime = 0; while (sampleSplit == null && needMoreSplit()) { try { - assignLock.wait(waitIntervalTimeMillis); + readinessLock.wait(waitIntervalTimeMillis); } catch (InterruptedException e) { throw new UserException(e.getMessage(), e); } @@ -121,15 +143,15 @@ private void appendBatch(Multimap batch) throws UserException { for (Backend backend : batch.keySet()) { Collection locations = createScanRanges(backend, batch.get(backend)); while (needMoreSplit()) { - BlockingQueue> queue = - getOrCreateAssignmentQueue(backend); + AssignmentQueue queue = getOrCreateAssignmentQueue(backend); try { - if (queue.offer(locations, 100, TimeUnit.MILLISECONDS)) { + if (queue.batches.offer(locations, 100, TimeUnit.MILLISECONDS)) { + queue.signalReaders(); break; } } catch (InterruptedException e) { Thread.currentThread().interrupt(); - addUserException(new UserException("Failed to offer batch split by interrupted", e)); + setException(new UserException("Failed to offer batch split by interrupted", e)); } } } @@ -137,7 +159,7 @@ private void appendBatch(Multimap batch) throws UserException { Multimap assignBatch(List splits) throws UserException { Preconditions.checkArgument(splits.size() <= splitBatchSize, "Producer batch exceeds split batch size"); - synchronized (assignLock) { + synchronized (backendAssignmentLock) { return backendPolicy.computeScanRangeAssignment(splits); } } @@ -150,48 +172,64 @@ Collection createScanRanges(Backend backend, Collection locations, Runnable resume) { - // Coordinate the failed offer and listener registration with consumer notifications so - // capacity released just before suspension cannot be missed. + if (!needMoreSplit()) { + return false; + } + AssignmentQueue queue = getOrCreateAssignmentQueue(backend); + if (queue.batches.offer(locations)) { + queue.signalReaders(); + removeQueueSpaceListener(backend, resume); + return true; + } 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)) { + queueSpaceListeners.computeIfAbsent(backend.getId(), key -> new HashSet<>()).add(resume); + // A consumer with no registered waiter skips this monitor. Retry after registration + // so capacity released between the first offer and registration cannot be missed. + if (queue.batches.offer(locations)) { + queue.signalReaders(); removeQueueSpaceListener(backend, resume); return true; } - queueSpaceListeners.computeIfAbsent(backend, key -> new HashSet<>()).add(resume); return false; } } void removeQueueSpaceListener(Backend backend, Runnable listener) { + if (!queueSpaceListeners.containsKey(backend.getId())) { + return; + } synchronized (queueSpaceListeners) { - Set listeners = queueSpaceListeners.get(backend); + Set listeners = queueSpaceListeners.get(backend.getId()); if (listeners != null) { listeners.remove(listener); if (listeners.isEmpty()) { - queueSpaceListeners.remove(backend); + queueSpaceListeners.remove(backend.getId()); } } } } void notifyQueueSpace(Backend backend) { + if (!queueSpaceListeners.containsKey(backend.getId())) { + return; + } Set listeners; synchronized (queueSpaceListeners) { - listeners = queueSpaceListeners.remove(backend); + listeners = queueSpaceListeners.remove(backend.getId()); } if (listeners != null) { listeners.forEach(Runnable::run); @@ -224,20 +262,38 @@ public int getSplitBatchSize() { } private void notifyAssignment() { - synchronized (assignLock) { - assignLock.notify(); + synchronized (readinessLock) { + readinessLock.notifyAll(); } + assignment.values().forEach(AssignmentQueue::signalReaders); } - private BlockingQueue> getOrCreateAssignmentQueue(Backend backend) { - return assignment.computeIfAbsent(backend, be -> new LinkedBlockingQueue<>(queueCapacity)); + private AssignmentQueue getOrCreateAssignmentQueue(Backend backend) { + return assignment.computeIfAbsent(backend.getId(), be -> new AssignmentQueue(queueCapacity)); + } + + Collection awaitBatch(Backend backend, long waitNanos) throws InterruptedException { + AssignmentQueue queue = getOrCreateAssignmentQueue(backend); + synchronized (queue) { + queue.waitingReaders++; + try { + // Register before checking both predicates. An offer or terminal transition either + // precedes this check or signals after wait releases this channel's monitor. + if (queue.batches.isEmpty() && needMoreSplit()) { + TimeUnit.NANOSECONDS.timedWait(queue, waitNanos); + } + return queue.batches.poll(); + } finally { + queue.waitingReaders--; + } + } } public BlockingQueue> getAssignedSplits(Backend backend) throws UserException { if (exception != null) { throw exception; } - BlockingQueue> splits = getOrCreateAssignmentQueue(backend); + BlockingQueue> splits = getOrCreateAssignmentQueue(backend).batches; if (scheduleFinished.get() && splits.isEmpty() || isStopped.get()) { return null; } @@ -255,11 +311,17 @@ public void setException(UserException e) { listeners.forEach(Runnable::run); } - private synchronized void addUserException(UserException e) { - if (exception != null) { - exception.addSuppressed(e); - } else { - exception = e; + private void addUserException(UserException e) { + synchronized (readinessLock) { + // Cancellation can interrupt an active source; preserve failures published before stop. + if (isStopped.get()) { + return; + } + if (exception != null) { + exception.addSuppressed(e); + } else { + exception = e; + } } } @@ -269,16 +331,19 @@ public void finishSchedule() { } public void stop() { - if (!isStopped.compareAndSet(false, true)) { - return; + synchronized (readinessLock) { + if (!isStopped.compareAndSet(false, true)) { + return; + } } + notifyAssignment(); List resources; synchronized (closeableResources) { resources = new ArrayList<>(closeableResources); closeableResources.clear(); + moreSplitsRequestedCallback = () -> {}; } resources.forEach(this::closeResource); - notifyAssignment(); if (exception != null) { throw new RuntimeException(exception); } @@ -288,6 +353,21 @@ public boolean isStop() { return isStopped.get(); } + /** Registered by the streaming producer before publishing its first split. */ + public void setMoreSplitsRequestedCallback(Runnable callback) { + synchronized (closeableResources) { + if (!isStopped.get()) { + moreSplitsRequestedCallback = callback; + } + } + } + + public void onMoreSplitsRequested() { + if (needMoreSplit()) { + moreSplitsRequestedCallback.run(); + } + } + public void addCloseable(Closeable resource) { synchronized (closeableResources) { if (!isStop()) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java index 6d2f0e26207070..8fe29b198c4982 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/split/SplitProducer.java @@ -18,6 +18,7 @@ package org.apache.doris.datasource.split; import org.apache.doris.common.UserException; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.spi.Split; import org.apache.doris.system.Backend; import org.apache.doris.thrift.TScanRangeLocations; @@ -33,15 +34,21 @@ import java.util.List; import java.util.Map; import java.util.concurrent.Executor; +import java.util.concurrent.TimeUnit; import java.util.function.Consumer; +import java.util.function.LongSupplier; /** - * Serially advances one producer batch per task, yielding the shared executor when a BE queue is + * Serially advances a bounded number of producer batches per task, yielding 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); + // Bound both cheap cached batches and slow listing work so another query gets a turn. + static final int MAX_BATCHES_PER_TASK = 16; + static final long MAX_TASK_NANOS = TimeUnit.MILLISECONDS.toNanos(1); + public interface BatchSource extends Closeable { List getNextBatch() throws Exception; @@ -54,7 +61,9 @@ default void close() throws IOException { private final Executor executor; private final BatchSource source; private final Consumer completion; + private final LongSupplier nanoTime; private final Runnable resume = this::schedule; + private final Runnable retryAdmission = this::submit; private final Object scheduleLock = new Object(); private boolean scheduled; private boolean wakeupRequested; @@ -66,6 +75,12 @@ default void close() throws IOException { public SplitProducer(SplitAssignment assignment, Executor executor, BatchSource source, Consumer completion) { + this(assignment, executor, source, completion, System::nanoTime); + } + + SplitProducer(SplitAssignment assignment, Executor executor, BatchSource source, + Consumer completion, LongSupplier nanoTime) { + this.nanoTime = nanoTime; this.assignment = assignment; this.executor = executor; this.source = source; @@ -92,8 +107,23 @@ private void schedule() { } private void submit() { + submit(false); + } + + private void submit(boolean continuation) { try { - executor.execute(this); + if (executor instanceof ConnectorTaskScheduler) { + ConnectorTaskScheduler scheduler = (ConnectorTaskScheduler) executor; + Consumer onFailure = failure -> complete( + cancelled || !assignment.needMoreSplit() ? null : new Exception(failure)); + if (continuation) { + scheduler.tryExecuteContinuation(this, onFailure, retryAdmission); + } else { + scheduler.tryExecute(this, onFailure, retryAdmission); + } + } else { + executor.execute(this); + } } catch (RuntimeException e) { complete(cancelled || !assignment.needMoreSplit() ? null : e); } @@ -102,11 +132,20 @@ private void submit() { @Override public void run() { synchronized (scheduleLock) { + if (completed) { + return; + } wakeupRequested = false; } boolean more = false; try { - more = produceBatch(); + long start = nanoTime.getAsLong(); + for (int batches = 0; batches < MAX_BATCHES_PER_TASK; batches++) { + more = produceBatch(); + if (!more || nanoTime.getAsLong() - start >= MAX_TASK_NANOS) { + break; + } + } } catch (Exception e) { complete(e); } @@ -116,7 +155,7 @@ public void run() { return; } } - submit(); + submit(true); } private boolean produceBatch() throws Exception { @@ -126,7 +165,8 @@ private boolean produceBatch() throws Exception { } if (pendingRanges == null && !pending.hasNext()) { List batch = source.getNextBatch(); - if (batch.isEmpty()) { + // Cancellation may arrive while the source is listing files. Do not assign that batch. + if (batch.isEmpty() || cancelled || !assignment.needMoreSplit()) { complete(null); return false; } @@ -160,7 +200,13 @@ private void complete(Exception failure) { } completed = true; } - assignment.removeQueueSpaceListener(pendingBackend, resume); + if (executor instanceof ConnectorTaskScheduler) { + ((ConnectorTaskScheduler) executor).removeSpaceListener(retryAdmission); + } + // Cancellation may complete before this producer has assigned its first backend. + if (pendingBackend != null) { + assignment.removeQueueSpaceListener(pendingBackend, resume); + } assignment.removeCloseable(this); try { completion.accept(failure); 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 b1be85133b26b6..1ca7614242a6a1 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 @@ -45,16 +45,23 @@ */ public class SplitSource { private static final AtomicLong UNIQUE_ID_GENERATOR = new AtomicLong(0); - private static final long WAIT_TIME_OUT = 100; // 100ms private final long uniqueId; private final Backend backend; private final SplitAssignment splitAssignment; private final AtomicBoolean isLastBatch; + private final AtomicBoolean firstBatchRequested = new AtomicBoolean(); private final long maxWaitTime; + private final int initialBatchSize; private Iterator pendingSplits = Collections.emptyIterator(); public SplitSource(Backend backend, SplitAssignment splitAssignment, long maxWaitTime) { + this(backend, splitAssignment, maxWaitTime, Integer.MAX_VALUE); + } + + public SplitSource(Backend backend, SplitAssignment splitAssignment, long maxWaitTime, int initialBatchSize) { + Preconditions.checkArgument(initialBatchSize > 0, "Initial split batch size must be positive"); + this.initialBatchSize = initialBatchSize; this.uniqueId = UNIQUE_ID_GENERATOR.getAndIncrement(); this.backend = backend; this.splitAssignment = splitAssignment; @@ -75,6 +82,12 @@ public synchronized List getNextBatch(int maxBatchSize) thr if (isLastBatch.get()) { return Collections.emptyList(); } + if (firstBatchRequested.compareAndSet(false, true)) { + // Limit startup RPC work independently of producer batches and buffered lookahead. + maxBatchSize = Math.min(maxBatchSize, initialBatchSize); + } else { + splitAssignment.onMoreSplitsRequested(); + } List scanRanges = new ArrayList<>( Math.min(maxBatchSize, splitAssignment.getSplitBatchSize())); long startTime = System.nanoTime(); @@ -91,7 +104,7 @@ public synchronized List getNextBatch(int maxBatchSize) thr scanRanges.add(pendingSplits.next()); } if (scanRanges.size() == maxBatchSize) { - return scanRanges; + return FileScanRangeCompactor.compactSplitBatch(scanRanges); } pendingSplits = Collections.emptyIterator(); // A finished, empty queue is EOF only after the pending batch has been drained. @@ -99,6 +112,27 @@ public synchronized List getNextBatch(int maxBatchSize) thr isLastBatch.set(true); break; } + // Each queued batch contains at most splitBatchSize ranges. Drain only batches + // that fit this response, bounding the local buffer and amortizing the queue lock. + int batchesToDrain = Math.min(64, + (maxBatchSize - scanRanges.size()) / splitAssignment.getSplitBatchSize()); + if (batchesToDrain > 1 && splits.size() > 1) { + List> readyBatches = new ArrayList<>(batchesToDrain); + splits.drainTo(readyBatches, batchesToDrain); + splitAssignment.notifyQueueSpace(backend); + // Resuming a producer may publish an error or cancellation before we return. + splitAssignment.getAssignedSplits(backend); + if (splitAssignment.isStop()) { + isLastBatch.set(true); + break; + } + for (Collection readyBatch : readyBatches) { + for (TScanRangeLocations range : readyBatch) { + scanRanges.add(range); + } + } + continue; + } Collection batch = splits.poll(); if (batch != null) { splitAssignment.notifyQueueSpace(backend); @@ -113,8 +147,7 @@ public synchronized List getNextBatch(int maxBatchSize) thr throw new UserException("Timeout. Max wait time(ms): " + maxWaitTime); } try { - batch = splits.poll(Math.min(remainingNanos, TimeUnit.MILLISECONDS.toNanos(WAIT_TIME_OUT)), - TimeUnit.NANOSECONDS); + batch = splitAssignment.awaitBatch(backend, remainingNanos); if (batch != null) { splitAssignment.notifyQueueSpace(backend); pendingSplits = batch.iterator(); @@ -124,6 +157,6 @@ public synchronized List getNextBatch(int maxBatchSize) thr throw new UserException(e.getMessage(), e); } } - return scanRanges; + return FileScanRangeCompactor.compactSplitBatch(scanRanges); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java index 80bd559d339b6a..898f1bc12fa99c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java @@ -882,6 +882,13 @@ private synchronized void endConnectorPlanning() { } } + /** Closes a scope created during this attempt even if no scan registered a query-finish callback. */ + public synchronized void closeConnectorStatementScope() { + if (this.connectorStatementScope != null) { + this.connectorStatementScope.closeAll(); + } + } + /** * Closes (deterministically) then drops the connector scope so the next statement execution starts fresh. * Prepared-statement EXECUTE reuses one StatementContext across executions (see @@ -892,9 +899,7 @@ private synchronized void endConnectorPlanning() { * execution/attempt has finished (no off-thread pump still running), so close-before-drop is safe. */ public synchronized void resetConnectorStatementScope() { - if (this.connectorStatementScope != null) { - this.connectorStatementScope.closeAll(); - } + closeConnectorStatementScope(); this.connectorStatementScope = null; this.connectorWriteSchemas.clear(); this.connectorWriteMetadataIdentities.clear(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java index 99521b61ad0921..f0f598ea3fc2c3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java @@ -38,6 +38,7 @@ import org.apache.doris.common.util.ListUtil; import org.apache.doris.datasource.scan.ExternalScanNode; import org.apache.doris.datasource.scan.FileQueryScanNode; +import org.apache.doris.datasource.split.FileScanRangeCompactor; import org.apache.doris.load.loadv2.LoadJob; import org.apache.doris.metric.MetricRepo; import org.apache.doris.nereids.NereidsPlanner; @@ -3663,7 +3664,8 @@ Map toThrift(int backendNum) { TPipelineInstanceParams localParams = new TPipelineInstanceParams(); localParams.setFragmentInstanceId(instanceExecParam.instanceId); - localParams.setPerNodeScanRanges(scanRanges); + localParams.setPerNodeScanRanges(FileScanRangeCompactor.compactInstanceScanRanges( + scanRanges, fileScanRangeParamsMap, queryOptions)); localParams.setSenderId(i); localParams.setBackendNum(backendNum++); localParams.setRuntimeFilterParams(new TRuntimeFilterParams()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessor.java index 39a56deed6977b..4d95da465f8ba9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessor.java @@ -45,6 +45,13 @@ public interface QeProcessor { */ void registerQueryFinishCallback(String queryId, Runnable callback); + /** + * Release resources of a finished attempt, including planning-only attempts without a coordinator. + * Coordinator-backed queries also call this from unregisterQuery; repeated calls are harmless. + */ + void finishQueryCallbacks(TUniqueId queryId); + + Map getQueryStatistics(); String getCurrentQueryByQueryId(TUniqueId queryId); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java index 77a4e0b4538657..babfb4b2fcbe4d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java @@ -219,6 +219,11 @@ public void unregisterQuery(TUniqueId queryId) { // Run connector-registered query-finish callbacks (e.g. committing a hive // read transaction). Connector-agnostic: fe-core names no source here. + finishQueryCallbacks(queryId); + } + + @Override + public void finishQueryCallbacks(TUniqueId queryId) { queryFinishCallbackRegistry.runAndClear(DebugUtil.printId(queryId)); } 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 995dbe86223c4a..3b608d336f5a8b 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 @@ -592,11 +592,14 @@ public String toString() { public static final String ENABLE_ICEBERG_INTERNAL_PARTITION_PRUNING = "enable_iceberg_internal_partition_pruning"; + public static final String ICEBERG_MANIFEST_LOAD_CONCURRENCY = "iceberg_manifest_load_concurrency"; + public static final String NUM_PARTITIONS_IN_BATCH_MODE = "num_partitions_in_batch_mode"; public static final String NUM_FILES_IN_BATCH_MODE = "num_files_in_batch_mode"; public static final String SPLIT_BATCH_SIZE = "split_batch_size"; + public static final String SPLIT_INITIAL_BATCH_SIZE = "split_initial_batch_size"; public static final String SPLIT_QUEUE_CAPACITY = "split_queue_capacity"; @@ -2636,6 +2639,15 @@ public Map getForceEagerAggHintMap() { description = "Maximum number of queued split batches per backend.") public int splitQueueCapacity = 32; + @VarAttrDef.VarAttr(name = SPLIT_INITIAL_BATCH_SIZE, checker = "checkSplitInitialBatchSize", needForward = true, + description = "Maximum number of splits returned by the first fetch per backend; " + + "does not limit background split production or subsequent fetches.") + public int splitInitialBatchSize = 64; + + public void checkSplitInitialBatchSize(String value) throws Exception { + checkFieldValue(SPLIT_INITIAL_BATCH_SIZE, 1, value); + } + public void checkSplitBatchSize(String value) throws Exception { checkFieldValue(SPLIT_BATCH_SIZE, 1, value); } @@ -3737,6 +3749,19 @@ public void checkAnnIndexCandidateRowsPercentThreshold(String value) { public static final String ICEBERG_PARTITION_OVERWRITE_MODE = "iceberg_partition_overwrite_mode"; public static final String ICEBERG_PARTITION_OVERWRITE_MODE_DYNAMIC = "dynamic"; public static final String ICEBERG_PARTITION_OVERWRITE_MODE_STATIC = "static"; + @VarAttrDef.VarAttr(name = ICEBERG_MANIFEST_LOAD_CONCURRENCY, needForward = true, + checker = "checkIcebergManifestLoadConcurrency", + description = "Maximum concurrent Iceberg manifest-cache reads per query and prefetch window per scan. " + + "Set to 1 for sequential loading.") + public int icebergManifestLoadConcurrency = 8; + + public void checkIcebergManifestLoadConcurrency(String value) { + int concurrency = Integer.parseInt(value); + if (concurrency < 1 || concurrency > 64) { + throw new IllegalArgumentException("iceberg_manifest_load_concurrency must be between 1 and 64"); + } + } + @VarAttrDef.VarAttr(name = ICEBERG_PARTITION_OVERWRITE_MODE, needForward = true, checker = "checkIcebergPartitionOverwriteMode", setter = "setIcebergPartitionOverwriteMode", description = "Scope of INSERT OVERWRITE on an Iceberg table when no PARTITION clause is given. " @@ -5001,6 +5026,10 @@ public int getSplitBatchSize() { return splitBatchSize; } + public int getSplitInitialBatchSize() { + return splitInitialBatchSize; + } + public int getSplitQueueCapacity() { return splitQueueCapacity; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index d317da22da3be4..054b8904799e07 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -166,6 +166,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; @@ -247,6 +248,7 @@ private AuditBackendSnapshot(TUniqueId queryId, Set backendIds) { // queries) register it here so Ctrl+C / KILL QUERY / statement timeout on the outer // statement also cancel the inner work, which those paths cannot address directly. private volatile Consumer cancelDelegate = null; + private final List queryCancelCallbacks = new CopyOnWriteArrayList<>(); private RedirectStatus redirectStatus = null; private Planner planner; private boolean isProxy; @@ -769,6 +771,11 @@ private void executeWithReplanRetries(TUniqueId queryId) throws Exception { if (finishProfileInQueryRetry) { profile.clearExecutionProfiles(); } + // QueryFinishCleanup closed the failed attempt's scope. Replanning must not register + // fresh task groups in that closed scope, whose closeAll() is already complete. + if (context.getStatementContext() != null) { + context.getStatementContext().resetConnectorStatementScope(); + } // Drop the old execution state before replanning. planner = null; setCoord(null); @@ -887,6 +894,22 @@ public void execute(TUniqueId queryId) throws Exception { } finally { sessionVariable.setIsSingleSetVar(false); sessionVariable.clearSessionOriginValue(); + // Planning can register scan cleanup before a coordinator exists. EXPLAIN and + // planning errors never reach unregisterQuery(), but still finish this attempt. + // A deferred Flight query owns its resources until its later finalization. + if (!deferredForArrowFlight) { + try { + QeProcessorImpl.INSTANCE.finishQueryCallbacks(queryId); + } finally { + try { + if (context.getStatementContext() != null) { + context.getStatementContext().closeConnectorStatementScope(); + } + } finally { + queryCancelCallbacks.clear(); + } + } + } } } } @@ -1244,6 +1267,7 @@ void finalizeQuery(boolean willRetry) { updateProfile(true); } QeProcessorImpl.INSTANCE.unregisterQuery(queryId()); + queryCancelCallbacks.clear(); } // The id this query runs under: the context's, until the query is deferred for Arrow Flight and @@ -1589,8 +1613,31 @@ public void clearCancelDelegate() { this.cancelDelegate = null; } + /** Register before starting asynchronous query work. A concurrent or earlier cancellation runs it once. */ + public void registerQueryCancellationCallback(Runnable callback) { + queryCancelCallbacks.add(callback); + if (pendingCoordinatorCancelReason.get() != null && queryCancelCallbacks.remove(callback)) { + runQueryCancellationCallback(callback); + } + } + + private void runQueryCancellationCallback(Runnable callback) { + try { + callback.run(); + } catch (Exception e) { + LOG.warn("Query cancellation callback failed", e); + } + } + public void cancel(Status cancelReason, boolean needWaitCancelComplete) { pendingCoordinatorCancelReason.compareAndSet(null, cancelReason); + // Query-owned work may already exist during planning, before there is a coordinator to cancel. + // Remove before invoking: registration and repeated cancellation must not run it twice. + for (Runnable callback : queryCancelCallbacks) { + if (queryCancelCallbacks.remove(callback)) { + runQueryCancellationCallback(callback); + } + } Status coordinatorCancelReason = pendingCoordinatorCancelReason.get(); Consumer delegate = cancelDelegate; if (delegate != null) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java index 64813f1b7b58e9..e36b0978827f10 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java @@ -24,6 +24,7 @@ import org.apache.doris.catalog.Resource; import org.apache.doris.common.Config; import org.apache.doris.datasource.scan.FileQueryScanNode; +import org.apache.doris.datasource.split.FileScanRangeCompactor; import org.apache.doris.nereids.StatementContext; import org.apache.doris.nereids.trees.plans.distribute.DistributedPlan; import org.apache.doris.nereids.trees.plans.distribute.PipelineDistributedPlan; @@ -578,7 +579,9 @@ private static void setScanSourceParam( scanParams = computeDefaultScanSourceParam((DefaultScanSource) scanSource); } // perNodeScanRanges is required - instanceParams.setPerNodeScanRanges(scanParams.perNodeScanRanges); + instanceParams.setPerNodeScanRanges(FileScanRangeCompactor.compactInstanceScanRanges( + scanParams.perNodeScanRanges, currentFragmentParam.getFileScanParams(), + currentFragmentParam.getQueryOptions())); } // local shuffle has two functions: diff --git a/fe/fe-core/src/test/java/org/apache/doris/common/ThreadPoolManagerTest.java b/fe/fe-core/src/test/java/org/apache/doris/common/ThreadPoolManagerTest.java index a00c37d880509f..22c1c075751bce 100755 --- a/fe/fe-core/src/test/java/org/apache/doris/common/ThreadPoolManagerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/common/ThreadPoolManagerTest.java @@ -24,6 +24,49 @@ public class ThreadPoolManagerTest { + @Test + public void callerHelpDoesNotBlockOnFullPoolAndExecutesExactlyOnce() throws Exception { + ThreadPoolExecutor pool = ThreadPoolManager.newDaemonFixedThreadPool(1, 1, "help_full", false); + java.util.concurrent.atomic.AtomicInteger calls = new java.util.concurrent.atomic.AtomicInteger(); + try { + pool.submit(() -> { + pool.execute(() -> { }); // fill the one queue slot while this parent owns the worker + java.util.concurrent.RunnableFuture task = + ThreadPoolManager.submitWithCallerHelp(pool, calls::incrementAndGet); + task.run(); + task.run(); + try { + Assertions.assertEquals(1, task.get().intValue()); + } catch (Exception e) { + throw new RuntimeException(e); + } + }).get(5, java.util.concurrent.TimeUnit.SECONDS); + Assertions.assertEquals(1, calls.get()); + } finally { + pool.shutdownNow(); + } + } + + @Test + public void helpedAndCancelledTasksDoNotAccumulateInTheSharedQueue() throws Exception { + ThreadPoolExecutor pool = ThreadPoolManager.newDaemonFixedThreadPool(1, 1, "help_remove", false); + try { + pool.submit(() -> { + java.util.concurrent.RunnableFuture task = + ThreadPoolManager.submitWithCallerHelp(pool, () -> 42); + Assertions.assertEquals(1, pool.getQueue().size()); + task.run(); + Assertions.assertTrue(pool.getQueue().isEmpty()); + java.util.concurrent.RunnableFuture unused = + ThreadPoolManager.submitWithCallerHelp(pool, () -> 43); + unused.cancel(true); + Assertions.assertTrue(pool.getQueue().isEmpty()); + }).get(5, java.util.concurrent.TimeUnit.SECONDS); + } finally { + pool.shutdownNow(); + } + } + @Test public void testNormal() throws InterruptedException { ThreadPoolExecutor testCachedPool = ThreadPoolManager.newDaemonCacheThreadPool(2, "test_cache_pool", true); diff --git a/fe/fe-core/src/test/java/org/apache/doris/common/util/FairTaskExecutorTest.java b/fe/fe-core/src/test/java/org/apache/doris/common/util/FairTaskExecutorTest.java new file mode 100644 index 00000000000000..38cbff04c7e55b --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/common/util/FairTaskExecutorTest.java @@ -0,0 +1,570 @@ +// 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.common.util; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +@Timeout(30) +class FairTaskExecutorTest { + private FairTaskExecutor executor(int workers) { + return new FairTaskExecutor(workers, task -> { + Thread thread = new Thread(task); + thread.setDaemon(true); + return thread; + }); + } + + private static void await(CountDownLatch latch) { + try { + Assertions.assertTrue(latch.await(10, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + throw new AssertionError(e); + } + } + + @Test + void groupsRoundRobinInsteadOfDrainingTheFirstQuery() throws Exception { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group first = pool.newGroup(100); + FairTaskExecutor.Group second = pool.newGroup(100); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + first.execute(() -> { + started.countDown(); + await(release); + }); + await(started); + List order = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch done = new CountDownLatch(20); + for (int i = 0; i < 10; i++) { + first.execute(() -> { + order.add("A"); + done.countDown(); + }); + second.execute(() -> { + order.add("B"); + done.countDown(); + }); + } + release.countDown(); + await(done); + for (int i = 0; i < 20; i++) { + Assertions.assertEquals(i % 2 == 0 ? "A" : "B", order.get(i)); + } + } + } + + @Test + void oneGroupCanUseAllWorkers() { + try (FairTaskExecutor pool = executor(4)) { + FairTaskExecutor.Group group = pool.newGroup(4); + CountDownLatch started = new CountDownLatch(4); + CountDownLatch release = new CountDownLatch(1); + for (int i = 0; i < 4; i++) { + group.execute(() -> { + started.countDown(); + await(release); + }); + } + try { + await(started); + } finally { + release.countDown(); + } + } + } + + @Test + void independentSingleWorkerPoolsCannotStarveChildren() throws Exception { + try (FairTaskExecutor parents = executor(1); FairTaskExecutor children = executor(1)) { + FairTaskExecutor.Group parent = parents.newGroup(1); + FairTaskExecutor.Group child = children.newGroup(1); + CompletableFuture result = new CompletableFuture<>(); + parent.execute(() -> { + CompletableFuture value = new CompletableFuture<>(); + child.tryExecute(() -> value.complete(42), value::completeExceptionally, null); + result.complete(value.join()); + }); + Assertions.assertEquals(42, result.get(10, TimeUnit.SECONDS)); + } + } + + @Test + void fullQueueResumesWithoutLosingWakeupsOrDuplicatingTasks() { + try (FairTaskExecutor pool = executor(4)) { + FairTaskExecutor.Group group = pool.newGroup(1); + AtomicReference failure = new AtomicReference<>(); + int lanes = 12; + int rounds = 2000; + CountDownLatch done = new CountDownLatch(lanes); + AtomicInteger executions = new AtomicInteger(); + class Lane implements Runnable { + int remaining = rounds; + final Runnable admission = this::submit; + + void submit() { + group.tryExecute(this, failure::set, admission); + } + + @Override + public void run() { + executions.incrementAndGet(); + if (--remaining == 0) { + done.countDown(); + } else { + submit(); + } + } + } + + for (int i = 0; i < lanes; i++) { + new Lane().submit(); + } + await(done); + Assertions.assertNull(failure.get()); + Assertions.assertEquals(lanes * rounds, executions.get()); + Assertions.assertEquals(0, group.queuedTasks()); + } + } + + @Test + void closeRejectsQueuedWorkAndWakesCapacityWaitersExactlyOnce() { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group group = pool.newGroup(1); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + group.execute(() -> { + started.countDown(); + await(release); + }); + await(started); + AtomicInteger rejected = new AtomicInteger(); + AtomicInteger resumed = new AtomicInteger(); + Assertions.assertTrue(group.tryExecute(() -> Assertions.fail("cancelled task ran"), + failure -> rejected.incrementAndGet(), null)); + Assertions.assertFalse(group.tryExecute(() -> Assertions.fail("unaccepted task ran"), + failure -> Assertions.fail(failure), resumed::incrementAndGet)); + try { + group.close(); + group.close(); + Assertions.assertEquals(1, rejected.get()); + Assertions.assertEquals(1, resumed.get()); + Assertions.assertEquals(0, group.queuedTasks()); + Assertions.assertThrows(java.util.concurrent.RejectedExecutionException.class, + () -> group.execute(() -> { })); + } finally { + release.countDown(); + } + } + } + + @Test + void admissionRacingCloseHasExactlyOneTerminalOutcome() throws Exception { + try (FairTaskExecutor pool = executor(4)) { + for (int round = 0; round < 100; round++) { + FairTaskExecutor.Group group = pool.newGroup(8); + AtomicInteger accepted = new AtomicInteger(); + AtomicInteger finished = new AtomicInteger(); + AtomicReference failure = new AtomicReference<>(); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch admitted = new CountDownLatch(1); + List submitters = new ArrayList<>(); + for (int lane = 0; lane < 4; lane++) { + Thread submitter = new Thread(() -> { + await(start); + for (int task = 0; task < 100; task++) { + java.util.concurrent.atomic.AtomicBoolean terminal = + new java.util.concurrent.atomic.AtomicBoolean(); + Runnable complete = () -> { + if (!terminal.compareAndSet(false, true)) { + failure.set(new AssertionError("task completed twice")); + } + finished.incrementAndGet(); + }; + try { + if (group.tryExecute(complete, error -> complete.run(), null)) { + accepted.incrementAndGet(); + admitted.countDown(); + } + } catch (java.util.concurrent.RejectedExecutionException e) { + return; + } + } + }); + submitter.start(); + submitters.add(submitter); + } + start.countDown(); + await(admitted); + group.close(); + for (Thread submitter : submitters) { + submitter.join(5000); + Assertions.assertFalse(submitter.isAlive()); + } + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (finished.get() != accepted.get() && System.nanoTime() < deadline) { + Thread.yield(); + } + Assertions.assertNull(failure.get()); + Assertions.assertEquals(accepted.get(), finished.get()); + Assertions.assertEquals(0, group.queuedTasks()); + } + } + } + + @Test + void failureIsDeliveredAndDoesNotRemoveTheWorker() throws Exception { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group group = pool.newGroup(2); + CompletableFuture failure = new CompletableFuture<>(); + AssertionError expected = new AssertionError("task failure"); + group.tryExecute(() -> { + throw expected; + }, failure::complete, null); + Assertions.assertSame(expected, failure.get(10, TimeUnit.SECONDS)); + CompletableFuture next = new CompletableFuture<>(); + group.execute(() -> next.complete(true)); + Assertions.assertTrue(next.get(10, TimeUnit.SECONDS)); + } + } + + @Test + void closeCompletesRunningFuturesWithoutInterruptingSharedIo() throws Exception { + FairTaskExecutor pool = executor(1); + FairTaskExecutor.Group group = pool.newGroup(1); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + FutureTask task = new FutureTask<>(() -> { + started.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + failure.set(e); + } + return 1; + }); + try { + group.execute(task); + await(started); + group.close(); + Assertions.assertThrows(java.util.concurrent.CancellationException.class, + () -> task.get(1, TimeUnit.SECONDS)); + Assertions.assertNull(failure.get()); + pool.close(); + Assertions.assertFalse(pool.awaitTermination(1, TimeUnit.MILLISECONDS), "shared IO is still running"); + } finally { + release.countDown(); + pool.close(); + Assertions.assertTrue(pool.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + void idleWorkersWakeForRepeatedIsolatedArrivals() throws Exception { + try (FairTaskExecutor pool = executor(4)) { + FairTaskExecutor.Group group = pool.newGroup(1); + for (int i = 0; i < 1000; i++) { + CompletableFuture result = new CompletableFuture<>(); + group.execute(() -> result.complete(true)); + Assertions.assertTrue(result.get(10, TimeUnit.SECONDS)); + } + } + } + + @Test + void serialContinuationsRetainTheWorkerWithoutRunningConcurrently() throws Exception { + List workers = new ArrayList<>(); + try (FairTaskExecutor pool = new FairTaskExecutor(4, task -> { + Thread thread = new Thread(task); + thread.setDaemon(true); + workers.add(thread); + return thread; + })) { + FairTaskExecutor.Group group = pool.newGroup(1); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CompletableFuture done = new CompletableFuture<>(); + AtomicReference owner = new AtomicReference<>(); + AtomicInteger runs = new AtomicInteger(); + AtomicInteger active = new AtomicInteger(); + Runnable task = new Runnable() { + @Override + public void run() { + Assertions.assertEquals(1, active.incrementAndGet()); + try { + if (runs.getAndIncrement() == 0) { + owner.set(Thread.currentThread()); + started.countDown(); + await(release); + } + Assertions.assertSame(owner.get(), Thread.currentThread()); + } finally { + active.decrementAndGet(); + } + if (runs.get() == 1000) { + done.complete(runs.get()); + } else { + Assertions.assertTrue(group.tryExecuteContinuation(this, done::completeExceptionally, null)); + } + } + }; + group.tryExecute(task, done::completeExceptionally, null); + await(started); + try { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (workers.stream().anyMatch(thread -> thread != owner.get() + && thread.getState() != Thread.State.WAITING) && System.nanoTime() < deadline) { + Thread.yield(); + } + Assertions.assertTrue(workers.stream().allMatch(thread -> thread == owner.get() + || thread.getState() == Thread.State.WAITING)); + } finally { + release.countDown(); + } + Assertions.assertEquals(1000, done.get(10, TimeUnit.SECONDS)); + } + } + + @Test + void continuationGoesBehindOtherReadyQueries() throws Exception { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group first = pool.newGroup(1); + FairTaskExecutor.Group second = pool.newGroup(1); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CompletableFuture done = new CompletableFuture<>(); + List order = new ArrayList<>(); + first.tryExecute(new Runnable() { + int runs; + + @Override + public void run() { + order.add("A"); + if (++runs == 1) { + started.countDown(); + await(release); + first.tryExecuteContinuation(this, done::completeExceptionally, null); + } else { + done.complete(true); + } + } + }, done::completeExceptionally, null); + await(started); + second.execute(() -> order.add("B")); + release.countDown(); + Assertions.assertTrue(done.get(10, TimeUnit.SECONDS)); + Assertions.assertEquals(java.util.Arrays.asList("A", "B", "A"), order); + } + } + + @Test + void continuationAdmissionFromAnotherPoolStillWakesAWorker() throws Exception { + try (FairTaskExecutor parents = executor(1); FairTaskExecutor children = executor(1)) { + FairTaskExecutor.Group parent = parents.newGroup(1); + FairTaskExecutor.Group child = children.newGroup(1); + CompletableFuture result = new CompletableFuture<>(); + parent.execute(() -> { + CompletableFuture value = new CompletableFuture<>(); + child.tryExecuteContinuation(() -> value.complete(42), value::completeExceptionally, null); + result.complete(value.join()); + }); + Assertions.assertEquals(42, result.get(10, TimeUnit.SECONDS)); + } + } + + @Test + void fullQueueContinuationResumesExactlyOnce() throws Exception { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group group = pool.newGroup(1); + CompletableFuture done = new CompletableFuture<>(); + AtomicInteger resumed = new AtomicInteger(); + group.tryExecute(new Runnable() { + int runs; + final Runnable retry = () -> { + resumed.incrementAndGet(); + group.tryExecute(this, done::completeExceptionally, null); + }; + + @Override + public void run() { + if (++runs == 1) { + group.execute(() -> { }); + Assertions.assertFalse(group.tryExecuteContinuation(this, done::completeExceptionally, retry)); + } else { + done.complete(runs); + } + } + }, done::completeExceptionally, null); + Assertions.assertEquals(2, done.get(10, TimeUnit.SECONDS)); + Assertions.assertEquals(1, resumed.get()); + Assertions.assertEquals(0, group.queuedTasks()); + } + } + + @Test + void closeRejectsAnAcceptedContinuationExactlyOnce() throws Exception { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group group = pool.newGroup(1); + CountDownLatch accepted = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger runs = new AtomicInteger(); + AtomicInteger rejected = new AtomicInteger(); + group.execute(new Runnable() { + @Override + public void run() { + runs.incrementAndGet(); + group.tryExecuteContinuation(this, failure -> rejected.incrementAndGet(), null); + accepted.countDown(); + // Simulate preemption between the final submission and return. + await(release); + } + }); + await(accepted); + try { + group.close(); + group.close(); + Assertions.assertEquals(1, rejected.get()); + Assertions.assertEquals(1, runs.get()); + Assertions.assertEquals(0, group.queuedTasks()); + } finally { + release.countDown(); + } + } + } + + + @Test + void manyGroupsWithBoundedContinuationsMakeProgress() { + try (FairTaskExecutor pool = executor(4)) { + int groups = 8; + int lanes = 4; + int rounds = 1000; + CountDownLatch done = new CountDownLatch(groups * lanes); + AtomicReference failure = new AtomicReference<>(); + AtomicInteger executions = new AtomicInteger(); + for (int i = 0; i < groups; i++) { + FairTaskExecutor.Group group = pool.newGroup(1); + class Lane implements Runnable { + int remaining = rounds; + final Runnable retry = this::submit; + + void submit() { + group.tryExecute(this, failure::set, retry); + } + + @Override + public void run() { + executions.incrementAndGet(); + if (--remaining == 0) { + done.countDown(); + } else { + group.tryExecuteContinuation(this, failure::set, retry); + } + } + } + + for (int lane = 0; lane < lanes; lane++) { + Lane task = new Lane(); + group.tryExecute(task, failure::set, task.retry); + } + } + await(done); + Assertions.assertNull(failure.get()); + Assertions.assertEquals(groups * lanes * rounds, executions.get()); + } + } + + private static final class AsyncTask extends CompletableFuture implements Runnable { + private final CompletableFuture shared; + private final AtomicInteger cancellations = new AtomicInteger(); + + private AsyncTask(CompletableFuture shared) { + this.shared = shared; + } + + @Override + public void run() { + shared.whenComplete((value, failure) -> { + if (failure == null) { + complete(value); + } else { + completeExceptionally(failure); + } + }); + } + + @Override + public boolean cancel(boolean interrupt) { + cancellations.incrementAndGet(); + return super.cancel(interrupt); + } + } + + @Test + void suspendedFutureReleasesWorkerButRemainsOwnedByItsQuery() throws Exception { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group first = pool.newGroup(2); + FairTaskExecutor.Group second = pool.newGroup(2); + CompletableFuture shared = new CompletableFuture<>(); + AsyncTask task = new AsyncTask(shared); + first.execute(task); + CompletableFuture marker = new CompletableFuture<>(); + // The sole worker must leave the suspended first query and run another query. + second.execute(() -> marker.complete(true)); + Assertions.assertTrue(marker.get(5, TimeUnit.SECONDS)); + Assertions.assertFalse(task.isDone()); + first.close(); + Assertions.assertTrue(task.isCancelled(), "returned asynchronous work must still be cancelled"); + Assertions.assertEquals(1, task.cancellations.get()); + Assertions.assertFalse(shared.isDone(), "cancelling a query must not cancel its shared load"); + shared.complete(42); + Assertions.assertTrue(task.isCancelled()); + } + } + + @Test + void completedAsyncTaskIsNotRetainedUntilQueryClose() throws Exception { + try (FairTaskExecutor pool = executor(1)) { + FairTaskExecutor.Group group = pool.newGroup(2); + CompletableFuture shared = new CompletableFuture<>(); + AsyncTask task = new AsyncTask(shared); + group.execute(task); + CompletableFuture returned = new CompletableFuture<>(); + group.execute(() -> returned.complete(true)); + Assertions.assertTrue(returned.get(5, TimeUnit.SECONDS)); + shared.complete(42); + Assertions.assertEquals(42, task.get(5, TimeUnit.SECONDS)); + group.close(); + Assertions.assertEquals(0, task.cancellations.get(), "completed tasks must leave the ownership set"); + } + } + +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/common/util/LocationPathParsingTest.java b/fe/fe-core/src/test/java/org/apache/doris/common/util/LocationPathParsingTest.java new file mode 100644 index 00000000000000..24abe5f7551c5b --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/common/util/LocationPathParsingTest.java @@ -0,0 +1,105 @@ +// 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.common.util; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.net.URI; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.util.Random; + +class LocationPathParsingTest { + @Test + void preserveFileSystemIdentityAndOpaquePath() { + String[] locations = { + "s3://bucket/a//b/../c/", "s3://bucket/a b+%2F?#中文😀.parquet", + "hdfs://namenode:8020/p=1/file", "hdfs:///path", "hdfs:/path", "file:///tmp/file", + "abfss://container@account.dfs.core.windows.net/dir/file", + "abfss://workspace@onelake.dfs.fabric.microsoft.com/lake/Files/file", + "oss://bucket.oss-dls.aliyuncs.com/file", "s3://bucket", "s3://bucket/", + "s3://bucket name/file", "s3://bucket+name/file", "s3://bucket%20name/file", + "hdfs://[::1]:8020/file", "connector://virtual", "S3://bucket/file", + "/local/file", "relative/file", "//host/file", "a/b://c/d", "a:b/c", + "s3:///", "s3:////a", "s3://bucket/?#", "s3://bucket/\u0000file", "", + "s3://", "1s3://bucket/file", "s+3://bucket/file", "://bucket/file" + }; + for (String location : locations) { + assertEquivalent(location); + } + } + + @Test + void preserveIdentityForSpecialCharactersInAuthorityAndPath() { + Random random = new Random(20261014); + String alphabet = "abcXYZ019 /:+%?#@[]*._-中文é\u0000" + + Character.MIN_HIGH_SURROGATE + Character.MIN_LOW_SURROGATE; + for (int i = 0; i < 2000; i++) { + StringBuilder location = new StringBuilder(i % 2 == 0 ? "s3://" : "hdfs://"); + for (int j = 0; j < 8; j++) { + location.append(alphabet.charAt(random.nextInt(alphabet.length()))); + } + location.append('/'); + for (int j = 0; j < 40; j++) { + location.append(alphabet.charAt(random.nextInt(alphabet.length()))); + } + assertEquivalent(location.toString()); + } + } + + @Test + void reuseFilesystemIdentityAcrossFilesWithoutChangingPathBytes() { + LocationPath.Factory factory = new LocationPath.Factory(); + LocationPath first = factory.create("s3://bucket/a"); + LocationPath second = factory.create("s3://bucket/b //中文%2F+?#"); + Assertions.assertSame(first.getFsIdentifier(), second.getFsIdentifier()); + Assertions.assertEquals("s3://bucket/b //中文%2F+?#", second.getNormalizedLocation()); + Assertions.assertSame(second, factory.create(second.getNormalizedLocation())); + for (String location : new String[] {"s3://bucket-other/a", "hdfs://bucket/a", "hdfs:///a", + "hdfs:///b", "s3://bucket name/a", "s3://bucket name/b", "hdfs:/a", "hdfs:/b", + "file:///tmp/a", "file:///tmp/b", "/local/a", "/local/b", "s3://bucket", "s3://bucket/a"}) { + LocationPath expected = LocationPath.of(location); + LocationPath actual = factory.create(location); + Assertions.assertEquals(expected.getFsIdentifier(), actual.getFsIdentifier(), location); + Assertions.assertEquals(expected.getSchema(), actual.getSchema(), location); + Assertions.assertEquals(expected.getTFileTypeForBE(), actual.getTFileTypeForBE(), location); + Assertions.assertEquals(location, actual.getNormalizedLocation()); + } + Assertions.assertThrows(IllegalArgumentException.class, () -> factory.create("1s3://bucket/file")); + Assertions.assertEquals("s3://bucket", factory.create("s3://bucket/after-error").getFsIdentifier()); + } + + private static void assertEquivalent(String location) { + // The old full-path parser is the compatibility oracle: URL metacharacters in an object + // key are literal bytes, and authority escaping (including spaces becoming '+') is unchanged. + URI expected; + try { + expected = URI.create(URLEncoder.encode(location, StandardCharsets.UTF_8) + .replace("%2F", "/").replace("%3A", ":")); + } catch (IllegalArgumentException e) { + Assertions.assertThrows(IllegalArgumentException.class, () -> LocationPath.of(location), location); + return; + } + LocationPath actual = LocationPath.of(location); + Assertions.assertEquals((expected.getScheme() == null ? "" : expected.getScheme()) + "://" + + (expected.getAuthority() == null ? "" : expected.getAuthority()), actual.getFsIdentifier(), location); + Assertions.assertEquals(location, actual.getNormalizedLocation()); + Assertions.assertEquals(location, actual.toStorageLocation().toString()); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/connector/QueryTaskGroupsTest.java b/fe/fe-core/src/test/java/org/apache/doris/connector/QueryTaskGroupsTest.java new file mode 100644 index 00000000000000..c9c7ac67d657f1 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/connector/QueryTaskGroupsTest.java @@ -0,0 +1,76 @@ +// 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; + +import org.apache.doris.common.util.FairTaskExecutor; +import org.apache.doris.connector.spi.ConnectorSession; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CancellationException; +import java.util.concurrent.RejectedExecutionException; + +class QueryTaskGroupsTest { + @Test + void cancellationBeforeFirstSchedulerAcquisitionDoesNotCreateAGroup() { + ConnectorStatementScopeImpl scope = new ConnectorStatementScopeImpl(); + scope.computeIfAbsent(QueryTaskGroups.SCOPE_KEY, () -> new QueryTaskGroups(stage -> { + throw new AssertionError("a cancelled query must not create a task group"); + })); + ConnectorSession session = ConnectorSessionBuilder.create().withQueryId("cancel-before-acquisition") + .withCatalogId(1).withStatementScope(scope).build(); + scope.closeAll(); + for (ConnectorTaskScheduler.Stage stage : ConnectorTaskScheduler.Stage.values()) { + Assertions.assertThrows(CancellationException.class, () -> session.getTaskScheduler(stage)); + } + } + + @Test + void scansAndCatalogsShareBothGroupsAndCannotResurrectFinishedQuery() { + try (FairTaskExecutor split = new FairTaskExecutor(1, Thread::new); + FairTaskExecutor metadata = new FairTaskExecutor(1, Thread::new)) { + ConnectorStatementScopeImpl scope = new ConnectorStatementScopeImpl(); + scope.computeIfAbsent(QueryTaskGroups.SCOPE_KEY, () -> new QueryTaskGroups( + stage -> stage == ConnectorTaskScheduler.Stage.SPLIT ? split : metadata)); + ConnectorSession first = ConnectorSessionBuilder.create().withQueryId("q") + .withCatalogId(1).withStatementScope(scope).build(); + ConnectorSession second = ConnectorSessionBuilder.create().withQueryId("q") + .withCatalogId(2).withStatementScope(scope).build(); + for (ConnectorTaskScheduler.Stage stage : ConnectorTaskScheduler.Stage.values()) { + Assertions.assertSame(first.getTaskScheduler(stage), second.getTaskScheduler(stage)); + } + ConnectorTaskScheduler captured = first.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT); + Assertions.assertFalse(captured.isClosed()); + ConnectorTaskScheduler capturedMetadata = first.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA); + Assertions.assertNotSame(captured, capturedMetadata); + Assertions.assertFalse(capturedMetadata.isClosed()); + scope.closeAll(); + scope.closeAll(); + Assertions.assertTrue(captured.isClosed()); + Assertions.assertTrue(capturedMetadata.isClosed()); + Assertions.assertThrows(RejectedExecutionException.class, + () -> captured.execute(() -> Assertions.fail("finished query ran"))); + Assertions.assertThrows(RejectedExecutionException.class, + () -> captured.tryExecuteContinuation(() -> { }, failure -> Assertions.fail(failure), null)); + Assertions.assertThrows(CancellationException.class, + () -> second.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA)); + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeBatchModeTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeBatchModeTest.java index ba546ecda19c5d..e2bbaae189b97f 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeBatchModeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeBatchModeTest.java @@ -57,7 +57,7 @@ */ public class PluginDrivenScanNodeBatchModeTest { - private static final int THRESHOLD = 1024; // num_partitions_in_batch_mode default; pinned (it is fuzzy at runtime) + private static final int THRESHOLD = 1024; // Explicit threshold, independent of session defaults. private static SelectedPartitions pruned(int count) { Map items = new LinkedHashMap<>(); @@ -169,10 +169,40 @@ private static PluginDrivenScanNode streamingNode(ConnectorScanPlanProvider prov Mockito.doReturn(TPushAggOp.NONE).when(node).getPushDownAggNoGroupingOp(); SessionVariable sv = Mockito.mock(SessionVariable.class); Mockito.when(sv.getNumPartitionsInBatchMode()).thenReturn(THRESHOLD); + Mockito.when(sv.getEnableExternalTableBatchMode()).thenReturn(true); Deencapsulation.setField(node, "sessionVariable", sv); return node; } + @Test + void disabledBatchModeSkipsBothPartitionAndStreamingProviders() { + for (SelectedPartitions partitions : new SelectedPartitions[] {pruned(1), SelectedPartitions.NOT_PRUNED}) { + ConnectorScanPlanProvider provider = Mockito.mock(ConnectorScanPlanProvider.class); + PluginDrivenScanNode node = streamingNode(provider); + SessionVariable defaults = new SessionVariable(); + defaults.enableExternalTableBatchMode = false; + defaults.numPartitionsInBatchMode = 1; + Deencapsulation.setField(node, "sessionVariable", defaults); + Deencapsulation.setField(node, "selectedPartitions", partitions); + Assertions.assertFalse(node.isBatchMode()); + Mockito.verifyNoInteractions(provider); + } + } + + @Test + void enabledBatchModeAcceptsOnePartitionThroughTheEngineGate() { + ConnectorScanPlanProvider provider = Mockito.mock(ConnectorScanPlanProvider.class); + Mockito.when(provider.supportsBatchScan(Mockito.any(), Mockito.any())).thenReturn(true); + PluginDrivenScanNode node = streamingNode(provider); + SessionVariable enabled = new SessionVariable(); + enabled.numPartitionsInBatchMode = 1; + Deencapsulation.setField(node, "sessionVariable", enabled); + Deencapsulation.setField(node, "selectedPartitions", pruned(1)); + Assertions.assertTrue(node.isBatchMode()); + Mockito.verify(provider, Mockito.never()).streamingSplitEstimate( + Mockito.any(), Mockito.any(), Mockito.any(), Mockito.anyBoolean()); + } + @Test public void testComputeBatchModePrefersStreamingWhenEstimateNonNegative() { // A connector that returns a non-negative streamingSplitEstimate enters the streaming flavor of batch 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 index 0be56ce2953520..8e85a2f2321df7 100644 --- 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 @@ -26,6 +26,7 @@ import org.apache.doris.common.Config; import org.apache.doris.common.jmockit.Deencapsulation; import org.apache.doris.common.util.DebugUtil; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.connector.hive.HiveCatalogProperties; import org.apache.doris.connector.hive.HiveFileListingCache; import org.apache.doris.connector.hive.HiveReadTransactionManager; @@ -41,7 +42,7 @@ 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.connector.spi.ConnectorTaskScheduler; import org.apache.doris.datasource.split.PluginDrivenSplit; import org.apache.doris.datasource.split.SplitAssignment; import org.apache.doris.datasource.split.SplitGenerator; @@ -56,7 +57,12 @@ import org.apache.doris.qe.SessionVariable; import org.apache.doris.spi.Split; import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TExternalScanRange; +import org.apache.doris.thrift.TFileRangeDesc; +import org.apache.doris.thrift.TFileScanRange; +import org.apache.doris.thrift.TFileScanRangeParams; import org.apache.doris.thrift.TPushAggOp; +import org.apache.doris.thrift.TScanRange; import org.apache.doris.thrift.TScanRangeLocation; import org.apache.doris.thrift.TScanRangeLocations; import org.apache.doris.thrift.TUniqueId; @@ -81,8 +87,6 @@ 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; @@ -147,7 +151,7 @@ void cancellationStopsHivePartitionBatchesBeforeFurtherListing() throws Exceptio Assertions.assertTrue(fixture.entered.await(10, TimeUnit.SECONDS)); fixture.assignment.stop(); fixture.release.countDown(); - fixture.executor.shutdown(); + fixture.executor.close(); Assertions.assertTrue(fixture.executor.awaitTermination(10, TimeUnit.SECONDS)); Assertions.assertEquals(2, fixture.listings.get()); } @@ -160,10 +164,10 @@ private static class Fixture implements AutoCloseable { final AtomicInteger listings = new AtomicInteger(); final AtomicInteger activeListings = new AtomicInteger(); final AtomicInteger peakListings = new AtomicInteger(); - final ExecutorService executor = Executors.newFixedThreadPool(4); + final FairTaskExecutor executor = new FairTaskExecutor(4, Thread::new); 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 int previousParallelism = Config.external_scan_split_threads; final MockedStatic envStatic; final PluginDrivenScanNode node; final SplitAssignment assignment; @@ -235,6 +239,18 @@ public void close() { 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); + FairTaskExecutor.Group group = executor.newGroup(1); + ConnectorTaskScheduler scheduler = Mockito.mock(ConnectorTaskScheduler.class, Mockito.CALLS_REAL_METHODS); + Mockito.when(scheduler.tryExecute(Mockito.any(), Mockito.any(), Mockito.any())) + .thenAnswer(call -> group.tryExecute(call.getArgument(0), call.getArgument(1), call.getArgument(2))); + Mockito.when(scheduler.tryExecuteContinuation(Mockito.any(), Mockito.any(), Mockito.any())) + .thenAnswer(call -> group.tryExecuteContinuation( + call.getArgument(0), call.getArgument(1), call.getArgument(2))); + Mockito.doAnswer(call -> { + group.removeSpaceListener(call.getArgument(0)); + return null; + }).when(scheduler).removeSpaceListener(Mockito.any()); + Mockito.when(session.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT)).thenReturn(scheduler); 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. @@ -247,6 +263,7 @@ public void close() { Mockito.when(desc.getSlots()).thenReturn(new ArrayList<>(Collections.singletonList( Mockito.mock(SlotDescriptor.class)))); SessionVariable variables = Mockito.mock(SessionVariable.class); + Mockito.when(variables.getEnableExternalTableBatchMode()).thenReturn(true); Mockito.when(variables.getNumPartitionsInBatchMode()).thenReturn(2); Mockito.when(variables.getSplitBatchSize()).thenReturn(2); Mockito.when(variables.getSplitQueueCapacity()).thenReturn(1); @@ -273,18 +290,19 @@ public void close() { (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))); + TFileScanRange files = new TFileScanRange().setParams(new TFileScanRangeParams()) + .setRanges(Collections.singletonList(new TFileRangeDesc().setPath("file-" + identity))); + return new TScanRangeLocations().setScanRange(new TScanRange().setExtScanRange( + new TExternalScanRange().setFileScanRange(files))) + .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; + Config.external_scan_split_threads = 2; } void start() { @@ -315,12 +333,12 @@ public void close() throws Exception { assignment.stop(); release.countDown(); QeProcessorImpl.INSTANCE.unregisterQuery(queryId); - executor.shutdown(); + executor.close(); Assertions.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); } finally { - executor.shutdownNow(); + executor.close(); envStatic.close(); - Config.max_external_cache_loader_thread_pool_size = previousParallelism; + Config.external_scan_split_threads = previousParallelism; } } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeReadTxnReleaseTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeReadTxnReleaseTest.java index 8c8da92cdce952..1ffa92f7faec7b 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeReadTxnReleaseTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeReadTxnReleaseTest.java @@ -23,6 +23,7 @@ import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.ConnectorStatementScope; import org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider; +import org.apache.doris.connector.spi.scan.ConnectorSplitSource; import org.apache.doris.datasource.split.SplitAssignment; import org.apache.doris.qe.QeProcessorImpl; import org.apache.doris.thrift.TUniqueId; @@ -52,6 +53,26 @@ */ public class PluginDrivenScanNodeReadTxnReleaseTest { + @Test + public void demandCallbackPinsAndRestoresThePluginLoader() throws Exception { + ConnectorScanPlanProvider provider = Mockito.mock(ConnectorScanPlanProvider.class); + ConnectorSplitSource source = Mockito.mock(ConnectorSplitSource.class); + Mockito.doAnswer(invocation -> { + Assertions.assertSame(provider.getClass().getClassLoader(), + Thread.currentThread().getContextClassLoader()); + return null; + }).when(source).onMoreSplitsRequested(); + ClassLoader original = Thread.currentThread().getContextClassLoader(); + try (URLClassLoader sentinel = new URLClassLoader(new URL[0], null)) { + Thread.currentThread().setContextClassLoader(sentinel); + PluginDrivenScanNode.buildSplitSourceDemandCallback(provider, source).run(); + Mockito.verify(source).onMoreSplitsRequested(); + Assertions.assertSame(sentinel, Thread.currentThread().getContextClassLoader()); + } finally { + Thread.currentThread().setContextClassLoader(original); + } + } + @Test public void callbackReleasesReadTransactionForTheQuery() { ConnectorScanPlanProvider provider = Mockito.mock(ConnectorScanPlanProvider.class); diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeRedispatchTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeRedispatchTest.java index a7d73f7bb5e42e..a4f824ce5d667b 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeRedispatchTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/scan/PluginDrivenScanNodeRedispatchTest.java @@ -19,6 +19,7 @@ import org.apache.doris.common.jmockit.Deencapsulation; import org.apache.doris.connector.spi.scan.ConnectorScanRange; +import org.apache.doris.datasource.split.PluginDrivenSplit; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -65,16 +66,17 @@ public boolean isSingleUse() { public void planWithSingleUseRangeCannotBeRedispatched() { PluginDrivenScanNode node = Mockito.mock(PluginDrivenScanNode.class, Mockito.CALLS_REAL_METHODS); Assertions.assertFalse(node.cannotBeRedispatched()); + PluginDrivenSplit.Factory splitFactory = new PluginDrivenSplit.Factory(); // Ranges the source serves afresh on every read (a file, a statement run when it is read) keep the // same-plan retry. - Deencapsulation.invoke(node, "toSplit", range(false)); + Deencapsulation.invoke(node, "toSplit", range(false), splitFactory); Assertions.assertFalse(node.cannotBeRedispatched()); // One range that can be read only once is enough, wherever it comes in the plan. - Deencapsulation.invoke(node, "toSplit", range(true)); + Deencapsulation.invoke(node, "toSplit", range(true), splitFactory); Assertions.assertTrue(node.cannotBeRedispatched()); - Deencapsulation.invoke(node, "toSplit", range(false)); + Deencapsulation.invoke(node, "toSplit", range(false), splitFactory); Assertions.assertTrue(node.cannotBeRedispatched()); } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/FileScanRangeCompactorTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/FileScanRangeCompactorTest.java new file mode 100644 index 00000000000000..0d5f81607d5293 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/FileScanRangeCompactorTest.java @@ -0,0 +1,269 @@ +// 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.thrift.TExternalScanRange; +import org.apache.doris.thrift.TFetchSplitBatchResult; +import org.apache.doris.thrift.TFileFormatType; +import org.apache.doris.thrift.TFileRangeDesc; +import org.apache.doris.thrift.TFileScanRange; +import org.apache.doris.thrift.TFileScanRangeParams; +import org.apache.doris.thrift.THudiFileDesc; +import org.apache.doris.thrift.TIcebergDeleteFileDesc; +import org.apache.doris.thrift.TIcebergFileDesc; +import org.apache.doris.thrift.TLanceFileDesc; +import org.apache.doris.thrift.TMaxComputeFileDesc; +import org.apache.doris.thrift.TNetworkAddress; +import org.apache.doris.thrift.TPaimonDeletionFileDesc; +import org.apache.doris.thrift.TPaimonFileDesc; +import org.apache.doris.thrift.TPaimonReaderType; +import org.apache.doris.thrift.TQueryOptions; +import org.apache.doris.thrift.TRemoteDorisFileDesc; +import org.apache.doris.thrift.TScanRange; +import org.apache.doris.thrift.TScanRangeLocation; +import org.apache.doris.thrift.TScanRangeLocations; +import org.apache.doris.thrift.TScanRangeParams; +import org.apache.doris.thrift.TSplitSource; +import org.apache.doris.thrift.TTableFormatFileDesc; +import org.apache.doris.thrift.TTransactionalHiveDeleteDeltaDesc; +import org.apache.doris.thrift.TTransactionalHiveDesc; +import org.apache.doris.thrift.TTrinoConnectorFileDesc; + +import org.apache.thrift.TDeserializer; +import org.apache.thrift.TSerializer; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +public class FileScanRangeCompactorTest { + @Test + void preservesEveryConnectorDescriptorThroughSplitBatchSerialization() throws Exception { + List expected = connectorRanges(); + List locations = new ArrayList<>(); + for (TFileRangeDesc file : expected) { + locations.add(new TScanRangeLocations().setScanRange(scanRange(file)) + .setLocations(List.of(new TScanRangeLocation().setBackendId(1) + .setServer(new TNetworkAddress("backend", 9060))))); + } + TFetchSplitBatchResult result = new TFetchSplitBatchResult() + .setSplits(FileScanRangeCompactor.compactSplitBatch(locations)); + TFetchSplitBatchResult decoded = new TFetchSplitBatchResult(); + new TDeserializer().deserialize(decoded, new TSerializer().serialize(result)); + + Assertions.assertEquals(1, decoded.getSplitsSize()); + Assertions.assertEquals(expected, decoded.getSplits().get(0).getScanRange() + .getExtScanRange().getFileScanRange().getRanges()); + for (TScanRangeLocations original : locations) { + Assertions.assertEquals(1, original.getScanRange().getExtScanRange().getFileScanRange().getRangesSize()); + } + } + + @Test + void compactsAssignedFileRangesWithoutReducingScannerConcurrency() { + List original = ranges(1000); + Map> input = Map.of(1, original); + TQueryOptions options = new TQueryOptions().setMaxFileScannersConcurrency(16).setParallelInstance(7); + + List compacted = FileScanRangeCompactor.compactInstanceScanRanges( + input, Map.of(1, new TFileScanRangeParams()), options).get(1); + + Assertions.assertTrue(compacted.size() >= 112); + Assertions.assertTrue(compacted.size() < original.size()); + Assertions.assertEquals(flatten(original), flatten(compacted)); + Assertions.assertEquals(1000, input.get(1).size()); + Assertions.assertTrue(original.stream().allMatch(r -> files(r).getRangesSize() == 1)); + // Whatever the BE thread-pool/operator limit is, its final scanner count remains identical. + for (int backendLimit = 1; backendLimit <= 112; backendLimit++) { + Assertions.assertEquals(Math.min(backendLimit, original.size()), + Math.min(backendLimit, compacted.size())); + } + } + + @Test + void usesBackendDefaultScannerLimitForUnsetOrNonPositiveSessionValue() { + for (int configured : List.of(0, -1)) { + List original = ranges(100); + List compacted = FileScanRangeCompactor.compactInstanceScanRanges( + Map.of(1, original), Map.of(1, new TFileScanRangeParams()), + new TQueryOptions().setMaxFileScannersConcurrency(configured).setParallelInstance(1)).get(1); + Assertions.assertTrue(compacted.size() >= 16); + Assertions.assertTrue(compacted.size() < 100); + Assertions.assertEquals(flatten(original), flatten(compacted)); + } + } + + @Test + void leavesSmallFileScansNonFileScansAndBatchSourceReferencesUnchanged() { + List small = ranges(3); + List nativeRanges = List.of(new TScanRangeParams().setScanRange(new TScanRange())); + TFileScanRange source = new TFileScanRange().setSplitSource(new TSplitSource().setSplitSourceId(99)); + List batch = List.of(new TScanRangeParams().setScanRange( + new TScanRange().setExtScanRange(new TExternalScanRange().setFileScanRange(source)))); + Map> result = FileScanRangeCompactor.compactInstanceScanRanges( + Map.of(1, small, 2, nativeRanges, 3, batch), + Map.of(1, new TFileScanRangeParams(), 3, new TFileScanRangeParams()), new TQueryOptions()); + + Assertions.assertSame(small, result.get(1)); + Assertions.assertSame(nativeRanges, result.get(2)); + Assertions.assertSame(batch, result.get(3)); + } + + @Test + void skipsQueriesWithoutFileScans() { + Map> nativeScans = Map.of(1, + List.of(new TScanRangeParams().setScanRange(new TScanRange()))); + Assertions.assertSame(nativeScans, FileScanRangeCompactor.compactInstanceScanRanges( + nativeScans, Map.of(), new TQueryOptions())); + } + + @Test + void neverCombinesIndependentSplitSourceStreams() { + List sources = new ArrayList<>(); + for (int i = 0; i < 40; i++) { + sources.add(new TScanRangeParams().setScanRange(new TScanRange().setExtScanRange( + new TExternalScanRange().setFileScanRange(new TFileScanRange() + .setSplitSource(new TSplitSource().setSplitSourceId(i)))))); + } + List result = FileScanRangeCompactor.compactInstanceScanRanges( + Map.of(1, sources), Map.of(1, new TFileScanRangeParams()), + new TQueryOptions().setMaxFileScannersConcurrency(1).setParallelInstance(1)).get(1); + Assertions.assertEquals(sources, result); + } + + @Test + void preservesVolumeIdsParameterBoundariesAndInputObjects() { + List original = ranges(20); + for (int i = 0; i < 20; i++) { + original.get(i).setVolumeId(i < 10 ? 1 : 2); + files(original.get(i)).setParams(new TFileScanRangeParams() + .setFormatType(i < 5 ? TFileFormatType.FORMAT_PARQUET : TFileFormatType.FORMAT_JNI)); + } + List compacted = FileScanRangeCompactor.compactInstanceScanRanges( + Map.of(1, original), Map.of(1, new TFileScanRangeParams()), + new TQueryOptions().setMaxFileScannersConcurrency(1).setParallelInstance(1)).get(1); + + Assertions.assertEquals(3, compacted.size()); + Assertions.assertEquals(List.of(5, 5, 10), compacted.stream() + .map(r -> files(r).getRangesSize()).collect(Collectors.toList())); + Assertions.assertEquals(List.of(1, 1, 2), compacted.stream() + .map(TScanRangeParams::getVolumeId).collect(Collectors.toList())); + Assertions.assertEquals(flatten(original), flatten(compacted)); + Assertions.assertTrue(original.stream().allMatch(r -> files(r).getRangesSize() == 1)); + } + + @Test + void keepsConnectorFieldsAndShrinksSynchronousSerializedRequests() throws Exception { + List expected = connectorRanges(); + List original = new ArrayList<>(); + for (int i = 0; i < 30; i++) { + for (TFileRangeDesc file : expected) { + original.add(new TScanRangeParams().setScanRange(scanRange(file))); + } + } + List compacted = FileScanRangeCompactor.compactInstanceScanRanges( + Map.of(1, original), Map.of(1, new TFileScanRangeParams()), new TQueryOptions()).get(1); + TSerializer serializer = new TSerializer(); + int beforeBytes = 0; + int afterBytes = 0; + List decoded = new ArrayList<>(); + for (TScanRangeParams range : original) { + beforeBytes += serializer.serialize(range).length; + } + for (TScanRangeParams range : compacted) { + byte[] bytes = serializer.serialize(range); + afterBytes += bytes.length; + TScanRangeParams copy = new TScanRangeParams(); + new TDeserializer().deserialize(copy, bytes); + decoded.add(copy); + } + Assertions.assertEquals(flatten(original), flatten(decoded)); + Assertions.assertTrue(afterBytes < beforeBytes); + } + + private static List ranges(int count) { + List ranges = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ranges.add(new TScanRangeParams().setScanRange(scanRange( + new TFileRangeDesc().setPath("file" + i).setStartOffset(i).setSize(10)))); + } + return ranges; + } + + private static TScanRange scanRange(TFileRangeDesc file) { + return new TScanRange().setExtScanRange(new TExternalScanRange() + .setFileScanRange(new TFileScanRange().setRanges(new ArrayList<>(List.of(file))))); + } + + private static TFileScanRange files(TScanRangeParams range) { + return range.getScanRange().getExtScanRange().getFileScanRange(); + } + + private static List flatten(List ranges) { + return ranges.stream().flatMap(r -> files(r).getRanges().stream()).collect(Collectors.toList()); + } + + private static List connectorRanges() { + List formats = List.of( + new TTableFormatFileDesc().setTableFormatType("hive").setTransactionalHiveParams( + new TTransactionalHiveDesc().setPartition("p=1").setDeleteDeltas(List.of( + new TTransactionalHiveDeleteDeltaDesc().setDirectoryLocation("delete_delta") + .setFileNames(List.of("bucket_00000"))))), + new TTableFormatFileDesc().setTableFormatType("iceberg").setIcebergParams( + new TIcebergFileDesc().setFormatVersion(3).setPartitionSpecId(7) + .setPartitionDataJson("{partition:1}").setFirstRowId(100).setSerializedSplit("split") + .setDeleteFiles(List.of(new TIcebergDeleteFileDesc().setPath("position-delete") + .setContent(1), new TIcebergDeleteFileDesc().setPath("equality-delete") + .setContent(2).setFieldIds(List.of(1))))), + new TTableFormatFileDesc().setTableFormatType("paimon").setPaimonParams( + new TPaimonFileDesc().setReaderType(TPaimonReaderType.PAIMON_NATIVE).setSchemaId(8) + .setDeletionFile(new TPaimonDeletionFileDesc().setPath("deletion-vector") + .setOffset(9).setLength(10))), + new TTableFormatFileDesc().setTableFormatType("paimon").setPaimonParams( + new TPaimonFileDesc().setReaderType(TPaimonReaderType.PAIMON_JNI).setPaimonSplit("jni-split")), + new TTableFormatFileDesc().setTableFormatType("hudi").setHudiParams( + new THudiFileDesc().setInstantTime("20261013").setDataFilePath("base.parquet") + .setDeltaLogs(List.of("log.1", "log.2")).setSchemaId(11)), + new TTableFormatFileDesc().setTableFormatType("max_compute").setMaxComputeParams( + new TMaxComputeFileDesc().setSessionId("session").setTableBatchReadSession("read-session")), + new TTableFormatFileDesc().setTableFormatType("trino_connector").setTrinoConnectorParams( + new TTrinoConnectorFileDesc().setTrinoConnectorSplit("trino-split")), + new TTableFormatFileDesc().setTableFormatType("jdbc").setJdbcParams(Map.of("query", "SELECT 1")), + new TTableFormatFileDesc().setTableFormatType("es").setEsParams(Map.of("shard_id", "2")), + new TTableFormatFileDesc().setTableFormatType("adbc").setAdbcParams(Map.of("partition_b64", "AAAA")), + new TTableFormatFileDesc().setTableFormatType("lance").setLanceParams( + new TLanceFileDesc().setDatasetUri("dataset").setFragmentIds(List.of(1L, 2L)).setVersion(3)), + new TTableFormatFileDesc().setTableFormatType("fluss").setFlussParams( + Map.of("range_type", "LOG", "start_offset", "5")), + new TTableFormatFileDesc().setTableFormatType("remote_doris").setRemoteDorisParams( + new TRemoteDorisFileDesc().setTicket(new byte[] {1, 2, 3}).setLocationUri("grpc://backend"))); + List result = new ArrayList<>(); + for (int i = 0; i < formats.size(); i++) { + result.add(new TFileRangeDesc().setPath("file" + i).setStartOffset(i).setSize(10 + i) + .setFileSize(100 + i).setModificationTime(1000 + i).setFsName("filesystem" + i) + .setFormatType(i % 2 == 0 ? TFileFormatType.FORMAT_PARQUET : TFileFormatType.FORMAT_JNI) + .setSelfSplitWeight(i + 1).setFileCacheAdmission(i % 2 == 0) + .setColumnsFromPath(List.of("partition" + i)).setColumnsFromPathKeys(List.of("p")) + .setColumnsFromPathIsNull(List.of(i % 2 == 0)).setTableFormatParams(formats.get(i))); + } + return result; + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/PluginDrivenSplitPathTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/PluginDrivenSplitPathTest.java new file mode 100644 index 00000000000000..aa1a6e3e77c658 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/PluginDrivenSplitPathTest.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.datasource.split; + +import org.apache.doris.connector.spi.scan.ConnectorScanRange; +import org.apache.doris.thrift.TFileType; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.Map; +import java.util.Optional; + +class PluginDrivenSplitPathTest { + @Test + void reuseFilePathWhilePreservingPerSplitOffsets() { + PluginDrivenSplit.Factory factory = new PluginDrivenSplit.Factory(); + String path = "s3://bucket/a//p=中文%2F+ b/file.parquet"; + PluginDrivenSplit first = factory.create(range(path, 0)); + PluginDrivenSplit second = factory.create(range(new String(path), 100)); + Assertions.assertSame(first.getPath(), second.getPath()); + Assertions.assertEquals(0, first.getStart()); + Assertions.assertEquals(100, second.getStart()); + Assertions.assertEquals(path, second.getPathString()); + Assertions.assertEquals(TFileType.FILE_S3, second.getLocationType()); + } + + @Test + void changedFileAndIndependentProducersDoNotShareMutableState() { + PluginDrivenSplit.Factory factory = new PluginDrivenSplit.Factory(); + PluginDrivenSplit first = factory.create(range("s3://bucket/a", 0)); + PluginDrivenSplit other = factory.create(range("hdfs://namenode:8020/b", 0)); + Assertions.assertNotSame(first.getPath(), other.getPath()); + Assertions.assertEquals("s3://bucket/a", first.getPathString()); + Assertions.assertEquals("hdfs://namenode:8020", other.getPath().getFsIdentifier()); + Assertions.assertEquals(TFileType.FILE_HDFS, other.getLocationType()); + Assertions.assertNotSame(first.getPath(), factory.create(range("s3://bucket/a", 100)).getPath()); + Assertions.assertNotSame(other.getPath(), new PluginDrivenSplit.Factory() + .create(range("hdfs://namenode:8020/b", 100)).getPath()); + } + + @Test + void pathlessRangesKeepVirtualPath() { + PluginDrivenSplit.Factory factory = new PluginDrivenSplit.Factory(); + PluginDrivenSplit first = factory.create(range(null, 0)); + PluginDrivenSplit second = factory.create(range(null, 100)); + Assertions.assertEquals("connector://virtual", first.getPathString()); + Assertions.assertSame(first.getPath(), second.getPath()); + Assertions.assertNotSame(first.getConnectorScanRange(), second.getConnectorScanRange()); + } + + private static ConnectorScanRange range(String path, long start) { + return new ConnectorScanRange() { + @Override + public Map getProperties() { + return Collections.emptyMap(); + } + + @Override + public Optional getPath() { + return Optional.ofNullable(path); + } + + @Override + public long getStart() { + return start; + } + }; + } +} 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 2a40fce91dcb0f..267737b2a4b9f7 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 @@ -86,6 +86,97 @@ void setUp() { ); } + @Test + void firstConvertedSplitDoesNotWaitForAnotherBackendAssignment() throws Exception { + CountDownLatch assignmentEntered = new CountDownLatch(1); + CountDownLatch releaseAssignment = new CountDownLatch(1); + CountDownLatch initStarted = new CountDownLatch(1); + Mockito.doAnswer(invocation -> { + initStarted.countDown(); + return null; + }).when(mockSplitGenerator).startSplit(Mockito.anyInt()); + Mockito.when(mockBackendPolicy.computeScanRangeAssignment(Mockito.any())).thenAnswer(invocation -> { + assignmentEntered.countDown(); + Assertions.assertTrue(releaseAssignment.await(30, TimeUnit.SECONDS)); + return ArrayListMultimap.create(); + }); + Mockito.when(mockSplitToScanRange.getScanRange(Mockito.any(), Mockito.any(), Mockito.any(), + Mockito.any(), Mockito.anyBoolean())).thenReturn(mockScanRangeLocations); + ExecutorService workers = Executors.newFixedThreadPool(3); + try { + Future init = workers.submit(() -> { + splitAssignment.init(); + return null; + }); + Assertions.assertTrue(initStarted.await(5, TimeUnit.SECONDS)); + Future assigning = workers.submit(() -> splitAssignment.assignBatch( + Collections.singletonList(mockSplit))); + Assertions.assertTrue(assignmentEntered.await(5, TimeUnit.SECONDS)); + Future converted = workers.submit(() -> splitAssignment.createScanRanges( + mockBackend, Collections.singletonList(mockSplit))); + converted.get(5, TimeUnit.SECONDS); + init.get(5, TimeUnit.SECONDS); + Assertions.assertSame(mockSplit, splitAssignment.getSampleSplit()); + Assertions.assertFalse(assigning.isDone(), "readiness must not wait for unrelated assignment work"); + } finally { + releaseAssignment.countDown(); + workers.shutdownNow(); + Assertions.assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void demandCallbackStopsWithTheAssignment() { + java.util.concurrent.atomic.AtomicInteger requested = new java.util.concurrent.atomic.AtomicInteger(); + splitAssignment.setMoreSplitsRequestedCallback(requested::incrementAndGet); + splitAssignment.onMoreSplitsRequested(); + Assertions.assertEquals(1, requested.get()); + splitAssignment.stop(); + splitAssignment.onMoreSplitsRequested(); + splitAssignment.setMoreSplitsRequestedCallback(requested::incrementAndGet); + splitAssignment.onMoreSplitsRequested(); + Assertions.assertEquals(1, requested.get()); + } + + @Test + void asyncResourceRegisteredAfterStopIsClosedImmediately() { + java.util.concurrent.atomic.AtomicInteger closes = new java.util.concurrent.atomic.AtomicInteger(); + splitAssignment.stop(); + splitAssignment.addCloseable(closes::incrementAndGet); + Assertions.assertEquals(1, closes.get()); + splitAssignment.stop(); + Assertions.assertEquals(1, closes.get()); + } + + @Test + void stopClosesRegisteredResourceOutsideRegistryLock() { + java.util.concurrent.atomic.AtomicInteger closes = new java.util.concurrent.atomic.AtomicInteger(); + splitAssignment.addCloseable(() -> { + closes.incrementAndGet(); + splitAssignment.addCloseable(closes::incrementAndGet); + }); + splitAssignment.stop(); + Assertions.assertEquals(2, closes.get()); + } + + @Test + void closingActiveProducerDoesNotTurnSuccessfulStopIntoFailure() throws Exception { + splitAssignment.addCloseable(() -> splitAssignment.setException(new UserException("source closed"))); + Assertions.assertDoesNotThrow(() -> splitAssignment.stop()); + splitAssignment.setException(new UserException("producer interrupted after stop")); + Assertions.assertNull(splitAssignment.getAssignedSplits(mockBackend)); + } + + @Test + void stopPreservesFailurePublishedBeforeCancellation() { + UserException failure = new UserException("manifest read failed"); + splitAssignment.setException(failure); + splitAssignment.addCloseable(() -> splitAssignment.setException(new UserException("source closed"))); + RuntimeException actual = Assertions.assertThrows(RuntimeException.class, () -> splitAssignment.stop()); + Assertions.assertSame(failure, actual.getCause()); + Assertions.assertEquals(0, failure.getSuppressed().length); + } + // ==================== init() method tests ==================== @Test diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java index 1dfe1b89c91db9..890b9075f2a526 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java @@ -18,10 +18,17 @@ package org.apache.doris.datasource.split; import org.apache.doris.common.UserException; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.common.util.LocationPath; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.datasource.scan.FederationBackendPolicy; import org.apache.doris.spi.Split; import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TExternalScanRange; +import org.apache.doris.thrift.TFileRangeDesc; +import org.apache.doris.thrift.TFileScanRange; +import org.apache.doris.thrift.TFileScanRangeParams; +import org.apache.doris.thrift.TScanRange; import org.apache.doris.thrift.TScanRangeLocation; import org.apache.doris.thrift.TScanRangeLocations; @@ -48,8 +55,164 @@ import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.LongSupplier; class SplitProducerTest { + private ConnectorTaskScheduler scheduler(FairTaskExecutor.Group group) { + return new ConnectorTaskScheduler() { + @Override + public boolean tryExecute(Runnable task, java.util.function.Consumer onFailure, + Runnable onSpace) { + return group.tryExecute(task, onFailure, onSpace); + } + + @Override + public boolean tryExecuteContinuation(Runnable task, java.util.function.Consumer onFailure, + Runnable onSpace) { + return group.tryExecuteContinuation(task, onFailure, onSpace); + } + + @Override + public void removeSpaceListener(Runnable onSpace) { + group.removeSpaceListener(onSpace); + } + + @Override + public void execute(Runnable task) { + group.execute(task); + } + }; + } + + @Test + void producersFillConfiguredBufferBeforeAnyBackendFetch() throws Exception { + Fixture fixture = new Fixture(1, 3); + ManualExecutor executor = new ManualExecutor(); + RecordingSource source = new RecordingSource(1, 0, 2, 4, 6, 8); + fixture.start(executor, source); + executor.runUntilIdle(); + Assertions.assertEquals(3, fixture.assignment.getAssignedSplits(fixture.backend).size()); + Assertions.assertEquals(0, fixture.completions.get()); + SplitSource consumer = fixture.source(); + List actual = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + actual.addAll(identities(consumer.getNextBatch(1))); + executor.runUntilIdle(); + } + Assertions.assertEquals(Arrays.asList(0L, 2L, 4L, 6L, 8L), actual); + Assertions.assertTrue(consumer.getNextBatch(1).isEmpty()); + Assertions.assertEquals(1, fixture.completions.get()); + } + + @Test + void queryCloseReleasesAProducerWaitingForTaskQueueCapacity() throws Exception { + Fixture fixture = new Fixture(1, 1); + RecordingSource source = new RecordingSource(1, 0, 2, 4); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + try (FairTaskExecutor pool = new FairTaskExecutor(1, Thread::new)) { + FairTaskExecutor.Group group = pool.newGroup(1); + group.execute(() -> { + entered.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }); + try { + Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS)); + group.execute(() -> { }); + fixture.start(scheduler(group), source); + fixture.assignment.stop(); + group.close(); + Assertions.assertTrue(source.closed.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + } finally { + release.countDown(); + pool.close(); + Assertions.assertTrue(pool.awaitTermination(5, TimeUnit.SECONDS)); + } + } + } + + @Test + void fairExecutorRetainsBackendBackpressureAndProducesEverySplitOnce() throws Exception { + Fixture fixture = new Fixture(1, 1); + RecordingSource source = new RecordingSource(1, 0, 2, 4, 6, 8); + try (FairTaskExecutor pool = new FairTaskExecutor(2, Thread::new)) { + fixture.start(scheduler(pool.newGroup(1)), source); + Assertions.assertEquals(Arrays.asList(0L, 2L, 4L, 6L, 8L), drain(fixture.source(), 1)); + Assertions.assertTrue(source.closed.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(1, fixture.completions.get()); + pool.close(); + Assertions.assertTrue(pool.awaitTermination(5, TimeUnit.SECONDS)); + } finally { + fixture.assignment.stop(); + } + } + + @Test + void fairContinuationYieldsToAnotherQueryAndPreservesAllSplits() throws Exception { + Fixture first = new Fixture(1, 512); + Fixture second = new Fixture(1, 512); + long[] ids = new long[200]; + List expected = new ArrayList<>(); + for (int i = 0; i < ids.length; i++) { + ids[i] = i * 2L; + expected.add(ids[i]); + } + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + RecordingSource firstSource = new RecordingSource(1, ids) { + boolean firstRead = true; + + @Override + public List getNextBatch() throws IOException { + if (firstRead) { + firstRead = false; + started.countDown(); + try { + Assertions.assertTrue(release.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + } + return super.getNextBatch(); + } + }; + AtomicInteger firstReadsWhenSecondStarts = new AtomicInteger(); + RecordingSource secondSource = new RecordingSource(1, 1000) { + @Override + public List getNextBatch() throws IOException { + firstReadsWhenSecondStarts.compareAndSet(0, firstSource.next); + return super.getNextBatch(); + } + }; + try (FairTaskExecutor pool = new FairTaskExecutor(1, Thread::new)) { + first.start(scheduler(pool.newGroup(1)), firstSource, () -> 0L); + try { + Assertions.assertTrue(started.await(5, TimeUnit.SECONDS)); + second.start(scheduler(pool.newGroup(1)), secondSource); + } finally { + release.countDown(); + } + Assertions.assertTrue(firstSource.closed.await(5, TimeUnit.SECONDS)); + Assertions.assertTrue(secondSource.closed.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(SplitProducer.MAX_BATCHES_PER_TASK, firstReadsWhenSecondStarts.get()); + Assertions.assertEquals(expected, drain(first.source(), 1)); + Assertions.assertEquals(Collections.singletonList(1000L), drain(second.source(), 1)); + Assertions.assertEquals(1, firstSource.closeCalls.get()); + Assertions.assertEquals(1, secondSource.closeCalls.get()); + } finally { + first.assignment.stop(); + second.assignment.stop(); + } + } + @Test void twoScansInitializeOnOneWorkerBeforeAnyConsumerStarts() throws Exception { Fixture first = new Fixture(1, 1); @@ -188,6 +351,7 @@ public List getNextBatch() throws IOException { Assertions.assertTrue(source.closed.await(5, TimeUnit.SECONDS)); Assertions.assertEquals(1, source.closeCalls.get()); Assertions.assertEquals(1, fixture.completions.get()); + Mockito.verify(fixture.policy, Mockito.never()).computeScanRangeAssignment(Mockito.any()); } finally { release.countDown(); executor.shutdownNow(); @@ -291,6 +455,70 @@ void failureRacingWaiterRegistrationDoesNotLeaveTheSourceSuspended() throws Exce } } + @Test + void productiveTaskYieldsToAnotherQueryBeforeFinishing() throws Exception { + Fixture first = new Fixture(1, 1000); + Fixture second = new Fixture(1, 1000); + ManualExecutor executor = new ManualExecutor(); + long[] ids = new long[200]; + List expected = new ArrayList<>(); + for (int i = 0; i < ids.length; i++) { + ids[i] = i * 2L; + expected.add(ids[i]); + } + RecordingSource source = new RecordingSource(1, ids); + first.start(executor, source, () -> 0L); + second.start(executor, new RecordingSource(1, 1000)); + executor.runNext(); + Assertions.assertEquals(SplitProducer.MAX_BATCHES_PER_TASK, source.next); + Assertions.assertEquals(0, first.completions.get()); + executor.runNext(); + Assertions.assertEquals(1, second.assignment.getAssignedSplits(second.backend).size()); + executor.runUntilIdle(); + Assertions.assertEquals(expected, drain(first.source(), 1)); + Assertions.assertEquals(1, source.closeCalls.get()); + } + + @Test + void slowBatchYieldsEvenBeforeBatchBudgetIsExhausted() throws Exception { + Fixture fixture = new Fixture(1, 100); + ManualExecutor executor = new ManualExecutor(); + AtomicLong clock = new AtomicLong(); + RecordingSource source = new RecordingSource(1, 0, 2) { + @Override + public List getNextBatch() throws IOException { + clock.addAndGet(SplitProducer.MAX_TASK_NANOS); + return super.getNextBatch(); + } + }; + fixture.start(executor, source, clock::get); + executor.runNext(); + Assertions.assertEquals(1, source.next); + Assertions.assertEquals(1, executor.tasks.size()); + fixture.assignment.stop(); + executor.runUntilIdle(); + Assertions.assertEquals(1, source.next); + Assertions.assertEquals(1, source.closeCalls.get()); + } + + @Test + void cancellationBetweenInlineBatchesStopsFurtherReads() throws Exception { + Fixture fixture = new Fixture(1, 100); + fixture.assignment = Mockito.spy(fixture.assignment); + Mockito.doAnswer(call -> { + boolean result = (boolean) call.callRealMethod(); + fixture.assignment.stop(); + return result; + }).when(fixture.assignment).offerBatch(Mockito.any(), Mockito.any(), Mockito.any()); + ManualExecutor executor = new ManualExecutor(); + RecordingSource source = new RecordingSource(1, 0, 2, 4); + fixture.start(executor, source); + executor.runUntilIdle(); + Assertions.assertEquals(1, source.next); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + } + private static List drain(SplitSource source, int batchSize) throws Exception { List result = new ArrayList<>(); while (true) { @@ -389,19 +617,27 @@ private static class Fixture { }); 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()))); + TFileScanRange range = new TFileScanRange().setParams(new TFileScanRangeParams()) + .setRanges(Collections.singletonList(new TFileRangeDesc())); + return new TScanRangeLocations().setScanRange(new TScanRange().setExtScanRange( + new TExternalScanRange().setFileScanRange(range))) + .setLocations(Collections.singletonList( + new TScanRangeLocation().setBackendId(split.getStart()))); }, Collections.emptyMap(), Collections.emptyList(), true, batchSize, queueCapacity); } void start(Executor executor, RecordingSource source) { + start(executor, source, System::nanoTime); + } + + void start(Executor executor, RecordingSource source, LongSupplier nanoTime) { new SplitProducer(assignment, executor, source, failure -> { if (failure != null) { assignment.setException(new UserException(failure.getMessage(), failure)); } completions.incrementAndGet(); assignment.finishSchedule(); - }).start(); + }, nanoTime).start(); } SplitSource source() { 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 index 5af23f18954e53..5ed216e9b6a495 100644 --- 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 @@ -19,17 +19,29 @@ import org.apache.doris.common.UserException; import org.apache.doris.datasource.scan.FederationBackendPolicy; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.VarAttrDef; import org.apache.doris.spi.Split; import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TExternalScanRange; +import org.apache.doris.thrift.TFileRangeDesc; +import org.apache.doris.thrift.TFileScanRange; +import org.apache.doris.thrift.TFileScanRangeParams; +import org.apache.doris.thrift.TScanRange; import org.apache.doris.thrift.TScanRangeLocations; import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.Multimap; +import org.awaitility.Awaitility; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.Mockito; +import java.lang.management.ManagementFactory; +import java.lang.reflect.Field; import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; @@ -40,11 +52,14 @@ import java.util.Set; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; 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; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; public class SplitSourceTest { private Backend backend; @@ -68,7 +83,12 @@ void setUp() throws Exception { IdentityHashMap locations = new IdentityHashMap<>(); for (int i = 0; i < 8; i++) { Split split = Mockito.mock(Split.class); - TScanRangeLocations range = Mockito.mock(TScanRangeLocations.class); + // Keep envelopes incompatible so the queue tests still observe one envelope per split. + TScanRangeLocations range = new TScanRangeLocations().setScanRange(new TScanRange() + .setExtScanRange(new TExternalScanRange().setFileScanRange(new TFileScanRange() + .setRanges(new ArrayList<>(Collections.singletonList( + new TFileRangeDesc().setPath("file-" + i)))) + .setParams(new TFileScanRangeParams().setSrcTupleId(i))))); splits.add(split); ranges.add(range); locations.put(split, range); @@ -78,6 +98,187 @@ void setUp() throws Exception { Collections.emptyMap(), Collections.emptyList(), true, 3, 3); } + private Object capacityMonitor() throws Exception { + Field field = SplitAssignment.class.getDeclaredField("queueSpaceListeners"); + field.setAccessible(true); + return field.get(assignment); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void uncontendedChannelDoesNotAcquireCapacityMonitor(boolean produce) throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + synchronized (capacityMonitor()) { + Future task = executor.submit(() -> { + if (produce) { + Assertions.assertTrue(assignment.offerBatch(backend, + Collections.singletonList(ranges.get(0)), () -> {})); + } else { + assignment.notifyQueueSpace(backend); + } + }); + task.get(5, TimeUnit.SECONDS); + } + } finally { + executor.shutdownNow(); + } + } + + @Test + void capacityReleasedBeforeWaiterRegistrationIsRetried() throws Exception { + BlockingQueue> queue = assignment.getAssignedSplits(backend); + while (queue.remainingCapacity() > 0) { + queue.add(Collections.singletonList(ranges.get(0))); + } + AtomicInteger resumed = new AtomicInteger(); + AtomicReference failure = new AtomicReference<>(); + AtomicReference offered = new AtomicReference<>(); + Thread producer = new Thread(() -> { + try { + offered.set(assignment.offerBatch(backend, Collections.singletonList(ranges.get(1)), + resumed::incrementAndGet)); + } catch (Throwable t) { + failure.set(t); + } + }); + Object monitor = capacityMonitor(); + synchronized (monitor) { + producer.start(); + Awaitility.await().atMost(Duration.ofSeconds(5)).until(() -> { + java.lang.management.ThreadInfo info = ManagementFactory.getThreadMXBean() + .getThreadInfo(producer.getId()); + return info != null && info.getThreadState() == Thread.State.BLOCKED + && info.getLockInfo().getIdentityHashCode() == System.identityHashCode(monitor); + }); + Assertions.assertNotNull(queue.poll()); + assignment.notifyQueueSpace(backend); + } + producer.join(5000); + Assertions.assertFalse(producer.isAlive()); + Assertions.assertNull(failure.get()); + Assertions.assertEquals(Boolean.TRUE, offered.get()); + Assertions.assertEquals(0, queue.remainingCapacity()); + assignment.notifyQueueSpace(backend); + Assertions.assertEquals(0, resumed.get(), "the successful retry must remove its capacity callback"); + } + + @Test + void backendChannelsUseStableIdsAcrossEndpointChanges() throws Exception { + Backend first = new Backend(11L, "127.0.0.1", 9050); + Backend second = new Backend(12L, "127.0.0.2", 9050); + BlockingQueue> firstQueue = assignment.getAssignedSplits(first); + BlockingQueue> secondQueue = assignment.getAssignedSplits(second); + AtomicInteger resumed = new AtomicInteger(); + Runnable callback = resumed::incrementAndGet; + int capacity = firstQueue.remainingCapacity(); + for (int i = 0; i < capacity; i++) { + Assertions.assertTrue(assignment.offerBatch(first, Collections.singletonList(ranges.get(0)), callback)); + } + Assertions.assertFalse(assignment.offerBatch(first, Collections.singletonList(ranges.get(0)), callback)); + first.setAlive(true); + first.setBePort(9061); + Backend sameId = new Backend(11L, "127.0.0.3", 9150); + Assertions.assertSame(firstQueue, assignment.getAssignedSplits(sameId)); + Assertions.assertSame(firstQueue, assignment.getAssignedSplits(first)); + Assertions.assertNotSame(firstQueue, secondQueue); + assignment.notifyQueueSpace(second); + Assertions.assertEquals(0, resumed.get()); + Assertions.assertNotNull(firstQueue.poll()); + assignment.notifyQueueSpace(sameId); + Assertions.assertEquals(1, resumed.get(), "capacity wakeups must find the same backend by ID"); + Assertions.assertEquals(capacity - 1, firstQueue.size()); + assignment.notifyQueueSpace(first); + Assertions.assertEquals(1, resumed.get(), "a registered callback is removed after notification"); + } + + @Test + void bulkDrainRespectsRequestAndResumesBackpressuredProducer() throws Exception { + assignment.addToQueue(splits); + AtomicInteger resumed = new AtomicInteger(); + Assertions.assertFalse(assignment.offerBatch(backend, Collections.singletonList(ranges.get(0)), + resumed::incrementAndGet)); + SplitSource source = new SplitSource(backend, assignment, 1000); + Assertions.assertEquals(ranges.subList(0, 6), source.getNextBatch(6)); + Assertions.assertEquals(1, resumed.get()); + Assertions.assertEquals(1, assignment.getAssignedSplits(backend).size(), + "batches beyond the BE request must stay in the bounded queue"); + assignment.finishSchedule(); + Assertions.assertEquals(ranges.subList(6, 8), source.getNextBatch(6)); + Assertions.assertTrue(source.getNextBatch(6).isEmpty()); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void bulkDrainObservesTerminalStatePublishedByCapacityCallback(boolean fail) throws Exception { + assignment.addToQueue(splits); + Assertions.assertFalse(assignment.offerBatch(backend, Collections.singletonList(ranges.get(0)), () -> { + if (fail) { + assignment.setException(new UserException("bulk drain failure")); + } else { + assignment.stop(); + } + })); + SplitSource source = new SplitSource(backend, assignment, 1000); + if (fail) { + UserException error = Assertions.assertThrows(UserException.class, () -> source.getNextBatch(6)); + Assertions.assertTrue(error.getMessage().contains("bulk drain failure")); + } else { + Assertions.assertTrue(source.getNextBatch(6).isEmpty()); + Assertions.assertTrue(source.getNextBatch(6).isEmpty()); + } + } + + @Test + void initialFetchLimitIsIndependentOfProductionAndSubsequentFetches() throws Exception { + // All eight splits are already produced without a BE request, including a partially drained batch. + assignment.addToQueue(splits); + assignment.finishSchedule(); + SplitSource source = new SplitSource(backend, assignment, 1000, 2); + Assertions.assertEquals(ranges.subList(0, 2), source.getNextBatch(100)); + Assertions.assertEquals(ranges.subList(2, 8), source.getNextBatch(100)); + Assertions.assertTrue(source.getNextBatch(100).isEmpty()); + } + + @Test + void firstFetchNeverExceedsBackendRequest() throws Exception { + assignment.addToQueue(splits); + assignment.finishSchedule(); + SplitSource source = new SplitSource(backend, assignment, 1000, 5); + Assertions.assertEquals(ranges.subList(0, 1), source.getNextBatch(1)); + Assertions.assertEquals(ranges.subList(1, 8), source.getNextBatch(100)); + } + + @Test + void initialLimitIsAppliedBeforeRangeCompaction() throws Exception { + for (TScanRangeLocations range : ranges) { + range.getScanRange().getExtScanRange().getFileScanRange().setParams(new TFileScanRangeParams()); + } + assignment.addToQueue(splits); + assignment.finishSchedule(); + SplitSource source = new SplitSource(backend, assignment, 1000, 2); + List first = source.getNextBatch(100); + Assertions.assertEquals(1, first.size()); + Assertions.assertEquals(2, first.get(0).getScanRange().getExtScanRange().getFileScanRange().getRangesSize()); + List rest = source.getNextBatch(100); + Assertions.assertEquals(1, rest.size()); + Assertions.assertEquals(6, rest.get(0).getScanRange().getExtScanRange().getFileScanRange().getRangesSize()); + } + + @Test + void initialFetchSizeValidationAndForwarding() throws Exception { + SessionVariable variables = new SessionVariable(); + Assertions.assertEquals(64, variables.getSplitInitialBatchSize()); + variables.checkSplitInitialBatchSize("1"); + Assertions.assertThrows(Exception.class, () -> variables.checkSplitInitialBatchSize("0")); + Assertions.assertThrows(Exception.class, () -> variables.checkSplitInitialBatchSize("-1")); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new SplitSource(backend, assignment, 1000, 0)); + VarAttrDef.VarAttr annotation = SessionVariable.class + .getField("splitInitialBatchSize").getAnnotation(VarAttrDef.VarAttr.class); + Assertions.assertTrue(annotation.needForward()); + } + @Test void drainPendingBatchBeforeFinishedQueueBecomesEof() throws Exception { assignment.addToQueue(splits); @@ -198,10 +399,233 @@ void cancellationUnblocksProducerAndClosesLateResources() throws Exception { } } + @ParameterizedTest + @ValueSource(strings = {"finish", "stop", "failure", "data"}) + void queueStateChangesWakeBlockedConsumer(String change) throws Exception { + SplitSource source = new SplitSource(backend, assignment, 30000); + AtomicReference waiting = new AtomicReference<>(); + CountDownLatch started = new CountDownLatch(1); + ExecutorService executor = Executors.newSingleThreadExecutor(); + UserException failure = new UserException("manifest read failed"); + try { + Future> result = executor.submit(() -> { + waiting.set(Thread.currentThread()); + started.countDown(); + return source.getNextBatch(1); + }); + Assertions.assertTrue(started.await(5, TimeUnit.SECONDS)); + awaitConsumerWait(waiting.get()); + switch (change) { + case "finish": + assignment.finishSchedule(); + break; + case "stop": + assignment.stop(); + break; + case "failure": + assignment.setException(failure); + break; + case "data": + assignment.addToQueue(splits.subList(0, 1)); + break; + default: + throw new AssertionError(change); + } + if ("failure".equals(change)) { + ExecutionException actual = Assertions.assertThrows(ExecutionException.class, + () -> result.get(5, TimeUnit.SECONDS)); + Assertions.assertSame(failure, actual.getCause()); + } else { + Assertions.assertEquals("data".equals(change) ? ranges.subList(0, 1) : Collections.emptyList(), + result.get(5, TimeUnit.SECONDS)); + } + } finally { + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void backendWaitChannelsAreIndependentAndCompletionWakesBoth() throws Exception { + Backend otherBackend = new Backend(2L, "127.0.0.2", 9050); + SplitSource first = new SplitSource(backend, assignment, 30000); + SplitSource second = new SplitSource(otherBackend, assignment, 30000); + AtomicReference firstThread = new AtomicReference<>(); + AtomicReference secondThread = new AtomicReference<>(); + CountDownLatch started = new CountDownLatch(2); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future> firstResult = executor.submit(() -> { + firstThread.set(Thread.currentThread()); + started.countDown(); + return first.getNextBatch(1); + }); + Future> secondResult = executor.submit(() -> { + secondThread.set(Thread.currentThread()); + started.countDown(); + return second.getNextBatch(1); + }); + Assertions.assertTrue(started.await(5, TimeUnit.SECONDS)); + awaitConsumerWait(firstThread.get()); + awaitConsumerWait(secondThread.get()); + Assertions.assertTrue(assignment.offerBatch(backend, ranges.subList(0, 1), () -> {})); + Assertions.assertEquals(ranges.subList(0, 1), firstResult.get(5, TimeUnit.SECONDS)); + Assertions.assertFalse(secondResult.isDone()); + assignment.finishSchedule(); + Assertions.assertTrue(secondResult.get(5, TimeUnit.SECONDS).isEmpty()); + } finally { + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void completionRacingConsumerRegistrationDoesNotLoseWakeup() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int round = 0; round < 200; round++) { + SplitAssignment current = new SplitAssignment(Mockito.mock(FederationBackendPolicy.class), Mockito.mock(SplitGenerator.class), + (be, properties, split, keys, admission) -> ranges.get(0), + Collections.emptyMap(), Collections.emptyList(), true, 1, 1); + CountDownLatch start = new CountDownLatch(1); + Future> result = executor.submit(() -> { + start.await(); + return new SplitSource(backend, current, 30000).getNextBatch(1); + }); + Future finished = executor.submit(() -> { + start.await(); + current.finishSchedule(); + return null; + }); + start.countDown(); + finished.get(5, TimeUnit.SECONDS); + Assertions.assertTrue(result.get(5, TimeUnit.SECONDS).isEmpty()); + } + } finally { + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void interruptedProducerPublishesFailureToWaitingConsumer() throws Exception { + assignment.addToQueue(splits); + SplitSource source = new SplitSource(new Backend(2L, "127.0.0.2", 9050), assignment, 30000); + AtomicReference producerThread = new AtomicReference<>(); + AtomicReference consumerThread = new AtomicReference<>(); + CountDownLatch started = new CountDownLatch(2); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future producer = executor.submit(() -> { + producerThread.set(Thread.currentThread()); + started.countDown(); + assignment.addToQueue(splits.subList(0, 1)); + return null; + }); + Future consumer = executor.submit(() -> { + consumerThread.set(Thread.currentThread()); + started.countDown(); + return source.getNextBatch(1); + }); + Assertions.assertTrue(started.await(5, TimeUnit.SECONDS)); + awaitConsumerWait(producerThread.get()); + awaitConsumerWait(consumerThread.get()); + producerThread.get().interrupt(); + producer.get(5, TimeUnit.SECONDS); + ExecutionException actual = Assertions.assertThrows(ExecutionException.class, + () -> consumer.get(5, TimeUnit.SECONDS)); + Assertions.assertInstanceOf(UserException.class, actual.getCause()); + Assertions.assertTrue(actual.getCause().getMessage().contains("interrupted")); + } finally { + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void cancellationWakesConsumerBeforeClosingResources() throws Exception { + java.util.concurrent.CompletableFuture allowClose = new java.util.concurrent.CompletableFuture<>(); + CountDownLatch closing = new CountDownLatch(1); + assignment.addCloseable(() -> { + closing.countDown(); + allowClose.join(); + }); + AtomicReference consumerThread = new AtomicReference<>(); + CountDownLatch started = new CountDownLatch(1); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future> consumer = executor.submit(() -> { + consumerThread.set(Thread.currentThread()); + started.countDown(); + return new SplitSource(backend, assignment, 30000).getNextBatch(1); + }); + Assertions.assertTrue(started.await(5, TimeUnit.SECONDS)); + awaitConsumerWait(consumerThread.get()); + Future stopping = executor.submit(assignment::stop); + Assertions.assertTrue(closing.await(5, TimeUnit.SECONDS)); + Assertions.assertTrue(consumer.get(5, TimeUnit.SECONDS).isEmpty()); + Assertions.assertFalse(stopping.isDone()); + allowClose.complete(null); + stopping.get(5, TimeUnit.SECONDS); + } finally { + allowClose.complete(null); + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static void awaitConsumerWait(Thread thread) { + Assertions.assertTimeoutPreemptively(Duration.ofSeconds(5), () -> { + while (thread.getState() != Thread.State.TIMED_WAITING) { + Assertions.assertTrue(thread.isAlive()); + Thread.yield(); + } + }); + } + @Test void invalidRequestSizeDoesNotReturnFalseEof() { SplitSource source = new SplitSource(backend, assignment, 1000); Assertions.assertThrows(IllegalArgumentException.class, () -> source.getNextBatch(0)); Assertions.assertThrows(IllegalArgumentException.class, () -> source.getNextBatch(-1)); } + + @Test + void demandHintStartsWithTheSecondRequestAndStopsOnCancellation() throws Exception { + AtomicInteger requests = new AtomicInteger(); + assignment.setMoreSplitsRequestedCallback(requests::incrementAndGet); + assignment.addToQueue(splits); + SplitSource source = new SplitSource(backend, assignment, 1000); + source.getNextBatch(2); + Assertions.assertEquals(0, requests.get()); + source.getNextBatch(2); + Assertions.assertEquals(1, requests.get()); + assignment.stop(); + Assertions.assertTrue(source.getNextBatch(2).isEmpty()); + Assertions.assertEquals(1, requests.get()); + } + + @Test + void compactedResponsesPreservePendingSplitsAndRequestLimits() throws Exception { + ranges.forEach(range -> range.getScanRange().getExtScanRange().getFileScanRange() + .setParams(new TFileScanRangeParams().setSrcTupleId(1))); + assignment.addToQueue(splits); + assignment.finishSchedule(); + SplitSource source = new SplitSource(backend, assignment, 1000); + List actual = new ArrayList<>(); + for (int requested : Arrays.asList(2, 3, 3)) { + List response = source.getNextBatch(requested); + Assertions.assertEquals(1, response.size()); + List files = response.get(0).getScanRange().getExtScanRange() + .getFileScanRange().getRanges(); + Assertions.assertEquals(requested, files.size()); + actual.addAll(files); + } + List expected = new ArrayList<>(); + ranges.forEach(range -> expected.addAll(range.getScanRange().getExtScanRange() + .getFileScanRange().getRanges())); + Assertions.assertEquals(expected, actual); + Assertions.assertEquals(8, expected.size(), "compaction must not mutate queued envelopes"); + Assertions.assertTrue(source.getNextBatch(2).isEmpty()); + } } 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 index 40c09b31d96082..f51636717d65a6 100644 --- 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 @@ -21,6 +21,11 @@ import org.apache.doris.datasource.scan.FederationBackendPolicy; import org.apache.doris.spi.Split; import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TExternalScanRange; +import org.apache.doris.thrift.TFileRangeDesc; +import org.apache.doris.thrift.TFileScanRange; +import org.apache.doris.thrift.TFileScanRangeParams; +import org.apache.doris.thrift.TScanRange; import org.apache.doris.thrift.TScanRangeLocation; import org.apache.doris.thrift.TScanRangeLocations; @@ -81,8 +86,12 @@ private void verifyConcurrentStreaming(boolean consumerFirst) throws Exception { queuesFilled.countDown(); } // Preserve a unique split identity through assignment and residual batches. - return new TScanRangeLocations().setLocations(Collections.singletonList( - new TScanRangeLocation().setBackendId(split.getStart()))); + TFileScanRange files = new TFileScanRange().setParams(new TFileScanRangeParams()) + .setRanges(Collections.singletonList(new TFileRangeDesc().setPath("file-" + split.getStart()))); + return new TScanRangeLocations().setScanRange(new TScanRange().setExtScanRange( + new TExternalScanRange().setFileScanRange(files))) + .setLocations(Collections.singletonList( + new TScanRangeLocation().setBackendId(split.getStart()))); }, Collections.emptyMap(), Collections.emptyList(), true, producerBatchSize, queueCapacity); List splits = new ArrayList<>(); LocationPath path = LocationPath.of("s3://bucket/file.parquet"); diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java index a7d92cc14b39fd..6532d28bde46e5 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/planner/FederationBackendPolicyTest.java @@ -24,6 +24,7 @@ import org.apache.doris.datasource.scan.FederationBackendPolicy; import org.apache.doris.datasource.scan.NodeSelectionStrategy; import org.apache.doris.datasource.split.FileSplit; +import org.apache.doris.datasource.split.SplitWeight; import org.apache.doris.resource.computegroup.ComputeGroupMgr; import org.apache.doris.spi.Split; import org.apache.doris.system.Backend; @@ -46,6 +47,7 @@ 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.Objects; @@ -107,6 +109,52 @@ public void testRoundRobinUsesCandidateOrderAcrossBatches() throws UserException } } + @ParameterizedTest + @ValueSource(ints = {1, 3}) + public void testRedistributionAcrossStreamingBatches(int variance) throws UserException { + SystemInfoService service = new SystemInfoService(); + List backends = List.of(new Backend(12001L, "192.0.2.1", 9050), + new Backend(12002L, "192.0.2.2", 9050), new Backend(12003L, "192.0.2.3", 9050)); + for (Backend backend : backends) { + backend.setAlive(true); + service.addBackend(backend); + } + mockedEnvStatic.when(Env::getCurrentSystemInfo).thenReturn(service); + Mockito.when(env.getComputeGroupMgr()).thenReturn(new ComputeGroupMgr(service)); + int previousVariance = Config.split_assigner_max_split_num_variance; + Config.split_assigner_max_split_num_variance = variance; + try { + FederationBackendPolicy policy = new FederationBackendPolicy(); + policy.init(); + Map expectedWeights = new HashMap<>(); + backends.forEach(backend -> expectedWeights.put(backend, 0L)); + Set assigned = Collections.newSetFromMap(new IdentityHashMap<>()); + int splitCount = 0; + // Locality initially sends every split to the same BE. Repeated small batches cross the + // variance boundary in both directions; a larger batch also exercises repeated heap updates. + for (int batchSize : List.of(1, 1, 1, 1, 3, 128, 1, 1, 1)) { + List splits = new ArrayList<>(); + for (int i = 0; i < batchSize; i++) { + splits.add(new FileSplit(LocationPath.of("hdfs://example/streaming/" + splitCount++), + 0, 1, 1, 0, new String[]{backends.get(0).getHost()}, Collections.emptyList())); + } + Multimap assignment = policy.computeScanRangeAssignment(splits); + Assertions.assertEquals(batchSize, assignment.size()); + for (Map.Entry entry : assignment.entries()) { + Assertions.assertTrue(assigned.add(entry.getValue())); + expectedWeights.merge(entry.getKey(), entry.getValue().getSplitWeight().getRawValue(), Long::sum); + } + Assertions.assertEquals(expectedWeights, policy.getAssignedWeightPerBackend()); + long min = Collections.min(expectedWeights.values()); + long max = Collections.max(expectedWeights.values()); + Assertions.assertTrue(max - min <= SplitWeight.rawValueForStandardSplitCount(variance)); + } + Assertions.assertEquals(splitCount, assigned.size()); + } finally { + Config.split_assigner_max_split_num_variance = previousVariance; + } + } + @Test public void testRoundRobinPreservesOrderForBackendsOnSameHost() throws UserException { SystemInfoService service = new SystemInfoService(); @@ -365,6 +413,74 @@ public void testConsistentHashSingleSplitTieIsStableAcrossPolicies() throws User } } + @ParameterizedTest + @ValueSource(ints = {1, 2, 3, 8, 32, 128}) + public void testConsistentHashCandidateLimit(int backendCount) throws UserException { + checkConsistentHashCandidateLimit(backendCount); + } + + private void checkConsistentHashCandidateLimit(int backendCount) throws UserException { + int originalCandidates = Config.split_assigner_min_consistent_hash_candidate_num; + int originalVirtualNodes = Config.split_assigner_virtual_node_number; + try { + Config.split_assigner_virtual_node_number = 256; + SystemInfoService service = new SystemInfoService(); + List backends = new ArrayList<>(); + for (int i = 0; i < backendCount; i++) { + HashCountingBackend backend = new HashCountingBackend(80000L + backendCount * 10 + i, + "172.30.1." + i, 9050); + backend.setAlive(true); + service.addBackend(backend); + backends.add(backend); + } + mockedEnvStatic.when(Env::getCurrentSystemInfo).thenReturn(service); + Mockito.when(env.getComputeGroupMgr()).thenReturn(new ComputeGroupMgr(service)); + List splits = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + splits.add(new FileSplit(LocationPath.of("s3://bucket/part-" + i + ".parquet"), + 0, 1024, 1024, 0, null, Collections.emptyList())); + } + + Config.split_assigner_min_consistent_hash_candidate_num = backendCount; + FederationBackendPolicy expectedPolicy = new FederationBackendPolicy(NodeSelectionStrategy.CONSISTENT_HASHING); + expectedPolicy.init(); + backends.forEach(backend -> backend.hashCalls = 0); + Multimap expected = expectedPolicy.computeScanRangeAssignment(new ArrayList<>(splits)); + int expectedHashCalls = backends.stream().mapToInt(backend -> backend.hashCalls).sum(); + + Config.split_assigner_min_consistent_hash_candidate_num = backendCount + 1; + FederationBackendPolicy actualPolicy = new FederationBackendPolicy(NodeSelectionStrategy.CONSISTENT_HASHING); + actualPolicy.init(); + backends.forEach(backend -> backend.hashCalls = 0); + Multimap actual = actualPolicy.computeScanRangeAssignment(new ArrayList<>(splits)); + int actualHashCalls = backends.stream().mapToInt(backend -> backend.hashCalls).sum(); + + Assertions.assertEquals(expected, actual); + Assertions.assertEquals(splits.size(), actual.size()); + // Count operations rather than elapsed time so the regression test is deterministic. + Assertions.assertTrue(actualHashCalls <= expectedHashCalls * 2, + "Excess candidates must not trigger a full virtual-node traversal: " + actualHashCalls + + " hash calls versus " + expectedHashCalls); + } finally { + Config.split_assigner_min_consistent_hash_candidate_num = originalCandidates; + Config.split_assigner_virtual_node_number = originalVirtualNodes; + } + } + + private static class HashCountingBackend extends Backend { + private int hashCalls; + + HashCountingBackend(long id, String host, int heartbeatPort) { + super(id, host, heartbeatPort); + } + + @Override + public int hashCode() { + hashCalls++; + return super.hashCode(); + } + } + public static void sortSplits(List splits) { splits.sort((split1, split2) -> { int pathComparison = split1.getPathString().compareTo(split2.getPathString()); diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java index aed8110a7b3b4c..7344b9cae5c081 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java @@ -36,9 +36,14 @@ import org.apache.doris.common.profile.RuntimeProfile; import org.apache.doris.common.profile.SummaryProfile; import org.apache.doris.common.util.DebugUtil; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.common.util.TimeUtils; +import org.apache.doris.connector.ConnectorSessionBuilder; +import org.apache.doris.connector.spi.ConnectorSession; +import org.apache.doris.connector.spi.ConnectorTaskScheduler; import org.apache.doris.mysql.authenticate.TestLogAppender; import org.apache.doris.nereids.NereidsPlanner; +import org.apache.doris.nereids.StatementContext; import org.apache.doris.nereids.glue.LogicalPlanAdapter; import org.apache.doris.nereids.properties.PhysicalProperties; import org.apache.doris.nereids.trees.plans.LimitPhase; @@ -78,6 +83,7 @@ import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.api.parallel.ResourceLock; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.EnumSource; import org.junit.jupiter.params.provider.ValueSource; import org.mockito.MockedConstruction; @@ -134,6 +140,140 @@ public void testCommittedTsoErrorSurvivesPlannerWrapping() throws Exception { } } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @ResourceLock("global") + public void testQueryCallbacksWithoutCoordinator(boolean planningFailure) throws Exception { + ConnectContext context = Mockito.spy(connectContext); + context.getState().reset(); + Mockito.doReturn(Mockito.mock(ResultSender.class)).when(context).getResultSender(); + TUniqueId queryId = new TUniqueId(0x22050L, planningFailure ? 1L : 2L); + AtomicInteger cleanup = new AtomicInteger(); + QeProcessorImpl.INSTANCE.registerQueryFinishCallback(DebugUtil.printId(queryId), cleanup::incrementAndGet); + try { + StmtExecutor executor = new StmtExecutor(context, "explain select 1"); + if (planningFailure) { + try (MockedConstruction planners = Mockito.mockConstruction(NereidsPlanner.class, + (planner, construction) -> Mockito.doThrow(new NereidsException("planning failed", + new UserException("planning failed"))) + .when(planner).plan(Mockito.any(StatementBase.class), Mockito.any(TQueryOptions.class)))) { + executor.execute(queryId); + Assertions.assertEquals(1, planners.constructed().size()); + } + Assertions.assertEquals(QueryState.MysqlStateType.ERR, context.getState().getStateType()); + } else { + executor.execute(queryId); + Assertions.assertNotEquals(QueryState.MysqlStateType.ERR, context.getState().getStateType(), + context.getState().getErrorMessage()); + } + Assertions.assertNull(QeProcessorImpl.INSTANCE.getCoordinator(queryId)); + Assertions.assertEquals(1, cleanup.get(), "planning-only attempts must release registered callbacks"); + QeProcessorImpl.INSTANCE.unregisterQuery(queryId); + Assertions.assertEquals(1, cleanup.get(), "later finalization must not repeat cleanup"); + } finally { + QeProcessorImpl.INSTANCE.unregisterQuery(queryId); + } + } + + @Test + @ResourceLock("global") + public void testPlanningFailureClosesTaskGroupBeforeScanRegistration() throws Exception { + ConnectContext context = Mockito.spy(connectContext); + context.getState().reset(); + Mockito.doReturn(Mockito.mock(ResultSender.class)).when(context).getResultSender(); + StmtExecutor executor = new StmtExecutor(context, "explain select 1"); + ConnectorSession session = ConnectorSessionBuilder.from(context).withCatalogId(1).build(); + FairTaskExecutor pool = context.getEnv().getExtMetaCacheMgr() + .getQueryTaskExecutor(ConnectorTaskScheduler.Stage.METADATA); + int groupsBefore = pool.getGroupCount(); + ConnectorTaskScheduler group = session.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA); + Assertions.assertEquals(groupsBefore + 1, pool.getGroupCount()); + TUniqueId queryId = new TUniqueId(0x22050L, 3L); + try (MockedConstruction planners = Mockito.mockConstruction(NereidsPlanner.class, + (planner, construction) -> Mockito.doThrow(new NereidsException("planning failed", + new UserException("planning failed"))) + .when(planner).plan(Mockito.any(StatementBase.class), Mockito.any(TQueryOptions.class)))) { + executor.execute(queryId); + Assertions.assertEquals(1, planners.constructed().size()); + Assertions.assertEquals(QueryState.MysqlStateType.ERR, context.getState().getStateType()); + Assertions.assertTrue(group.isClosed()); + Assertions.assertEquals(groupsBefore, pool.getGroupCount()); + } finally { + session.getStatementScope().closeAll(); + QeProcessorImpl.INSTANCE.unregisterQuery(queryId); + } + } + + @ParameterizedTest + @CsvSource({"false, CANCELLED", "true, CANCELLED", "false, TIMEOUT", "true, TIMEOUT"}) + @ResourceLock("global") + public void testPlanningCancellationClosesCapturedTaskGroups(boolean cancelBeforeCapture, + TStatusCode statusCode) throws Exception { + ConnectContext context = Mockito.spy(connectContext); + context.setStatementContext(new StatementContext(context, new OriginStatement("select 1", 0))); + context.setQueryId(new TUniqueId(0x22051L, statusCode.getValue() * 2L + (cancelBeforeCapture ? 1 : 0))); + StmtExecutor executor = new StmtExecutor(context, "select 1"); + context.setExecutor(executor); + Status reason = new Status(statusCode, "cancel before coordinator publication"); + if (cancelBeforeCapture) { + executor.cancel(reason, false); + } + ConnectorSession session = ConnectorSessionBuilder.from(context).withCatalogId(1).build(); + try { + if (cancelBeforeCapture) { + Assertions.assertThrows(java.util.concurrent.CancellationException.class, + () -> session.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT)); + } else { + ConnectorTaskScheduler split = session.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT); + ConnectorTaskScheduler metadata = session.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA); + Assertions.assertFalse(split.isClosed()); + Assertions.assertFalse(metadata.isClosed()); + executor.cancel(reason, false); + Assertions.assertTrue(split.isClosed()); + Assertions.assertTrue(metadata.isClosed()); + } + } finally { + session.getStatementScope().closeAll(); + } + } + + @Test + @ResourceLock("global") + public void testCancellationRegistrationWhileCallbackIsRunning() throws Exception { + StmtExecutor executor = new StmtExecutor(connectContext, "select 1"); + AtomicInteger first = new AtomicInteger(); + AtomicInteger late = new AtomicInteger(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + executor.registerQueryCancellationCallback(() -> { + first.incrementAndGet(); + entered.countDown(); + try { + Assertions.assertTrue(release.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }); + ExecutorService worker = Executors.newSingleThreadExecutor(); + try { + Future cancelling = worker.submit(() -> executor.cancel( + new Status(TStatusCode.CANCELLED, "cancel during callback registration"), false)); + Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS)); + executor.registerQueryCancellationCallback(late::incrementAndGet); + Assertions.assertEquals(1, late.get()); + release.countDown(); + cancelling.get(5, TimeUnit.SECONDS); + executor.cancel(new Status(TStatusCode.TIMEOUT, "later cancellation"), false); + Assertions.assertEquals(1, first.get()); + Assertions.assertEquals(1, late.get()); + } finally { + release.countDown(); + worker.shutdownNow(); + Assertions.assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS)); + } + } + @Test public void testShow() throws Exception { StmtExecutor stmtExecutor = new StmtExecutor(connectContext, ""); @@ -587,6 +727,56 @@ public void testSuccessfulRetryRendersMergedProfileAndActualRows(@TempDir Path p } } + @Test + @ResourceLock("global") + public void outerReplanRetryReplacesClosedConnectorScope() throws Exception { + int oldRetryTime = Config.max_query_retry_time; + ConnectContext context = Mockito.spy(connectContext); + context.setStartTime(); + context.setQueryId(new TUniqueId(0x33040L, 1L)); + StmtExecutor executor = Mockito.spy(new StmtExecutor(context, "select 1")); + context.setExecutor(executor); + Method parse = StmtExecutor.class.getDeclaredMethod("parseByNereids"); + parse.setAccessible(true); + parse.invoke(executor); + AtomicInteger attempts = new AtomicInteger(); + List groups = Lists.newArrayList(); + FairTaskExecutor splitPool = context.getEnv().getExtMetaCacheMgr() + .getQueryTaskExecutor(ConnectorTaskScheduler.Stage.SPLIT); + FairTaskExecutor metadataPool = context.getEnv().getExtMetaCacheMgr() + .getQueryTaskExecutor(ConnectorTaskScheduler.Stage.METADATA); + int splitGroupsBefore = splitPool.getGroupCount(); + int metadataGroupsBefore = metadataPool.getGroupCount(); + Mockito.doAnswer(invocation -> { + context.setQueryId(invocation.getArgument(0)); + ConnectorSession session = ConnectorSessionBuilder.from(context).withCatalogId(1).build(); + groups.add(session.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT)); + groups.add(session.getTaskScheduler(ConnectorTaskScheduler.Stage.METADATA)); + try { + if (attempts.incrementAndGet() == 1) { + throw new UserException("E-230"); + } + } finally { + session.getStatementScope().closeAll(); + } + return null; + }).when(executor).execute(Mockito.any(TUniqueId.class)); + try { + Config.max_query_retry_time = 1; + Method execute = StmtExecutor.class.getDeclaredMethod("executeWithReplanRetries", TUniqueId.class); + execute.setAccessible(true); + execute.invoke(executor, context.queryId()); + Assertions.assertEquals(2, attempts.get()); + Assertions.assertEquals(4, groups.size()); + Assertions.assertTrue(groups.stream().allMatch(ConnectorTaskScheduler::isClosed)); + Assertions.assertEquals(splitGroupsBefore, splitPool.getGroupCount()); + Assertions.assertEquals(metadataGroupsBefore, metadataPool.getGroupCount()); + } finally { + Config.max_query_retry_time = oldRetryTime; + context.getStatementContext().resetConnectorStatementScope(); + } + } + private interface CoordinatorConfigurer { void configure(Coordinator coordinator) throws Exception; } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java index 62a83172f94194..c7706d9e19f3f2 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java @@ -19,25 +19,39 @@ import org.apache.doris.analysis.TupleDescriptor; import org.apache.doris.analysis.TupleId; +import org.apache.doris.common.jmockit.Deencapsulation; +import org.apache.doris.datasource.scan.FileQueryScanNode; import org.apache.doris.nereids.trees.plans.distribute.worker.DistributedPlanWorker; import org.apache.doris.nereids.trees.plans.distribute.worker.job.AssignedJob; import org.apache.doris.nereids.trees.plans.distribute.worker.job.DefaultScanSource; import org.apache.doris.nereids.trees.plans.distribute.worker.job.LocalShuffleAssignedJob; +import org.apache.doris.nereids.trees.plans.distribute.worker.job.ScanRanges; import org.apache.doris.nereids.trees.plans.distribute.worker.job.UnassignedJob; import org.apache.doris.planner.PlanNodeId; import org.apache.doris.planner.RecursiveCteScanNode; import org.apache.doris.planner.ScanNode; import org.apache.doris.planner.SortNode; +import org.apache.doris.thrift.TExternalScanRange; +import org.apache.doris.thrift.TFileRangeDesc; +import org.apache.doris.thrift.TFileScanRange; +import org.apache.doris.thrift.TFileScanRangeParams; +import org.apache.doris.thrift.TPipelineFragmentParams; +import org.apache.doris.thrift.TPipelineInstanceParams; +import org.apache.doris.thrift.TQueryOptions; import org.apache.doris.thrift.TRecCTETarget; +import org.apache.doris.thrift.TScanRange; +import org.apache.doris.thrift.TScanRangeParams; import org.apache.doris.thrift.TUniqueId; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.mockito.Mockito; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; public class ThriftPlansBuilderTest { @Test @@ -108,6 +122,35 @@ public void testBuildRecCTETargetsKeepsOneLocalShuffleInstancePerBackend() { Assertions.assertEquals(11, targets.get(1).getNodeId()); } + @Test + public void compactsOnlyTheAssignedInstanceTransportCopy() { + FileQueryScanNode node = Mockito.mock(FileQueryScanNode.class); + Mockito.when(node.getId()).thenReturn(new PlanNodeId(1)); + ScanRanges assigned = new ScanRanges(); + for (int i = 0; i < 1000; i++) { + TFileScanRange files = new TFileScanRange().setRanges(new ArrayList<>(List.of( + new TFileRangeDesc().setPath("file" + i).setStartOffset(0).setSize(10)))); + assigned.addScanRange(new TScanRangeParams().setScanRange(new TScanRange() + .setExtScanRange(new TExternalScanRange().setFileScanRange(files))), 10); + } + AssignedJob job = Mockito.mock(AssignedJob.class); + Mockito.when(job.getScanSource()).thenReturn(new DefaultScanSource(Map.of(node, assigned))); + TPipelineFragmentParams fragment = new TPipelineFragmentParams() + .setFileScanParams(Map.of(1, new TFileScanRangeParams())) + .setQueryOptions(new TQueryOptions().setMaxFileScannersConcurrency(16).setParallelInstance(1)); + TPipelineInstanceParams instance = new TPipelineInstanceParams(); + + Deencapsulation.invoke(ThriftPlansBuilder.class, "setScanSourceParam", fragment, job, instance); + + List transported = instance.getPerNodeScanRanges().get(1); + Assertions.assertTrue(transported.size() >= 16); + Assertions.assertTrue(transported.size() < 1000); + Assertions.assertEquals(1000, transported.stream().mapToInt(r -> r.getScanRange() + .getExtScanRange().getFileScanRange().getRangesSize()).sum()); + Assertions.assertEquals(1000, assigned.params.size()); + Assertions.assertEquals(10000, assigned.totalBytes); + } + private AssignedJob newAssignedJob(DistributedPlanWorker worker, TUniqueId instanceId) { AssignedJob assignedJob = Mockito.mock(AssignedJob.class); Mockito.when(assignedJob.getAssignedWorker()).thenReturn(worker); 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 index cd6e775b10bdc4..5f524016caf2e6 100644 --- 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 @@ -60,6 +60,61 @@ -- !synchronous_union_limit -- 10 +-- !streaming_1_4096_unpartitioned_aggregate -- +120 7140 120 660 + +-- !streaming_1_4096_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_1_4096_unpartitioned_limit -- +10 + +-- !streaming_1_4096_partitioned_aggregate -- +120 7140 120 660 240 + +-- !streaming_1_4096_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_1_4096_partition_range -- +80 5160 80 200 + +-- !streaming_1_4096_single_partition_limit -- +10 + +-- !streaming_1_4096_empty_partitions -- +0 \N + +-- !streaming_1_4096_join_aggregate -- +120 7140 240 + +-- !streaming_1_4096_union_aggregate -- +240 14280 + +-- !streaming_1_4096_union_limit -- +10 + -- !streaming_3_1_unpartitioned_aggregate -- 120 7140 120 660 diff --git a/regression-test/data/external_table_p0/iceberg/test_iceberg_scheduled_splits.out b/regression-test/data/external_table_p0/iceberg/test_iceberg_scheduled_splits.out new file mode 100644 index 00000000000000..5b1dbbf893e720 --- /dev/null +++ b/regression-test/data/external_table_p0/iceberg/test_iceberg_scheduled_splits.out @@ -0,0 +1,49 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !1_cold_limit -- +10 + +-- !1_warm_limit -- +10 + +-- !1_cold_full -- +192 192 18336 + +-- !1_warm_full -- +192 192 18336 + +-- !1_cold_multi_scan -- +384 36672 + +-- !1_cold_multi_scan_limit -- +10 + +-- !1_warm_multi_scan -- +384 36672 + +-- !1_warm_multi_scan_limit -- +10 + +-- !4096_cold_limit -- +10 + +-- !4096_warm_limit -- +10 + +-- !4096_cold_full -- +192 192 18336 + +-- !4096_warm_full -- +192 192 18336 + +-- !4096_cold_multi_scan -- +384 36672 + +-- !4096_cold_multi_scan_limit -- +10 + +-- !4096_warm_multi_scan -- +384 36672 + +-- !4096_warm_multi_scan_limit -- +10 + 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 index 7a0fca1119397b..6584101f101de0 100644 --- 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 @@ -31,7 +31,8 @@ suite("test_hive_streaming_splits", "p0,external") { '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"] + "num_partitions_in_batch_mode", "split_batch_size", "split_queue_capacity", + "split_initial_batch_size"] def saved = settings.collectEntries { name -> [(name): sql("show variables like '${name}'")[0][1]] } try { sql "create database if not exists ${catalog}.hive_streaming_splits" @@ -100,9 +101,10 @@ suite("test_hive_streaming_splits", "p0,external") { 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]]) { + for (def limits : [[1, 4096, 64], [3, 1, 1], [7, 2, 2], [128, 32, 3]]) { sql "set split_batch_size=${limits[0]}" sql "set split_queue_capacity=${limits[1]}" + sql "set split_initial_batch_size=${limits[2]}" explain { sql "select * from streaming_unpartitioned limit 10" contains "(approximate)inputSplitNum" @@ -121,6 +123,10 @@ suite("test_hive_streaming_splits", "p0,external") { 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_initial_batch_size=0" + exception "split_initial_batch_size value should greater than or equal 1" + } test { sql "set split_batch_size=0" exception "split_batch_size value should greater than or equal 1" diff --git a/regression-test/suites/external_table_p0/iceberg/test_iceberg_scheduled_splits.groovy b/regression-test/suites/external_table_p0/iceberg/test_iceberg_scheduled_splits.groovy new file mode 100644 index 00000000000000..417f29b7a094cc --- /dev/null +++ b/regression-test/suites/external_table_p0/iceberg/test_iceberg_scheduled_splits.groovy @@ -0,0 +1,78 @@ +// 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_iceberg_scheduled_splits", "p0,external,iceberg") { + if (!"true".equalsIgnoreCase(context.config.otherConfigs.get("enableIcebergTest"))) { + return + } + String catalog = context.config.otherConfigs.get("icebergSplitCatalog") + if (!catalog) { + catalog = "iceberg_scheduled_splits" + String host = context.config.otherConfigs.get("externalEnvIp") + String restPort = context.config.otherConfigs.get("iceberg_rest_uri_port") + String minioPort = context.config.otherConfigs.get("iceberg_minio_port") + sql "drop catalog if exists ${catalog}" + sql """create catalog ${catalog} properties ( + 'type'='iceberg', 'iceberg.catalog.type'='rest', 'uri'='http://${host}:${restPort}', + 's3.access_key'='admin', 's3.secret_key'='password', + 's3.endpoint'='http://${host}:${minioPort}', 's3.region'='us-east-1', + 'meta.cache.iceberg.manifest.enable'='true')""" + } + def settings = ["enable_external_table_batch_mode", "num_files_in_batch_mode", + "split_batch_size", "split_queue_capacity", "enable_sql_cache"] + def saved = settings.collectEntries { name -> [(name): sql("show variables like '${name}'")[0][1]] } + try { + sql "create database if not exists ${catalog}.iceberg_scheduled_splits" + sql "use ${catalog}.iceberg_scheduled_splits" + sql "drop table if exists scheduled_manifest_scan" + sql "create table scheduled_manifest_scan (id bigint) engine=iceberg" + // Separate commits produce more manifests than one prefetch window can retain. + for (int batch = 0; batch < 12; batch++) { + sql """insert into scheduled_manifest_scan + select number + ${batch * 16} from numbers('number'='16')""" + } + sql "set enable_external_table_batch_mode=true" + sql "set num_files_in_batch_mode=1" + sql "set enable_sql_cache=false" + sql "set split_batch_size=1" + for (int capacity : [1, 4096]) { + sql "set split_queue_capacity=${capacity}" + for (String kind : ["limit", "full", "multi_scan"]) { + sql "refresh catalog ${catalog}" + for (String cache : ["cold", "warm"]) { + if (kind == "limit") { + quickTest("${capacity}_${cache}_limit", """select count(*) from + (select id from scheduled_manifest_scan limit 10) limited""", true) + } else if (kind == "full") { + quickTest("${capacity}_${cache}_full", """select count(*), count(distinct id), sum(id) + from scheduled_manifest_scan""", true) + } else { + // Both scan nodes share the query's task group, including on early stop. + quickTest("${capacity}_${cache}_multi_scan", """select count(*), sum(id) from + (select id from scheduled_manifest_scan union all + select id from scheduled_manifest_scan) combined""", true) + quickTest("${capacity}_${cache}_multi_scan_limit", """select count(*) from + (select id from scheduled_manifest_scan union all + select id from scheduled_manifest_scan limit 10) limited""", true) + } + } + } + } + } finally { + saved.each { name, value -> sql "set ${name}=${value}" } + } +}