This is an automated email from the ASF dual-hosted git repository. Wei-hao-Li pushed a commit to branch deviceEntrySpill-devPr in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 99f76a0a3cf5148ff44c615b6380c39cc16bbd90 Author: Weihao Li <[email protected]> AuthorDate: Tue Aug 18 23:52:14 2026 +0800 fix & add IT --- .../it/env/cluster/config/MppDataNodeConfig.java | 6 + .../it/env/remote/config/RemoteDataNodeConfig.java | 5 + .../apache/iotdb/itbase/env/DataNodeConfig.java | 2 + .../it/query/recent/IoTDBDeviceEntrySpillIT.java | 143 ++++++++++ .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../fragment/FragmentInstanceContext.java | 15 +- .../queryengine/plan/execution/QueryExecution.java | 12 + .../metadata/fetcher/DeviceEntryFetchContext.java | 71 +++++ .../metadata/fetcher/TableDeviceSchemaFetcher.java | 298 ++++++++++----------- .../spill/AbstractDeviceEntryMaterializer.java | 33 ++- .../metadata/spill/DeviceEntryMaterializer.java | 22 +- .../spill/DeviceEntrySortedMaterializer.java | 88 +++++- .../distribute/TableDistributedPlanGenerator.java | 76 +++++- .../relational/sql/ast/AbstractTraverseDevice.java | 5 +- .../spill/DeviceEntryMaterializerTest.java | 27 ++ 16 files changed, 611 insertions(+), 196 deletions(-) diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java index 8399955a5c5..8c5f56ad9f1 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java @@ -180,4 +180,10 @@ public class MppDataNodeConfig extends MppBaseConfig implements DataNodeConfig { setProperty("dn_multi_dir_strategy", multiDirStrategy); return this; } + + @Override + public DataNodeConfig setTableQueryDeviceEntryBatchSizeInBytes(long batchSizeInBytes) { + setProperty("table_query_device_entry_batch_size_in_bytes", String.valueOf(batchSizeInBytes)); + return this; + } } diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java b/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java index a76608e4851..c97e2b5065a 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java @@ -125,4 +125,9 @@ public class RemoteDataNodeConfig implements DataNodeConfig { public DataNodeConfig setDnMultiDirStrategy(String multiDirStrategy) { return this; } + + @Override + public DataNodeConfig setTableQueryDeviceEntryBatchSizeInBytes(long batchSizeInBytes) { + return this; + } } diff --git a/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java b/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java index bc045c9ba2f..fb969c778ea 100644 --- a/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java +++ b/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java @@ -65,4 +65,6 @@ public interface DataNodeConfig { DataNodeConfig setDnDataDirs(String dnDataDirs); DataNodeConfig setDnMultiDirStrategy(String multiDirStrategy); + + DataNodeConfig setTableQueryDeviceEntryBatchSizeInBytes(long batchSizeInBytes); } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBDeviceEntrySpillIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBDeviceEntrySpillIT.java new file mode 100644 index 00000000000..39bfbb9dd0a --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBDeviceEntrySpillIT.java @@ -0,0 +1,143 @@ +/* + * 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.relational.it.query.recent; + +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; +import org.apache.iotdb.itbase.env.BaseEnv; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; + +import static org.junit.Assert.assertEquals; + +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class IoTDBDeviceEntrySpillIT { + + @BeforeClass + public static void setUp() throws Exception { + EnvFactory.getEnv().getConfig().getDataNodeConfig().setTableQueryDeviceEntryBatchSizeInBytes(1); + EnvFactory.getEnv().initClusterEnvironment(); + try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); + Statement statement = connection.createStatement()) { + statement.execute("CREATE DATABASE spill_test"); + statement.execute( + "CREATE TABLE spill_test.device_data (tag1 STRING TAG, tag2 STRING TAG, " + + "value INT32 FIELD)"); + statement.execute( + "INSERT INTO spill_test.device_data(tag1, tag2, time, value) " + + "VALUES ('a', 'x', 1, 10), ('a', 'x', 2, 20), " + + "('b', 'y', 1, 30), ('c', 'z', 1, 40)"); + } + } + + @AfterClass + public static void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testRawFullTableQueryWithSpill() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); + Statement statement = connection.createStatement(); + ResultSet resultSet = + statement.executeQuery("SELECT tag1, tag2, value FROM spill_test.device_data")) { + int rowCount = 0; + while (resultSet.next()) { + rowCount++; + } + assertEquals(4, rowCount); + } + } + + @Test + public void testRawQueriesWithTimeFilterProjectionFilterLimitAndOrdering() throws Exception { + String[] queries = { + "SELECT * FROM spill_test.device_data WHERE time >= 1 AND time < 3", + "SELECT time, value FROM spill_test.device_data WHERE time >= 1 AND time < 3", + "SELECT tag1, tag2, value FROM spill_test.device_data " + + "WHERE time >= 1 AND time < 3 AND value > 10", + "SELECT * FROM spill_test.device_data WHERE time >= 1 AND time < 3 LIMIT 2", + "SELECT * FROM spill_test.device_data WHERE time >= 1 AND time < 3 ORDER BY time ASC", + "SELECT * FROM spill_test.device_data WHERE time >= 1 AND time < 3 " + "ORDER BY tag1, time" + }; + for (String query : queries) { + assertRowCount(query, query.contains("LIMIT 2") ? 2 : query.contains("value > 10") ? 2 : 4); + } + } + + @Test + public void testAggregationQueryWithSpill() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); + Statement statement = connection.createStatement(); + ResultSet resultSet = + statement.executeQuery("SELECT count(value) FROM spill_test.device_data")) { + assertEquals(true, resultSet.next()); + assertEquals(4, resultSet.getLong(1)); + assertEquals(false, resultSet.next()); + } + } + + @Test + public void testGroupedAggregationAcrossSpillSegments() throws Exception { + assertRowCount("SELECT tag1, count(*) FROM spill_test.device_data GROUP BY tag1", 3); + assertRowCount( + "SELECT tag1, tag2, count(*), sum(value) FROM spill_test.device_data " + + "GROUP BY tag1, tag2", + 3); + assertRowCount( + "SELECT date_bin(1s, time), count(*) FROM spill_test.device_data " + + "GROUP BY date_bin(1s, time)", + 2); + } + + @Test + public void testOrPredicateDoesNotDuplicateDeviceRows() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); + Statement statement = connection.createStatement(); + ResultSet resultSet = + statement.executeQuery( + "SELECT count(*) FROM spill_test.device_data " + + "WHERE tag1 = 'a' OR tag2 = 'x'")) { + assertEquals(true, resultSet.next()); + assertEquals(2, resultSet.getLong(1)); + assertEquals(false, resultSet.next()); + } + } + + private void assertRowCount(String sql, int expectedRowCount) throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); + Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(sql)) { + int rowCount = 0; + while (resultSet.next()) { + rowCount++; + } + assertEquals(expectedRowCount, rowCount); + } + } +} diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 5cb613debef..6821e295157 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -3821,5 +3821,7 @@ public final class DataNodeQueryMessages { public static final String LOG_FAILED_TO_CLEAN_DEVICEENTRY_DATA_SET_ASYNCHRONOUSLY_QUERYID_ARG_PLANNODEID_ARG_9106C4C5 = "Failed to clean DeviceEntry data set asynchronously: queryId=%s, planNodeId=%s"; + public static final String LOG_FAILED_TO_CLEAN_DEVICEENTRY_SPILL_DIRECTORY_QUERYID_ARG_ADF95D63 = + "Failed to clean DeviceEntry spill directory for query %s"; } diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 89bafbd34ed..b1ec4c71a93 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -4578,5 +4578,7 @@ public final class DataNodeQueryMessages { public static final String LOG_FAILED_TO_CLEAN_DEVICEENTRY_DATA_SET_ASYNCHRONOUSLY_QUERYID_ARG_PLANNODEID_ARG_9106C4C5 = "异步清理 DeviceEntry 数据集失败：queryId=%s，planNodeId=%s"; + public static final String LOG_FAILED_TO_CLEAN_DEVICEENTRY_SPILL_DIRECTORY_QUERYID_ARG_ADF95D63 = + "清理 query %s 的 DeviceEntry spill 目录失败"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index ad363016808..1a0af9bb579 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -819,14 +819,13 @@ public class FragmentInstanceContext extends QueryContext { return unfinishedResultSupplier.get(); } QueryDataSource dataSource = - (QueryDataSource) - dataRegion.query( - sourcePaths, - singleDeviceId, - this, - globalTimeFilter != null ? globalTimeFilter.copy() : null, - timePartitions, - waitForLockTime); + dataRegion.query( + sourcePaths, + singleDeviceId, + this, + globalTimeFilter != null ? globalTimeFilter.copy() : null, + timePartitions, + waitForLockTime); if (dataSource == null) { return unfinishedResultSupplier.get(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 035bac7b5d8..18578e625fd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -48,6 +48,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.DistributedQueryPlan; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeUtil; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntrySpillManager; import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -437,6 +438,17 @@ public class QueryExecution implements IQueryExecution { cleanUpResultHandle(); } context.releaseExternalTsFileQueryResources(); + if (t != null) { + try { + DeviceEntrySpillManager.getInstance().deregisterQuery(context.getQueryId().getId()); + } catch (Exception e) { + LOGGER.warn( + DataNodeQueryMessages + .LOG_FAILED_TO_CLEAN_DEVICEENTRY_SPILL_DIRECTORY_QUERYID_ARG_ADF95D63, + context.getQueryId().getId(), + e); + } + } } /** diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/DeviceEntryFetchContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/DeviceEntryFetchContext.java new file mode 100644 index 00000000000..47c01459154 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/DeviceEntryFetchContext.java @@ -0,0 +1,71 @@ +/* + * 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.fetcher; + +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.spill.AbstractDeviceEntryMaterializer; + +/** Query-scoped state shared while fetching and materializing DeviceEntry objects. */ +public final class DeviceEntryFetchContext { + + private final MPPQueryContext queryContext; + private final PlanNodeId planNodeId; + private AbstractDeviceEntryMaterializer materializer; + private boolean mayContainDuplicateDevice; + private boolean containsNonAlignedDevice; + + public DeviceEntryFetchContext(final MPPQueryContext queryContext, final PlanNodeId planNodeId) { + this.queryContext = queryContext; + this.planNodeId = planNodeId; + } + + public MPPQueryContext getQueryContext() { + return queryContext; + } + + public PlanNodeId getPlanNodeId() { + return planNodeId; + } + + public AbstractDeviceEntryMaterializer getMaterializer() { + return materializer; + } + + public void setMaterializer(final AbstractDeviceEntryMaterializer materializer) { + this.materializer = materializer; + } + + public boolean mayContainDuplicateDevice() { + return mayContainDuplicateDevice; + } + + public void setMayContainDuplicateDevice(final boolean mayContainDuplicateDevice) { + this.mayContainDuplicateDevice = mayContainDuplicateDevice; + } + + public boolean containsNonAlignedDevice() { + return containsNonAlignedDevice; + } + + public void markContainsNonAlignedDevice() { + containsNonAlignedDevice = true; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/TableDeviceSchemaFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/TableDeviceSchemaFetcher.java index 1ac6e1789a4..32fcfdc546c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/TableDeviceSchemaFetcher.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/TableDeviceSchemaFetcher.java @@ -50,9 +50,11 @@ import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.cache.De import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.cache.IDeviceSchema; import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.cache.TableDeviceSchemaCache; import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.cache.TreeDeviceNormalSchema; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.AbstractDeviceEntryMaterializer; import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryDataSet; import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryDataSetResult; import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntryMaterializer; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntrySortedMaterializer; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.AbstractTraverseDevice; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.FetchDevice; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ShowDevice; @@ -74,6 +76,7 @@ import java.io.UncheckedIOException; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.Comparator; import java.util.HashMap; import java.util.LinkedHashSet; import java.util.List; @@ -228,7 +231,7 @@ public class TableDeviceSchemaFetcher { final MPPQueryContext queryContext) { final Map<String, List<DeviceEntry>> deviceEntryMap = new HashMap<>(); final TsTable tableInstance = DataNodeTableCache.getInstance().getTable(database, table); - final AtomicBoolean mayContainDuplicateDevice = new AtomicBoolean(false); + final DeviceEntryFetchContext fetchContext = new DeviceEntryFetchContext(queryContext, null); if (!TreeViewSchema.isTreeViewTable(tableInstance)) { deviceEntryMap.put(database, new ArrayList<>()); } @@ -241,17 +244,16 @@ public class TableDeviceSchemaFetcher { statement, deviceEntryMap, attributeColumns, - queryContext, - mayContainDuplicateDevice, + fetchContext, false)) { fetchMissingDeviceSchemaForQuery( - database, tableInstance, attributeColumns, statement, deviceEntryMap, null, queryContext); + database, tableInstance, attributeColumns, statement, deviceEntryMap, fetchContext); } // TODO table metadata: implement deduplicate during schemaRegion execution // TODO table metadata: need further process on input predicates and transform them into // disjoint sets - return mayContainDuplicateDevice.get() + return fetchContext.mayContainDuplicateDevice() ? deviceEntryMap.entrySet().stream() .collect( Collectors.toMap( @@ -268,119 +270,79 @@ public class TableDeviceSchemaFetcher { final MPPQueryContext queryContext, final PlanNodeId planNodeId) { final TsTable tableInstance = DataNodeTableCache.getInstance().getTable(database, table); - if (TreeViewSchema.isTreeViewTable(tableInstance)) { - final Map<String, List<DeviceEntry>> deviceEntryMap = new HashMap<>(); - final AtomicBoolean mayContainDuplicateDevice = new AtomicBoolean(false); - boolean containsNonAlignedDevice = false; - final ShowDevice statement = new ShowDevice(database, table); - try (DeviceEntryMaterializer materializer = - new DeviceEntryMaterializer( - queryContext.getQueryId().getId(), - planNodeId, - CONFIG.getTableQueryDeviceEntryBatchSizeInBytes(), - true, - queryContext)) { - final boolean needRemoteFetch = - parseFilter4TraverseDevice( - tableInstance, - expressionList, - statement, - deviceEntryMap, - attributeColumns, - queryContext, - mayContainDuplicateDevice, - false); - for (List<DeviceEntry> entries : deviceEntryMap.values()) { - for (DeviceEntry entry : entries) { - appendToMaterializer(materializer, entry, queryContext, true); - if (entry instanceof NonAlignedDeviceEntry) { - containsNonAlignedDevice = true; - } - } - entries.clear(); - } - if (needRemoteFetch) { - containsNonAlignedDevice |= - fetchMissingDeviceSchemaForQuery( - database, - tableInstance, - attributeColumns, - statement, - deviceEntryMap, - materializer, - queryContext); - } - if (deviceEntryMap.size() > 1) { - throw new SemanticException( - DataNodeQueryMessages.TREE_DEVICE_VIEW_WITH_MULTIPLE_DATABASES - + deviceEntryMap.keySet() - + DataNodeQueryMessages.IS_UNSUPPORTED_YET); - } - final String resultDatabase = - deviceEntryMap.isEmpty() ? null : deviceEntryMap.keySet().iterator().next(); - return new DeviceEntryDataSetResult( - resultDatabase, materializer.finish(), containsNonAlignedDevice); - } catch (IOException e) { - throw new UncheckedIOException(e); - } - } - - final Map<String, List<DeviceEntry>> cachedEntries = new HashMap<>(); - cachedEntries.put(database, new ArrayList<>()); - final AtomicBoolean mayContainDuplicateDevice = new AtomicBoolean(false); + final DeviceEntryFetchContext fetchContext = + new DeviceEntryFetchContext(queryContext, planNodeId); final ShowDevice statement = new ShowDevice(database, table); - final boolean needRemoteFetch = - parseFilter4TraverseDevice( - tableInstance, - expressionList, - statement, - cachedEntries, - attributeColumns, - queryContext, - mayContainDuplicateDevice, - false); - - if (mayContainDuplicateDevice.get()) { + final Map<String, List<DeviceEntry>> deviceEntryMap = new HashMap<>(); + if (!TreeViewSchema.isTreeViewTable(tableInstance)) { + deviceEntryMap.put(database, new ArrayList<>()); + } + Throwable failure = null; + try { + final boolean needRemoteFetch = + parseFilter4TraverseDevice( + tableInstance, + expressionList, + statement, + deviceEntryMap, + attributeColumns, + fetchContext, + false); + final AbstractDeviceEntryMaterializer materializer = fetchContext.getMaterializer(); if (needRemoteFetch) { fetchMissingDeviceSchemaForQuery( - database, - tableInstance, - attributeColumns, - statement, - cachedEntries, - null, - queryContext); + database, tableInstance, attributeColumns, statement, deviceEntryMap, fetchContext); } - cachedEntries.put( - database, new ArrayList<>(new LinkedHashSet<>(cachedEntries.get(database)))); - } - - try (DeviceEntryMaterializer materializer = - new DeviceEntryMaterializer( - queryContext.getQueryId().getId(), - planNodeId, - CONFIG.getTableQueryDeviceEntryBatchSizeInBytes(), - true, - queryContext)) { - for (DeviceEntry entry : cachedEntries.get(database)) { - appendToMaterializer(materializer, entry, queryContext, true); + if (deviceEntryMap.size() > 1) { + throw new SemanticException( + DataNodeQueryMessages.TREE_DEVICE_VIEW_WITH_MULTIPLE_DATABASES + + deviceEntryMap.keySet() + + DataNodeQueryMessages.IS_UNSUPPORTED_YET); } - cachedEntries.get(database).clear(); - if (needRemoteFetch && !mayContainDuplicateDevice.get()) { - fetchMissingDeviceSchemaForQuery( - database, - tableInstance, - attributeColumns, - statement, - cachedEntries, - materializer, - queryContext); - } - final DeviceEntryDataSet dataSet = materializer.finish(); - return new DeviceEntryDataSetResult(database, dataSet, false); + + DeviceEntryDataSet dataSet; + dataSet = materializer.finish(); + + return new DeviceEntryDataSetResult( + deviceEntryMap.isEmpty() ? null : deviceEntryMap.keySet().iterator().next(), + dataSet, + fetchContext.containsNonAlignedDevice()); } catch (IOException e) { + failure = e; throw new UncheckedIOException(e); + } catch (Exception e) { + failure = e; + throw e; + } finally { + final AbstractDeviceEntryMaterializer materializer = fetchContext.getMaterializer(); + if (materializer != null) { + try { + materializer.close(); + } catch (IOException e) { + if (failure != null) { + failure.addSuppressed(e); + } else { + throw new UncheckedIOException(e); + } + } + } + } + } + + private AbstractDeviceEntryMaterializer createDataSetMaterializer( + MPPQueryContext queryContext, PlanNodeId planNodeId, boolean distinct) { + long batchSize = CONFIG.getTableQueryDeviceEntryBatchSizeInBytes(); + if (distinct) { + return new DeviceEntrySortedMaterializer( + queryContext.getQueryId().getId(), + planNodeId, + batchSize, + Comparator.comparing(entry -> entry.getDeviceID().toString()), + true, + queryContext); } + return new DeviceEntryMaterializer( + queryContext.getQueryId().getId(), planNodeId, batchSize, true, queryContext); } // Used by show/count device and update device. @@ -391,9 +353,9 @@ public class TableDeviceSchemaFetcher { final AbstractTraverseDevice statement, final Map<String, List<DeviceEntry>> deviceEntryMap, final List<String> attributeColumns, - final MPPQueryContext queryContext, - final AtomicBoolean mayContainDuplicateDevice, + final DeviceEntryFetchContext fetchContext, final boolean isDirectDeviceQuery) { + final MPPQueryContext queryContext = fetchContext.getQueryContext(); final Pair<List<Expression>, List<Expression>> separatedExpression = SchemaPredicateUtil.separateTagDeterminedPredicate( expressionList, tableInstance, queryContext, isDirectDeviceQuery); @@ -405,9 +367,20 @@ public class TableDeviceSchemaFetcher { // Each element represents one batch of possible devices // expressions inner each element are and-concat representing conditions of different column + final AtomicBoolean mayContainDuplicateDevice = new AtomicBoolean(false); final List<Map<Integer, List<SchemaFilter>>> index2FilterMapList = SchemaPredicateUtil.convertTagPredicateToOrConcatList( tagDeterminedPredicateList, tableInstance, mayContainDuplicateDevice); + fetchContext.setMayContainDuplicateDevice(mayContainDuplicateDevice.get()); + + if (fetchContext.getPlanNodeId() != null) { + fetchContext.setMaterializer( + createDataSetMaterializer( + queryContext, + fetchContext.getPlanNodeId(), + fetchContext.mayContainDuplicateDevice())); + } + // If a predicate branch contains comparisons for all tag columns and can use SchemaCache, we // store its index. final List<Integer> tagSingleMatchIndexList = @@ -455,7 +428,7 @@ public class TableDeviceSchemaFetcher { attributeColumns, fetchPaths, isDirectDeviceQuery, - queryContext)) { + fetchContext)) { tagSingleMatchPredicateNotInCache.add(index); } } @@ -521,7 +494,7 @@ public class TableDeviceSchemaFetcher { final List<String> attributeColumns, final List<IDeviceID> fetchPaths, final boolean isDirectDeviceQuery, - final MPPQueryContext queryContext) { + final DeviceEntryFetchContext fetchContext) { final String[] tagValues = new String[tableInstance.getTagNum()]; for (final List<SchemaFilter> schemaFilters : tagFilters.values()) { final TagFilter tagFilter = (TagFilter) schemaFilters.get(0); @@ -539,8 +512,9 @@ public class TableDeviceSchemaFetcher { fetchPaths, isDirectDeviceQuery, tagValues, - queryContext) - : tryGetTreeDeviceInCache(deviceEntryMap, tableInstance, check, fetchPaths, tagValues); + fetchContext) + : tryGetTreeDeviceInCache( + deviceEntryMap, tableInstance, check, fetchPaths, tagValues, fetchContext); } private boolean tryGetTableDeviceInCache( @@ -552,7 +526,9 @@ public class TableDeviceSchemaFetcher { final List<IDeviceID> fetchPaths, final boolean isDirectDeviceQuery, final String[] tagValues, - final MPPQueryContext queryContext) { + final DeviceEntryFetchContext fetchContext) { + final MPPQueryContext queryContext = fetchContext.getQueryContext(); + final AbstractDeviceEntryMaterializer materializer = fetchContext.getMaterializer(); final IDeviceID deviceID = convertTagValuesToDeviceID(tableInstance.getTableName(), tagValues); final Map<String, Binary> attributeMap = cache.getDeviceAttribute(database, deviceID); @@ -572,13 +548,16 @@ public class TableDeviceSchemaFetcher { // TODO table metadata: process cases that selected attr columns different from those used for // predicate if (check.test(deviceEntry)) { - deviceEntryList.add(deviceEntry); + queryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); + if (materializer == null) { + deviceEntryList.add(deviceEntry); + } else { + appendToMaterializer(materializer, deviceEntry, queryContext); + } // If we partially hit cache in direct device query, we must fetch for all the predicates // because now we do not support combining memory source and other sources if (isDirectDeviceQuery) { fetchPaths.add(deviceID); - } else { - queryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); } } return true; @@ -589,7 +568,10 @@ public class TableDeviceSchemaFetcher { final TsTable tableInstance, final Predicate<AlignedDeviceEntry> check, final List<IDeviceID> fetchPaths, - final String[] tagValues) { + final String[] tagValues, + final DeviceEntryFetchContext fetchContext) { + final MPPQueryContext queryContext = fetchContext.getQueryContext(); + final AbstractDeviceEntryMaterializer materializer = fetchContext.getMaterializer(); final IDeviceID deviceID = DataNodeTreeViewSchemaUtils.convertToIDeviceID(tableInstance, tagValues); final IDeviceSchema schema = TableDeviceSchemaCache.getInstance().getDeviceSchema(deviceID); @@ -601,12 +583,19 @@ public class TableDeviceSchemaFetcher { return false; } database = ((TreeDeviceNormalSchema) schema).getDatabase(); - deviceEntryMap - .computeIfAbsent(database, k -> new ArrayList<>()) - .add( - ((TreeDeviceNormalSchema) schema).isAligned() - ? new AlignedDeviceEntry(deviceID, new Binary[0]) - : new NonAlignedDeviceEntry(deviceID, new Binary[0])); + final DeviceEntry deviceEntry = + ((TreeDeviceNormalSchema) schema).isAligned() + ? new AlignedDeviceEntry(deviceID, new Binary[0]) + : new NonAlignedDeviceEntry(deviceID, new Binary[0]); + queryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); + if (materializer == null) { + deviceEntryMap.computeIfAbsent(database, k -> new ArrayList<>()).add(deviceEntry); + } else { + appendToMaterializer(materializer, deviceEntry, queryContext); + } + if (deviceEntry instanceof NonAlignedDeviceEntry) { + fetchContext.markContainsNonAlignedDevice(); + } return true; } @@ -619,16 +608,16 @@ public class TableDeviceSchemaFetcher { return IDeviceID.Factory.DEFAULT_FACTORY.create(deviceIdNodes); } - private boolean fetchMissingDeviceSchemaForQuery( + private void fetchMissingDeviceSchemaForQuery( final String database, final TsTable tableInstance, final List<String> attributeColumns, final ShowDevice statement, final Map<String, List<DeviceEntry>> deviceEntryMap, - final DeviceEntryMaterializer materializer, - final MPPQueryContext mppQueryContext) { + final DeviceEntryFetchContext fetchContext) { Throwable t = null; - boolean containsNonAlignedDevice = false; + final AbstractDeviceEntryMaterializer materializer = fetchContext.getMaterializer(); + final MPPQueryContext mppQueryContext = fetchContext.getQueryContext(); final long queryId = SessionManager.getInstance().requestQueryId(); // For the correctness of attribute remote update @@ -702,14 +691,8 @@ public class TableDeviceSchemaFetcher { deviceEntryMap.get(database), materializer); } else { - containsNonAlignedDevice |= - constructTreeResults( - tsBlock.get(), - columnHeaderList, - tableInstance, - mppQueryContext, - deviceEntryMap, - materializer); + constructTreeResults( + tsBlock.get(), columnHeaderList, tableInstance, deviceEntryMap, fetchContext); } } } else { @@ -722,7 +705,6 @@ public class TableDeviceSchemaFetcher { TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode()); } } - return containsNonAlignedDevice; } catch (final Throwable throwable) { t = throwable; throw throwable; @@ -743,7 +725,7 @@ public class TableDeviceSchemaFetcher { final MPPQueryContext mppQueryContext, final List<String> attributeColumns, final List<DeviceEntry> deviceEntryList, - final DeviceEntryMaterializer materializer) { + final AbstractDeviceEntryMaterializer materializer) { final Column[] columns = tsBlock.getValueColumns(); for (int i = 0; i < tsBlock.getPositionCount(); i++) { final String[] nodes = new String[tableInstance.getTagNum() + 1]; @@ -760,11 +742,11 @@ public class TableDeviceSchemaFetcher { final AlignedDeviceEntry deviceEntry = new AlignedDeviceEntry( deviceID, attributeColumns.stream().map(attributeMap::get).toArray(Binary[]::new)); + mppQueryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); if (materializer == null) { - mppQueryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); deviceEntryList.add(deviceEntry); } else { - appendToMaterializer(materializer, deviceEntry, mppQueryContext, false); + appendToMaterializer(materializer, deviceEntry, mppQueryContext); } // Only cache those exact device query // Fetch paths is null iff there are fuzzy queries related to id columns @@ -775,33 +757,27 @@ public class TableDeviceSchemaFetcher { } private static void appendToMaterializer( - DeviceEntryMaterializer materializer, + AbstractDeviceEntryMaterializer materializer, DeviceEntry deviceEntry, - MPPQueryContext queryContext, - boolean memoryAlreadyReserved) { + MPPQueryContext queryContext) { try { long releasedRamBytes = materializer.appendWithMemoryControl(deviceEntry); if (releasedRamBytes > 0) { queryContext.releaseMemoryReservedForFrontEnd(releasedRamBytes); } - if (memoryAlreadyReserved && materializer.isSpilled()) { - queryContext.releaseMemoryReservedForFrontEnd(deviceEntry.ramBytesUsed()); - } else if (!memoryAlreadyReserved && !materializer.isSpilled()) { - queryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); - } } catch (IOException e) { throw new UncheckedIOException(e); } } - private boolean constructTreeResults( + private void constructTreeResults( final TsBlock tsBlock, final List<ColumnHeader> columnHeaderList, final TsTable tableInstance, - final MPPQueryContext mppQueryContext, final Map<String, List<DeviceEntry>> deviceEntryMap, - final DeviceEntryMaterializer materializer) { - boolean containsNonAlignedDevice = false; + final DeviceEntryFetchContext fetchContext) { + final MPPQueryContext mppQueryContext = fetchContext.getQueryContext(); + final AbstractDeviceEntryMaterializer materializer = fetchContext.getMaterializer(); final Column[] columns = tsBlock.getValueColumns(); for (int i = 0; i < tsBlock.getPositionCount(); i++) { final String[] nodes = new String[tableInstance.getTagNum()]; @@ -809,23 +785,25 @@ public class TableDeviceSchemaFetcher { Collections.emptyMap(), nodes, null, columnHeaderList, columns, tableInstance, i); final IDeviceID deviceID = DataNodeTreeViewSchemaUtils.convertToIDeviceID(tableInstance, nodes); + final boolean isAligned = columns[columns.length - 2].getBoolean(i); final DeviceEntry deviceEntry = - columns[columns.length - 2].getBoolean(i) + isAligned ? new AlignedDeviceEntry(deviceID, new Binary[0]) : new NonAlignedDeviceEntry(deviceID, new Binary[0]); - containsNonAlignedDevice |= deviceEntry instanceof NonAlignedDeviceEntry; + if (!isAligned) { + fetchContext.markContainsNonAlignedDevice(); + } + mppQueryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); final List<DeviceEntry> deviceEntries = deviceEntryMap.computeIfAbsent( columns[columns.length - 1].getBinary(i).getStringValue(TSFileConfig.STRING_CHARSET), k -> new ArrayList<>()); if (materializer == null) { - mppQueryContext.reserveMemoryForFrontEnd(deviceEntry.ramBytesUsed()); deviceEntries.add(deviceEntry); } else { - appendToMaterializer(materializer, deviceEntry, mppQueryContext, false); + appendToMaterializer(materializer, deviceEntry, mppQueryContext); } } - return containsNonAlignedDevice; } private void constructNodesArrayAndAttributeMap( 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 92886b0036a..e2b5e4f70f6 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 @@ -36,9 +36,10 @@ public abstract class AbstractDeviceEntryMaterializer implements AutoCloseable { private final String queryId; private final PlanNodeId planNodeId; private final long thresholdInBytes; - private final List<DeviceEntry> bufferedEntries = new ArrayList<>(); + private List<DeviceEntry> bufferedEntries = new ArrayList<>(); private int entryCount; + private long bufferedRamBytes; private Path ownerDirectory; private boolean ownerRegistered; private boolean finished; @@ -60,6 +61,19 @@ public abstract class AbstractDeviceEntryMaterializer implements AutoCloseable { */ public abstract void append(DeviceEntry entry) throws IOException; + /** + * Appends a DeviceEntry while controlling this materializer's in-memory buffer. + * + * @return RAM bytes released when buffered entries are spilled + */ + public abstract long appendWithMemoryControl(DeviceEntry entry) throws IOException; + + public abstract boolean isSpilled(); + + public final long getBufferedRamBytes() { + return bufferedRamBytes; + } + public abstract DeviceEntryDataSet finish() throws IOException; protected final String queryId() { @@ -75,6 +89,14 @@ public abstract class AbstractDeviceEntryMaterializer implements AutoCloseable { entryCount++; } + protected final void addBufferedRamBytes(long ramBytes) { + bufferedRamBytes += ramBytes; + } + + protected final void clearBufferedRamBytes() { + bufferedRamBytes = 0; + } + protected final void incrementEntryCount() { entryCount++; } @@ -91,6 +113,11 @@ public abstract class AbstractDeviceEntryMaterializer implements AutoCloseable { return new ArrayList<>(bufferedEntries); } + protected final void replaceBufferedEntries(List<DeviceEntry> entries) { + bufferedEntries = entries; + entryCount = entries.size(); + } + protected final void sortBufferedEntries(Comparator<DeviceEntry> comparator) { bufferedEntries.sort(comparator); } @@ -116,6 +143,10 @@ public abstract class AbstractDeviceEntryMaterializer implements AutoCloseable { return entryCount; } + protected final void setEntryCount(int entryCount) { + this.entryCount = entryCount; + } + protected final void clearBuffer() { bufferedEntries.clear(); } 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 82f53bb2caf..950c34d22f3 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 @@ -30,8 +30,6 @@ public final class DeviceEntryMaterializer extends AbstractDeviceEntryMaterializ private final boolean rawSegment; private DeviceEntryDiskSpiller spiller; - // Only be used in fetchDeviceSchema, manages memory itself - private long rawBufferedRamBytes; public DeviceEntryMaterializer( String queryId, PlanNodeId planNodeId, long thresholdInBytes, boolean rawSegment) { @@ -56,17 +54,20 @@ public final class DeviceEntryMaterializer extends AbstractDeviceEntryMaterializ } /** Returns the RAM bytes released when Coordinator Raw Fetch switches to spill mode. */ + @Override public long appendWithMemoryControl(DeviceEntry entry) throws IOException { checkNotFinished(); - long ramBytesUsed = entry.ramBytesUsed(); - if (spiller == null && rawBufferedRamBytes + ramBytesUsed <= thresholdInBytes()) { - appendToBuffer(entry); - rawBufferedRamBytes += ramBytesUsed; - return 0; + if (spiller == null) { + long ramBytesUsed = entry.ramBytesUsed(); + if (getBufferedRamBytes() + ramBytesUsed <= thresholdInBytes()) { + appendToBuffer(entry); + addBufferedRamBytes(ramBytesUsed); + return 0; + } } - long releasedRamBytes = rawBufferedRamBytes; + long releasedRamBytes = getBufferedRamBytes(); ensureSpiller(); - rawBufferedRamBytes = 0; + clearBufferedRamBytes(); spiller.append(entry.serializeToBytes()); incrementEntryCount(); return releasedRamBytes; @@ -78,9 +79,10 @@ public final class DeviceEntryMaterializer extends AbstractDeviceEntryMaterializ if (spiller == null && !isBufferEmpty()) { ensureSpiller(); } - rawBufferedRamBytes = 0; + clearBufferedRamBytes(); } + @Override public boolean isSpilled() { return spiller != null; } 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 index d135139f700..bc087ad6efb 100644 --- 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 @@ -38,6 +38,7 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate private static final int MAX_MERGE_FAN_IN = 32; private final Comparator<DeviceEntry> comparator; + private final boolean distinct; private final List<List<Path>> sortedRuns = new ArrayList<>(); private Path runDirectory; @@ -47,8 +48,18 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate PlanNodeId planNodeId, long bufferSizeInBytes, Comparator<DeviceEntry> comparator) { + this(queryId, planNodeId, bufferSizeInBytes, comparator, false); + } + + public DeviceEntrySortedMaterializer( + String queryId, + PlanNodeId planNodeId, + long bufferSizeInBytes, + Comparator<DeviceEntry> comparator, + boolean distinct) { super(queryId, planNodeId, bufferSizeInBytes); this.comparator = comparator; + this.distinct = distinct; } public DeviceEntrySortedMaterializer( @@ -61,12 +72,44 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate setQueryContext(queryContext); } + public DeviceEntrySortedMaterializer( + String queryId, + PlanNodeId planNodeId, + long bufferSizeInBytes, + Comparator<DeviceEntry> comparator, + boolean distinct, + MPPQueryContext queryContext) { + this(queryId, planNodeId, bufferSizeInBytes, comparator, distinct); + setQueryContext(queryContext); + } + @Override public void append(DeviceEntry entry) throws IOException { checkNotFinished(); appendToBuffer(entry); } + @Override + public long appendWithMemoryControl(DeviceEntry entry) throws IOException { + checkNotFinished(); + long entryRamBytes = entry.ramBytesUsed(); + if (!isBufferEmpty() && getBufferedRamBytes() + entryRamBytes > thresholdInBytes()) { + long releasedRamBytes = getBufferedRamBytes(); + flushRun(); + appendToBuffer(entry); + addBufferedRamBytes(entryRamBytes); + return releasedRamBytes; + } + appendToBuffer(entry); + addBufferedRamBytes(entryRamBytes); + return 0; + } + + @Override + public boolean isSpilled() { + return !sortedRuns.isEmpty(); + } + @Override public void forceSpill() throws IOException { checkNotFinished(); @@ -83,6 +126,9 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate } if (sortedRuns.isEmpty()) { sortBufferedEntries(comparator); + if (distinct) { + deduplicateBufferedEntries(); + } DeviceEntryDataSet dataSet = new InMemoryDeviceEntryDataSet(copyBufferedEntries()); markFinished(); return dataSet; @@ -93,15 +139,17 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate List<List<Path>> finalRuns = compactRuns(new ArrayList<>(sortedRuns)); Path finalDirectory = ownerDirectory().resolve("fi"); List<Path> finalSegments; + int finalEntryCount; try (DeviceEntryDiskSpiller outputSpiller = new DeviceEntryDiskSpiller(finalDirectory, thresholdInBytes(), ioContext())) { if (finalRuns.size() == 1) { - copyRun(finalRuns.get(0), outputSpiller); + finalEntryCount = copyRun(finalRuns.get(0), outputSpiller); } else { - mergeRuns(finalRuns, outputSpiller); + finalEntryCount = mergeRuns(finalRuns, outputSpiller); } finalSegments = outputSpiller.finish(); } + setEntryCount(finalEntryCount); DeviceEntryDataSet dataSet = new SpilledDeviceEntryDataSet(queryId(), ownerDirectory(), finalSegments, entryCount()); markFinished(); @@ -132,6 +180,7 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate sortedRuns.add(runSpiller.finish()); } clearBuffer(); + clearBufferedRamBytes(); } private void ensureSpillDirectory() throws IOException { @@ -142,13 +191,21 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate runDirectory = ensureOwnerDirectory().resolve("sort-run"); } - private void copyRun(List<Path> run, DeviceEntryDiskSpiller outputSpiller) throws IOException { + private int copyRun(List<Path> run, DeviceEntryDiskSpiller outputSpiller) throws IOException { + int outputCount = 0; + DeviceEntry previous = null; try (DeviceEntryFileSpillerReader reader = new DeviceEntryFileSpillerReader(run, true, ioContext())) { while (reader.hasNext()) { - outputSpiller.append(reader.next().serializeToBytes()); + DeviceEntry entry = reader.next(); + if (!distinct || previous == null || comparator.compare(previous, entry) != 0) { + outputSpiller.append(entry.serializeToBytes()); + previous = entry; + outputCount++; + } } } + return outputCount; } private void deleteRunDirectoryBestEffort() { @@ -186,7 +243,7 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate return runs; } - private void mergeRuns(List<List<Path>> runs, DeviceEntryDiskSpiller outputSpiller) + private int mergeRuns(List<List<Path>> runs, DeviceEntryDiskSpiller outputSpiller) throws IOException { List<DeviceEntryFileSpillerReader> readers = new ArrayList<>(runs.size()); PriorityQueue<MergeElement> queue = @@ -196,6 +253,8 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate return result != 0 ? result : Integer.compare(left.readerIndex, right.readerIndex); }); Throwable failure = null; + int outputCount = 0; + DeviceEntry previous = null; try { for (int i = 0; i < runs.size(); i++) { DeviceEntryFileSpillerReader reader = @@ -207,7 +266,11 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate } while (!queue.isEmpty()) { MergeElement element = queue.poll(); - outputSpiller.append(element.entry.serializeToBytes()); + if (!distinct || previous == null || comparator.compare(previous, element.entry) != 0) { + outputSpiller.append(element.entry.serializeToBytes()); + previous = element.entry; + outputCount++; + } DeviceEntryFileSpillerReader reader = readers.get(element.readerIndex); if (reader.hasNext()) { queue.add(new MergeElement(reader.next(), element.readerIndex)); @@ -237,6 +300,19 @@ public final class DeviceEntrySortedMaterializer extends AbstractDeviceEntryMate } } } + return outputCount; + } + + private void deduplicateBufferedEntries() { + List<DeviceEntry> distinctEntries = new ArrayList<>(); + DeviceEntry previous = null; + for (DeviceEntry entry : bufferedEntries()) { + if (previous == null || comparator.compare(previous, entry) != 0) { + distinctEntries.add(entry); + previous = entry; + } + } + replaceBufferedEntries(distinctEntries); } private static final class MergeElement { 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 7e4024efe87..623d6e3f7ae 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 @@ -1311,6 +1311,7 @@ public class TableDistributedPlanGenerator node.getTreeDBName(), node.getMeasurementColumnNameMap()); scanNode.setRegionReplicaSet(regionReplicaSet); + scanNode.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); pair.left = scanNode; } @@ -1333,6 +1334,7 @@ public class TableDistributedPlanGenerator node.getTreeDBName(), node.getMeasurementColumnNameMap()); scanNode.setRegionReplicaSet(regionReplicaSet); + scanNode.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); pair.right = scanNode; } @@ -1723,9 +1725,10 @@ 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 - && node.getAggregations().values().stream() + && physicalAggregation.getAggregations().values().stream() .noneMatch(aggregation -> aggregation.isDistinct() || aggregation.hasMask())) { UnionNode unionNode = (UnionNode) childrenNodes.get(0); List<PlanNode> children = unionNode.getChildren(); @@ -1746,7 +1749,8 @@ public class TableDistributedPlanGenerator } // 2. split the aggregation into partial and final - Pair<AggregationNode, AggregationNode> splitResult = split(node, symbolAllocator, queryId); + Pair<AggregationNode, AggregationNode> splitResult = + split(physicalAggregation, symbolAllocator, queryId); AggregationNode intermediate = splitResult.right; // 3. add the aggregation node above the project node @@ -1781,7 +1785,7 @@ public class TableDistributedPlanGenerator return Collections.singletonList(splitResult.left); } - return Collections.singletonList(node); + return Collections.singletonList(physicalAggregation); } // We cannot do multi-stage Aggregate if any aggregation-function is distinct. @@ -1789,10 +1793,12 @@ public class TableDistributedPlanGenerator // MarkDistinctNode will merge all data from different child. if (node.getAggregations().values().stream() .anyMatch(aggregation -> aggregation.isDistinct() || aggregation.hasMask())) { - node.setChild( + PlanNode physicalChild = mergeChildrenViaCollectOrMergeSort( - nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes)); - return Collections.singletonList(node); + nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes); + node.setChild(physicalChild); + return Collections.singletonList( + withRateFunctionInputOrdering(node, nodeOrderingMap.get(physicalChild.getPlanNodeId()))); } Pair<AggregationNode, AggregationNode> splitResult = split(node, symbolAllocator, queryId); AggregationNode intermediate = splitResult.right; @@ -1824,6 +1830,58 @@ 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()) { @@ -1848,7 +1906,8 @@ public class TableDistributedPlanGenerator return Collections.singletonList(node); } - if (node.getCoordinatorDeviceEntryDataSet().isSpilled()) { + if (node.getCoordinatorDeviceEntryDataSet() != null + && node.getCoordinatorDeviceEntryDataSet().isSpilled()) { return constructSpilledAggregationTableScanByRegionReplicaSet( node, context, dataPartition, dbName); } @@ -2113,7 +2172,8 @@ public class TableDistributedPlanGenerator node.getMeasurementColumnNameMap())); } - if (node.getCoordinatorDeviceEntryDataSet().isSpilled()) { + if (node.getCoordinatorDeviceEntryDataSet() != null + && node.getCoordinatorDeviceEntryDataSet().isSpilled()) { return constructSpilledAggregationTreeDeviceViewScanByRegionReplicaSet( node, context, dataPartition, dbName); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AbstractTraverseDevice.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AbstractTraverseDevice.java index 310da433060..d807d565cc2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AbstractTraverseDevice.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AbstractTraverseDevice.java @@ -34,6 +34,7 @@ import org.apache.iotdb.commons.schema.filter.SchemaFilter; import org.apache.iotdb.commons.schema.table.TsTable; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.plan.relational.metadata.DeviceEntry; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.DeviceEntryFetchContext; import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.TableDeviceSchemaFetcher; import org.apache.iotdb.db.queryengine.plan.relational.planner.ir.ExtractCommonPredicatesExpressionRewriter; @@ -46,7 +47,6 @@ import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.concurrent.atomic.AtomicBoolean; import static org.apache.iotdb.db.queryengine.plan.relational.sql.ast.AbstractQueryDeviceWithCache.getDeviceColumnHeaderList; @@ -158,8 +158,7 @@ public abstract class AbstractTraverseDevice extends Statement { this, entries, attributeColumns, - context, - new AtomicBoolean(false), + new DeviceEntryFetchContext(context, null), true); } 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 20e17d195fd..12b0ab43c2a 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 @@ -196,6 +196,33 @@ public class DeviceEntryMaterializerTest { assertEquals(expected, actual); } + @Test + public void testDistinctSortedMaterializerDeduplicatesAcrossRuns() throws Exception { + List<DeviceEntry> expected = createEntries(20); + List<DeviceEntry> input = new ArrayList<>(expected); + input.addAll(expected); + Comparator<DeviceEntry> comparator = + Comparator.comparing(entry -> entry.getDeviceID().toString()); + try (DeviceEntrySortedMaterializer materializer = + new DeviceEntrySortedMaterializer( + "q-distinct", new PlanNodeId("scan-0"), 128, comparator, true)) { + for (DeviceEntry entry : input) { + materializer.appendWithMemoryControl(entry); + } + List<DeviceEntry> actual = new ArrayList<>(); + try (DeviceEntryDataSet dataSet = materializer.finish(); + DeviceEntryReader reader = dataSet.openReader()) { + assertTrue(dataSet.isSpilled()); + assertEquals(expected.size(), dataSet.getEntryCount()); + while (reader.hasNext()) { + actual.add(reader.next()); + } + } + expected.sort(comparator); + assertEquals(expected, actual); + } + } + private static List<DeviceEntry> createEntries(int count) { List<DeviceEntry> entries = new ArrayList<>(count); for (int i = 0; i < count; i++) {
