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 0e54ff317f30eb..4bfa681adc82c8 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 @@ -2472,6 +2472,17 @@ public class Config extends ConfigBase { @ConfField(mutable = false, masterOnly = false) public static int max_external_cache_loader_thread_pool_size = 64; + @ConfField(callback = PositiveIntConfHandler.class, description = "Positive worker count for fair split production") + public static int external_scan_split_threads = 64; + + @ConfField(callback = PositiveIntConfHandler.class, + description = "Positive worker count for independent metadata tasks") + public static int external_scan_metadata_threads = 64; + + @ConfField(callback = PositiveIntConfHandler.class, + description = "Positive 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-common/src/main/java/org/apache/doris/common/ConfigBase.java b/fe/fe-common/src/main/java/org/apache/doris/common/ConfigBase.java index 2c65e083285bef..51d57dbf685540 100644 --- a/fe/fe-common/src/main/java/org/apache/doris/common/ConfigBase.java +++ b/fe/fe-common/src/main/java/org/apache/doris/common/ConfigBase.java @@ -104,6 +104,21 @@ private static String checkedOption(Field field, String confVal) throws ConfigEx } } + public static class PositiveIntConfHandler implements ConfHandler { + @Override + public void handle(Field field, String confVal) throws Exception { + setConfigField(field, checkedValue(field, confVal)); + } + + private static String checkedValue(Field field, String confVal) throws ConfigException { + int value = Integer.parseInt(confVal.trim()); + if (value <= 0) { + throw new ConfigException("Config '" + field.getName() + "' must be greater than 0"); + } + return Integer.toString(value); + } + } + static class CommaSeparatedIntersectConfHandler implements ConfHandler { @Override public void handle(Field field, String newVal) throws Exception { @@ -296,13 +311,16 @@ private static void setFields(Properties props, boolean isLdapConfig) throws Exc } try { - // Runtime callbacks can have side effects; only options validation is safe here. + // Runtime callbacks can have side effects; only pure value validators run here. if (anno.callback() == OptionsConfHandler.class) { confVal = OptionsConfHandler.checkedOption(f, confVal); } if (confVal.isEmpty()) { continue; } + if (anno.callback() == PositiveIntConfHandler.class) { + confVal = PositiveIntConfHandler.checkedValue(f, confVal); + } setConfigField(f, confVal); } catch (Exception e) { String msg = String.format("Failed to set config, name: %s, value: %s", f.getName(), confVal); diff --git a/fe/fe-common/src/test/java/org/apache/doris/common/ConfigTest.java b/fe/fe-common/src/test/java/org/apache/doris/common/ConfigTest.java index 9f1029bdaa2310..5cd5251e563873 100644 --- a/fe/fe-common/src/test/java/org/apache/doris/common/ConfigTest.java +++ b/fe/fe-common/src/test/java/org/apache/doris/common/ConfigTest.java @@ -22,12 +22,17 @@ import org.junit.Test; import org.junit.jupiter.api.Assertions; +import java.lang.reflect.Field; import java.nio.file.Files; import java.nio.file.Path; import java.util.List; import java.util.Map; public class ConfigTest { + private static final String[] EXTERNAL_SCAN_SCHEDULER_CONFIGS = { + "external_scan_split_threads", "external_scan_metadata_threads", "external_scan_task_queue_capacity" + }; + @BeforeClass public static void setUp() throws Exception { Config config = new Config(); @@ -153,6 +158,64 @@ public void testOptionsConfigIsValidatedInCustomStartupFile() throws Exception { } } + @Test + public void testExternalScanSchedulerRejectsNonPositiveStartupValues() throws Exception { + assertExternalScanSchedulerRejectsNonPositiveValues(false); + } + + @Test + public void testExternalScanSchedulerRejectsNonPositiveCustomStartupValues() throws Exception { + assertExternalScanSchedulerRejectsNonPositiveValues(true); + } + + private void assertExternalScanSchedulerRejectsNonPositiveValues(boolean custom) throws Exception { + Config config = new Config(); + config.init(configFile("").toString()); + for (String name : EXTERNAL_SCAN_SCHEDULER_CONFIGS) { + Field field = Config.class.getField(name); + int original = field.getInt(null); + try { + for (String value : new String[] {"0", "-1"}) { + Path invalid = configFile(name + "=" + value + "\n"); + IllegalArgumentException error = Assertions.assertThrows(IllegalArgumentException.class, () -> { + if (custom) { + config.initCustom(invalid.toString()); + } else { + config.init(invalid.toString()); + } + }); + Assertions.assertTrue(error.getMessage().contains(name)); + Assertions.assertInstanceOf(ConfigException.class, error.getCause()); + Assertions.assertTrue(error.getCause().getMessage().contains("must be greater than 0")); + Assertions.assertEquals(original, field.getInt(null)); + } + } finally { + field.setInt(null, original); + } + } + } + + @Test + public void testExternalScanSchedulerAcceptsPositiveStartupValuesAndStaysImmutable() throws Exception { + Config config = new Config(); + for (String name : EXTERNAL_SCAN_SCHEDULER_CONFIGS) { + Field field = Config.class.getField(name); + int original = field.getInt(null); + try { + config.init(configFile(name + "=1\n").toString()); + Assertions.assertEquals(1, field.getInt(null)); + config.initCustom(configFile(name + "= 128 \n").toString()); + Assertions.assertEquals(128, field.getInt(null)); + config.init(configFile(name + "=\n").toString()); + Assertions.assertEquals(128, field.getInt(null)); + Assertions.assertThrows(ConfigException.class, () -> ConfigBase.setMutableConfig(name, "8")); + Assertions.assertEquals(128, field.getInt(null)); + } finally { + field.setInt(null, original); + } + } + } + private static Path configFile(String contents) throws Exception { Path file = Files.createTempFile("fe_ut_options_", ".conf"); file.toFile().deleteOnExit(); 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 69e35520280d1a..979c8c401c9e3c 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 @@ -19,6 +19,7 @@ import org.apache.doris.common.security.authentication.ExecutionAuthenticator; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.metric.Metric; import org.apache.doris.metric.Metric.MetricUnit; import org.apache.doris.metric.MetricLabel; @@ -75,6 +76,7 @@ public class ThreadPoolManager { private static final Logger LOG = LogManager.getLogger(ThreadPoolManager.class); private static Map nameToThreadPoolMap = Maps.newConcurrentMap(); + private static final Map nameToFairTaskExecutorMap = Maps.newConcurrentMap(); private static String[] poolMetricTypes = {"pool_size", "active_thread_num", "task_in_queue"}; @@ -85,6 +87,21 @@ public static void registerAllThreadPoolMetric() { registerThreadPoolMetric(entry.getKey(), entry.getValue()); } nameToThreadPoolMap.clear(); + for (Map.Entry entry : nameToFairTaskExecutorMap.entrySet()) { + registerThreadPoolMetric(entry.getKey(), entry.getValue()); + } + nameToFairTaskExecutorMap.clear(); + } + + public static void registerThreadPoolMetric(String poolName, FairTaskExecutor executor) { + Metric.MetricType gauge = Metric.MetricType.GAUGE; + MetricUnit nounit = MetricUnit.NOUNIT; + registerMetric(poolName, "pool_size", gauge, nounit, executor::getPoolSize); + registerMetric(poolName, "active_thread_num", gauge, nounit, executor::getActiveCount); + registerMetric(poolName, "task_in_queue", gauge, nounit, executor::getQueuedTaskCount); + registerMetric(poolName, "completed_task_count", Metric.MetricType.COUNTER, nounit, + executor::getCompletedTaskCount); + registerMetric(poolName, "task_group_num", gauge, nounit, executor::getGroupCount); } public static void registerThreadPoolMetric(String poolName, ThreadPoolExecutor threadPool) { @@ -117,6 +134,15 @@ public T getValue() { MetricRepo.DORIS_METRIC_REGISTER.addMetrics(gauge); } + public static FairTaskExecutor newDaemonFairTaskExecutor(int parallelism, String poolName, + boolean needRegisterMetric) { + FairTaskExecutor executor = new FairTaskExecutor(parallelism, namedThreadFactory(poolName)); + if (needRegisterMetric) { + nameToFairTaskExecutorMap.put(poolName, executor); + } + return executor; + } + public static ThreadPoolExecutor newDaemonCacheThreadPool(int maxNumThread, String poolName, boolean needRegisterMetric) { return newDaemonThreadPool(0, maxNumThread, KEEP_ALIVE_TIME, 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/datasource/ExternalMetaCacheMgr.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalMetaCacheMgr.java index 300f2c6c3bf911..91266c74747e7f 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,6 +21,7 @@ import org.apache.doris.common.Config; import org.apache.doris.common.Pair; import org.apache.doris.common.ThreadPoolManager; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.datasource.doris.DorisExternalMetaCache; import org.apache.doris.datasource.hive.HiveExternalMetaCache; import org.apache.doris.datasource.hudi.HudiExternalMetaCache; @@ -57,6 +58,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.locks.Lock; import java.util.function.Consumer; +import java.util.function.Supplier; import java.util.stream.Collectors; import javax.annotation.Nullable; @@ -101,6 +103,8 @@ public class ExternalMetaCacheMgr { private ExecutorService fileListingExecutor; // This executor is used to schedule the getting split tasks private ExecutorService scheduleExecutor; + private final FairTaskExecutor splitTaskExecutor; + private final FairTaskExecutor metadataTaskExecutor; private final ExternalMetaCacheRegistry cacheRegistry; private final ExternalMetaCacheRouteResolver routeResolver; private final LegacyMetaCacheFactory legacyMetaCacheFactory; @@ -138,6 +142,12 @@ public ExternalMetaCacheMgr(boolean isCheckpointCatalog) { Config.max_external_cache_loader_thread_pool_size * 1000, "scheduleExecutor", 10, true); + String executorNamePrefix = isCheckpointCatalog ? "Checkpoint" : "NotCheckpoint"; + splitTaskExecutor = ThreadPoolManager.newDaemonFairTaskExecutor(Config.external_scan_split_threads, + executorNamePrefix + "ExternalSplit", true); + metadataTaskExecutor = ThreadPoolManager.newDaemonFairTaskExecutor(Config.external_scan_metadata_threads, + executorNamePrefix + "ExternalMetadata", true); + fsCache = new FileSystemCache(); rowCountCache = new ExternalRowCountCache(rowCountRefreshExecutor); cacheRegistry = new ExternalMetaCacheRegistry(); @@ -166,6 +176,10 @@ public ExecutorService getFileListingExecutor() { return fileListingExecutor; } + public FairTaskExecutor getScanTaskExecutor(boolean metadata) { + return metadata ? metadataTaskExecutor : splitTaskExecutor; + } + public ExecutorService getScheduleExecutor() { return scheduleExecutor; } @@ -404,7 +418,7 @@ public void invalidateCatalogByEngine(long catalogId, String engine) { * and group retirement. The stripe lock is reentrant, so fenced actions may call back into * removeCatalog/rollbackCatalogProperties for the same catalog. */ - public T withCatalogLifecycleLock(long catalogId, java.util.function.Supplier action) { + public T withCatalogLifecycleLock(long catalogId, Supplier action) { Lock lifecycleLock = catalogLifecycleLocks.get(catalogId); lifecycleLock.lock(); try { diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java index e64aebe1c51cba..c9bcaf85e1bbec 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java @@ -79,6 +79,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import java.io.Closeable; import java.net.URI; import java.util.Arrays; import java.util.Collection; @@ -89,6 +90,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicReference; import java.util.function.ToLongFunction; import java.util.stream.Collectors; @@ -123,6 +125,7 @@ public abstract class FileQueryScanNode extends FileScanNode { protected FileSplitter fileSplitter; protected SummaryProfile summaryProfile; + private ScanTaskGroups.Lease scanTaskLease; // The data cache function only works for queries on Hive, Iceberg, Hudi(via HMS), and Paimon tables. // See: https://doris.incubator.apache.org/docs/dev/lakehouse/data-cache @@ -503,9 +506,8 @@ public void createScanRangeLocations() throws UserException { Preconditions.checkArgument(maxWaitTime > 0, "fetch_splits_max_wait_time_ms must be positive"); // File splits are generated lazily, and fetched by backends while scanning. // Only provide the unique ID of split source to backend. - splitAssignment = new SplitAssignment(backendPolicy, this, this::splitToScanRange, - locationProperties, pathPartitionKeys, admissionResult, - sessionVariable.splitBatchSize, sessionVariable.splitQueueCapacity); + getScanTaskLease(); + splitAssignment = createSplitAssignment(locationProperties, pathPartitionKeys, admissionResult); splitAssignment.init(); if (executor != null) { executor.getSummaryProfile().setGetSplitsFinishTime(); @@ -526,7 +528,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.splitInitialBatchSize); splitSources.add(splitSource); Env.getCurrentEnv().getSplitSourceManager().registerSplitSource(splitSource); TScanRangeLocations curLocations = newLocations(); @@ -587,7 +590,7 @@ public void createScanRangeLocations() throws UserException { } } - private TScanRangeLocations splitToScanRange( + protected TScanRangeLocations splitToScanRange( Backend backend, Map locationProperties, Split split, @@ -903,13 +906,60 @@ protected String getHiveParquetTimeZone() throws UserException { // the new structure of storage parameters expected from MS. protected abstract Map getLocationProperties() throws UserException; + protected Executor getSplitTaskExecutor() { + return scanTaskLease == null + ? Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor() + : scanTaskLease.group(false); + } + + protected boolean isScanTaskLeaseClosed() { + return scanTaskLease != null && scanTaskLease.isClosed(); + } + + protected ScanTaskGroups.Lease getScanTaskLease() { + if (scanTaskLease == null) { + ConnectContext context = ConnectContext.get(); + if (context == null || context.getStatementContext() == null) { + return null; + } + ScanTaskGroups groups = context.getStatementContext() + .getOrRegisterStatementResource(ScanTaskGroups.class, ScanTaskGroups::new); + scanTaskLease = groups.acquire(); + context.getStatementContext().getOrRegisterStatementResource(this, () -> (Closeable) () -> { + if (context.getExecutor() == null || !context.getExecutor().isDeferredForArrowFlight()) { + stop(); + } + }); + } + return scanTaskLease; + } + + protected SplitAssignment createSplitAssignment(Map properties, + List partitionKeys, boolean admission) { + return new SplitAssignment(backendPolicy, this, this::splitToScanRange, + properties, partitionKeys, admission, sessionVariable.getSplitBatchSize(), + sessionVariable.getSplitQueueCapacity()); + } + @Override public void stop() { + if (splitAssignment == null) { + if (scanTaskLease != null) { + scanTaskLease.close(); + } + return; + } if (splitAssignment != null) { - splitAssignment.stop(); - SplitSourceManager manager = Env.getCurrentEnv().getSplitSourceManager(); - for (Long sourceId : splitAssignment.getSources()) { - manager.removeSplitSource(sourceId); + try { + splitAssignment.stop(); + } finally { + if (scanTaskLease != null) { + scanTaskLease.close(); + } + SplitSourceManager manager = Env.getCurrentEnv().getSplitSourceManager(); + for (Long sourceId : splitAssignment.getSources()) { + manager.removeSplitSource(sourceId); + } } } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileSplitter.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileSplitter.java index b4e417b4f9184f..a8cede6dd6e592 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileSplitter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileSplitter.java @@ -27,18 +27,18 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import org.apache.hadoop.fs.BlockLocation; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; import java.io.IOException; +import java.io.UncheckedIOException; import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; import java.util.List; +import java.util.NoSuchElementException; import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; public class FileSplitter { - private static final Logger LOG = LogManager.getLogger(FileSplitter.class); - // If the number of files is larger than parallel instances * num of backends, // we don't need to split the file. // Otherwise, split the file to avoid local shuffle. @@ -53,7 +53,6 @@ public static boolean needSplitForCountPushdown(int parallelism, int numBackends private int maxInitialSplitNum; private final AtomicInteger remainingInitialSplitNum; - private long currentMaxSplitSize; public long getMaxInitialSplitSize() { return maxInitialSplitSize; @@ -87,132 +86,112 @@ public FileSplitter(long maxInitialSplitSize, long maxSplitSize, int maxInitialS this.maxInitialSplitSize = maxInitialSplitSize; this.maxSplitSize = maxSplitSize; this.maxInitialSplitNum = maxInitialSplitNum; - currentMaxSplitSize = maxInitialSplitSize; remainingInitialSplitNum = new AtomicInteger(maxInitialSplitNum); } - public List splitFile( - LocationPath path, - long specifiedFileSplitSize, - BlockLocation[] blockLocations, - long length, - long modificationTime, - boolean splittable, - List partitionValues, - SplitCreator splitCreator) - throws IOException { - if (length <= 0) { - // Zero-length files contain no data; skip to avoid sending empty splits to BE. - return Lists.newArrayList(); - } - // Pass splitCreator.create() to set target file split size to calculate split weight. - long targetFileSplitSize = specifiedFileSplitSize > 0 ? specifiedFileSplitSize : maxSplitSize; - if (blockLocations == null) { - blockLocations = new BlockLocation[1]; - blockLocations[0] = new BlockLocation(null, null, 0L, length); - } + public List splitFile(LocationPath path, long specifiedFileSplitSize, + BlockLocation[] blockLocations, long length, long modificationTime, boolean splittable, + List partitionValues, SplitCreator splitCreator) throws IOException { List result = Lists.newArrayList(); - TFileCompressType compressType = Util.inferFileCompressTypeByPath(path.getNormalizedLocation()); - if (!splittable || compressType != TFileCompressType.PLAIN) { - if (LOG.isDebugEnabled()) { - LOG.debug("Path {} is not splittable.", path); - } - String[] hosts = blockLocations.length == 0 ? null : blockLocations[0].getHosts(); - result.add(splitCreator.create(path, 0, length, length, - targetFileSplitSize, - modificationTime, hosts, partitionValues)); - updateCurrentMaxSplitSize(); - return result; + try { + splitFileIterator(path, specifiedFileSplitSize, blockLocations, length, modificationTime, + splittable, partitionValues, splitCreator).forEachRemaining(result::add); + } catch (UncheckedIOException e) { + throw e.getCause(); } + return result; + } - // if specified split size is not zero, split file by specified size - if (specifiedFileSplitSize > 0) { - long bytesRemaining; - for (bytesRemaining = length; (double) bytesRemaining / (double) specifiedFileSplitSize > 1.1D; - bytesRemaining -= specifiedFileSplitSize) { - int location = getBlockIndex(blockLocations, length - bytesRemaining); - String[] hosts = location == -1 ? null : blockLocations[location].getHosts(); - result.add(splitCreator.create(path, length - bytesRemaining, specifiedFileSplitSize, - length, specifiedFileSplitSize, modificationTime, hosts, partitionValues)); - } - if (bytesRemaining != 0L) { - int location = getBlockIndex(blockLocations, length - bytesRemaining); - String[] hosts = location == -1 ? null : blockLocations[location].getHosts(); - result.add(splitCreator.create(path, length - bytesRemaining, bytesRemaining, - length, specifiedFileSplitSize, modificationTime, hosts, partitionValues)); - } - return result; + /** Uses the existing file boundaries while constructing only ranges actually consumed. */ + public Iterator splitFileIterator(LocationPath path, long specifiedFileSplitSize, + BlockLocation[] blockLocations, long length, long modificationTime, boolean splittable, + List partitionValues, SplitCreator splitCreator) throws IOException { + if (length <= 0) { + return Collections.emptyIterator(); } - - // split file by block - long start = 0; + long targetFileSplitSize = specifiedFileSplitSize > 0 ? specifiedFileSplitSize : maxSplitSize; + BlockLocation[] locations = blockLocations == null + ? new BlockLocation[] {new BlockLocation(null, null, 0L, length)} : blockLocations; + boolean canSplit = splittable + && Util.inferFileCompressTypeByPath(path.getNormalizedLocation()) == TFileCompressType.PLAIN; ImmutableList.Builder blockBuilder = ImmutableList.builder(); - for (BlockLocation blockLocation : blockLocations) { - // clamp the block range - long blockStart = Math.max(start, blockLocation.getOffset()); - long blockEnd = Math.min(start + length, blockLocation.getOffset() + blockLocation.getLength()); - if (blockStart > blockEnd) { - // block is outside split range - continue; - } - if (blockStart == blockEnd && !(blockStart == start && blockEnd == start + length)) { - // skip zero-width block, except in the special circumstance: - // slice is empty, and the block covers the empty slice interval. - continue; + if (canSplit && specifiedFileSplitSize <= 0) { + for (BlockLocation block : locations) { + long blockStart = Math.max(0, block.getOffset()); + long blockEnd = Math.min(length, block.getOffset() + block.getLength()); + if (blockStart > blockEnd || blockStart == blockEnd && length > 0) { + continue; + } + blockBuilder.add(new InternalBlock(blockStart, blockEnd, block.getHosts())); } - blockBuilder.add(new InternalBlock(blockStart, blockEnd, blockLocation.getHosts())); } List blocks = blockBuilder.build(); - if (blocks.isEmpty()) { - result.add(splitCreator.create(path, 0, length, length, - targetFileSplitSize, modificationTime, null, - partitionValues)); - updateCurrentMaxSplitSize(); - return result; - } + return new Iterator() { + private long offset; + private int blockIndex; - long splitStart = start; - int currentBlockIdx = 0; - while (splitStart < start + length) { - updateCurrentMaxSplitSize(); - long splitBytes; - long remainingBlockBytes = blocks.get(currentBlockIdx).getEnd() - splitStart; - if (remainingBlockBytes <= currentMaxSplitSize) { - splitBytes = remainingBlockBytes; - } else if (currentMaxSplitSize * 2 >= remainingBlockBytes) { - // Second to last split in this block, generate two evenly sized splits - splitBytes = remainingBlockBytes / 2; - } else { - splitBytes = currentMaxSplitSize; + @Override + public boolean hasNext() { + return offset < length; } - result.add(splitCreator.create(path, splitStart, splitBytes, - length, targetFileSplitSize, modificationTime, blocks.get(currentBlockIdx).getHosts(), - partitionValues)); - splitStart += splitBytes; - if (splitStart == blocks.get(currentBlockIdx).getEnd()) { - currentBlockIdx++; - if (currentBlockIdx != blocks.size()) { - Verify.verify(splitStart == blocks.get(currentBlockIdx).getStart()); + + @Override + public Split next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + try { + long bytes; + String[] hosts; + if (!canSplit) { + bytes = length; + hosts = locations.length == 0 ? null : locations[0].getHosts(); + nextMaxSplitSize(); + } else if (specifiedFileSplitSize > 0) { + long remaining = length - offset; + bytes = (double) remaining / specifiedFileSplitSize > 1.1D + ? specifiedFileSplitSize : remaining; + int index = getBlockIndex(locations, offset); + hosts = index == -1 ? null : locations[index].getHosts(); + } else if (blocks.isEmpty()) { + bytes = length; + hosts = null; + nextMaxSplitSize(); + } else { + long maxSize = nextMaxSplitSize(); + InternalBlock block = blocks.get(blockIndex); + long remaining = block.getEnd() - offset; + bytes = remaining <= maxSize ? remaining + : maxSize * 2 >= remaining ? remaining / 2 : maxSize; + hosts = block.getHosts(); + } + Split result = splitCreator.create(path, offset, bytes, length, targetFileSplitSize, + modificationTime, hosts, partitionValues); + offset += bytes; + if (canSplit && specifiedFileSplitSize <= 0 && !blocks.isEmpty() + && offset == blocks.get(blockIndex).getEnd()) { + blockIndex++; + if (blockIndex < blocks.size()) { + Verify.verify(offset == blocks.get(blockIndex).getStart()); + } + } + return result; + } catch (IOException e) { + throw new UncheckedIOException(e); } } - } - - if (LOG.isDebugEnabled()) { - LOG.debug("Path {} includes {} splits.", path, result.size()); - } - return result; + }; } - private void updateCurrentMaxSplitSize() { - currentMaxSplitSize = maxSplitSize; + private long nextMaxSplitSize() { int cur = remainingInitialSplitNum.get(); while (cur > 0) { if (remainingInitialSplitNum.compareAndSet(cur, cur - 1)) { - currentMaxSplitSize = maxInitialSplitSize; - break; + return maxInitialSplitSize; } cur = remainingInitialSplitNum.get(); } + return maxSplitSize; } private int getBlockIndex(BlockLocation[] blkLocations, long offset) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/ScanTaskGroups.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/ScanTaskGroups.java new file mode 100644 index 00000000000000..f2076587a312d1 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/ScanTaskGroups.java @@ -0,0 +1,125 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource; + +import org.apache.doris.catalog.Env; +import org.apache.doris.common.Config; +import org.apache.doris.common.util.FairTaskExecutor; + +import java.io.Closeable; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CancellationException; + +/** Statement-owned queues. A deferred Flight scan retains its lease until coordinator cleanup. */ +public final class ScanTaskGroups implements Closeable { + private final FairTaskExecutor.Group[] groups = new FairTaskExecutor.Group[2]; + private int references = 1; + private boolean ownerClosed; + + public synchronized Lease acquire() { + if (ownerClosed) { + throw new CancellationException("Statement scan task groups are closed"); + } + references++; + return new Lease(); + } + + private synchronized FairTaskExecutor.Group group(boolean metadata) { + int index = metadata ? 1 : 0; + if (references == 0) { + throw new CancellationException("Statement scan task groups are closed"); + } + if (groups[index] == null) { + groups[index] = Env.getCurrentEnv().getExtMetaCacheMgr() + .getScanTaskExecutor(metadata).newGroup(Config.external_scan_task_queue_capacity); + } + return groups[index]; + } + + private void release() { + FairTaskExecutor.Group[] retired = null; + synchronized (this) { + if (--references == 0) { + retired = groups.clone(); + } + } + if (retired != null) { + for (FairTaskExecutor.Group group : retired) { + if (group != null) { + group.close(); + } + } + } + } + + @Override + public void close() { + synchronized (this) { + if (ownerClosed) { + return; + } + ownerClosed = true; + } + release(); + } + + public final class Lease implements Closeable { + private boolean closed; + private final List closeActions = new ArrayList<>(); + + public synchronized FairTaskExecutor.Group group(boolean metadata) { + if (closed) { + throw new CancellationException("Scan task lease is closed"); + } + return ScanTaskGroups.this.group(metadata); + } + + public synchronized boolean isClosed() { + return closed; + } + + public void onClose(Runnable action) { + synchronized (this) { + if (!closed) { + closeActions.add(action); + return; + } + } + action.run(); + } + + @Override + public void close() { + List actions; + synchronized (this) { + if (closed) { + return; + } + closed = true; + actions = new ArrayList<>(closeActions); + closeActions.clear(); + } + try { + actions.forEach(Runnable::run); + } finally { + release(); + } + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java index 4f8a62a888ab75..0d9dd746f97f02 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitAssignment.java @@ -28,8 +28,11 @@ import org.apache.logging.log4j.Logger; import java.io.Closeable; +import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Collection; +import java.util.Deque; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -65,6 +68,16 @@ public class SplitAssignment { private volatile UserException exception = null; private final List closeableResources = new ArrayList<>(); + private final Map> queueSpaceListeners = new HashMap<>(); + + // Preserve branch-4.2 defaults for generators without explicit session controls. + // Hive and Hudi supply their configured batch and queue limits through the overload below. + public SplitAssignment(FederationBackendPolicy backendPolicy, SplitGenerator splitGenerator, + SplitToScanRange splitToScanRange, Map locationProperties, + List pathPartitionKeys, boolean fileCacheAdmission) { + this(backendPolicy, splitGenerator, splitToScanRange, locationProperties, pathPartitionKeys, + fileCacheAdmission, 8, 1024); + } public SplitAssignment( FederationBackendPolicy backendPolicy, @@ -72,23 +85,19 @@ public SplitAssignment( SplitToScanRange splitToScanRange, Map locationProperties, List pathPartitionKeys, - boolean fileCacheAdmission) { - this(backendPolicy, splitGenerator, splitToScanRange, locationProperties, - pathPartitionKeys, fileCacheAdmission, 8, 1024); - } - - public SplitAssignment(FederationBackendPolicy backendPolicy, SplitGenerator splitGenerator, - SplitToScanRange splitToScanRange, Map locationProperties, - List pathPartitionKeys, boolean fileCacheAdmission, int splitBatchSize, int queueCapacity) { - Preconditions.checkArgument(splitBatchSize > 0 && queueCapacity > 0); - this.splitBatchSize = splitBatchSize; - this.queueCapacity = queueCapacity; + boolean fileCacheAdmission, + int splitBatchSize, + int queueCapacity) { + Preconditions.checkArgument(splitBatchSize > 0, "Split batch size must be positive"); + Preconditions.checkArgument(queueCapacity > 0, "Split queue capacity must be positive"); this.backendPolicy = backendPolicy; this.splitGenerator = splitGenerator; this.splitToScanRange = splitToScanRange; this.locationProperties = locationProperties; this.pathPartitionKeys = pathPartitionKeys; this.fileCacheAdmission = fileCacheAdmission; + this.splitBatchSize = splitBatchSize; + this.queueCapacity = queueCapacity; } public void init() throws UserException { @@ -121,30 +130,13 @@ public boolean needMoreSplit() { private void appendBatch(Multimap batch) throws UserException { for (Backend backend : batch.keySet()) { - if (stoppedBackends.contains(backend.getId())) { - continue; - } - Collection splits = batch.get(backend); - List locations = new ArrayList<>(splits.size()); - for (Split split : splits) { - locations.add(splitToScanRange.getScanRange(backend, locationProperties, split, pathPartitionKeys, - fileCacheAdmission)); - } - BlockingQueue queue = assignment.computeIfAbsent( - backend.getId(), be -> new LinkedBlockingQueue<>(queueCapacity)); + Collection locations = createScanRanges(backend, batch.get(backend)); for (TScanRangeLocations location : locations) { while (needMoreSplit() && !stoppedBackends.contains(backend.getId())) { try { - if (queue.offer(location, 100, TimeUnit.MILLISECONDS)) { + if (getOrCreateAssignmentQueue(backend).offer(location, 100, TimeUnit.MILLISECONDS)) { if (stoppedBackends.contains(backend.getId())) { - queue.clear(); - break; - } - synchronized (assignLock) { - if (sampleSplit == null) { - sampleSplit = splits.iterator().next(); - assignLock.notifyAll(); - } + getOrCreateAssignmentQueue(backend).clear(); } break; } @@ -157,18 +149,96 @@ 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) { + return backendPolicy.computeScanRangeAssignment(splits); + } + } + + Deque createScanRanges(Backend backend, Collection splits) + throws UserException { + Deque locations = new ArrayDeque<>(splits.size()); + for (Split split : splits) { + locations.add(splitToScanRange.getScanRange(backend, locationProperties, split, pathPartitionKeys, + fileCacheAdmission)); + } + // Publish readiness after range conversion initializes the scan-level parameters. + synchronized (assignLock) { + if (sampleSplit == null) { + sampleSplit = splits.iterator().next(); + assignLock.notifyAll(); + } + } + return locations; + } + + boolean offerBatch(Backend backend, Deque pending, Runnable resume) { + // Coordinate the failed offer and listener registration with consumer notifications so + // capacity released just before suspension cannot be missed. + synchronized (queueSpaceListeners) { + // An error/stop may race range conversion. Do not register a waiter after its wakeup. + if (!needMoreSplit()) { + return false; + } + if (stoppedBackends.contains(backend.getId())) { + pending.clear(); + removeQueueSpaceListener(backend, resume); + return true; + } + BlockingQueue queue = getOrCreateAssignmentQueue(backend); + while (!pending.isEmpty()) { + if (!queue.offer(pending.peekFirst())) { + queueSpaceListeners.computeIfAbsent(backend.getId(), key -> new HashSet<>()).add(resume); + return false; + } + pending.removeFirst(); + } + removeQueueSpaceListener(backend, resume); + return true; + } + } + public void stopBackend(long backendId) { - if (stoppedBackends.add(backendId)) { + Set listeners; + synchronized (queueSpaceListeners) { + stoppedBackends.add(backendId); BlockingQueue queue = assignment.get(backendId); if (queue != null) { queue.clear(); } - if (stoppedBackends.size() == backendPolicy.numBackends()) { - stop(); + listeners = queueSpaceListeners.remove(backendId); + } + if (listeners != null) { + listeners.forEach(Runnable::run); + } + if (stoppedBackends.size() == backendPolicy.numBackends()) { + stop(); + } + } + + void removeQueueSpaceListener(Backend backend, Runnable listener) { + synchronized (queueSpaceListeners) { + Set listeners = queueSpaceListeners.get(backend == null ? null : backend.getId()); + if (listeners != null) { + listeners.remove(listener); + if (listeners.isEmpty()) { + queueSpaceListeners.remove(backend == null ? null : backend.getId()); + } } } } + void notifyQueueSpace(Backend backend) { + Set listeners; + synchronized (queueSpaceListeners) { + listeners = queueSpaceListeners.remove(backend == null ? null : backend.getId()); + } + if (listeners != null) { + listeners.forEach(Runnable::run); + } + } + public void registerSource(long uniqueId) { sources.add(uniqueId); } @@ -182,32 +252,38 @@ public Split getSampleSplit() { } public void addToQueue(List splits) throws UserException { - if (splits.isEmpty()) { - return; - } for (int begin = 0; begin < splits.size() && needMoreSplit();) { int end = begin + Math.min(splitBatchSize, splits.size() - begin); - Multimap batch; - synchronized (assignLock) { - batch = backendPolicy.computeScanRangeAssignment(new ArrayList<>(splits.subList(begin, end))); - } - appendBatch(batch); + List batchSplits = new ArrayList<>(splits.subList(begin, end)); + appendBatch(assignBatch(batchSplits)); begin = end; } } + public int getSplitBatchSize() { + return splitBatchSize; + } + private void notifyAssignment() { synchronized (assignLock) { - assignLock.notify(); + assignLock.notifyAll(); + } + for (BlockingQueue queue : assignment.values()) { + synchronized (queue) { + queue.notifyAll(); + } } } + private BlockingQueue getOrCreateAssignmentQueue(Backend backend) { + return assignment.computeIfAbsent(backend.getId(), be -> new LinkedBlockingQueue<>(queueCapacity)); + } + public BlockingQueue getAssignedSplits(Backend backend) throws UserException { if (exception != null) { throw exception; } - BlockingQueue splits = assignment.computeIfAbsent(backend.getId(), - be -> new LinkedBlockingQueue<>(queueCapacity)); + BlockingQueue splits = getOrCreateAssignmentQueue(backend); if (scheduleFinished.get() && splits.isEmpty() || isStopped.get() || stoppedBackends.contains(backend.getId())) { return null; @@ -218,9 +294,15 @@ public BlockingQueue getAssignedSplits(Backend backend) thr public void setException(UserException e) { addUserException(e); notifyAssignment(); + Set listeners = new HashSet<>(); + synchronized (queueSpaceListeners) { + queueSpaceListeners.values().forEach(listeners::addAll); + queueSpaceListeners.clear(); + } + listeners.forEach(Runnable::run); } - private void addUserException(UserException e) { + private synchronized void addUserException(UserException e) { if (exception != null) { exception.addSuppressed(e); } else { @@ -234,23 +316,15 @@ public void finishSchedule() { } public void stop() { + if (!isStopped.compareAndSet(false, true)) { + return; + } List resources; synchronized (closeableResources) { - if (isStop()) { - return; - } - isStopped.set(true); resources = new ArrayList<>(closeableResources); closeableResources.clear(); } - resources.forEach((closeable) -> { - try { - closeable.close(); - } catch (Exception e) { - LOG.warn("close resource error:{}", e.getMessage(), e); - // ignore - } - }); + resources.forEach(this::closeResource); notifyAssignment(); } @@ -265,11 +339,7 @@ public void addCloseable(Closeable resource) { return; } } - try { - resource.close(); - } catch (Exception e) { - LOG.warn("close resource registered after stop error:{}", e.getMessage(), e); - } + closeResource(resource); } public void removeCloseable(Closeable resource) { @@ -277,4 +347,12 @@ public void removeCloseable(Closeable resource) { closeableResources.remove(resource); } } + + private void closeResource(Closeable resource) { + try { + resource.close(); + } catch (Exception e) { + LOG.warn("close resource error:{}", e.getMessage(), e); + } + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitProducer.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitProducer.java new file mode 100644 index 00000000000000..dff653ad9b6299 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitProducer.java @@ -0,0 +1,230 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource; + +import org.apache.doris.common.UserException; +import org.apache.doris.common.util.FairTaskExecutor; +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TScanRangeLocations; + +import java.io.Closeable; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Deque; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Executor; +import java.util.function.Consumer; + +/** + * Serially advances one producer batch per task, yielding the shared executor when a BE queue is + * full. Each producer retains at most one assigned batch and resumes when a consumer frees space. + */ +public class SplitProducer implements Runnable, Closeable { + private static final int INITIAL_BATCH_CAPACITY = 128; + + /** Allocate for actual work without eagerly reserving an arbitrarily large configured maximum. */ + public static List newBatch(int batchSize) { + return new ArrayList<>(Math.min(batchSize, INITIAL_BATCH_CAPACITY)); + } + + public interface BatchSource extends Closeable { + List getNextBatch() throws Exception; + + @Override + default void close() throws IOException { + } + } + + private final SplitAssignment assignment; + private final Executor executor; + private final BatchSource source; + private final Consumer completion; + private final Runnable resume = this::schedule; + private final Runnable taskSpaceAvailable = this::submit; + private Thread runner; + private final Object scheduleLock = new Object(); + private boolean scheduled; + private boolean wakeupRequested; + private boolean completed; + private volatile boolean cancelled; + private Iterator>> pending = Collections.emptyIterator(); + private Backend pendingBackend; + private Deque pendingRanges; + + public SplitProducer(SplitAssignment assignment, Executor executor, BatchSource source, + Consumer completion) { + this.assignment = assignment; + this.executor = executor; + this.source = source; + this.completion = completion; + } + + public void start() { + assignment.addCloseable(this); + schedule(); + } + + private void schedule() { + synchronized (scheduleLock) { + if (completed) { + return; + } + wakeupRequested = true; + if (scheduled) { + return; + } + scheduled = true; + } + submit(); + } + + private void submit() { + try { + if (executor instanceof FairTaskExecutor.Group) { + ((FairTaskExecutor.Group) executor).tryExecuteContinuation(this, + failure -> complete(new RuntimeException(failure)), taskSpaceAvailable); + } else { + executor.execute(this); + } + } catch (RuntimeException e) { + complete(cancelled || !assignment.needMoreSplit() ? null : e); + } + } + + @Override + public void run() { + synchronized (scheduleLock) { + if (completed) { + return; + } + runner = Thread.currentThread(); + wakeupRequested = false; + } + boolean more = false; + try { + more = produceBatch(); + } catch (Exception e) { + complete(cancelled ? null : e); + } finally { + synchronized (scheduleLock) { + runner = null; + } + } + synchronized (scheduleLock) { + if (completed || (!more && !wakeupRequested)) { + scheduled = false; + return; + } + } + submit(); + } + + private boolean produceBatch() throws Exception { + if (cancelled || !assignment.needMoreSplit()) { + complete(null); + return false; + } + if (pendingRanges == null && !pending.hasNext()) { + List batch = source.getNextBatch(); + if (batch.isEmpty()) { + complete(null); + return false; + } + pending = assignment.assignBatch(batch).asMap().entrySet().iterator(); + } + while (!cancelled && assignment.needMoreSplit()) { + if (pendingRanges == null) { + if (!pending.hasNext()) { + return true; + } + Map.Entry> entry = pending.next(); + pendingBackend = entry.getKey(); + pendingRanges = assignment.createScanRanges(pendingBackend, entry.getValue()); + } + if (!assignment.offerBatch(pendingBackend, pendingRanges, resume)) { + if (cancelled || !assignment.needMoreSplit()) { + complete(null); + } + return false; + } + pendingRanges = null; + } + complete(null); + return false; + } + + private void complete(Exception failure) { + synchronized (scheduleLock) { + if (completed) { + return; + } + completed = true; + } + finishCompletion(failure); + } + + private void finishCompletion(Exception failure) { + assignment.removeQueueSpaceListener(pendingBackend, resume); + if (executor instanceof FairTaskExecutor.Group) { + ((FairTaskExecutor.Group) executor).removeSpaceListener(taskSpaceAvailable); + } + assignment.removeCloseable(this); + pending = Collections.emptyIterator(); + pendingRanges = null; + try { + source.close(); + } catch (Exception e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + try { + completion.accept(failure); + } catch (Exception e) { + assignment.setException(new UserException(e.getMessage(), e)); + } + } + + @Override + public void close() { + Thread active; + synchronized (scheduleLock) { + cancelled = true; + if (completed) { + return; + } + active = runner; + if (active == null) { + // A queued run must observe completion before source cleanup begins. + completed = true; + } + } + if (active == null) { + finishCompletion(null); + } else { + active.interrupt(); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java index 2fcddcec8b6699..e733b0dc139b7f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/SplitSource.java @@ -22,8 +22,8 @@ import org.apache.doris.thrift.TScanRangeLocations; import com.google.common.base.Preconditions; -import com.google.common.collect.Lists; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.concurrent.BlockingQueue; @@ -50,8 +50,16 @@ public class SplitSource { private final SplitAssignment splitAssignment; private final AtomicBoolean isLastBatch; private final long maxWaitTime; + private final int initialBatchSize; + private boolean firstBatch = true; 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 batch size must be positive"); + this.initialBatchSize = initialBatchSize; Preconditions.checkArgument(maxWaitTime > 0, "fetch_splits_max_wait_time_ms must be positive"); this.uniqueId = UNIQUE_ID_GENERATOR.getAndIncrement(); this.backend = backend; @@ -83,7 +91,10 @@ public synchronized List getNextBatch(int maxBatchSize) thr if (isLastBatch.get()) { return Collections.emptyList(); } - List scanRanges = Lists.newArrayListWithExpectedSize(maxBatchSize); + if (firstBatch) { + maxBatchSize = Math.min(maxBatchSize, initialBatchSize); + } + List scanRanges = new ArrayList<>(Math.min(maxBatchSize, 128)); long started = System.nanoTime(); long waitNanos = TimeUnit.MILLISECONDS.toNanos(maxWaitTime); while (scanRanges.size() < maxBatchSize) { @@ -94,6 +105,7 @@ public synchronized List getNextBatch(int maxBatchSize) thr } TScanRangeLocations split = splits.poll(); if (split != null) { + splitAssignment.notifyQueueSpace(backend); scanRanges.add(split); continue; } @@ -108,6 +120,7 @@ public synchronized List getNextBatch(int maxBatchSize) thr split = splits.poll(Math.min(wait, TimeUnit.MILLISECONDS.toNanos(WAIT_TIME_OUT)), TimeUnit.NANOSECONDS); if (split != null) { + splitAssignment.notifyQueueSpace(backend); scanRanges.add(split); } } catch (InterruptedException e) { @@ -115,6 +128,9 @@ public synchronized List getNextBatch(int maxBatchSize) thr throw new UserException("Interrupted while waiting for scan splits", e); } } + if (!scanRanges.isEmpty()) { + firstBatch = false; + } return scanRanges; } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HiveExternalMetaCache.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HiveExternalMetaCache.java index ce288bc37aefd0..6432f1ebf13614 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HiveExternalMetaCache.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HiveExternalMetaCache.java @@ -84,6 +84,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -579,6 +580,37 @@ public HivePartitionValues getPartitionValues(PartitionValueCacheKey key) { return partitionValuesEntry.get(key.nameMapping.getCtlId()).get(key); } + /** Cheap scheduling hint only: do not list directories or initialize a cache on a miss. */ + public long getCachedFileCount(List partitions) { + if (partitions.isEmpty()) { + return 0; + } + HivePartition first = partitions.get(0); + long catalogId = first.getNameMapping().getCtlId(); + MetaCacheEntry entry = fileEntryIfInitialized(catalogId); + if (entry == null) { + return -1; + } + long fileId = Util.genIdByName(first.getNameMapping().getLocalDbName(), + first.getNameMapping().getLocalTblName()); + long count = 0; + for (HivePartition partition : partitions) { + FileCacheKey key = partition.isDummyPartition() + ? FileCacheKey.createDummyCacheKey(catalogId, fileId, + partition.getPath(), partition.getInputFormat()) + : new FileCacheKey(catalogId, fileId, partition.getPath(), + partition.getInputFormat(), partition.getPartitionValues()); + FileCacheValue value = entry.peekIfPresent(key); + if (value == null) { + return -1; + } + if (value.getFiles() != null) { + count += value.getFiles().size(); + } + } + return count; + } + public List getFilesByPartitions(List partitions, boolean withCache, boolean concurrent, @@ -836,7 +868,7 @@ private void refreshAffectedPartitionsCache(HMSExternalTable table, mergedPartNames.addAll(newPartNames); if (!mergedPartNames.isEmpty()) { addPartitionsCache(table.getOrBuildNameMapping(), mergedPartNames, - table.getPartitionColumnTypes(java.util.Optional.empty())); + table.getPartitionColumnTypes(Optional.empty())); } LOG.info("Refreshed cache for table {}: {} modified partitions, {} new partitions", @@ -1350,8 +1382,8 @@ private static int inferPartitionColumnCount(Map> values) { return first == null ? 0 : first.size(); } - public java.util.Optional> getSortedPartitionRanges() { - return java.util.Optional.ofNullable(sortedPartitionRanges); + public Optional> getSortedPartitionRanges() { + return Optional.ofNullable(sortedPartitionRanges); } private SortedPartitionRanges buildSortedPartitionRanges() { diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java index b1fa1addd338c1..69b021c5509f2c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/source/HiveScanNode.java @@ -34,6 +34,8 @@ import org.apache.doris.datasource.FileQueryScanNode; import org.apache.doris.datasource.FileSplit; import org.apache.doris.datasource.FileSplitter; +import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.SplitProducer; import org.apache.doris.datasource.TableFormatType; import org.apache.doris.datasource.hive.AcidInfo; import org.apache.doris.datasource.hive.AcidInfo.DeleteDeltaInfo; @@ -54,6 +56,7 @@ import org.apache.doris.planner.ScanContext; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; import org.apache.doris.spi.Split; import org.apache.doris.statistics.StatisticalType; import org.apache.doris.thrift.TFileAttributes; @@ -74,19 +77,20 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Random; -import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Executor; -import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; public class HiveScanNode extends FileQueryScanNode { @@ -101,9 +105,10 @@ public class HiveScanNode extends FileQueryScanNode { private DirectoryLister directoryLister; private boolean partitionInit = false; - private final AtomicReference batchException = new AtomicReference<>(null); private List prunedPartitions; - private final Semaphore splittersOnFlight = new Semaphore(NUM_SPLITTERS_ON_FLIGHT); + private Boolean batchMode; + private boolean fileCountBatch; + private long estimatedFileCount; private final AtomicInteger numSplitsPerPartition = new AtomicInteger(NUM_SPLITS_PER_PARTITION); private boolean skipCheckingAcidVersionFile = false; @@ -231,73 +236,240 @@ public List getSplits(int numBackends) throws UserException { } } - private UserException createSplitPlanningException(Throwable failure) { + private UserException createSplitPlanningException(Throwable cause) { return new UserException("get file split failed for table: " + hmsTable.getName() - + ", err: " + Util.getRootCauseMessage(failure), failure); + + ", err: " + Util.getRootCauseMessage(cause), cause); } @Override - public void startSplit(int numBackends) { + protected SplitAssignment createSplitAssignment(Map properties, + List partitionKeys, boolean admission) { + return new SplitAssignment(backendPolicy, this, this::splitToScanRange, properties, + partitionKeys, admission, sessionVariable.getSplitBatchSize(), + sessionVariable.getSplitQueueCapacity()); + } + + @Override + public void startSplit(int numBackends) throws UserException { + ConnectContext context = ConnectContext.get(); + StatementContext statement = context == null ? null : context.getStatementContext(); + if (statement != null) { + StmtExecutor executor = context.getExecutor(); + statement.getOrRegisterStatementResource(this, () -> (Closeable) () -> { + // A deferred Flight query still fetches splits after statement handling ends. + // Its coordinator owns cleanup; abandoned planning and EXPLAIN stop here. + if (executor == null || !executor.isDeferredForArrowFlight()) { + stop(); + } + }); + } if (prunedPartitions.isEmpty()) { splitAssignment.finishSchedule(); return; } HiveExternalMetaCache cache = Env.getCurrentEnv().getExtMetaCacheMgr() .hive(hmsTable.getCatalog().getId()); - Executor scheduleExecutor = Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor(); - String bindBrokerName = hmsTable.getCatalog().bindBrokerName(); - AtomicInteger numFinishedPartitions = new AtomicInteger(0); - long startTime = System.currentTimeMillis(); - List> partitionBatches = Lists.partition( - prunedPartitions, sessionVariable.splitPartitionBatchSize); - CompletableFuture.runAsync(() -> { - for (List partitions : partitionBatches) { - if (batchException.get() != null || splitAssignment.isStop()) { - break; + Executor executor = getSplitTaskExecutor(); + String broker = hmsTable.getCatalog().bindBrokerName(); + // Without the file cache, the eager path listed partitions concurrently. Keep that overlap. + boolean parallelListing = !fileCountBatch || Config.max_external_file_cache_num <= 0; + int producers = parallelListing + ? Math.min(prunedPartitions.size(), Config.max_external_cache_loader_thread_pool_size) : 1; + AtomicInteger nextPartition = new AtomicInteger(); + long started = System.currentTimeMillis(); + if (!usesCountFileSplitPolicy()) { + startSplitProducers(cache, executor, broker, numBackends, producers, nextPartition, null, started); + return; + } + int parallelism = sessionVariable.getParallelExecInstanceNum(scanContext.getClusterName()); + CountSplitPlan countPlan = new CountSplitPlan(); + splitAssignment.addCloseable(countPlan); + AtomicInteger remaining = new AtomicInteger(producers); + // The count policy needs only a scan-wide lower bound. Reuse the same cancellable workers + // and singleton partition loads, then publish one decision before starting range producers. + for (int i = 0; i < producers; i++) { + new SplitProducer(splitAssignment, executor, () -> { + prefetchCountFiles(countPlan, cache, broker, parallelism, numBackends, nextPartition); + return Collections.emptyList(); + }, failure -> { + if (failure != null) { + splitAssignment.setException(createSplitPlanningException(failure)); } - try { - splittersOnFlight.acquire(); - CompletableFuture.runAsync(() -> { + if (remaining.decrementAndGet() == 0) { + if (splitAssignment.needMoreSplit()) { + countPlan.needSplit = FileSplitter.needSplitForCountPushdown( + parallelism, numBackends, countPlan.fileCount.get()); + startSplitProducers(cache, executor, broker, numBackends, + producers, nextPartition, countPlan, started); + } else { + closeCountSplitPlan(countPlan); + splitAssignment.finishSchedule(); + } + } + }).start(); + } + } + + private void prefetchCountFiles(CountSplitPlan plan, HiveExternalMetaCache cache, String broker, + int parallelism, int numBackends, AtomicInteger nextPartition) + throws IOException, UserException { + while (splitAssignment.needMoreSplit() + && FileSplitter.needSplitForCountPushdown(parallelism, numBackends, plan.fileCount.get())) { + int index = nextPartition.getAndIncrement(); + if (index >= prunedPartitions.size()) { + return; + } + List loaded = getFileCachesByPartitions(cache, + Collections.singletonList(prunedPartitions.get(index)), broker, true); + for (FileCacheValue value : loaded) { + if (value.getFiles() != null && !value.getFiles().isEmpty()) { + plan.prefetched.add(value); + plan.fileCount.addAndGet(value.getFiles().size()); + } + } + } + } + + private static class CountSplitPlan implements Closeable { + private final ConcurrentLinkedQueue prefetched = new ConcurrentLinkedQueue<>(); + private final AtomicLong fileCount = new AtomicLong(); + private boolean needSplit; + + @Override + public void close() { + prefetched.clear(); + } + } + + private void closeCountSplitPlan(CountSplitPlan plan) { + splitAssignment.removeCloseable(plan); + plan.close(); + } + + private void startSplitProducers(HiveExternalMetaCache cache, Executor executor, String broker, + int numBackends, int producers, AtomicInteger nextPartition, CountSplitPlan countPlan, long started) { + AtomicInteger remaining = new AtomicInteger(producers); + for (int i = 0; i < producers; i++) { + new SplitProducer(splitAssignment, executor, + createSplitBatchSource(cache, nextPartition, broker, numBackends, countPlan), failure -> { try { - List allFiles = Lists.newArrayList(); - getFileSplitByPartitions( - cache, partitions, allFiles, bindBrokerName, - numBackends, true); - updateNumSplitsPerPartition(allFiles); - if (splitAssignment.needMoreSplit()) { - splitAssignment.addToQueue(allFiles); + if (failure != null) { + splitAssignment.setException(createSplitPlanningException(failure)); } - } catch (Exception e) { - batchException.set(createSplitPlanningException(e)); } finally { - splittersOnFlight.release(); - if (batchException.get() != null) { - splitAssignment.setException(batchException.get()); - } - if (numFinishedPartitions.incrementAndGet() == partitionBatches.size()) { + if (remaining.decrementAndGet() == 0) { + if (countPlan != null) { + closeCountSplitPlan(countPlan); + } if (getSummaryProfile() != null) { getSummaryProfile().addExternalTableGetPartitionFilesTime( - System.currentTimeMillis() - startTime); + System.currentTimeMillis() - started); } splitAssignment.finishSchedule(); } } - }, scheduleExecutor); - } catch (Exception e) { - // When submitting a task, an exception will be thrown if the task pool(scheduleExecutor) is full - batchException.set(createSplitPlanningException(e)); - break; + }).start(); + } + } + + SplitProducer.BatchSource createSplitBatchSource(HiveExternalMetaCache cache, AtomicInteger nextPartition, + String broker, int numBackends) { + return createSplitBatchSource(cache, nextPartition, broker, numBackends, null); + } + + private SplitProducer.BatchSource createSplitBatchSource(HiveExternalMetaCache cache, AtomicInteger nextPartition, + String broker, int numBackends, CountSplitPlan countPlan) { + return new SplitProducer.BatchSource() { + private Iterator fileCaches = Collections.emptyIterator(); + private Iterator files = Collections.emptyIterator(); + private Iterator ranges = Collections.emptyIterator(); + private FileCacheValue currentCache; + private long targetSize; + private boolean needSplit; + private long partitionSplitCount; + + @Override + public List getNextBatch() throws IOException, UserException { + List batch = SplitProducer.newBatch(splitAssignment.getSplitBatchSize()); + while (batch.size() < splitAssignment.getSplitBatchSize() && splitAssignment.needMoreSplit()) { + if (ranges.hasNext()) { + batch.add(ranges.next()); + partitionSplitCount++; + continue; + } + if (!loadNextFileRanges()) { + break; + } + } + updatePartitionSplitEstimate(); + return batch; + } + + private void updatePartitionSplitEstimate() { + numSplitsPerPartition.accumulateAndGet( + (int) Math.min(Integer.MAX_VALUE, partitionSplitCount), Math::max); + } + + private boolean loadNextFileRanges() throws IOException, UserException { + while (splitAssignment.needMoreSplit()) { + if (files.hasNext()) { + HiveExternalMetaCache.HiveFileStatus file = files.next(); + ranges = fileSplitter.splitFileIterator(file.getPath(), targetSize, + file.getBlockLocations(), file.getLength(), file.getModificationTime(), + currentCache.isSplittable() && needSplit, currentCache.getPartitionValues(), + new HiveSplitCreator(currentCache.getAcidInfo())); + return true; + } + if (fileCaches.hasNext()) { + currentCache = fileCaches.next(); + files = currentCache.getFiles() == null ? Collections.emptyIterator() + : currentCache.getFiles().iterator(); + continue; + } + // One producer batch can span partitions, and one partition can span batches. + // Publish its accumulated count before moving to the next selected partition. + updatePartitionSplitEstimate(); + partitionSplitCount = 0; + FileCacheValue prefetched = countPlan == null ? null : countPlan.prefetched.poll(); + List loaded; + if (prefetched != null) { + loaded = Collections.singletonList(prefetched); + } else { + int index = nextPartition.getAndIncrement(); + if (index >= prunedPartitions.size()) { + return false; + } + loaded = getFileCachesByPartitions(cache, + Collections.singletonList(prunedPartitions.get(index)), broker, true); + } + targetSize = determineTargetFileSplitSize(loaded, true); + needSplit = countPlan == null ? needSplitFiles(loaded, numBackends) : countPlan.needSplit; + fileCaches = loaded.iterator(); } + return false; } - if (batchException.get() != null) { - splitAssignment.setException(batchException.get()); + + @Override + public void close() { + ranges = Collections.emptyIterator(); + files = Collections.emptyIterator(); + fileCaches = Collections.emptyIterator(); + currentCache = null; } - }); + }; } @Override public boolean isBatchMode() { - if (!isBatchModeEnabled()) { + if (batchMode == null) { + batchMode = selectBatchMode(); + } + return batchMode; + } + + private boolean selectBatchMode() { + if (!isBatchModeEnabled() + || hmsTable.isHiveTransactionalTable()) { return false; } if (!partitionInit) { @@ -308,13 +480,30 @@ public boolean isBatchMode() { } partitionInit = true; } - int numPartitions = sessionVariable.getNumPartitionsInBatchMode(); - return numPartitions >= 0 && prunedPartitions.size() >= numPartitions; + int partitionThreshold = sessionVariable.getNumPartitionsInBatchMode(); + if (partitionThreshold >= 0 && prunedPartitions.size() >= partitionThreshold) { + return true; + } + int fileThreshold = sessionVariable.getNumFilesInBatchMode(); + if (fileThreshold < 0) { + return false; + } + HiveExternalMetaCache cache = Env.getCurrentEnv().getExtMetaCacheMgr() + .hive(hmsTable.getCatalog().getId()); + estimatedFileCount = cache.getCachedFileCount(prunedPartitions); + if (estimatedFileCount < 0 || estimatedFileCount >= fileThreshold) { + fileCountBatch = true; + estimatedFileCount = estimatedFileCount < 0 ? Math.max(1, fileThreshold) : estimatedFileCount; + return true; + } + return false; } @Override public int numApproximateSplits() { - return numSplitsPerPartition.get() * prunedPartitions.size(); + long estimate = fileCountBatch ? estimatedFileCount + : (long) numSplitsPerPartition.get() * prunedPartitions.size(); + return (int) Math.min(Integer.MAX_VALUE, estimate); } private void updateNumSplitsPerPartition(List splits) { @@ -331,6 +520,47 @@ private void updateNumSplitsPerPartition(List splits) { private void getFileSplitByPartitions(HiveExternalMetaCache cache, List partitions, List allFiles, String bindBrokerName, int numBackends, boolean isBatchMode) throws IOException, UserException { + List fileCaches = getFileCachesByPartitions(cache, partitions, bindBrokerName, isBatchMode); + long targetFileSplitSize = determineTargetFileSplitSize(fileCaches, isBatchMode); + if (tableSample != null) { + List hiveFileStatuses = selectFiles(fileCaches); + splitAllFiles(allFiles, hiveFileStatuses, targetFileSplitSize); + return; + } + + /** + * If a table-level COUNT(*) is pushed down, + * we don't need to split the file because for parquet/orc format, only metadata is read. + * If we split the file, we will read metadata of a file multiple times, which is not efficient. + * + * - Hive Full Acid Transactional Table may need merge on read, so do not apply this optimization. + * - If the file format is not parquet/orc, eg, text, we need to split the file to increase the parallelism. + */ + boolean needSplit = needSplitFiles(fileCaches, numBackends); + + for (HiveExternalMetaCache.FileCacheValue fileCacheValue : fileCaches) { + if (fileCacheValue.getFiles() == null) { + continue; + } + boolean isSplittable = fileCacheValue.isSplittable(); + + for (HiveExternalMetaCache.HiveFileStatus status : fileCacheValue.getFiles()) { + allFiles.addAll(fileSplitter.splitFile( + status.getPath(), + targetFileSplitSize, + status.getBlockLocations(), + status.getLength(), + status.getModificationTime(), + isSplittable && needSplit, + fileCacheValue.getPartitionValues(), + new HiveSplitCreator(fileCacheValue.getAcidInfo()))); + } + } + } + + private List getFileCachesByPartitions(HiveExternalMetaCache cache, + List partitions, String bindBrokerName, boolean isBatchMode) + throws IOException, UserException { List fileCaches; long startTime = System.currentTimeMillis(); if (hiveTransaction != null) { @@ -374,24 +604,12 @@ private void getFileSplitByPartitions(HiveExternalMetaCache cache, List hiveFileStatuses = selectFiles(fileCaches); - splitAllFiles(allFiles, hiveFileStatuses, targetFileSplitSize); - return; - } + return fileCaches; + } - /** - * If a table-level COUNT(*) is pushed down, - * we don't need to split the file because for parquet/orc format, only metadata is read. - * If we split the file, we will read metadata of a file multiple times, which is not efficient. - * - * - Hive Full Acid Transactional Table may need merge on read, so do not apply this optimization. - * - If the file format is not parquet/orc, eg, text, we need to split the file to increase the parallelism. - */ + private boolean needSplitFiles(List fileCaches, int numBackends) throws UserException { boolean needSplit = true; - if (isTableLevelCountStarPushdown() - && !(hmsTable.isHiveTransactionalTable() && hmsTable.isFullAcidTable())) { + if (usesCountFileSplitPolicy()) { int totalFileNum = 0; for (FileCacheValue fileCacheValue : fileCaches) { if (fileCacheValue.getFiles() != null) { @@ -401,25 +619,12 @@ private void getFileSplitByPartitions(HiveExternalMetaCache cache, List fileCaches) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java index feba79913f59ec..faf731d9c184ea 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/hudi/source/HudiScanNode.java @@ -37,6 +37,7 @@ import org.apache.doris.datasource.FileQueryScanNode; import org.apache.doris.datasource.NameMapping; import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.SplitProducer; import org.apache.doris.datasource.TableFormatType; import org.apache.doris.datasource.hive.HMSExternalCatalog; import org.apache.doris.datasource.hive.HivePartition; @@ -54,6 +55,7 @@ import org.apache.doris.planner.ScanContext; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; import org.apache.doris.spi.Split; import org.apache.doris.statistics.StatisticalType; import org.apache.doris.thrift.TExplainLevel; @@ -89,6 +91,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Locale; import java.util.Map; @@ -704,25 +707,33 @@ public void startSplit(int numBackends) { releaseFsViewOnce(); return; } + Executor splitExecutor = getSplitTaskExecutor(); acquireFsViewLease(); ExecutorService scheduleExecutor = Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor(); Executor producerExecutor = Env.getCurrentEnv().getExtMetaCacheMgr().getFileListingExecutor(); long startTime = System.currentTimeMillis(); - BatchFsViewOwner createdOwner = new BatchFsViewOwner( + BatchFsViewOwner batchOwner = new BatchFsViewOwner( splitAssignment, fsViewLease, scheduleExecutor, producerExecutor); - BatchFsViewOwner batchOwner = createdOwner; ConnectContext connectContext = ConnectContext.get(); StatementContext statementContext = connectContext == null ? null : connectContext.getStatementContext(); if (statementContext != null) { + StmtExecutor executor = connectContext.getExecutor(); + Closeable statementCleanup = () -> { + // GetFlightInfo ends the statement before the BE fetches deferred DoGet splits. + // Capture this executor: processor cleanup can clear it from the connection context. + if (executor == null || !executor.isDeferredForArrowFlight()) { + batchOwner.close(); + } + }; try { - batchOwner = statementContext.getOrRegisterStatementResource( - batchFsViewResourceKey, () -> createdOwner); - if (batchOwner != createdOwner) { - createdOwner.finish(); + Closeable registered = statementContext.getOrRegisterStatementResource( + batchFsViewResourceKey, () -> statementCleanup); + if (registered != statementCleanup) { + batchOwner.finish(); throw new IllegalStateException("Hudi batch split owner was registered twice"); } } catch (RuntimeException e) { - createdOwner.finish(); + batchOwner.finish(); throw e; } } @@ -731,8 +742,8 @@ public void startSplit(int numBackends) { splitAssignment.addCloseable(finalBatchOwner); // Publish the owner before the blocking timeline sync: cancellation can now signal the // sync thread, and the owner keeps the lease until the sync actually terminates. - finalBatchOwner.beginSync(Thread.currentThread()); try { + finalBatchOwner.beginSync(Thread.currentThread()); syncFsView(); } catch (RuntimeException e) { finalBatchOwner.finish(); @@ -769,6 +780,7 @@ public void startSplit(int numBackends) { break; } pendingTasks.incrementAndGet(); + AtomicBoolean deliverySubmitted = new AtomicBoolean(); TerminalTask partitionTask = terminalTask(() -> { try { ensureHmsRuntimeGeneration(); @@ -779,13 +791,28 @@ public void startSplit(int numBackends) { } if (splitAssignment.needMoreSplit()) { ensureHmsRuntimeGeneration(); - splitAssignment.addToQueue(allFiles); + pendingTasks.incrementAndGet(); + deliverySubmitted.set(true); + new SplitProducer(splitAssignment, splitExecutor, createSplitBatchSource(allFiles), + failure -> { + try { + if (failure != null) { + recordBatchException(failure); + } + } finally { + splittersOnFlight.release(); + taskFinished.run(); + } + }).start(); } } catch (Throwable t) { recordBatchException(t); } }, () -> { - splittersOnFlight.release(); + // A parked delivery still owns this partition's buffered splits and admission slot. + if (!deliverySubmitted.get()) { + splittersOnFlight.release(); + } taskFinished.run(); }); try { @@ -810,6 +837,19 @@ public void startSplit(int numBackends) { } } + @VisibleForTesting + public SplitProducer.BatchSource createSplitBatchSource(List splits) { + Iterator iterator = splits.iterator(); + return () -> { + List batch = SplitProducer.newBatch(splitAssignment.getSplitBatchSize()); + while (batch.size() < splitAssignment.getSplitBatchSize() + && splitAssignment.needMoreSplit() && iterator.hasNext()) { + batch.add(iterator.next()); + } + return batch; + }; + } + private TerminalTask terminalTask(Runnable task, Runnable taskFinished) { return new TerminalTask(task, taskFinished); } @@ -921,7 +961,11 @@ void finish() { } void beginSync(Thread thread) { + // Publish before checking cancellation so close either sees this thread or is observed here. syncThread = thread; + if (stopping.get()) { + throw new CancellationException("Hudi filesystem-view sync was cancelled"); + } } void endSync() { @@ -1057,7 +1101,11 @@ void discardBeforeSubmission() { } void beginSync(Thread thread) { + // Publish before checking cancellation so close either sees this thread or is observed here. syncThread = thread; + if (stopping.get()) { + throw new CancellationException("Hudi filesystem-view sync was cancelled"); + } } void endSync() { diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/IcebergExternalMetaCache.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/IcebergExternalMetaCache.java index 23bc3bf17fa334..3875b7d4a97865 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/IcebergExternalMetaCache.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/IcebergExternalMetaCache.java @@ -1344,8 +1344,10 @@ private ManifestCacheValue loadDeleteFiles( org.apache.iceberg.ManifestFile manifest, Table table, boolean accountRetainedSize) throws IOException { ManifestCacheValue.Builder builder = ManifestCacheValue.deleteFilesBuilder(accountRetainedSize); + // Cache the manifest's physical partition types, just as loadDataFiles does. Current table + // specs may promote a DATE to nanos or erase a dropped field before query-local adaptation. try (ManifestReader reader = ManifestFiles.readDeleteManifest(manifest, - table.io(), table.specs())) { + table.io(), null)) { for (org.apache.iceberg.DeleteFile deleteFile : reader) { builder.addDeleteFile(deleteFile.copy()); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/cache/IcebergPartitionValueAdapter.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/cache/IcebergPartitionValueAdapter.java new file mode 100644 index 00000000000000..e31bf947527733 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/cache/IcebergPartitionValueAdapter.java @@ -0,0 +1,351 @@ +// 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.iceberg.cache; + +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.io.Serializable; +import java.math.BigDecimal; +import java.math.BigInteger; +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 static final long MICROS_PER_DAY = 86_400_000_000L; + private static final long NANOS_PER_DAY = 86_400_000_000_000L; + 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(); + } + if (source.typeId() == Type.TypeID.DATE) { + // A dropped DATE may lie outside the long-nanos range. Keep the key lossless. + return BigInteger.valueOf((Integer) value).multiply(BigInteger.valueOf(NANOS_PER_DAY)); + } + if (source.typeId() == Type.TypeID.TIMESTAMP) { + return BigInteger.valueOf((Long) value).multiply(BigInteger.valueOf(1000L)); + } + if (source.typeId() == Type.TypeID.TIMESTAMP_NANO) { + return BigInteger.valueOf((Long) value); + } + // 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.DATE && target.equals(Types.TimestampNanoType.withoutZone())) { + return Math.multiplyExact(((Integer) value).longValue(), NANOS_PER_DAY); + } + if (source.equals(Types.TimestampNanoType.withoutZone()) && target.typeId() == Type.TypeID.DATE) { + long nanos = (Long) value; + if (nanos % NANOS_PER_DAY != 0) { + throw new IllegalArgumentException("Partition value cannot be represented as date: " + value); + } + return Math.toIntExact(nanos / NANOS_PER_DAY); + } + if (source.typeId() == Type.TypeID.DATE && target.equals(Types.TimestampType.withoutZone())) { + return Math.multiplyExact(((Integer) value).longValue(), MICROS_PER_DAY); + } + if (source.equals(Types.TimestampType.withoutZone()) && target.typeId() == Type.TypeID.DATE) { + long micros = (Long) value; + if (micros % MICROS_PER_DAY != 0) { + throw new IllegalArgumentException("Partition value cannot be represented as date: " + value); + } + return Math.toIntExact(micros / MICROS_PER_DAY); + } + 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, Serializable { + private static final long serialVersionUID = 1L; + 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 { + private static final long serialVersionUID = 1L; + + 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 { + private static final long serialVersionUID = 1L; + + 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-core/src/main/java/org/apache/doris/datasource/iceberg/cache/ManifestCacheValue.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/cache/ManifestCacheValue.java index 617d2aa6fb53b6..8cb90121fae468 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/cache/ManifestCacheValue.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/cache/ManifestCacheValue.java @@ -20,10 +20,12 @@ import org.apache.doris.datasource.metacache.MetaCacheWeightUtils; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; import org.apache.iceberg.ContentFile; import org.apache.iceberg.DataFile; import org.apache.iceberg.DeleteFile; import org.apache.iceberg.PartitionData; +import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.StructLike; import org.apache.iceberg.types.Types; @@ -106,6 +108,39 @@ public static Builder deleteFilesBuilder(boolean accountRetainedSize) { return new Builder(false, accountRetainedSize); } + /** Query-local views preserve the immutable shared cache payload and its accounting. */ + public ScanFiles forScan(PartitionSpec spec) { + if (dataFiles.isEmpty() && deleteFiles.isEmpty()) { + return new ScanFiles(dataFiles, deleteFiles); + } + PartitionData partition = (PartitionData) (dataFiles.isEmpty() + ? deleteFiles.get(0).partition() : dataFiles.get(0).partition()); + Types.StructType target = spec.rawPartitionType(); + if (partition.getPartitionType().equals(target)) { + return new ScanFiles(dataFiles, deleteFiles); + } + IcebergPartitionValueAdapter adapter = new IcebergPartitionValueAdapter(partition.getPartitionType(), target); + return new ScanFiles(Lists.transform(dataFiles, adapter::adapt), Lists.transform(deleteFiles, adapter::adapt)); + } + + public static final class ScanFiles { + private final List dataFiles; + private final List deleteFiles; + + private ScanFiles(List dataFiles, List deleteFiles) { + this.dataFiles = dataFiles; + this.deleteFiles = deleteFiles; + } + + public List getDataFiles() { + return dataFiles; + } + + public List getDeleteFiles() { + return deleteFiles; + } + } + public List getDataFiles() { return dataFiles; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java index 743736f649a696..575df9274c0e49 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java @@ -30,17 +30,21 @@ import org.apache.doris.catalog.StructField; import org.apache.doris.catalog.StructType; import org.apache.doris.catalog.TableIf; +import org.apache.doris.common.Config; import org.apache.doris.common.Pair; import org.apache.doris.common.UserException; import org.apache.doris.common.profile.SummaryProfile; import org.apache.doris.common.security.authentication.ExecutionAuthenticator; +import org.apache.doris.common.util.FairTaskExecutor; import org.apache.doris.common.util.LocationPath; import org.apache.doris.common.util.Util; import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.ExternalTable; import org.apache.doris.datasource.ExternalUtil; import org.apache.doris.datasource.FileQueryScanNode; +import org.apache.doris.datasource.ScanTaskGroups; import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.SplitProducer; import org.apache.doris.datasource.TableFormatType; import org.apache.doris.datasource.credentials.CredentialUtils; import org.apache.doris.datasource.credentials.VendedCredentialsFactory; @@ -56,7 +60,6 @@ import org.apache.doris.datasource.iceberg.IcebergTimestampUtils; import org.apache.doris.datasource.iceberg.IcebergUtils; import org.apache.doris.datasource.iceberg.cache.IcebergManifestCacheLoader; -import org.apache.doris.datasource.iceberg.cache.ManifestCacheValue; import org.apache.doris.datasource.iceberg.profile.IcebergMetricsReporter; import org.apache.doris.datasource.iceberg.source.IcebergDeleteFileFilter.EqualityDelete; import org.apache.doris.datasource.mvcc.MvccSnapshot; @@ -166,10 +169,10 @@ import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.CancellationException; +import java.util.concurrent.Executor; import java.util.concurrent.FutureTask; import java.util.concurrent.ThreadPoolExecutor; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Supplier; import java.util.stream.Collectors; @@ -1739,65 +1742,65 @@ public void startSplit(int numBackends) throws UserException { public void doStartSplit() throws UserException { TableScan scan = createTableScan(); - ExecutorService executor = Env.getCurrentEnv().getExtMetaCacheMgr().getScheduleExecutor(); + Executor executor = getSplitTaskExecutor(); Closeable generationLease = IcebergUtils.retainTableGenerationForAsyncPlanning( source.getTargetTable(), frozenGenerationSource); - AsyncPlanningTask planningTask = new AsyncPlanningTask(executor, splitAssignment, generationLease, () -> { - AtomicReference> taskRef = new AtomicReference<>(); - try { - preExecutionAuthenticator.execute( - () -> { - long startTime = System.currentTimeMillis(); - try { - CloseableIterable fileScanTasks = planFileScanTask(scan); - taskRef.set(fileScanTasks); - CloseableIterator iterator = fileScanTasks.iterator(); - while (splitAssignment.needMoreSplit() && iterator.hasNext()) { - try { - splitAssignment.addToQueue( - Lists.newArrayList(createIcebergSplit(iterator.next()))); - } catch (UserException e) { - throw new RuntimeException(e); - } - } - } finally { - if (getSummaryProfile() != null) { - getSummaryProfile().addExternalTableGetFileScanTasksTime( - System.currentTimeMillis() - startTime); - } - } - } - ); - splitAssignment.finishSchedule(); - recordManifestCacheProfile(); - } catch (Exception e) { - Optional opt = checkNotSupportedException(e); - if (opt.isPresent()) { - splitAssignment.setException(new UserException(opt.get().getMessage(), opt.get())); - } else { - splitAssignment.setException(new UserException(e.getMessage(), e)); - } - } finally { - if (taskRef.get() != null) { - try { - taskRef.get().close(); - } catch (IOException e) { - // ignore + long started = System.currentTimeMillis(); + SplitProducer.BatchSource batches = new SplitProducer.BatchSource() { + private CloseableIterable tasks; + private CloseableIterator iterator; + + @Override + public List getNextBatch() throws Exception { + return preExecutionAuthenticator.execute(() -> { + if (tasks == null) { + tasks = planFileScanTask(scan); + iterator = tasks.iterator(); + } + List batch = SplitProducer.newBatch(splitAssignment.getSplitBatchSize()); + while (batch.size() < splitAssignment.getSplitBatchSize() + && splitAssignment.needMoreSplit() && iterator.hasNext()) { + batch.add(createIcebergSplit(iterator.next())); } + return batch; + }); + } + + @Override + public void close() throws IOException { + try { + if (tasks != null) { + tasks.close(); + } + } finally { + generationLease.close(); } } - }); - splitAssignment.addCloseable(planningTask); - planningTask.submit(); + }; + new SplitProducer(splitAssignment, executor, batches, failure -> { + if (failure != null) { + Optional unsupported = checkNotSupportedException(failure); + splitAssignment.setException(new UserException( + unsupported.isPresent() ? unsupported.get().getMessage() : failure.getMessage(), failure)); + } else { + splitAssignment.finishSchedule(); + } + if (getSummaryProfile() != null) { + getSummaryProfile().addExternalTableGetFileScanTasksTime(System.currentTimeMillis() - started); + } + recordManifestCacheProfile(); + }).start(); } /** Owns the planner generation until the worker actually exits, including cancellation races. */ @VisibleForTesting static final class AsyncPlanningTask implements Runnable, Closeable { - private final ExecutorService executor; + private final Executor executor; private final SplitAssignment splitAssignment; private final Closeable generationLease; private final Runnable planning; + private final Runnable termination; + private final Runnable admissionAvailable = this::submit; private final FutureTask submission; private boolean submitted; private boolean started; @@ -1805,8 +1808,14 @@ static final class AsyncPlanningTask implements Runnable, Closeable { private boolean finished; private Thread runner; - AsyncPlanningTask(ExecutorService executor, SplitAssignment splitAssignment, + AsyncPlanningTask(Executor executor, SplitAssignment splitAssignment, Closeable generationLease, Runnable planning) { + this(executor, splitAssignment, generationLease, planning, () -> { }); + } + + AsyncPlanningTask(Executor executor, SplitAssignment splitAssignment, + Closeable generationLease, Runnable planning, Runnable termination) { + this.termination = termination; this.executor = Objects.requireNonNull(executor, "executor is null"); this.splitAssignment = Objects.requireNonNull(splitAssignment, "splitAssignment is null"); this.generationLease = Objects.requireNonNull(generationLease, "generationLease is null"); @@ -1822,7 +1831,12 @@ void submit() { } submitted = true; } - executor.execute(submission); + if (executor instanceof FairTaskExecutor.Group) { + ((FairTaskExecutor.Group) executor).tryExecute( + submission, failure -> close(), admissionAvailable); + } else { + executor.execute(submission); + } synchronized (this) { if (finished && executor instanceof ThreadPoolExecutor) { ((ThreadPoolExecutor) executor).remove(submission); @@ -1891,12 +1905,17 @@ private void finish() { finished = true; runner = null; } + if (executor instanceof FairTaskExecutor.Group) { + ((FairTaskExecutor.Group) executor).removeSpaceListener(admissionAvailable); + } splitAssignment.removeCloseable(this); try { generationLease.close(); } catch (IOException e) { splitAssignment.setException( new UserException("Failed to release Iceberg planning generation", e)); + } finally { + termination.run(); } } } @@ -2086,6 +2105,11 @@ CloseableIterable planFileScanTaskWithoutReuse(TableScan scan) { try { return planFileScanTaskWithManifestCache(scan); } catch (Exception e) { + if (isScanTaskLeaseClosed() || Thread.currentThread().isInterrupted() + || (splitAssignment != null && splitAssignment.isStop()) + || e instanceof CancellationException) { + throw new CancellationException("Iceberg scan planning was cancelled"); + } manifestCacheFailures++; LOG.warn("Plan with manifest cache failed, fallback to original scan: " + e.getMessage(), e); return splitFiles(scan); @@ -2389,10 +2413,11 @@ private static final class SerializationSizeLimitException extends IOException { } private CloseableIterable planFileScanTaskWithManifestCache(TableScan scan) throws IOException { - if (sessionVariable.getFileSplitSize() > 0 || isBatchMode() || tableLevelPushDownCount) { - // These paths intentionally stream Iceberg tasks to keep FE memory bounded. The - // manifest-cache planner materializes its result, and retaining that list in the - // statement cache would defeat the streaming contract. + if (sessionVariable.getFileSplitSize() > 0) { + return splitFiles(scan); + } + boolean batchMode = isBatchMode(); + if (tableLevelPushDownCount) { return splitFiles(scan); } // Get the snapshot from the scan; return empty if no snapshot exists @@ -2400,6 +2425,10 @@ private CloseableIterable planFileScanTaskWithManifestCache(TableS if (snapshot == null) { return CloseableIterable.withNoopClose(Collections.emptyList()); } + if (batchMode) { + targetSplitSize = sessionVariable.getMaxSplitSize(); + return splitFiles(streamFileScanTasksWithManifestCache(scan, snapshot), targetSplitSize); + } List tasks; try { tasks = getOrPlanSerializedIcebergTasks( @@ -2417,12 +2446,18 @@ private CloseableIterable planFileScanTaskWithManifestCache(TableS @VisibleForTesting protected List loadFileScanTasksWithManifestCache( TableScan scan, Snapshot snapshot) throws IOException { + try (CloseableIterable tasks = streamFileScanTasksWithManifestCache(scan, snapshot)) { + return Lists.newArrayList(tasks); + } + } + + @VisibleForTesting + protected CloseableIterable streamFileScanTasksWithManifestCache( + TableScan scan, Snapshot snapshot) throws IOException { // Initialize manifest cache for efficient manifest file access - IcebergExternalMetaCache cache = Env.getCurrentEnv().getExtMetaCacheMgr().iceberg(source.getCatalog().getId()); if (!(source.getTargetTable() instanceof ExternalTable)) { throw new RuntimeException("Iceberg scan target table is not an external table"); } - ExternalTable targetExternalTable = (ExternalTable) source.getTargetTable(); Schema scanSchema = scan.schema(); Set unsafeNanosFields = IcebergTimestampUtils.unsafeNanosFields(icebergTable, scanSchema); @@ -2450,6 +2485,7 @@ protected List loadFileScanTasksWithManifestCache( // ========== Phase 1: Load delete files from delete manifests ========== List deleteFiles = new ArrayList<>(); List deleteManifests = snapshot.deleteManifests(icebergTable.io()); + List matchingDeletes = new ArrayList<>(); for (ManifestFile manifest : deleteManifests) { // Skip non-delete manifests if (manifest.content() != ManifestContent.DELETES) { @@ -2467,10 +2503,13 @@ protected List loadFileScanTasksWithManifestCache( if (!evaluator.eval(manifest)) { continue; } - // Load delete files from cache (or from storage if not cached) - ManifestCacheValue value = IcebergManifestCacheLoader.loadDeleteFilesWithCache(cache, - targetExternalTable, manifest, icebergTable, runtimeContext, this::recordManifestCacheAccess); - deleteFiles.addAll(value.getDeleteFiles()); + matchingDeletes.add(manifest); + } + try (NativeManifestPrefetcher prefetch = prefetchManifests(matchingDeletes, true)) { + for (NativeManifestPrefetcher.Result result : prefetch) { + deleteFiles.addAll(result.value.forScan(specsById.get(result.manifest.partitionSpecId())) + .getDeleteFiles()); + } } // Build delete file index for efficient lookup of deletes applicable to each data file @@ -2484,57 +2523,99 @@ protected List loadFileScanTasksWithManifestCache( } DeleteFileIndex deleteIndex = deleteIndexBuilder.build(); - // ========== Phase 2: Load data files and create scan tasks ========== - List tasks = new ArrayList<>(); - try (CloseableIterable dataManifests = - IcebergUtils.getMatchingManifest(snapshot.dataManifests(icebergTable.io()), - specsById, filterExpr)) { - for (ManifestFile manifest : dataManifests) { - // Skip non-data manifests - if (manifest.content() != ManifestContent.DATA) { - continue; - } - // Get the partition spec for this manifest - PartitionSpec spec = specsById.get(manifest.partitionSpecId()); - if (spec == null) { - continue; - } - // Get the residual evaluator for this partition spec - ResidualEvaluator residualEvaluator = residualEvaluators.get(manifest.partitionSpecId()); - if (residualEvaluator == null) { - continue; - } - - // Load data files from cache (or from storage if not cached) - ManifestCacheValue value = IcebergManifestCacheLoader.loadDataFilesWithCache(cache, - targetExternalTable, manifest, icebergTable, runtimeContext, this::recordManifestCacheAccess); - - // Process each data file in the manifest - for (org.apache.iceberg.DataFile dataFile : value.getDataFiles()) { - // Skip file if column statistics indicate no matching rows (metrics-based pruning) - if (metricsEvaluator != null && !metricsEvaluator.eval(dataFile)) { + // A batch scan retains only the current manifest view and the bounded prefetch window. + // Errors after delivery must fail the query; restarting a full SDK scan could duplicate rows. + CloseableIterable dataManifests = IcebergUtils.getMatchingManifest( + snapshot.dataManifests(icebergTable.io()), specsById, filterExpr); + NativeManifestPrefetcher prefetch; + try { + prefetch = prefetchManifests(dataManifests, false); + } catch (RuntimeException | Error e) { + dataManifests.close(); + throw e; + } + String schemaJson = SchemaParser.toJson(scanSchema); + Map specJson = new HashMap<>(); + specsById.forEach((id, spec) -> specJson.put(id, PartitionSpecParser.toJson(spec))); + Iterable tasks = () -> new com.google.common.collect.AbstractIterator() { + private final java.util.Iterator manifests = prefetch.iterator(); + private java.util.Iterator files = Collections.emptyIterator(); + private PartitionSpec spec; + private ResidualEvaluator residual; + + @Override + protected FileScanTask computeNext() { + while (true) { + while (files.hasNext()) { + org.apache.iceberg.DataFile file = files.next(); + if (!metricsEvaluator.eval(file) + || residual.residualFor(file.partition()).equals(Expressions.alwaysFalse())) { + continue; + } + DeleteFile[] deletes = deleteIndex.forDataFile(file.dataSequenceNumber(), file); + return new BaseFileScanTask(file, deletes, schemaJson, specJson.get(spec.specId()), residual); + } + if (!manifests.hasNext()) { + return endOfData(); + } + NativeManifestPrefetcher.Result result = manifests.next(); + if (result.manifest.content() != ManifestContent.DATA) { continue; } - // Skip file if partition values don't match the residual filter - if (residualEvaluator.residualFor(dataFile.partition()).equals(Expressions.alwaysFalse())) { + spec = specsById.get(result.manifest.partitionSpecId()); + if (spec == null) { continue; } - // Find all delete files that apply to this data file based on sequence number - List deletes = Arrays.asList( - deleteIndex.forDataFile(dataFile.dataSequenceNumber(), dataFile)); - - // Create a FileScanTask containing the data file, associated deletes, and metadata - tasks.add(new BaseFileScanTask( - dataFile, - deletes.toArray(new DeleteFile[0]), - SchemaParser.toJson(scanSchema), - PartitionSpecParser.toJson(spec), - residualEvaluator)); + residual = residualEvaluators.get(spec.specId()); + files = result.value.forScan(spec).getDataFiles().iterator(); } } + }; + return CloseableIterable.combine(tasks, () -> { + try { + prefetch.close(); + } finally { + dataManifests.close(); + } + }); + } + + private NativeManifestPrefetcher prefetchManifests(Iterable manifests, boolean deletes) { + Executor metadata = getManifestReadExecutor(); + SplitAssignment assignment = splitAssignment == null + ? new SplitAssignment(backendPolicy, this, this::splitToScanRange, + Collections.emptyMap(), Collections.emptyList(), true) : splitAssignment; + if (splitAssignment == null && metadata instanceof FairTaskExecutor.Group) { + ScanTaskGroups.Lease lease = getScanTaskLease(); + if (lease != null) { + lease.onClose(assignment::stop); + } } + ExternalTable target = (ExternalTable) source.getTargetTable(); + IcebergExternalMetaCache cache = Env.getCurrentEnv().getExtMetaCacheMgr().iceberg(source.getCatalog().getId()); + return new NativeManifestPrefetcher(manifests, metadata, assignment, + metadata instanceof FairTaskExecutor.Group && getScanTaskLease() == null, + Math.max(1, Config.external_scan_metadata_threads), + () -> IcebergUtils.retainTableGenerationForAsyncPlanning(target, frozenGenerationSource), + manifest -> { + try { + return preExecutionAuthenticator.execute(() -> deletes + ? IcebergManifestCacheLoader.loadDeleteFilesWithCache(cache, target, manifest, + icebergTable, runtimeContext, this::recordManifestCacheAccess) + : IcebergManifestCacheLoader.loadDataFilesWithCache(cache, target, manifest, + icebergTable, runtimeContext, this::recordManifestCacheAccess)); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } - return tasks; + protected Executor getManifestReadExecutor() { + ScanTaskGroups.Lease lease = getScanTaskLease(); + return lease == null + ? Env.getCurrentEnv().getExtMetaCacheMgr().getScanTaskExecutor(true) + .newGroup(Config.external_scan_task_queue_capacity) + : lease.group(true); } /** @@ -3189,7 +3270,7 @@ public List getPathPartitionKeys() throws UserException { return getOrderedPathPartitionKeys(); } - private void recordManifestCacheAccess(boolean cacheHit) { + private synchronized void recordManifestCacheAccess(boolean cacheHit) { if (cacheHit) { manifestCacheHits++; } else { diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/NativeManifestPrefetcher.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/NativeManifestPrefetcher.java new file mode 100644 index 00000000000000..f7a8c0a35fbf68 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/NativeManifestPrefetcher.java @@ -0,0 +1,162 @@ +// 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.iceberg.source; + +import org.apache.doris.common.util.FairTaskExecutor; +import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.iceberg.cache.ManifestCacheValue; + +import org.apache.iceberg.ManifestFile; + +import java.io.Closeable; +import java.util.Iterator; +import java.util.NoSuchElementException; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Function; + +/** Bounded completion-order reads. Every admitted read owns its generation until its worker exits. */ +final class NativeManifestPrefetcher implements Iterable, Closeable { + static final class Result { + final ManifestFile manifest; + ManifestCacheValue value; + Throwable failure; + IcebergScanNode.AsyncPlanningTask task; + + Result(ManifestFile manifest) { + this.manifest = manifest; + } + } + + private final Iterator manifests; + private final Executor executor; + private final SplitAssignment assignment; + private final boolean ownsExecutor; + private final Callable retain; + private final Function loader; + private final BlockingQueue completed = new LinkedBlockingQueue<>(); + private final Set pending = ConcurrentHashMap.newKeySet(); + private final AtomicBoolean closed = new AtomicBoolean(); + private int inFlight; + + NativeManifestPrefetcher(Iterable manifests, Executor executor, + SplitAssignment assignment, boolean ownsExecutor, int width, + Callable retain, Function loader) { + this.manifests = manifests.iterator(); + this.executor = executor; + this.assignment = assignment; + this.ownsExecutor = ownsExecutor; + this.retain = retain; + this.loader = loader; + assignment.addCloseable(this); + try { + for (int i = 0; i < width && this.manifests.hasNext(); i++) { + submitNext(); + } + } catch (RuntimeException | Error e) { + close(); + throw e; + } + } + + private void submitNext() { + if (closed.get() || assignment.isStop()) { + throw new CancellationException("Iceberg manifest prefetch was cancelled"); + } + Result result = new Result(manifests.next()); + Closeable generation; + try { + generation = retain.call(); + } catch (Exception e) { + throw new RuntimeException(e); + } + result.task = new IcebergScanNode.AsyncPlanningTask(executor, assignment, generation, () -> { + try { + if (closed.get() || assignment.isStop()) { + throw new CancellationException("Iceberg manifest prefetch was cancelled"); + } + result.value = loader.apply(result.manifest); + } catch (Throwable failure) { + result.failure = failure; + } + }, () -> { + pending.remove(result.task); + completed.add(result); + }); + inFlight++; + pending.add(result.task); + assignment.addCloseable(result.task); + result.task.submit(); + } + + @Override + public Iterator iterator() { + return new Iterator() { + @Override + public boolean hasNext() { + return inFlight != 0; + } + + @Override + public Result next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + Result result; + try { + result = completed.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new CancellationException("Iceberg manifest prefetch was interrupted"); + } + if (closed.get() || assignment.isStop() || result.value == null) { + if (result.failure != null) { + throw new RuntimeException(result.failure); + } + throw new CancellationException("Iceberg manifest prefetch was cancelled"); + } + inFlight--; + if (manifests.hasNext()) { + submitNext(); + } + return result; + } + }; + } + + @Override + public void close() { + if (!closed.compareAndSet(false, true)) { + return; + } + assignment.removeCloseable(this); + for (IcebergScanNode.AsyncPlanningTask task : pending) { + task.close(); + } + completed.add(new Result(null)); + if (ownsExecutor) { + ((FairTaskExecutor.Group) executor).close(); + } + } +} 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 de60dd1e10c7bc..766cd190080c0f 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 @@ -609,6 +609,7 @@ public String toString() { 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_INITIAL_BATCH_SIZE = "split_initial_batch_size"; public static final String FETCH_SPLITS_MAX_WAIT_TIME = "fetch_splits_max_wait_time_ms"; @@ -2768,6 +2769,21 @@ public void checkSplitPartitionBatchSize(String value) throws Exception { needForward = true) public int numFilesInBatchMode = 1; + @VariableMgr.VarAttr(name = SPLIT_INITIAL_BATCH_SIZE, checker = "checkSplitInitialBatchSize", needForward = true) + public int splitInitialBatchSize = 64; + + public void checkSplitInitialBatchSize(String value) throws Exception { + checkFieldValue(SPLIT_INITIAL_BATCH_SIZE, 1, value); + } + + public int getSplitBatchSize() { + return splitBatchSize; + } + + public int getSplitQueueCapacity() { + return splitQueueCapacity; + } + @VariableMgr.VarAttr( name = FETCH_SPLITS_MAX_WAIT_TIME, checker = "checkFetchSplitsMaxWaitTime", 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 bbabe512bd066c..d1b731b358d501 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 @@ -17,17 +17,107 @@ package org.apache.doris.common; +import org.apache.doris.common.util.FairTaskExecutor; +import org.apache.doris.datasource.ExternalMetaCacheMgr; +import org.apache.doris.metric.Metric; +import org.apache.doris.metric.MetricRepo; + import org.junit.Assert; import org.junit.Test; +import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.FutureTask; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; public class ThreadPoolManagerTest { + @Test(timeout = 30000) + public void testExternalScanSchedulersExportMetrics() throws Exception { + int splitThreads = Config.external_scan_split_threads; + int metadataThreads = Config.external_scan_metadata_threads; + try { + Config.external_scan_split_threads = 1; + Config.external_scan_metadata_threads = 1; + for (boolean checkpoint : new boolean[] {false, true}) { + ExternalMetaCacheMgr manager = new ExternalMetaCacheMgr(checkpoint); + try (FairTaskExecutor split = manager.getScanTaskExecutor(false); + FairTaskExecutor metadata = manager.getScanTaskExecutor(true)) { + ThreadPoolManager.registerAllThreadPoolMetric(); + String prefix = checkpoint ? "Checkpoint" : "NotCheckpoint"; + verifyFairTaskMetrics(split, prefix + "ExternalSplit"); + verifyFairTaskMetrics(metadata, prefix + "ExternalMetadata"); + } + } + } finally { + Config.external_scan_split_threads = splitThreads; + Config.external_scan_metadata_threads = metadataThreads; + } + } + + private void verifyFairTaskMetrics(FairTaskExecutor executor, String name) throws Exception { + Assert.assertEquals(0, metricValue(name, "pool_size")); + Assert.assertEquals(0, metricValue(name, "active_thread_num")); + Assert.assertEquals(0, metricValue(name, "task_in_queue")); + Assert.assertEquals(0, metricValue(name, "completed_task_count")); + Assert.assertEquals(0, metricValue(name, "task_group_num")); + Assert.assertEquals(Metric.MetricType.GAUGE, metric(name, "task_in_queue").getType()); + Assert.assertEquals(Metric.MetricType.COUNTER, metric(name, "completed_task_count").getType()); + + FairTaskExecutor.Group group = executor.newGroup(2); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch completed = new CountDownLatch(1); + try { + group.execute(() -> { + started.countDown(); + try { + Assert.assertTrue(release.await(10, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + }); + Assert.assertTrue(started.await(10, TimeUnit.SECONDS)); + group.execute(completed::countDown); + Assert.assertEquals(1, metricValue(name, "pool_size")); + Assert.assertEquals(1, metricValue(name, "active_thread_num")); + Assert.assertEquals(1, metricValue(name, "task_in_queue")); + Assert.assertEquals(1, metricValue(name, "task_group_num")); + Assert.assertEquals(0, metricValue(name, "completed_task_count")); + release.countDown(); + Assert.assertTrue(completed.await(10, TimeUnit.SECONDS)); + } finally { + release.countDown(); + executor.close(); + Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + Assert.assertEquals(0, metricValue(name, "active_thread_num")); + Assert.assertEquals(0, metricValue(name, "task_in_queue")); + Assert.assertEquals(0, metricValue(name, "task_group_num")); + Assert.assertEquals(2, metricValue(name, "completed_task_count")); + } + + private long metricValue(String name, String type) { + return ((Number) metric(name, type).getValue()).longValue(); + } + + private Metric metric(String name, String type) { + List> metrics = MetricRepo.getMetricsByName("thread_pool").stream() + .>map(metric -> metric) + .filter(metric -> metric.getLabels().stream() + .anyMatch(label -> label.getKey().equals("name") && label.getValue().equals(name))) + .filter(metric -> metric.getLabels().stream() + .anyMatch(label -> label.getKey().equals("type") && label.getValue().equals(type))) + .collect(Collectors.toList()); + Assert.assertEquals(name + "/" + type, 1, metrics.size()); + return metrics.get(0); + } + @Test public void testBlockedPolicySkipsCancelledTask() { LinkedBlockingQueue queue = new LinkedBlockingQueue<>(1); 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..ce33fe8ee8972c --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/common/util/FairTaskExecutorTest.java @@ -0,0 +1,571 @@ +// 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.CancellationException; +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(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/datasource/FileSplitterTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/FileSplitterTest.java index 5a39135103825d..b7c48b41e8f137 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/FileSplitterTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/FileSplitterTest.java @@ -25,7 +25,9 @@ import org.junit.Test; import java.util.Collections; +import java.util.Iterator; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; public class FileSplitterTest { @@ -237,4 +239,47 @@ public void testSmallFileNoSplit() throws Exception { FileSplit s = (FileSplit) splits.get(0); Assert.assertEquals(2 * MB, s.getLength()); } + + @Test + public void testMillionRangesAreConstructedOnlyWhenConsumed() throws Exception { + FileSplitter splitter = new FileSplitter(32, 64, 2); + AtomicInteger created = new AtomicInteger(); + Iterator ranges = splitter.splitFileIterator(LocationPath.of("hdfs://test/million"), + 1, null, 1000000, 0, true, Collections.emptyList(), + (path, start, length, fileLength, splitSize, modificationTime, hosts, partitions) -> { + created.incrementAndGet(); + return FileSplit.FileSplitCreator.DEFAULT.create(path, start, length, fileLength, + splitSize, modificationTime, hosts, partitions); + }); + Assert.assertEquals(0, created.get()); + for (int offset = 0; offset < 7; offset++) { + Assert.assertTrue(ranges.hasNext()); + Assert.assertTrue(ranges.hasNext()); + Split split = ranges.next(); + Assert.assertEquals(offset, split.getStart()); + Assert.assertEquals(1, split.getLength()); + } + Assert.assertEquals(7, created.get()); + Assert.assertTrue(ranges.hasNext()); + } + + @Test + public void testExplicitSizeKeepsSlopTailAndPartitionValues() throws Exception { + FileSplitter splitter = new FileSplitter(32, 64, 2); + Iterator ranges = splitter.splitFileIterator(LocationPath.of("hdfs://test/slop"), + 10, null, 31, 123, true, Collections.singletonList("p"), FileSplit.FileSplitCreator.DEFAULT); + long offset = 0; + for (long length : new long[] {10, 10, 11}) { + FileSplit split = (FileSplit) ranges.next(); + Assert.assertEquals(offset, split.getStart()); + Assert.assertEquals(length, split.getLength()); + Assert.assertEquals(Collections.singletonList("p"), split.getPartitionValues()); + Assert.assertEquals(123, split.getModificationTime()); + offset += length; + } + Assert.assertFalse(ranges.hasNext()); + Assert.assertThrows(java.util.NoSuchElementException.class, ranges::next); + Assert.assertEquals(2, splitter.getRemainingInitialSplitNum()); + } + } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/HudiSplitAssignmentTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/HudiSplitAssignmentTest.java new file mode 100644 index 00000000000000..5dc56f8e652e77 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/HudiSplitAssignmentTest.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.datasource; + +import org.apache.doris.common.UserException; +import org.apache.doris.common.util.LocationPath; +import org.apache.doris.datasource.hudi.source.HudiScanNode; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TScanRangeLocation; +import org.apache.doris.thrift.TScanRangeLocations; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.Multimap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +class HudiSplitAssignmentTest { + @Test + void hudiEagerFactoriesInitializeTwoScansOnOneWorkerBeforeConsumption() throws Exception { + Backend backend = new Backend(1, "127.0.0.1", 9050); + List splits = new ArrayList<>(); + for (int i = 0; i < 4097; i++) { + splits.add(new FileSplit(LocationPath.of("s3://bucket/file.parquet"), i, 1, + 4097, 0, null, Collections.emptyList())); + } + ExecutorService worker = Executors.newSingleThreadExecutor(); + ExecutorService planner = Executors.newSingleThreadExecutor(); + SplitAssignment first = eagerAssignment(backend, worker, splits); + SplitAssignment second = eagerAssignment(backend, worker, splits); + try { + Future initialized = planner.submit(() -> { + first.init(); + second.init(); + return null; + }); + initialized.get(5, TimeUnit.SECONDS); + + for (SplitAssignment assignment : new SplitAssignment[] {first, second}) { + Assertions.assertEquals(7, assignment.getSplitBatchSize()); + Assertions.assertTrue(assignment.getAssignedSplits(backend).size() <= 1); + SplitSource source = new SplitSource(backend, assignment, 1000); + Set actual = new HashSet<>(); + while (true) { + List batch = source.getNextBatch(7); + Assertions.assertTrue(batch.size() <= 7); + if (batch.isEmpty()) { + break; + } + for (TScanRangeLocations range : batch) { + Assertions.assertTrue(actual.add(range.getLocations().get(0).getBackendId())); + } + } + Assertions.assertEquals(4097, actual.size()); + for (long i = 0; i < 4097; i++) { + Assertions.assertTrue(actual.contains(i)); + } + } + } finally { + first.stop(); + second.stop(); + planner.shutdownNow(); + worker.shutdownNow(); + Assertions.assertTrue(planner.awaitTermination(5, TimeUnit.SECONDS)); + Assertions.assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private SplitAssignment eagerAssignment(Backend backend, Executor worker, List splits) throws Exception { + FileQueryScanNode node = Mockito.mock(HudiScanNode.class, Mockito.CALLS_REAL_METHODS); + node.sessionVariable = new SessionVariable(); + node.sessionVariable.splitBatchSize = 7; + node.sessionVariable.splitQueueCapacity = 1; + FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + Mockito.when(policy.numBackends()).thenReturn(1); + Mockito.when(policy.computeScanRangeAssignment(Mockito.anyList())).thenAnswer(call -> { + Multimap batch = ArrayListMultimap.create(); + for (Split split : call.>getArgument(0)) { + batch.put(backend, split); + } + return batch; + }); + Field policyField = ExternalScanNode.class.getDeclaredField("backendPolicy"); + policyField.setAccessible(true); + policyField.set(node, policy); + Mockito.doAnswer(call -> new TScanRangeLocations().setLocations(Collections.singletonList( + new TScanRangeLocation().setBackendId(call.getArgument(2).getStart())))) + .when(node).splitToScanRange(Mockito.any(), Mockito.anyMap(), Mockito.any(), + Mockito.anyList(), Mockito.anyBoolean()); + // Exercise virtual dispatch through the actual Hudi factory; replace only remote listing + // with a fixed partition listing delivered through Hudi's production batch source. + SplitAssignment assignment = node.createSplitAssignment( + Collections.emptyMap(), Collections.emptyList(), true); + Field assignmentField = org.apache.doris.planner.ScanNode.class.getDeclaredField("splitAssignment"); + assignmentField.setAccessible(true); + assignmentField.set(node, assignment); + Mockito.doAnswer(call -> { + new SplitProducer(assignment, worker, ((HudiScanNode) node).createSplitBatchSource(splits), failure -> { + if (failure != null) { + assignment.setException(new UserException(failure.getMessage(), failure)); + } + assignment.finishSchedule(); + }).start(); + return null; + }).when(node).startSplit(Mockito.anyInt()); + return assignment; + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/ScanTaskGroupsTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/ScanTaskGroupsTest.java new file mode 100644 index 00000000000000..01b2be79186496 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/ScanTaskGroupsTest.java @@ -0,0 +1,62 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource; + +import org.apache.doris.catalog.Env; +import org.apache.doris.common.util.FairTaskExecutor; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.util.concurrent.CancellationException; +import java.util.concurrent.atomic.AtomicInteger; + +class ScanTaskGroupsTest { + @Test + void scansShareQueryAdmissionAndDeferredLeaseOutlivesStatementOwner() throws Exception { + Env env = Mockito.mock(Env.class); + ExternalMetaCacheMgr manager = Mockito.mock(ExternalMetaCacheMgr.class); + Mockito.when(env.getExtMetaCacheMgr()).thenReturn(manager); + try (FairTaskExecutor pool = new FairTaskExecutor(1, task -> new Thread(task)); + MockedStatic mocked = Mockito.mockStatic(Env.class)) { + Mockito.when(manager.getScanTaskExecutor(Mockito.anyBoolean())).thenReturn(pool); + mocked.when(Env::getCurrentEnv).thenReturn(env); + ScanTaskGroups groups = new ScanTaskGroups(); + ScanTaskGroups.Lease first = groups.acquire(); + ScanTaskGroups.Lease flight = groups.acquire(); + FairTaskExecutor.Group split = first.group(false); + Assertions.assertSame(split, flight.group(false)); + Assertions.assertNotSame(split, flight.group(true)); + first.close(); + groups.close(); + Assertions.assertFalse(split.isClosed()); + Assertions.assertThrows(CancellationException.class, groups::acquire); + AtomicInteger cleanups = new AtomicInteger(); + flight.onClose(cleanups::incrementAndGet); + flight.close(); + flight.close(); + Assertions.assertEquals(1, cleanups.get()); + flight.onClose(cleanups::incrementAndGet); + Assertions.assertEquals(2, cleanups.get()); + Assertions.assertTrue(split.isClosed()); + Assertions.assertThrows(CancellationException.class, () -> flight.group(false)); + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitInitialDeliveryTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitInitialDeliveryTest.java new file mode 100644 index 00000000000000..8efaca28c2dc55 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitInitialDeliveryTest.java @@ -0,0 +1,62 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource; + +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TScanRangeLocations; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.Multimap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +class SplitInitialDeliveryTest { + @Test + void initialResponseHasIndependentBoundAndBackendMetadataDoesNotChangeChannelIdentity() throws Exception { + Backend original = new Backend(1, "localhost", 0); + Backend refreshed = new Backend(1, "new-host", 9000); + FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + Mockito.when(policy.computeScanRangeAssignment(Mockito.anyList())).thenAnswer(call -> { + Multimap assigned = ArrayListMultimap.create(); + for (Split split : call.>getArgument(0)) { + assigned.put(original, split); + } + return assigned; + }); + SplitAssignment assignment = new SplitAssignment(policy, new SplitGenerator() { }, + (be, properties, split, keys, admission) -> new TScanRangeLocations(), + Collections.emptyMap(), Collections.emptyList(), true, 128, 200); + List splits = new ArrayList<>(); + for (int i = 0; i < 200; i++) { + splits.add(Mockito.mock(Split.class)); + } + assignment.addToQueue(splits); + assignment.finishSchedule(); + SplitSource source = new SplitSource(refreshed, assignment, 1000, 64); + Assertions.assertEquals(64, source.getNextBatch(100).size()); + Assertions.assertEquals(100, source.getNextBatch(100).size()); + Assertions.assertEquals(36, source.getNextBatch(100).size()); + Assertions.assertTrue(source.getNextBatch(100).isEmpty()); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitProducerTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitProducerTest.java new file mode 100644 index 00000000000000..6d7abf1e7696e1 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitProducerTest.java @@ -0,0 +1,459 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource; + +import org.apache.doris.common.UserException; +import org.apache.doris.common.util.LocationPath; +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TScanRangeLocation; +import org.apache.doris.thrift.TScanRangeLocations; + +import com.google.common.collect.LinkedListMultimap; +import com.google.common.collect.Multimap; +import com.google.common.util.concurrent.Uninterruptibles; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +class SplitProducerTest { + @Test + void closingOneBackendResumesTheOtherBoundedQueue() throws Exception { + Backend closed = new Backend(1, "closed", 0); + Backend live = new Backend(2, "live", 0); + FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + Mockito.when(policy.numBackends()).thenReturn(2); + Mockito.when(policy.computeScanRangeAssignment(Mockito.anyList())).thenAnswer(call -> { + Multimap assigned = LinkedListMultimap.create(); + for (Split split : call.>getArgument(0)) { + assigned.put(split.getStart() < 2 ? closed : live, split); + } + return assigned; + }); + SplitAssignment assignment = new SplitAssignment(policy, Mockito.mock(SplitGenerator.class), + (backend, properties, split, keys, admission) -> new TScanRangeLocations().setLocations( + Collections.singletonList(new TScanRangeLocation().setBackendId(split.getStart()))), + Collections.emptyMap(), Collections.emptyList(), true, 4, 1); + RecordingSource source = new RecordingSource(4, 0, 1, 2, 3); + new SplitProducer(assignment, Runnable::run, source, failure -> { + Assertions.assertNull(failure); + assignment.finishSchedule(); + }).start(); + new SplitSource(closed, assignment, 1000).close(); + SplitSource consumer = new SplitSource(live, assignment, 1000); + List ranges = consumer.getNextBatch(2); + Assertions.assertEquals(2, ranges.size()); + Assertions.assertEquals(2, ranges.get(0).getLocations().get(0).getBackendId()); + Assertions.assertEquals(3, ranges.get(1).getLocations().get(0).getBackendId()); + Assertions.assertTrue(consumer.getNextBatch(1).isEmpty()); + Assertions.assertTrue(consumer.isLastBatch()); + Assertions.assertEquals(1, source.closeCalls.get()); + assignment.stop(); + } + + @Test + void twoScansInitializeOnOneWorkerBeforeAnyConsumerStarts() throws Exception { + Fixture first = new Fixture(1, 1); + Fixture second = new Fixture(1, 1); + ExecutorService worker = Executors.newSingleThreadExecutor(); + ExecutorService planner = Executors.newSingleThreadExecutor(); + RecordingSource firstSource = new RecordingSource(1, 0, 2, 4, 6, 8); + RecordingSource secondSource = new RecordingSource(1, 10, 12, 14, 16, 18); + Mockito.doAnswer(call -> { + first.start(worker, firstSource); + return null; + }).when(first.generator).startSplit(Mockito.anyInt()); + Mockito.doAnswer(call -> { + second.start(worker, secondSource); + return null; + }).when(second.generator).startSplit(Mockito.anyInt()); + try { + Future initialized = planner.submit(() -> { + first.assignment.init(); + second.assignment.init(); + return null; + }); + initialized.get(5, TimeUnit.SECONDS); + worker.submit(() -> { }).get(5, TimeUnit.SECONDS); + Assertions.assertEquals(1, first.assignment.getAssignedSplits(first.backend).size()); + Assertions.assertEquals(1, second.assignment.getAssignedSplits(second.backend).size()); + Assertions.assertEquals(Arrays.asList(0L, 2L, 4L, 6L, 8L), drain(first.source(), 2)); + Assertions.assertEquals(Arrays.asList(10L, 12L, 14L, 16L, 18L), drain(second.source(), 1)); + Assertions.assertTrue(firstSource.closed.await(5, TimeUnit.SECONDS)); + Assertions.assertTrue(secondSource.closed.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(1, firstSource.closeCalls.get()); + Assertions.assertEquals(1, secondSource.closeCalls.get()); + } finally { + first.assignment.stop(); + second.assignment.stop(); + planner.shutdownNow(); + worker.shutdownNow(); + Assertions.assertTrue(planner.awaitTermination(5, TimeUnit.SECONDS)); + Assertions.assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void partialMultiBackendSubmissionDoesNotReplayOrReassignTheBatch() throws Exception { + Fixture fixture = new Fixture(2, 1); + ManualExecutor executor = new ManualExecutor(); + fixture.assignment.addToQueue(Collections.singletonList(split(1))); + RecordingSource source = new RecordingSource(2, 0, 3); + fixture.start(executor, source); + executor.runNext(); + Assertions.assertTrue(executor.tasks.isEmpty()); + Assertions.assertEquals(0, fixture.completions.get()); + Assertions.assertEquals(Collections.singletonList(0L), identities(fixture.source().getNextBatch(1))); + Assertions.assertTrue(executor.tasks.isEmpty()); + SplitSource other = new SplitSource(fixture.otherBackend, fixture.assignment, 1000); + Assertions.assertEquals(Collections.singletonList(1L), identities(other.getNextBatch(1))); + executor.runUntilIdle(); + Assertions.assertEquals(Collections.singletonList(3L), identities(other.getNextBatch(1))); + Assertions.assertTrue(other.getNextBatch(1).isEmpty()); + Assertions.assertTrue(fixture.source().getNextBatch(1).isEmpty()); + Mockito.verify(fixture.policy, Mockito.times(2)).computeScanRangeAssignment(Mockito.any()); + Assertions.assertEquals(1, fixture.conversions.get(0L)); + Assertions.assertEquals(1, fixture.conversions.get(3L)); + Assertions.assertEquals(1, fixture.completions.get()); + Assertions.assertEquals(1, source.closeCalls.get()); + } + + @Test + void capacityReleasedBeforeTheProducerSuspendsIsNotLost() throws Exception { + Fixture fixture = new Fixture(1, 1); + fixture.assignment = Mockito.spy(fixture.assignment); + fixture.assignment.addToQueue(Collections.singletonList(split(0))); + SplitSource consumer = fixture.source(); + Mockito.doAnswer(call -> { + boolean offered = (boolean) call.callRealMethod(); + if (!offered) { + Assertions.assertEquals(Collections.singletonList(0L), identities(consumer.getNextBatch(1))); + } + return offered; + }).when(fixture.assignment).offerBatch(Mockito.any(), Mockito.any(), Mockito.any()); + ManualExecutor executor = new ManualExecutor(); + RecordingSource source = new RecordingSource(1, 2); + fixture.start(executor, source); + executor.runNext(); + Assertions.assertEquals(1, executor.tasks.size()); + executor.runUntilIdle(); + Assertions.assertEquals(Collections.singletonList(2L), identities(consumer.getNextBatch(1))); + Assertions.assertTrue(consumer.getNextBatch(1).isEmpty()); + Assertions.assertEquals(1, source.closeCalls.get()); + } + + @Test + void cancellationClosesASuspendedProducerWithoutAConsumer() throws Exception { + Fixture fixture = new Fixture(1, 1); + fixture.assignment.addToQueue(Collections.singletonList(split(0))); + ManualExecutor executor = new ManualExecutor(); + RecordingSource source = new RecordingSource(1, 2); + fixture.start(executor, source); + executor.runNext(); + Assertions.assertTrue(executor.tasks.isEmpty()); + Assertions.assertEquals(0, source.closeCalls.get()); + fixture.assignment.stop(); + executor.runUntilIdle(); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + Assertions.assertTrue(fixture.source().getNextBatch(1).isEmpty()); + fixture.assignment.stop(); + Assertions.assertTrue(executor.tasks.isEmpty()); + } + + @Test + void cancellationDoesNotCloseTheSourceDuringAnActiveRead() throws Exception { + Fixture fixture = new Fixture(1, 1); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch finishClose = new CountDownLatch(1); + RecordingSource source = new RecordingSource(1, 0) { + @Override + public List getNextBatch() throws IOException { + reading.countDown(); + // Keep the read active until the test explicitly releases it, even after cancellation. + Assertions.assertTrue(Uninterruptibles.awaitUninterruptibly(release, 5, TimeUnit.SECONDS)); + return super.getNextBatch(); + } + + @Override + public void close() { + super.close(); + // source.closed signals close entry, before the producer invokes its completion callback. + Assertions.assertTrue(Uninterruptibles.awaitUninterruptibly(finishClose, 5, TimeUnit.SECONDS)); + } + }; + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + fixture.start(executor, source); + Assertions.assertTrue(reading.await(5, TimeUnit.SECONDS)); + fixture.assignment.stop(); + Assertions.assertEquals(0, source.closeCalls.get()); + release.countDown(); + Assertions.assertTrue(source.closed.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(0, fixture.completions.get()); + finishClose.countDown(); + Assertions.assertTrue(fixture.completed.await(5, TimeUnit.SECONDS)); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + } finally { + release.countDown(); + finishClose.countDown(); + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void rejectedCancellationCleanupStillClosesWithoutAddingAnError() throws Exception { + Fixture fixture = new Fixture(1, 1); + fixture.assignment.addToQueue(Collections.singletonList(split(0))); + ManualExecutor executor = new ManualExecutor(); + RecordingSource source = new RecordingSource(1, 2); + fixture.start(executor, source); + executor.runNext(); + executor.reject = true; + fixture.assignment.stop(); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + Assertions.assertTrue(fixture.source().getNextBatch(1).isEmpty()); + } + + @Test + void rejectedInitialSubmissionPublishesFailureAndClosesTheSource() throws Exception { + Fixture fixture = new Fixture(1, 1); + ManualExecutor executor = new ManualExecutor(); + executor.reject = true; + RecordingSource source = new RecordingSource(1, 0); + fixture.start(executor, source); + UserException failure = Assertions.assertThrows(UserException.class, + () -> fixture.assignment.getAssignedSplits(fixture.backend)); + Assertions.assertInstanceOf(RejectedExecutionException.class, failure.getCause()); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + } + + @Test + void rejectedResumptionClosesThePendingSource() throws Exception { + Fixture fixture = new Fixture(1, 1); + fixture.assignment.addToQueue(Collections.singletonList(split(0))); + ManualExecutor executor = new ManualExecutor(); + RecordingSource source = new RecordingSource(1, 2); + fixture.start(executor, source); + executor.runNext(); + executor.reject = true; + UserException failure = Assertions.assertThrows(UserException.class, () -> { + SplitSource consumer = fixture.source(); + consumer.getNextBatch(1); + consumer.getNextBatch(1); + }); + Assertions.assertInstanceOf(RejectedExecutionException.class, failure.getCause()); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + } + + @Test + void sourceFailureWakesOtherSuspendedProducers() throws Exception { + Fixture fixture = new Fixture(1, 1); + fixture.assignment.addToQueue(Collections.singletonList(split(0))); + ManualExecutor executor = new ManualExecutor(); + RecordingSource pending = new RecordingSource(1, 2); + fixture.start(executor, pending); + executor.runNext(); + RecordingSource failed = new RecordingSource(1) { + @Override + public List getNextBatch() throws IOException { + throw new IOException("listing failed"); + } + }; + fixture.start(executor, failed); + executor.runUntilIdle(); + UserException failure = Assertions.assertThrows(UserException.class, + () -> fixture.assignment.getAssignedSplits(fixture.backend)); + Assertions.assertInstanceOf(IOException.class, failure.getCause()); + Assertions.assertEquals(1, pending.closeCalls.get()); + Assertions.assertEquals(1, failed.closeCalls.get()); + Assertions.assertEquals(2, fixture.completions.get()); + } + + @Test + void failureRacingWaiterRegistrationDoesNotLeaveTheSourceSuspended() throws Exception { + for (boolean beforeOffer : Arrays.asList(true, false)) { + Fixture fixture = new Fixture(1, 1); + fixture.assignment = Mockito.spy(fixture.assignment); + fixture.assignment.addToQueue(Collections.singletonList(split(0))); + UserException expected = new UserException("another producer failed"); + Mockito.doAnswer(call -> { + if (beforeOffer) { + fixture.assignment.setException(expected); + } + boolean offered = (boolean) call.callRealMethod(); + if (!beforeOffer) { + fixture.assignment.setException(expected); + } + return offered; + }).when(fixture.assignment).offerBatch(Mockito.any(), Mockito.any(), Mockito.any()); + ManualExecutor executor = new ManualExecutor(); + RecordingSource source = new RecordingSource(1, 2); + fixture.start(executor, source); + executor.runUntilIdle(); + Assertions.assertSame(expected, Assertions.assertThrows(UserException.class, + () -> fixture.assignment.getAssignedSplits(fixture.backend))); + Assertions.assertEquals(1, source.closeCalls.get()); + Assertions.assertEquals(1, fixture.completions.get()); + } + } + + private static List drain(SplitSource source, int batchSize) throws Exception { + List result = new ArrayList<>(); + while (true) { + List batch = source.getNextBatch(batchSize); + Assertions.assertTrue(batch.size() <= batchSize); + if (batch.isEmpty()) { + return result; + } + result.addAll(identities(batch)); + } + } + + private static List identities(List ranges) { + List result = new ArrayList<>(); + for (TScanRangeLocations range : ranges) { + result.add(range.getLocations().get(0).getBackendId()); + } + return result; + } + + private static Split split(long identity) { + return new FileSplit(LocationPath.of("s3://bucket/file.parquet"), identity, 1, 100, 0, + null, Collections.emptyList()); + } + + private static class RecordingSource implements SplitProducer.BatchSource { + private final List splits = new ArrayList<>(); + private final int batchSize; + private int next; + final AtomicInteger closeCalls = new AtomicInteger(); + final CountDownLatch closed = new CountDownLatch(1); + + RecordingSource(int batchSize, long... identities) { + this.batchSize = batchSize; + for (long identity : identities) { + splits.add(split(identity)); + } + } + + @Override + public List getNextBatch() throws IOException { + int end = Math.min(splits.size(), next + batchSize); + List batch = splits.subList(next, end); + next = end; + return batch; + } + + @Override + public void close() { + closeCalls.incrementAndGet(); + closed.countDown(); + } + } + + private static class ManualExecutor implements Executor { + final Queue tasks = new ArrayDeque<>(); + boolean reject; + + @Override + public void execute(Runnable task) { + if (reject) { + throw new RejectedExecutionException("executor stopped"); + } + tasks.add(task); + } + + void runNext() { + Assertions.assertFalse(tasks.isEmpty()); + tasks.remove().run(); + } + + void runUntilIdle() { + for (int i = 0; i < 100 && !tasks.isEmpty(); i++) { + runNext(); + } + Assertions.assertTrue(tasks.isEmpty(), "Producer must yield on a full queue or terminate at EOF"); + } + } + + private static class Fixture { + final Backend backend = new Backend(0, "127.0.0.1", 9050); + final Backend otherBackend = new Backend(1, "127.0.0.1", 9051); + final SplitGenerator generator = Mockito.mock(SplitGenerator.class); + final FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + final AtomicInteger completions = new AtomicInteger(); + final CountDownLatch completed = new CountDownLatch(1); + final Map conversions = new HashMap<>(); + SplitAssignment assignment; + + Fixture(int batchSize, int queueCapacity) throws Exception { + Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(call -> { + Multimap result = LinkedListMultimap.create(); + for (Split split : call.>getArgument(0)) { + result.put(split.getStart() % 2 == 0 ? backend : otherBackend, split); + } + return result; + }); + assignment = new SplitAssignment(policy, generator, (be, properties, split, keys, admission) -> { + conversions.merge(split.getStart(), 1, Integer::sum); + return new TScanRangeLocations().setLocations(Collections.singletonList( + new TScanRangeLocation().setBackendId(split.getStart()))); + }, Collections.emptyMap(), Collections.emptyList(), true, batchSize, queueCapacity); + } + + void start(Executor executor, RecordingSource source) { + new SplitProducer(assignment, executor, source, failure -> { + if (failure != null) { + assignment.setException(new UserException(failure.getMessage(), failure)); + } + completions.incrementAndGet(); + assignment.finishSchedule(); + completed.countDown(); + }).start(); + } + + SplitSource source() { + return new SplitSource(backend, assignment, 1000); + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitSourceTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitSourceTest.java new file mode 100644 index 00000000000000..ca2c5fb5f7d918 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitSourceTest.java @@ -0,0 +1,204 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource; + +import org.apache.doris.common.UserException; +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TScanRangeLocations; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.Multimap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +public class SplitSourceTest { + private Backend backend; + private SplitAssignment assignment; + private List splits; + private List ranges; + + @BeforeEach + void setUp() throws Exception { + backend = Mockito.mock(Backend.class); + FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(invocation -> { + Multimap batch = ArrayListMultimap.create(); + for (Split split : invocation.>getArgument(0)) { + batch.put(backend, split); + } + return batch; + }); + splits = new ArrayList<>(); + ranges = new ArrayList<>(); + IdentityHashMap locations = new IdentityHashMap<>(); + for (int i = 0; i < 8; i++) { + Split split = Mockito.mock(Split.class); + TScanRangeLocations range = Mockito.mock(TScanRangeLocations.class); + splits.add(split); + ranges.add(range); + locations.put(split, range); + } + assignment = new SplitAssignment(policy, Mockito.mock(SplitGenerator.class), + (be, properties, split, keys, admission) -> locations.get(split), + Collections.emptyMap(), Collections.emptyList(), true, 3, 8); + } + + @Test + void drainPendingBatchBeforeFinishedQueueBecomesEof() throws Exception { + assignment.addToQueue(splits); + assignment.finishSchedule(); + SplitSource source = new SplitSource(backend, assignment, 1000); + + List actual = new ArrayList<>(); + for (int requested : Arrays.asList(2, 1, 4, 1)) { + List batch = source.getNextBatch(requested); + Assertions.assertEquals(requested, batch.size()); + actual.addAll(batch); + } + Assertions.assertEquals(ranges, actual); + Assertions.assertTrue(source.getNextBatch(2).isEmpty()); + Assertions.assertTrue(source.getNextBatch(2).isEmpty()); + } + + @Test + void producerBatchesAndConsumerCreatedQueueRespectConfiguredLimits() throws Exception { + BlockingQueue queue = assignment.getAssignedSplits(backend); + Assertions.assertEquals(8, queue.remainingCapacity()); + assignment.addToQueue(splits); + Assertions.assertEquals(8, queue.size()); + Assertions.assertEquals(ranges, new ArrayList<>(queue)); + } + + @Test + void fullQueueResumesAfterConsumerFetch() throws Exception { + assignment.addToQueue(splits); + CountDownLatch started = new CountDownLatch(1); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future producer = executor.submit(() -> { + started.countDown(); + assignment.addToQueue(Collections.singletonList(splits.get(0))); + return null; + }); + Assertions.assertTrue(started.await(5, TimeUnit.SECONDS)); + Assertions.assertThrows(TimeoutException.class, () -> producer.get(200, TimeUnit.MILLISECONDS)); + SplitSource source = new SplitSource(backend, assignment, 1000); + Assertions.assertEquals(ranges.subList(0, 1), source.getNextBatch(1)); + producer.get(5, TimeUnit.SECONDS); + Assertions.assertEquals(8, assignment.getAssignedSplits(backend).size()); + } finally { + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void availablePartialBatchDoesNotWaitForProducer() throws Exception { + assignment.addToQueue(splits.subList(0, 1)); + SplitSource source = new SplitSource(backend, assignment, 5000); + Assertions.assertTimeoutPreemptively(Duration.ofSeconds(1), () -> + Assertions.assertEquals(ranges.subList(0, 1), source.getNextBatch(100))); + } + + @Test + void liveEmptyPollIsIncompleteRatherThanEof() { + SplitSource source = new SplitSource(backend, assignment, 20); + Assertions.assertDoesNotThrow(() -> Assertions.assertTrue(source.getNextBatch(1).isEmpty())); + Assertions.assertFalse(source.isLastBatch()); + } + + @Test + void errorIsCheckedWhilePendingSplitsRemain() throws Exception { + assignment.addToQueue(splits.subList(0, 3)); + SplitSource source = new SplitSource(backend, assignment, 1000); + source.getNextBatch(1); + UserException failure = new UserException("failed to generate splits"); + assignment.setException(failure); + Assertions.assertSame(failure, Assertions.assertThrows(UserException.class, () -> source.getNextBatch(1))); + } + + @Test + void concurrentRequestsDoNotDuplicateOrLosePendingSplits() throws Exception { + assignment.addToQueue(splits); + assignment.finishSchedule(); + SplitSource source = new SplitSource(backend, assignment, 1000); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future> first = executor.submit(() -> source.getNextBatch(2)); + Future> second = executor.submit(() -> source.getNextBatch(2)); + List actual = new ArrayList<>(first.get(5, TimeUnit.SECONDS)); + actual.addAll(second.get(5, TimeUnit.SECONDS)); + actual.addAll(source.getNextBatch(8)); + Set unique = Collections.newSetFromMap(new IdentityHashMap<>()); + unique.addAll(actual); + Assertions.assertEquals(8, actual.size()); + Assertions.assertEquals(8, unique.size()); + Assertions.assertTrue(unique.containsAll(ranges)); + } finally { + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void cancellationUnblocksProducerAndClosesLateResources() throws Exception { + assignment.addToQueue(splits); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future producer = executor.submit(() -> { + assignment.addToQueue(Collections.singletonList(splits.get(0))); + return null; + }); + Assertions.assertThrows(TimeoutException.class, () -> producer.get(200, TimeUnit.MILLISECONDS)); + assignment.stop(); + producer.get(5, TimeUnit.SECONDS); + CountDownLatch closed = new CountDownLatch(1); + assignment.addCloseable(closed::countDown); + Assertions.assertEquals(0, closed.getCount()); + Assertions.assertTrue(new SplitSource(backend, assignment, 1000).getNextBatch(1).isEmpty()); + } finally { + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void invalidRequestSizeDoesNotReturnFalseEof() { + SplitSource source = new SplitSource(backend, assignment, 1000); + Assertions.assertThrows(IllegalArgumentException.class, () -> source.getNextBatch(0)); + Assertions.assertThrows(IllegalArgumentException.class, () -> source.getNextBatch(-1)); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitStreamingConcurrencyTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitStreamingConcurrencyTest.java new file mode 100644 index 00000000000000..373eb247ef04e3 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/SplitStreamingConcurrencyTest.java @@ -0,0 +1,159 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource; + +import org.apache.doris.common.util.LocationPath; +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TScanRangeLocation; +import org.apache.doris.thrift.TScanRangeLocations; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.Multimap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +class SplitStreamingConcurrencyTest { + @Test + void consumerCreatedQueuesStayBoundedAcrossBackends() throws Exception { + verifyConcurrentStreaming(true); + } + + @Test + void producerCreatedQueuesStayBoundedAcrossBackends() throws Exception { + verifyConcurrentStreaming(false); + } + + private void verifyConcurrentStreaming(boolean consumerFirst) throws Exception { + int producerBatchSize = 7; + int queueCapacity = 7; + int splitCount = 10003; + int[] requestedSizes = {1, 10, 16}; + List backends = new ArrayList<>(); + for (int i = 0; i < requestedSizes.length; i++) { + backends.add(new Backend(i, "127.0.0.1", 9050 + i)); + } + FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + Mockito.when(policy.computeScanRangeAssignment(Mockito.any())).thenAnswer(invocation -> { + List batch = invocation.getArgument(0); + Assertions.assertTrue(batch.size() <= producerBatchSize); + Multimap assigned = ArrayListMultimap.create(); + for (Split split : batch) { + assigned.put(backends.get((int) (split.getStart() % backends.size())), split); + } + return assigned; + }); + CountDownLatch queuesFilled = new CountDownLatch(1); + SplitAssignment assignment = new SplitAssignment(policy, Mockito.mock(SplitGenerator.class), + (backend, properties, split, keys, admission) -> { + // The next range follows one full queue of individual ranges per backend. + if (split.getStart() >= queueCapacity * backends.size()) { + queuesFilled.countDown(); + } + // Preserve a unique split identity through assignment and residual batches. + return new TScanRangeLocations().setLocations(Collections.singletonList( + new TScanRangeLocation().setBackendId(split.getStart()))); + }, Collections.emptyMap(), Collections.emptyList(), true, producerBatchSize, queueCapacity); + List splits = new ArrayList<>(); + LocationPath path = LocationPath.of("s3://bucket/file.parquet"); + for (int i = 0; i < splitCount; i++) { + splits.add(new FileSplit(path, i, 1, splitCount, 0, null, Collections.emptyList())); + } + List> queues = new ArrayList<>(); + if (consumerFirst) { + for (Backend backend : backends) { + queues.add(assignment.getAssignedSplits(backend)); + } + } + ExecutorService executor = Executors.newFixedThreadPool(backends.size() + 1); + try { + Future producer = executor.submit(() -> { + assignment.addToQueue(splits); + assignment.finishSchedule(); + return null; + }); + Assertions.assertTrue(queuesFilled.await(5, TimeUnit.SECONDS)); + if (!consumerFirst) { + for (Backend backend : backends) { + queues.add(assignment.getAssignedSplits(backend)); + } + } + for (BlockingQueue queue : queues) { + Assertions.assertEquals(queueCapacity, queue.size()); + Assertions.assertEquals(0, queue.remainingCapacity()); + verifyQueue(queue, producerBatchSize, queueCapacity); + } + List>> consumers = new ArrayList<>(); + for (int i = 0; i < backends.size(); i++) { + int backendIndex = i; + SplitSource source = new SplitSource(backends.get(i), assignment, 10000); + consumers.add(executor.submit(() -> { + Set identities = new HashSet<>(); + while (true) { + verifyQueue(queues.get(backendIndex), producerBatchSize, queueCapacity); + List batch = source.getNextBatch(requestedSizes[backendIndex]); + Assertions.assertTrue(batch.size() <= requestedSizes[backendIndex]); + if (batch.isEmpty()) { + Assertions.assertTrue(source.getNextBatch(1).isEmpty()); + return identities; + } + for (TScanRangeLocations range : batch) { + long identity = range.getLocations().get(0).getBackendId(); + Assertions.assertEquals(backendIndex, identity % backends.size()); + Assertions.assertTrue(identities.add(identity), "Duplicate split " + identity); + } + } + })); + } + producer.get(30, TimeUnit.SECONDS); + Set actual = new HashSet<>(); + for (Future> consumer : consumers) { + actual.addAll(consumer.get(30, TimeUnit.SECONDS)); + } + Assertions.assertEquals(splitCount, actual.size()); + for (long i = 0; i < splitCount; i++) { + Assertions.assertTrue(actual.contains(i), "Missing split " + i); + } + System.out.printf("Split streaming: consumerFirst=%s, backends=%d, splits=%d, " + + "producerBatchLimit=%d, observedQueuePeak=%d, requests=[1,10,16]%n", + consumerFirst, backends.size(), splitCount, producerBatchSize, queueCapacity); + } finally { + assignment.stop(); + executor.shutdownNow(); + Assertions.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private void verifyQueue(BlockingQueue queue, + int producerBatchSize, int queueCapacity) { + Assertions.assertTrue(queue.size() <= queueCapacity); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HiveMetaStoreCacheTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HiveMetaStoreCacheTest.java index f067c879705377..14c65454354161 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HiveMetaStoreCacheTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HiveMetaStoreCacheTest.java @@ -48,6 +48,41 @@ public class HiveMetaStoreCacheTest { + @Test + public void testCachedFileCountUsesOnlyPresentEntriesAndHonorsInvalidation() { + ThreadPoolExecutor executor = ThreadPoolManager.newDaemonFixedThreadPool(1, 1, "hint-refresh", 1, false); + ThreadPoolExecutor listExecutor = ThreadPoolManager.newDaemonFixedThreadPool(1, 1, "hint-file", 1, false); + try { + HiveExternalMetaCache cache = new HiveExternalMetaCache(executor, listExecutor); + NameMapping name = NameMapping.createForTest("db", "files"); + HivePartition partition = new HivePartition(name, false, "parquet", "hdfs://test/p=1", + Collections.singletonList("1"), Collections.emptyMap()); + List partitions = Collections.singletonList(partition); + Assertions.assertEquals(0, cache.getCachedFileCount(Collections.emptyList())); + Assertions.assertEquals(-1, cache.getCachedFileCount(partitions)); + cache.initCatalog(0, new HashMap<>()); + MetaCacheEntry entry = + cache.entry(0, HiveExternalMetaCache.ENTRY_FILE, + HiveExternalMetaCache.FileCacheKey.class, HiveExternalMetaCache.FileCacheValue.class); + HiveExternalMetaCache.FileCacheKey key = new HiveExternalMetaCache.FileCacheKey(0, + Util.genIdByName("db", "files"), partition.getPath(), "parquet", partition.getPartitionValues()); + HiveExternalMetaCache.FileCacheValue files = new HiveExternalMetaCache.FileCacheValue(); + for (int i = 0; i < 2; i++) { + HiveExternalMetaCache.HiveFileStatus file = new HiveExternalMetaCache.HiveFileStatus(); + file.setPath(org.apache.doris.common.util.LocationPath.of("hdfs://test/p=1/" + i)); + file.setLength(10); + files.getFiles().add(file); + } + entry.put(key, files); + Assertions.assertEquals(2, cache.getCachedFileCount(partitions)); + cache.invalidateTable(0, "db", "files"); + Assertions.assertEquals(-1, cache.getCachedFileCount(partitions)); + } finally { + executor.shutdownNow(); + listExecutor.shutdownNow(); + } + } + @Test public void testPartitionValueWeightScalesLinearlyToOneHundredThousandPartitions() { HiveExternalMetaCache.PartitionValueCacheKey key = new HiveExternalMetaCache.PartitionValueCacheKey( diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java index c3e05409d3e40a..56fcfb59bdf103 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java @@ -76,6 +76,8 @@ public void testNativeBatchAndSynchronousFailuresKeepRootCause() throws Exceptio partitions.set(node, Collections.singletonList(Mockito.mock(HivePartition.class))); org.apache.doris.datasource.SplitAssignment assignment = Mockito.mock(org.apache.doris.datasource.SplitAssignment.class); + Mockito.when(assignment.needMoreSplit()).thenReturn(true); + Mockito.when(assignment.getSplitBatchSize()).thenReturn(8); Field splitAssignment = org.apache.doris.planner.ScanNode.class.getDeclaredField("splitAssignment"); splitAssignment.setAccessible(true); splitAssignment.set(node, assignment); diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveSplitEstimateTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveSplitEstimateTest.java index 01a9c94df6f1d0..f3ef1c69dd619d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveSplitEstimateTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveSplitEstimateTest.java @@ -26,9 +26,11 @@ import org.apache.doris.catalog.Type; import org.apache.doris.common.util.LocationPath; import org.apache.doris.datasource.ExternalMetaCacheMgr; +import org.apache.doris.datasource.FederationBackendPolicy; import org.apache.doris.datasource.FileQueryScanNode; import org.apache.doris.datasource.FileSplitter; import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.SplitSource; import org.apache.doris.datasource.hive.HMSExternalCatalog; import org.apache.doris.datasource.hive.HMSExternalTable; import org.apache.doris.datasource.hive.HiveExternalMetaCache; @@ -40,8 +42,13 @@ import org.apache.doris.planner.ScanNode; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.SessionVariable; +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; import org.apache.doris.thrift.TFileFormatType; +import org.apache.doris.thrift.TScanRangeLocations; +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.Multimap; import org.junit.Assert; import org.junit.Test; import org.mockito.MockedStatic; @@ -58,7 +65,6 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; public class HiveSplitEstimateTest { @@ -157,13 +163,28 @@ private int runSplitPlanning(int batchSize, TableSample sample, boolean choosePl Field initialized = HiveScanNode.class.getDeclaredField("partitionInit"); initialized.setAccessible(true); initialized.set(node, true); - SplitAssignment assignment = Mockito.mock(SplitAssignment.class); - Mockito.when(assignment.needMoreSplit()).thenReturn(true); + Backend backend = new Backend(1, "127.0.0.1", 9050); + FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + Mockito.when(policy.numBackends()).thenReturn(1); + Mockito.when(policy.computeScanRangeAssignment(Mockito.anyList())).thenAnswer(call -> { + Multimap assigned = ArrayListMultimap.create(); + for (Split split : call.>getArgument(0)) { + assigned.put(backend, split); + } + return assigned; + }); + // Retain the complete small fixture so scheduling can finish before consumption. + // Exercise the real producer/assignment path rather than the removed addToQueue call. + SplitAssignment assignment = Mockito.spy(new SplitAssignment(policy, node, + (be, properties, split, keys, admission) -> new TScanRangeLocations(), + Collections.emptyMap(), Collections.emptyList(), true, + variables.getSplitBatchSize(), Math.max(1, totalFiles))); Field assignmentField = ScanNode.class.getDeclaredField("splitAssignment"); assignmentField.setAccessible(true); assignmentField.set(node, assignment); CompletableFuture finished = new CompletableFuture<>(); Mockito.doAnswer(call -> { + call.callRealMethod(); finished.complete(null); return null; }).when(assignment).finishSchedule(); @@ -171,11 +192,6 @@ private int runSplitPlanning(int batchSize, TableSample sample, boolean choosePl finished.completeExceptionally(call.getArgument(0)); return null; }).when(assignment).setException(Mockito.any()); - AtomicInteger delivered = new AtomicInteger(); - Mockito.doAnswer(call -> { - delivered.addAndGet(call.>getArgument(0).size()); - return null; - }).when(assignment).addToQueue(Mockito.anyList()); HiveExternalMetaCache cache = Mockito.mock(HiveExternalMetaCache.class); Mockito.when(cache.getFilesByPartitions(Mockito.anyList(), Mockito.anyBoolean(), Mockito.anyBoolean(), Mockito.isNull(), Mockito.same(table))).thenAnswer(call -> call.>getArgument(0) @@ -203,12 +219,18 @@ private int runSplitPlanning(int batchSize, TableSample sample, boolean choosePl node.startSplit(1); finished.get(5, TimeUnit.SECONDS); Mockito.verify(assignment, Mockito.never()).setException(Mockito.any()); + int delivered = 0; + SplitSource source = new SplitSource(backend, assignment, 1000); + while (!source.isLastBatch()) { + delivered += source.getNextBatch(variables.getSplitBatchSize()).size(); + } if (choosePlanningMode) { - return delivered.get(); + return delivered; } - Assert.assertEquals(totalFiles, delivered.get()); + Assert.assertEquals(totalFiles, delivered); return node.numApproximateSplits(); } finally { + assignment.stop(); executor.shutdownNow(); Assert.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveStreamingSplitTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveStreamingSplitTest.java new file mode 100644 index 00000000000000..ed3a7d414584e9 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveStreamingSplitTest.java @@ -0,0 +1,642 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource.hive.source; + +import org.apache.doris.analysis.TupleDescriptor; +import org.apache.doris.analysis.TupleId; +import org.apache.doris.blockrule.SqlBlockRuleMgr; +import org.apache.doris.common.Config; +import org.apache.doris.common.UserException; +import org.apache.doris.common.util.FairTaskExecutor; +import org.apache.doris.common.util.LocationPath; +import org.apache.doris.datasource.CacheException; +import org.apache.doris.datasource.ExternalMetaCacheMgr; +import org.apache.doris.datasource.FederationBackendPolicy; +import org.apache.doris.datasource.FileSplit; +import org.apache.doris.datasource.FileSplitter; +import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.SplitProducer; +import org.apache.doris.datasource.SplitSource; +import org.apache.doris.datasource.SplitSourceManager; +import org.apache.doris.datasource.hive.HMSExternalCatalog; +import org.apache.doris.datasource.hive.HMSExternalTable; +import org.apache.doris.datasource.hive.HiveExternalMetaCache; +import org.apache.doris.datasource.hive.HivePartition; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.planner.PlanNodeId; +import org.apache.doris.planner.ScanContext; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.spi.Split; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TExplainLevel; +import org.apache.doris.thrift.TFileFormatType; +import org.apache.doris.thrift.TFileScanRangeParams; +import org.apache.doris.thrift.TPushAggOp; +import org.apache.doris.thrift.TScanRangeLocation; +import org.apache.doris.thrift.TScanRangeLocations; + +import com.google.common.collect.LinkedListMultimap; +import com.google.common.collect.Multimap; +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +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; + +public class HiveStreamingSplitTest { + @Test + void realModeGateAndDispatchPreservePartitionOverlapAndEveryRange() throws Exception { + try (Fixture f = new Fixture(8, 2)) { + // Both gates qualify: partition routing must win and preserve parallel listing. + Mockito.when(f.cache.getCachedFileCount(Mockito.anyList())).thenReturn(10000L); + Assertions.assertTrue(f.node.isBatchMode()); + Mockito.verify(f.cache, Mockito.never()).getCachedFileCount(Mockito.anyList()); + f.node.startSplit(1); + Assertions.assertTrue(f.entered.await(10, TimeUnit.SECONDS), + "Two actual partition listings must overlap through startSplit"); + Assertions.assertEquals(2, f.listings.get()); + f.release.countDown(); + List actual = f.drain(); + Set expected = expectedRanges(8); + Assertions.assertEquals(expected.size(), actual.size()); + Assertions.assertEquals(expected, new HashSet<>(actual)); + } + } + + @Test + void partitionEstimateAccumulatesAcrossSmallProducerBatches() throws Exception { + try (Fixture f = new Fixture(1, 1, 1)) { + Assertions.assertTrue(f.node.isBatchMode()); + f.node.startSplit(1); + f.release.countDown(); + Assertions.assertEquals(17, f.drain().size()); + Assertions.assertEquals(17, f.node.numApproximateSplits()); + } + } + + @Test + void partitionEstimateDoesNotCountSeveralPartitionsAsOne() throws Exception { + try (Fixture f = new Fixture(3, 1, 64)) { + f.fileLengths[0] = 17; + f.fileLengths[1] = 4; + f.fileLengths[2] = 8; + Assertions.assertTrue(f.node.isBatchMode()); + f.release.countDown(); + SplitProducer.BatchSource source = f.node.createSplitBatchSource(f.cache, new AtomicInteger(), "", 1); + try { + Assertions.assertEquals(29, source.getNextBatch().size()); + Assertions.assertTrue(source.getNextBatch().isEmpty()); + Assertions.assertEquals(51, f.node.numApproximateSplits()); + } finally { + source.close(); + } + } + } + + @Test + void countPushdownUsesOneScanWideDecisionWithoutRelisting() throws Exception { + long fileCacheLimit = Config.max_external_file_cache_num; + try { + Config.max_external_file_cache_num = 0; + for (TFileFormatType format : new TFileFormatType[] { + TFileFormatType.FORMAT_PARQUET, TFileFormatType.FORMAT_ORC}) { + for (boolean scannerV2 : new boolean[] {false, true}) { + assertCountUsesOneScanWideDecision(format, scannerV2); + } + } + } finally { + Config.max_external_file_cache_num = fileCacheLimit; + } + } + + private void assertCountUsesOneScanWideDecision(TFileFormatType format, boolean scannerV2) throws Exception { + try (Fixture f = new Fixture(100, 1)) { + Arrays.fill(f.fileLengths, 32); + f.variables.setFileSplitSize(4); + f.variables.parallelPipelineTaskNum = 32; + f.variables.enableFileScannerV2 = scannerV2; + Mockito.when(f.table.getFileFormatType(Mockito.any())).thenReturn(format); + f.node.setPushDownAggNoGrouping(TPushAggOp.COUNT); + Assertions.assertTrue(f.node.isBatchMode()); + f.release.countDown(); + f.node.startSplit(1); + Assertions.assertEquals(100, f.drain().size(), "one metadata range per selected file"); + int prefetched = f.listedAtFirstAssignment.get(); + Assertions.assertTrue(prefetched >= 32 && prefetched <= 33, + "stop at the parallelism threshold, allowing one other in-flight partition: " + prefetched); + Assertions.assertEquals(100, f.listedPartitions.get(), + "reuse looked-ahead metadata even with the global file cache disabled"); + } + } + + @Test + void countPushdownStillSplitsWhenTheWholeScanHasTooFewFiles() throws Exception { + try (Fixture f = new Fixture(2, 1)) { + Arrays.fill(f.fileLengths, 32); + f.variables.setFileSplitSize(4); + f.variables.parallelPipelineTaskNum = 32; + f.node.setPushDownAggNoGrouping(TPushAggOp.COUNT); + Assertions.assertTrue(f.node.isBatchMode()); + f.release.countDown(); + f.node.startSplit(1); + Assertions.assertEquals(16, f.drain().size()); + Assertions.assertEquals(2, f.listedPartitions.get()); + } + } + + @Test + void maximumBatchSettingAllocatesOnlyForActualSplits() throws Exception { + try (Fixture f = new Fixture(1, 1, Integer.MAX_VALUE)) { + f.fileLengths[0] = 257; + Assertions.assertTrue(f.node.isBatchMode()); + f.release.countDown(); + try (SplitProducer.BatchSource source = f.node.createSplitBatchSource( + f.cache, new AtomicInteger(), "", 1)) { + Assertions.assertEquals(257, source.getNextBatch().size(), + "the initial capacity must not cap the configured batch maximum"); + Assertions.assertTrue(source.getNextBatch().isEmpty()); + } + } + } + + @Test + void countLookaheadFailurePreservesItsCauseAndFinishesPlanning() throws Exception { + try (Fixture f = new Fixture(100, 1)) { + f.variables.parallelPipelineTaskNum = 32; + f.node.setPushDownAggNoGrouping(TPushAggOp.COUNT); + UserException missing = new UserException("Partition location does not exist: hdfs://test/missing"); + CacheException failure = new CacheException("failed to list count files", missing); + Mockito.doThrow(failure).when(f.cache).getFilesByPartitions(Mockito.anyList(), + Mockito.anyBoolean(), Mockito.anyBoolean(), Mockito.isNull(), Mockito.eq(f.table)); + UserException error = Assertions.assertThrows(UserException.class, f.node::createScanRangeLocations); + Assertions.assertTrue(error.getMessage().contains("hdfs://test/missing")); + Assertions.assertSame(failure, error.getCause()); + } + } + + @Test + void cancellingCountLookaheadInterruptsListingAndStartsNoRangeProducers() throws Exception { + try (Fixture f = new Fixture(100, 1)) { + f.variables.parallelPipelineTaskNum = 32; + f.node.setPushDownAggNoGrouping(TPushAggOp.COUNT); + Assertions.assertTrue(f.node.isBatchMode()); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + Mockito.doAnswer(call -> { + entered.countDown(); + try { + new CountDownLatch(1).await(); + throw new AssertionError("listing must be interrupted"); + } catch (InterruptedException e) { + interrupted.countDown(); + throw e; + } + }).when(f.cache).getFilesByPartitions(Mockito.anyList(), Mockito.anyBoolean(), Mockito.anyBoolean(), + Mockito.isNull(), Mockito.eq(f.table)); + f.node.startSplit(1); + Assertions.assertTrue(entered.await(5, TimeUnit.SECONDS)); + f.node.stop(); + Assertions.assertTrue(interrupted.await(5, TimeUnit.SECONDS)); + f.executor.shutdown(); + Assertions.assertTrue(f.executor.awaitTermination(5, TimeUnit.SECONDS)); + Assertions.assertEquals(-1, f.listedAtFirstAssignment.get()); + Mockito.verify(f.cache, Mockito.atMost(2)).getFilesByPartitions(Mockito.anyList(), + Mockito.anyBoolean(), Mockito.anyBoolean(), Mockito.isNull(), Mockito.eq(f.table)); + } + } + + @Test + void fileCountGateUsesOneProducerAndHonorsSmallConsumerBatch() throws Exception { + try (Fixture f = new Fixture(2, 1024)) { + Mockito.when(f.cache.getCachedFileCount(Mockito.anyList())).thenReturn(10000L); + Assertions.assertTrue(f.node.isBatchMode()); + f.node.startSplit(1); + f.release.countDown(); + List actual = f.drain(); + Assertions.assertEquals(34, actual.size()); + Assertions.assertEquals(34, new HashSet<>(actual).size()); + Assertions.assertEquals(1, f.peak.get()); + Mockito.verify(f.cache).getCachedFileCount(Mockito.anyList()); + } + } + + @Test + void unknownFileCountKeepsParallelListingsWhenFileCacheIsDisabled() throws Exception { + long previousFileCacheLimit = Config.max_external_file_cache_num; + try (Fixture f = new Fixture(8, 1024)) { + Config.max_external_file_cache_num = 0; + Mockito.when(f.cache.getCachedFileCount(Mockito.anyList())).thenReturn(-1L); + Assertions.assertTrue(f.node.isBatchMode()); + f.node.startSplit(1); + Assertions.assertTrue(f.entered.await(10, TimeUnit.SECONDS), + "The file-only route must preserve concurrent uncached partition listings"); + Assertions.assertEquals(2, f.listings.get()); + Assertions.assertEquals(2, f.peak.get()); + f.release.countDown(); + List actual = f.drain(); + Set expected = expectedRanges(8); + Assertions.assertEquals(expected.size(), actual.size()); + Assertions.assertEquals(expected, new HashSet<>(actual)); + Assertions.assertTrue(f.peak.get() <= 2); + Mockito.verify(f.cache, Mockito.atLeastOnce()).getFilesByPartitions(Mockito.anyList(), + Mockito.eq(false), Mockito.eq(false), Mockito.isNull(), Mockito.eq(f.table)); + } finally { + Config.max_external_file_cache_num = previousFileCacheLimit; + } + } + + @Test + void coldSmallTableKeepsBatchRangesWhenRenderingVerboseExplain() throws Exception { + try (Fixture f = new Fixture(1, 1024)) { + Mockito.when(f.cache.getCachedFileCount(Mockito.anyList())) + .thenAnswer(call -> f.listings.get() == 0 ? -1L : 1L); + f.release.countDown(); + f.node.createScanRangeLocations(); + Assertions.assertEquals(1, f.listings.get()); + Assertions.assertTrue(f.node.getNodeExplainString("", TExplainLevel.VERBOSE) + .contains("(approximate)inputSplitNum")); + Assertions.assertEquals(17, f.drain().size()); + Mockito.verify(f.cache).getCachedFileCount(Mockito.anyList()); + } + } + + @Test + void synchronousModeStaysSelectedAfterFileCacheEviction() throws Exception { + try (Fixture f = new Fixture(1, 1024)) { + Mockito.when(f.cache.getCachedFileCount(Mockito.anyList())).thenReturn(1L, -1L); + Assertions.assertFalse(f.node.isBatchMode()); + Assertions.assertFalse(f.node.isBatchMode()); + Mockito.verify(f.cache).getCachedFileCount(Mockito.anyList()); + } + } + + @Test + void statementCleanupStopsAbandonedPlanningWithoutListingAllPartitions() throws Exception { + try (Fixture f = new Fixture(100, 2)) { + f.context.setExecutor(Mockito.mock(StmtExecutor.class)); + Assertions.assertTrue(f.node.isBatchMode()); + f.node.startSplit(1); + Assertions.assertTrue(f.entered.await(10, TimeUnit.SECONDS)); + SplitAssignment assignment = (SplitAssignment) f.getAssignment(); + SplitSource source = new SplitSource(f.backend, assignment, 10000); + f.manager.registerSplitSource(source); + f.statement.close(); + Assertions.assertTrue(assignment.isStop()); + Assertions.assertNull(f.manager.getSplitSource(source.getUniqueId())); + f.release.countDown(); + f.executor.shutdown(); + Assertions.assertTrue(f.executor.awaitTermination(10, TimeUnit.SECONDS)); + Assertions.assertEquals(2, f.listings.get()); + } + } + + @Test + void deferredFlightScanSurvivesStatementCleanupAndContinuesFetching() throws Exception { + try (Fixture f = new Fixture(8, 2)) { + StmtExecutor owner = Mockito.mock(StmtExecutor.class); + f.context.setExecutor(owner); + Assertions.assertTrue(f.node.isBatchMode()); + f.node.startSplit(1); + Assertions.assertTrue(f.entered.await(10, TimeUnit.SECONDS)); + SplitAssignment assignment = (SplitAssignment) f.getAssignment(); + SplitSource source = new SplitSource(f.backend, assignment, 10000); + f.manager.registerSplitSource(source); + // Deferral is decided after planning, before GetFlightInfo returns. + Mockito.when(owner.isDeferredForArrowFlight()).thenReturn(true); + f.statement.close(); + Assertions.assertFalse(assignment.isStop()); + Assertions.assertSame(source, f.manager.getSplitSource(source.getUniqueId())); + f.release.countDown(); + List actual = f.drain(source); + Set expected = expectedRanges(8); + Assertions.assertEquals(expected.size(), actual.size()); + Assertions.assertEquals(expected, new HashSet<>(actual)); + f.node.stop(); + Assertions.assertNull(f.manager.getSplitSource(source.getUniqueId())); + } + } + + @Test + void sampledTransactionalAndDisabledScansKeepSynchronousPlanning() throws Exception { + try (Fixture f = new Fixture(2, 1)) { + f.variables.enableExternalTableBatchMode = false; + Assertions.assertFalse(f.node.isBatchMode()); + Mockito.verifyNoInteractions(f.cache); + } + try (Fixture f = new Fixture(2, 1)) { + Mockito.when(f.table.isHiveTransactionalTable()).thenReturn(true); + Assertions.assertFalse(f.node.isBatchMode()); + Mockito.verifyNoInteractions(f.cache); + } + try (Fixture f = new Fixture(2, 1)) { + f.node.setTableSample(new org.apache.doris.analysis.TableSample(true, 10L, 0L)); + Assertions.assertFalse(f.node.isBatchMode()); + Mockito.verifyNoInteractions(f.cache); + } + } + + @Test + void millionFilePartitionEnumeratesOnlyTheFirstConfiguredProducerBatch() throws Exception { + try (Fixture f = new Fixture(1, 1)) { + HiveExternalMetaCache.FileCacheValue files = Mockito.mock(HiveExternalMetaCache.FileCacheValue.class); + HiveExternalMetaCache.HiveFileStatus file = Mockito.spy(f.file(0)); + Mockito.when(files.getFiles()).thenReturn(Collections.nCopies(1000000, file)); + Mockito.when(files.isSplittable()).thenReturn(false); + Mockito.when(f.cache.getFilesByPartitions(Mockito.anyList(), Mockito.anyBoolean(), Mockito.anyBoolean(), + Mockito.isNull(), Mockito.eq(f.table))).thenReturn(Collections.singletonList(files)); + Assertions.assertTrue(f.node.isBatchMode()); + SplitProducer.BatchSource source = f.node.createSplitBatchSource(f.cache, new AtomicInteger(), "", 1); + Assertions.assertEquals(7, source.getNextBatch().size()); + Mockito.verify(file, Mockito.times(7)).getPath(); + source.close(); + } + } + + @Test + void coldStreamingScanPreservesMissingPartitionCauseDuringInitialization() throws Exception { + try (Fixture f = new Fixture(1, 1024)) { + Mockito.when(f.table.getName()).thenReturn("missing_partition_table"); + Mockito.when(f.cache.getCachedFileCount(Mockito.anyList())).thenReturn(-1L); + UserException missing = new UserException("Partition location does not exist: hdfs://test/missing"); + CacheException listingFailure = new CacheException("failed to get input splits", missing); + Mockito.doThrow(listingFailure).when(f.cache).getFilesByPartitions(Mockito.anyList(), + Mockito.anyBoolean(), Mockito.anyBoolean(), Mockito.isNull(), Mockito.eq(f.table)); + + UserException error = Assertions.assertThrows(UserException.class, f.node::createScanRangeLocations); + Assertions.assertTrue(error.getMessage().contains("Partition location does not exist: hdfs://test/missing"), + error.getMessage()); + Assertions.assertTrue(error.getMessage().contains("missing_partition_table"), error.getMessage()); + Assertions.assertSame(listingFailure, error.getCause()); + } + } + + @Test + void synchronousScanPreservesTheSameMissingPartitionCause() throws Exception { + try (Fixture f = new Fixture(1, 1024)) { + Mockito.when(f.table.getName()).thenReturn("missing_partition_table"); + UserException missing = new UserException("Partition location does not exist: hdfs://test/missing"); + CacheException listingFailure = new CacheException("failed to get input splits", missing); + Mockito.doThrow(listingFailure).when(f.cache).getFilesByPartitions(Mockito.anyList(), + Mockito.anyBoolean(), Mockito.anyBoolean(), Mockito.isNull(), Mockito.eq(f.table)); + + UserException error = Assertions.assertThrows(UserException.class, () -> f.node.getSplits(1)); + Assertions.assertTrue(error.getMessage().contains("Partition location does not exist: hdfs://test/missing"), + error.getMessage()); + Assertions.assertTrue(error.getMessage().contains("missing_partition_table"), error.getMessage()); + Assertions.assertSame(listingFailure, error.getCause()); + } + } + + private static Set expectedRanges(int partitions) { + Set expected = new HashSet<>(); + for (int p = 0; p < partitions; p++) { + for (int offset = 0; offset < 17; offset++) { + expected.add(p * 100L + offset); + } + } + return expected; + } + + private static void field(Object target, String name, Object value) throws Exception { + for (Class type = target.getClass(); type != null; type = type.getSuperclass()) { + try { + Field field = type.getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + return; + } catch (NoSuchFieldException e) { + // This test initializes inherited scan-level environment fields, not mode-selection logic. + } + } + throw new NoSuchFieldException(name); + } + + private static class TestNode extends HiveScanNode { + private final List partitions; + + TestNode(HMSExternalTable table, SessionVariable variables, List partitions) { + super(new PlanNodeId(0), descriptor(table), false, variables, null, ScanContext.EMPTY); + this.partitions = partitions; + fileSplitter = new FileSplitter(8, 8, 0); + params = new TFileScanRangeParams(); + } + + private static TupleDescriptor descriptor(HMSExternalTable table) { + TupleDescriptor desc = new TupleDescriptor(new TupleId(0)); + desc.setTable(table); + return desc; + } + + @Override + protected List getPartitions() { + return partitions; + } + + @Override + public List getPathPartitionKeys() { + return Collections.emptyList(); + } + + void initializeAssignment() { + splitAssignment = createSplitAssignment(Collections.emptyMap(), Collections.emptyList(), true); + } + + @Override + protected TScanRangeLocations splitToScanRange(Backend backend, java.util.Map properties, + Split split, List keys, boolean admission) { + FileSplit file = (FileSplit) split; + int partition = Integer.parseInt(file.getPathString().substring(file.getPathString().lastIndexOf('/') + 1)); + return new TScanRangeLocations().setLocations(Collections.singletonList( + new TScanRangeLocation().setBackendId(partition * 100L + split.getStart()))); + } + } + + private static class Fixture implements AutoCloseable { + final SessionVariable variables = new SessionVariable(); + final HMSExternalTable table = Mockito.mock(HMSExternalTable.class); + final HiveExternalMetaCache cache = Mockito.mock(HiveExternalMetaCache.class); + final SplitSourceManager manager = new SplitSourceManager(); + final Backend backend = new Backend(1, "localhost", 0); + final ExecutorService executor = Executors.newFixedThreadPool(2); + final FairTaskExecutor scheduler = + new FairTaskExecutor(2, + new ThreadFactoryBuilder().setDaemon(true).build()); + final CountDownLatch entered = new CountDownLatch(2); + final CountDownLatch release = new CountDownLatch(1); + final AtomicInteger listings = new AtomicInteger(); + final AtomicInteger listedPartitions = new AtomicInteger(); + final AtomicInteger listedAtFirstAssignment = new AtomicInteger(-1); + final AtomicInteger active = new AtomicInteger(); + final AtomicInteger peak = new AtomicInteger(); + final int previousParallelism = Config.max_external_cache_loader_thread_pool_size; + final ConnectContext previousContext = ConnectContext.get(); + final ConnectContext context = new ConnectContext(); + final StatementContext statement = new StatementContext(context, null); + final MockedStatic mockedEnv; + final TestNode node; + final long[] fileLengths; + + Fixture(int partitionCount, int partitionThreshold) throws Exception { + this(partitionCount, partitionThreshold, 7); + } + + Fixture(int partitionCount, int partitionThreshold, int producerBatchSize) throws Exception { + fileLengths = new long[partitionCount]; + Arrays.fill(fileLengths, 17); + Config.max_external_cache_loader_thread_pool_size = 2; + variables.numPartitionsInBatchMode = partitionThreshold; + variables.numFilesInBatchMode = 1024; + variables.splitBatchSize = producerBatchSize; + variables.splitQueueCapacity = 1; + variables.enableFileScannerV2 = false; + variables.setFileSplitSize(1); + context.setSessionVariable(variables); + context.setStatementContext(statement); + context.setThreadLocalInfo(); + HMSExternalCatalog catalog = Mockito.mock(HMSExternalCatalog.class); + Mockito.when(table.getCatalog()).thenReturn(catalog); + Mockito.when(catalog.bindBrokerName()).thenReturn(""); + Mockito.when(table.getFileFormatType(Mockito.any())).thenReturn(TFileFormatType.FORMAT_PARQUET); + org.apache.doris.catalog.Env env = Mockito.mock(org.apache.doris.catalog.Env.class); + ExternalMetaCacheMgr mgr = Mockito.mock(ExternalMetaCacheMgr.class); + Mockito.when(env.getExtMetaCacheMgr()).thenReturn(mgr); + Mockito.when(env.getSqlBlockRuleMgr()).thenReturn(Mockito.mock(SqlBlockRuleMgr.class)); + Mockito.when(env.getSplitSourceManager()).thenReturn(manager); + Mockito.when(mgr.hive(Mockito.anyLong())).thenReturn(cache); + Mockito.when(mgr.getScheduleExecutor()).thenReturn(executor); + Mockito.when(mgr.getScanTaskExecutor(Mockito.anyBoolean())).thenReturn(scheduler); + mockedEnv = Mockito.mockStatic(org.apache.doris.catalog.Env.class); + mockedEnv.when(org.apache.doris.catalog.Env::getCurrentEnv).thenReturn(env); + List partitions = new ArrayList<>(); + for (int p = 0; p < partitionCount; p++) { + partitions.add(new HivePartition(null, false, "parquet", "hdfs://test/" + p, + Collections.singletonList(Integer.toString(p)), Collections.emptyMap())); + } + node = new TestNode(table, variables, partitions); + FederationBackendPolicy policy = Mockito.mock(FederationBackendPolicy.class); + Mockito.when(policy.numBackends()).thenReturn(1); + Mockito.when(policy.getBackends()).thenReturn(Collections.singletonList(backend)); + Mockito.when(policy.computeScanRangeAssignment(Mockito.anyList())).thenAnswer(call -> { + listedAtFirstAssignment.compareAndSet(-1, listedPartitions.get()); + Multimap assigned = LinkedListMultimap.create(); + for (Split split : call.>getArgument(0)) { + assigned.put(backend, split); + } + return assigned; + }); + field(node, "backendPolicy", policy); + node.initializeAssignment(); + Mockito.when(cache.getFilesByPartitions(Mockito.anyList(), Mockito.anyBoolean(), Mockito.anyBoolean(), + Mockito.isNull(), Mockito.eq(table))).thenAnswer(call -> { + int count = active.incrementAndGet(); + peak.accumulateAndGet(count, Math::max); + listings.incrementAndGet(); + entered.countDown(); + try { + Assertions.assertTrue(release.await(10, TimeUnit.SECONDS)); + List selected = call.getArgument(0); + listedPartitions.addAndGet(selected.size()); + List loaded = new ArrayList<>(); + for (HivePartition p : selected) { + int number = Integer.parseInt(p.getPath().substring(p.getPath().lastIndexOf('/') + 1)); + HiveExternalMetaCache.FileCacheValue value = new HiveExternalMetaCache.FileCacheValue(); + value.setSplittable(true); + value.setPartitionValues(p.getPartitionValues()); + value.getFiles().add(file(number)); + loaded.add(value); + } + return loaded; + } finally { + active.decrementAndGet(); + } + }); + } + + HiveExternalMetaCache.HiveFileStatus file(int partition) { + HiveExternalMetaCache.HiveFileStatus status = new HiveExternalMetaCache.HiveFileStatus(); + status.setPath(LocationPath.of("hdfs://test/" + partition)); + status.setLength(fileLengths[partition]); + return status; + } + + List drain() throws Exception { + SplitAssignment assignment = (SplitAssignment) getAssignment(); + return drain(new SplitSource(backend, assignment, 10000)); + } + + List drain(SplitSource source) throws Exception { + List identities = new ArrayList<>(); + while (true) { + List batch = source.getNextBatch(3); + Assertions.assertTrue(batch.size() <= 3); + if (batch.isEmpty()) { + return identities; + } + for (TScanRangeLocations range : batch) { + identities.add(range.getLocations().get(0).getBackendId()); + } + } + } + + Object getAssignment() throws Exception { + for (Class type = node.getClass(); type != null; type = type.getSuperclass()) { + try { + Field field = type.getDeclaredField("splitAssignment"); + field.setAccessible(true); + return field.get(node); + } catch (NoSuchFieldException e) { + // Continue to the owning superclass. + } + } + throw new NoSuchFieldException("splitAssignment"); + } + + @Override + public void close() throws Exception { + node.stop(); + release.countDown(); + executor.shutdown(); + Assertions.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + try { + statement.close(); + } finally { + mockedEnv.close(); + Config.max_external_cache_loader_thread_pool_size = previousParallelism; + ConnectContext.remove(); + if (previousContext != null) { + previousContext.setThreadLocalInfo(); + } + } + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/hudi/source/HudiBatchFsViewOwnerTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/hudi/source/HudiBatchFsViewOwnerTest.java index 4c615717917b72..a40100e1b77adc 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/hudi/source/HudiBatchFsViewOwnerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hudi/source/HudiBatchFsViewOwnerTest.java @@ -17,17 +17,39 @@ package org.apache.doris.datasource.hudi.source; +import org.apache.doris.analysis.TupleDescriptor; +import org.apache.doris.analysis.TupleId; +import org.apache.doris.catalog.Env; import org.apache.doris.common.ThreadPoolManager; +import org.apache.doris.common.security.authentication.ExecutionAuthenticator; +import org.apache.doris.datasource.ExternalMetaCacheMgr; import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.SplitSource; +import org.apache.doris.datasource.SplitSourceManager; +import org.apache.doris.datasource.hive.HMSExternalCatalog; +import org.apache.doris.datasource.hive.HMSExternalTable; +import org.apache.doris.datasource.hive.HivePartition; +import org.apache.doris.datasource.hudi.HudiExternalMetaCache; import org.apache.doris.datasource.hudi.HudiFsViewCacheValue; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.planner.PlanNodeId; +import org.apache.doris.planner.ScanContext; +import org.apache.doris.planner.ScanNode; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.system.Backend; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; import org.mockito.Mockito; +import java.lang.reflect.Field; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Optional; import java.util.concurrent.CancellationException; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; @@ -36,12 +58,191 @@ import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.RejectedExecutionHandler; +import java.util.concurrent.Semaphore; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; class HudiBatchFsViewOwnerTest { + @Test + void deferredBatchSourceSurvivesStatementCloseUntilCoordinatorStopsIt() throws Exception { + try (BatchScanFixture f = new BatchScanFixture()) { + f.node.startSplit(1); + Assertions.assertEquals(1, f.submitted.size()); + SplitSource source = new SplitSource(new Backend(1, "localhost", 0), f.assignment, 1000); + f.manager.registerSplitSource(source); + // Deferral is selected after planning. Processor cleanup can also clear the context's executor. + Mockito.when(f.executor.isDeferredForArrowFlight()).thenReturn(true); + f.context.setExecutor(null); + f.statement.close(); + + Assertions.assertFalse(f.assignment.isStop(), "DoGet still needs the deferred batch source"); + Assertions.assertSame(source, f.manager.getSplitSource(source.getUniqueId())); + Assertions.assertFalse(f.producer().isCancelled()); + Mockito.verify(f.lease, Mockito.never()).close(); + + f.node.stop(); + + Assertions.assertTrue(f.assignment.isStop()); + Assertions.assertNull(f.manager.getSplitSource(source.getUniqueId())); + Assertions.assertTrue(f.producer().isCancelled()); + Mockito.verify(f.lease).close(); + } + } + + @Test + void cancelledBatchAdmissionDoesNotAcquireFsView() throws Exception { + try (BatchScanFixture f = new BatchScanFixture()) { + f.rejectAdmission.set(true); + Assertions.assertThrows(CancellationException.class, () -> f.node.startSplit(1)); + Mockito.verify(f.generation, Mockito.never()).acquireFsView(Mockito.any()); + Mockito.verify(f.lease, Mockito.never()).close(); + } + } + + @Test + void partitionAdmissionIsHeldUntilQueuedDeliveryTerminates() throws Exception { + try (BatchScanFixture f = new BatchScanFixture()) { + Semaphore admission = new Semaphore(1); + HivePartition partition = Mockito.mock(HivePartition.class); + Mockito.when(partition.isDummyPartition()).thenReturn(true); + f.variables.setForceJniScanner(false); + f.variables.setEnableRuntimeFilterPartitionPrune(false); + org.apache.hudi.common.table.view.HoodieTableFileSystemView view = + Mockito.mock(org.apache.hudi.common.table.view.HoodieTableFileSystemView.class); + Mockito.when(view.getLatestBaseFilesBeforeOrOn(Mockito.any(), Mockito.any())) + .thenAnswer(call -> Stream.empty()); + Mockito.when(f.lease.get()).thenReturn(view); + BatchScanFixture.setField(f.node, HudiScanNode.class, "isCowTable", true); + BatchScanFixture.setField(f.node, HudiScanNode.class, "splittersOnFlight", admission); + BatchScanFixture.setField(f.node, HudiScanNode.class, "prunedPartitions", + Collections.singletonList(partition)); + f.node.startSplit(1); + f.producer().run(); + Assertions.assertEquals(2, f.submitted.size()); + f.submitted.get(1).run(); + Assertions.assertEquals(3, f.submitted.size()); + Assertions.assertEquals(0, admission.availablePermits()); + Mockito.verify(f.lease, Mockito.never()).close(); + f.node.stop(); + Assertions.assertEquals(1, admission.availablePermits()); + Mockito.verify(f.lease).close(); + } + } + + @Test + void ordinaryBatchStatementCloseStillCancelsQueuedPlanning() throws Exception { + try (BatchScanFixture f = new BatchScanFixture()) { + f.node.startSplit(1); + f.statement.close(); + + Assertions.assertTrue(f.assignment.isStop()); + Assertions.assertTrue(f.producer().isCancelled()); + Mockito.verify(f.lease).close(); + } + } + + @Test + void abandonedBatchWithoutExecutorStillStopsAtStatementClose() throws Exception { + try (BatchScanFixture f = new BatchScanFixture()) { + f.context.setExecutor(null); + f.node.startSplit(1); + f.statement.close(); + + Assertions.assertTrue(f.assignment.isStop()); + Assertions.assertTrue(f.producer().isCancelled()); + Mockito.verify(f.lease).close(); + } + } + + private static class BatchScanFixture implements AutoCloseable { + private final ConnectContext previous = ConnectContext.get(); + private final ConnectContext context = new ConnectContext(); + private final SessionVariable variables = new SessionVariable(); + private final StatementContext statement = new StatementContext(context, null); + private final StmtExecutor executor = Mockito.mock(StmtExecutor.class); + private final SplitAssignment assignment = new SplitAssignment( + null, null, null, Collections.emptyMap(), Collections.emptyList(), false); + private final SplitSourceManager manager = new SplitSourceManager(); + private final HudiFsViewCacheValue.Lease lease = Mockito.mock(HudiFsViewCacheValue.Lease.class); + private final AtomicBoolean rejectAdmission = new AtomicBoolean(); + private final HudiExternalMetaCache.FsViewGeneration generation = + Mockito.mock(HudiExternalMetaCache.FsViewGeneration.class); + private final List submitted = new ArrayList<>(); + private final MockedStatic mockedEnv; + private final HudiScanNode node; + + BatchScanFixture() throws Exception { + context.setStatementContext(statement); + context.setExecutor(executor); + context.setThreadLocalInfo(); + Env env = Mockito.mock(Env.class); + ExternalMetaCacheMgr cache = Mockito.mock(ExternalMetaCacheMgr.class); + ExecutorService worker = Mockito.mock(ExecutorService.class); + Mockito.doAnswer(call -> { + submitted.add(call.getArgument(0)); + return null; + }).when(worker).execute(Mockito.any(Runnable.class)); + Mockito.when(env.getExtMetaCacheMgr()).thenReturn(cache); + Mockito.when(env.getSplitSourceManager()).thenReturn(manager); + Mockito.when(cache.getScheduleExecutor()).thenReturn(worker); + Mockito.when(cache.getFileListingExecutor()).thenReturn(worker); + mockedEnv = Mockito.mockStatic(Env.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + HMSExternalTable table = Mockito.mock(HMSExternalTable.class); + HMSExternalCatalog catalog = Mockito.mock(HMSExternalCatalog.class); + HMSExternalCatalog.HudiScanRuntimeContext runtime = + Mockito.mock(HMSExternalCatalog.HudiScanRuntimeContext.class); + Mockito.when(table.getCatalog()).thenReturn(catalog); + Mockito.when(catalog.getHudiScanRuntimeContext()).thenReturn(runtime); + Mockito.when(runtime.getAuthenticator()).thenReturn(new ExecutionAuthenticator() {}); + Mockito.when(runtime.getFsViewGeneration()).thenReturn(generation); + Mockito.when(generation.acquireFsView(Mockito.any())).thenReturn(lease); + TupleDescriptor desc = new TupleDescriptor(new TupleId(0)); + desc.setTable(table); + node = new HudiScanNode(new PlanNodeId(0), desc, false, Optional.empty(), Optional.empty(), + variables, null, ScanContext.EMPTY) { + @Override + protected Executor getSplitTaskExecutor() { + if (rejectAdmission.get()) { + throw new CancellationException("Scan task lease is closed"); + } + return super.getSplitTaskExecutor(); + } + }; + setField(node, HudiScanNode.class, "prunedPartitions", + Collections.singletonList(Mockito.mock(HivePartition.class))); + setField(node, ScanNode.class, "splitAssignment", assignment); + } + + private static void setField(Object target, Class owner, String name, Object value) throws Exception { + Field field = owner.getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } + + HudiScanNode.TerminalTask producer() { + return (HudiScanNode.TerminalTask) submitted.get(0); + } + + @Override + public void close() { + try { + node.stop(); + statement.close(); + } finally { + mockedEnv.close(); + ConnectContext.remove(); + if (previous != null) { + previous.setThreadLocalInfo(); + } + } + } + } + @Test void statementCloseReturnsWhileRunningTaskKeepsLeasePinned() throws Exception { SplitAssignment assignment = Mockito.mock(SplitAssignment.class); @@ -392,6 +593,41 @@ void synchronousListingCloseDuringImmediateSubmissionStopsTask() { Assertions.assertThrows(CancellationException.class, owner::awaitCompletion); } + @Test + void stoppedAssignmentRejectsBatchSyncBeforeThreadPublication() { + SplitAssignment assignment = new SplitAssignment(null, null, null, + Collections.emptyMap(), Collections.emptyList(), true); + HudiFsViewCacheValue.Lease lease = Mockito.mock(HudiFsViewCacheValue.Lease.class); + HudiScanNode.BatchFsViewOwner owner = new HudiScanNode.BatchFsViewOwner( + assignment, lease, Runnable::run, Runnable::run); + assignment.stop(); + // Mirrors startSplit adding its owner after statement cleanup stopped the assignment. + assignment.addCloseable(owner); + try { + Assertions.assertThrows(CancellationException.class, () -> owner.beginSync(Thread.currentThread())); + Mockito.verify(lease, Mockito.never()).close(); + } finally { + owner.endSync(); + owner.finish(); + } + Mockito.verify(lease).close(); + } + + @Test + void closedListingOwnerRejectsSyncBeforeThreadPublication() { + HudiFsViewCacheValue.Lease lease = Mockito.mock(HudiFsViewCacheValue.Lease.class); + HudiScanNode.ListingFsViewOwner owner = new HudiScanNode.ListingFsViewOwner(lease, Runnable::run); + owner.close(); + try { + Assertions.assertThrows(CancellationException.class, () -> owner.beginSync(Thread.currentThread())); + Mockito.verify(lease, Mockito.never()).close(); + } finally { + owner.endSync(); + owner.discardBeforeSubmission(); + } + Mockito.verify(lease).close(); + } + @Test void closeInterruptsBlockedBatchSyncAndKeepsLeaseUntilTerminal() throws Exception { SplitAssignment assignment = Mockito.mock(SplitAssignment.class); diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/cache/IcebergPartitionValueAdapterTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/cache/IcebergPartitionValueAdapterTest.java new file mode 100644 index 00000000000000..f3bd1eec953e57 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/cache/IcebergPartitionValueAdapterTest.java @@ -0,0 +1,178 @@ +// 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.iceberg.cache; + +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())); + } + + @Test + public void dateAndNanosPartitionConversionsAreExact() { + for (int day : new int[] {-19724, -1, 0, 19724}) { + long nanos = day * 86_400_000_000_000L; + Assertions.assertEquals(nanos, + convert(Types.DateType.get(), day, Types.TimestampNanoType.withoutZone())); + Assertions.assertEquals(day, + convert(Types.TimestampNanoType.withoutZone(), nanos, Types.DateType.get())); + assertUnknownPartitionKeysEqual(Types.DateType.get(), day, + Types.TimestampNanoType.withoutZone(), nanos); + } + Assertions.assertThrows(ArithmeticException.class, + () -> convert(Types.DateType.get(), Integer.MAX_VALUE, Types.TimestampNanoType.withoutZone())); + Assertions.assertThrows(IllegalArgumentException.class, + () -> convert(Types.TimestampNanoType.withoutZone(), -1L, Types.DateType.get())); + Assertions.assertNull(convert(Types.DateType.get(), null, Types.TimestampNanoType.withoutZone())); + } + + @Test + public void dateAndMicrosPartitionConversionsAreExact() { + for (int day : new int[] {-19724, -1, 0, 19724}) { + long micros = day * 86_400_000_000L; + Assertions.assertEquals(micros, + convert(Types.DateType.get(), day, Types.TimestampType.withoutZone())); + Assertions.assertEquals(day, + convert(Types.TimestampType.withoutZone(), micros, Types.DateType.get())); + assertUnknownPartitionKeysEqual(Types.DateType.get(), day, Types.TimestampType.withoutZone(), micros); + assertUnknownPartitionKeysEqual(Types.TimestampType.withoutZone(), micros, + Types.TimestampNanoType.withoutZone(), micros * 1000L); + } + Assertions.assertThrows(ArithmeticException.class, + () -> convert(Types.DateType.get(), Integer.MAX_VALUE, Types.TimestampType.withoutZone())); + Assertions.assertThrows(IllegalArgumentException.class, + () -> convert(Types.TimestampType.withoutZone(), -1L, Types.DateType.get())); + Assertions.assertNull(convert(Types.DateType.get(), null, Types.TimestampType.withoutZone())); + Assertions.assertThrows(IllegalArgumentException.class, + () -> convert(Types.DateType.get(), 1, Types.TimestampType.withZone())); + Assertions.assertThrows(IllegalArgumentException.class, + () -> convert(Types.DateType.get(), 1, Types.TimestampNanoType.withZone())); + } + + 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-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java index 5408d4869d280a..8a3a1d71d6d6c5 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/IcebergScanNodeTest.java @@ -55,6 +55,7 @@ import org.apache.doris.datasource.SplitToScanRange; import org.apache.doris.datasource.TableFormatType; import org.apache.doris.datasource.iceberg.IcebergExternalCatalog; +import org.apache.doris.datasource.iceberg.IcebergExternalMetaCache; import org.apache.doris.datasource.iceberg.IcebergExternalTable; import org.apache.doris.datasource.iceberg.IcebergMvccSnapshot; import org.apache.doris.datasource.iceberg.IcebergPartitionInfo; @@ -108,6 +109,7 @@ import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileMetadata; import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.HasTableOperations; import org.apache.iceberg.ManifestFile; import org.apache.iceberg.ManifestFiles; import org.apache.iceberg.MetadataColumns; @@ -161,6 +163,7 @@ import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.Callable; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.RejectedExecutionException; @@ -371,6 +374,12 @@ boolean isRealBatchMode() { return super.isBatchMode(); } + @Override + protected Executor getManifestReadExecutor() { + // This fixture mocks static manifest IO on the calling thread; concurrency has separate tests. + return Runnable::run; + } + @Override public boolean isBatchMode() { return batchMode; @@ -447,6 +456,13 @@ private static class ManifestPlanningIcebergScanNode extends TestIcebergScanNode this.manifestLoadCount = manifestLoadCount; } + @Override + protected CloseableIterable streamFileScanTasksWithManifestCache( + TableScan scan, Snapshot snapshot) { + manifestLoadCount.incrementAndGet(); + return CloseableIterable.withNoopClose(Collections.emptyList()); + } + @Override protected List loadFileScanTasksWithManifestCache( TableScan scan, Snapshot snapshot) { @@ -1169,7 +1185,7 @@ public void testManifestPlanningBypassesStatementCacheForStreamingModes() throws setIcebergSource(batchNode, mockIcebergSource(10L, 20L)); try (CloseableIterable ignored = planFileScanTaskWithManifestCache( batchNode, mockTableScanWithPlanCounter(batchPlanCalls))) { - // Batch mode keeps the SDK iterable lazy instead of materializing manifest tasks. + // Batch mode uses the streaming manifest path without the serialized task cache. } SessionVariable explicitSizeVariable = new SessionVariable(); @@ -1182,8 +1198,8 @@ explicitSizeNode, mockTableScanWithPlanCounter(explicitSizePlanCalls))) { // Explicit split size also keeps the SDK iterable lazy. } - Assert.assertEquals(0, manifestLoadCount.get()); - Assert.assertEquals(1, batchPlanCalls.get()); + Assert.assertEquals(1, manifestLoadCount.get()); + Assert.assertEquals(0, batchPlanCalls.get()); Assert.assertEquals(1, explicitSizePlanCalls.get()); } @@ -3860,6 +3876,12 @@ public void testManifestCachePlanningResolvesDroppedEqualityDeleteField() throws setPrivateField(node, "icebergTable", table); Env env = Mockito.mock(Env.class); Mockito.when(env.getExtMetaCacheMgr()).thenReturn(Mockito.mock(ExternalMetaCacheMgr.class)); + setPreExecutionAuthenticator(node, new ExecutionAuthenticator() { }); + IcebergSnapshotCacheValue planningGeneration = new IcebergSnapshotCacheValue( + new IcebergPartitionInfo(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap()), + new IcebergSnapshot(table.currentSnapshot().snapshotId(), + table.schema().schemaId())).bindSourceGeneration(new IcebergTableCacheValue(table)); + setPrivateField(node, "frozenGenerationSource", planningGeneration); try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class); MockedStatic loader = Mockito.mockStatic(IcebergManifestCacheLoader.class)) { @@ -3926,51 +3948,101 @@ public void testHistoricalDateDeleteBoundsAfterDroppingPromotedKey() throws Exce checkDateToNanosEqualityDeletes(true, true, true); } + @Test + public void testStreamingManifestCacheKeepsDateIdentityPartitionAfterPromotion() throws Exception { + checkDateToNanosEqualityDeletes(false, true, false, true); + } + + @Test + public void testStreamingManifestCacheKeepsDroppedDateIdentityPartitionAfterPromotion() throws Exception { + checkDateToNanosEqualityDeletes(true, true, false, true); + } + private void checkDateToNanosEqualityDeletes(boolean dropped, boolean manifestCache, boolean dateDelete) throws Exception { + checkDateToNanosEqualityDeletes(dropped, manifestCache, dateDelete, false); + } + + @Test + public void testStreamingManifestCacheKeepsDateIdentityPartitionAfterMicrosPromotion() throws Exception { + checkDateToTimestampEqualityDeletes(false, true, false, true, Types.TimestampType.withoutZone()); + } + + @Test + public void testStreamingManifestCacheKeepsDroppedDateIdentityPartitionAfterMicrosPromotion() throws Exception { + checkDateToTimestampEqualityDeletes(true, true, false, true, Types.TimestampType.withoutZone()); + } + + private void checkDateToNanosEqualityDeletes(boolean dropped, boolean manifestCache, + boolean dateDelete, boolean partitioned) throws Exception { + checkDateToTimestampEqualityDeletes(dropped, manifestCache, dateDelete, partitioned, + Types.TimestampNanoType.withoutZone()); + } + + private void checkDateToTimestampEqualityDeletes(boolean dropped, boolean manifestCache, + boolean dateDelete, boolean partitioned, org.apache.iceberg.types.Type.PrimitiveType timestampType) + throws Exception { Schema schema = new Schema(Types.NestedField.required(1, "ts", Types.DateType.get()), Types.NestedField.optional(2, "id", Types.IntegerType.get())); + long timestampValue = timestampType.typeId() == org.apache.iceberg.types.Type.TypeID.TIMESTAMP_NANO + ? 1704153600000000000L : 1704153600000000L; String tableLocation = temporaryFolder.getRoot().toPath() - .resolve("manifest_cache_date_to_nanos_" + manifestCache + "_" + dateDelete).toUri().toString(); + .resolve("manifest_cache_date_to_timestamp_" + manifestCache + "_" + dateDelete + + "_" + partitioned + "_" + timestampType.typeId()).toUri().toString(); Table table = new HadoopTables(new Configuration()).create( - schema, PartitionSpec.unpartitioned(), SortOrder.unsorted(), + schema, partitioned ? PartitionSpec.builderFor(schema).identity("ts").build() + : PartitionSpec.unpartitioned(), SortOrder.unsorted(), ImmutableMap.of(TableProperties.FORMAT_VERSION, "3"), tableLocation); Map dateBounds = ImmutableMap.of(1, org.apache.iceberg.types.Conversions.toByteBuffer(Types.DateType.get(), 19724)); - Map nanosBounds = ImmutableMap.of(1, + Map timestampBounds = ImmutableMap.of(1, org.apache.iceberg.types.Conversions.toByteBuffer( - Types.TimestampNanoType.withoutZone(), 1704153600000000000L)); + timestampType, timestampValue)); + PartitionData datePartition = new PartitionData(table.spec().rawPartitionType()); + if (partitioned) { + datePartition.set(0, 19724); + } DataFile dataFile = DataFiles.builder(table.spec()) + .withPartition(datePartition) .withPath(tableLocation + "/data/old.parquet") .withFormat(FileFormat.PARQUET) .withFileSizeInBytes(10) .withMetrics(new org.apache.iceberg.Metrics(1L, null, ImmutableMap.of(1, 1L), ImmutableMap.of(1, 0L), null, - dateDelete ? nanosBounds : dateBounds, dateDelete ? nanosBounds : dateBounds)) + dateDelete ? timestampBounds : dateBounds, dateDelete ? timestampBounds : dateBounds)) .build(); if (!dateDelete) { table.newFastAppend().appendFile(dataFile).commit(); } - TableOperations operations = ((org.apache.iceberg.HasTableOperations) table).operations(); + TableOperations operations = ((HasTableOperations) table).operations(); TableMetadata before = operations.current(); operations.commit(before, TableMetadata.buildFrom(before).setCurrentSchema(new Schema( - Types.NestedField.required(1, "ts", Types.TimestampNanoType.withoutZone()), + Types.NestedField.required(1, "ts", timestampType), Types.NestedField.optional(2, "id", Types.IntegerType.get())), 2).build()); table.refresh(); if (dateDelete) { table.newFastAppend().appendFile(dataFile).commit(); } + PartitionData timestampPartition = new PartitionData(table.spec().rawPartitionType()); + if (partitioned) { + timestampPartition.set(0, timestampValue); + } + // Partitioned cases isolate partition-key matching; unpartitioned cases cover promoted equality-key bounds. table.newRowDelta().addDeletes(FileMetadata.deleteFileBuilder(table.spec()) - .ofEqualityDeletes(1) + .withPartition(timestampPartition) + .ofEqualityDeletes(partitioned ? 2 : 1) .withPath(tableLocation + "/data/eq-delete.parquet") .withFormat(FileFormat.PARQUET) .withFileSizeInBytes(10) .withMetrics(new org.apache.iceberg.Metrics(1L, null, ImmutableMap.of(1, 1L), ImmutableMap.of(1, 0L), null, - dateDelete ? dateBounds : nanosBounds, dateDelete ? dateBounds : nanosBounds)) + dateDelete ? dateBounds : timestampBounds, dateDelete ? dateBounds : timestampBounds)) .build()).commit(); if (dropped) { + if (partitioned) { + table.updateSpec().removeField("ts").commit(); + } table.updateSchema().deleteColumn("ts").commit(); } @@ -3982,6 +4054,12 @@ private void checkDateToNanosEqualityDeletes(boolean dropped, boolean manifestCa setPrivateField(node, "icebergTable", table); Env env = Mockito.mock(Env.class); Mockito.when(env.getExtMetaCacheMgr()).thenReturn(Mockito.mock(ExternalMetaCacheMgr.class)); + setPreExecutionAuthenticator(node, new ExecutionAuthenticator() { }); + IcebergSnapshotCacheValue planningGeneration = new IcebergSnapshotCacheValue( + new IcebergPartitionInfo(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap()), + new IcebergSnapshot(table.currentSnapshot().snapshotId(), + table.schema().schemaId())).bindSourceGeneration(new IcebergTableCacheValue(table)); + setPrivateField(node, "frozenGenerationSource", planningGeneration); try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class); MockedStatic loader = Mockito.mockStatic(IcebergManifestCacheLoader.class)) { @@ -3989,29 +4067,41 @@ private void checkDateToNanosEqualityDeletes(boolean dropped, boolean manifestCa loader.when(() -> IcebergManifestCacheLoader.loadDeleteFilesWithCache(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any())) .thenAnswer(invocation -> { - List files = new ArrayList<>(); - try (CloseableIterable reader = ManifestFiles.readDeleteManifest( - invocation.getArgument(2), table.io(), table.specs())) { - reader.forEach(file -> files.add(file.copy())); + ManifestCacheValue value = readPhysicalManifest( + invocation.getArgument(2), table, true); + if (partitioned) { + PartitionData partition = (PartitionData) value.getDeleteFiles().get(0).partition(); + Assert.assertEquals(timestampType, + partition.getPartitionType().fields().get(0).type()); } - return ManifestCacheValue.forDeleteFiles(files); + return value; }); loader.when(() -> IcebergManifestCacheLoader.loadDataFilesWithCache(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any())) .thenAnswer(invocation -> { - List files = new ArrayList<>(); - try (CloseableIterable reader = ManifestFiles.read( - invocation.getArgument(2), table.io(), table.specs())) { - reader.forEach(file -> files.add(file.copy())); + ManifestCacheValue value = readPhysicalManifest( + invocation.getArgument(2), table, false); + if (partitioned) { + PartitionData partition = (PartitionData) value.getDataFiles().get(0).partition(); + Assert.assertEquals(Types.DateType.get(), + partition.getPartitionType().fields().get(0).type()); } - return ManifestCacheValue.forDataFiles(files); + return value; }); TableScan scan = org.apache.iceberg.DorisDataTableScan.wrap(table.newScan().project(table.schema())); Assert.assertEquals(dropped, scan.schema().findField("ts") == null); List tasks; if (manifestCache) { - tasks = node.loadFileScanTasksWithManifestCache(scan, table.currentSnapshot()); + if (partitioned) { + tasks = new ArrayList<>(); + try (CloseableIterable planned = node.streamFileScanTasksWithManifestCache( + scan, table.currentSnapshot())) { + planned.forEach(tasks::add); + } + } else { + tasks = node.loadFileScanTasksWithManifestCache(scan, table.currentSnapshot()); + } } else { tasks = new ArrayList<>(); try (CloseableIterable planned = scan.planFiles()) { @@ -4019,10 +4109,10 @@ private void checkDateToNanosEqualityDeletes(boolean dropped, boolean manifestCa } } - // Historical DATE bounds are days, not nanos; they must not exclude the later delete. + // Historical DATE bounds use epoch days; they must not exclude the later delete. Assert.assertEquals(1, tasks.size()); Assert.assertEquals(1, tasks.get(0).deletes().size()); - if (!dateDelete) { + if (!dateDelete && !partitioned) { Map otherNanos = ImmutableMap.of(1, org.apache.iceberg.types.Conversions.toByteBuffer( Types.TimestampNanoType.withoutZone(), 1704240000000000000L)); @@ -4038,6 +4128,81 @@ private void checkDateToNanosEqualityDeletes(boolean dropped, boolean manifestCa } } + @Test + public void testNonBatchManifestReuseSerializesAdaptedDataAndDeleteFiles() throws Exception { + Schema schema = new Schema(Types.NestedField.optional(1, "p", Types.IntegerType.get()), + Types.NestedField.optional(2, "id", Types.IntegerType.get())); + PartitionSpec spec = PartitionSpec.builderFor(schema).identity("p").build(); + String location = temporaryFolder.newFolder("adapted-task-reuse").toURI().toString(); + Table table = new HadoopTables(new Configuration()).create(schema, spec, SortOrder.unsorted(), + ImmutableMap.of(TableProperties.FORMAT_VERSION, "2"), location); + PartitionData partition = new PartitionData(spec.rawPartitionType()); + partition.set(0, 7); + table.newFastAppend().appendFile(DataFiles.builder(spec) + .withPath(location + "/data.parquet").withPartition(partition) + .withRecordCount(1).withFileSizeInBytes(10).build()).commit(); + table.newRowDelta().addDeletes(FileMetadata.deleteFileBuilder(spec).ofEqualityDeletes(2) + .withPath(location + "/delete.parquet").withFormat(FileFormat.PARQUET).withPartition(partition) + .withRecordCount(1).withFileSizeInBytes(10).build()).commit(); + table.updateSchema().updateColumn("p", Types.LongType.get()).commit(); + table.refresh(); + StatementContext statement = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statement); + context.setThreadLocalInfo(); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getExtMetaCacheMgr()).thenReturn(Mockito.mock(ExternalMetaCacheMgr.class)); + AtomicInteger dataLoads = new AtomicInteger(); + AtomicInteger deleteLoads = new AtomicInteger(); + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class); + MockedStatic loader = + Mockito.mockStatic(IcebergManifestCacheLoader.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + loader.when(() -> IcebergManifestCacheLoader.loadDataFilesWithCache(Mockito.any(), Mockito.any(), + Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any())).thenAnswer(call -> { + dataLoads.incrementAndGet(); + return readPhysicalManifest(call.getArgument(2), table, false); + }); + loader.when(() -> IcebergManifestCacheLoader.loadDeleteFilesWithCache(Mockito.any(), Mockito.any(), + Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any())).thenAnswer(call -> { + deleteLoads.incrementAndGet(); + return readPhysicalManifest(call.getArgument(2), table, true); + }); + for (int relation = 0; relation < 2; relation++) { + TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()); + setIcebergSource(node, mockIcebergSource(10L, 20L)); + setIcebergTable(node, table); + setPreExecutionAuthenticator(node, new ExecutionAuthenticator() {}); + setPrivateField(node, "frozenGenerationSource", new IcebergSnapshotCacheValue( + new IcebergPartitionInfo(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap()), + new IcebergSnapshot(table.currentSnapshot().snapshotId(), + table.schema().schemaId())).bindSourceGeneration(new IcebergTableCacheValue(table))); + TableScan scan = DorisDataTableScan.wrap(table.newScan().project(table.schema())); + try (CloseableIterable planned = planFileScanTaskWithManifestCache(node, scan)) { + List tasks = new ArrayList<>(); + planned.forEach(tasks::add); + Assert.assertEquals(1, tasks.size()); + Assert.assertEquals(7L, tasks.get(0).file().partition().get(0, Long.class).longValue()); + Assert.assertEquals(1, tasks.get(0).deletes().size()); + Assert.assertEquals(7L, tasks.get(0).deletes().get(0).partition().get(0, Long.class).longValue()); + } + } + Assert.assertEquals("equivalent relations reuse the adapted task payload", 1, dataLoads.get()); + Assert.assertEquals(1, deleteLoads.get()); + } finally { + statement.close(); + ConnectContext.remove(); + } + } + + private ManifestCacheValue readPhysicalManifest(ManifestFile manifest, Table table, boolean deletes) + throws Exception { + Method loader = IcebergExternalMetaCache.class.getDeclaredMethod( + deletes ? "loadDeleteFiles" : "loadDataFiles", ManifestFile.class, Table.class, boolean.class); + loader.setAccessible(true); + return (ManifestCacheValue) loader.invoke(Mockito.mock(IcebergExternalMetaCache.class), manifest, table, false); + } + @Test public void testPinnedBranchUsesFrozenSnapshotWithCurrentSchema() throws Exception { Schema snapshotSchema = new Schema(11, ImmutableList.of( @@ -4246,6 +4411,74 @@ public void testPinnedGenerationUsesFrozenSchemaMappingOptions() throws Exceptio Assert.assertFalse(node.getEnableMappingTimestampTz()); } + @Test + public void testMaximumBatchSettingClosesEmptyTasksAndGeneration() throws Exception { + List queued = new ArrayList<>(); + SplitAssignment assignment = new SplitAssignment( + Mockito.mock(FederationBackendPolicy.class), Mockito.mock(SplitGenerator.class), + Mockito.mock(SplitToScanRange.class), Collections.emptyMap(), Collections.emptyList(), + false, Integer.MAX_VALUE, 10000); + TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()) { + { + splitAssignment = assignment; + } + + @Override + protected Executor getSplitTaskExecutor() { + return queued::add; + } + + @Override + CloseableIterable planFileScanTaskWithoutReuse(TableScan scan) { + return CloseableIterable.empty(); + } + }; + node.setTableScan(Mockito.mock(TableScan.class)); + setPreExecutionAuthenticator(node, new ExecutionAuthenticator() {}); + IcebergExternalTable table = Mockito.mock(IcebergExternalTable.class); + IcebergSource source = Mockito.mock(IcebergSource.class); + Mockito.when(source.getTargetTable()).thenReturn(table); + setIcebergSource(node, source); + AtomicInteger closed = new AtomicInteger(); + try (MockedStatic utils = Mockito.mockStatic(IcebergUtils.class)) { + utils.when(() -> IcebergUtils.retainTableGenerationForAsyncPlanning(table, null)) + .thenReturn((Closeable) closed::incrementAndGet); + node.doStartSplit(); + Assert.assertEquals(1, queued.size()); + queued.remove(0).run(); + Assert.assertEquals(1, closed.get()); + Assert.assertFalse(assignment.needMoreSplit()); + assignment.init(); + } + } + + @Test + public void testCancelledAdmissionDoesNotLeakAnIndependentGenerationLease() throws Exception { + AtomicInteger references = new AtomicInteger(1); + TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()) { + @Override + protected Executor getSplitTaskExecutor() { + // Model statement cleanup retiring its original generation reference before admission. + references.decrementAndGet(); + throw new java.util.concurrent.CancellationException("Scan task lease is closed"); + } + }; + node.setTableScan(Mockito.mock(TableScan.class)); + IcebergExternalTable table = Mockito.mock(IcebergExternalTable.class); + IcebergSource source = Mockito.mock(IcebergSource.class); + Mockito.when(source.getTargetTable()).thenReturn(table); + setIcebergSource(node, source); + try (MockedStatic utils = Mockito.mockStatic(IcebergUtils.class, Mockito.CALLS_REAL_METHODS)) { + utils.when(() -> IcebergUtils.retainTableGenerationForAsyncPlanning( + Mockito.same(table), Mockito.isNull())).thenAnswer(call -> { + references.incrementAndGet(); + return (Closeable) () -> references.decrementAndGet(); + }); + Assert.assertThrows(java.util.concurrent.CancellationException.class, node::doStartSplit); + Assert.assertEquals("an independent planning lease must not outlive failed admission", 0, references.get()); + } + } + @Test public void testFrozenGenerationIsCapturedForAsyncPlanning() throws Exception { Table frozenTable = Mockito.mock(Table.class); diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/NativeManifestPrefetcherTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/NativeManifestPrefetcherTest.java new file mode 100644 index 00000000000000..9a72063c8a3171 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/iceberg/source/NativeManifestPrefetcherTest.java @@ -0,0 +1,160 @@ +// 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.iceberg.source; + +import org.apache.doris.common.util.FairTaskExecutor; +import org.apache.doris.datasource.SplitAssignment; +import org.apache.doris.datasource.iceberg.cache.ManifestCacheValue; + +import org.apache.iceberg.ManifestFile; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.mockito.Mockito; + +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +@Timeout(30) +class NativeManifestPrefetcherTest { + private FairTaskExecutor pool(int workers) { + return new FairTaskExecutor(workers, task -> { + Thread thread = new Thread(task); + thread.setDaemon(true); + return thread; + }); + } + + private void await(CountDownLatch latch) { + try { + Assertions.assertTrue(latch.await(10, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + } + + @Test + void boundedWindowReturnsCompletedManifestBeforeBlockedFirstManifest() throws Exception { + ManifestFile first = Mockito.mock(ManifestFile.class); + ManifestFile second = Mockito.mock(ManifestFile.class); + ManifestFile third = Mockito.mock(ManifestFile.class); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch initialStarted = new CountDownLatch(2); + CountDownLatch thirdStarted = new CountDownLatch(1); + AtomicInteger leases = new AtomicInteger(); + try (FairTaskExecutor pool = pool(2); + NativeManifestPrefetcher prefetch = new NativeManifestPrefetcher(Arrays.asList(first, second, third), + pool.newGroup(8), Mockito.mock(SplitAssignment.class), false, 2, + () -> { + leases.incrementAndGet(); + return leases::decrementAndGet; + }, manifest -> { + if (manifest == third) { + thirdStarted.countDown(); + } else { + initialStarted.countDown(); + } + if (manifest == first) { + await(release); + } + return ManifestCacheValue.forDataFiles(Collections.emptyList()); + })) { + await(initialStarted); + Assertions.assertEquals(1L, thirdStarted.getCount()); + Iterator iterator = prefetch.iterator(); + Assertions.assertSame(second, iterator.next().manifest); + await(thirdStarted); + Assertions.assertSame(third, iterator.next().manifest); + release.countDown(); + Assertions.assertSame(first, iterator.next().manifest); + Assertions.assertFalse(iterator.hasNext()); + Assertions.assertEquals(0, leases.get()); + } finally { + release.countDown(); + } + } + + @Test + void cancellationRetiresQueuedLeaseButKeepsRunningLeaseUntilWorkerActuallyExits() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch retired = new CountDownLatch(1); + AtomicInteger leases = new AtomicInteger(); + try (FairTaskExecutor pool = pool(1)) { + NativeManifestPrefetcher prefetch = new NativeManifestPrefetcher( + Arrays.asList(Mockito.mock(ManifestFile.class), Mockito.mock(ManifestFile.class)), + pool.newGroup(8), Mockito.mock(SplitAssignment.class), false, 2, + () -> { + leases.incrementAndGet(); + return () -> { + if (leases.decrementAndGet() == 0) { + retired.countDown(); + } + }; + }, manifest -> { + started.countDown(); + boolean done = false; + while (!done) { + try { + release.await(); + done = true; + } catch (InterruptedException ignored) { + // Simulate an IO operation that does not exit on interruption. + } + } + return ManifestCacheValue.forDataFiles(Collections.emptyList()); + }); + await(started); + prefetch.close(); + Assertions.assertEquals(1, leases.get()); + release.countDown(); + await(retired); + Assertions.assertEquals(0, leases.get()); + } finally { + release.countDown(); + } + } + + @Test + void rejectedAdmissionReleasesGenerationAndLoaderFailureRemainsAnError() throws Exception { + AtomicInteger leases = new AtomicInteger(); + try (FairTaskExecutor pool = pool(1)) { + FairTaskExecutor.Group closed = pool.newGroup(1); + closed.close(); + Assertions.assertThrows(RuntimeException.class, () -> new NativeManifestPrefetcher( + Collections.singletonList(Mockito.mock(ManifestFile.class)), closed, + Mockito.mock(SplitAssignment.class), false, 1, + () -> { + leases.incrementAndGet(); + return leases::decrementAndGet; + }, manifest -> null)); + Assertions.assertEquals(0, leases.get()); + try (NativeManifestPrefetcher prefetch = new NativeManifestPrefetcher( + Collections.singletonList(Mockito.mock(ManifestFile.class)), pool.newGroup(1), + Mockito.mock(SplitAssignment.class), false, 1, () -> () -> { }, manifest -> { + throw new IllegalStateException("read failed"); + })) { + Assertions.assertThrows(RuntimeException.class, () -> prefetch.iterator().next()); + } + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java index b49e70c8c38bf3..dc984606dcba12 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java @@ -377,4 +377,25 @@ public void testHyperscanFallbackPropagatesToBackends() throws Exception { Assertions.assertFalse(variable.toThrift().isEnableHyperscanFallback()); } + + @Test + public void testSplitBatchAndQueueLimitsUseTheirOwnVariableCheckers() throws Exception { + SessionVariable variable = new SessionVariable(); + for (String name : new String[] {SessionVariable.SPLIT_BATCH_SIZE, + SessionVariable.SPLIT_QUEUE_CAPACITY, SessionVariable.SPLIT_INITIAL_BATCH_SIZE}) { + Field field = VariableMgr.getVarContext(name).getField(); + for (int value : new int[] {1, 9}) { + VariableMgr.setVar(variable, new SetVar(SetType.SESSION, name, new IntLiteral(value))); + Assertions.assertEquals(value, field.getInt(variable)); + } + for (int value : new int[] {0, -1}) { + DdlException failure = Assertions.assertThrows(DdlException.class, + () -> VariableMgr.setVar(variable, new SetVar(SetType.SESSION, name, new IntLiteral(value)))); + Assertions.assertTrue(failure.getMessage().contains(name + " value should greater than or equal 1")); + Assertions.assertEquals(9, field.getInt(variable)); + } + } + } + + } diff --git a/regression-test/data/external_table_p0/hive/test_hive_streaming_splits.out b/regression-test/data/external_table_p0/hive/test_hive_streaming_splits.out new file mode 100644 index 00000000000000..cd6e775b10bdc4 --- /dev/null +++ b/regression-test/data/external_table_p0/hive/test_hive_streaming_splits.out @@ -0,0 +1,236 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !empty_unpartitioned -- +0 \N + +-- !empty_partitioned -- +0 \N + +-- !synchronous_unpartitioned_aggregate -- +120 7140 120 660 + +-- !synchronous_unpartitioned_files -- +0 10 10 45 0 9 +1 10 10 145 10 19 +10 10 10 1045 100 109 +11 10 10 1145 110 119 +2 10 10 245 20 29 +3 10 10 345 30 39 +4 10 10 445 40 49 +5 10 10 545 50 59 +6 10 10 645 60 69 +7 10 10 745 70 79 +8 10 10 845 80 89 +9 10 10 945 90 99 + +-- !synchronous_unpartitioned_limit -- +10 + +-- !synchronous_partitioned_aggregate -- +120 7140 120 660 240 + +-- !synchronous_partitioned_files -- +1 0 10 10 45 0 9 +1 3 10 10 345 30 39 +1 6 10 10 645 60 69 +1 9 10 10 945 90 99 +2 1 10 10 145 10 19 +2 10 10 10 1045 100 109 +2 4 10 10 445 40 49 +2 7 10 10 745 70 79 +3 11 10 10 1145 110 119 +3 2 10 10 245 20 29 +3 5 10 10 545 50 59 +3 8 10 10 845 80 89 + +-- !synchronous_partition_range -- +80 5160 80 200 + +-- !synchronous_single_partition_limit -- +10 + +-- !synchronous_empty_partitions -- +0 \N + +-- !synchronous_join_aggregate -- +120 7140 240 + +-- !synchronous_union_aggregate -- +240 14280 + +-- !synchronous_union_limit -- +10 + +-- !streaming_3_1_unpartitioned_aggregate -- +120 7140 120 660 + +-- !streaming_3_1_unpartitioned_files -- +0 10 10 45 0 9 +1 10 10 145 10 19 +10 10 10 1045 100 109 +11 10 10 1145 110 119 +2 10 10 245 20 29 +3 10 10 345 30 39 +4 10 10 445 40 49 +5 10 10 545 50 59 +6 10 10 645 60 69 +7 10 10 745 70 79 +8 10 10 845 80 89 +9 10 10 945 90 99 + +-- !streaming_3_1_unpartitioned_limit -- +10 + +-- !streaming_3_1_partitioned_aggregate -- +120 7140 120 660 240 + +-- !streaming_3_1_partitioned_files -- +1 0 10 10 45 0 9 +1 3 10 10 345 30 39 +1 6 10 10 645 60 69 +1 9 10 10 945 90 99 +2 1 10 10 145 10 19 +2 10 10 10 1045 100 109 +2 4 10 10 445 40 49 +2 7 10 10 745 70 79 +3 11 10 10 1145 110 119 +3 2 10 10 245 20 29 +3 5 10 10 545 50 59 +3 8 10 10 845 80 89 + +-- !streaming_3_1_partition_range -- +80 5160 80 200 + +-- !streaming_3_1_single_partition_limit -- +10 + +-- !streaming_3_1_empty_partitions -- +0 \N + +-- !streaming_3_1_join_aggregate -- +120 7140 240 + +-- !streaming_3_1_union_aggregate -- +240 14280 + +-- !streaming_3_1_union_limit -- +10 + +-- !streaming_7_2_unpartitioned_aggregate -- +120 7140 120 660 + +-- !streaming_7_2_unpartitioned_files -- +0 10 10 45 0 9 +1 10 10 145 10 19 +10 10 10 1045 100 109 +11 10 10 1145 110 119 +2 10 10 245 20 29 +3 10 10 345 30 39 +4 10 10 445 40 49 +5 10 10 545 50 59 +6 10 10 645 60 69 +7 10 10 745 70 79 +8 10 10 845 80 89 +9 10 10 945 90 99 + +-- !streaming_7_2_unpartitioned_limit -- +10 + +-- !streaming_7_2_partitioned_aggregate -- +120 7140 120 660 240 + +-- !streaming_7_2_partitioned_files -- +1 0 10 10 45 0 9 +1 3 10 10 345 30 39 +1 6 10 10 645 60 69 +1 9 10 10 945 90 99 +2 1 10 10 145 10 19 +2 10 10 10 1045 100 109 +2 4 10 10 445 40 49 +2 7 10 10 745 70 79 +3 11 10 10 1145 110 119 +3 2 10 10 245 20 29 +3 5 10 10 545 50 59 +3 8 10 10 845 80 89 + +-- !streaming_7_2_partition_range -- +80 5160 80 200 + +-- !streaming_7_2_single_partition_limit -- +10 + +-- !streaming_7_2_empty_partitions -- +0 \N + +-- !streaming_7_2_join_aggregate -- +120 7140 240 + +-- !streaming_7_2_union_aggregate -- +240 14280 + +-- !streaming_7_2_union_limit -- +10 + +-- !streaming_128_32_unpartitioned_aggregate -- +120 7140 120 660 + +-- !streaming_128_32_unpartitioned_files -- +0 10 10 45 0 9 +1 10 10 145 10 19 +10 10 10 1045 100 109 +11 10 10 1145 110 119 +2 10 10 245 20 29 +3 10 10 345 30 39 +4 10 10 445 40 49 +5 10 10 545 50 59 +6 10 10 645 60 69 +7 10 10 745 70 79 +8 10 10 845 80 89 +9 10 10 945 90 99 + +-- !streaming_128_32_unpartitioned_limit -- +10 + +-- !streaming_128_32_partitioned_aggregate -- +120 7140 120 660 240 + +-- !streaming_128_32_partitioned_files -- +1 0 10 10 45 0 9 +1 3 10 10 345 30 39 +1 6 10 10 645 60 69 +1 9 10 10 945 90 99 +2 1 10 10 145 10 19 +2 10 10 10 1045 100 109 +2 4 10 10 445 40 49 +2 7 10 10 745 70 79 +3 11 10 10 1145 110 119 +3 2 10 10 245 20 29 +3 5 10 10 545 50 59 +3 8 10 10 845 80 89 + +-- !streaming_128_32_partition_range -- +80 5160 80 200 + +-- !streaming_128_32_single_partition_limit -- +10 + +-- !streaming_128_32_empty_partitions -- +0 \N + +-- !streaming_128_32_join_aggregate -- +120 7140 240 + +-- !streaming_128_32_union_aggregate -- +240 14280 + +-- !streaming_128_32_union_limit -- +10 + +-- !ctas -- +120 7140 120 660 + +-- !statement_hint -- +120 7140 1 1 + +-- !settings_after_hint -- +128 32 + diff --git a/regression-test/suites/external_table_p0/hive/test_hive_streaming_splits.groovy b/regression-test/suites/external_table_p0/hive/test_hive_streaming_splits.groovy new file mode 100644 index 00000000000000..4832e91ec55b10 --- /dev/null +++ b/regression-test/suites/external_table_p0/hive/test_hive_streaming_splits.groovy @@ -0,0 +1,144 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +suite("test_hive_streaming_splits", "p0,external") { + if (!"true".equalsIgnoreCase(context.config.otherConfigs.get("enableHiveTest"))) { + return + } + + // An existing catalog also allows this suite to run against an isolated S3-backed HMS fixture. + String catalog = context.config.otherConfigs.get("hiveSplitCatalog") + if (!catalog) { + catalog = "hive_streaming_splits" + String host = context.config.otherConfigs.get("externalEnvIp") + String port = context.config.otherConfigs.get("hive3HmsPort") + sql "drop catalog if exists ${catalog}" + sql """create catalog ${catalog} properties ( + 'type' = 'hms', 'hive.metastore.uris' = 'thrift://${host}:${port}')""" + } + def settings = ["enable_external_table_batch_mode", "num_files_in_batch_mode", + "num_partitions_in_batch_mode", "split_batch_size", "split_queue_capacity"] + def saved = settings.collectEntries { name -> [(name): sql("show variables like '${name}'")[0][1]] } + try { + sql "create database if not exists ${catalog}.hive_streaming_splits" + sql "use ${catalog}.hive_streaming_splits" + // Truncate while the previous partition metadata still exists, so repeated runs clear its files. + if (!sql("show tables like 'streaming_unpartitioned'").isEmpty()) { + sql "truncate table streaming_unpartitioned" + } + if (!sql("show tables like 'streaming_partitioned'").isEmpty()) { + sql "truncate table streaming_partitioned" + } + if (!sql("show tables like 'streaming_ctas'").isEmpty()) { + sql "truncate table streaming_ctas" + } + sql "drop table if exists streaming_unpartitioned" + sql "drop table if exists streaming_partitioned" + sql "drop table if exists streaming_ctas" + sql """create table streaming_unpartitioned (id bigint, file_id int) + engine=hive properties ('file_format'='parquet')""" + sql """create table streaming_partitioned (id bigint, file_id int, p int) + engine=hive partition by list (p) () properties ('file_format'='parquet')""" + sql "set enable_external_table_batch_mode=true" + sql "set num_files_in_batch_mode=1" + sql "set num_partitions_in_batch_mode=1" + qt_empty_unpartitioned "select count(*), sum(id) from streaming_unpartitioned" + qt_empty_partitioned "select count(*), sum(id) from streaming_partitioned" + // Separate inserts create enough files to cross several small producer batches. + for (int file = 0; file < 12; file++) { + sql """insert into streaming_unpartitioned + select number + ${file * 10}, ${file} from numbers('number'='10')""" + sql """insert into streaming_partitioned + select number + ${file * 10}, ${file}, ${file % 3 + 1} from numbers('number'='10')""" + } + // A cold file-count hint selects streaming even when the loaded table is below the threshold. + sql "set num_files_in_batch_mode=1024" + sql "set num_partitions_in_batch_mode=1024" + sql "refresh table streaming_unpartitioned" + explain { + sql "select * from streaming_unpartitioned limit 10" + contains "(approximate)inputSplitNum" + } + sql "set num_files_in_batch_mode=1" + def checkResults = { String label -> + quickTest("${label}_unpartitioned_aggregate", """select count(*), sum(id), count(distinct id), sum(file_id) + from streaming_unpartitioned""", true) + quickTest("${label}_unpartitioned_files", """select file_id, count(*), count(distinct id), + sum(id), min(id), max(id) from streaming_unpartitioned group by file_id order by file_id""", true) + quickTest("${label}_unpartitioned_limit", """select count(*) from + (select id from streaming_unpartitioned limit 10) limited""", true) + quickTest("${label}_partitioned_aggregate", """select count(*), sum(id), count(distinct id), sum(file_id), sum(p) + from streaming_partitioned""", true) + quickTest("${label}_partitioned_files", """select p, file_id, count(*), count(distinct id), + sum(id), min(id), max(id) from streaming_partitioned group by p, file_id order by p, file_id""", true) + quickTest("${label}_partition_range", """select count(*), sum(id), count(distinct id), sum(p) + from streaming_partitioned where p between 2 and 3""", true) + quickTest("${label}_single_partition_limit", """select count(*) from + (select id from streaming_partitioned where p=1 limit 10) limited""", true) + quickTest("${label}_empty_partitions", """select count(*), sum(id) + from streaming_partitioned where p>3""", true) + quickTest("${label}_join_aggregate", """select count(*), sum(u.id), sum(p.p) + from streaming_unpartitioned u join streaming_partitioned p on u.id=p.id""", true) + quickTest("${label}_union_aggregate", """select count(*), sum(id) from + (select id from streaming_unpartitioned union all + select id from streaming_partitioned) combined""", true) + quickTest("${label}_union_limit", """select count(*) from + (select id from streaming_unpartitioned union all + select id from streaming_partitioned limit 10) limited""", true) + } + sql "set enable_external_table_batch_mode=false" + sql "set num_partitions_in_batch_mode=0" + explain { + sql "select * from streaming_partitioned limit 10" + notContains "(approximate)inputSplitNum" + } + checkResults("synchronous") + sql "set enable_external_table_batch_mode=true" + sql "set num_partitions_in_batch_mode=1" + for (def limits : [[3, 1], [7, 2], [128, 32]]) { + sql "set split_batch_size=${limits[0]}" + sql "set split_queue_capacity=${limits[1]}" + explain { + sql "select * from streaming_unpartitioned limit 10" + contains "(approximate)inputSplitNum" + } + explain { + sql "select * from streaming_partitioned where p=1 limit 10" + contains "(approximate)inputSplitNum" + } + checkResults("streaming_${limits[0]}_${limits[1]}") + } + // CTAS probes an absent target before creating and resolving it in the same statement. + sql """create table streaming_ctas properties ('file_format'='parquet') + as select id, file_id from streaming_unpartitioned""" + qt_ctas "select count(*), sum(id), count(distinct id), sum(file_id) from streaming_ctas" + qt_statement_hint """select /*+ SET_VAR(split_batch_size=1, split_queue_capacity=1) */ + count(*), sum(id), max(@@split_batch_size), max(@@split_queue_capacity) + from streaming_unpartitioned""" + qt_settings_after_hint "select @@split_batch_size, @@split_queue_capacity" + test { + sql "set split_batch_size=0" + exception "split_batch_size value should greater than or equal 1" + } + test { + sql "set split_queue_capacity=-1" + exception "split_queue_capacity value should greater than or equal 1" + } + } finally { + saved.each { name, value -> sql "set ${name}=${value}" } + } +}