This is an automated email from the ASF dual-hosted git repository. shuwenwei pushed a commit to branch transferSchemaTreeInBatches-1.3 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit f03ab0962d372072501ef4758fc38bdf7282821e Author: shuwenwei <[email protected]> AuthorDate: Thu Jul 31 13:58:41 2025 +0800 Transfer schema tree in batches & add memory control for schema tree --- .../apache/iotdb/db/it/orderBy/IoTDBOrderByIT.java | 4 + .../db/queryengine/common/MPPQueryContext.java | 33 ++++ .../common/schematree/ClusterSchemaTree.java | 192 ++++++++++++++++----- .../queryengine/common/schematree/ISchemaTree.java | 3 +- .../common/schematree/node/SchemaEntityNode.java | 23 +++ .../common/schematree/node/SchemaInternalNode.java | 17 ++ .../schematree/node/SchemaMeasurementNode.java | 24 +++ .../common/schematree/node/SchemaNode.java | 6 +- .../operator/schema/SchemaFetchScanOperator.java | 104 ++++++++--- .../queryengine/plan/analyze/AnalyzeVisitor.java | 2 +- .../db/queryengine/plan/analyze/Analyzer.java | 14 +- .../analyze/schema/ClusterSchemaFetchExecutor.java | 29 +++- .../iotdb/db/schemaengine/template/Template.java | 15 +- .../schema/SchemaFetchScanOperatorTest.java | 52 ++++-- pom.xml | 2 +- 15 files changed, 430 insertions(+), 90 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/orderBy/IoTDBOrderByIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/orderBy/IoTDBOrderByIT.java index 097fa98281f..f35bd2bb978 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/orderBy/IoTDBOrderByIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/orderBy/IoTDBOrderByIT.java @@ -107,6 +107,10 @@ public class IoTDBOrderByIT { @BeforeClass public static void setUp() throws Exception { EnvFactory.getEnv().getConfig().getDataNodeCommonConfig().setSortBufferSize(1024 * 1024L); + EnvFactory.getEnv() + .getConfig() + .getDataNodeCommonConfig() + .setQueryMemoryProportion("1:100:200:50:200:400:200:50"); EnvFactory.getEnv().initClusterEnvironment(); insertData(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java index 63970762de7..0035f726c4d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java @@ -37,6 +37,7 @@ import java.util.HashSet; import java.util.LinkedList; import java.util.List; import java.util.Set; +import java.util.function.LongConsumer; /** * This class is used to record the context of a query including QueryId, query statement, session @@ -88,6 +89,10 @@ public class MPPQueryContext { // the updateScanNum process in distributed planning can be skipped. private boolean needUpdateScanNumForLastQuery = false; + private long reservedMemoryCostForSchemaTree = 0; + private boolean releaseSchemaTreeAfterAnalyzing = true; + private LongConsumer reserveMemoryForSchemaTreeFunc = null; + private boolean userQuery = false; public MPPQueryContext(QueryId queryId) { @@ -128,6 +133,34 @@ public class MPPQueryContext { this.initResultNodeContext(); } + public void setReserveMemoryForSchemaTreeFunc(LongConsumer reserveMemoryForSchemaTreeFunc) { + this.reserveMemoryForSchemaTreeFunc = reserveMemoryForSchemaTreeFunc; + } + + public void reserveMemoryForSchemaTree(long memoryCost) { + if (reserveMemoryForSchemaTreeFunc == null) { + return; + } + reserveMemoryForSchemaTreeFunc.accept(memoryCost); + this.reservedMemoryCostForSchemaTree += memoryCost; + } + + public void setReleaseSchemaTreeAfterAnalyzing(boolean releaseSchemaTreeAfterAnalyzing) { + this.releaseSchemaTreeAfterAnalyzing = releaseSchemaTreeAfterAnalyzing; + } + + public boolean releaseSchemaTreeAfterAnalyzing() { + return releaseSchemaTreeAfterAnalyzing; + } + + public void releaseMemoryForSchemaTree() { + if (reservedMemoryCostForSchemaTree <= 0) { + return; + } + this.memoryReservationManager.releaseMemoryCumulatively(reservedMemoryCostForSchemaTree); + reservedMemoryCostForSchemaTree = 0; + } + public void prepareForRetry() { this.initResultNodeContext(); this.releaseAllMemoryReservedForFrontEnd(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ClusterSchemaTree.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ClusterSchemaTree.java index dc9776ab54a..40a49c368d9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ClusterSchemaTree.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ClusterSchemaTree.java @@ -39,6 +39,7 @@ import org.apache.iotdb.db.schemaengine.template.ClusterTemplateManager; import org.apache.iotdb.db.schemaengine.template.Template; import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.utils.RamUsageEstimator; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.schema.IMeasurementSchema; @@ -49,8 +50,10 @@ import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Deque; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.NoSuchElementException; import java.util.Set; import static org.apache.iotdb.commons.conf.IoTDBConstant.PATH_ROOT; @@ -60,6 +63,8 @@ import static org.apache.iotdb.db.queryengine.common.schematree.node.SchemaNode. import static org.apache.iotdb.db.queryengine.common.schematree.node.SchemaNode.SCHEMA_MEASUREMENT_NODE; public class ClusterSchemaTree implements ISchemaTree { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ClusterSchemaTree.class); private static final ClusterTemplateManager templateManager = ClusterTemplateManager.getInstance(); @@ -75,6 +80,8 @@ public class ClusterSchemaTree implements ISchemaTree { private Map<Integer, Template> templateMap = new HashMap<>(); + private long ramBytesUsed; + public ClusterSchemaTree() { root = new SchemaInternalNode(PATH_ROOT); } @@ -484,59 +491,158 @@ public class ClusterSchemaTree implements ISchemaTree { root.serialize(outputStream); } - public static ClusterSchemaTree deserialize(InputStream inputStream) throws IOException { + public Iterator<SchemaNode> getIteratorForSerialize() { + return new SchemaNodePostOrderIterator(root); + } - byte nodeType; - int childNum; - Deque<SchemaNode> stack = new ArrayDeque<>(); - SchemaNode child; - boolean hasLogicalView = false; - boolean hasNormalTimeSeries = false; - Map<Integer, Template> templateMap = new HashMap<>(); - - while (inputStream.available() > 0) { - nodeType = ReadWriteIOUtils.readByte(inputStream); - if (nodeType == SCHEMA_MEASUREMENT_NODE) { - SchemaMeasurementNode measurementNode = SchemaMeasurementNode.deserialize(inputStream); - stack.push(measurementNode); - if (measurementNode.isLogicalView()) { - hasLogicalView = true; + @Override + public long ramBytesUsed() { + if (ramBytesUsed > 0) { + return ramBytesUsed; + } + ramBytesUsed = + root.ramBytesUsed() + + SHALLOW_SIZE + + RamUsageEstimator.sizeOfMapWithKnownShallowSize( + templateMap, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP_ENTRY); + return ramBytesUsed; + } + + public void setRamBytesUsed(long ramBytesUsed) { + this.ramBytesUsed = ramBytesUsed; + } + + private static class SchemaNodePostOrderIterator implements Iterator<SchemaNode> { + // This class is likely to be faster than Stack when used as a stack + private final Deque<Pair<SchemaNode, Iterator<SchemaNode>>> stack = new ArrayDeque<>(); + private SchemaNode nextNode; + + public SchemaNodePostOrderIterator(SchemaNode root) { + stack.push(new Pair<>(root, root.getChildrenIterator())); + prepareNext(); + } + + @Override + public boolean hasNext() { + return nextNode != null; + } + + @Override + public SchemaNode next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + SchemaNode result = nextNode; + prepareNext(); + return result; + } + + private void prepareNext() { + nextNode = null; + while (!stack.isEmpty()) { + Pair<SchemaNode, Iterator<SchemaNode>> pair = stack.peek(); + SchemaNode currentNode = pair.getLeft(); + Iterator<SchemaNode> childrenIterator = pair.getRight(); + if (childrenIterator.hasNext()) { + SchemaNode child = childrenIterator.next(); + stack.push(new Pair<>(child, child.getChildrenIterator())); + } else { + stack.pop(); + nextNode = currentNode; + return; } - hasNormalTimeSeries = true; - } else { - SchemaInternalNode internalNode; - if (nodeType == SCHEMA_ENTITY_NODE) { - internalNode = SchemaEntityNode.deserialize(inputStream); - int templateId = internalNode.getAsEntityNode().getTemplateId(); - if (templateId != NON_TEMPLATE) { - templateMap.putIfAbsent(templateId, templateManager.getTemplate(templateId)); + } + } + } + + public static class SchemaNodeBatchDeserializer { + private byte nodeType; + private int childNum; + // This class is likely to be faster than Stack when used as a stack + private final Deque<SchemaNode> stack = new ArrayDeque<>(); + private SchemaNode child; + private boolean hasLogicalView = false; + private boolean hasNormalTimeSeries = false; + private Map<Integer, Template> templateMap = new HashMap<>(); + private boolean isFirstBatch = true; + + public boolean isFirstBatch() { + return isFirstBatch; + } + + public void deserializeFromBatch(InputStream inputStream) throws IOException { + isFirstBatch = false; + while (inputStream.available() > 0) { + nodeType = ReadWriteIOUtils.readByte(inputStream); + if (nodeType == SCHEMA_MEASUREMENT_NODE) { + SchemaMeasurementNode measurementNode = SchemaMeasurementNode.deserialize(inputStream); + stack.push(measurementNode); + if (measurementNode.isLogicalView()) { + hasLogicalView = true; } + hasNormalTimeSeries = true; } else { - internalNode = SchemaInternalNode.deserialize(inputStream); - } + SchemaInternalNode internalNode; + if (nodeType == SCHEMA_ENTITY_NODE) { + internalNode = SchemaEntityNode.deserialize(inputStream); + int templateId = internalNode.getAsEntityNode().getTemplateId(); + if (templateId != NON_TEMPLATE) { + templateMap.putIfAbsent(templateId, templateManager.getTemplate(templateId)); + } + } else { + internalNode = SchemaInternalNode.deserialize(inputStream); + } - childNum = ReadWriteIOUtils.readInt(inputStream); - while (childNum > 0) { - child = stack.pop(); - internalNode.addChild(child.getName(), child); - if (child.isMeasurement()) { - SchemaMeasurementNode measurementNode = child.getAsMeasurementNode(); - if (measurementNode.getAlias() != null) { - internalNode - .getAsEntityNode() - .addAliasChild(measurementNode.getAlias(), measurementNode); + childNum = ReadWriteIOUtils.readInt(inputStream); + while (childNum > 0) { + child = stack.pop(); + internalNode.addChild(child.getName(), child); + if (child.isMeasurement()) { + SchemaMeasurementNode measurementNode = child.getAsMeasurementNode(); + if (measurementNode.getAlias() != null) { + internalNode + .getAsEntityNode() + .addAliasChild(measurementNode.getAlias(), measurementNode); + } } + childNum--; } - childNum--; + stack.push(internalNode); } - stack.push(internalNode); } } - ClusterSchemaTree result = new ClusterSchemaTree(stack.poll()); - result.templateMap = templateMap; - result.hasLogicalMeasurementPath = hasLogicalView; - result.hasNormalTimeSeries = hasNormalTimeSeries; - return result; + + public ClusterSchemaTree finish() { + try { + ClusterSchemaTree result = new ClusterSchemaTree(stack.poll()); + result.templateMap = templateMap; + result.hasLogicalMeasurementPath = hasLogicalView; + result.hasNormalTimeSeries = hasNormalTimeSeries; + return result; + } finally { + reset(); + } + } + + private void reset() { + nodeType = 0; + childNum = 0; + stack.clear(); + child = null; + hasLogicalView = false; + hasNormalTimeSeries = false; + // templateMap is set to the returned schema tree, so we should create a new one + templateMap = new HashMap<>(); + isFirstBatch = true; + } + } + + public static ClusterSchemaTree deserialize(InputStream inputStream) throws IOException { + SchemaNodeBatchDeserializer schemaNodeBatchDeserializer = new SchemaNodeBatchDeserializer(); + schemaNodeBatchDeserializer.deserializeFromBatch(inputStream); + return schemaNodeBatchDeserializer.finish(); } /** diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ISchemaTree.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ISchemaTree.java index 0288033b695..34cb2a09786 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ISchemaTree.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/ISchemaTree.java @@ -24,12 +24,13 @@ import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.schemaengine.template.Template; +import org.apache.tsfile.utils.Accountable; import org.apache.tsfile.utils.Pair; import java.util.List; import java.util.Set; -public interface ISchemaTree { +public interface ISchemaTree extends Accountable { /** * Return all measurement paths for given path pattern and filter the result by slimit and offset. * diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaEntityNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaEntityNode.java index b877dc9d7a2..cf111e61b29 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaEntityNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaEntityNode.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.common.schematree.node; +import org.apache.tsfile.utils.RamUsageEstimator; import org.apache.tsfile.utils.ReadWriteIOUtils; import java.io.IOException; @@ -31,6 +32,9 @@ import static org.apache.iotdb.commons.schema.SchemaConstant.NON_TEMPLATE; public class SchemaEntityNode extends SchemaInternalNode { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(SchemaEntityNode.class); + private boolean isAligned; private Map<String, SchemaMeasurementNode> aliasChildren; @@ -117,6 +121,11 @@ public class SchemaEntityNode extends SchemaInternalNode { @Override public void serialize(OutputStream outputStream) throws IOException { serializeChildren(outputStream); + this.serializeNodeOwnContent(outputStream); + } + + @Override + public void serializeNodeOwnContent(OutputStream outputStream) throws IOException { ReadWriteIOUtils.write(getType(), outputStream); ReadWriteIOUtils.write(name, outputStream); ReadWriteIOUtils.write(isAligned, outputStream); @@ -133,4 +142,18 @@ public class SchemaEntityNode extends SchemaInternalNode { entityNode.setTemplateId(templateId); return entityNode; } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + + RamUsageEstimator.sizeOf(name) + + RamUsageEstimator.sizeOfMapWithKnownShallowSize( + children, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP_ENTRY) + + RamUsageEstimator.sizeOfMapWithKnownShallowSize( + aliasChildren, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP_ENTRY); + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaInternalNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaInternalNode.java index 5c6de241baf..84ccf5b35fa 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaInternalNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaInternalNode.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.common.schematree.node; +import org.apache.tsfile.utils.RamUsageEstimator; import org.apache.tsfile.utils.ReadWriteIOUtils; import java.io.IOException; @@ -30,6 +31,9 @@ import java.util.Map; public class SchemaInternalNode extends SchemaNode { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(SchemaInternalNode.class); + protected Map<String, SchemaNode> children = new HashMap<>(); public SchemaInternalNode(String name) { @@ -85,7 +89,10 @@ public class SchemaInternalNode extends SchemaNode { public void serialize(OutputStream outputStream) throws IOException { serializeChildren(outputStream); + serializeNodeOwnContent(outputStream); + } + public void serializeNodeOwnContent(OutputStream outputStream) throws IOException { ReadWriteIOUtils.write(getType(), outputStream); ReadWriteIOUtils.write(name, outputStream); ReadWriteIOUtils.write(children.size(), outputStream); @@ -102,4 +109,14 @@ public class SchemaInternalNode extends SchemaNode { return new SchemaInternalNode(name); } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + + RamUsageEstimator.sizeOf(name) + + RamUsageEstimator.sizeOfMapWithKnownShallowSize( + children, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP_ENTRY); + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaMeasurementNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaMeasurementNode.java index b1eaa7ff5c1..6f5bd07f08f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaMeasurementNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaMeasurementNode.java @@ -22,6 +22,7 @@ package org.apache.iotdb.db.queryengine.common.schematree.node; import org.apache.iotdb.commons.schema.view.LogicalViewSchema; import org.apache.iotdb.db.queryengine.common.schematree.IMeasurementSchemaInfo; +import org.apache.tsfile.utils.RamUsageEstimator; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.schema.IMeasurementSchema; import org.apache.tsfile.write.schema.MeasurementSchema; @@ -34,6 +35,9 @@ import java.util.Map; public class SchemaMeasurementNode extends SchemaNode implements IMeasurementSchemaInfo { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(SchemaMeasurementNode.class); + private String alias; private IMeasurementSchema schema; private Map<String, String> tagMap; @@ -44,6 +48,22 @@ public class SchemaMeasurementNode extends SchemaNode implements IMeasurementSch this.schema = schema; } + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + + RamUsageEstimator.sizeOf(name) + + RamUsageEstimator.sizeOf(alias) + + schema.ramBytesUsed() + + RamUsageEstimator.sizeOfMapWithKnownShallowSize( + tagMap, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP_ENTRY) + + RamUsageEstimator.sizeOfMapWithKnownShallowSize( + attributeMap, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP, + RamUsageEstimator.SHALLOW_SIZE_OF_HASHMAP_ENTRY); + } + public String getAlias() { return alias; } @@ -143,6 +163,10 @@ public class SchemaMeasurementNode extends SchemaNode implements IMeasurementSch @Override public void serialize(OutputStream outputStream) throws IOException { + serializeNodeOwnContent(outputStream); + } + + public void serializeNodeOwnContent(OutputStream outputStream) throws IOException { ReadWriteIOUtils.write(getType(), outputStream); ReadWriteIOUtils.write(name, outputStream); ReadWriteIOUtils.write(alias, outputStream); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaNode.java index e2625cd97ac..dd4e5e57a2e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/node/SchemaNode.java @@ -21,13 +21,15 @@ package org.apache.iotdb.db.queryengine.common.schematree.node; import org.apache.iotdb.commons.schema.tree.ITreeNode; +import org.apache.tsfile.utils.Accountable; + import java.io.IOException; import java.io.OutputStream; import java.util.Collections; import java.util.Iterator; import java.util.Map; -public abstract class SchemaNode implements ITreeNode { +public abstract class SchemaNode implements ITreeNode, Accountable { public static final byte SCHEMA_INTERNAL_NODE = 0; public static final byte SCHEMA_ENTITY_NODE = 1; @@ -80,4 +82,6 @@ public abstract class SchemaNode implements ITreeNode { public abstract byte getType(); public abstract void serialize(OutputStream outputStream) throws IOException; + + public abstract void serializeNodeOwnContent(OutputStream outputStream) throws IOException; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperator.java index f4b40c69b79..6cc8e06878c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperator.java @@ -23,7 +23,9 @@ import org.apache.iotdb.commons.exception.MetadataException; import org.apache.iotdb.commons.exception.runtime.SchemaExecutionException; import org.apache.iotdb.commons.path.PathPatternTree; import org.apache.iotdb.commons.schema.SchemaConstant; +import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.queryengine.common.schematree.ClusterSchemaTree; +import org.apache.iotdb.db.queryengine.common.schematree.node.SchemaNode; import org.apache.iotdb.db.queryengine.execution.MemoryEstimationHelper; import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext; import org.apache.iotdb.db.queryengine.execution.operator.source.SourceOperator; @@ -36,12 +38,12 @@ import org.apache.tsfile.read.common.block.TsBlock; import org.apache.tsfile.read.common.block.column.BinaryColumn; import org.apache.tsfile.read.common.block.column.TimeColumn; import org.apache.tsfile.utils.Binary; +import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.RamUsageEstimator; import org.apache.tsfile.utils.ReadWriteIOUtils; -import java.io.ByteArrayOutputStream; -import java.io.IOException; import java.util.Collections; +import java.util.Iterator; import java.util.Map; import java.util.NoSuchElementException; import java.util.Optional; @@ -62,7 +64,12 @@ public class SchemaFetchScanOperator implements SourceOperator { private boolean isFinished = false; private final PathPatternTree authorityScope; - private static final int DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES = + private Iterator<SchemaNode> schemaNodeIteratorForSerialize = null; + private long schemaTreeMemCost; + private PublicBAOS baos = null; + // Reserve some bytes to avoid capacity grow + private static final int EXTRA_SIZE_TO_AVOID_GROW = 1024; + private static int DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); private static final long INSTANCE_SIZE = @@ -152,12 +159,33 @@ public class SchemaFetchScanOperator implements SourceOperator { if (!hasNext()) { throw new NoSuchElementException(); } - isFinished = true; - try { - return fetchSchema(); - } catch (MetadataException e) { - throw new SchemaExecutionException(e); + + boolean isFirstBatch = schemaNodeIteratorForSerialize == null; + prepareSchemaNodeIteratorForSerialize(); + // to indicate this binary data is a part of schema tree, and the remaining parts will be sent + // later + ReadWriteIOUtils.write((byte) 2, baos); + // the estimated mem cost to deserialize the total schema tree + if (isFirstBatch) { + ReadWriteIOUtils.write(schemaTreeMemCost, baos); + } + while (schemaNodeIteratorForSerialize.hasNext() + && baos.size() < DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES) { + SchemaNode node = schemaNodeIteratorForSerialize.next(); + node.serializeNodeOwnContent(baos); } + byte[] currentBatch = baos.toByteArray(); + baos.reset(); + isFinished = !schemaNodeIteratorForSerialize.hasNext(); + if (isFinished) { + // indicate all continuous binary data is finished + currentBatch[0] = 3; + releaseSchemaTree(); + baos = null; + } + return new TsBlock( + new TimeColumn(1, new long[] {0}), + new BinaryColumn(1, Optional.empty(), new Binary[] {new Binary(currentBatch)})); } @Override @@ -172,7 +200,8 @@ public class SchemaFetchScanOperator implements SourceOperator { @Override public void close() throws Exception { - // do nothing + releaseSchemaTree(); + baos = null; } @Override @@ -180,26 +209,28 @@ public class SchemaFetchScanOperator implements SourceOperator { return sourceId; } - private TsBlock fetchSchema() throws MetadataException { - ClusterSchemaTree schemaTree = - fetchDevice - ? schemaRegion.fetchDeviceSchema(patternTree, authorityScope) - : schemaRegion.fetchSeriesSchema( - patternTree, templateMap, withTags, withAttributes, withTemplate, withAliasForce); - - ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + private void prepareSchemaNodeIteratorForSerialize() { + if (schemaNodeIteratorForSerialize != null) { + return; + } try { - // to indicate this binary data is database info - ReadWriteIOUtils.write((byte) 1, outputStream); - - schemaTree.serialize(outputStream); - } catch (IOException e) { - // Totally memory operation. This case won't happen. + ClusterSchemaTree schemaTree = + fetchDevice + ? schemaRegion.fetchDeviceSchema(patternTree, authorityScope) + : schemaRegion.fetchSeriesSchema( + patternTree, templateMap, withTags, withAttributes, withTemplate, withAliasForce); + schemaNodeIteratorForSerialize = schemaTree.getIteratorForSerialize(); + baos = new PublicBAOS(DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES + EXTRA_SIZE_TO_AVOID_GROW); + if (operatorContext != null) { + schemaTreeMemCost = schemaTree.ramBytesUsed(); + operatorContext + .getInstanceContext() + .getMemoryReservationContext() + .reserveMemoryCumulatively(schemaTreeMemCost); + } + } catch (MetadataException e) { + throw new SchemaExecutionException(e); } - return new TsBlock( - new TimeColumn(1, new long[] {0}), - new BinaryColumn( - 1, Optional.empty(), new Binary[] {new Binary(outputStream.toByteArray())})); } @Override @@ -221,6 +252,25 @@ public class SchemaFetchScanOperator implements SourceOperator { public long ramBytesUsed() { return INSTANCE_SIZE + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext) + + DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES + + EXTRA_SIZE_TO_AVOID_GROW + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(sourceId); } + + private void releaseSchemaTree() { + if (schemaTreeMemCost <= 0 || operatorContext == null) { + return; + } + operatorContext + .getInstanceContext() + .getMemoryReservationContext() + .releaseMemoryCumulatively(schemaTreeMemCost); + schemaTreeMemCost = 0; + schemaNodeIteratorForSerialize = null; + } + + @TestOnly + public static void setDefaultMaxTsBlockSizeInBytes(int newSize) { + DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES = newSize; + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java index 9ec064245f4..dab31639205 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java @@ -535,7 +535,6 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext> context.setFetchSchemaCost(schemaFetchCost); QueryPlanCostMetricSet.getInstance().recordPlanCost(SCHEMA_FETCHER, schemaFetchCost); } - analysis.setSchemaTree(schemaTree); return schemaTree; } @@ -3610,6 +3609,7 @@ public class AnalyzeVisitor extends StatementVisitor<Analysis, MPPQueryContext> } } analysis.setSchemaTree(schemaTree); + context.setReleaseSchemaTreeAfterAnalyzing(false); Map<String, List<DataPartitionQueryParam>> sgNameToQueryParamsMap = new HashMap<>(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java index f7ca04f7d05..34b17003f1b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java @@ -43,8 +43,18 @@ public class Analyzer { public Analysis analyze(Statement statement) { long startTime = System.nanoTime(); - Analysis analysis = - new AnalyzeVisitor(partitionFetcher, schemaFetcher).process(statement, context); + AnalyzeVisitor visitor = new AnalyzeVisitor(partitionFetcher, schemaFetcher); + Analysis analysis = null; + context.setReserveMemoryForSchemaTreeFunc(context::reserveMemoryForFrontEnd); + try { + analysis = visitor.process(statement, context); + } finally { + if (analysis != null && context.releaseSchemaTreeAfterAnalyzing()) { + analysis.setSchemaTree(null); + context.releaseMemoryForSchemaTree(); + } + context.setReserveMemoryForSchemaTreeFunc(null); + } if (statement.isQuery()) { QueryPlanCostMetricSet.getInstance().recordPlanCost(ANALYZER, System.nanoTime() - startTime); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java index d62b5640d9d..e71b3299612 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java @@ -252,6 +252,8 @@ class ClusterSchemaFetchExecutor { } try (SetThreadName threadName = new SetThreadName(executionResult.queryId.getId())) { ClusterSchemaTree result = new ClusterSchemaTree(); + ClusterSchemaTree.SchemaNodeBatchDeserializer deserializer = + new ClusterSchemaTree.SchemaNodeBatchDeserializer(); Set<String> databaseSet = new HashSet<>(); while (coordinator.getQueryExecution(queryId).hasNextResult()) { // The query will be transited to FINISHED when invoking getBatchResult() at the last time @@ -268,7 +270,7 @@ class ClusterSchemaFetchExecutor { } Column column = tsBlock.get().getColumn(0); for (int i = 0; i < column.getPositionCount(); i++) { - parseFetchedData(column.getBinary(i), result, databaseSet); + parseFetchedData(column.getBinary(i), result, deserializer, databaseSet, context); } } result.setDatabases(databaseSet); @@ -283,7 +285,11 @@ class ClusterSchemaFetchExecutor { } private void parseFetchedData( - Binary data, ClusterSchemaTree resultSchemaTree, Set<String> databaseSet) { + Binary data, + ClusterSchemaTree resultSchemaTree, + ClusterSchemaTree.SchemaNodeBatchDeserializer deserializer, + Set<String> databaseSet, + MPPQueryContext context) { InputStream inputStream = new ByteArrayInputStream(data.getValues()); try { byte type = ReadWriteIOUtils.readByte(inputStream); @@ -293,7 +299,24 @@ class ClusterSchemaFetchExecutor { databaseSet.add(ReadWriteIOUtils.readString(inputStream)); } } else if (type == 1) { - resultSchemaTree.mergeSchemaTree(ClusterSchemaTree.deserialize(inputStream)); + // for data from old version + ClusterSchemaTree deserializedSchemaTree = ClusterSchemaTree.deserialize(inputStream); + if (context != null) { + context.reserveMemoryForSchemaTree(deserializedSchemaTree.ramBytesUsed()); + } + resultSchemaTree.mergeSchemaTree(deserializedSchemaTree); + } else if (type == 2 || type == 3) { + if (deserializer.isFirstBatch()) { + long memCost = ReadWriteIOUtils.readLong(inputStream); + if (context != null) { + context.reserveMemoryForSchemaTree(memCost); + } + } + deserializer.deserializeFromBatch(inputStream); + if (type == 3) { + // 'type == 3' indicates this batch is finished + resultSchemaTree.mergeSchemaTree(deserializer.finish()); + } } else { throw new RuntimeException( new MetadataException("Failed to fetch schema because of unrecognized data")); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/template/Template.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/template/Template.java index 8fe86a9abfe..6639427b59b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/template/Template.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/template/Template.java @@ -25,6 +25,8 @@ import org.apache.commons.lang3.builder.HashCodeBuilder; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.utils.Accountable; +import org.apache.tsfile.utils.RamUsageEstimator; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.schema.IMeasurementSchema; import org.apache.tsfile.write.schema.MeasurementSchema; @@ -39,8 +41,9 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; -public class Template implements Serializable { +public class Template implements Serializable, Accountable { + private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(Template.class); private int id; private String name; private boolean isDirectAligned; @@ -226,4 +229,14 @@ public class Template implements Serializable { public int hashCode() { return new HashCodeBuilder(17, 37).append(name).append(schemaMap).toHashCode(); } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + + RamUsageEstimator.sizeOf(name) + + RamUsageEstimator.sizeOfMapWithKnownShallowSize( + schemaMap, + RamUsageEstimator.SHALLOW_SIZE_OF_CONCURRENT_HASHMAP, + RamUsageEstimator.SHALLOW_SIZE_OF_CONCURRENT_HASHMAP_ENTRY); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperatorTest.java index 50249bc55d5..2ac62e4e633 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperatorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/schema/SchemaFetchScanOperatorTest.java @@ -24,9 +24,9 @@ import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.commons.path.PathPatternTree; import org.apache.iotdb.db.queryengine.common.schematree.ClusterSchemaTree; import org.apache.iotdb.db.queryengine.common.schematree.DeviceSchemaInfo; -import org.apache.iotdb.db.queryengine.common.schematree.ISchemaTree; import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion; +import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; @@ -36,19 +36,42 @@ import org.apache.tsfile.utils.Pair; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.schema.IMeasurementSchema; import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; import org.junit.Assert; import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; import org.mockito.Mockito; import java.io.ByteArrayInputStream; import java.io.InputStream; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.stream.Collectors; +@RunWith(Parameterized.class) public class SchemaFetchScanOperatorTest { + @Parameterized.Parameters(name = "{0}") + public static Collection<Object[]> data() { + return Arrays.asList( + new Object[][] { + {1024}, {128 * 1024}, {16}, + }); + } + + public SchemaFetchScanOperatorTest(int maxTsBlockSizeInBytes) { + SchemaFetchScanOperator.setDefaultMaxTsBlockSizeInBytes(maxTsBlockSizeInBytes); + } + + @After + public void tearDown() { + SchemaFetchScanOperator.setDefaultMaxTsBlockSizeInBytes( + TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes()); + } + @Test public void testSchemaFetchResult() throws Exception { ISchemaRegion schemaRegion = mockSchemaRegion(); @@ -70,16 +93,25 @@ public class SchemaFetchScanOperatorTest { true, false); - Assert.assertTrue(schemaFetchScanOperator.hasNext()); - - TsBlock tsBlock = schemaFetchScanOperator.next(); - + ClusterSchemaTree.SchemaNodeBatchDeserializer deserializer = + new ClusterSchemaTree.SchemaNodeBatchDeserializer(); + byte type = 0; + while (schemaFetchScanOperator.hasNext()) { + TsBlock tsBlock = schemaFetchScanOperator.next(); + Binary binary = tsBlock.getColumn(0).getBinary(0); + InputStream inputStream = new ByteArrayInputStream(binary.getValues()); + if (!deserializer.isFirstBatch()) { + Assert.assertEquals(2, type); + } + type = ReadWriteIOUtils.readByte(inputStream); + if (deserializer.isFirstBatch()) { + ReadWriteIOUtils.readLong(inputStream); + } + deserializer.deserializeFromBatch(inputStream); + } + Assert.assertEquals(3, type); Assert.assertFalse(schemaFetchScanOperator.hasNext()); - - Binary binary = tsBlock.getColumn(0).getBinary(0); - InputStream inputStream = new ByteArrayInputStream(binary.getValues()); - Assert.assertEquals(1, ReadWriteIOUtils.readByte(inputStream)); - ISchemaTree schemaTree = ClusterSchemaTree.deserialize(inputStream); + ClusterSchemaTree schemaTree = deserializer.finish(); DeviceSchemaInfo deviceSchemaInfo = schemaTree.searchDeviceSchemaInfo( diff --git a/pom.xml b/pom.xml index 2ae4f3a04e4..58bd5e89bca 100644 --- a/pom.xml +++ b/pom.xml @@ -175,7 +175,7 @@ <thrift.version>0.14.1</thrift.version> <xz.version>1.9</xz.version> <zstd-jni.version>1.5.6-3</zstd-jni.version> - <tsfile.version>1.1.2-250725-SNAPSHOT</tsfile.version> + <tsfile.version>1.1.2-250801-SNAPSHOT</tsfile.version> </properties> <!-- if we claim dependencies in dependencyManagement, then we do not claim
