This is an automated email from the ASF dual-hosted git repository. Wei-hao-Li pushed a commit to branch deviceEntrySpill-dev2 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit d05e9e2e6748211cfe9d1958778d9eda818d8dd9 Author: Weihao Li <[email protected]> AuthorDate: Fri Aug 14 09:54:08 2026 +0800 module-2 Signed-off-by: Weihao Li <[email protected]> --- .../spill/AbstractDeviceEntryMaterializer.java | 4 +- .../metadata/spill/DeviceEntryDataSet.java | 2 +- .../metadata/spill/DeviceEntryDataSetHandle.java | 107 +++ ...DeviceEntryMaterializationMemoryController.java | 79 ++ .../metadata/spill/DeviceEntryMaterializer.java | 2 +- .../spill/DeviceEntrySortedMaterializer.java | 251 ++++++ .../metadata/spill/DeviceEntrySpillManager.java | 3 +- .../metadata/spill/InMemoryDeviceEntryDataSet.java | 2 +- .../metadata/spill/SpilledDeviceEntryDataSet.java | 20 +- .../distribute/TableDistributedPlanGenerator.java | 915 +++++++++++++++++++-- .../planner/node/DeviceTableScanNode.java | 82 +- .../spill/DeviceEntryMaterializerTest.java | 53 +- 12 files changed, 1407 insertions(+), 113 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/AbstractDeviceEntryMaterializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/AbstractDeviceEntryMaterializer.java index 76aaa903675..92886b0036a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/AbstractDeviceEntryMaterializer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/AbstractDeviceEntryMaterializer.java @@ -38,7 +38,7 @@ public abstract class AbstractDeviceEntryMaterializer implements AutoCloseable { private final long thresholdInBytes; private final List<DeviceEntry> bufferedEntries = new ArrayList<>(); - private long entryCount; + private int entryCount; private Path ownerDirectory; private boolean ownerRegistered; private boolean finished; @@ -112,7 +112,7 @@ public abstract class AbstractDeviceEntryMaterializer implements AutoCloseable { return ioContext; } - protected final long entryCount() { + protected final int entryCount() { return entryCount; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSet.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSet.java index 92a9caf6654..9642acdb004 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSet.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSet.java @@ -26,7 +26,7 @@ import java.util.List; public interface DeviceEntryDataSet extends AutoCloseable { - long getEntryCount(); + int getEntryCount(); boolean isSpilled(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSetHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSetHandle.java new file mode 100644 index 00000000000..d731aaeb778 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryDataSetHandle.java @@ -0,0 +1,107 @@ +/* + * 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.iotdb.db.queryengine.plan.relational.metadata.spill; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.utils.ThriftCommonsSerDeUtils; + +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; + +public final class DeviceEntryDataSetHandle { + + private final String queryId; + private final PlanNodeId planNodeId; + private final TEndPoint coordinatorEndPoint; + private final int segmentCount; + private final int entryCount; + private final boolean ordered; + + public DeviceEntryDataSetHandle( + String queryId, + PlanNodeId planNodeId, + TEndPoint coordinatorEndPoint, + int segmentCount, + int entryCount, + boolean ordered) { + this.queryId = queryId; + this.planNodeId = planNodeId; + this.coordinatorEndPoint = coordinatorEndPoint; + this.segmentCount = segmentCount; + this.entryCount = entryCount; + this.ordered = ordered; + } + + public String getQueryId() { + return queryId; + } + + public PlanNodeId getPlanNodeId() { + return planNodeId; + } + + public TEndPoint getCoordinatorEndPoint() { + return coordinatorEndPoint; + } + + public int getSegmentCount() { + return segmentCount; + } + + public int getEntryCount() { + return entryCount; + } + + public boolean isOrdered() { + return ordered; + } + + public void serialize(ByteBuffer byteBuffer) { + ReadWriteIOUtils.write(queryId, byteBuffer); + ReadWriteIOUtils.write(planNodeId.getId(), byteBuffer); + ThriftCommonsSerDeUtils.serializeTEndPoint(coordinatorEndPoint, byteBuffer); + ReadWriteIOUtils.write(segmentCount, byteBuffer); + ReadWriteIOUtils.write(entryCount, byteBuffer); + ReadWriteIOUtils.write(ordered, byteBuffer); + } + + public void serialize(DataOutputStream stream) throws IOException { + ReadWriteIOUtils.write(queryId, stream); + ReadWriteIOUtils.write(planNodeId.getId(), stream); + ThriftCommonsSerDeUtils.serializeTEndPoint(coordinatorEndPoint, stream); + ReadWriteIOUtils.write(segmentCount, stream); + ReadWriteIOUtils.write(entryCount, stream); + ReadWriteIOUtils.write(ordered, stream); + } + + public static DeviceEntryDataSetHandle deserialize(ByteBuffer byteBuffer) { + return new DeviceEntryDataSetHandle( + ReadWriteIOUtils.readString(byteBuffer), + new PlanNodeId(ReadWriteIOUtils.readString(byteBuffer)), + ThriftCommonsSerDeUtils.deserializeTEndPoint(byteBuffer), + ReadWriteIOUtils.readInt(byteBuffer), + ReadWriteIOUtils.readLong(byteBuffer), + ReadWriteIOUtils.readBool(byteBuffer)); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializationMemoryController.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializationMemoryController.java new file mode 100644 index 00000000000..ef991aadfa1 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializationMemoryController.java @@ -0,0 +1,79 @@ +/* + * 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.iotdb.db.queryengine.plan.relational.metadata.spill; + +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import java.io.IOException; +import java.util.IdentityHashMap; +import java.util.Map; + +/** Controls the total in-memory DeviceEntry buffers owned by all Region materializers. */ +public final class DeviceEntryMaterializationMemoryController { + + private final long memoryLimitInBytes; + private final Map<AbstractDeviceEntryMaterializer, Long> retainedBytesByMaterializer = + new IdentityHashMap<>(); + private long retainedBytes; + + public DeviceEntryMaterializationMemoryController(long memoryLimitInBytes) { + if (memoryLimitInBytes <= 0) { + throw new IllegalArgumentException(); + } + this.memoryLimitInBytes = memoryLimitInBytes; + } + + public void append(AbstractDeviceEntryMaterializer materializer, DeviceEntry deviceEntry) + throws IOException { + materializer.append(deviceEntry); + long entryRamBytes = deviceEntry.ramBytesUsed(); + retainedBytesByMaterializer.merge(materializer, entryRamBytes, Long::sum); + retainedBytes += entryRamBytes; + enforceMemoryLimit(); + } + + public long getRetainedBytes() { + return retainedBytes; + } + + public long getMemoryLimitInBytes() { + return memoryLimitInBytes; + } + + private void enforceMemoryLimit() throws IOException { + while (retainedBytes > memoryLimitInBytes) { + AbstractDeviceEntryMaterializer largest = null; + long largestRetainedBytes = 0; + for (Map.Entry<AbstractDeviceEntryMaterializer, Long> entry : + retainedBytesByMaterializer.entrySet()) { + if (entry.getValue() > largestRetainedBytes) { + largest = entry.getKey(); + largestRetainedBytes = entry.getValue(); + } + } + if (largest == null) { + throw new IllegalStateException(); + } + largest.forceSpill(); + retainedBytesByMaterializer.put(largest, 0L); + retainedBytes -= largestRetainedBytes; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializer.java index 0a660369501..82f53bb2caf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializer.java @@ -94,7 +94,7 @@ public final class DeviceEntryMaterializer extends AbstractDeviceEntryMaterializ } else { dataSet = new SpilledDeviceEntryDataSet( - queryId(), ownerDirectory(), spiller.finish(), entryCount(), true); + queryId(), ownerDirectory(), spiller.finish(), entryCount()); } if (rawSegment && getQueryContext() != null) { getQueryContext().recordDeviceEntryCount(entryCount()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySortedMaterializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySortedMaterializer.java new file mode 100644 index 00000000000..d135139f700 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySortedMaterializer.java @@ -0,0 +1,251 @@ +/* + * 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.iotdb.db.queryengine.plan.relational.metadata.spill; + +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; + +import org.apache.tsfile.external.commons.io.FileUtils; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.PriorityQueue; + +/** Materializes a sorted data set in memory or through sorted runs and a K-way merge. */ +public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMaterializer { + + private static final int MAX_MERGE_FAN_IN = 32; + + private final Comparator<DeviceEntry> comparator; + private final List<List<Path>> sortedRuns = new ArrayList<>(); + + private Path runDirectory; + + public DeviceEntrySortedMaterializer( + String queryId, + PlanNodeId planNodeId, + long bufferSizeInBytes, + Comparator<DeviceEntry> comparator) { + super(queryId, planNodeId, bufferSizeInBytes); + this.comparator = comparator; + } + + public DeviceEntrySortedMaterializer( + String queryId, + PlanNodeId planNodeId, + long bufferSizeInBytes, + Comparator<DeviceEntry> comparator, + MPPQueryContext queryContext) { + this(queryId, planNodeId, bufferSizeInBytes, comparator); + setQueryContext(queryContext); + } + + @Override + public void append(DeviceEntry entry) throws IOException { + checkNotFinished(); + appendToBuffer(entry); + } + + @Override + public void forceSpill() throws IOException { + checkNotFinished(); + flushRun(); + } + + @Override + public DeviceEntryDataSet finish() throws IOException { + checkNotFinished(); + if (entryCount() == 0) { + DeviceEntryDataSet dataSet = new InMemoryDeviceEntryDataSet(copyBufferedEntries()); + markFinished(); + return dataSet; + } + if (sortedRuns.isEmpty()) { + sortBufferedEntries(comparator); + DeviceEntryDataSet dataSet = new InMemoryDeviceEntryDataSet(copyBufferedEntries()); + markFinished(); + return dataSet; + } + + try { + flushRun(); + List<List<Path>> finalRuns = compactRuns(new ArrayList<>(sortedRuns)); + Path finalDirectory = ownerDirectory().resolve("fi"); + List<Path> finalSegments; + try (DeviceEntryDiskSpiller outputSpiller = + new DeviceEntryDiskSpiller(finalDirectory, thresholdInBytes(), ioContext())) { + if (finalRuns.size() == 1) { + copyRun(finalRuns.get(0), outputSpiller); + } else { + mergeRuns(finalRuns, outputSpiller); + } + finalSegments = outputSpiller.finish(); + } + DeviceEntryDataSet dataSet = + new SpilledDeviceEntryDataSet(queryId(), ownerDirectory(), finalSegments, entryCount()); + markFinished(); + deleteRunDirectoryBestEffort(); + return dataSet; + } catch (IOException | RuntimeException e) { + try { + cleanupOwnerDirectory(); + } catch (IOException cleanupException) { + e.addSuppressed(cleanupException); + } + throw e; + } + } + + private void flushRun() throws IOException { + if (isBufferEmpty()) { + return; + } + ensureSpillDirectory(); + sortBufferedEntries(comparator); + Path currentRunDirectory = runDirectory.resolve(String.format("run-%06d", sortedRuns.size())); + try (DeviceEntryDiskSpiller runSpiller = + new DeviceEntryDiskSpiller(currentRunDirectory, thresholdInBytes(), ioContext())) { + for (DeviceEntry entry : bufferedEntries()) { + runSpiller.append(entry.serializeToBytes()); + } + sortedRuns.add(runSpiller.finish()); + } + clearBuffer(); + } + + private void ensureSpillDirectory() throws IOException { + if (ownerDirectory() != null) { + return; + } + createIOContextOnSpill(false); + runDirectory = ensureOwnerDirectory().resolve("sort-run"); + } + + private void copyRun(List<Path> run, DeviceEntryDiskSpiller outputSpiller) throws IOException { + try (DeviceEntryFileSpillerReader reader = + new DeviceEntryFileSpillerReader(run, true, ioContext())) { + while (reader.hasNext()) { + outputSpiller.append(reader.next().serializeToBytes()); + } + } + } + + private void deleteRunDirectoryBestEffort() { + try { + FileUtils.deleteDirectory(runDirectory.toFile()); + } catch (IOException ignored) { + // Query cleanup removes the published data set and any remaining runs. + } + } + + private List<List<Path>> compactRuns(List<List<Path>> runs) throws IOException { + int level = 1; + while (runs.size() > MAX_MERGE_FAN_IN) { + List<List<Path>> nextRuns = new ArrayList<>(); + for (int from = 0, group = 0; from < runs.size(); from += MAX_MERGE_FAN_IN, group++) { + int to = Math.min(from + MAX_MERGE_FAN_IN, runs.size()); + List<List<Path>> runGroup = new ArrayList<>(runs.subList(from, to)); + if (runGroup.size() == 1) { + nextRuns.add(runGroup.get(0)); + continue; + } + Path outputDirectory = + runDirectory + .resolve(String.format("level-%06d", level)) + .resolve(String.format("run-%06d", group)); + try (DeviceEntryDiskSpiller outputSpiller = + new DeviceEntryDiskSpiller(outputDirectory, thresholdInBytes(), ioContext())) { + mergeRuns(runGroup, outputSpiller); + nextRuns.add(outputSpiller.finish()); + } + } + runs = nextRuns; + level++; + } + return runs; + } + + private void mergeRuns(List<List<Path>> runs, DeviceEntryDiskSpiller outputSpiller) + throws IOException { + List<DeviceEntryFileSpillerReader> readers = new ArrayList<>(runs.size()); + PriorityQueue<MergeElement> queue = + new PriorityQueue<>( + (left, right) -> { + int result = comparator.compare(left.entry, right.entry); + return result != 0 ? result : Integer.compare(left.readerIndex, right.readerIndex); + }); + Throwable failure = null; + try { + for (int i = 0; i < runs.size(); i++) { + DeviceEntryFileSpillerReader reader = + new DeviceEntryFileSpillerReader(runs.get(i), true, ioContext()); + readers.add(reader); + if (reader.hasNext()) { + queue.add(new MergeElement(reader.next(), i)); + } + } + while (!queue.isEmpty()) { + MergeElement element = queue.poll(); + outputSpiller.append(element.entry.serializeToBytes()); + DeviceEntryFileSpillerReader reader = readers.get(element.readerIndex); + if (reader.hasNext()) { + queue.add(new MergeElement(reader.next(), element.readerIndex)); + } + } + } catch (IOException | RuntimeException | Error e) { + failure = e; + throw e; + } finally { + IOException closeException = null; + for (DeviceEntryFileSpillerReader reader : readers) { + try { + reader.close(); + } catch (IOException e) { + if (closeException == null) { + closeException = e; + } else { + closeException.addSuppressed(e); + } + } + } + if (closeException != null) { + if (failure != null) { + failure.addSuppressed(closeException); + } else { + throw closeException; + } + } + } + } + + private static final class MergeElement { + private final DeviceEntry entry; + private final int readerIndex; + + private MergeElement(DeviceEntry entry, int readerIndex) { + this.entry = entry; + this.readerIndex = readerIndex; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java index 65ea3fd517f..5580c0886b2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntrySpillManager.java @@ -32,6 +32,7 @@ import java.util.Comparator; import java.util.List; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; public final class DeviceEntrySpillManager { @@ -77,7 +78,7 @@ public final class DeviceEntrySpillManager { .sorted( Comparator.comparingInt((Path path) -> path.getFileName().toString().length()) .thenComparing(path -> path.getFileName().toString())) - .toList(); + .collect(Collectors.toList()); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/InMemoryDeviceEntryDataSet.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/InMemoryDeviceEntryDataSet.java index 831cbe351f3..7c5564fcbf1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/InMemoryDeviceEntryDataSet.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/InMemoryDeviceEntryDataSet.java @@ -34,7 +34,7 @@ public final class InMemoryDeviceEntryDataSet implements DeviceEntryDataSet { } @Override - public long getEntryCount() { + public int getEntryCount() { return entries.size(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SpilledDeviceEntryDataSet.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SpilledDeviceEntryDataSet.java index 0380f210f34..99803c5c0fb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SpilledDeviceEntryDataSet.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/SpilledDeviceEntryDataSet.java @@ -19,8 +19,6 @@ package org.apache.iotdb.db.queryengine.plan.relational.metadata.spill; -import org.apache.tsfile.external.commons.io.FileUtils; - import java.io.IOException; import java.nio.file.Path; import java.util.List; @@ -30,24 +28,18 @@ public final class SpilledDeviceEntryDataSet implements DeviceEntryDataSet { private final String queryId; private final Path ownerDirectory; private final List<Path> segments; - private final long entryCount; - private final boolean managedBySpillManager; + private final int entryCount; public SpilledDeviceEntryDataSet( - String queryId, - Path ownerDirectory, - List<Path> segments, - long entryCount, - boolean managedBySpillManager) { + String queryId, Path ownerDirectory, List<Path> segments, int entryCount) { this.queryId = queryId; this.ownerDirectory = ownerDirectory; this.segments = segments; this.entryCount = entryCount; - this.managedBySpillManager = managedBySpillManager; } @Override - public long getEntryCount() { + public int getEntryCount() { return entryCount; } @@ -76,10 +68,6 @@ public final class SpilledDeviceEntryDataSet implements DeviceEntryDataSet { @Override public void close() throws IOException { - if (managedBySpillManager) { - DeviceEntrySpillManager.getInstance().deregisterOwner(queryId, ownerDirectory); - } else { - FileUtils.deleteDirectory(ownerDirectory.toFile()); - } + DeviceEntrySpillManager.getInstance().deregisterOwner(queryId, ownerDirectory); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java index 25845758a08..7e4024efe87 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java @@ -90,6 +90,13 @@ import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis; import org.apache.iotdb.db.queryengine.plan.relational.function.tvf.read_tsfile.ExternalTsFileQueryResource.DeviceTaskPartition; import org.apache.iotdb.db.queryengine.plan.relational.metadata.AlignedDeviceEntry; import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryDataSet; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryDataSetHandle; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryMaterializationMemoryController; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryMaterializer; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryReader; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntrySortedMaterializer; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.SpilledDeviceEntryDataSet; import org.apache.iotdb.db.queryengine.plan.relational.planner.SymbolAllocator; import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode; import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTreeDeviceViewScanNode; @@ -129,6 +136,8 @@ import org.apache.tsfile.utils.Pair; import javax.annotation.Nonnull; +import java.io.IOException; +import java.io.UncheckedIOException; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -147,10 +156,6 @@ import java.util.stream.IntStream; import static com.google.common.collect.ImmutableList.toImmutableList; import static org.apache.iotdb.calc.utils.constant.SqlConstant.COUNT; -import static org.apache.iotdb.calc.utils.constant.SqlConstant.DELTA; -import static org.apache.iotdb.calc.utils.constant.SqlConstant.INCREASE; -import static org.apache.iotdb.calc.utils.constant.SqlConstant.IRATE; -import static org.apache.iotdb.calc.utils.constant.SqlConstant.RATE; import static org.apache.iotdb.commons.partition.DataPartition.NOT_ASSIGNED; import static org.apache.iotdb.commons.queryengine.plan.relational.function.FunctionKind.AGGREGATE; import static org.apache.iotdb.commons.queryengine.plan.relational.metadata.FunctionNullability.getAggregationFunctionNullability; @@ -975,6 +980,11 @@ public class TableDistributedPlanGenerator String.format(DataNodeQueryMessages.GIVEN_QUERIED_DATABASE_S_IS_NOT_EXIST, dbName)); } + if (node.getCoordinatorDeviceEntryDataSet().isSpilled()) { + return constructSpilledDeviceTableScanByRegionReplicaSet( + node, context, dataPartition, seriesSlotMap); + } + final Map<TRegionReplicaSet, DeviceTableScanNode> tableScanNodeMap = new HashMap<>(); Map<Integer, List<TRegionReplicaSet>> cachedSeriesSlotWithRegions = new HashMap<>(); @@ -1048,6 +1058,197 @@ public class TableDistributedPlanGenerator return resultTableScanNodeList; } + private List<PlanNode> constructSpilledDeviceTableScanByRegionReplicaSet( + DeviceTableScanNode node, + PlanContext context, + DataPartition dataPartition, + Map<TSeriesPartitionSlot, Map<TTimePartitionSlot, List<TRegionReplicaSet>>> seriesSlotMap) { + Optional<SortPropertyContext> sortPropertyContext = + context.hasSortProperty ? analyzeSortProperty(node, context) : Optional.empty(); + Comparator<DeviceEntry> comparator = + sortPropertyContext.map(property -> property.comparator).orElse(null); + long batchSize = + IoTDBDescriptor.getInstance().getConfig().getTableQueryDeviceEntryBatchSizeInBytes(); + Map<TRegionReplicaSet, DeviceTableScanNode> scanNodes = new HashMap<>(); + Map<TRegionReplicaSet, DeviceEntryMaterializer> materializers = new HashMap<>(); + Map<TRegionReplicaSet, DeviceEntrySortedMaterializer> sortedMaterializers = new HashMap<>(); + Map<TRegionReplicaSet, Integer> regionEntryCounts = new HashMap<>(); + Map<Integer, List<TRegionReplicaSet>> cachedSeriesSlotWithRegions = new HashMap<>(); + DeviceEntryMaterializationMemoryController memoryController = + new DeviceEntryMaterializationMemoryController(batchSize); + + try (DeviceEntryReader reader = node.getCoordinatorDeviceEntryDataSet().openConsumingReader()) { + while (reader.hasNext()) { + DeviceEntry deviceEntry = reader.next(); + List<TRegionReplicaSet> regionReplicaSets = + getDeviceReplicaSets( + dataPartition, + seriesSlotMap, + deviceEntry.getDeviceID(), + node.getTimeFilter(), + cachedSeriesSlotWithRegions); + if (regionReplicaSets.size() > 1) { + context.deviceCrossRegion = true; + } + for (TRegionReplicaSet regionReplicaSet : regionReplicaSets) { + DeviceTableScanNode scanNode = + scanNodes.computeIfAbsent( + regionReplicaSet, + ignored -> createRegionDeviceTableScanNode(node, regionReplicaSet)); + PlanNodeId ownerId = scanNode.getPlanNodeId(); + if (comparator == null) { + DeviceEntryMaterializer materializer = + materializers.computeIfAbsent( + regionReplicaSet, + ignored -> + new DeviceEntryMaterializer( + queryId.getId(), ownerId, batchSize, false, queryContext)); + memoryController.append(materializer, deviceEntry); + } else { + DeviceEntrySortedMaterializer sortedMaterializer = + sortedMaterializers.get(regionReplicaSet); + if (sortedMaterializer == null) { + sortedMaterializer = + new DeviceEntrySortedMaterializer( + queryId.getId(), ownerId, batchSize, comparator, queryContext); + sortedMaterializers.put(regionReplicaSet, sortedMaterializer); + } + memoryController.append(sortedMaterializer, deviceEntry); + } + regionEntryCounts.merge(regionReplicaSet, 1, Integer::sum); + } + } + + for (Map.Entry<TRegionReplicaSet, DeviceTableScanNode> entry : scanNodes.entrySet()) { + DeviceEntryDataSet dataSet = + comparator == null + ? materializers.get(entry.getKey()).finish() + : sortedMaterializers.get(entry.getKey()).finish(); + installDataSet(entry.getValue(), dataSet, comparator != null); + } + } catch (IOException e) { + closeSpillWriters(materializers.values(), sortedMaterializers.values()); + throw new UncheckedIOException(e); + } + + if (scanNodes.isEmpty()) { + node.setRegionReplicaSet(NOT_ASSIGNED); + return Collections.singletonList(node); + } + + List<PlanNode> result = new ArrayList<>(); + TRegionReplicaSet mostUsedRegion = null; + int maxEntryCount = -1; + for (Map.Entry<TRegionReplicaSet, DeviceTableScanNode> entry : + topology.filterReachableCandidates(scanNodes.entrySet())) { + result.add(entry.getValue()); + int entryCount = regionEntryCounts.getOrDefault(entry.getKey(), 0); + if (entryCount > maxEntryCount) { + mostUsedRegion = entry.getKey(); + maxEntryCount = entryCount; + } + } + if (mostUsedRegion == null) { + throw new RootFIPlacementException(scanNodes.keySet()); + } + context.mostUsedRegion = mostUsedRegion; + sortPropertyContext.ifPresent(property -> applySortProperty(node, result, property, false)); + return result; + } + + private DeviceTableScanNode createRegionDeviceTableScanNode( + DeviceTableScanNode node, TRegionReplicaSet regionReplicaSet) { + DeviceTableScanNode scanNode = + new DeviceTableScanNode( + queryId.genPlanNodeId(), + node.getQualifiedObjectName(), + node.getOutputSymbols(), + node.getAssignments(), + new ArrayList<>(), + node.getTagAndAttributeIndexMap(), + node.getScanOrder(), + node.getTimePredicate().orElse(null), + node.getPushDownPredicate(), + node.getPushDownLimit(), + node.getPushDownOffset(), + node.isPushLimitToEachDevice(), + node.containsNonAlignedDevice()); + scanNode.setRegionReplicaSet(regionReplicaSet); + return scanNode; + } + + private void installDataSet( + DeviceTableScanNode scanNode, DeviceEntryDataSet dataSet, boolean ordered) { + scanNode.setCoordinatorDeviceEntryDataSet(dataSet); + if (!dataSet.isSpilled()) { + return; + } + SpilledDeviceEntryDataSet spilled = (SpilledDeviceEntryDataSet) dataSet; + scanNode.setDeviceEntryDataSetHandle( + new DeviceEntryDataSetHandle( + queryId.getId(), + scanNode.getPlanNodeId(), + DataNodeEndPoints.getLocalDataNodeLocation().getInternalEndPoint(), + spilled.getSegments().size(), + spilled.getEntryCount(), + ordered)); + } + + private void closeSpillWriters( + Collection<DeviceEntryMaterializer> materializers, + Collection<DeviceEntrySortedMaterializer> sortedMaterializers) { + for (DeviceEntryMaterializer writer : materializers) { + try { + writer.close(); + } catch (Exception ignored) { + // The original planning exception is more useful than a cleanup failure. + } + } + for (DeviceEntrySortedMaterializer writer : sortedMaterializers) { + try { + writer.close(); + } catch (Exception ignored) { + // The original planning exception is more useful than a cleanup failure. + } + } + } + + private DeviceEntryDataSet finishRegionStagingDataSet( + DeviceEntryDataSet stagingDataSet, + PlanNodeId ownerId, + long batchSize, + Comparator<DeviceEntry> comparator) + throws IOException { + if (comparator == null) { + return stagingDataSet; + } + + try (DeviceEntrySortedMaterializer sortedMaterializer = + new DeviceEntrySortedMaterializer( + queryId.getId(), ownerId, batchSize, comparator, queryContext); + DeviceEntryReader reader = + stagingDataSet.isSpilled() + ? stagingDataSet.openConsumingReader() + : stagingDataSet.openReader()) { + DeviceEntryMaterializationMemoryController memoryController = + new DeviceEntryMaterializationMemoryController(batchSize); + while (reader.hasNext()) { + memoryController.append(sortedMaterializer, reader.next()); + } + return sortedMaterializer.finish(); + } + } + + private void closeDeviceEntryDataSets(Collection<DeviceEntryDataSet> dataSets) { + for (DeviceEntryDataSet dataSet : dataSets) { + try { + dataSet.close(); + } catch (Exception ignored) { + // The original planning exception is more useful than a cleanup failure. + } + } + } + @Override public List<PlanNode> visitTreeDeviceViewScan(TreeDeviceViewScanNode node, PlanContext context) { DataPartition dataPartition = analysis.getDataPartitionInfo(); @@ -1066,6 +1267,11 @@ public class TableDistributedPlanGenerator String.format(DataNodeQueryMessages.GIVEN_QUERIED_DATABASE_S_IS_NOT_EXIST, dbName)); } + if (node.getCoordinatorDeviceEntryDataSet().isSpilled()) { + return constructSpilledTreeDeviceViewScanByRegionReplicaSet( + node, context, dataPartition, seriesSlotMap); + } + Map<TRegionReplicaSet, Pair<TreeAlignedDeviceViewScanNode, TreeNonAlignedDeviceViewScanNode>> tableScanNodeMap = new HashMap<>(); Map<Integer, List<TRegionReplicaSet>> cachedSeriesSlotWithRegions = new HashMap<>(); @@ -1105,7 +1311,6 @@ public class TableDistributedPlanGenerator node.getTreeDBName(), node.getMeasurementColumnNameMap()); scanNode.setRegionReplicaSet(regionReplicaSet); - scanNode.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); pair.left = scanNode; } @@ -1128,7 +1333,6 @@ public class TableDistributedPlanGenerator node.getTreeDBName(), node.getMeasurementColumnNameMap()); scanNode.setRegionReplicaSet(regionReplicaSet); - scanNode.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); pair.right = scanNode; } @@ -1186,6 +1390,201 @@ public class TableDistributedPlanGenerator return resultTableScanNodeList; } + private List<PlanNode> constructSpilledTreeDeviceViewScanByRegionReplicaSet( + TreeDeviceViewScanNode node, + PlanContext context, + DataPartition dataPartition, + Map<TSeriesPartitionSlot, Map<TTimePartitionSlot, List<TRegionReplicaSet>>> seriesSlotMap) { + Optional<SortPropertyContext> sortPropertyContext = + context.hasSortProperty ? analyzeSortProperty(node, context) : Optional.empty(); + Comparator<DeviceEntry> comparator = + sortPropertyContext.map(property -> property.comparator).orElse(null); + long batchSize = + IoTDBDescriptor.getInstance().getConfig().getTableQueryDeviceEntryBatchSizeInBytes(); + Map<TRegionReplicaSet, Pair<TreeAlignedDeviceViewScanNode, TreeNonAlignedDeviceViewScanNode>> + scanNodes = new HashMap<>(); + Map<DeviceTableScanNode, DeviceEntryMaterializer> materializers = new HashMap<>(); + Map<DeviceTableScanNode, DeviceEntrySortedMaterializer> sortedMaterializers = new HashMap<>(); + Map<TRegionReplicaSet, Integer> regionEntryCounts = new HashMap<>(); + Map<Integer, List<TRegionReplicaSet>> cachedSeriesSlotWithRegions = new HashMap<>(); + DeviceEntryMaterializationMemoryController memoryController = + new DeviceEntryMaterializationMemoryController(batchSize); + + try (DeviceEntryReader reader = node.getCoordinatorDeviceEntryDataSet().openConsumingReader()) { + while (reader.hasNext()) { + DeviceEntry deviceEntry = reader.next(); + List<TRegionReplicaSet> regionReplicaSets = + getDeviceReplicaSets( + dataPartition, + seriesSlotMap, + deviceEntry.getDeviceID(), + node.getTimeFilter(), + cachedSeriesSlotWithRegions); + if (regionReplicaSets.size() > 1) { + context.deviceCrossRegion = true; + } + boolean aligned = deviceEntry instanceof AlignedDeviceEntry; + for (TRegionReplicaSet regionReplicaSet : regionReplicaSets) { + Pair<TreeAlignedDeviceViewScanNode, TreeNonAlignedDeviceViewScanNode> pair = + scanNodes.computeIfAbsent(regionReplicaSet, ignored -> new Pair<>(null, null)); + DeviceTableScanNode scanNode; + if (aligned) { + if (pair.left == null) { + pair.left = createTreeAlignedScanNode(node, regionReplicaSet); + } + scanNode = pair.left; + } else { + if (pair.right == null) { + pair.right = createTreeNonAlignedScanNode(node, regionReplicaSet); + } + scanNode = pair.right; + } + appendToRegionDataSet( + scanNode, + deviceEntry, + comparator, + batchSize, + materializers, + sortedMaterializers, + memoryController); + regionEntryCounts.merge(regionReplicaSet, 1, Integer::sum); + } + } + + for (Pair<TreeAlignedDeviceViewScanNode, TreeNonAlignedDeviceViewScanNode> pair : + scanNodes.values()) { + if (pair.left != null) { + finishAndInstallDataSet(pair.left, comparator, materializers, sortedMaterializers); + } + if (pair.right != null) { + finishAndInstallDataSet(pair.right, comparator, materializers, sortedMaterializers); + } + } + } catch (IOException e) { + closeSpillWriters(materializers.values(), sortedMaterializers.values()); + throw new UncheckedIOException(e); + } + + if (scanNodes.isEmpty()) { + node.setRegionReplicaSet(NOT_ASSIGNED); + node.setTreeDBName(null); + return Collections.singletonList(node); + } + + List<PlanNode> result = new ArrayList<>(); + TRegionReplicaSet mostUsedRegion = null; + int maxEntryCount = -1; + for (Map.Entry< + TRegionReplicaSet, + Pair<TreeAlignedDeviceViewScanNode, TreeNonAlignedDeviceViewScanNode>> + entry : topology.filterReachableCandidates(scanNodes.entrySet())) { + if (entry.getValue().left != null) { + result.add(entry.getValue().left); + } + if (entry.getValue().right != null) { + result.add(entry.getValue().right); + } + int entryCount = regionEntryCounts.getOrDefault(entry.getKey(), 0); + if (entryCount > maxEntryCount) { + mostUsedRegion = entry.getKey(); + maxEntryCount = entryCount; + } + } + if (mostUsedRegion == null) { + throw new RootFIPlacementException(scanNodes.keySet()); + } + context.mostUsedRegion = mostUsedRegion; + sortPropertyContext.ifPresent(property -> applySortProperty(node, result, property, false)); + return result; + } + + private TreeAlignedDeviceViewScanNode createTreeAlignedScanNode( + TreeDeviceViewScanNode node, TRegionReplicaSet regionReplicaSet) { + TreeAlignedDeviceViewScanNode scanNode = + new TreeAlignedDeviceViewScanNode( + queryId.genPlanNodeId(), + node.getQualifiedObjectName(), + node.getOutputSymbols(), + node.getAssignments(), + new ArrayList<>(), + node.getTagAndAttributeIndexMap(), + node.getScanOrder(), + node.getTimePredicate().orElse(null), + node.getPushDownPredicate(), + node.getPushDownLimit(), + node.getPushDownOffset(), + node.isPushLimitToEachDevice(), + node.containsNonAlignedDevice(), + node.getTreeDBName(), + node.getMeasurementColumnNameMap()); + scanNode.setRegionReplicaSet(regionReplicaSet); + return scanNode; + } + + private TreeNonAlignedDeviceViewScanNode createTreeNonAlignedScanNode( + TreeDeviceViewScanNode node, TRegionReplicaSet regionReplicaSet) { + TreeNonAlignedDeviceViewScanNode scanNode = + new TreeNonAlignedDeviceViewScanNode( + queryId.genPlanNodeId(), + node.getQualifiedObjectName(), + node.getOutputSymbols(), + node.getAssignments(), + new ArrayList<>(), + node.getTagAndAttributeIndexMap(), + node.getScanOrder(), + node.getTimePredicate().orElse(null), + node.getPushDownPredicate(), + node.getPushDownLimit(), + node.getPushDownOffset(), + node.isPushLimitToEachDevice(), + node.containsNonAlignedDevice(), + node.getTreeDBName(), + node.getMeasurementColumnNameMap()); + scanNode.setRegionReplicaSet(regionReplicaSet); + return scanNode; + } + + private void appendToRegionDataSet( + DeviceTableScanNode scanNode, + DeviceEntry deviceEntry, + Comparator<DeviceEntry> comparator, + long batchSize, + Map<DeviceTableScanNode, DeviceEntryMaterializer> materializers, + Map<DeviceTableScanNode, DeviceEntrySortedMaterializer> sortedMaterializers, + DeviceEntryMaterializationMemoryController memoryController) + throws IOException { + if (comparator == null) { + materializers.computeIfAbsent( + scanNode, + ignored -> + new DeviceEntryMaterializer( + queryId.getId(), scanNode.getPlanNodeId(), batchSize, false, queryContext)); + memoryController.append(materializers.get(scanNode), deviceEntry); + return; + } + DeviceEntrySortedMaterializer sortedMaterializer = sortedMaterializers.get(scanNode); + if (sortedMaterializer == null) { + sortedMaterializer = + new DeviceEntrySortedMaterializer( + queryId.getId(), scanNode.getPlanNodeId(), batchSize, comparator, queryContext); + sortedMaterializers.put(scanNode, sortedMaterializer); + } + memoryController.append(sortedMaterializer, deviceEntry); + } + + private void finishAndInstallDataSet( + DeviceTableScanNode scanNode, + Comparator<DeviceEntry> comparator, + Map<DeviceTableScanNode, DeviceEntryMaterializer> materializers, + Map<DeviceTableScanNode, DeviceEntrySortedMaterializer> sortedMaterializers) + throws IOException { + DeviceEntryDataSet dataSet = + comparator == null + ? materializers.get(scanNode).finish() + : sortedMaterializers.get(scanNode).finish(); + installDataSet(scanNode, dataSet, comparator != null); + } + @Override public List<PlanNode> visitInformationSchemaTableScan( InformationSchemaTableScanNode node, PlanContext context) { @@ -1324,10 +1723,9 @@ public class TableDistributedPlanGenerator // push down aggregation if the child of aggregation node only has the union Node if (childrenNodes.size() == 1) { node.setChild(childrenNodes.get(0)); - AggregationNode physicalAggregation = withRateFunctionInputOrdering(node, childOrdering); if (childrenNodes.get(0) instanceof UnionNode - && physicalAggregation.getAggregations().values().stream() + && node.getAggregations().values().stream() .noneMatch(aggregation -> aggregation.isDistinct() || aggregation.hasMask())) { UnionNode unionNode = (UnionNode) childrenNodes.get(0); List<PlanNode> children = unionNode.getChildren(); @@ -1348,8 +1746,7 @@ public class TableDistributedPlanGenerator } // 2. split the aggregation into partial and final - Pair<AggregationNode, AggregationNode> splitResult = - split(physicalAggregation, symbolAllocator, queryId); + Pair<AggregationNode, AggregationNode> splitResult = split(node, symbolAllocator, queryId); AggregationNode intermediate = splitResult.right; // 3. add the aggregation node above the project node @@ -1368,7 +1765,7 @@ public class TableDistributedPlanGenerator intermediate.getStep(), intermediate.getHashSymbol(), intermediate.getGroupIdSymbol()); - if (physicalAggregation.isStreamable() && childOrdering != null) { + if (node.isStreamable() && childOrdering != null) { nodeOrderingMap.put(planNodeId, expectedOrderingSchema); } return aggregationNode; @@ -1384,7 +1781,7 @@ public class TableDistributedPlanGenerator return Collections.singletonList(splitResult.left); } - return Collections.singletonList(physicalAggregation); + return Collections.singletonList(node); } // We cannot do multi-stage Aggregate if any aggregation-function is distinct. @@ -1392,12 +1789,10 @@ public class TableDistributedPlanGenerator // MarkDistinctNode will merge all data from different child. if (node.getAggregations().values().stream() .anyMatch(aggregation -> aggregation.isDistinct() || aggregation.hasMask())) { - PlanNode physicalChild = + node.setChild( mergeChildrenViaCollectOrMergeSort( - nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes); - node.setChild(physicalChild); - return Collections.singletonList( - withRateFunctionInputOrdering(node, nodeOrderingMap.get(physicalChild.getPlanNodeId()))); + nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes)); + return Collections.singletonList(node); } Pair<AggregationNode, AggregationNode> splitResult = split(node, symbolAllocator, queryId); AggregationNode intermediate = splitResult.right; @@ -1429,58 +1824,6 @@ public class TableDistributedPlanGenerator return Collections.singletonList(splitResult.left); } - private static AggregationNode withRateFunctionInputOrdering( - AggregationNode node, OrderingScheme childOrdering) { - Map<Symbol, AggregationNode.Aggregation> aggregations = new LinkedHashMap<>(); - node.getAggregations() - .forEach( - (symbol, aggregation) -> - aggregations.put( - symbol, - new AggregationNode.Aggregation( - aggregation.getResolvedFunction(), - aggregation.getArguments(), - aggregation.isDistinct(), - aggregation.getFilter(), - aggregation.getOrderingScheme(), - aggregation.getMask(), - isInputOrderedByTimeAscending( - aggregation, node.getStep(), node.getGroupingKeys(), childOrdering)))); - return AggregationNode.builderFrom(node).setAggregations(aggregations).build(); - } - - static boolean isInputOrderedByTimeAscending( - AggregationNode.Aggregation aggregation, - AggregationNode.Step step, - List<Symbol> groupingKeys, - OrderingScheme childOrdering) { - String functionName = aggregation.getResolvedFunction().getSignature().getName(); - if (step != SINGLE - || childOrdering == null - || aggregation.getArguments().size() < 2 - || !(RATE.equalsIgnoreCase(functionName) - || INCREASE.equalsIgnoreCase(functionName) - || IRATE.equalsIgnoreCase(functionName) - || DELTA.equalsIgnoreCase(functionName))) { - return false; - } - - Symbol timeSymbol = Symbol.from(aggregation.getArguments().get(1)); - List<Symbol> orderBy = childOrdering.getOrderBy(); - int timeIndex = orderBy.indexOf(timeSymbol); - if (timeIndex < 0 || !childOrdering.getOrdering(timeSymbol).isAscending()) { - return false; - } - - Set<Symbol> groupingKeySet = new HashSet<>(groupingKeys); - for (int i = 0; i < timeIndex; i++) { - if (!groupingKeySet.contains(orderBy.get(i))) { - return false; - } - } - return true; - } - private boolean prefixMatched(OrderingScheme childOrdering, List<Symbol> preGroupedSymbols) { List<Symbol> orderKeys = childOrdering.getOrderBy(); if (orderKeys.size() < preGroupedSymbols.size()) { @@ -1505,6 +1848,11 @@ public class TableDistributedPlanGenerator return Collections.singletonList(node); } + if (node.getCoordinatorDeviceEntryDataSet().isSpilled()) { + return constructSpilledAggregationTableScanByRegionReplicaSet( + node, context, dataPartition, dbName); + } + AggregationDistributionInfo distributionInfo = prepareAggregationDistribution(node, dbName, dataPartition, context); @@ -1554,6 +1902,185 @@ public class TableDistributedPlanGenerator return resultTableScanNodeList; } + private List<PlanNode> constructSpilledAggregationTableScanByRegionReplicaSet( + AggregationTableScanNode node, + PlanContext context, + DataPartition dataPartition, + String dbName) { + Map<TSeriesPartitionSlot, Map<TTimePartitionSlot, List<TRegionReplicaSet>>> seriesSlotMap = + dataPartition.getDataPartitionMap().get(dbName); + if (seriesSlotMap == null) { + throw new SemanticException( + String.format(DataNodeQueryMessages.GIVEN_QUERIED_DATABASE_S_IS_NOT_EXIST, dbName)); + } + + long batchSize = + IoTDBDescriptor.getInstance().getConfig().getTableQueryDeviceEntryBatchSizeInBytes(); + Map<Integer, List<TRegionReplicaSet>> cachedSeriesSlotWithRegions = new HashMap<>(); + Map<DeviceEntry, Integer> crossRegionDeviceCounts = + node.mayUseLastCache() ? new HashMap<>() : Collections.emptyMap(); + Map<TRegionReplicaSet, PlanNodeId> regionPlanNodeIds = new HashMap<>(); + Map<TRegionReplicaSet, DeviceEntryMaterializer> stagingMaterializers = new HashMap<>(); + Map<TRegionReplicaSet, DeviceEntryDataSet> stagingDataSets = new HashMap<>(); + Map<TRegionReplicaSet, Integer> regionEntryCounts = new HashMap<>(); + DeviceEntryMaterializationMemoryController memoryController = + new DeviceEntryMaterializationMemoryController(batchSize); + boolean hasCrossRegionDevice = false; + try (DeviceEntryReader reader = node.getCoordinatorDeviceEntryDataSet().openConsumingReader()) { + while (reader.hasNext()) { + DeviceEntry deviceEntry = reader.next(); + List<TRegionReplicaSet> regions = + getDeviceReplicaSets( + dataPartition, + seriesSlotMap, + deviceEntry.getDeviceID(), + node.getTimeFilter(), + cachedSeriesSlotWithRegions); + if (regions.size() > 1) { + hasCrossRegionDevice = true; + context.deviceCrossRegion = true; + if (node.mayUseLastCache()) { + crossRegionDeviceCounts.put(deviceEntry, regions.size()); + } + } + for (TRegionReplicaSet region : regions) { + PlanNodeId regionPlanNodeId = + regionPlanNodeIds.computeIfAbsent(region, ignored -> queryId.genPlanNodeId()); + stagingMaterializers.computeIfAbsent( + region, + ignored -> + new DeviceEntryMaterializer( + queryId.getId(), regionPlanNodeId, batchSize, false, queryContext)); + memoryController.append(stagingMaterializers.get(region), deviceEntry); + regionEntryCounts.merge(region, 1, Integer::sum); + } + } + for (Map.Entry<TRegionReplicaSet, DeviceEntryMaterializer> entry : + stagingMaterializers.entrySet()) { + stagingDataSets.put(entry.getKey(), entry.getValue().finish()); + } + } catch (IOException e) { + closeSpillWriters(stagingMaterializers.values(), Collections.emptyList()); + throw new UncheckedIOException(e); + } catch (RuntimeException e) { + closeSpillWriters(stagingMaterializers.values(), Collections.emptyList()); + throw e; + } + + boolean needSplit = hasCrossRegionDevice && node.getStep() == SINGLE; + AggregationTableScanNode templateNode = node; + AggregationNode finalAggregation = null; + try { + if (needSplit) { + Pair<AggregationNode, AggregationTableScanNode> splitResult = + split(node, symbolAllocator, queryId); + finalAggregation = splitResult.left; + templateNode = splitResult.right; + if (!context.hasSortProperty && finalAggregation.isStreamable()) { + context.setExpectedOrderingScheme(constructOrderingSchema(node.getPreGroupedSymbols())); + } + } + } catch (RuntimeException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw e; + } + if (hasCrossRegionDevice && node.mayUseLastCache()) { + queryContext.setNeedUpdateScanNumForLastQuery(true); + } + + Optional<SortPropertyContext> sortPropertyContext; + try { + sortPropertyContext = + context.hasSortProperty ? analyzeSortProperty(node, context) : Optional.empty(); + } catch (RuntimeException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw e; + } + Comparator<DeviceEntry> comparator = + sortPropertyContext.map(property -> property.comparator).orElse(null); + Map<TRegionReplicaSet, AggregationTableScanNode> scanNodes = new HashMap<>(); + try { + for (Map.Entry<TRegionReplicaSet, DeviceEntryDataSet> entry : stagingDataSets.entrySet()) { + TRegionReplicaSet region = entry.getKey(); + PlanNodeId regionPlanNodeId = regionPlanNodeIds.get(region); + AggregationTableScanNode scanNode = + createAggregationScanNode(templateNode, regionPlanNodeId, region); + DeviceEntryDataSet dataSet = + finishRegionStagingDataSet(entry.getValue(), regionPlanNodeId, batchSize, comparator); + installDataSet(scanNode, dataSet, comparator != null); + if (!crossRegionDeviceCounts.isEmpty()) { + scanNode.setDeviceCountMap(crossRegionDeviceCounts); + } + scanNodes.put(region, scanNode); + } + } catch (IOException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw new UncheckedIOException(e); + } catch (RuntimeException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw e; + } + + List<PlanNode> result = new ArrayList<>(); + TRegionReplicaSet mostUsedRegion = null; + int maxEntryCount = -1; + for (Map.Entry<TRegionReplicaSet, AggregationTableScanNode> entry : + topology.filterReachableCandidates(scanNodes.entrySet())) { + result.add(entry.getValue()); + int entryCount = regionEntryCounts.getOrDefault(entry.getKey(), 0); + if (entryCount > maxEntryCount) { + mostUsedRegion = entry.getKey(); + maxEntryCount = entryCount; + } + } + if (mostUsedRegion == null) { + throw new RootFIPlacementException(scanNodes.keySet()); + } + context.mostUsedRegion = mostUsedRegion; + sortPropertyContext.ifPresent(property -> applySortProperty(node, result, property, false)); + + if (needSplit) { + if (result.size() == 1) { + finalAggregation.setChild(result.get(0)); + } else { + finalAggregation.setChild( + mergeChildrenViaCollectOrMergeSort( + nodeOrderingMap.get(result.get(0).getPlanNodeId()), result)); + } + return Collections.singletonList(finalAggregation); + } + return result; + } + + private AggregationTableScanNode createAggregationScanNode( + AggregationTableScanNode template, + PlanNodeId planNodeId, + TRegionReplicaSet regionReplicaSet) { + AggregationTableScanNode scanNode = + new AggregationTableScanNode( + planNodeId, + template.getQualifiedObjectName(), + template.getOutputSymbols(), + template.getAssignments(), + new ArrayList<>(), + template.getTagAndAttributeIndexMap(), + template.getScanOrder(), + template.getTimePredicate().orElse(null), + template.getPushDownPredicate(), + template.getPushDownLimit(), + template.getPushDownOffset(), + template.isPushLimitToEachDevice(), + template.containsNonAlignedDevice(), + template.getProjection(), + template.getAggregations(), + template.getGroupingSets(), + template.getPreGroupedSymbols(), + template.getStep(), + template.getGroupIdSymbol()); + scanNode.setRegionReplicaSet(regionReplicaSet); + return scanNode; + } + @Override public List<PlanNode> visitAggregationTreeDeviceViewScan( AggregationTreeDeviceViewScanNode node, PlanContext context) { @@ -1586,6 +2113,11 @@ public class TableDistributedPlanGenerator node.getMeasurementColumnNameMap())); } + if (node.getCoordinatorDeviceEntryDataSet().isSpilled()) { + return constructSpilledAggregationTreeDeviceViewScanByRegionReplicaSet( + node, context, dataPartition, dbName); + } + AggregationDistributionInfo distributionInfo = prepareAggregationDistribution(node, dbName, dataPartition, context); @@ -1737,6 +2269,249 @@ public class TableDistributedPlanGenerator return resultTableScanNodeList; } + private List<PlanNode> constructSpilledAggregationTreeDeviceViewScanByRegionReplicaSet( + AggregationTreeDeviceViewScanNode node, + PlanContext context, + DataPartition dataPartition, + String dbName) { + Map<TSeriesPartitionSlot, Map<TTimePartitionSlot, List<TRegionReplicaSet>>> seriesSlotMap = + dataPartition.getDataPartitionMap().get(dbName); + if (seriesSlotMap == null) { + throw new SemanticException( + String.format(DataNodeQueryMessages.GIVEN_QUERIED_DATABASE_S_IS_NOT_EXIST, dbName)); + } + + long batchSize = + IoTDBDescriptor.getInstance().getConfig().getTableQueryDeviceEntryBatchSizeInBytes(); + boolean hasCrossRegionDevice = false; + Map<Integer, List<TRegionReplicaSet>> cachedSeriesSlotWithRegions = new HashMap<>(); + Map<TRegionReplicaSet, Pair<PlanNodeId, PlanNodeId>> regionPlanNodeIds = new HashMap<>(); + Map<PlanNodeId, DeviceEntryMaterializer> stagingMaterializers = new HashMap<>(); + Map<PlanNodeId, DeviceEntryDataSet> stagingDataSets = new HashMap<>(); + Map<TRegionReplicaSet, Integer> regionEntryCounts = new HashMap<>(); + DeviceEntryMaterializationMemoryController memoryController = + new DeviceEntryMaterializationMemoryController(batchSize); + try (DeviceEntryReader reader = node.getCoordinatorDeviceEntryDataSet().openConsumingReader()) { + while (reader.hasNext()) { + DeviceEntry deviceEntry = reader.next(); + List<TRegionReplicaSet> regions = + getDeviceReplicaSets( + dataPartition, + seriesSlotMap, + deviceEntry.getDeviceID(), + node.getTimeFilter(), + cachedSeriesSlotWithRegions); + if (regions.size() > 1) { + hasCrossRegionDevice = true; + context.deviceCrossRegion = true; + } + boolean aligned = deviceEntry instanceof AlignedDeviceEntry; + for (TRegionReplicaSet region : regions) { + Pair<PlanNodeId, PlanNodeId> planNodeIds = + regionPlanNodeIds.computeIfAbsent(region, ignored -> new Pair<>(null, null)); + PlanNodeId planNodeId; + if (aligned) { + if (planNodeIds.left == null) { + planNodeIds.left = queryId.genPlanNodeId(); + } + planNodeId = planNodeIds.left; + } else { + if (planNodeIds.right == null) { + planNodeIds.right = queryId.genPlanNodeId(); + } + planNodeId = planNodeIds.right; + } + stagingMaterializers.computeIfAbsent( + planNodeId, + ignored -> + new DeviceEntryMaterializer( + queryId.getId(), planNodeId, batchSize, false, queryContext)); + memoryController.append(stagingMaterializers.get(planNodeId), deviceEntry); + regionEntryCounts.merge(region, 1, Integer::sum); + } + } + for (Map.Entry<PlanNodeId, DeviceEntryMaterializer> entry : stagingMaterializers.entrySet()) { + stagingDataSets.put(entry.getKey(), entry.getValue().finish()); + } + } catch (IOException e) { + closeSpillWriters(stagingMaterializers.values(), Collections.emptyList()); + throw new UncheckedIOException(e); + } catch (RuntimeException e) { + closeSpillWriters(stagingMaterializers.values(), Collections.emptyList()); + throw e; + } + + boolean needSplit = hasCrossRegionDevice && node.getStep() == SINGLE; + AggregationTableScanNode templateNode = node; + AggregationNode finalAggregation = null; + try { + if (needSplit) { + Pair<AggregationNode, AggregationTableScanNode> splitResult = + split(node, symbolAllocator, queryId); + finalAggregation = splitResult.left; + templateNode = splitResult.right; + if (!context.hasSortProperty && finalAggregation.isStreamable()) { + context.setExpectedOrderingScheme(constructOrderingSchema(node.getPreGroupedSymbols())); + } + } + } catch (RuntimeException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw e; + } + + Optional<SortPropertyContext> sortPropertyContext; + try { + sortPropertyContext = + context.hasSortProperty ? analyzeSortProperty(node, context) : Optional.empty(); + } catch (RuntimeException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw e; + } + Comparator<DeviceEntry> comparator = + sortPropertyContext.map(property -> property.comparator).orElse(null); + Map< + TRegionReplicaSet, + Pair< + AlignedAggregationTreeDeviceViewScanNode, + NonAlignedAggregationTreeDeviceViewScanNode>> + scanNodes = new HashMap<>(); + try { + for (Map.Entry<TRegionReplicaSet, Pair<PlanNodeId, PlanNodeId>> entry : + regionPlanNodeIds.entrySet()) { + TRegionReplicaSet region = entry.getKey(); + Pair<AlignedAggregationTreeDeviceViewScanNode, NonAlignedAggregationTreeDeviceViewScanNode> + scanNodePair = new Pair<>(null, null); + if (entry.getValue().left != null) { + PlanNodeId planNodeId = entry.getValue().left; + scanNodePair.left = + createAlignedAggregationTreeScanNode(node, templateNode, planNodeId, region); + DeviceEntryDataSet dataSet = + finishRegionStagingDataSet( + stagingDataSets.get(planNodeId), planNodeId, batchSize, comparator); + installDataSet(scanNodePair.left, dataSet, comparator != null); + } + if (entry.getValue().right != null) { + PlanNodeId planNodeId = entry.getValue().right; + scanNodePair.right = + createNonAlignedAggregationTreeScanNode(node, templateNode, planNodeId, region); + DeviceEntryDataSet dataSet = + finishRegionStagingDataSet( + stagingDataSets.get(planNodeId), planNodeId, batchSize, comparator); + installDataSet(scanNodePair.right, dataSet, comparator != null); + } + scanNodes.put(region, scanNodePair); + } + } catch (IOException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw new UncheckedIOException(e); + } catch (RuntimeException e) { + closeDeviceEntryDataSets(stagingDataSets.values()); + throw e; + } + + List<PlanNode> result = new ArrayList<>(); + TRegionReplicaSet mostUsedRegion = null; + int maxEntryCount = -1; + for (Map.Entry< + TRegionReplicaSet, + Pair< + AlignedAggregationTreeDeviceViewScanNode, + NonAlignedAggregationTreeDeviceViewScanNode>> + entry : topology.filterReachableCandidates(scanNodes.entrySet())) { + if (entry.getValue().left != null) { + result.add(entry.getValue().left); + } + if (entry.getValue().right != null) { + result.add(entry.getValue().right); + } + int entryCount = regionEntryCounts.getOrDefault(entry.getKey(), 0); + if (entryCount > maxEntryCount) { + mostUsedRegion = entry.getKey(); + maxEntryCount = entryCount; + } + } + if (mostUsedRegion == null) { + throw new RootFIPlacementException(scanNodes.keySet()); + } + context.mostUsedRegion = mostUsedRegion; + sortPropertyContext.ifPresent(property -> applySortProperty(node, result, property, false)); + if (needSplit) { + if (result.size() == 1) { + finalAggregation.setChild(result.get(0)); + } else { + finalAggregation.setChild( + mergeChildrenViaCollectOrMergeSort( + nodeOrderingMap.get(result.get(0).getPlanNodeId()), result)); + } + return Collections.singletonList(finalAggregation); + } + return result; + } + + private AlignedAggregationTreeDeviceViewScanNode createAlignedAggregationTreeScanNode( + AggregationTreeDeviceViewScanNode source, + AggregationTableScanNode template, + PlanNodeId planNodeId, + TRegionReplicaSet region) { + AlignedAggregationTreeDeviceViewScanNode scanNode = + new AlignedAggregationTreeDeviceViewScanNode( + planNodeId, + template.getQualifiedObjectName(), + template.getOutputSymbols(), + template.getAssignments(), + new ArrayList<>(), + template.getTagAndAttributeIndexMap(), + template.getScanOrder(), + template.getTimePredicate().orElse(null), + template.getPushDownPredicate(), + template.getPushDownLimit(), + template.getPushDownOffset(), + template.isPushLimitToEachDevice(), + template.containsNonAlignedDevice(), + template.getProjection(), + template.getAggregations(), + template.getGroupingSets(), + template.getPreGroupedSymbols(), + template.getStep(), + template.getGroupIdSymbol(), + source.getTreeDBName(), + source.getMeasurementColumnNameMap()); + scanNode.setRegionReplicaSet(region); + return scanNode; + } + + private NonAlignedAggregationTreeDeviceViewScanNode createNonAlignedAggregationTreeScanNode( + AggregationTreeDeviceViewScanNode source, + AggregationTableScanNode template, + PlanNodeId planNodeId, + TRegionReplicaSet region) { + NonAlignedAggregationTreeDeviceViewScanNode scanNode = + new NonAlignedAggregationTreeDeviceViewScanNode( + planNodeId, + template.getQualifiedObjectName(), + template.getOutputSymbols(), + template.getAssignments(), + new ArrayList<>(), + template.getTagAndAttributeIndexMap(), + template.getScanOrder(), + template.getTimePredicate().orElse(null), + template.getPushDownPredicate(), + template.getPushDownLimit(), + template.getPushDownOffset(), + template.isPushLimitToEachDevice(), + template.containsNonAlignedDevice(), + template.getProjection(), + template.getAggregations(), + template.getGroupingSets(), + template.getPreGroupedSymbols(), + template.getStep(), + template.getGroupIdSymbol(), + source.getTreeDBName(), + source.getMeasurementColumnNameMap()); + scanNode.setRegionReplicaSet(region); + return scanNode; + } + private static class AggregationDistributionInfo { private final List<List<TRegionReplicaSet>> regionReplicaSetsList; private final AggregationTableScanNode templateNode; @@ -2162,7 +2937,7 @@ public class TableDistributedPlanGenerator sortPropertyContext.sortOrders, sortPropertyContext.lastIsTimeRelated, resultTableScanNodeList.size() == 1 - && ((DeviceTableScanNode) resultTableScanNodeList.get(0)).getDeviceEntries().size() + && ((DeviceTableScanNode) resultTableScanNodeList.get(0)).getDeviceEntryCount() == 1); for (final PlanNode planNode : resultTableScanNodeList) { final DeviceTableScanNode scanNode = (DeviceTableScanNode) planNode; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/DeviceTableScanNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/DeviceTableScanNode.java index d3d8af862bd..50a8e59bd17 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/DeviceTableScanNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/DeviceTableScanNode.java @@ -31,6 +31,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; import org.apache.iotdb.db.queryengine.plan.relational.metadata.AlignedDeviceEntry; import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryDataSet; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryDataSetHandle; import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering; import org.apache.tsfile.read.filter.basic.Filter; @@ -52,7 +53,10 @@ public class DeviceTableScanNode extends TableScanNode { protected List<DeviceEntry> deviceEntries = Collections.emptyList(); - @Nullable protected transient DeviceEntryDataSet deviceEntryDataSet; + @Nullable protected DeviceEntryDataSetHandle deviceEntryDataSetHandle; + + // Only used on the FE before distributed planning and is not serialized to the BE. + protected transient DeviceEntryDataSet coordinatorDeviceEntryDataSet; // Indicates the respective index order of tag and attribute columns in DeviceEntry. // For example, for DeviceEntry `table1.tag1.tag2.attribute1.attribute2.s1.s2`, the content of @@ -152,16 +156,23 @@ public class DeviceTableScanNode extends TableScanNode { pushLimitToEachDevice, containsNonAlignedDevice); cloned.topKRuntimeFilterSourceId = topKRuntimeFilterSourceId; - return copyDeviceEntryDataSetTo(cloned); + cloned.deviceEntryDataSetHandle = deviceEntryDataSetHandle; + cloned.coordinatorDeviceEntryDataSet = coordinatorDeviceEntryDataSet; + return cloned; } protected static void serializeMemberVariables( DeviceTableScanNode node, ByteBuffer byteBuffer, boolean serializeOutputSymbols) { TableScanNode.serializeMemberVariables(node, byteBuffer, serializeOutputSymbols); - ReadWriteIOUtils.write(node.deviceEntries.size(), byteBuffer); - for (DeviceEntry entry : node.deviceEntries) { - entry.serialize(byteBuffer); + ReadWriteIOUtils.write(node.deviceEntryDataSetHandle != null, byteBuffer); + if (node.deviceEntryDataSetHandle != null) { + node.deviceEntryDataSetHandle.serialize(byteBuffer); + } else { + ReadWriteIOUtils.write(node.deviceEntries.size(), byteBuffer); + for (DeviceEntry entry : node.deviceEntries) { + entry.serialize(byteBuffer); + } } ReadWriteIOUtils.write(node.tagAndAttributeIndexMap.size(), byteBuffer); @@ -189,9 +200,14 @@ public class DeviceTableScanNode extends TableScanNode { throws IOException { TableScanNode.serializeMemberVariables(node, stream, serializeOutputSymbols); - ReadWriteIOUtils.write(node.deviceEntries.size(), stream); - for (DeviceEntry entry : node.deviceEntries) { - entry.serialize(stream); + ReadWriteIOUtils.write(node.deviceEntryDataSetHandle != null, stream); + if (node.deviceEntryDataSetHandle != null) { + node.deviceEntryDataSetHandle.serialize(stream); + } else { + ReadWriteIOUtils.write(node.deviceEntries.size(), stream); + for (DeviceEntry entry : node.deviceEntries) { + entry.serialize(stream); + } } ReadWriteIOUtils.write(node.tagAndAttributeIndexMap.size(), stream); @@ -218,12 +234,18 @@ public class DeviceTableScanNode extends TableScanNode { ByteBuffer byteBuffer, DeviceTableScanNode node, boolean deserializeOutputSymbols) { TableScanNode.deserializeMemberVariables(byteBuffer, node, deserializeOutputSymbols); - int size = ReadWriteIOUtils.readInt(byteBuffer); - List<DeviceEntry> deviceEntries = new ArrayList<>(size); - while (size-- > 0) { - deviceEntries.add(AlignedDeviceEntry.deserialize(byteBuffer)); + int size; + if (ReadWriteIOUtils.readBool(byteBuffer)) { + node.deviceEntryDataSetHandle = DeviceEntryDataSetHandle.deserialize(byteBuffer); + node.deviceEntries = new ArrayList<>(); + } else { + size = ReadWriteIOUtils.readInt(byteBuffer); + List<DeviceEntry> deviceEntries = new ArrayList<>(size); + while (size-- > 0) { + deviceEntries.add(AlignedDeviceEntry.deserialize(byteBuffer)); + } + node.deviceEntries = deviceEntries; } - node.deviceEntries = deviceEntries; size = ReadWriteIOUtils.readInt(byteBuffer); Map<Symbol, Integer> tagAndAttributeIndexMap = new HashMap<>(size); @@ -269,28 +291,48 @@ public class DeviceTableScanNode extends TableScanNode { public void setDeviceEntries(List<DeviceEntry> deviceEntries) { this.deviceEntries = deviceEntries; + this.deviceEntryDataSetHandle = null; } public void setDeviceEntryDataSet(final DeviceEntryDataSet deviceEntryDataSet) { - this.deviceEntryDataSet = deviceEntryDataSet; + this.coordinatorDeviceEntryDataSet = deviceEntryDataSet; this.deviceEntries = deviceEntryDataSet.isSpilled() ? Collections.emptyList() : deviceEntryDataSet.getInlineEntries(); } - @Nullable - public DeviceEntryDataSet getDeviceEntryDataSet() { - return deviceEntryDataSet; + public void setDeviceEntryDataSetHandle(DeviceEntryDataSetHandle deviceEntryDataSetHandle) { + this.deviceEntryDataSetHandle = deviceEntryDataSetHandle; + this.deviceEntries = Collections.emptyList(); + } + + public Optional<DeviceEntryDataSetHandle> getDeviceEntryDataSetHandle() { + return Optional.ofNullable(deviceEntryDataSetHandle); + } + + public boolean hasSpilledDeviceEntries() { + return deviceEntryDataSetHandle != null; } public <T extends DeviceTableScanNode> T copyDeviceEntryDataSetTo(final T target) { - target.deviceEntryDataSet = deviceEntryDataSet; + target.deviceEntryDataSetHandle = deviceEntryDataSetHandle; + target.coordinatorDeviceEntryDataSet = coordinatorDeviceEntryDataSet; return target; } - public long getDeviceEntryCount() { - return deviceEntryDataSet == null ? deviceEntries.size() : deviceEntryDataSet.getEntryCount(); + public int getDeviceEntryCount() { + return deviceEntryDataSetHandle == null + ? deviceEntries.size() + : deviceEntryDataSetHandle.getEntryCount(); + } + + public void setCoordinatorDeviceEntryDataSet(DeviceEntryDataSet dataSet) { + setDeviceEntryDataSet(dataSet); + } + + public DeviceEntryDataSet getCoordinatorDeviceEntryDataSet() { + return coordinatorDeviceEntryDataSet; } public Map<Symbol, Integer> getTagAndAttributeIndexMap() { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializerTest.java index d643a91c831..5350341698d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/spill/DeviceEntryMaterializerTest.java @@ -35,7 +35,9 @@ import java.nio.ByteBuffer; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; +import java.util.Comparator; import java.util.List; +import java.util.stream.Collectors; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -118,7 +120,10 @@ public class DeviceEntryMaterializerTest { List<Path> segments; try (java.util.stream.Stream<Path> stream = Files.list(rawDirectory)) { segments = - stream.filter(path -> path.getFileName().toString().endsWith(".bin")).sorted().toList(); + stream + .filter(path -> path.getFileName().toString().endsWith(".bin")) + .sorted() + .collect(Collectors.toList()); } assertTrue(segments.size() > 1); assertTrue(Files.size(segments.get(0)) > 0); @@ -144,6 +149,52 @@ public class DeviceEntryMaterializerTest { dataSet.close(); } + @Test + public void testMemoryControllerSpillsLargestMaterializer() throws Exception { + DeviceEntry first = createEntries(1).get(0); + DeviceEntry second = createEntries(2).get(1); + try (DeviceEntryMaterializer firstMaterializer = + new DeviceEntryMaterializer("q-controller", new PlanNodeId("scan-0"), 128, false); + DeviceEntryMaterializer secondMaterializer = + new DeviceEntryMaterializer("q-controller", new PlanNodeId("scan-1"), 128, false)) { + DeviceEntryMaterializationMemoryController controller = + new DeviceEntryMaterializationMemoryController( + first.ramBytesUsed() + second.ramBytesUsed() - 1); + controller.append(firstMaterializer, first); + controller.append(secondMaterializer, second); + + assertTrue(firstMaterializer.isSpilled() || secondMaterializer.isSpilled()); + } + } + + @Test + public void testSortedMaterializerMergesRunsInOrder() throws Exception { + List<DeviceEntry> input = createEntries(40); + input.sort(Comparator.comparing(entry -> entry.getDeviceID().toString()).reversed()); + List<DeviceEntry> actual = new ArrayList<>(); + try (DeviceEntrySortedMaterializer materializer = + new DeviceEntrySortedMaterializer( + "q-sorted", + new PlanNodeId("scan-0"), + 128, + Comparator.comparing(entry -> entry.getDeviceID().toString()))) { + DeviceEntryMaterializationMemoryController controller = + new DeviceEntryMaterializationMemoryController(128); + for (DeviceEntry entry : input) { + controller.append(materializer, entry); + } + try (DeviceEntryDataSet dataSet = materializer.finish(); + DeviceEntryReader reader = dataSet.openReader()) { + while (reader.hasNext()) { + actual.add(reader.next()); + } + } + } + List<DeviceEntry> expected = new ArrayList<>(input); + expected.sort(Comparator.comparing(entry -> entry.getDeviceID().toString())); + assertEquals(expected, actual); + } + private static List<DeviceEntry> createEntries(int count) { List<DeviceEntry> entries = new ArrayList<>(count); for (int i = 0; i < count; i++) {
