diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalScanTaskCacheKey.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalScanTaskCacheKey.java new file mode 100644 index 00000000000000..f08363045d3dea --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/ExternalScanTaskCacheKey.java @@ -0,0 +1,29 @@ +// 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; + +/** + * A statement-scoped cache key for connector-native scan tasks. + * + *

Implementations must include every scan property that can change the planned task list, + * such as the table generation, snapshot, projected schema, predicates, and connector options. + * The type parameter ties a connector key to its native task type without exposing that type + * in StatementContext. + */ +public interface ExternalScanTaskCacheKey { +} 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 a7a0e7b381a6b0..f45ac1634689e5 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 @@ -41,6 +41,7 @@ import org.apache.doris.datasource.mvcc.MvccSnapshot; import org.apache.doris.datasource.mvcc.MvccTable; import org.apache.doris.datasource.mvcc.MvccUtil; +import org.apache.doris.nereids.StatementContext; import org.apache.doris.planner.PlanNodeId; import org.apache.doris.planner.ScanContext; import org.apache.doris.qe.ConnectContext; @@ -85,6 +86,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.Callable; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; @@ -97,6 +99,7 @@ public abstract class FileQueryScanNode extends FileScanNode { protected Map destSlotDescByName; protected TFileScanRangeParams params; + private final StatementContext.ExternalScanTaskCache externalScanTaskCache; @Getter protected TableSample tableSample; @@ -136,6 +139,18 @@ public FileQueryScanNode(PlanNodeId id, TupleDescriptor desc, String planNodeNam StatisticalType statisticalType, ScanContext scanContext, boolean needCheckColumnPriv, SessionVariable sv) { super(id, desc, planNodeName, statisticalType, scanContext, needCheckColumnPriv); this.sessionVariable = sv; + ConnectContext context = ConnectContext.get(); + StatementContext statementContext = context == null ? null : context.getStatementContext(); + this.externalScanTaskCache = statementContext == null + ? null : statementContext.getExternalScanTaskCache(); + } + + protected List getOrLoadExternalScanTasks( + ExternalScanTaskCacheKey key, Callable> loader) throws Exception { + if (externalScanTaskCache == null) { + return loader.call(); + } + return externalScanTaskCache.getOrLoad(key, loader); } /** 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 f5acd7834b6402..93ef36811070b4 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 @@ -29,6 +29,7 @@ import org.apache.doris.common.UserException; import org.apache.doris.common.util.DebugUtil; import org.apache.doris.common.util.Util; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.FileQueryScanNode; import org.apache.doris.datasource.FileSplit; import org.apache.doris.datasource.FileSplitter; @@ -78,6 +79,7 @@ import java.util.Collections; 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.Executor; @@ -320,8 +322,17 @@ private void getFileSplitByPartitions(HiveExternalMetaCache cache, List 0; - fileCaches = cache.getFilesByPartitions(partitions, withCache, partitions.size() > 1, - directoryLister, hmsTable); + HiveFileScanTaskCacheKey cacheKey = new HiveFileScanTaskCacheKey( + hmsTable.getCatalog().getId(), hmsTable.getId(), partitions); + try { + fileCaches = getOrLoadExternalScanTasks(cacheKey, + () -> cache.getFilesByPartitions(partitions, withCache, partitions.size() > 1, + directoryLister, hmsTable)); + } catch (IOException | UserException e) { + throw e; + } catch (Exception e) { + throw new IOException("Failed to list Hive files", e); + } } if (!isBatchMode && getSummaryProfile() != null) { getSummaryProfile().addExternalTableGetPartitionFilesTime(System.currentTimeMillis() - startTime); @@ -428,11 +439,10 @@ private List selectFiles(List selectFiles(List { + private final long catalogId; + private final long tableId; + private final List partitions; + + private HiveFileScanTaskCacheKey(long catalogId, long tableId, List partitions) { + this.catalogId = catalogId; + this.tableId = tableId; + this.partitions = partitions.stream() + .map(HivePartitionCacheKey::new) + .collect(Collectors.toList()); + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof HiveFileScanTaskCacheKey)) { + return false; + } + HiveFileScanTaskCacheKey that = (HiveFileScanTaskCacheKey) object; + return catalogId == that.catalogId + && tableId == that.tableId + && partitions.equals(that.partitions); + } + + @Override + public int hashCode() { + return Objects.hash(catalogId, tableId, partitions); + } + } + + private static final class HivePartitionCacheKey { + private final String inputFormat; + private final String path; + private final List partitionValues; + + private HivePartitionCacheKey(HivePartition partition) { + this.inputFormat = partition.getInputFormat(); + this.path = partition.getPath(); + this.partitionValues = partition.getPartitionValues() == null + ? null : Collections.unmodifiableList(new ArrayList<>(partition.getPartitionValues())); + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof HivePartitionCacheKey)) { + return false; + } + HivePartitionCacheKey that = (HivePartitionCacheKey) object; + return Objects.equals(inputFormat, that.inputFormat) + && Objects.equals(path, that.path) + && Objects.equals(partitionValues, that.partitionValues); + } + + @Override + public int hashCode() { + return Objects.hash(inputFormat, path, partitionValues); + } + } + private List getFileSplitByTransaction(HiveExternalMetaCache cache, List partitions, String bindBrokerName) { for (HivePartition partition : partitions) { 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 6c6888e3792d69..286d09d84fd0b9 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 @@ -30,6 +30,7 @@ import org.apache.doris.common.util.BrokerUtil; import org.apache.doris.common.util.FileFormatUtils; import org.apache.doris.common.util.LocationPath; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.ExternalTable; import org.apache.doris.datasource.ExternalUtil; import org.apache.doris.datasource.NameMapping; @@ -52,6 +53,8 @@ import org.apache.doris.thrift.THudiFileDesc; import org.apache.doris.thrift.TTableFormatFileDesc; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import org.apache.commons.lang3.exception.ExceptionUtils; @@ -79,6 +82,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; @@ -380,29 +384,68 @@ private List getPrunedPartitions(HoodieTableMetaClient metaClient private List getIncrementalSplits() { long startTime = System.currentTimeMillis(); - if (canUseNativeReader()) { - List splits = incrementalRelation.collectSplits(); - noLogsSplitNum.addAndGet(splits.size()); + try { + HudiIncrementalScanTaskCacheKey cacheKey = new HudiIncrementalScanTaskCacheKey( + hmsTable.getCatalog().getId(), hmsTable.getId(), + incrementalRelation.getStartTs(), incrementalRelation.getEndTs(), + canUseNativeReader(), incrementalRelation.getHoodieParams()); + List plannedSplits = getOrLoadExternalScanTasks(cacheKey, () -> { + if (canUseNativeReader()) { + return incrementalRelation.collectSplits().stream() + .map(split -> { + Preconditions.checkState(split instanceof HudiSplit, + "Hudi COW incremental relation must produce HudiSplit"); + return (HudiSplit) split; + }) + .collect(Collectors.toList()); + } + Option partitionColumns = hudiClient.getTableConfig().getPartitionFields(); + List partitionNames = partitionColumns.isPresent() + ? Arrays.asList(partitionColumns.get()) : Collections.emptyList(); + return incrementalRelation.collectFileSlices().stream() + .map(fileSlice -> generateHudiSplit(fileSlice, + HudiPartitionUtils.parsePartitionValues( + partitionNames, fileSlice.getPartitionPath()), + incrementalRelation.getEndTs())) + .collect(Collectors.toList()); + }); + List splits = plannedSplits.stream() + .map(HudiScanNode::copyHudiSplit) + .collect(Collectors.toList()); + for (HudiSplit split : plannedSplits) { + if (canUseNativeReader() + || (!sessionVariable.isForceJniScanner() && split.getHudiDeltaLogs().isEmpty())) { + noLogsSplitNum.incrementAndGet(); + } + } + return splits; + } catch (Exception e) { + throw new RuntimeException("Failed to plan Hudi incremental scan tasks", e); + } finally { if (getSummaryProfile() != null) { getSummaryProfile().addExternalTableGetFileScanTasksTime(System.currentTimeMillis() - startTime); } - return splits; } - Option partitionColumns = hudiClient.getTableConfig().getPartitionFields(); - List partitionNames = partitionColumns.isPresent() ? Arrays.asList(partitionColumns.get()) - : Collections.emptyList(); - List splits = incrementalRelation.collectFileSlices().stream() - .map(fileSlice -> generateHudiSplit(fileSlice, - HudiPartitionUtils.parsePartitionValues(partitionNames, fileSlice.getPartitionPath()), - incrementalRelation.getEndTs())) - .collect(Collectors.toList()); - if (getSummaryProfile() != null) { - getSummaryProfile().addExternalTableGetFileScanTasksTime(System.currentTimeMillis() - startTime); + } + + private void getPartitionSplits(HivePartition partition, List splits) throws Exception { + HudiFileScanTaskCacheKey cacheKey = new HudiFileScanTaskCacheKey( + hmsTable.getCatalog().getId(), hmsTable.getId(), queryInstant, + canUseNativeReader(), sessionVariable.isEnableRuntimeFilterPartitionPrune(), partition); + List plannedSplits = getOrLoadExternalScanTasks( + cacheKey, () -> planPartitionSplits(partition)); + for (HudiSplit plannedSplit : plannedSplits) { + HudiSplit split = copyHudiSplit(plannedSplit); + if (canUseNativeReader() + || (!sessionVariable.isForceJniScanner() && split.getHudiDeltaLogs().isEmpty())) { + noLogsSplitNum.incrementAndGet(); + } + splits.add(split); } - return splits; } - private void getPartitionSplits(HivePartition partition, List splits) throws IOException { + private List planPartitionSplits(HivePartition partition) throws IOException { + List splits = new ArrayList<>(); String partitionName; if (partition.isDummyPartition()) { partitionName = ""; @@ -417,7 +460,6 @@ private void getPartitionSplits(HivePartition partition, List splits) thr if (canUseNativeReader()) { fsView.getLatestBaseFilesBeforeOrOn(partitionName, queryInstant).forEach(baseFile -> { - noLogsSplitNum.incrementAndGet(); String filePath = baseFile.getPath(); long fileSize = baseFile.getFileSize(); @@ -436,6 +478,7 @@ private void getPartitionSplits(HivePartition partition, List splits) thr .forEach(fileSlice -> splits.add( generateHudiSplit(fileSlice, partition.getPartitionValues(), queryInstant))); } + return splits; } private void getPartitionsSplits(List partitions, List splits) { @@ -581,10 +624,6 @@ private HudiSplit generateHudiSplit(FileSlice fileSlice, List partitionV List logs = fileSlice.getLogFiles().map(HoodieLogFile::getPath) .map(StoragePath::toString) .collect(Collectors.toList()); - if (logs.isEmpty() && !sessionVariable.isForceJniScanner()) { - noLogsSplitNum.incrementAndGet(); - } - // no base file, use log file to parse file type String agencyPath = filePath.isEmpty() ? logs.get(0) : filePath; LocationPath locationPath = LocationPath.of(agencyPath, hmsTable.getStoragePropertiesMap()); @@ -602,6 +641,133 @@ private HudiSplit generateHudiSplit(FileSlice fileSlice, List partitionV return split; } + @VisibleForTesting + static HudiSplit copyHudiSplit(HudiSplit sourceSplit) { + HudiSplit copy = new HudiSplit( + sourceSplit.getPath(), + sourceSplit.getStart(), + sourceSplit.getLength(), + sourceSplit.getFileLength(), + Arrays.copyOf(sourceSplit.getHosts(), sourceSplit.getHosts().length), + copyList(sourceSplit.getPartitionValues())); + copy.setModificationTime(sourceSplit.getModificationTime()); + copy.setTableFormatType(sourceSplit.getTableFormatType()); + copy.setAlternativeHosts(copyList(sourceSplit.getAlternativeHosts())); + copy.selfSplitWeight = sourceSplit.selfSplitWeight; + copy.setTargetSplitSize(sourceSplit.getTargetSplitSize()); + copy.setInstantTime(sourceSplit.getInstantTime()); + copy.setSerde(sourceSplit.getSerde()); + copy.setInputFormat(sourceSplit.getInputFormat()); + copy.setBasePath(sourceSplit.getBasePath()); + copy.setDataFilePath(sourceSplit.getDataFilePath()); + copy.setHudiDeltaLogs(copyList(sourceSplit.getHudiDeltaLogs())); + copy.setHudiColumnNames(copyList(sourceSplit.getHudiColumnNames())); + copy.setHudiColumnTypes(copyList(sourceSplit.getHudiColumnTypes())); + copy.setNestedFields(copyList(sourceSplit.getNestedFields())); + copy.setHudiPartitionValues(sourceSplit.getHudiPartitionValues() == null + ? null : new HashMap<>(sourceSplit.getHudiPartitionValues())); + return copy; + } + + private static List copyList(List values) { + return values == null ? null : new ArrayList<>(values); + } + + private static final class HudiFileScanTaskCacheKey + implements ExternalScanTaskCacheKey { + private final long catalogId; + private final long tableId; + private final String queryInstant; + private final boolean nativeReader; + private final boolean runtimePartitionPrune; + private final String inputFormat; + private final String path; + private final List partitionValues; + + private HudiFileScanTaskCacheKey( + long catalogId, long tableId, String queryInstant, boolean nativeReader, + boolean runtimePartitionPrune, HivePartition partition) { + this.catalogId = catalogId; + this.tableId = tableId; + this.queryInstant = queryInstant; + this.nativeReader = nativeReader; + this.runtimePartitionPrune = runtimePartitionPrune; + this.inputFormat = partition.getInputFormat(); + this.path = partition.getPath(); + this.partitionValues = partition.getPartitionValues() == null + ? null : Collections.unmodifiableList(new ArrayList<>(partition.getPartitionValues())); + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof HudiFileScanTaskCacheKey)) { + return false; + } + HudiFileScanTaskCacheKey that = (HudiFileScanTaskCacheKey) object; + return catalogId == that.catalogId + && tableId == that.tableId + && nativeReader == that.nativeReader + && runtimePartitionPrune == that.runtimePartitionPrune + && Objects.equals(queryInstant, that.queryInstant) + && Objects.equals(inputFormat, that.inputFormat) + && Objects.equals(path, that.path) + && Objects.equals(partitionValues, that.partitionValues); + } + + @Override + public int hashCode() { + return Objects.hash( + catalogId, tableId, queryInstant, nativeReader, runtimePartitionPrune, + inputFormat, path, partitionValues); + } + } + + private static final class HudiIncrementalScanTaskCacheKey + implements ExternalScanTaskCacheKey { + private final long catalogId; + private final long tableId; + private final String startTs; + private final String endTs; + private final boolean nativeReader; + private final Map hoodieParams; + + private HudiIncrementalScanTaskCacheKey( + long catalogId, long tableId, String startTs, String endTs, + boolean nativeReader, Map hoodieParams) { + this.catalogId = catalogId; + this.tableId = tableId; + this.startTs = startTs; + this.endTs = endTs; + this.nativeReader = nativeReader; + this.hoodieParams = Collections.unmodifiableMap(new HashMap<>(hoodieParams)); + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof HudiIncrementalScanTaskCacheKey)) { + return false; + } + HudiIncrementalScanTaskCacheKey that = (HudiIncrementalScanTaskCacheKey) object; + return catalogId == that.catalogId + && tableId == that.tableId + && nativeReader == that.nativeReader + && Objects.equals(startTs, that.startTs) + && Objects.equals(endTs, that.endTs) + && hoodieParams.equals(that.hoodieParams); + } + + @Override + public int hashCode() { + return Objects.hash(catalogId, tableId, startTs, endTs, nativeReader, hoodieParams); + } + } + @Override public String getNodeExplainString(String prefix, TExplainLevel detailLevel) { if (isBatchMode()) { 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 957ab6ed55e193..0cd4b4837ae322 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 @@ -31,6 +31,7 @@ import org.apache.doris.common.security.authentication.ExecutionAuthenticator; 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; @@ -134,6 +135,7 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; public class IcebergScanNode extends FileQueryScanNode { @@ -879,6 +881,12 @@ private CloseableIterable splitFiles(TableScan scan) { // Non Batch Mode // Materialize planFiles() into a list to avoid iterating the CloseableIterable twice. // RISK: It will cost memory if the table is large. + List fileScanTaskList = getOrPlanFileScanTasks(scan, () -> materializeFileScanTasks(scan)); + targetSplitSize = determineTargetFileSplitSize(fileScanTaskList); + return TableScanUtil.splitFiles(CloseableIterable.withNoopClose(fileScanTaskList), targetSplitSize); + } + + private List materializeFileScanTasks(TableScan scan) { List fileScanTaskList = new ArrayList<>(); try (CloseableIterable scanTasksIter = scan.planFiles()) { for (FileScanTask task : scanTasksIter) { @@ -887,9 +895,28 @@ private CloseableIterable splitFiles(TableScan scan) { } catch (Exception e) { throw new RuntimeException("Failed to materialize file scan tasks", e); } + return fileScanTaskList; + } - targetSplitSize = determineTargetFileSplitSize(fileScanTaskList); - return TableScanUtil.splitFiles(CloseableIterable.withNoopClose(fileScanTaskList), targetSplitSize); + @VisibleForTesting + List getOrPlanFileScanTasks(TableScan scan, Supplier> planner) { + try { + return getOrLoadExternalScanTasks(createFileScanTaskCacheKey(scan), planner::get); + } catch (Exception e) { + throw new RuntimeException("Failed to plan Iceberg file scan tasks", e); + } + } + + private IcebergScanTaskCacheKey createFileScanTaskCacheKey(TableScan scan) { + Snapshot snapshot = scan.snapshot(); + return new IcebergScanTaskCacheKey<>( + source.getCatalog().getId(), + source.getTargetTable().getId(), + snapshot == null ? null : snapshot.snapshotId(), + scan.schema().schemaId(), + scan.filter(), + scan.isCaseSensitive(), + FileScanTask.class.getName()); } private long determineTargetFileSplitSize(Iterable> tasks) { @@ -917,13 +944,76 @@ private long determinePositionDeleteTargetSplitSize(Iterable + implements ExternalScanTaskCacheKey { + private final long catalogId; + private final long tableId; + private final Long snapshotId; + private final int schemaId; + private final byte[] serializedFilter; + private final boolean caseSensitive; + private final String taskType; + + private IcebergScanTaskCacheKey( + long catalogId, long tableId, Long snapshotId, int schemaId, + Expression filter, boolean caseSensitive, String taskType) { + this.catalogId = catalogId; + this.tableId = tableId; + this.snapshotId = snapshotId; + this.schemaId = schemaId; + this.serializedFilter = filter == null ? null : SerializationUtil.serializeToBytes(filter); + this.caseSensitive = caseSensitive; + this.taskType = taskType; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof IcebergScanTaskCacheKey)) { + return false; + } + IcebergScanTaskCacheKey that = (IcebergScanTaskCacheKey) object; + return catalogId == that.catalogId + && tableId == that.tableId + && schemaId == that.schemaId + && caseSensitive == that.caseSensitive + && Objects.equals(snapshotId, that.snapshotId) + && taskType.equals(that.taskType) + && Arrays.equals(serializedFilter, that.serializedFilter); + } + + @Override + public int hashCode() { + return 31 * Objects.hash(catalogId, tableId, snapshotId, schemaId, caseSensitive, taskType) + + Arrays.hashCode(serializedFilter); + } + } + private CloseableIterable planFileScanTaskWithManifestCache(TableScan scan) throws IOException { // Get the snapshot from the scan; return empty if no snapshot exists Snapshot snapshot = scan.snapshot(); if (snapshot == null) { return CloseableIterable.withNoopClose(Collections.emptyList()); } + List tasks; + try { + tasks = getOrLoadExternalScanTasks( + createFileScanTaskCacheKey(scan), + () -> loadFileScanTasksWithManifestCache(scan, snapshot)); + } catch (IOException e) { + throw e; + } catch (Exception e) { + throw new IOException("Failed to plan Iceberg scan tasks with manifest cache", e); + } + targetSplitSize = determineTargetFileSplitSize(tasks); + return TableScanUtil.splitFiles(CloseableIterable.withNoopClose(tasks), targetSplitSize); + } + @VisibleForTesting + protected List loadFileScanTasksWithManifestCache( + 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)) { @@ -1035,9 +1125,7 @@ private CloseableIterable planFileScanTaskWithManifestCache(TableS } } - // Split tasks into smaller chunks based on target split size for parallel processing - targetSplitSize = determineTargetFileSplitSize(tasks); - return TableScanUtil.splitFiles(CloseableIterable.withNoopClose(tasks), targetSplitSize); + return tasks; } /** @@ -1367,9 +1455,11 @@ private List doGetSystemTableSplits() throws UserException { List splits = new ArrayList<>(); TableScan scan = createTableScan(); long startTime = System.currentTimeMillis(); - try (CloseableIterable fileScanTasks = scan.planFiles()) { + try { + List fileScanTasks = getOrLoadExternalScanTasks( + createFileScanTaskCacheKey(scan), () -> materializeFileScanTasks(scan)); fileScanTasks.forEach(task -> splits.add(createIcebergSysSplit(task))); - } catch (IOException e) { + } catch (Exception e) { throw new UserException(e.getMessage(), e); } finally { if (getSummaryProfile() != null) { @@ -1416,14 +1506,30 @@ private List doGetPositionDeletesSystemTableSplits() throws UserException long startTime = System.currentTimeMillis(); scan = scan.planWith(source.getCatalog().getThreadPoolWithPreAuth()); - try (CloseableIterable scanTasks = scan.planFiles()) { - for (ScanTask task : scanTasks) { - if (!(task instanceof PositionDeletesScanTask)) { - throw new UserException("Unexpected Iceberg position_deletes scan task: " + task); + BatchScan plannedScan = scan; + Snapshot snapshot = plannedScan.snapshot(); + IcebergScanTaskCacheKey cacheKey = new IcebergScanTaskCacheKey<>( + source.getCatalog().getId(), + source.getTargetTable().getId(), + snapshot == null ? null : snapshot.snapshotId(), + plannedScan.schema().schemaId(), + plannedScan.filter(), + plannedScan.isCaseSensitive(), + PositionDeletesScanTask.class.getName()); + try { + positionDeleteTasks = getOrLoadExternalScanTasks(cacheKey, () -> { + List tasks = new ArrayList<>(); + try (CloseableIterable scanTasks = plannedScan.planFiles()) { + for (ScanTask task : scanTasks) { + if (!(task instanceof PositionDeletesScanTask)) { + throw new UserException("Unexpected Iceberg position_deletes scan task: " + task); + } + tasks.add((PositionDeletesScanTask) task); + } } - positionDeleteTasks.add((PositionDeletesScanTask) task); - } - } catch (IOException e) { + return tasks; + }); + } catch (Exception e) { throw new UserException(e.getMessage(), e); } finally { if (getSummaryProfile() != null) { @@ -1614,6 +1720,9 @@ private void recordManifestCacheProfile() { if (!IcebergUtils.isManifestCacheEnabled(source.getCatalog())) { return; } + if (manifestCacheHits == 0 && manifestCacheMisses == 0 && manifestCacheFailures == 0) { + return; + } SummaryProfile summaryProfile = SummaryProfile.getSummaryProfile(ConnectContext.get()); if (summaryProfile == null || summaryProfile.getExecutionSummary() == null) { return; diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/paimon/source/PaimonScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/paimon/source/PaimonScanNode.java index bcf5be93649c68..a54758d38369da 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/paimon/source/PaimonScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/paimon/source/PaimonScanNode.java @@ -26,6 +26,7 @@ import org.apache.doris.common.UserException; import org.apache.doris.common.util.FileFormatUtils; import org.apache.doris.common.util.LocationPath; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.ExternalTable; import org.apache.doris.datasource.ExternalUtil; import org.apache.doris.datasource.FileQueryScanNode; @@ -87,6 +88,7 @@ import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.OptionalInt; import java.util.OptionalLong; @@ -741,45 +743,25 @@ public List getPaimonSplitFromAPI() throws if (PaimonScanParams.isPinnedEmptyScan(resolvedOptions)) { return Collections.emptyList(); } - Optional fileCreationTime = PaimonScanParams.getPinnedFileCreationTime(resolvedOptions); - if (fileCreationTime.isPresent()) { - if (!(paimonTable instanceof FileStoreTable)) { - throw new UserException("Paimon file-creation OPTIONS require a data table."); + int[] projectedColumns = new int[0]; + if (!PaimonScanParams.getPinnedFileCreationTime(resolvedOptions).isPresent()) { + List fieldNames = paimonTable.rowType().getFieldNames(); + projectedColumns = desc.getSlots().stream().mapToInt( + slot -> getFieldIndex(fieldNames, slot.getColumn().getName())) + .toArray(); + if (Arrays.stream(projectedColumns).anyMatch(index -> index < 0)) { + throw new UserException("Paimon scan schema does not contain all bound Doris columns."); } - FileStoreTable fileStoreTable = (FileStoreTable) paimonTable; - SnapshotReader snapshotReader = fileStoreTable.newSnapshotReader() - .withMode(ScanMode.ALL) - .withSnapshot(Long.parseLong( - paimonTable.options().get(CoreOptions.SCAN_SNAPSHOT_ID.key()))) - .withManifestEntryFilter(entry -> - entry.file().creationTimeEpochMillis() >= fileCreationTime.get()); - preserveBatchScanFilters(fileStoreTable, snapshotReader); - if (predicates != null) { - predicates.forEach(snapshotReader::withFilter); - } - return snapshotReader.read().splits(); - } - List fieldNames = paimonTable.rowType().getFieldNames(); - int[] projected = desc.getSlots().stream().mapToInt( - slot -> getFieldIndex(fieldNames, slot.getColumn().getName())) - .toArray(); - if (Arrays.stream(projected).anyMatch(index -> index < 0)) { - throw new UserException("Paimon scan schema does not contain all bound Doris columns."); - } - ReadBuilder readBuilder = paimonTable.newReadBuilder(); - TableScan scan = readBuilder.withFilter(predicates) - .withProjection(projected) - .newScan(); - PaimonMetricRegistry registry = new PaimonMetricRegistry(); - if (scan instanceof InnerTableScan) { - scan = ((InnerTableScan) scan).withMetricRegistry(registry); - } - List splits = scan.plan().splits(); - PaimonScanMetricsReporter.report(source.getTargetTable(), paimonTable.name(), registry); - if (!registry.getAllGroups().isEmpty()) { - registry.clear(); } - return splits; + int[] projected = projectedColumns; + PaimonSplitTaskCacheKey cacheKey = createPaimonSplitTaskCacheKey( + relationSnapshot, paimonTable, resolvedOptions, projected); + return getOrLoadExternalScanTasks(cacheKey, + () -> planPaimonSplits(paimonTable, resolvedOptions, projected)); + } catch (UserException e) { + throw e; + } catch (Exception e) { + throw new UserException("Failed to plan Paimon scan tasks", e); } finally { if (getSummaryProfile() != null) { getSummaryProfile().addExternalTableGetFileScanTasksTime(System.currentTimeMillis() - startTime); @@ -787,6 +769,121 @@ public List getPaimonSplitFromAPI() throws } } + private List planPaimonSplits( + Table paimonTable, Map resolvedOptions, int[] projected) throws UserException { + Optional fileCreationTime = PaimonScanParams.getPinnedFileCreationTime(resolvedOptions); + if (fileCreationTime.isPresent()) { + if (!(paimonTable instanceof FileStoreTable)) { + throw new UserException("Paimon file-creation OPTIONS require a data table."); + } + FileStoreTable fileStoreTable = (FileStoreTable) paimonTable; + SnapshotReader snapshotReader = fileStoreTable.newSnapshotReader() + .withMode(ScanMode.ALL) + .withSnapshot(Long.parseLong( + paimonTable.options().get(CoreOptions.SCAN_SNAPSHOT_ID.key()))) + .withManifestEntryFilter(entry -> + entry.file().creationTimeEpochMillis() >= fileCreationTime.get()); + preserveBatchScanFilters(fileStoreTable, snapshotReader); + if (predicates != null) { + predicates.forEach(snapshotReader::withFilter); + } + return snapshotReader.read().splits(); + } + ReadBuilder readBuilder = paimonTable.newReadBuilder(); + TableScan scan = readBuilder.withFilter(predicates) + .withProjection(projected) + .newScan(); + PaimonMetricRegistry registry = new PaimonMetricRegistry(); + if (scan instanceof InnerTableScan) { + scan = ((InnerTableScan) scan).withMetricRegistry(registry); + } + List splits = scan.plan().splits(); + PaimonScanMetricsReporter.report(source.getTargetTable(), paimonTable.name(), registry); + if (!registry.getAllGroups().isEmpty()) { + registry.clear(); + } + return splits; + } + + private PaimonSplitTaskCacheKey createPaimonSplitTaskCacheKey( + Optional relationSnapshot, Table paimonTable, + Map resolvedOptions, int[] projected) { + Long snapshotId = null; + Long schemaId = null; + if (relationSnapshot.isPresent() && relationSnapshot.get() instanceof PaimonMvccSnapshot) { + PaimonSnapshot snapshot = ((PaimonMvccSnapshot) relationSnapshot.get()) + .getSnapshotCacheValue().getSnapshot(); + snapshotId = snapshot.getSnapshotId(); + schemaId = snapshot.getSchemaId(); + } + return new PaimonSplitTaskCacheKey( + source.getCatalog().getId(), + source.getExternalTable().getId(), + source.getTargetTable().getId(), + snapshotId, + schemaId, + resolvedOptions, + paimonTable.options(), + projected, + PaimonUtil.encodeObjectToString(predicates)); + } + + private static final class PaimonSplitTaskCacheKey + implements ExternalScanTaskCacheKey { + private final long catalogId; + private final long relationTableId; + private final long targetTableId; + private final Long snapshotId; + private final Long schemaId; + private final Map resolvedOptions; + private final Map tableOptions; + private final int[] projected; + private final String serializedPredicates; + + private PaimonSplitTaskCacheKey( + long catalogId, long relationTableId, long targetTableId, Long snapshotId, Long schemaId, + Map resolvedOptions, Map tableOptions, int[] projected, + String serializedPredicates) { + this.catalogId = catalogId; + this.relationTableId = relationTableId; + this.targetTableId = targetTableId; + this.snapshotId = snapshotId; + this.schemaId = schemaId; + this.resolvedOptions = Collections.unmodifiableMap(new HashMap<>(resolvedOptions)); + this.tableOptions = Collections.unmodifiableMap(new HashMap<>(tableOptions)); + this.projected = Arrays.copyOf(projected, projected.length); + this.serializedPredicates = serializedPredicates; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof PaimonSplitTaskCacheKey)) { + return false; + } + PaimonSplitTaskCacheKey that = (PaimonSplitTaskCacheKey) object; + return catalogId == that.catalogId + && relationTableId == that.relationTableId + && targetTableId == that.targetTableId + && Objects.equals(snapshotId, that.snapshotId) + && Objects.equals(schemaId, that.schemaId) + && resolvedOptions.equals(that.resolvedOptions) + && tableOptions.equals(that.tableOptions) + && Arrays.equals(projected, that.projected) + && serializedPredicates.equals(that.serializedPredicates); + } + + @Override + public int hashCode() { + return 31 * Objects.hash( + catalogId, relationTableId, targetTableId, snapshotId, schemaId, + resolvedOptions, tableOptions, serializedPredicates) + + Arrays.hashCode(projected); + } + } + private void preserveBatchScanFilters(FileStoreTable table, SnapshotReader snapshotReader) { CoreOptions options = table.coreOptions(); // This direct reader bypasses DataTableBatchScan, so preserve its correctness filters for diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java index ba0a6d23261784..bd1c7386a2f51e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java @@ -30,6 +30,7 @@ import org.apache.doris.common.Id; import org.apache.doris.common.IdGenerator; import org.apache.doris.common.Pair; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.ExternalTable; import org.apache.doris.datasource.mvcc.MvccSnapshot; import org.apache.doris.datasource.mvcc.MvccTable; @@ -99,6 +100,10 @@ import java.util.Set; import java.util.Stack; import java.util.TreeMap; +import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import javax.annotation.concurrent.GuardedBy; @@ -322,6 +327,7 @@ public enum TableFrom { // IcebergScanNode // TODO: better solution? private List icebergRewriteFileScanTasks = null; + private volatile ExternalScanTaskCache externalScanTaskCache = new ExternalScanTaskCache(); // For Iceberg rewrite operations: control whether to use GATHER distribution // When true, data will be collected to a single node to avoid generating too many small files private boolean useGatherForIcebergRewrite = false; @@ -910,6 +916,7 @@ protected void finalize() throws Throwable { @Override public void close() { + clearExternalScanTasks(); releasePlannerResources(); } @@ -1013,6 +1020,9 @@ public void resetMvccSnapshots() { latestSnapshotFences.clear(); resolvedSnapshotScanParams.clear(); tableMetadataSnapshots.clear(); + ExternalScanTaskCache oldCache = externalScanTaskCache; + externalScanTaskCache = new ExternalScanTaskCache(); + oldCache.invalidate(); // PREPARE keeps preload candidates, but completion belongs to one analysis pass and must // not suppress preloading after the next EXECUTE resets its snapshot generation. externalMetadataPreloadResult = null; @@ -1366,6 +1376,87 @@ public List getAndClearIcebergRewriteFileScanTa return tasks; } + public ExternalScanTaskCache getExternalScanTaskCache() { + return externalScanTaskCache; + } + + /** + * Release scan tasks at the end of one execution without closing reusable prepared-statement + * state. Delayed scan work retains only the invalidated generation and cannot repopulate this + * StatementContext. + */ + public void clearExternalScanTasks() { + externalScanTaskCache.invalidate(); + } + + /** + * One execution generation of statement-scoped external scan tasks. + * + *

Scan nodes capture this object when they are constructed. Resetting a prepared statement + * swaps the generation before invalidating the old one, so a delayed asynchronous scan from + * the previous execution cannot insert tasks into the next execution's cache. + */ + public static final class ExternalScanTaskCache { + private final Map, CompletableFuture>> tasks = + new ConcurrentHashMap<>(); + private boolean invalidated; + + /** + * Return the tasks for {@code key}, loading and publishing an immutable result once per + * cache generation. + */ + @SuppressWarnings("unchecked") + public List getOrLoad( + ExternalScanTaskCacheKey key, Callable> loader) throws Exception { + CompletableFuture> newLoad = new CompletableFuture<>(); + CompletableFuture> load; + boolean cacheable; + synchronized (this) { + cacheable = !invalidated; + if (cacheable) { + load = tasks.putIfAbsent(key, newLoad); + } else { + load = null; + } + } + if (!cacheable) { + return immutableCopy(loader.call()); + } + if (load == null) { + try { + List loadedTasks = immutableCopy(loader.call()); + newLoad.complete(loadedTasks); + return loadedTasks; + } catch (Exception | Error throwable) { + newLoad.completeExceptionally(throwable); + tasks.remove(key, newLoad); + throw throwable; + } + } + try { + return (List) load.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw e; + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception) { + throw (Exception) cause; + } + throw (Error) cause; + } + } + + private synchronized void invalidate() { + invalidated = true; + tasks.clear(); + } + + private static List immutableCopy(List loadedTasks) { + return Collections.unmodifiableList(new ArrayList<>(loadedTasks)); + } + } + /** * Set whether to use GATHER distribution for Iceberg rewrite operations. * When enabled, data will be collected to a single node to minimize output files. diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java index fca09750e51ae7..5bda3026fa6e49 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java @@ -190,11 +190,15 @@ protected void handleExecute(PrepareCommand prepareCommand, long stmtId, Prepare ctx.getState().setError(ErrorCode.ERR_UNKNOWN_ERROR, e.getClass().getSimpleName() + ", msg: " + e.getMessage()); } - if (ctx.getSessionVariable().isEnablePreparedStmtAuditLog()) { - auditAfterExec(stmtStr, executor.getParsedStmt(), executor.getQueryStatisticsForAuditLog(), true); - } else { - // When audit log is disabled for prepared statements, still update QPS metrics. - AuditLogHelper.updateMetrics(ctx); + try { + if (ctx.getSessionVariable().isEnablePreparedStmtAuditLog()) { + auditAfterExec(stmtStr, executor.getParsedStmt(), executor.getQueryStatisticsForAuditLog(), true); + } else { + // When audit log is disabled for prepared statements, still update QPS metrics. + AuditLogHelper.updateMetrics(ctx); + } + } finally { + prepCtx.statementContext.clearExternalScanTasks(); } } 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 38aaceda1739d0..ca61c58c692648 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/source/HiveScanNodeTest.java @@ -17,14 +17,20 @@ package org.apache.doris.datasource.hive.source; +import org.apache.doris.analysis.TableSample; import org.apache.doris.analysis.TupleDescriptor; import org.apache.doris.analysis.TupleId; import org.apache.doris.datasource.TableFormatType; +import org.apache.doris.datasource.hive.HMSCachedClient; 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.datasource.hive.HiveTransaction; +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.thrift.TFileScanRangeParams; import org.apache.doris.thrift.TFileTextScanRangeParams; @@ -33,13 +39,138 @@ import org.junit.Test; import org.mockito.Mockito; +import java.lang.reflect.Field; import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; public class HiveScanNodeTest { private static final long MB = 1024L * 1024L; + @Test + public void testStatementCacheReusesListingOnlyForSamePartitionIdentity() throws Exception { + ConnectContext previousContext = ConnectContext.get(); + ConnectContext context = new ConnectContext(); + StatementContext statementContext = new StatementContext(context, null); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + HMSExternalTable table = Mockito.mock(HMSExternalTable.class); + HMSExternalCatalog catalog = Mockito.mock(HMSExternalCatalog.class); + Mockito.when(table.getCatalog()).thenReturn(catalog); + Mockito.when(table.getId()).thenReturn(2L); + Mockito.when(catalog.getId()).thenReturn(1L); + Mockito.when(catalog.bindBrokerName()).thenReturn(""); + HiveScanNode firstNode = createHiveScanNode(0, table); + HiveScanNode secondNode = createHiveScanNode(1, table); + HiveExternalMetaCache cache = Mockito.mock(HiveExternalMetaCache.class); + Mockito.when(cache.getFilesByPartitions( + Mockito.anyList(), Mockito.anyBoolean(), Mockito.anyBoolean(), + Mockito.isNull(), Mockito.eq(table))).thenReturn(Collections.emptyList()); + List firstPartitions = Collections.singletonList(new HivePartition( + null, false, "parquet", "hdfs://warehouse/t/p=1", + Collections.singletonList("1"), Collections.emptyMap())); + List equivalentPartitions = Collections.singletonList(new HivePartition( + null, false, "parquet", "hdfs://warehouse/t/p=1", + Collections.singletonList("1"), Collections.emptyMap())); + List differentPartitions = Collections.singletonList(new HivePartition( + null, false, "parquet", "hdfs://warehouse/t/p=2", + Collections.singletonList("2"), Collections.emptyMap())); + + invokeGetFileSplitByPartitions(firstNode, cache, firstPartitions); + invokeGetFileSplitByPartitions(secondNode, cache, equivalentPartitions); + invokeGetFileSplitByPartitions(secondNode, cache, differentPartitions); + + Mockito.verify(cache).getFilesByPartitions( + Mockito.same(firstPartitions), Mockito.anyBoolean(), Mockito.eq(false), + Mockito.isNull(), Mockito.eq(table)); + Mockito.verify(cache).getFilesByPartitions( + Mockito.same(differentPartitions), Mockito.anyBoolean(), Mockito.eq(false), + Mockito.isNull(), Mockito.eq(table)); + Mockito.verifyNoMoreInteractions(cache); + } finally { + statementContext.close(); + ConnectContext.remove(); + if (previousContext != null) { + previousContext.setThreadLocalInfo(); + } + } + } + + @Test + public void testTableSampleDoesNotMutateCachedFileStatus() throws Exception { + HiveScanNode node = createHiveScanNode(); + node.setTableSample(new TableSample(true, 100L, 0L)); + HiveExternalMetaCache.HiveFileStatus cachedStatus = + new HiveExternalMetaCache.HiveFileStatus(); + cachedStatus.setLength(10L); + HiveExternalMetaCache.FileCacheValue cacheValue = + new HiveExternalMetaCache.FileCacheValue(); + cacheValue.setSplittable(true); + cacheValue.setPartitionValues(Arrays.asList("2026", "08")); + cacheValue.getFiles().add(cachedStatus); + + Method method = HiveScanNode.class.getDeclaredMethod("selectFiles", List.class); + method.setAccessible(true); + @SuppressWarnings("unchecked") + List sampled = + (List) method.invoke( + node, Collections.singletonList(cacheValue)); + + Assert.assertEquals(1, sampled.size()); + Assert.assertNotSame(cachedStatus, sampled.get(0)); + Assert.assertTrue(sampled.get(0).isSplittable()); + Assert.assertEquals(Arrays.asList("2026", "08"), sampled.get(0).getPartitionValues()); + Assert.assertFalse(cachedStatus.isSplittable()); + Assert.assertNull(cachedStatus.getPartitionValues()); + } + + @Test + public void testTransactionalListingBypassesStatementCachePath() throws Exception { + ConnectContext previousContext = ConnectContext.get(); + ConnectContext context = new ConnectContext(); + StatementContext statementContext = new StatementContext(context, null); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + HMSExternalTable table = Mockito.mock(HMSExternalTable.class); + HMSExternalCatalog catalog = Mockito.mock(HMSExternalCatalog.class); + Mockito.when(table.getCatalog()).thenReturn(catalog); + Mockito.when(catalog.bindBrokerName()).thenReturn(""); + Mockito.when(catalog.getClient()).thenReturn(Mockito.mock(HMSCachedClient.class)); + HiveScanNode node = createHiveScanNode(0, table); + HiveTransaction transaction = Mockito.mock(HiveTransaction.class); + Map validWriteIds = + Collections.singletonMap("db.table", "valid-write-ids"); + Mockito.when(transaction.getValidWriteIds(Mockito.any())).thenReturn(validWriteIds); + Mockito.when(transaction.isFullAcid()).thenReturn(true); + Field transactionField = HiveScanNode.class.getDeclaredField("hiveTransaction"); + transactionField.setAccessible(true); + transactionField.set(node, transaction); + + HiveExternalMetaCache cache = Mockito.mock(HiveExternalMetaCache.class); + Mockito.when(cache.getFilesByTransaction( + Collections.emptyList(), validWriteIds, true, null)) + .thenReturn(Collections.emptyList()); + + invokeGetFileSplitByPartitions(node, cache, Collections.emptyList()); + invokeGetFileSplitByPartitions(node, cache, Collections.emptyList()); + + Mockito.verify(cache, Mockito.times(2)).getFilesByTransaction( + Collections.emptyList(), validWriteIds, true, null); + Mockito.verifyNoMoreInteractions(cache); + } finally { + statementContext.close(); + ConnectContext.remove(); + if (previousContext != null) { + previousContext.setThreadLocalInfo(); + } + } + } + @Test public void testDetermineTargetFileSplitSizeHonorsMaxFileSplitNum() throws Exception { SessionVariable sv = new SessionVariable(); @@ -125,4 +256,21 @@ private HiveScanNode createHiveScanNode(boolean partitioned) { desc.setTable(table); return new HiveScanNode(new PlanNodeId(0), desc, false, sv, null, ScanContext.EMPTY); } + + private HiveScanNode createHiveScanNode(int id, HMSExternalTable table) { + TupleDescriptor desc = new TupleDescriptor(new TupleId(id)); + desc.setTable(table); + return new HiveScanNode( + new PlanNodeId(id), desc, false, new SessionVariable(), null, ScanContext.EMPTY); + } + + private void invokeGetFileSplitByPartitions( + HiveScanNode node, HiveExternalMetaCache cache, List partitions) + throws Exception { + Method method = HiveScanNode.class.getDeclaredMethod( + "getFileSplitByPartitions", HiveExternalMetaCache.class, List.class, + List.class, String.class, int.class, boolean.class); + method.setAccessible(true); + method.invoke(node, cache, partitions, new ArrayList<>(), null, 1, false); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/hudi/source/HudiScanNodeTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/hudi/source/HudiScanNodeTest.java new file mode 100644 index 00000000000000..354abf8bc411e3 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hudi/source/HudiScanNodeTest.java @@ -0,0 +1,327 @@ +// 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.hudi.source; + +import org.apache.doris.common.profile.SummaryProfile; +import org.apache.doris.common.util.LocationPath; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; +import org.apache.doris.datasource.FileQueryScanNode; +import org.apache.doris.datasource.TableFormatType; +import org.apache.doris.datasource.hive.HMSExternalTable; +import org.apache.doris.datasource.hive.HivePartition; +import org.apache.doris.datasource.hive.source.HiveScanNode; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.qe.SessionVariable; +import org.apache.doris.spi.Split; + +import com.google.common.collect.ImmutableMap; +import org.apache.hudi.common.model.HoodieBaseFile; +import org.apache.hudi.common.table.HoodieTableMetaClient; +import org.apache.hudi.common.table.view.HoodieTableFileSystemView; +import org.apache.hudi.storage.StoragePath; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Answers; +import org.mockito.Mockito; + +import java.lang.reflect.Constructor; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Stream; + +public class HudiScanNodeTest { + + @Test + public void testCopyHudiSplitIsolatesMutableState() { + HudiSplit source = new HudiSplit( + LocationPath.of("hdfs://host/table/file.parquet"), + 1, 2, 3, new String[] {"host-1"}, new ArrayList<>(Collections.singletonList("p1"))); + source.setModificationTime(4); + source.setTableFormatType(TableFormatType.HUDI); + source.setAlternativeHosts(new ArrayList<>(Collections.singletonList("host-2"))); + source.setSelfSplitWeight(5L); + source.setTargetSplitSize(6L); + source.setHudiDeltaLogs(new ArrayList<>(Collections.singletonList("log-1"))); + source.setHudiColumnNames(new ArrayList<>(Collections.singletonList("column-1"))); + source.setHudiColumnTypes(new ArrayList<>(Collections.singletonList("type-1"))); + source.setNestedFields(new ArrayList<>(Collections.singletonList("nested-1"))); + source.setHudiPartitionValues(new java.util.HashMap<>(ImmutableMap.of("key", "value"))); + + HudiSplit copy = HudiScanNode.copyHudiSplit(source); + copy.getHosts()[0] = "changed-host"; + copy.getPartitionValues().set(0, "changed-partition"); + copy.getAlternativeHosts().set(0, "changed-alternative-host"); + copy.getHudiDeltaLogs().set(0, "changed-log"); + copy.getHudiColumnNames().set(0, "changed-column"); + copy.getHudiColumnTypes().set(0, "changed-type"); + copy.getNestedFields().set(0, "changed-nested"); + copy.getHudiPartitionValues().put("key", "changed-value"); + copy.setTargetSplitSize(7L); + + Assertions.assertNotSame(source.getHosts(), copy.getHosts()); + Assertions.assertEquals(Arrays.asList("host-1"), Arrays.asList(source.getHosts())); + Assertions.assertEquals(Collections.singletonList("p1"), source.getPartitionValues()); + Assertions.assertEquals(Collections.singletonList("host-2"), source.getAlternativeHosts()); + Assertions.assertEquals(Collections.singletonList("log-1"), source.getHudiDeltaLogs()); + Assertions.assertEquals(Collections.singletonList("column-1"), source.getHudiColumnNames()); + Assertions.assertEquals(Collections.singletonList("type-1"), source.getHudiColumnTypes()); + Assertions.assertEquals(Collections.singletonList("nested-1"), source.getNestedFields()); + Assertions.assertEquals(ImmutableMap.of("key", "value"), source.getHudiPartitionValues()); + Assertions.assertEquals(6L, source.getTargetSplitSize()); + } + + @Test + public void testPartitionPlanningCacheHitsAndReturnsIndependentCopies() throws Exception { + StatementContext.ExternalScanTaskCache cache = new StatementContext.ExternalScanTaskCache(); + HivePartition partition = partition("file:///table/p=1", Collections.singletonList("1")); + HoodieTableFileSystemView firstView = fileSystemView("file:///table/p=1/file.parquet"); + HoodieTableFileSystemView duplicateView = fileSystemView("file:///should-not-be-planned.parquet"); + HudiScanNode firstNode = partitionScanNode(cache, firstView, "100", true, false); + HudiScanNode duplicateNode = partitionScanNode(cache, duplicateView, "100", true, false); + + List first = invokeGetPartitionSplits(firstNode, partition); + List duplicate = invokeGetPartitionSplits(duplicateNode, partition); + + Mockito.verify(firstView, Mockito.times(1)).getLatestBaseFilesBeforeOrOn("p=1", "100"); + Mockito.verify(duplicateView, Mockito.never()).getLatestBaseFilesBeforeOrOn(Mockito.any(), Mockito.any()); + Assertions.assertEquals(1, first.size()); + Assertions.assertEquals(1, duplicate.size()); + Assertions.assertNotSame(first.get(0), duplicate.get(0)); + Assertions.assertEquals(first.get(0).getPathString(), duplicate.get(0).getPathString()); + + HudiSplit firstSplit = (HudiSplit) first.get(0); + HudiSplit duplicateSplit = (HudiSplit) duplicate.get(0); + firstSplit.getPartitionValues().set(0, "changed"); + firstSplit.setTargetSplitSize(123L); + Assertions.assertEquals(Collections.singletonList("1"), duplicateSplit.getPartitionValues()); + Assertions.assertNull(duplicateSplit.getTargetSplitSize()); + } + + @Test + public void testPartitionPlanningCacheMissesForInstantAndPartition() throws Exception { + StatementContext.ExternalScanTaskCache cache = new StatementContext.ExternalScanTaskCache(); + HoodieTableFileSystemView firstView = fileSystemView("file:///table/p=1/first.parquet"); + HoodieTableFileSystemView instantView = fileSystemView("file:///table/p=1/instant.parquet"); + HoodieTableFileSystemView partitionView = fileSystemView("file:///table/p=2/partition.parquet"); + HivePartition firstPartition = partition("file:///table/p=1", Collections.singletonList("1")); + HivePartition secondPartition = partition("file:///table/p=2", Collections.singletonList("2")); + + invokeGetPartitionSplits(partitionScanNode(cache, firstView, "100", true, false), firstPartition); + invokeGetPartitionSplits(partitionScanNode(cache, instantView, "101", true, false), firstPartition); + invokeGetPartitionSplits(partitionScanNode(cache, partitionView, "100", true, false), secondPartition); + + Mockito.verify(firstView, Mockito.times(1)).getLatestBaseFilesBeforeOrOn("p=1", "100"); + Mockito.verify(instantView, Mockito.times(1)).getLatestBaseFilesBeforeOrOn("p=1", "101"); + Mockito.verify(partitionView, Mockito.times(1)).getLatestBaseFilesBeforeOrOn("p=2", "100"); + } + + @Test + public void testPartitionCacheKeySeparatesReaderAndRuntimePruneModes() throws Exception { + HivePartition partition = partition("file:///table/p=1", Collections.singletonList("1")); + Object nativeKey = newPartitionCacheKey("100", true, false, partition); + Object sameKey = newPartitionCacheKey("100", true, false, partition); + Object jniKey = newPartitionCacheKey("100", false, false, partition); + Object runtimePruneKey = newPartitionCacheKey("100", true, true, partition); + + assertCacheHitsOnlyEquivalentKeys(nativeKey, sameKey, jniKey, runtimePruneKey); + } + + @Test + public void testIncrementalPlanningCacheUsesConnectorWrapperAndCopiesSplits() throws Exception { + StatementContext.ExternalScanTaskCache cache = new StatementContext.ExternalScanTaskCache(); + Map options = ImmutableMap.of("hoodie.datasource.query.type", "incremental"); + AtomicInteger firstLoads = new AtomicInteger(); + AtomicInteger duplicateLoads = new AtomicInteger(); + AtomicInteger differentStartLoads = new AtomicInteger(); + AtomicInteger differentOptionsLoads = new AtomicInteger(); + IncrementalRelation firstRelation = + incrementalRelation("10", "20", options, firstLoads, "first.parquet"); + IncrementalRelation duplicateRelation = + incrementalRelation("10", "20", new HashMap<>(options), duplicateLoads, "duplicate.parquet"); + IncrementalRelation differentStartRelation = + incrementalRelation("11", "20", options, differentStartLoads, "different-start.parquet"); + IncrementalRelation differentOptionsRelation = incrementalRelation( + "10", "20", ImmutableMap.of("hoodie.datasource.query.type", "snapshot"), + differentOptionsLoads, "different-options.parquet"); + + List first = invokeGetIncrementalSplits( + incrementalScanNode(cache, firstRelation, true)); + List duplicate = invokeGetIncrementalSplits( + incrementalScanNode(cache, duplicateRelation, true)); + List differentStart = invokeGetIncrementalSplits( + incrementalScanNode(cache, differentStartRelation, true)); + List differentOptions = invokeGetIncrementalSplits( + incrementalScanNode(cache, differentOptionsRelation, true)); + + Assertions.assertEquals(1, firstLoads.get()); + Assertions.assertEquals(0, duplicateLoads.get()); + Assertions.assertEquals(1, differentStartLoads.get()); + Assertions.assertEquals(1, differentOptionsLoads.get()); + Assertions.assertEquals(1, first.size()); + Assertions.assertEquals(1, duplicate.size()); + Assertions.assertEquals(1, differentStart.size()); + Assertions.assertEquals(1, differentOptions.size()); + Assertions.assertNotSame(first.get(0), duplicate.get(0)); + Assertions.assertEquals(first.get(0).getPathString(), duplicate.get(0).getPathString()); + + HudiSplit firstSplit = (HudiSplit) first.get(0); + HudiSplit duplicateSplit = (HudiSplit) duplicate.get(0); + firstSplit.getPartitionValues().set(0, "changed"); + firstSplit.setTargetSplitSize(123L); + Assertions.assertEquals(Collections.singletonList("p=1"), duplicateSplit.getPartitionValues()); + Assertions.assertNull(duplicateSplit.getTargetSplitSize()); + } + + private static HudiScanNode partitionScanNode( + StatementContext.ExternalScanTaskCache cache, HoodieTableFileSystemView fsView, + String queryInstant, boolean nativeReader, boolean runtimePrune) throws Exception { + HudiScanNode node = Mockito.mock(HudiScanNode.class, Answers.CALLS_REAL_METHODS); + HMSExternalTable table = Mockito.mock(HMSExternalTable.class, Answers.RETURNS_DEEP_STUBS); + Mockito.when(table.getCatalog().getId()).thenReturn(1L); + Mockito.when(table.getId()).thenReturn(2L); + Mockito.when(table.getStoragePropertiesMap()).thenReturn(Collections.emptyMap()); + SessionVariable sessionVariable = new SessionVariable(); + sessionVariable.setForceJniScanner(!nativeReader); + sessionVariable.setEnableRuntimeFilterPartitionPrune(runtimePrune); + HoodieTableMetaClient metaClient = Mockito.mock(HoodieTableMetaClient.class); + Mockito.when(metaClient.getBasePath()).thenReturn(new StoragePath("file:///table")); + + setField(node, FileQueryScanNode.class, "externalScanTaskCache", cache); + setField(node, HiveScanNode.class, "hmsTable", table); + setField(node, FileQueryScanNode.class, "sessionVariable", sessionVariable); + setField(node, HudiScanNode.class, "isCowTable", true); + setField(node, HudiScanNode.class, "queryInstant", queryInstant); + setField(node, HudiScanNode.class, "hudiClient", metaClient); + setField(node, HudiScanNode.class, "fsView", fsView); + setField(node, HudiScanNode.class, "noLogsSplitNum", new AtomicLong()); + return node; + } + + private static HoodieTableFileSystemView fileSystemView(String filePath) { + HoodieBaseFile baseFile = Mockito.mock(HoodieBaseFile.class); + Mockito.when(baseFile.getPath()).thenReturn(filePath); + Mockito.when(baseFile.getFileSize()).thenReturn(10L); + HoodieTableFileSystemView fsView = Mockito.mock(HoodieTableFileSystemView.class); + Mockito.when(fsView.getLatestBaseFilesBeforeOrOn(Mockito.any(), Mockito.any())) + .thenAnswer(invocation -> Stream.of(baseFile)); + return fsView; + } + + private static HudiScanNode incrementalScanNode( + StatementContext.ExternalScanTaskCache cache, IncrementalRelation relation, + boolean nativeReader) throws Exception { + HudiScanNode node = Mockito.mock(HudiScanNode.class, Answers.CALLS_REAL_METHODS); + HMSExternalTable table = Mockito.mock(HMSExternalTable.class, Answers.RETURNS_DEEP_STUBS); + Mockito.when(table.getCatalog().getId()).thenReturn(1L); + Mockito.when(table.getId()).thenReturn(2L); + SessionVariable sessionVariable = new SessionVariable(); + sessionVariable.setForceJniScanner(!nativeReader); + + setField(node, FileQueryScanNode.class, "externalScanTaskCache", cache); + setField(node, HiveScanNode.class, "hmsTable", table); + setField(node, FileQueryScanNode.class, "sessionVariable", sessionVariable); + setField(node, FileQueryScanNode.class, "summaryProfile", Mockito.mock(SummaryProfile.class)); + setField(node, HudiScanNode.class, "isCowTable", true); + setField(node, HudiScanNode.class, "incrementalRelation", relation); + setField(node, HudiScanNode.class, "noLogsSplitNum", new AtomicLong()); + return node; + } + + private static IncrementalRelation incrementalRelation( + String start, String end, Map options, + AtomicInteger loads, String plannedPath) { + IncrementalRelation relation = Mockito.mock(IncrementalRelation.class); + Mockito.when(relation.getStartTs()).thenReturn(start); + Mockito.when(relation.getEndTs()).thenReturn(end); + Mockito.when(relation.getHoodieParams()).thenReturn(options); + Mockito.when(relation.collectSplits()).thenAnswer(invocation -> { + loads.incrementAndGet(); + HudiSplit split = new HudiSplit( + LocationPath.of("file:///table/" + plannedPath), + 0, 10, 10, new String[0], + new ArrayList<>(Collections.singletonList("p=1"))); + split.setHudiDeltaLogs(Collections.emptyList()); + return Collections.singletonList(split); + }); + return relation; + } + + private static HivePartition partition(String path, List values) { + return new HivePartition(null, false, "parquet", path, new ArrayList<>(values), Collections.emptyMap()); + } + + @SuppressWarnings("unchecked") + private static List invokeGetPartitionSplits(HudiScanNode node, HivePartition partition) + throws Exception { + Method method = HudiScanNode.class.getDeclaredMethod( + "getPartitionSplits", HivePartition.class, List.class); + method.setAccessible(true); + List splits = new ArrayList<>(); + method.invoke(node, partition, splits); + return splits; + } + + @SuppressWarnings("unchecked") + private static List invokeGetIncrementalSplits(HudiScanNode node) throws Exception { + Method method = HudiScanNode.class.getDeclaredMethod("getIncrementalSplits"); + method.setAccessible(true); + return (List) method.invoke(node); + } + + private static Object newPartitionCacheKey( + String instant, boolean nativeReader, boolean runtimePrune, HivePartition partition) + throws Exception { + Class keyClass = Class.forName(HudiScanNode.class.getName() + "$HudiFileScanTaskCacheKey"); + Constructor constructor = keyClass.getDeclaredConstructor( + long.class, long.class, String.class, boolean.class, boolean.class, HivePartition.class); + constructor.setAccessible(true); + return constructor.newInstance(1L, 2L, instant, nativeReader, runtimePrune, partition); + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private static void assertCacheHitsOnlyEquivalentKeys(Object first, Object same, Object... different) + throws Exception { + StatementContext.ExternalScanTaskCache cache = new StatementContext.ExternalScanTaskCache(); + AtomicInteger loads = new AtomicInteger(); + List firstResult = cache.getOrLoad((ExternalScanTaskCacheKey) first, + () -> Collections.singletonList("load-" + loads.incrementAndGet())); + List sameResult = cache.getOrLoad((ExternalScanTaskCacheKey) same, + () -> Collections.singletonList("load-" + loads.incrementAndGet())); + Assertions.assertSame(firstResult, sameResult); + Assertions.assertEquals(1, loads.get()); + for (Object key : different) { + cache.getOrLoad((ExternalScanTaskCacheKey) key, + () -> Collections.singletonList("load-" + loads.incrementAndGet())); + } + Assertions.assertEquals(1 + different.length, loads.get()); + } + + 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); + } +} 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 670754614ddbad..2056eccc518e6d 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 @@ -61,9 +61,12 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.AppendFiles; import org.apache.iceberg.BaseMetadataTable; import org.apache.iceberg.BaseTable; import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; import org.apache.iceberg.DeleteFile; import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; @@ -79,10 +82,15 @@ import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableProperties; import org.apache.iceberg.TableScan; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.hadoop.HadoopTables; +import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.types.Types; import org.apache.iceberg.util.ScanTaskUtil; import org.junit.Assert; +import org.junit.Rule; import org.junit.Test; +import org.junit.rules.TemporaryFolder; import org.mockito.Mockito; import java.lang.reflect.Field; @@ -98,10 +106,14 @@ import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; public class IcebergScanNodeTest { private static final long MB = 1024L * 1024L; + @Rule + public TemporaryFolder temporaryFolder = new TemporaryFolder(); + @SuppressWarnings("unchecked") private static Optional>> extractNameMapping( IcebergScanNode node) throws Exception { @@ -116,6 +128,23 @@ private static Table useFrozenTableGeneration(IcebergScanNode node, Table table) return (Table) method.invoke(node, table); } + @SuppressWarnings("unchecked") + private static CloseableIterable planFileScanTaskWithManifestCache( + IcebergScanNode node, TableScan scan) throws Exception { + Method method = IcebergScanNode.class.getDeclaredMethod( + "planFileScanTaskWithManifestCache", TableScan.class); + method.setAccessible(true); + return (CloseableIterable) method.invoke(node, scan); + } + + @SuppressWarnings("unchecked") + private static CloseableIterable splitFiles( + IcebergScanNode node, TableScan scan) throws Exception { + Method method = IcebergScanNode.class.getDeclaredMethod("splitFiles", TableScan.class); + method.setAccessible(true); + return (CloseableIterable) method.invoke(node, scan); + } + private static class TestIcebergScanNode extends IcebergScanNode { private final boolean enableMappingVarbinary; private final boolean batchMode; @@ -209,6 +238,22 @@ TFileScanRangeParams initializeAndGetIcebergSchemaInfo() throws UserException { } } + private static class ManifestPlanningIcebergScanNode extends TestIcebergScanNode { + private final AtomicInteger manifestLoadCount; + + ManifestPlanningIcebergScanNode(SessionVariable sv, AtomicInteger manifestLoadCount) { + super(sv); + this.manifestLoadCount = manifestLoadCount; + } + + @Override + protected List loadFileScanTasksWithManifestCache( + TableScan scan, Snapshot snapshot) { + manifestLoadCount.incrementAndGet(); + return Collections.emptyList(); + } + } + @Test public void testEmitsCurrentIcebergScanSemanticsCapability() { TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()); @@ -556,6 +601,337 @@ public void testCountStarVariantCompatibilityExemptionRequiresSnapshotCount() th } } + @Test + public void testSameIcebergScanReusesPlannedFileTasksWithinStatement() throws Exception { + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + TestIcebergScanNode firstNode = new TestIcebergScanNode(new SessionVariable()); + TestIcebergScanNode secondNode = new TestIcebergScanNode(new SessionVariable()); + setIcebergSource(firstNode, mockIcebergSource(10L, 20L)); + setIcebergSource(secondNode, mockIcebergSource(10L, 20L)); + TableScan firstScan = mockTableScan(30L, 40, Expressions.equal("id", 1)); + TableScan secondScan = mockTableScan(30L, 40, Expressions.equal("id", 1)); + FileScanTask task = Mockito.mock(FileScanTask.class); + AtomicInteger planCalls = new AtomicInteger(); + + List firstTasks = firstNode.getOrPlanFileScanTasks(firstScan, () -> { + planCalls.incrementAndGet(); + return Collections.singletonList(task); + }); + List secondTasks = secondNode.getOrPlanFileScanTasks(secondScan, () -> { + planCalls.incrementAndGet(); + return Collections.emptyList(); + }); + + Assert.assertEquals(1, planCalls.get()); + Assert.assertSame(firstTasks, secondTasks); + Assert.assertEquals(Collections.singletonList(task), secondTasks); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + @Test + public void testRepeatedActualIcebergPlanFilesUsesStatementCache() throws Exception { + Schema schema = new Schema( + Types.NestedField.optional(1, "id", Types.IntegerType.get())); + Table table = new HadoopTables(new Configuration()).create( + schema, PartitionSpec.unpartitioned(), + temporaryFolder.newFolder("repeated_scan_table").toURI().toString()); + AppendFiles append = table.newFastAppend(); + int fileCount = 100; + for (int i = 0; i < fileCount; i++) { + append.appendFile(DataFiles.builder(table.spec()) + .withPath("file:/tmp/repeated-scan-" + i + ".parquet") + .withFileSizeInBytes(1024) + .withRecordCount(1) + .withFormat(FileFormat.PARQUET) + .build()); + } + append.commit(); + + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()); + setIcebergSource(node, mockIcebergSource(10L, 20L)); + AtomicInteger planCalls = new AtomicInteger(); + + List tasks = Collections.emptyList(); + for (int i = 0; i < 3; i++) { + TableScan scan = table.newScan().filter(Expressions.equal("id", 1)); + tasks = node.getOrPlanFileScanTasks(scan, () -> { + planCalls.incrementAndGet(); + return materializeTasks(scan); + }); + } + + Assert.assertEquals(1, planCalls.get()); + Assert.assertEquals(fileCount, tasks.size()); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + @Test + public void testManifestPlanningPathUsesStatementCache() throws Exception { + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + AtomicInteger manifestLoadCount = new AtomicInteger(); + try { + ManifestPlanningIcebergScanNode firstNode = + new ManifestPlanningIcebergScanNode(new SessionVariable(), manifestLoadCount); + ManifestPlanningIcebergScanNode secondNode = + new ManifestPlanningIcebergScanNode(new SessionVariable(), manifestLoadCount); + setIcebergSource(firstNode, mockIcebergSource(10L, 20L)); + setIcebergSource(secondNode, mockIcebergSource(10L, 20L)); + + try (CloseableIterable ignored = planFileScanTaskWithManifestCache( + firstNode, mockTableScan(30L, 40, Expressions.equal("id", 1)))) { + // The test only verifies the production manifest-planning cache wrapper. + } + try (CloseableIterable ignored = planFileScanTaskWithManifestCache( + secondNode, mockTableScan(30L, 40, Expressions.equal("id", 1)))) { + // The second equivalent relation must reuse the first relation's task list. + } + + Assert.assertEquals(1, manifestLoadCount.get()); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + @Test + public void testSplitFilesUsesStatementCacheOutsideStreamingModes() throws Exception { + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + AtomicInteger planCalls = new AtomicInteger(); + try { + TestIcebergScanNode firstNode = new TestIcebergScanNode(new SessionVariable()); + TestIcebergScanNode secondNode = new TestIcebergScanNode(new SessionVariable()); + setIcebergSource(firstNode, mockIcebergSource(10L, 20L)); + setIcebergSource(secondNode, mockIcebergSource(10L, 20L)); + TableScan firstScan = mockTableScanWithPlanCounter(planCalls); + TableScan secondScan = mockTableScanWithPlanCounter(planCalls); + + try (CloseableIterable ignored = splitFiles(firstNode, firstScan)) { + // Materialization happens before splitFiles returns. + } + try (CloseableIterable ignored = splitFiles(secondNode, secondScan)) { + // The second equivalent relation must reuse the materialized native tasks. + } + + Assert.assertEquals(1, planCalls.get()); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + @Test + public void testSplitFilesBypassesStatementCacheForStreamingModes() throws Exception { + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + AtomicInteger batchPlanCalls = new AtomicInteger(); + AtomicInteger explicitSizePlanCalls = new AtomicInteger(); + try { + TestIcebergScanNode firstBatchNode = + new TestIcebergScanNode(new SessionVariable(), false, true); + TestIcebergScanNode secondBatchNode = + new TestIcebergScanNode(new SessionVariable(), false, true); + setIcebergSource(firstBatchNode, mockIcebergSource(10L, 20L)); + setIcebergSource(secondBatchNode, mockIcebergSource(10L, 20L)); + try (CloseableIterable ignored = splitFiles( + firstBatchNode, mockTableScanWithPlanCounter(batchPlanCalls))) { + // Batch mode preserves lazy Iceberg planning. + } + try (CloseableIterable ignored = splitFiles( + secondBatchNode, mockTableScanWithPlanCounter(batchPlanCalls))) { + // Each batch relation must own its streaming iterable. + } + + SessionVariable explicitSizeVariable = new SessionVariable(); + explicitSizeVariable.setFileSplitSize(MB); + TestIcebergScanNode firstExplicitSizeNode = new TestIcebergScanNode(explicitSizeVariable); + TestIcebergScanNode secondExplicitSizeNode = new TestIcebergScanNode(explicitSizeVariable); + setIcebergSource(firstExplicitSizeNode, mockIcebergSource(10L, 20L)); + setIcebergSource(secondExplicitSizeNode, mockIcebergSource(10L, 20L)); + try (CloseableIterable ignored = splitFiles( + firstExplicitSizeNode, mockTableScanWithPlanCounter(explicitSizePlanCalls))) { + // Explicit split size also preserves lazy Iceberg planning. + } + try (CloseableIterable ignored = splitFiles( + secondExplicitSizeNode, mockTableScanWithPlanCounter(explicitSizePlanCalls))) { + // Each explicitly sized relation must own its streaming iterable. + } + + Assert.assertEquals(2, batchPlanCalls.get()); + Assert.assertEquals(2, explicitSizePlanCalls.get()); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + private static List materializeTasks(TableScan scan) { + List tasks = new ArrayList<>(); + try (CloseableIterable plannedTasks = scan.planFiles()) { + plannedTasks.forEach(tasks::add); + } catch (Exception e) { + throw new RuntimeException(e); + } + return tasks; + } + + @Test + public void testIcebergScanTaskCacheSeparatesSnapshotSchemaAndPredicate() throws Exception { + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()); + TestIcebergScanNode otherCatalogNode = new TestIcebergScanNode(new SessionVariable()); + TestIcebergScanNode otherTableNode = new TestIcebergScanNode(new SessionVariable()); + setIcebergSource(node, mockIcebergSource(10L, 20L)); + setIcebergSource(otherCatalogNode, mockIcebergSource(11L, 20L)); + setIcebergSource(otherTableNode, mockIcebergSource(10L, 21L)); + AtomicInteger planCalls = new AtomicInteger(); + + node.getOrPlanFileScanTasks( + mockTableScan(30L, 40, Expressions.equal("id", 1)), + () -> plannedTask(planCalls)); + node.getOrPlanFileScanTasks( + mockTableScan(31L, 40, Expressions.equal("id", 1)), + () -> plannedTask(planCalls)); + node.getOrPlanFileScanTasks( + mockTableScan(30L, 41, Expressions.equal("id", 1)), + () -> plannedTask(planCalls)); + node.getOrPlanFileScanTasks( + mockTableScan(30L, 40, Expressions.equal("id", 2)), + () -> plannedTask(planCalls)); + otherCatalogNode.getOrPlanFileScanTasks( + mockTableScan(30L, 40, Expressions.equal("id", 1)), + () -> plannedTask(planCalls)); + otherTableNode.getOrPlanFileScanTasks( + mockTableScan(30L, 40, Expressions.equal("id", 1)), + () -> plannedTask(planCalls)); + node.getOrPlanFileScanTasks( + mockTableScan(30L, 40, Expressions.equal("id", 1), true), + () -> plannedTask(planCalls)); + + Assert.assertEquals(7, planCalls.get()); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + @Test + public void testPreparedExecutionResetClearsIcebergScanTaskCache() throws Exception { + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()); + setIcebergSource(node, mockIcebergSource(10L, 20L)); + AtomicInteger planCalls = new AtomicInteger(); + TableScan scan = mockTableScan(30L, 40, Expressions.equal("id", 1)); + + node.getOrPlanFileScanTasks(scan, () -> plannedTask(planCalls)); + statementContext.resetMvccSnapshots(); + node.getOrPlanFileScanTasks(scan, () -> plannedTask(planCalls)); + + Assert.assertEquals(2, planCalls.get()); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + @Test + public void testStatementCloseClearsIcebergScanTaskCache() throws Exception { + StatementContext statementContext = new StatementContext(); + ConnectContext context = new ConnectContext(); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + TestIcebergScanNode node = new TestIcebergScanNode(new SessionVariable()); + setIcebergSource(node, mockIcebergSource(10L, 20L)); + AtomicInteger planCalls = new AtomicInteger(); + TableScan scan = mockTableScan(30L, 40, Expressions.equal("id", 1)); + + node.getOrPlanFileScanTasks(scan, () -> plannedTask(planCalls)); + statementContext.close(); + node.getOrPlanFileScanTasks(scan, () -> plannedTask(planCalls)); + + Assert.assertEquals(2, planCalls.get()); + } finally { + statementContext.close(); + ConnectContext.remove(); + } + } + + private static IcebergSource mockIcebergSource(long catalogId, long tableId) { + IcebergExternalCatalog catalog = Mockito.mock(IcebergExternalCatalog.class); + Mockito.when(catalog.getId()).thenReturn(catalogId); + IcebergExternalTable table = Mockito.mock(IcebergExternalTable.class); + Mockito.when(table.getId()).thenReturn(tableId); + IcebergSource source = Mockito.mock(IcebergSource.class); + Mockito.when(source.getCatalog()).thenReturn(catalog); + Mockito.when(source.getTargetTable()).thenReturn(table); + return source; + } + + private static TableScan mockTableScan( + long snapshotId, int schemaId, org.apache.iceberg.expressions.Expression filter) { + return mockTableScan(snapshotId, schemaId, filter, false); + } + + private static TableScan mockTableScan( + long snapshotId, int schemaId, org.apache.iceberg.expressions.Expression filter, + boolean caseSensitive) { + Snapshot snapshot = Mockito.mock(Snapshot.class); + Mockito.when(snapshot.snapshotId()).thenReturn(snapshotId); + Schema schema = new Schema(schemaId, + ImmutableList.of(Types.NestedField.optional(1, "id", Types.IntegerType.get()))); + TableScan scan = Mockito.mock(TableScan.class); + Mockito.when(scan.snapshot()).thenReturn(snapshot); + Mockito.when(scan.schema()).thenReturn(schema); + Mockito.when(scan.filter()).thenReturn(filter); + Mockito.when(scan.isCaseSensitive()).thenReturn(caseSensitive); + return scan; + } + + private static TableScan mockTableScanWithPlanCounter(AtomicInteger planCalls) { + TableScan scan = mockTableScan(30L, 40, Expressions.equal("id", 1)); + Mockito.when(scan.planFiles()).thenAnswer(invocation -> { + planCalls.incrementAndGet(); + return CloseableIterable.withNoopClose(Collections.emptyList()); + }); + return scan; + } + + private static List plannedTask(AtomicInteger planCalls) { + planCalls.incrementAndGet(); + return Collections.singletonList(Mockito.mock(FileScanTask.class)); + } + @Test public void testInitialDefaultMetadataUsesCurrentSchemaForOrdinaryScan() throws Exception { Schema snapshotSchema = new Schema(Types.NestedField.optional("historical_binary") diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/paimon/source/PaimonScanNodeTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/paimon/source/PaimonScanNodeTest.java index c933c627de66c2..cb5bdc90986fd5 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/paimon/source/PaimonScanNodeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/paimon/source/PaimonScanNodeTest.java @@ -24,6 +24,7 @@ import org.apache.doris.analysis.TupleId; import org.apache.doris.catalog.ArrayType; import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.DatabaseIf; import org.apache.doris.catalog.MapType; import org.apache.doris.catalog.StructField; import org.apache.doris.catalog.StructType; @@ -49,8 +50,10 @@ import org.apache.doris.datasource.paimon.PaimonUtils; import org.apache.doris.datasource.property.metastore.MetastoreProperties; import org.apache.doris.datasource.property.metastore.PaimonJdbcMetaStoreProperties; +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.thrift.TFileRangeDesc; import org.apache.doris.thrift.TFileScanRangeParams; @@ -65,6 +68,8 @@ import org.apache.paimon.fs.Path; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.stats.SimpleStats; import org.apache.paimon.table.AppendOnlyFileStoreTable; @@ -103,6 +108,7 @@ import java.util.Map; import java.util.Optional; import java.util.OptionalLong; +import java.util.concurrent.atomic.AtomicInteger; @RunWith(MockitoJUnitRunner.class) public class PaimonScanNodeTest { @@ -137,6 +143,120 @@ private void assertVariantProjectionRequiresVariantV2(Type variantType) throws U PaimonScanNode.checkVariantV2Enabled(desc, true); } + @Test + public void testStatementCacheIncludesSnapshotOptionsPredicateAndProjection() throws Exception { + ConnectContext previousContext = ConnectContext.get(); + ConnectContext context = new ConnectContext(); + StatementContext statementContext = new StatementContext(context, null); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + PaimonExternalCatalog catalog = Mockito.mock(PaimonExternalCatalog.class); + PaimonExternalTable relationTable = Mockito.mock(PaimonExternalTable.class); + PaimonExternalTable targetTable = Mockito.mock(PaimonExternalTable.class); + Mockito.when(catalog.getId()).thenReturn(7L); + Mockito.when(relationTable.getId()).thenReturn(11L); + Mockito.when(targetTable.getId()).thenReturn(13L); + + RowType rowType = new RowType(Arrays.asList( + new DataField(0, "id", DataTypes.INT()), + new DataField(1, "value", DataTypes.INT()))); + PredicateBuilder predicateBuilder = new PredicateBuilder(rowType); + List idEqualsOne = Collections.singletonList(predicateBuilder.equal(0, 1)); + List idEqualsTwo = Collections.singletonList(predicateBuilder.equal(0, 2)); + AtomicInteger planCount = new AtomicInteger(); + + Table baselineTable = mockPlanningTable(rowType, Collections.emptyMap(), planCount); + assertPlanCount(newCachedPlanningNode( + 0, relationTable, targetTable, catalog, baselineTable, + 101L, 3L, Collections.emptyMap(), idEqualsOne, "id"), planCount, 1); + assertPlanCount(newCachedPlanningNode( + 1, relationTable, targetTable, catalog, baselineTable, + 101L, 3L, Collections.emptyMap(), idEqualsOne, "id"), planCount, 1); + + assertPlanCount(newCachedPlanningNode( + 2, relationTable, targetTable, catalog, + mockPlanningTable(rowType, Collections.emptyMap(), planCount), + 102L, 3L, Collections.emptyMap(), idEqualsOne, "id"), planCount, 2); + assertPlanCount(newCachedPlanningNode( + 3, relationTable, targetTable, catalog, + mockPlanningTable(rowType, Collections.emptyMap(), planCount), + 101L, 3L, ImmutableMap.of("scan.mode", "delta"), idEqualsOne, "id"), planCount, 3); + assertPlanCount(newCachedPlanningNode( + 4, relationTable, targetTable, catalog, + mockPlanningTable(rowType, Collections.emptyMap(), planCount), + 101L, 3L, Collections.emptyMap(), idEqualsTwo, "id"), planCount, 4); + assertPlanCount(newCachedPlanningNode( + 5, relationTable, targetTable, catalog, + mockPlanningTable(rowType, Collections.emptyMap(), planCount), + 101L, 3L, Collections.emptyMap(), idEqualsOne, "value"), planCount, 5); + assertPlanCount(newCachedPlanningNode( + 6, relationTable, targetTable, catalog, + mockPlanningTable(rowType, ImmutableMap.of("bucket", "2"), planCount), + 101L, 3L, Collections.emptyMap(), idEqualsOne, "id"), planCount, 6); + } finally { + statementContext.close(); + ConnectContext.remove(); + if (previousContext != null) { + previousContext.setThreadLocalInfo(); + } + } + } + + @Test + public void testPinnedFileCreationCacheIgnoresUnusedProjection() throws Exception { + ConnectContext previousContext = ConnectContext.get(); + ConnectContext context = new ConnectContext(); + StatementContext statementContext = new StatementContext(context, null); + context.setStatementContext(statementContext); + context.setThreadLocalInfo(); + try { + PaimonExternalCatalog catalog = Mockito.mock(PaimonExternalCatalog.class); + PaimonExternalTable relationTable = Mockito.mock(PaimonExternalTable.class); + PaimonExternalTable targetTable = Mockito.mock(PaimonExternalTable.class); + Mockito.when(catalog.getId()).thenReturn(17L); + Mockito.when(relationTable.getId()).thenReturn(19L); + Mockito.when(targetTable.getId()).thenReturn(23L); + + FileStoreTable table = Mockito.mock(FileStoreTable.class); + CoreOptions coreOptions = Mockito.mock(CoreOptions.class); + SnapshotReader reader = Mockito.mock(SnapshotReader.class); + SnapshotReader.Plan plan = Mockito.mock(SnapshotReader.Plan.class); + AtomicInteger planCount = new AtomicInteger(); + Mockito.when(table.options()).thenReturn(ImmutableMap.of("scan.snapshot-id", "29")); + Mockito.when(table.primaryKeys()).thenReturn(Collections.emptyList()); + Mockito.when(table.coreOptions()).thenReturn(coreOptions); + Mockito.when(coreOptions.bucket()).thenReturn(1); + Mockito.when(table.newSnapshotReader()).thenReturn(reader); + Mockito.when(reader.withMode(ScanMode.ALL)).thenReturn(reader); + Mockito.when(reader.withSnapshot(29L)).thenReturn(reader); + Mockito.when(reader.withManifestEntryFilter(ArgumentMatchers.any())).thenReturn(reader); + Mockito.when(reader.read()).thenReturn(plan); + Mockito.when(plan.splits()).thenAnswer(invocation -> { + planCount.incrementAndGet(); + return Collections.emptyList(); + }); + Snapshot latestSnapshot = Mockito.mock(Snapshot.class); + Mockito.when(latestSnapshot.id()).thenReturn(29L); + Mockito.when(table.latestSnapshot()).thenReturn(Optional.of(latestSnapshot)); + Map options = PaimonScanParams.resolveOptions( + table, ImmutableMap.of("scan.file-creation-time-millis", "1234")); + + assertPlanCount(newCachedPlanningNode( + 7, relationTable, targetTable, catalog, table, + 29L, 4L, options, Collections.emptyList(), "projected_a"), planCount, 1); + assertPlanCount(newCachedPlanningNode( + 8, relationTable, targetTable, catalog, table, + 29L, 4L, options, Collections.emptyList(), "projected_b"), planCount, 1); + } finally { + statementContext.close(); + ConnectContext.remove(); + if (previousContext != null) { + previousContext.setThreadLocalInfo(); + } + } + } + @Test public void testSerializedTableCacheKeyIsStablePerScanNode() { PaimonScanNode first = newTestNode(new PlanNodeId(0), new TupleId(0), sv); @@ -725,6 +845,7 @@ public void testSystemTablePassesIncrementalOptionsToPaimonTable() throws Except public void testPinnedFileCreationScanPreservesBatchReaderFilters() throws Exception { PaimonScanNode node = newTestNode(new PlanNodeId(0), new TupleId(0), sv); PaimonSource source = Mockito.mock(PaimonSource.class); + PaimonExternalCatalog catalog = Mockito.mock(PaimonExternalCatalog.class); PaimonExternalTable externalTable = Mockito.mock(PaimonExternalTable.class); FileStoreTable table = Mockito.mock(FileStoreTable.class); Snapshot snapshot = Mockito.mock(Snapshot.class); @@ -734,7 +855,9 @@ public void testPinnedFileCreationScanPreservesBatchReaderFilters() throws Excep org.apache.paimon.options.Options configuration = new org.apache.paimon.options.Options(); configuration.set(CoreOptions.BATCH_SCAN_MODE, CoreOptions.BatchScanMode.NONE); + Mockito.when(source.getCatalog()).thenReturn(catalog); Mockito.when(source.getExternalTable()).thenReturn(externalTable); + Mockito.when(source.getTargetTable()).thenReturn(externalTable); Mockito.when(source.getPaimonTable()).thenReturn(table); Mockito.when(source.getPaimonTable(ArgumentMatchers.any(TableScanParams.class))).thenReturn(table); Mockito.when(snapshot.id()).thenReturn(23L); @@ -790,6 +913,7 @@ public void testBoundEmptySnapshotDoesNotReadLaterCommit() throws Exception { public void testBoundEmptyDataSnapshotStillPlansMetadataSystemTable() throws Exception { PaimonScanNode node = newTestNode(new PlanNodeId(0), new TupleId(0), sv); PaimonSource source = Mockito.mock(PaimonSource.class); + PaimonExternalCatalog catalog = Mockito.mock(PaimonExternalCatalog.class); PaimonSysExternalTable systemTable = Mockito.mock(PaimonSysExternalTable.class); Table paimonTable = Mockito.mock(Table.class); ReadBuilder readBuilder = Mockito.mock(ReadBuilder.class); @@ -798,7 +922,9 @@ public void testBoundEmptyDataSnapshotStillPlansMetadataSystemTable() throws Exc org.apache.paimon.table.source.Split schemaSplit = Mockito.mock(org.apache.paimon.table.source.Split.class); + Mockito.when(source.getCatalog()).thenReturn(catalog); Mockito.when(source.getExternalTable()).thenReturn(systemTable); + Mockito.when(source.getTargetTable()).thenReturn(systemTable); Mockito.when(source.getPaimonTable()).thenReturn(paimonTable); Mockito.when(systemTable.getSysTableType()).thenReturn("schemas"); Mockito.when(source.getPaimonTable((TableScanParams) null)).thenReturn(paimonTable); @@ -1424,6 +1550,68 @@ private void mockJniReader(PaimonScanNode spyNode) { Mockito.doReturn(false).when(spyNode).supportNativeReader(ArgumentMatchers.any(Optional.class)); } + private Table mockPlanningTable(RowType rowType, Map tableOptions, + AtomicInteger planCount) { + Table table = Mockito.mock(Table.class); + ReadBuilder readBuilder = Mockito.mock(ReadBuilder.class); + TableScan scan = Mockito.mock(TableScan.class); + TableScan.Plan plan = Mockito.mock(TableScan.Plan.class); + Mockito.when(table.rowType()).thenReturn(rowType); + Mockito.when(table.options()).thenReturn(tableOptions); + Mockito.when(table.newReadBuilder()).thenReturn(readBuilder); + Mockito.when(readBuilder.withFilter(ArgumentMatchers.anyList())).thenReturn(readBuilder); + Mockito.when(readBuilder.withProjection(ArgumentMatchers.any(int[].class))).thenReturn(readBuilder); + Mockito.when(readBuilder.newScan()).thenReturn(scan); + Mockito.when(scan.plan()).thenReturn(plan); + Mockito.when(plan.splits()).thenAnswer(invocation -> { + planCount.incrementAndGet(); + return Collections.emptyList(); + }); + return table; + } + + private PaimonScanNode newCachedPlanningNode(int id, PaimonExternalTable relationTable, + PaimonExternalTable targetTable, PaimonExternalCatalog catalog, Table paimonTable, + long snapshotId, long schemaId, Map resolvedOptions, + List predicates, String projectedColumn) throws Exception { + DatabaseIf database = Mockito.mock(DatabaseIf.class); + Mockito.when(database.getCatalog()).thenReturn(catalog); + Mockito.when(database.getFullName()).thenReturn("db"); + Mockito.when(relationTable.getDatabase()).thenReturn(database); + Mockito.when(relationTable.getName()).thenReturn("table"); + TupleDescriptor desc = new TupleDescriptor(new TupleId(id)); + desc.setTable(relationTable); + SlotDescriptor slot = new SlotDescriptor(new SlotId(id), desc); + slot.setColumn(new Column(projectedColumn, Type.INT)); + desc.addSlot(slot); + + PaimonScanNode node = + new PaimonScanNode(new PlanNodeId(id), desc, false, sv, ScanContext.EMPTY); + PaimonSource source = Mockito.mock(PaimonSource.class); + Mockito.when(source.getCatalog()).thenReturn(catalog); + Mockito.when(source.getExternalTable()).thenReturn(relationTable); + Mockito.when(source.getTargetTable()).thenReturn(targetTable); + node.setSource(source); + setField(PaimonScanNode.class, node, "processedTable", paimonTable); + setField(PaimonScanNode.class, node, "predicates", predicates); + node.setRelationSnapshot(Optional.of(new PaimonMvccSnapshot( + new PaimonSnapshotCacheValue(PaimonPartitionInfo.EMPTY, + new PaimonSnapshot(snapshotId, schemaId, paimonTable))))); + if (!resolvedOptions.isEmpty()) { + TableScanParams scanParams = new TableScanParams( + TableScanParams.OPTIONS, resolvedOptions, Collections.emptyList()); + scanParams.reuseResolvedMapParams(resolvedOptions); + node.setScanParams(scanParams); + } + return node; + } + + private void assertPlanCount(PaimonScanNode node, AtomicInteger planCount, int expected) + throws UserException { + Assert.assertTrue(node.getPaimonSplitFromAPI().isEmpty()); + Assert.assertEquals(expected, planCount.get()); + } + private PaimonScanNode newTestNode(PlanNodeId planNodeId, TupleId tupleId, SessionVariable sessionVariable) { TupleDescriptor desc = new TupleDescriptor(tupleId); PaimonExternalTable externalTable = Mockito.mock(PaimonExternalTable.class); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/StatementContextTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/StatementContextTest.java index 9499c188c25360..491e978cf550e0 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/StatementContextTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/StatementContextTest.java @@ -22,6 +22,7 @@ import org.apache.doris.catalog.DatabaseIf; import org.apache.doris.catalog.TableIf; import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.ExternalTable; import org.apache.doris.datasource.hive.HMSExternalTable; import org.apache.doris.datasource.hive.HMSExternalTable.DLAType; @@ -42,7 +43,14 @@ import org.mockito.Mockito; import java.util.Collections; +import java.util.List; import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; public class StatementContextTest { @@ -709,6 +717,209 @@ public void testInjectedSnapshotRemainsAuthoritativeForLatestRelation() { } } + @Test + public void testExternalScanTasksUseSingleFlight() throws Exception { + StatementContext statementContext = new StatementContext(); + StatementContext.ExternalScanTaskCache cache = statementContext.getExternalScanTaskCache(); + ExecutorService executor = Executors.newFixedThreadPool(2); + CountDownLatch loaderStarted = new CountDownLatch(1); + CountDownLatch waiterLookedUpKey = new CountDownLatch(1); + CountDownLatch releaseLoader = new CountDownLatch(1); + AtomicInteger loadCount = new AtomicInteger(); + AtomicInteger hashCalls = new AtomicInteger(); + ExternalScanTaskCacheKey observedKey = + new ObservableScanTaskCacheKey("same-scan", hashCalls, waiterLookedUpKey); + try { + Future> first = executor.submit( + () -> cache.getOrLoad(observedKey, () -> { + loadCount.incrementAndGet(); + loaderStarted.countDown(); + releaseLoader.await(); + return Collections.singletonList("task"); + })); + loaderStarted.await(); + Future> second = executor.submit( + () -> cache.getOrLoad(observedKey, () -> { + loadCount.incrementAndGet(); + return Collections.singletonList("duplicate"); + })); + + waiterLookedUpKey.await(); + releaseLoader.countDown(); + + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("task"), first.get()); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("task"), second.get()); + org.junit.jupiter.api.Assertions.assertEquals(1, loadCount.get()); + } finally { + executor.shutdownNow(); + statementContext.close(); + } + } + + @Test + public void testExternalScanTaskFailureCanRetry() throws Exception { + StatementContext statementContext = new StatementContext(); + StatementContext.ExternalScanTaskCache cache = statementContext.getExternalScanTaskCache(); + CountDownLatch loaderStarted = new CountDownLatch(1); + CountDownLatch waiterLookedUpKey = new CountDownLatch(1); + CountDownLatch releaseLoader = new CountDownLatch(1); + AtomicInteger hashCalls = new AtomicInteger(); + ExternalScanTaskCacheKey key = + new ObservableScanTaskCacheKey("retry", hashCalls, waiterLookedUpKey); + AtomicInteger loadCount = new AtomicInteger(); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future> owner = executor.submit( + () -> cache.getOrLoad(key, () -> { + loadCount.incrementAndGet(); + loaderStarted.countDown(); + releaseLoader.await(); + throw new IllegalStateException("load failed"); + })); + loaderStarted.await(); + Future> waiter = executor.submit( + () -> cache.getOrLoad(key, + () -> Collections.singletonList("duplicate"))); + waiterLookedUpKey.await(); + releaseLoader.countDown(); + + assertFutureFailedWith(owner, IllegalStateException.class, "load failed"); + assertFutureFailedWith(waiter, IllegalStateException.class, "load failed"); + List tasks = cache.getOrLoad(key, () -> { + loadCount.incrementAndGet(); + return Collections.singletonList("retry-task"); + }); + + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("retry-task"), tasks); + org.junit.jupiter.api.Assertions.assertEquals(2, loadCount.get()); + } finally { + executor.shutdownNow(); + statementContext.close(); + } + } + + @Test + public void testExternalScanTaskGenerationIsIsolatedByResetAndExecutionEnd() throws Exception { + StatementContext statementContext = new StatementContext(); + CountDownLatch loaderStarted = new CountDownLatch(1); + CountDownLatch waiterLookedUpKey = new CountDownLatch(1); + CountDownLatch releaseLoader = new CountDownLatch(1); + AtomicInteger hashCalls = new AtomicInteger(); + ExternalScanTaskCacheKey key = + new ObservableScanTaskCacheKey("lifecycle", hashCalls, waiterLookedUpKey); + AtomicInteger loadCount = new AtomicInteger(); + ExecutorService executor = Executors.newFixedThreadPool(2); + StatementContext.ExternalScanTaskCache oldGeneration = + statementContext.getExternalScanTaskCache(); + try { + Future> oldOwner = executor.submit( + () -> oldGeneration.getOrLoad(key, () -> { + int loadNumber = loadCount.incrementAndGet(); + loaderStarted.countDown(); + releaseLoader.await(); + return Collections.singletonList("old-" + loadNumber); + })); + loaderStarted.await(); + Future> oldWaiter = executor.submit( + () -> oldGeneration.getOrLoad(key, + () -> Collections.singletonList("duplicate"))); + waiterLookedUpKey.await(); + + statementContext.resetMvccSnapshots(); + StatementContext.ExternalScanTaskCache newGeneration = + statementContext.getExternalScanTaskCache(); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("new-2"), + newGeneration.getOrLoad(key, + () -> Collections.singletonList("new-" + loadCount.incrementAndGet()))); + releaseLoader.countDown(); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("old-1"), oldOwner.get()); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("old-1"), oldWaiter.get()); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("new-2"), + newGeneration.getOrLoad(key, + () -> Collections.singletonList("duplicate"))); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("old-3"), + oldGeneration.getOrLoad(key, + () -> Collections.singletonList("old-" + loadCount.incrementAndGet()))); + + statementContext.clearExternalScanTasks(); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("ended-4"), + newGeneration.getOrLoad(key, + () -> Collections.singletonList("ended-" + loadCount.incrementAndGet()))); + org.junit.jupiter.api.Assertions.assertEquals( + Collections.singletonList("ended-5"), + newGeneration.getOrLoad(key, + () -> Collections.singletonList("ended-" + loadCount.incrementAndGet()))); + + org.junit.jupiter.api.Assertions.assertEquals(5, loadCount.get()); + } finally { + executor.shutdownNow(); + statementContext.close(); + } + } + + private static void assertFutureFailedWith( + Future future, Class causeType, String message) { + ExecutionException exception = org.junit.jupiter.api.Assertions.assertThrows( + ExecutionException.class, future::get); + org.junit.jupiter.api.Assertions.assertInstanceOf(causeType, exception.getCause()); + org.junit.jupiter.api.Assertions.assertEquals(message, exception.getCause().getMessage()); + } + + private static final class TestScanTaskCacheKey implements ExternalScanTaskCacheKey { + private final String value; + + private TestScanTaskCacheKey(String value) { + this.value = value; + } + + @Override + public boolean equals(Object object) { + return object instanceof TestScanTaskCacheKey + && value.equals(((TestScanTaskCacheKey) object).value); + } + + @Override + public int hashCode() { + return value.hashCode(); + } + } + + private static final class ObservableScanTaskCacheKey implements ExternalScanTaskCacheKey { + private final String value; + private final AtomicInteger hashCalls; + private final CountDownLatch secondLookup; + + private ObservableScanTaskCacheKey( + String value, AtomicInteger hashCalls, CountDownLatch secondLookup) { + this.value = value; + this.hashCalls = hashCalls; + this.secondLookup = secondLookup; + } + + @Override + public boolean equals(Object object) { + return object instanceof ObservableScanTaskCacheKey + && value.equals(((ObservableScanTaskCacheKey) object).value); + } + + @Override + public int hashCode() { + if (hashCalls.incrementAndGet() == 2) { + secondLookup.countDown(); + } + return value.hashCode(); + } + } + @SuppressWarnings("unchecked") private DatabaseIf mockDatabase() { return Mockito.mock(DatabaseIf.class); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java index 5e1d2506b3bb11..6f32d39884e4da 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java @@ -21,6 +21,7 @@ import org.apache.doris.catalog.DatabaseIf; import org.apache.doris.catalog.TableIf; import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.mvcc.MvccSnapshot; import org.apache.doris.datasource.mvcc.MvccTable; import org.apache.doris.nereids.StatementContext; @@ -174,6 +175,66 @@ public void testMvccSnapshotsAreResetForEveryExecute() throws Exception { Mockito.verify(table, Mockito.times(2)).loadSnapshot(Optional.empty(), Optional.empty()); } + @Test + public void testExternalScanTasksUseANewGenerationForEveryExecute() throws Exception { + String sql = "select 1"; + LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + StatementContext statementContext = new StatementContext(); + PrepareCommand prepareCommand = new PrepareCommand( + "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); + PreparedStatementContext preparedStatement = new PreparedStatementContext( + prepareCommand, connectContext, statementContext, "stmt"); + StmtExecutor executor = Mockito.mock(StmtExecutor.class); + Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); + Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); + Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); + Mockito.when(executor.getContext()).thenReturn(connectContext); + ExternalScanTaskCacheKey key = new PreparedScanTaskCacheKey("same-scan"); + AtomicInteger loadCount = new AtomicInteger(); + + StatementContext.ExternalScanTaskCache preparedGeneration = + statementContext.getExternalScanTaskCache(); + preparedGeneration.getOrLoad(key, + () -> Collections.singletonList("prepared-" + loadCount.incrementAndGet())); + + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + StatementContext.ExternalScanTaskCache firstExecuteGeneration = + statementContext.getExternalScanTaskCache(); + Assertions.assertNotSame(preparedGeneration, firstExecuteGeneration); + Assertions.assertEquals(Collections.singletonList("execute-2"), + firstExecuteGeneration.getOrLoad(key, + () -> Collections.singletonList("execute-" + loadCount.incrementAndGet()))); + + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + StatementContext.ExternalScanTaskCache secondExecuteGeneration = + statementContext.getExternalScanTaskCache(); + Assertions.assertNotSame(firstExecuteGeneration, secondExecuteGeneration); + Assertions.assertEquals(Collections.singletonList("execute-3"), + secondExecuteGeneration.getOrLoad(key, + () -> Collections.singletonList("execute-" + loadCount.incrementAndGet()))); + Assertions.assertEquals(3, loadCount.get()); + } + + private static final class PreparedScanTaskCacheKey implements ExternalScanTaskCacheKey { + private final String value; + + private PreparedScanTaskCacheKey(String value) { + this.value = value; + } + + @Override + public boolean equals(Object object) { + return object instanceof PreparedScanTaskCacheKey + && value.equals(((PreparedScanTaskCacheKey) object).value); + } + + @Override + public int hashCode() { + return value.hashCode(); + } + } + private String resolveNextSnapshot(TableScanParams scanParams, AtomicInteger snapshotId) { return scanParams.getOrResolveMapParams(ignored -> ImmutableMap.of( "scan.snapshot-id", String.valueOf(snapshotId.incrementAndGet())))