diff --git a/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java b/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java index 98686d436be12..a9823b28572a4 100644 --- a/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java +++ b/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java @@ -40,6 +40,7 @@ import java.sql.Time; import java.sql.Timestamp; import java.util.Calendar; +import java.util.Collections; import java.util.List; import java.util.Map; @@ -68,6 +69,17 @@ public ClusterTestResultSet( this.env = env; } + /** + * Returns each node's result for assertions that do not require identical row order or plan text. + * Every returned result must be checked. The caller must close the original result set and must + * not mix per-node iteration with iteration through the cluster wrapper. + */ + public static List getNodeResultSets(ResultSet resultSet) { + return resultSet instanceof ClusterTestResultSet + ? Collections.unmodifiableList(((ClusterTestResultSet) resultSet).resultSets) + : Collections.singletonList(resultSet); + } + @Override public boolean next() throws SQLException { RequestDelegate delegate = createRequestDelegate(); diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java new file mode 100644 index 0000000000000..e53ad150ca431 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java @@ -0,0 +1,120 @@ +/* + * 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.it.utils; + +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Query; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QueryBody; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QuerySpecification; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.TableSubquery; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Explain; +import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser; +import org.apache.iotdb.itbase.runtime.ClusterTestResultSet; + +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/** Explicit SQL-aware assertions for queries affected by ready-input Collect consumption. */ +public final class CollectQueryAssertions { + private CollectQueryAssertions() {} + + /** Only queries without final ordering or an arbitrary top-level slice may use this path. */ + public static List getUnorderedResultSets(String sql, ResultSet resultSet) { + org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement statement = + new SqlParser().createStatement(sql, ZoneOffset.UTC, null); + assertTrue("Unordered assertions require a SELECT query: " + sql, statement instanceof Query); + Query query = (Query) statement; + while (true) { + assertFalse( + "Final ORDER BY requires ordered assertions: " + sql, query.getOrderBy().isPresent()); + assertFalse( + "LIMIT requires deterministic result selection: " + sql, query.getLimit().isPresent()); + assertFalse( + "OFFSET requires deterministic result selection: " + sql, query.getOffset().isPresent()); + QueryBody body = query.getQueryBody(); + if (body instanceof QuerySpecification) { + QuerySpecification specification = (QuerySpecification) body; + assertFalse( + "Final ORDER BY requires ordered assertions: " + sql, + specification.getOrderBy().isPresent()); + assertFalse( + "LIMIT requires deterministic result selection: " + sql, + specification.getLimit().isPresent()); + assertFalse( + "OFFSET requires deterministic result selection: " + sql, + specification.getOffset().isPresent()); + } + if (!(body instanceof TableSubquery)) { + break; + } + // Parentheses around the whole query retain its ordering; nested FROM/UNION branches do not. + query = ((TableSubquery) body).getQuery(); + } + return nodeResults(resultSet); + } + + /** EXPLAIN text may contain node-specific details; each node must satisfy the plan assertions. */ + public static List getExplainResultSets(String sql, ResultSet resultSet) { + assertTrue( + "Plan assertions require EXPLAIN: " + sql, + new SqlParser().createStatement(sql, ZoneOffset.UTC, null) instanceof Explain); + return nodeResults(resultSet); + } + + private static List nodeResults(ResultSet resultSet) { + return resultSet instanceof ClusterTestResultSet + ? ClusterTestResultSet.getNodeResultSets(resultSet) + : Collections.singletonList(resultSet); + } + + /** Compare every node independently while retaining duplicate multiplicity. */ + public static void assertUnorderedQueryResult( + String sql, ResultSet resultSet, String expectedHeader, List expectedRows) + throws SQLException { + List sortedExpected = new ArrayList<>(expectedRows); + Collections.sort(sortedExpected); + for (ResultSet nodeResult : getUnorderedResultSets(sql, resultSet)) { + ResultSetMetaData metadata = nodeResult.getMetaData(); + StringBuilder header = new StringBuilder(); + for (int i = 1; i <= metadata.getColumnCount(); i++) { + header.append(metadata.getColumnName(i)).append(","); + } + assertEquals(expectedHeader, header.toString()); + List actualRows = new ArrayList<>(); + while (nodeResult.next()) { + StringBuilder row = new StringBuilder(); + for (int i = 1; i <= metadata.getColumnCount(); i++) { + row.append(nodeResult.getString(i)).append(","); + } + actualRows.add(row.toString()); + } + Collections.sort(actualRows); + assertEquals(sortedExpected, actualRows); + } + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java index 2aa5b49f331f6..700d552f0c9a0 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java @@ -366,7 +366,11 @@ public static String generateHeaderResults() { } public static String getQuerySql(final String table) { - return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from " + table; + // ClusterTestResultSet compares replicas row by row. All tags plus time form a unique key, + // so require the same row order even when Collect consumes inputs in readiness order. + return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from " + + table + + " order by s0,s1,s2,s3,time"; } public static String getQueryCountSql(final String table) { diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java index 0712387775b9d..5d5313d9f891f 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java @@ -54,6 +54,8 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets; + @RunWith(IoTDBTestRunner.class) @Category({MultiClusterIT3.class}) public class IoTDBShowReceiversTripleIT { @@ -214,37 +216,66 @@ private boolean hasReceiversOnDifferentDataNodes( final String secondPipeName, final int secondReceiverDataNodeId) throws SQLException { - boolean hasFirstPipe = false; - boolean hasSecondPipe = false; - final Set receiverDataNodeIds = new HashSet<>(); try (final Connection connection = receiverEnv.getConnection( SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); final Statement statement = connection.createStatement(); final ResultSet resultSet = statement.executeQuery(sql)) { - while (resultSet.next()) { - final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); - if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) - || !"thrift".equals(getString(resultSet, "Protocol", "protocol")) - || pipeIds == null) { - continue; - } - final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id"); - if (receiverDataNodeId == firstReceiverDataNodeId - && pipeIds.contains(firstPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasFirstPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); - } - if (receiverDataNodeId == secondReceiverDataNodeId - && pipeIds.contains(secondPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasSecondPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); + // Each node must see both senders, but unordered table queries may return different row + // orders. + for (final ResultSet nodeResult : + BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) + ? getUnorderedResultSets(sql, resultSet) + : Collections.singletonList(resultSet)) { + if (!hasExpectedReceivers( + nodeResult, + sqlDialect, + firstPipeName, + firstReceiverDataNodeId, + secondPipeName, + secondReceiverDataNodeId)) { + return false; } } + return true; + } + } + + private static boolean hasExpectedReceivers( + final ResultSet resultSet, + final String sqlDialect, + final String firstPipeName, + final int firstReceiverDataNodeId, + final String secondPipeName, + final int secondReceiverDataNodeId) + throws SQLException { + boolean hasFirstPipe = false; + boolean hasSecondPipe = false; + final Set receiverDataNodeIds = new HashSet<>(); + while (resultSet.next()) { + final String pipeIds = getString(resultSet, sqlDialect, "PipeIDs", "pipe_ids"); + if (!"DataNode" + .equals(getString(resultSet, sqlDialect, "ReceiverNodeType", "receiver_node_type")) + || !"thrift".equals(getString(resultSet, sqlDialect, "Protocol", "protocol")) + || pipeIds == null) { + continue; + } + final int receiverDataNodeId = + getInt(resultSet, sqlDialect, "ReceiverNodeId", "receiver_node_id"); + if (receiverDataNodeId == firstReceiverDataNodeId + && pipeIds.contains(firstPipeName + "@") + && getInt(resultSet, sqlDialect, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, sqlDialect, "PipeCount", "pipe_count") >= 1) { + hasFirstPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } + if (receiverDataNodeId == secondReceiverDataNodeId + && pipeIds.contains(secondPipeName + "@") + && getInt(resultSet, sqlDialect, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, sqlDialect, "PipeCount", "pipe_count") >= 1) { + hasSecondPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } } return hasFirstPipe && hasSecondPipe && receiverDataNodeIds.size() >= 2; } @@ -266,22 +297,22 @@ private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception } private static String getString( - final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + final ResultSet resultSet, + final String sqlDialect, + final String treeColumnName, + final String tableColumnName) throws SQLException { - try { - return resultSet.getString(treeColumnName); - } catch (final SQLException ignored) { - return resultSet.getString(tableColumnName); - } + return resultSet.getString( + BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) ? tableColumnName : treeColumnName); } private static int getInt( - final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + final ResultSet resultSet, + final String sqlDialect, + final String treeColumnName, + final String tableColumnName) throws SQLException { - try { - return resultSet.getInt(treeColumnName); - } catch (final SQLException ignored) { - return resultSet.getInt(tableColumnName); - } + return resultSet.getInt( + BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) ? tableColumnName : treeColumnName); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java index 16d92c213436d..999f16e4b0806 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java @@ -264,7 +264,7 @@ public void testMultiDeviceQueryAndDelete() { } private void testSelectAll() { - String selectSql = "select * from t"; + String selectSql = "select * from t ORDER BY tag1, tag2, time"; try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); Statement statement = connection.createStatement()) { diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java index 9db3acb13e098..25c7176ca9fc8 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java @@ -39,6 +39,7 @@ import java.util.HashSet; import java.util.Set; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; import static org.junit.Assert.fail; @@ -200,37 +201,39 @@ public void testLargeResultIsSplitWithoutDataLoss() throws Exception { "large_result", "org.apache.iotdb.db.query.udf.example.relational.LargeResultTableFunction"); - Set returnedRows = new HashSet<>(); try (Connection connection = EnvFactory.getEnv().getTableConnection(); Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE_NAME); - try (ResultSet resultSet = - statement.executeQuery( - "SELECT * FROM large_result(vehicle, " - + LARGE_RESULT_REPEAT_COUNT - + ", " - + LARGE_RESULT_PAYLOAD_SIZE - + ")")) { - while (resultSet.next()) { - int repeatIndex = resultSet.getInt("repeat_index"); - long time = resultSet.getLong("time"); - Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT); - Assert.assertEquals( - repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE), - resultSet.getString("payload")); - assertPassThroughColumns(resultSet, time); - Assert.assertTrue( - "Duplicate result row for time " + time + " and repeat index " + repeatIndex, - returnedRows.add(time + ":" + repeatIndex)); - } - } - } + String sql = + "SELECT * FROM large_result(vehicle, " + + LARGE_RESULT_REPEAT_COUNT + + ", " + + LARGE_RESULT_PAYLOAD_SIZE + + ")"; + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getUnorderedResultSets(sql, clusterResult)) { + Set returnedRows = new HashSet<>(); + while (resultSet.next()) { + int repeatIndex = resultSet.getInt("repeat_index"); + long time = resultSet.getLong("time"); + Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT); + Assert.assertEquals( + repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE), + resultSet.getString("payload")); + assertPassThroughColumns(resultSet, time); + Assert.assertTrue( + "Duplicate result row for time " + time + " and repeat index " + repeatIndex, + returnedRows.add(time + ":" + repeatIndex)); + } - long[] inputTimes = new long[] {1, 2, 3, 5}; - Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size()); - for (long time : inputTimes) { - for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) { - Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex)); + long[] inputTimes = new long[] {1, 2, 3, 5}; + Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size()); + for (long time : inputTimes) { + for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) { + Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex)); + } + } + } } } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java index 174b2d4711ae8..6c3ce543300f0 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java @@ -279,7 +279,9 @@ public void testPartialColumn() throws SQLException { statement.execute("DELETE FROM vehicle1"); statement.execute( "INSERT INTO vehicle1(time, s2, manufacturer) SELECT time, s2, manufacturer FROM vehicle0"); - resultSet = statement.executeQuery("SELECT deviceId, manufacturer, s1, s2 FROM vehicle1"); + resultSet = + statement.executeQuery( + "SELECT deviceId, manufacturer, s1, s2 FROM vehicle1 ORDER BY manufacturer, s2"); while (resultSet.next()) { for (int i = 1; i <= 4; i++) { if (i % 2 == 1) { diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java new file mode 100644 index 0000000000000..332c69477bdad --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java @@ -0,0 +1,189 @@ +/* + * 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.db.it.utils.TestUtils; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; +import org.apache.iotdb.itbase.runtime.ClusterTestResultSet; + +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import javax.sql.rowset.CachedRowSet; +import javax.sql.rowset.RowSetMetaDataImpl; +import javax.sql.rowset.RowSetProvider; + +import java.lang.reflect.Proxy; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.assertUnorderedQueryResult; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getExplainResultSets; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +/** Deterministic coverage of cross-node assertions without depending on network arrival order. */ +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class IoTDBClusterResultSetIT { + + @Test + public void testUnorderedRowsCheckEveryNode() throws SQLException { + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "b,")); + } + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "c"})) { + assertThrows( + AssertionError.class, + () -> + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "b,"))); + } + } + + @Test + public void testUnorderedRowsPreserveDuplicateMultiplicity() throws SQLException { + try (ResultSet result = cluster(new String[] {"a", "a", "b"}, new String[] {"b", "a", "a"})) { + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "a,", "b,")); + } + for (String[] values : new String[][] {{"a", "b"}, {"a", "a", "a", "b"}}) { + try (ResultSet result = cluster(new String[] {"a", "a", "b"}, values)) { + assertThrows( + AssertionError.class, + () -> + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "a,", "b,"))); + } + } + } + + @Test + public void testFinalOrderingRejectsUnorderedAssertions() throws SQLException { + for (String sql : + new String[] { + "SELECT value FROM t ORDER BY value", + "(SELECT value FROM t ORDER BY value)", + "SELECT value FROM t UNION ALL SELECT value FROM t ORDER BY value" + }) { + try (ResultSet result = rows("a")) { + assertThrows(AssertionError.class, () -> getUnorderedResultSets(sql, result)); + } + } + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { + assertThrows( + AssertionError.class, + () -> TestUtils.assertResultSetEqual(result, "value,", Arrays.asList("a,", "b,"))); + } + } + + @Test + public void testNestedOrderingDoesNotOrderFinalResult() throws SQLException { + for (String sql : + new String[] { + "(SELECT value FROM t ORDER BY value LIMIT 1) UNION ALL (SELECT value FROM t ORDER BY value LIMIT 1)", + "SELECT * FROM (SELECT value FROM t ORDER BY value)", + "SELECT row_number() OVER (ORDER BY value) FROM t", + "SELECT 'order by' FROM t /* ORDER BY value */" + }) { + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { + assertUnorderedQueryResult(sql, result, "value,", Arrays.asList("a,", "b,")); + } + } + } + + @Test + public void testArbitrarySlicesNeedDeterministicSelection() throws SQLException { + for (String sql : + new String[] {"SELECT value FROM t LIMIT 1", "SELECT value FROM t OFFSET 1"}) { + try (ResultSet result = rows("a")) { + assertThrows(AssertionError.class, () -> getUnorderedResultSets(sql, result)); + } + } + } + + @Test + public void testSingleNodeAndNullValues() throws SQLException { + try (ResultSet result = rows("a", null, "a")) { + assertEquals(1, getUnorderedResultSets("SELECT value FROM t", result).size()); + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("null,", "a,", "a,")); + } + try (ResultSet result = cluster(new String[] {"a"})) { + assertEquals(1, getUnorderedResultSets("SELECT value FROM t", result).size()); + assertUnorderedQueryResult("SELECT value FROM t", result, "value,", Arrays.asList("a,")); + } + } + + @Test + public void testExplainChecksAreExplicit() throws SQLException { + try (ResultSet result = cluster(new String[] {"plan-a"}, new String[] {"plan-b"})) { + assertThrows(AssertionError.class, () -> getExplainResultSets("SELECT value FROM t", result)); + assertEquals( + 2, getExplainResultSets("EXPLAIN SELECT value FROM t ORDER BY value", result).size()); + } + } + + private static ResultSet cluster(String[]... nodeRows) throws SQLException { + List statements = new ArrayList<>(); + List endpoints = new ArrayList<>(); + for (String[] values : nodeRows) { + ResultSet result = rows(values); + Statement statement = + (Statement) + Proxy.newProxyInstance( + Statement.class.getClassLoader(), + new Class[] {Statement.class}, + (proxy, method, arguments) -> { + if (method.getName().equals("executeQuery")) { + return result; + } + throw new UnsupportedOperationException(method.getName()); + }); + statements.add(statement); + endpoints.add("node-" + endpoints.size()); + } + return new ClusterTestResultSet(statements, endpoints, "SELECT value", 10, null); + } + + private static CachedRowSet rows(String... values) throws SQLException { + CachedRowSet result = RowSetProvider.newFactory().createCachedRowSet(); + RowSetMetaDataImpl metadata = new RowSetMetaDataImpl(); + metadata.setColumnCount(1); + metadata.setColumnName(1, "value"); + metadata.setColumnType(1, Types.VARCHAR); + result.setMetaData(metadata); + for (String value : values) { + result.afterLast(); + result.moveToInsertRow(); + result.updateString(1, value); + result.insertRow(); + result.moveToCurrentRow(); + } + result.beforeFirst(); + return result; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java new file mode 100644 index 0000000000000..72e75a9cfb9c3 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java @@ -0,0 +1,362 @@ +/* + * 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.db.it.utils.TestUtils; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; +import org.apache.iotdb.pipe.it.dual.tablemodel.TableModelUtils; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.assertUnorderedQueryResult; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class IoTDBCollectFanInIT { + private static final String DATABASE = "collect_fan_in"; + private static final int DEVICES = 24; + private static final int ROWS = 8; + + @BeforeClass + public static void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setDegreeOfParallelism(1) + .setMaxTsBlockLineNumber(2) + .setSortBufferSize(128 * 1024) + .setDataRegionGroupExtensionPolicy("CUSTOM") + .setDataRegionPerDataNode(3); + EnvFactory.getEnv().initClusterEnvironment(); + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE DATABASE " + DATABASE); + statement.execute("USE " + DATABASE); + for (String table : new String[] {"readings1", "readings2"}) { + statement.execute( + "CREATE TABLE " + table + " (device STRING TAG, zone STRING TAG, value INT32 FIELD)"); + for (int device = 0; device < DEVICES; device++) { + for (int value = 1; value <= ROWS; value++) { + statement.execute( + "INSERT INTO " + + table + + " (time,device,zone,value) VALUES (" + + value + + ",'" + + device(device) + + "'," + + (zone(device) == null ? "NULL" : "'" + zone(device) + "'") + + "," + + value + + ")"); + } + } + } + statement.execute("FLUSH"); + } + } + + @Test + public void testUnorderedUnionPreservesEveryDuplicateOnEachNode() throws Exception { + List expectedRows = new ArrayList<>(); + for (int device = 0; device < DEVICES; device++) { + for (int value = 1; value <= ROWS; value++) { + String row = device(device) + "," + zone(device) + "," + value + ","; + expectedRows.add(row); + expectedRows.add(row); + } + } + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + String sql = + "SELECT device, zone, value FROM readings1 UNION ALL " + + "SELECT device, zone, value FROM readings1"; + try (ResultSet result = statement.executeQuery(sql)) { + assertUnorderedQueryResult(sql, result, "device,zone,value,", expectedRows); + } + } + } + + @Test + public void testPipeVerificationQueryUsesDeterministicPrimaryKeyOrder() throws Exception { + TableModelUtils.createDataBaseAndTable(EnvFactory.getEnv(), "pipe_readings", DATABASE); + TableModelUtils.insertData(DATABASE, "pipe_readings", 0, DEVICES * ROWS, EnvFactory.getEnv()); + List expected = + new ArrayList<>(TableModelUtils.generateExpectedResults(0, DEVICES * ROWS)); + Collections.sort(expected); + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery(TableModelUtils.getQuerySql("pipe_readings"))) { + TestUtils.assertResultSetEqual(result, TableModelUtils.generateHeaderResults(), expected); + } + } + } + + @Test + public void testLegacyTreeAndSharedTreeToTablePaths() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE DATABASE root.collect_tree"); + for (int device = 0; device < DEVICES; device++) { + String path = "root.collect_tree.z" + device % 3 + "." + device(device); + statement.execute("CREATE ALIGNED TIMESERIES " + path + "(value INT32)"); + for (int value = 1; value <= ROWS; value++) { + statement.execute( + "INSERT INTO " + path + "(time,value) ALIGNED VALUES (" + value + "," + value + ")"); + } + } + try (ResultSet result = + statement.executeQuery("SELECT value FROM root.collect_tree.*.* ALIGN BY DEVICE")) { + int rows = 0; + Set seenDevices = new HashSet<>(); + String currentDevice = null; + int deviceRows = 0; + while (result.next()) { + String nextDevice = result.getString("Device"); + if (!nextDevice.equals(currentDevice)) { + if (currentDevice != null) { + assertEquals(ROWS, deviceRows); + } + assertTrue("Tree results must retain contiguous devices", seenDevices.add(nextDevice)); + currentDevice = nextDevice; + deviceRows = 0; + } + assertEquals(++deviceRows, result.getLong("Time")); + assertEquals(result.getLong("Time"), result.getLong("value")); + rows++; + } + assertEquals(ROWS, deviceRows); + assertEquals(DEVICES, seenDevices.size()); + assertEquals(DEVICES * ROWS, rows); + } + } + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + statement.execute( + "CREATE VIEW tree_readings(zone STRING TAG, device STRING TAG, value INT32 FIELD) AS root.collect_tree.**"); + try (ResultSet result = + statement.executeQuery( + "SELECT device, SUM(value) FROM (" + + "SELECT device, value FROM tree_readings UNION ALL SELECT device, value FROM tree_readings" + + ") u GROUP BY device ORDER BY device")) { + for (int device = 0; device < DEVICES; device++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(72.0, result.getDouble(2), 0.0); + } + assertFalse(result.next()); + } + } + } + + @AfterClass + public static void tearDown() { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + private static String device(int id) { + return "d" + (id < 10 ? "0" : "") + id; + } + + private static String zone(int id) { + return id % 3 == 0 ? null : "z" + id % 3; + } + + @Test + public void testGroupedAggregationAfterGlobalWindow() throws Exception { + checkGroups( + "SELECT device, zone, value FROM (SELECT device, zone, value, " + + "COUNT(*) OVER () AS n FROM readings1) t WHERE n > 0", + 1); + } + + @Test + public void testUnionPartialFinalAndDistinctMask() throws Exception { + String input = + "SELECT device, zone, value FROM readings1 UNION ALL " + + "SELECT device, zone, value FROM readings2"; + checkGroups(input, 2); + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + // Without DISTINCT, this query exercises branch PARTIAL and post-Collect FINAL aggregation. + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, SUM(value) FROM (" + + input + + ") u GROUP BY device, zone ORDER BY device")) { + for (int device = 0; device < DEVICES; device++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(zone(device), result.getString(2)); + assertEquals(72.0, result.getDouble(3), 0.0); + } + assertFalse(result.next()); + } + } + } + + @Test + public void testOrderedFanInAndJoinDependencies() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + for (String limit : new String[] {"", " LIMIT 13"}) { + try (ResultSet result = + statement.executeQuery( + "SELECT device, value FROM (SELECT device, value FROM readings1 UNION ALL " + + "SELECT device, value FROM readings2) u ORDER BY device, value" + + limit)) { + int expectedRows = limit.isEmpty() ? DEVICES * ROWS * 2 : 13; + for (int row = 0; row < expectedRows; row++) { + assertTrue(result.next()); + assertEquals(device(row / (ROWS * 2)), result.getString(1)); + assertEquals(row / 2 % ROWS + 1, result.getInt(2)); + } + assertFalse(result.next()); + } + } + try (ResultSet result = + statement.executeQuery( + "SELECT COUNT(*) FROM readings1 a JOIN readings2 b " + + "ON a.device = b.device AND a.time = b.time WHERE a.value = b.value")) { + assertTrue(result.next()); + assertEquals(DEVICES * ROWS, result.getLong(1)); + assertFalse(result.next()); + } + } + } + + private void checkGroups(String input, int copies) throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, SUM(value), " + + "COUNT(*), COUNT(DISTINCT value) FROM (" + + input + + ") u GROUP BY device, zone ORDER BY device")) { + for (int device = 0; device < DEVICES; device++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(zone(device), result.getString(2)); + assertEquals(36.0 * copies, result.getDouble(3), 0.0); + assertEquals((long) ROWS * copies, result.getLong(4)); + assertEquals(ROWS, result.getLong(5)); + } + assertFalse(result.next()); + } + } + } + + @Test + public void testNestedWindowAfterCollectRestoresCompletePartitions() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, value, " + + "SUM(value) OVER (PARTITION BY device, zone ORDER BY value) AS running_sum FROM (" + + "SELECT device, zone, value, COUNT(*) OVER () AS n FROM readings1) t " + + "WHERE n > 0 ORDER BY device, value")) { + for (int device = 0; device < DEVICES; device++) { + for (int value = 1; value <= ROWS; value++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(zone(device), result.getString(2)); + assertEquals(value, result.getInt(3)); + assertEquals((double) value * (value + 1) / 2, result.getDouble(4), 0.0); + } + } + assertFalse(result.next()); + } + } + } + + @Test + public void testOrderByAfterGlobalWindowAndConstantFill() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, value, COUNT(*) OVER () AS n " + + "FROM readings1 FILL METHOD CONSTANT 0 ORDER BY device, value DESC")) { + for (int device = 0; device < DEVICES; device++) { + for (int value = ROWS; value > 0; value--) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(value, result.getInt(3)); + assertEquals((long) DEVICES * ROWS, result.getLong(4)); + } + } + assertFalse(result.next()); + } + } + } + + @Test + public void testExcludedTopKPlanRetainsQueryResults() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + int limit = DEVICES + 5; + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, value, COUNT(*) OVER () AS n FROM readings1 " + + "ORDER BY value DESC, device LIMIT " + + limit)) { + for (int row = 0; row < limit; row++) { + assertTrue(result.next()); + assertEquals(device(row % DEVICES), result.getString(1)); + assertEquals(zone(row % DEVICES), result.getString(2)); + assertEquals(ROWS - row / DEVICES, result.getInt(3)); + assertEquals((long) DEVICES * ROWS, result.getLong(4)); + } + assertFalse(result.next()); + } + } + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java index 8c0b52acc944b..7ada8c567bd6d 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java @@ -57,6 +57,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getExplainResultSets; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; import static org.junit.Assert.assertEquals; @@ -362,19 +363,22 @@ public void testNestExplain1() throws SQLException { statement.execute("USE testdb"); // explain - ResultSet resultSet = statement.executeQuery(sql); - ResultSetMetaData metaData = resultSet.getMetaData(); - assertEquals(metaData.getColumnCount(), 1); - assertEquals(metaData.getColumnName(1), "distribution plan"); - - StringBuilder sb = new StringBuilder(); - while (resultSet.next()) { - sb.append(resultSet.getString(1)).append(System.lineSeparator()); + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getExplainResultSets(sql, clusterResult)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + assertEquals(metaData.getColumnCount(), 1); + assertEquals(metaData.getColumnName(1), "distribution plan"); + + StringBuilder sb = new StringBuilder(); + while (resultSet.next()) { + sb.append(resultSet.getString(1)).append(System.lineSeparator()); + } + String result = sb.toString(); + assertFalse(result.contains("CTE Query : 'cte1'")); + assertTrue(result.contains("CTE Query : 'cte2'")); + assertTrue(result.contains("Main Query")); + } } - String result = sb.toString(); - assertFalse(result.contains("CTE Query : 'cte1'")); - assertTrue(result.contains("CTE Query : 'cte2'")); - assertTrue(result.contains("Main Query")); } } @@ -389,19 +393,22 @@ public void testNestExplain2() throws SQLException { statement.execute("USE testdb"); // explain - ResultSet resultSet = statement.executeQuery(sql); - ResultSetMetaData metaData = resultSet.getMetaData(); - assertEquals(metaData.getColumnCount(), 1); - assertEquals(metaData.getColumnName(1), "distribution plan"); - - StringBuilder sb = new StringBuilder(); - while (resultSet.next()) { - sb.append(resultSet.getString(1)).append(System.lineSeparator()); + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getExplainResultSets(sql, clusterResult)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + assertEquals(metaData.getColumnCount(), 1); + assertEquals(metaData.getColumnName(1), "distribution plan"); + + StringBuilder sb = new StringBuilder(); + while (resultSet.next()) { + sb.append(resultSet.getString(1)).append(System.lineSeparator()); + } + String result = sb.toString(); + assertTrue(result.contains("CTE Query : 'cte1'")); + assertTrue(result.contains("CTE Query : 'cte2'")); + assertTrue(result.contains("Main Query")); + } } - String result = sb.toString(); - assertTrue(result.contains("CTE Query : 'cte1'")); - assertTrue(result.contains("CTE Query : 'cte2'")); - assertTrue(result.contains("Main Query")); } } @@ -416,19 +423,22 @@ public void testExplainJoin() throws SQLException { statement.execute("USE testdb"); // explain - ResultSet resultSet = statement.executeQuery(sql); - ResultSetMetaData metaData = resultSet.getMetaData(); - assertEquals(metaData.getColumnCount(), 1); - assertEquals(metaData.getColumnName(1), "distribution plan"); - - StringBuilder sb = new StringBuilder(); - while (resultSet.next()) { - sb.append(resultSet.getString(1)).append(System.lineSeparator()); + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getExplainResultSets(sql, clusterResult)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + assertEquals(metaData.getColumnCount(), 1); + assertEquals(metaData.getColumnName(1), "distribution plan"); + + StringBuilder sb = new StringBuilder(); + while (resultSet.next()) { + sb.append(resultSet.getString(1)).append(System.lineSeparator()); + } + String result = sb.toString(); + assertFalse(result.contains("CTE Query : 'cte1'")); + assertTrue(result.contains("CTE Query : 'cte2'")); + assertTrue(result.contains("Main Query")); + } } - String result = sb.toString(); - assertFalse(result.contains("CTE Query : 'cte1'")); - assertTrue(result.contains("CTE Query : 'cte2'")); - assertTrue(result.contains("Main Query")); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java index 518b3da1fdd21..d707cd65bdc51 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java @@ -41,6 +41,7 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getExplainResultSets; import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; @@ -5902,19 +5903,23 @@ private void assertRateQueryUsesIntermediateAggregation(String query) { try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE_NAME); - try (ResultSet resultSet = statement.executeQuery("EXPLAIN (FORMAT JSON) " + query)) { - Assert.assertTrue(resultSet.next()); - String plan = resultSet.getString(1); - // The PARTIAL/FINAL split only exists when the scanned data spans multiple DataRegions. - // Under randomized partition allocation (e.g. the SHUFFLE strategy used by - // IoTDBTableAggregation2IT), all devices may occasionally land in one single region, - // where a SINGLE-step aggregation is the correct plan, so the check must be skipped. - if (countDistinctRegionIds(plan) < 2) { - return; + try (ResultSet clusterResult = statement.executeQuery("EXPLAIN (FORMAT JSON) " + query)) { + for (ResultSet resultSet : + getExplainResultSets("EXPLAIN (FORMAT JSON) " + query, clusterResult)) { + Assert.assertTrue(resultSet.next()); + String plan = resultSet.getString(1); + // The PARTIAL/FINAL split only exists when the scanned data spans multiple DataRegions. + // Under randomized partition allocation (e.g. the SHUFFLE strategy used by + // IoTDBTableAggregation2IT), all devices may occasionally land in one single region, + // where a SINGLE-step aggregation is the correct plan, so the check must be skipped. + if (countDistinctRegionIds(plan) < 2) { + continue; + } + Assert.assertTrue( + "Expected a PARTIAL aggregation in plan: " + plan, plan.contains("PARTIAL")); + Assert.assertTrue( + "Expected a FINAL aggregation in plan: " + plan, plan.contains("FINAL")); } - Assert.assertTrue( - "Expected a PARTIAL aggregation in plan: " + plan, plan.contains("PARTIAL")); - Assert.assertTrue("Expected a FINAL aggregation in plan: " + plan, plan.contains("FINAL")); } } catch (Exception exception) { throw new AssertionError(exception); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java index 7697bc81649bd..59499aa5db0de 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java @@ -140,7 +140,8 @@ public void unionAllSiblingTopKOrderByTimeLimit() { tableResultSetEqualTest( "(SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd1' ORDER BY time DESC LIMIT 1)" + " UNION ALL" - + " (SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd3' ORDER BY time DESC LIMIT 1)", + + " (SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd3' ORDER BY time DESC LIMIT 1)" + + " ORDER BY time, device_id, s1", expectedHeader, retArray, TABLE_DATABASE); diff --git a/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java b/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java index 750def8c4a3b0..33976ed731f22 100644 --- a/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java +++ b/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java @@ -582,4 +582,7 @@ private CalcMessages() {} EXCEPTION_AGGREGATE_FUNCTION_ARG_REQUIRES_ARG_ARGUMENTS_BUT_GOT_ARG_D6249DD2 = "Aggregate function [%s] requires %d arguments, but got %d"; + public static final String + EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4 = + "Ready input consumption must be enabled before Collect execution starts"; } diff --git a/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java b/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java index 06ac8d62cbbfd..19a56413d8b0f 100644 --- a/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java +++ b/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java @@ -558,4 +558,7 @@ private CalcMessages() {} public static final String EXCEPTION_AGGREGATE_FUNCTION_ARG_REQUIRES_ARG_ARGUMENTS_BUT_GOT_ARG_D6249DD2 = "聚合函数 [%s] 需要 %d 个参数，但实际传入 %d 个"; + public static final String + EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4 = + "必须在 Collect 开始执行前启用就绪输入消费"; } diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java new file mode 100644 index 0000000000000..c1ef2edb1b619 --- /dev/null +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java @@ -0,0 +1,278 @@ +/* + * 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.calc.execution.operator.process; + +import org.apache.iotdb.calc.execution.operator.Operator; + +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.utils.Accountable; +import org.apache.tsfile.utils.RamUsageEstimator; + +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; + +import static com.google.common.util.concurrent.Futures.getDone; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; + +/** + * Observes direct children without calling operators from completion callbacks. The driver owns + * selection/invalidation; the lock protects listener records shared with asynchronous completions. + */ +final class AnyChildBlocked implements Accountable { + private static final long INSTANCE_SIZE = + RamUsageEstimator.shallowSizeOfInstance(AnyChildBlocked.class); + private static final long LISTENER_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ChildFutureListener.class); + private static final long ROUND_SIZE = + RamUsageEstimator.shallowSizeOfInstance(WaitRound.class) + + RamUsageEstimator.shallowSizeOfInstance(SettableFuture.class); + + private final List children; + private final Object lock = new Object(); + private final ChildFutureListener[] listeners; + private ChildFutureListener readyHead; + private ChildFutureListener readyTail; + private WaitRound currentRound; + private Throwable firstFailure; + private boolean closed; + private boolean initialized; + private int childToObserve = -1; + + AnyChildBlocked(List children) { + this.children = children; + this.listeners = new ChildFutureListener[children.size()]; + } + + ListenableFuture getAnyChild() { + WaitRound round; + WaitRound selection; + boolean initialize; + int invalidatedChild; + synchronized (lock) { + if (currentRound != null) { + return currentRound.selected; + } + round = new WaitRound(); + currentRound = round; + initialize = !initialized; + initialized = true; + invalidatedChild = childToObserve; + childToObserve = -1; + selection = selectReadyChild(); + } + completeSelection(selection); + + if (initialize) { + // Activate every direct input once, even if an earlier input is already ready. Pending + // listeners enqueue their completion without being retargeted or polled on every block. + for (int childId = 0; childId < children.size(); childId++) { + observeChild(childId); + } + } else if (invalidatedChild >= 0) { + // Only advancing an input invalidates its readiness. Requeue it behind inputs that are + // already ready so that a continuously ready input cannot starve its siblings. + observeChild(invalidatedChild); + } + return round.selected; + } + + private void observeChild(int childId) { + Operator child; + synchronized (lock) { + if (closed) { + return; + } + child = children.get(childId); + } + if (child == null) { + return; + } + ListenableFuture observedFuture = child.isBlocked(); + ChildFutureListener listener = new ChildFutureListener(childId, observedFuture); + synchronized (lock) { + if (closed) { + return; + } + listeners[childId] = listener; + } + // Publish the record before registration, including when completion runs inline. Callbacks + // only update readiness; operator methods remain confined to the driver thread. + observedFuture.addListener(listener, directExecutor()); + } + + /** Called under lock; claiming a selection and removing its ready record are atomic. */ + private WaitRound selectReadyChild() { + if (currentRound == null || currentRound.claimed) { + return null; + } + if (closed || children.isEmpty()) { + currentRound.childId = -1; + } else if (firstFailure != null) { + currentRound.failure = firstFailure; + } else { + ChildFutureListener listener; + do { + listener = readyHead; + if (listener == null) { + return null; + } + readyHead = listener.nextReady; + listener.nextReady = null; + if (readyHead == null) { + readyTail = null; + } + // Retired records are discarded once, rather than scanning closed input slots per block. + } while (listeners[listener.childId] != listener); + currentRound.childId = listener.childId; + } + currentRound.claimed = true; + return currentRound; + } + + private static void completeSelection(WaitRound selection) { + if (selection == null) { + return; + } + // Completing the future can invoke downstream listeners synchronously. Do not hold lock. + if (selection.failure == null) { + selection.selected.set(selection.childId); + } else { + selection.selected.setException(selection.failure); + } + } + + void checkFailure() throws Exception { + Throwable failure; + synchronized (lock) { + failure = firstFailure; + } + if (failure instanceof Exception) { + throw (Exception) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure != null) { + throw new RuntimeException(failure); + } + } + + void invalidateChild(int childId) { + synchronized (lock) { + listeners[childId] = null; + currentRound = null; + childToObserve = childId; + } + } + + void finishChild(int childId) { + invalidateChild(childId); + } + + void close() { + WaitRound round; + synchronized (lock) { + closed = true; + round = currentRound; + currentRound = null; + childToObserve = -1; + while (readyHead != null) { + ChildFutureListener next = readyHead.nextReady; + readyHead.nextReady = null; + readyHead = next; + } + readyTail = null; + Arrays.fill(listeners, null); + } + if (round != null) { + round.selected.set(-1); + } + } + + @Override + public long ramBytesUsed() { + synchronized (lock) { + long size = + INSTANCE_SIZE + + RamUsageEstimator.shallowSizeOfInstance(Object.class) + + RamUsageEstimator.shallowSizeOf(listeners); + for (ChildFutureListener listener : listeners) { + if (listener != null) { + size += LISTENER_SIZE; + } + } + return size + (currentRound == null ? 0 : ROUND_SIZE); + } + } + + private static final class WaitRound { + private final SettableFuture selected = SettableFuture.create(); + private boolean claimed; + private int childId; + private Throwable failure; + } + + private final class ChildFutureListener implements Runnable { + private final int childId; + private final ListenableFuture observedFuture; + private ChildFutureListener nextReady; + private boolean completed; + + private ChildFutureListener(int childId, ListenableFuture observedFuture) { + this.childId = childId; + this.observedFuture = observedFuture; + } + + @Override + public void run() { + Throwable failure = null; + try { + getDone(observedFuture); + } catch (ExecutionException e) { + failure = e.getCause(); + } catch (CancellationException e) { + failure = e; + } + WaitRound selection; + synchronized (lock) { + if (closed || listeners[childId] != this || completed) { + return; + } + completed = true; + if (failure != null && firstFailure == null) { + firstFailure = failure; + } + if (failure == null) { + if (readyTail == null) { + readyHead = this; + } else { + readyTail.nextReady = this; + } + readyTail = this; + } + selection = selectReadyChild(); + } + completeSelection(selection); + } + } +} diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java index 921d273d04e8a..e94139ff0cdc2 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java @@ -21,6 +21,7 @@ import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.i18n.CalcMessages; import org.apache.iotdb.commons.queryengine.execution.MemoryEstimationHelper; import org.apache.iotdb.commons.utils.TestOnly; @@ -30,60 +31,145 @@ import java.util.List; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.util.concurrent.Futures.getDone; +import static com.google.common.util.concurrent.Futures.nonCancellationPropagating; + public class CollectOperator implements ProcessOperator { private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(CollectOperator.class); private final CommonOperatorContext operatorContext; protected final List children; - private boolean inited = false; - + private AnyChildBlocked anyChildBlocked; + private boolean inited; + private boolean executionStarted; protected int currentIndex; + private ListenableFuture selectionFuture; + private ListenableFuture exposedFuture; + private int remainingChildren; + private boolean closed; public CollectOperator(CommonOperatorContext operatorContext, List children) { this.operatorContext = operatorContext; this.children = children; - this.currentIndex = 0; + this.remainingChildren = children.size(); + } + + /** + * Enable only after the complete table FI has passed eligibility checks, before drivers start. + */ + public void enableReadyInputConsumption() { + if (anyChildBlocked == null) { + checkState( + !executionStarted, + CalcMessages + .EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4); + anyChildBlocked = new AnyChildBlocked(children); + } } @Override - public boolean hasNext() throws Exception { - return currentIndex < children.size(); + public boolean hasNext() { + return anyChildBlocked == null + ? currentIndex < children.size() + : !closed && remainingChildren > 0; } @Override public TsBlock next() throws Exception { - if (children.get(currentIndex).hasNextWithTimer()) { - return children.get(currentIndex).nextWithTimer(); - } else { + executionStarted = true; + if (anyChildBlocked == null) { + if (children.get(currentIndex).hasNextWithTimer()) { + TsBlock block = children.get(currentIndex).nextWithTimer(); + return block == null ? null : transformBlock(currentIndex, block); + } closeCurrentChild(currentIndex); currentIndex++; return null; } + anyChildBlocked.checkFailure(); + if (!hasNext()) { + return null; + } + if (selectionFuture == null) { + isBlocked(); + } + anyChildBlocked.checkFailure(); + if (!selectionFuture.isDone()) { + return null; + } + int childId = getDone(selectionFuture); + // Consume this selection before invoking any operator method. Another next must acquire a + // new readiness observation even when this advancement returns null or an empty block. + selectionFuture = null; + exposedFuture = null; + if (childId < 0 || closed || children.get(childId) == null) { + return null; + } + Operator child = children.get(childId); + try { + if (child.hasNextWithTimer()) { + TsBlock block = child.nextWithTimer(); + return block == null ? null : transformBlock(childId, block); + } + if (child.isFinished()) { + closeCurrentChild(childId); + } + return null; + } finally { + anyChildBlocked.invalidateChild(childId); + } + } + + protected TsBlock transformBlock(int childId, TsBlock block) { + return block; } protected void closeCurrentChild(int index) throws Exception { - children.get(index).close(); + if (anyChildBlocked == null) { + children.get(index).close(); + children.set(index, null); + return; + } + Operator child = children.get(index); + if (child == null) { + return; + } + // Retire the record even if closing its subtree throws. + anyChildBlocked.finishChild(index); children.set(index, null); + remainingChildren--; + child.close(); } @Override public ListenableFuture isBlocked() { - if (!inited) { - inited = true; - for (Operator child : children) { - child.isBlocked(); + executionStarted = true; + if (anyChildBlocked == null) { + if (!inited) { + inited = true; + for (Operator child : children) { + child.isBlocked(); + } } + return currentIndex >= children.size() ? NOT_BLOCKED : children.get(currentIndex).isBlocked(); } - if (currentIndex >= children.size()) { + if (!hasNext()) { return NOT_BLOCKED; } - return children.get(currentIndex).isBlocked(); + if (selectionFuture == null) { + selectionFuture = anyChildBlocked.getAnyChild(); + } + if (exposedFuture == null || exposedFuture.isCancelled()) { + exposedFuture = nonCancellationPropagating(selectionFuture); + } + return exposedFuture; } @Override - public boolean isFinished() throws Exception { - return currentIndex >= children.size(); + public boolean isFinished() { + return !hasNext(); } @Override @@ -93,35 +179,86 @@ public CommonOperatorContext getOperatorContext() { @Override public void close() throws Exception { - for (int i = currentIndex, n = children.size(); i < n; i++) { - Operator currentChild = children.get(i); - if (currentChild != null) { + executionStarted = true; + if (anyChildBlocked == null) { + for (int i = currentIndex; i < children.size(); i++) { + if (children.get(i) != null) { + closeCurrentChild(i); + } + } + return; + } + if (closed) { + return; + } + closed = true; + anyChildBlocked.close(); + selectionFuture = null; + exposedFuture = null; + Exception firstException = null; + for (int i = 0; i < children.size(); i++) { + try { closeCurrentChild(i); + } catch (Exception e) { + if (firstException == null) { + firstException = e; + } else if (firstException != e) { + firstException.addSuppressed(e); + } } } + if (firstException != null) { + throw firstException; + } } @Override public long calculateMaxPeekMemory() { - long maxPeekMemory = 0; + if (anyChildBlocked == null) { + long maxPeekMemory = 0; + for (Operator child : children) { + maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter()); + } + return maxPeekMemory; + } + long retained = calculateRetainedSizeAfterCallingNext(); + long peak = retained; for (Operator child : children) { - maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter()); + if (child != null) { + peak = + Math.max( + peak, + retained + - child.calculateRetainedSizeAfterCallingNext() + + child.calculateMaxPeekMemoryWithCounter()); + } } - return maxPeekMemory; + return peak; } @Override public long calculateMaxReturnSize() { long maxReturnSize = 0; for (Operator child : children) { - maxReturnSize = Math.max(maxReturnSize, child.calculateMaxReturnSize()); + if (child != null) { + maxReturnSize = Math.max(maxReturnSize, child.calculateMaxReturnSize()); + } } return maxReturnSize; } @Override public long calculateRetainedSizeAfterCallingNext() { - return 0L; + if (anyChildBlocked == null) { + return 0; + } + long retained = 0; + for (Operator child : children) { + if (child != null) { + retained += child.calculateRetainedSizeAfterCallingNext(); + } + } + return retained; } @TestOnly @@ -132,6 +269,7 @@ public List getChildren() { @Override public long ramBytesUsed() { return INSTANCE_SIZE + + (anyChildBlocked == null ? 0 : anyChildBlocked.ramBytesUsed()) + children.stream() .mapToLong(MemoryEstimationHelper::getEstimatedSizeOfAccountableObject) .sum() diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java index 8bc1c18e9e66a..02850bf0caf16 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java @@ -45,37 +45,43 @@ public MappingCollectOperator( List> mappings) { super(operatorContext, children); this.mappings = mappings; - outputColumnsCount = mappings.get(0).size(); + outputColumnsCount = mappings.isEmpty() ? 0 : mappings.get(0).size(); } @Override - public TsBlock next() throws Exception { - if (children.get(currentIndex).hasNextWithTimer()) { - TsBlock tsBlock = children.get(currentIndex).nextWithTimer(); - if (tsBlock == null) { - return null; - } else { - Column[] columns = new Column[outputColumnsCount]; - List mapping = mappings.get(currentIndex); - for (int i = 0; i < columns.length; i++) { - columns[i] = tsBlock.getColumn(mapping.get(i)); - } - return TsBlock.wrapBlocksWithoutCopy( - tsBlock.getPositionCount(), - new RunLengthEncodedColumn( - CommonOperatorUtils.TIME_COLUMN_TEMPLATE, tsBlock.getPositionCount()), - columns); - } - } else { - closeCurrentChild(currentIndex); - currentIndex++; - return null; + protected TsBlock transformBlock(int childId, TsBlock tsBlock) { + Column[] columns = new Column[outputColumnsCount]; + List mapping = mappings.get(childId); + for (int i = 0; i < columns.length; i++) { + columns[i] = tsBlock.getColumn(mapping.get(i)); } + return TsBlock.wrapBlocksWithoutCopy( + tsBlock.getPositionCount(), + new RunLengthEncodedColumn( + CommonOperatorUtils.TIME_COLUMN_TEMPLATE, tsBlock.getPositionCount()), + columns); } + @Override protected void closeCurrentChild(int index) throws Exception { - children.get(index).close(); - children.set(index, null); - mappings.set(index, null); + try { + super.closeCurrentChild(index); + } finally { + mappings.set(index, null); + } + } + + @Override + public long ramBytesUsed() { + long size = + super.ramBytesUsed() + + INSTANCE_SIZE + - RamUsageEstimator.shallowSizeOfInstance(CollectOperator.class); + for (List mapping : mappings) { + if (mapping != null) { + size += RamUsageEstimator.shallowSizeOf(mapping.toArray()); + } + } + return size; } } diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java new file mode 100644 index 0000000000000..ac05610712aaa --- /dev/null +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java @@ -0,0 +1,868 @@ +/* + * 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.calc.execution.operator.process; + +import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; +import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; + +import com.google.common.util.concurrent.ForwardingListenableFuture; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.junit.Test; + +import java.io.IOException; +import java.util.AbstractList; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class CollectOperatorTest { + private static TsBlock block(int... values) { + TsBlockBuilder builder = + new TsBlockBuilder(Collections.nCopies(values.length, TSDataType.INT32)); + builder.getTimeColumnBuilder().writeLong(0); + for (int i = 0; i < values.length; i++) { + builder.getColumnBuilder(i).writeInt(values[i]); + } + builder.declarePosition(); + return builder.build(); + } + + private static CollectOperator collect(Operator... children) { + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(children))); + collect.enableReadyInputConsumption(); + return collect; + } + + @Test + public void testReadyChildAndStableSelectionDuringBackpressure() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = collect(a, b); + ListenableFuture selected = collect.isBlocked(); + assertTrue(selected.isDone()); + a.ready(); + assertSame(selected, collect.isBlocked()); + assertSame(b.output, collect.next()); + assertEquals(0, a.advances); + assertEquals(1, b.advances); + assertTrue(collect.isBlocked().isDone()); + assertSame(a.output, collect.next()); + collect.close(); + } + + @Test + public void testThousandChildrenAndOneListenerPerPendingPeriod() throws Exception { + for (int count : new int[] {100, 1000}) { + List children = new ArrayList<>(); + for (int i = 0; i < count; i++) { + children.add(new Child(block(i))); + } + CollectOperator collect = new CollectOperator(new TestContext(), children); + collect.enableReadyInputConsumption(); + Child ready = (Child) children.get(count - 1); + for (int round = 0; round < 5; round++) { + ListenableFuture wait = collect.isBlocked(); + assertFalse(wait.isDone()); + for (int repeat = 0; repeat < 10; repeat++) { + assertSame(wait, collect.isBlocked()); + } + ready.ready(); + wait.get(5, TimeUnit.SECONDS); + assertSame(ready.output, collect.next()); + ready.reset(); + } + for (int i = 0; i < count - 1; i++) { + Child child = (Child) children.get(i); + assertEquals(1, child.waitCalls); + assertEquals(1, child.wait.registrations.get()); + } + collect.close(); + } + } + + @Test + public void testAllReadyChildrenHaveLinearSelectionCost() throws Exception { + int childCount = 1000; + int rounds = 5; + List children = new ArrayList<>(); + for (int i = 0; i < childCount; i++) { + Child child = new Child(block(i)); + child.ready(); + children.add(child); + } + CountingList countedChildren = new CountingList(children); + CollectOperator collect = new CollectOperator(new TestContext(), countedChildren); + collect.enableReadyInputConsumption(); + try { + for (int round = 0; round < rounds; round++) { + assertTrue(collect.isBlocked().isDone()); + // A completed future may notify more than once. It must occupy only one ready slot. + for (Child child : children) { + child.wait.listener.run(); + child.wait.listener.run(); + } + Set consumed = new HashSet<>(); + for (int i = 0; i < childCount; i++) { + assertTrue(collect.isBlocked().isDone()); + TsBlock result = collect.next(); + assertTrue( + "every ready child must be selected once per round", + consumed.add(result.getColumn(0).getInt(0))); + } + assertEquals(childCount, consumed.size()); + } + for (Child child : children) { + assertEquals(rounds, child.advances); + } + long probes = countedChildren.getCalls + countFutureProbes(children); + long linearLimit = 20L * (childCount + childCount * rounds); + assertTrue( + "selection must use O(children + blocks) probes, observed " + probes, + probes <= linearLimit); + } finally { + collect.close(); + } + } + + @Test + public void testClosedChildrenDoNotIncreaseRemainingChildSelectionCost() throws Exception { + int childCount = 1000; + int remainingBlocks = 3000; + List children = new ArrayList<>(); + for (int i = 0; i < childCount; i++) { + Child child = new Child(block(i)); + child.hasNext = i == childCount - 1; + child.finished = !child.hasNext; + child.ready(); + children.add(child); + } + CountingList countedChildren = new CountingList(children); + CollectOperator collect = new CollectOperator(new TestContext(), countedChildren); + collect.enableReadyInputConsumption(); + try { + // Retire the first 999 children and consume one block from the remaining child. + for (int i = 0; i < childCount; i++) { + assertTrue(collect.isBlocked().isDone()); + collect.next(); + } + for (int i = 0; i < childCount - 1; i++) { + assertEquals(1, children.get(i).closes); + } + Child remaining = children.get(childCount - 1); + assertEquals(1, remaining.advances); + long initialProbes = countedChildren.getCalls + countFutureProbes(children); + for (int i = 0; i < remainingBlocks; i++) { + assertTrue(collect.isBlocked().isDone()); + assertSame(remaining.output, collect.next()); + } + assertEquals(remainingBlocks + 1, remaining.advances); + long probes = countedChildren.getCalls + countFutureProbes(children) - initialProbes; + assertTrue( + "retired children must not be visited for each remaining block, observed " + probes, + probes <= 20L * remainingBlocks); + } finally { + collect.close(); + } + } + + private static long countFutureProbes(List children) { + long probes = 0; + for (Child child : children) { + probes += child.wait.probes.get(); + } + return probes; + } + + @Test + public void testPendingListenerReusedAndStaleCompletionIgnored() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + CollectOperator collect = collect(a, b); + ListenableFuture first = collect.isBlocked(); + b.ready(); + assertSame(b.output, collect.next()); + Runnable stale = b.wait.listener; + b.reset(); + ListenableFuture second = collect.isBlocked(); + assertFalse(second.isDone()); + stale.run(); + assertFalse(second.isDone()); + a.ready(); + second.get(5, TimeUnit.SECONDS); + assertEquals(1, a.wait.registrations.get()); + assertSame(a.output, collect.next()); + assertTrue(first.isDone()); + collect.close(); + } + + @Test + public void testCompletionBetweenSelectionsIsNotLost() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = collect(a, b); + try { + assertTrue(collect.isBlocked().isDone()); + assertSame(b.output, collect.next()); + b.reset(); + // No selection is waiting yet: retain this completion until the next driver observation. + a.ready(); + assertTrue(collect.isBlocked().isDone()); + assertSame(a.output, collect.next()); + a.reset(); + assertFalse(collect.isBlocked().isDone()); + assertEquals(1, a.advances); + assertEquals(1, b.advances); + } finally { + collect.close(); + } + } + + @Test + public void testNullEmptyAndUnfinishedChildRefreshReadiness() throws Exception { + Child child = new Child(null); + child.ready(); + CollectOperator collect = collect(child); + assertTrue(collect.isBlocked().isDone()); + assertNull(collect.next()); + child.reset(); + assertFalse(collect.isBlocked().isDone()); + child.output = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32)).build(); + child.ready(); + assertTrue(collect.next().isEmpty()); + child.reset(); + child.hasNext = false; + child.ready(); + assertNull(collect.next()); + assertFalse(collect.isFinished()); + assertEquals(0, child.closes); + child.finished = true; + assertNull(collect.next()); + assertTrue(collect.isFinished()); + assertEquals(1, child.closes); + assertTrue(collect.isBlocked().isDone()); + } + + @Test + public void testCancellationDoesNotPropagateToChild() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = collect(child); + ListenableFuture exposed = collect.isBlocked(); + assertTrue(exposed.cancel(true)); + assertFalse(child.wait.isCancelled()); + ListenableFuture replacement = collect.isBlocked(); + assertFalse(replacement.isCancelled()); + child.ready(); + replacement.get(5, TimeUnit.SECONDS); + assertSame(child.output, collect.next()); + collect.close(); + } + + @Test + public void testFailureOfUnselectedChildIsNotLost() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.ready(); + CollectOperator collect = collect(a, b); + assertTrue(collect.isBlocked().isDone()); + IOException expected = new IOException("child failure"); + b.wait.delegate.setException(expected); + try { + collect.next(); + fail("must propagate the failed child even when another child was selected"); + } catch (IOException e) { + assertSame(expected, e); + } finally { + collect.close(); + } + assertEquals(0, a.advances); + } + + @Test + public void testChildCancellationPropagates() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = collect(child); + ListenableFuture wait = collect.isBlocked(); + child.wait.cancel(false); + assertTrue(wait.isDone()); + try { + collect.next(); + fail("must propagate child cancellation"); + } catch (CancellationException expected) { + assertEquals(0, child.advances); + } finally { + collect.close(); + } + } + + @Test + public void testCompletionDuringRegistrationAndOnAnotherThread() throws Exception { + Child a = new Child(block(1)); + a.wait.beforeRegistration = a::ready; + CollectOperator collect = collect(a); + assertTrue(collect.isBlocked().isDone()); + assertSame(a.output, collect.next()); + a.reset(); + ListenableFuture pending = collect.isBlocked(); + Thread completion = new Thread(a::ready); + completion.start(); + pending.get(5, TimeUnit.SECONDS); + completion.join(); + assertSame(a.output, collect.next()); + collect.close(); + } + + @Test + public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.closeFailure = new IOException("close failure"); + b.closeFailure = new IOException("second close failure"); + CollectOperator collect = collect(a, b); + ListenableFuture pending = collect.isBlocked(); + try { + collect.close(); + fail("must propagate close failure"); + } catch (IOException e) { + assertSame(a.closeFailure, e); + assertEquals(1, e.getSuppressed().length); + assertSame(b.closeFailure, e.getSuppressed()[0]); + } + pending.get(5, TimeUnit.SECONDS); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + a.ready(); + b.ready(); + assertTrue(collect.isFinished()); + assertNull(collect.next()); + collect.close(); + assertEquals(1, b.closes); + } + + @Test + public void testDefaultCollectRetainsSequentialConsumption() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + assertFalse(collect.isBlocked().isDone()); + a.ready(); + assertSame(a.output, collect.next()); + a.hasNext = false; + a.finished = true; + assertNull(collect.next()); + assertSame(b.output, collect.next()); + assertEquals(0, collect.calculateRetainedSizeAfterCallingNext()); + collect.close(); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + } + + @Test + public void testDefaultMappingCollectRetainsSequentialConsumption() throws Exception { + Child a = new Child(block(10, 11)); + Child b = new Child(block(20, 21)); + b.ready(); + List> mappings = new ArrayList<>(); + mappings.add(Collections.singletonList(0)); + mappings.add(Collections.singletonList(1)); + MappingCollectOperator collect = + new MappingCollectOperator( + new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings); + assertFalse(collect.isBlocked().isDone()); + a.ready(); + assertEquals(10, collect.next().getColumn(0).getInt(0)); + a.hasNext = false; + a.finished = true; + assertNull(collect.next()); + assertEquals(21, collect.next().getColumn(0).getInt(0)); + collect.close(); + } + + @Test + public void testLegacyInitializesOtherInputsOnceAndFinishesAtLastChild() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + for (int i = 0; i < 10; i++) { + assertFalse(collect.isBlocked().isDone()); + } + assertEquals(11, a.waitCalls); + assertEquals(1, b.waitCalls); + assertEquals(0, b.advances); + a.ready(); + assertSame(a.output, collect.next()); + a.hasNext = false; + assertNull(collect.next()); + assertTrue(collect.hasNext()); + assertFalse(collect.isFinished()); + for (int i = 0; i < 10; i++) { + assertTrue(collect.isBlocked().isDone()); + } + assertEquals(11, a.waitCalls); + assertEquals(11, b.waitCalls); + assertSame(b.output, collect.next()); + b.hasNext = false; + assertNull(collect.next()); + assertFalse(collect.hasNext()); + assertTrue(collect.isFinished()); + assertTrue(collect.isBlocked().isDone()); + collect.close(); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + } + + @Test + public void testLegacyMemoryUsesLargestChildPeak() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.peak = 100; + a.retained = 30; + a.returnSize = 10; + b.peak = 80; + b.retained = 50; + b.returnSize = 20; + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + assertEquals(100, collect.calculateMaxPeekMemory()); + assertEquals(0, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(20, collect.calculateMaxReturnSize()); + collect.close(); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyInputIsRetired() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.hasNext = false; + b.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + assertNull(collect.next()); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + assertSame(b.output, collect.next()); + b.hasNext = false; + assertNull(collect.next()); + assertTrue(collect.isFinished()); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyReadinessObservation() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child))); + assertFalse(collect.isBlocked().isDone()); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + collect.close(); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyNextWithoutIsBlocked() throws Exception { + Child child = new Child(block(1)); + child.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child))); + assertSame(child.output, collect.next()); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + collect.close(); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyClose() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child))); + collect.close(); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + assertEquals(1, child.closes); + } + + @Test + public void testEnablingReadyConsumptionIsIdempotent() throws Exception { + Child child = new Child(block(1)); + child.ready(); + CollectOperator collect = collect(child); + collect.enableReadyInputConsumption(); + ListenableFuture selection = collect.isBlocked(); + collect.enableReadyInputConsumption(); + assertSame(selection, collect.isBlocked()); + assertSame(child.output, collect.next()); + collect.close(); + } + + @Test + public void testMappingUsesActualChildIdentity() throws Exception { + Child a = new Child(block(10, 11)); + Child b = new Child(block(20, 21)); + b.ready(); + List> mappings = new ArrayList<>(); + mappings.add(Collections.singletonList(0)); + mappings.add(Collections.singletonList(1)); + MappingCollectOperator collect = + new MappingCollectOperator( + new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings); + collect.enableReadyInputConsumption(); + assertTrue(collect.isBlocked().isDone()); + assertEquals(21, collect.next().getColumn(0).getInt(0)); + a.ready(); + assertEquals(10, collect.next().getColumn(0).getInt(0)); + collect.close(); + assertNull(mappings.get(0)); + assertNull(mappings.get(1)); + } + + @Test + public void testMemoryIncludesAllLiveRetainedStates() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.retained = 30; + a.peak = 100; + a.returnSize = 10; + b.retained = 50; + b.peak = 80; + b.returnSize = 20; + CollectOperator collect = collect(a, b); + assertEquals(80, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(150, collect.calculateMaxPeekMemory()); + assertEquals(20, collect.calculateMaxReturnSize()); + assertTrue(collect.ramBytesUsed() > 0); + a.hasNext = false; + a.finished = true; + a.ready(); + collect.next(); + assertEquals(50, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(80, collect.calculateMaxPeekMemory()); + collect.close(); + assertEquals(0, collect.calculateMaxPeekMemory()); + assertEquals(0, collect.calculateMaxReturnSize()); + assertEquals(0, collect.calculateRetainedSizeAfterCallingNext()); + } + + @Test + public void testEmptyCollectIsAlreadyFinished() throws Exception { + CollectOperator collect = collect(); + assertFalse(collect.hasNext()); + assertTrue(collect.isFinished()); + assertTrue(collect.isBlocked().isDone()); + assertNull(collect.next()); + assertEquals(0, collect.calculateMaxPeekMemory()); + assertEquals(0, collect.calculateMaxReturnSize()); + collect.close(); + collect.close(); + } + + @Test + public void testFailureWakesAllPendingChildren() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + CollectOperator collect = collect(a, b); + ListenableFuture pending = collect.isBlocked(); + assertFalse(pending.isDone()); + IOException expected = new IOException("blocked child failure"); + b.wait.delegate.setException(expected); + try { + pending.get(5, TimeUnit.SECONDS); + fail("child failure must wake the driver"); + } catch (ExecutionException e) { + assertSame(expected, e.getCause()); + } + try { + collect.next(); + fail("child failure must propagate without advancing other children"); + } catch (IOException e) { + assertSame(expected, e); + } finally { + collect.close(); + } + assertEquals(0, a.advances); + assertEquals(0, b.advances); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + } + + @Test + public void testConcurrentCompletionsConsumeEveryChildExactlyOnce() throws Exception { + int count = 16; + List children = new ArrayList<>(); + for (int i = 0; i < count; i++) children.add(new Child(block(i))); + CollectOperator collect = new CollectOperator(new TestContext(), new ArrayList<>(children)); + collect.enableReadyInputConsumption(); + ListenableFuture pending = collect.isBlocked(); + assertFalse(pending.isDone()); + ExecutorService executor = Executors.newFixedThreadPool(count); + CountDownLatch start = new CountDownLatch(1); + List> completions = new ArrayList<>(); + try { + for (Child child : children) { + completions.add( + executor.submit( + () -> { + start.await(); + child.ready(); + return null; + })); + } + start.countDown(); + pending.get(5, TimeUnit.SECONDS); + for (Future completion : completions) completion.get(5, TimeUnit.SECONDS); + Set consumed = new HashSet<>(); + for (int round = 0; round < 2 * count && collect.hasNext(); round++) { + assertTrue(collect.isBlocked().isDone()); + TsBlock result = collect.next(); + if (result != null) { + int childId = result.getColumn(0).getInt(0); + assertTrue("child emitted more than once", consumed.add(childId)); + children.get(childId).hasNext = false; + children.get(childId).finished = true; + } + } + assertEquals(count, consumed.size()); + assertTrue(collect.isFinished()); + assertNull(collect.next()); + for (Child child : children) { + assertEquals(1, child.advances); + assertEquals(1, child.closes); + } + } finally { + start.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + collect.close(); + } + } + + private static final class CountingList extends AbstractList { + private final List delegate; + private long getCalls; + + private CountingList(List children) { + delegate = new ArrayList<>(children); + } + + @Override + public Operator get(int index) { + getCalls++; + return delegate.get(index); + } + + @Override + public Operator set(int index, Operator child) { + return delegate.set(index, child); + } + + @Override + public int size() { + return delegate.size(); + } + } + + private static final class CountingFuture + extends ForwardingListenableFuture.SimpleForwardingListenableFuture { + private final SettableFuture delegate; + private final AtomicInteger registrations = new AtomicInteger(); + private final AtomicInteger probes = new AtomicInteger(); + private Runnable beforeRegistration; + private Runnable listener; + + private CountingFuture() { + this(SettableFuture.create()); + } + + private CountingFuture(SettableFuture delegate) { + super(delegate); + this.delegate = delegate; + } + + @Override + public boolean isDone() { + probes.incrementAndGet(); + return super.isDone(); + } + + @Override + public Void get() throws InterruptedException, ExecutionException { + probes.incrementAndGet(); + return super.get(); + } + + @Override + public Void get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + probes.incrementAndGet(); + return super.get(timeout, unit); + } + + @Override + public void addListener(Runnable listener, Executor executor) { + this.listener = listener; + registrations.incrementAndGet(); + if (beforeRegistration != null) { + beforeRegistration.run(); + } + super.addListener(listener, executor); + } + } + + private static final class TestContext extends CommonOperatorContext { + private TestContext() { + super(0, new PlanNodeId("collect"), "test"); + } + + @Override + public MemoryReservationManager getMemoryReservationContext() { + return null; + } + + @Override + public int getFragmentId() { + return 0; + } + + @Override + public int getPipelineId() { + return 0; + } + + @Override + public long ramBytesUsed() { + return 0; + } + } + + private static final class Child implements Operator { + private final Thread driver = Thread.currentThread(); + private final CommonOperatorContext context = new TestContext(); + private CountingFuture wait = new CountingFuture(); + private TsBlock output; + private boolean hasNext = true; + private boolean finished; + private int advances; + private int waitCalls; + private int closes; + private IOException closeFailure; + private long retained, peak, returnSize; + + private Child(TsBlock output) { + this.output = output; + } + + private void ready() { + wait.delegate.set(null); + } + + private void reset() { + wait = new CountingFuture(); + } + + private void checkDriver() { + assertSame(driver, Thread.currentThread()); + } + + @Override + public CommonOperatorContext getOperatorContext() { + return context; + } + + @Override + public ListenableFuture isBlocked() { + checkDriver(); + waitCalls++; + return wait; + } + + @Override + public boolean hasNext() { + checkDriver(); + return hasNext; + } + + @Override + public TsBlock next() { + checkDriver(); + advances++; + return output; + } + + @Override + public boolean isFinished() { + checkDriver(); + return finished; + } + + @Override + public void close() throws Exception { + checkDriver(); + closes++; + if (closeFailure != null) throw closeFailure; + } + + @Override + public long calculateMaxPeekMemory() { + return peak; + } + + @Override + public long calculateMaxReturnSize() { + return returnSize; + } + + @Override + public long calculateRetainedSizeAfterCallingNext() { + return retained; + } + + @Override + public long ramBytesUsed() { + return 0; + } + } +} 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 b030af64cfa77..965ae20034cb4 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 @@ -3941,4 +3941,7 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB = "Only InMemoryDeviceEntryDataSet supports get inline device entries"; + public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "Plan node %s is already registered on memory side %s, not %s"; + public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "Memory reservation is not registered for query %s, fragment %s, plan node %s"; + public static final String LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109 = "Shared exchange budget for FI {}: INPUT={}, OUTPUT={}"; } 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 5a504e2a0175f..290504411081d 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 @@ -4697,4 +4697,7 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB = "只有 InMemoryDeviceEntryDataSet 支持获取内存中的设备条目"; + public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "PlanNode %s 已注册到内存侧 %s，不能注册到 %s"; + public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "查询 %s、Fragment %s、PlanNode %s 尚未注册内存账目"; + public static final String LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109 = "FI {} 的共享 exchange 预算：INPUT={}，OUTPUT={}"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java index 5812d4e37cdd3..0d8423803618d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java @@ -27,6 +27,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.utils.CommonUtils; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -265,7 +266,10 @@ public ListenableFuture add(TsBlock tsBlock) { localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, + fullFragmentInstanceId, + localPlanNodeId, + MemorySide.OUTPUT); alreadyRegistered = true; } MemoryReservationResult reserveResult = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index d24a1552a5b22..4d592771ef23e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -30,6 +30,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SinkListener; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet; import org.apache.iotdb.db.queryengine.metric.DataExchangeCountMetricSet; import org.apache.iotdb.db.utils.SetThreadName; @@ -235,7 +236,10 @@ public SinkChannel( localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, + fullFragmentInstanceId, + localPlanNodeId, + MemorySide.OUTPUT); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index aed3f9533fb1d..8ee7ad601ef61 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -31,6 +31,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SourceHandleListener; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet; import org.apache.iotdb.db.queryengine.metric.DataExchangeCountMetricSet; import org.apache.iotdb.db.utils.SetThreadName; @@ -223,7 +224,10 @@ public SourceHandle( localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, + fullFragmentInstanceId, + localPlanNodeId, + MemorySide.INPUT); } @Override 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 5b8ee6ae1c1ae..16fecf6bc4983 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 @@ -120,6 +120,11 @@ public class FragmentInstanceContext extends QueryContext { private final MemoryReservationManager memoryReservationManager; + // Initialized after the complete local plan is generated, before drivers start. Both are zero + // when this FI uses independent per-plan-node exchange quotas. + private long inputMaxBytes; + private long outputMaxBytes; + protected IDataRegionForQuery dataRegion; private Filter globalTimeFilter; private List globalTimeFilterTimeRanges; @@ -551,6 +556,19 @@ public FragmentInstanceId getId() { return id; } + public void initializeExchangeMemoryBudget(long inputMaxBytes, long outputMaxBytes) { + this.inputMaxBytes = inputMaxBytes; + this.outputMaxBytes = outputMaxBytes; + } + + public long getInputMaxBytes() { + return inputMaxBytes; + } + + public long getOutputMaxBytes() { + return outputMaxBytes; + } + public void failed(Throwable cause) { stateMachine.failed(cause); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index b2d35568fa31c..3664a845d4f82 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -37,9 +37,9 @@ import java.util.Iterator; import java.util.List; import java.util.Map; -import java.util.Queue; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentSkipListMap; +import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; import static com.google.common.util.concurrent.Futures.immediateVoidFuture; @@ -55,10 +55,13 @@ public static class MemoryReservationFuture extends AbstractFuture { private final String fragmentInstanceId; private final String planNodeId; private final long bytesToReserve; + // Assigned once by the pool before publication. Zero denotes an untracked future. + private long reservationId; /** - * MemoryReservationFuture is created when SinkHandle or SourceHandle tries to reserve memory - * from pool. This field is max Bytes that SinkHandle or SourceHandle can reserve. + * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the + * pool. Local planning gives handles their FI-side limit in shared mode or their independent + * plan-node limit in legacy mode. Retain the supplied limit when retrying this request. */ private final long maxBytesCanReserve; @@ -152,14 +155,44 @@ public long getReservedBytes() { private final String id; private final IMemoryBlock memoryBlock; - private final long maxBytesPerFragmentInstance; - /** queryId -> fragmentInstanceId -> planNodeId -> bytesReserved. */ - private final Map>> queryMemoryReservations = + private static class PlanNodeReservation { + private final MemorySide side; + private volatile long bytesReserved; + + private PlanNodeReservation(MemorySide side) { + this.side = side; + } + + @Override + public String toString() { + return side + ":" + bytesReserved; + } + } + + private static class FragmentMemoryReservation { + private final Map planNodes = new ConcurrentHashMap<>(); + private final AtomicLong inputUsedBytes = new AtomicLong(); + private final AtomicLong outputUsedBytes = new AtomicLong(); + private final Map> pendingReservations = + new ConcurrentHashMap<>(); + // Enabled by local planning before drivers start; request limits are supplied by handles. + private volatile boolean sharedBudget; + + private AtomicLong getUsedBytes(MemorySide side) { + return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes; + } + } + + /** queryId -> fragmentInstanceId -> FI side totals and planNode details. */ + private final Map> queryMemoryReservations = new ConcurrentHashMap<>(); - private final Queue> memoryReservationFutures = - new ConcurrentLinkedQueue<>(); + // Sequence keys retain FIFO wakeup order and allow removal without scanning unrelated requests. + // Only blocked reservations use these indexes; successful admission keeps its existing fast path. + private final Map> memoryReservationFutures = + new ConcurrentSkipListMap<>(); + private final AtomicLong nextReservationId = new AtomicLong(); public MemoryPool(String id, MemoryManager memoryManager, long maxBytesPerFragmentInstance) { this.id = Validate.notNull(id, DataNodeQueryMessages.EXCEPTION_ID_CAN_NOT_BE_NULL_DOT_BDD2AD7D); @@ -176,7 +209,6 @@ public MemoryPool(String id, MemoryManager memoryManager, long maxBytesPerFragme .EXCEPTION_MAX_BYTES_PER_FI_SHOULD_BE_IN_LEFT_PAREN_0_COMMA_MAXBYTES_RIGHT_BRACKET_DOT_MAXB_4D37C457, maxBytesPerFragmentInstance, this.memoryBlock.getTotalMemorySizeInBytes()); - this.maxBytesPerFragmentInstance = maxBytesPerFragmentInstance; } public String getId() { @@ -204,12 +236,40 @@ public int getMemoryReservationSize() { * planNodeId to queryMemoryReservationsMap first. */ public void registerPlanNodeIdToQueryMemoryMap( - String queryId, String fragmentInstanceId, String planNodeId) { + String queryId, String fragmentInstanceId, String planNodeId, MemorySide side) { + java.util.Objects.requireNonNull(side); synchronized (queryMemoryReservations) { queryMemoryReservations .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) - .computeIfAbsent(fragmentInstanceId, x -> new ConcurrentHashMap<>()) - .putIfAbsent(planNodeId, 0L); + .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation()) + .planNodes + .compute( + planNodeId, + (key, reservation) -> { + if (reservation != null && reservation.side != side) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625, + planNodeId, + reservation.side, + side)); + } + return reservation == null ? new PlanNodeReservation(side) : reservation; + }); + } + } + + /** Enable shared accounting after the complete FI has been checked by local planning. */ + public void registerFragmentInstanceSharedBudget(String queryId, String fragmentInstanceId) { + synchronized (queryMemoryReservations) { + FragmentMemoryReservation fragment = + queryMemoryReservations + .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) + .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation()); + synchronized (fragment) { + fragment.sharedBudget = true; + } } } @@ -224,15 +284,19 @@ public void registerPlanNodeIdToQueryMemoryMap( */ public void deRegisterFragmentInstanceFromQueryMemoryMap( String queryId, String fragmentInstanceId, boolean forceDeregister) { - Map> queryRelatedMemory = queryMemoryReservations.get(queryId); + Map queryRelatedMemory = + queryMemoryReservations.get(queryId); if (queryRelatedMemory != null) { - Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + FragmentMemoryReservation fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); boolean hasPotentialMemoryLeak = false; // fragmentRelatedMemory could be null if the FI has not reserved any memory(For example, // next() of root operator returns no data) if (fragmentRelatedMemory != null) { hasPotentialMemoryLeak = - fragmentRelatedMemory.values().stream().anyMatch(value -> value != 0L); + fragmentRelatedMemory.inputUsedBytes.get() != 0L + || fragmentRelatedMemory.outputUsedBytes.get() != 0L + || fragmentRelatedMemory.planNodes.values().stream() + .anyMatch(value -> value.bytesReserved != 0L); } if (!forceDeregister && hasPotentialMemoryLeak) { // If hasPotentialMemoryLeak is true, it means that LocalSourceChannel/LocalSourceHandles @@ -241,6 +305,15 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap( // forceDeregister == true, after all LocalSourceChannel/LocalSourceHandles are closed. return; } + // Pending requests must not outlive their FI registration. Cancel before removing the + // account so another query's free() cannot retry a request against a missing FI. Do not + // hold the registration lock while notifying cancellation listeners. + if (fragmentRelatedMemory != null) { + for (MemoryReservationFuture future : + fragmentRelatedMemory.pendingReservations.values()) { + tryCancel(future); + } + } synchronized (queryMemoryReservations) { queryRelatedMemory.remove(fragmentInstanceId); // Another fragment may have recreated the query entry while this cleanup waited for the @@ -252,9 +325,9 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap( } if (hasPotentialMemoryLeak) { // hasPotentialMemoryLeak means that fragmentRelatedMemory is not null - List> invalidEntryList = - fragmentRelatedMemory.entrySet().stream() - .filter(entry -> entry.getValue() != 0L) + List> invalidEntryList = + fragmentRelatedMemory.planNodes.entrySet().stream() + .filter(entry -> entry.getValue().bytesReserved != 0L) .collect(Collectors.toList()); throw new MemoryLeakException( String.format( @@ -289,11 +362,10 @@ public MemoryReservationResult reserveWithPriority( Validate.notNull( planNodeId, DataNodeQueryMessages.EXCEPTION_PLANNODEID_CAN_NOT_BE_NULL_DOT_44027620); Validate.isTrue( - bytesToReserve > 0L && bytesToReserve <= maxBytesPerFragmentInstance, - DataNodeQueryMessages - .EXCEPTION_BYTESTORESERVE_SHOULD_BE_IN_LEFT_PAREN_0_COMMA_MAXBYTESPERFI_RIGHT_BRACKET_DOT_M_0753BB69, - maxBytesPerFragmentInstance, - bytesToReserve); + bytesToReserve > 0L, + DataNodeQueryMessages.EXCEPTION_BYTESTORESERVE_SHOULD_BE_GREATER_THAN_ZERO_DOT_56D15DE0); + // A request must fit both the caller's FI/handle limit and the physical global pool. + maxBytesCanReserve = Math.min(maxBytesCanReserve, getMaxBytes()); if (bytesToReserve > maxBytesCanReserve) { LOGGER.warn( DataNodeQueryMessages @@ -312,7 +384,6 @@ public MemoryReservationResult reserveWithPriority( if (tryReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve)) { return new MemoryReservationResult(immediateVoidFuture(), true, bytesToReserve); } else { - rollbackReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve); if (isHighestPriority) { // SHOW QUERIES: treat as success with zero bytes reserved from pool when insufficient. return new MemoryReservationResult(immediateVoidFuture(), true, 0L); @@ -321,10 +392,25 @@ public MemoryReservationResult reserveWithPriority( DataNodeQueryMessages.BLOCKED_RESERVE_REQUEST_ARG_BYTES_MEMORY_FOR_PLANNODEIDARG, bytesToReserve, planNodeId); - ListenableFuture result = + MemoryReservationFuture result = MemoryReservationFuture.create( queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve); - memoryReservationFutures.add((MemoryReservationFuture) result); + FragmentMemoryReservation fragment = + getFragmentReservation(queryId, fragmentInstanceId, planNodeId); + result.reservationId = nextReservationId.incrementAndGet(); + fragment.pendingReservations.put(result.reservationId, result); + memoryReservationFutures.put(result.reservationId, result); + // A concurrent free may have missed this request before it was published. Retry after + // publication, under the same future lock used by free and tryCancel to avoid charging twice. + try { + if (tryCompleteReservation(result)) { + removePendingReservation(result); + } + } catch (RuntimeException | Error e) { + // A concurrently removed registration must not leave a published request behind. + tryCancel(result); + throw e; + } return new MemoryReservationResult(result, false, bytesToReserve); } } @@ -343,18 +429,10 @@ public boolean tryReserveForTest( Validate.notNull( planNodeId, DataNodeQueryMessages.EXCEPTION_PLANNODEID_CAN_NOT_BE_NULL_DOT_44027620); Validate.isTrue( - bytesToReserve > 0L && bytesToReserve <= maxBytesPerFragmentInstance, - DataNodeQueryMessages - .EXCEPTION_BYTESTORESERVE_SHOULD_BE_IN_LEFT_PAREN_0_COMMA_MAXBYTESPERFI_RIGHT_BRACKET_DOT_M_0753BB69, - maxBytesPerFragmentInstance, - bytesToReserve); + bytesToReserve > 0L, + DataNodeQueryMessages.EXCEPTION_BYTESTORESERVE_SHOULD_BE_GREATER_THAN_ZERO_DOT_56D15DE0); - if (tryReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve)) { - return true; - } else { - rollbackReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve); - return false; - } + return tryReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve); } /** @@ -375,6 +453,9 @@ public synchronized long tryCancel(ListenableFuture future) { synchronized (future) { // If the future is not a MemoryReservationFuture, it must have been completed. if (future.isDone()) { + if (future instanceof MemoryReservationFuture) { + removePendingReservation((MemoryReservationFuture) future); + } return 0L; } Validate.isTrue( @@ -382,6 +463,7 @@ public synchronized long tryCancel(ListenableFuture future) { DataNodeQueryMessages.EXCEPTION_INVALID_FUTURE_TYPE_14507EF5 + future.getClass().getSimpleName()); future.cancel(true); + removePendingReservation((MemoryReservationFuture) future); } return ((MemoryReservationFuture) future).getBytesToReserve(); } @@ -390,66 +472,103 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l Validate.notNull(queryId, DataNodeQueryMessages.EXCEPTION_QUERYID_CAN_NOT_BE_NULL_DOT_16639DBE); Validate.isTrue(bytes > 0L); - try { - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .computeIfPresent( - planNodeId, - (k, reservedMemory) -> { - if (reservedMemory < bytes) { - throw new IllegalArgumentException( - DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); - } - return reservedMemory - bytes; - }); - } catch (NullPointerException e) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + if (fragment == null) { throw new IllegalArgumentException( DataNodeQueryMessages - .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2, - e); + .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2); } - - memoryBlock.release(bytes); + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + Object accountingLock = fragment.sharedBudget ? fragment : reservation; + synchronized (accountingLock) { + if (reservation.bytesReserved < bytes) { + throw new IllegalArgumentException( + DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); + } + reservation.bytesReserved -= bytes; + fragment.getUsedBytes(reservation.side).addAndGet(-bytes); + memoryBlock.release(bytes); + } + // Complete futures outside the accounting lock: listeners can reserve or free other inputs. if (memoryReservationFutures.isEmpty()) { return; } - Iterator> iterator = memoryReservationFutures.iterator(); + Iterator> iterator = memoryReservationFutures.values().iterator(); while (iterator.hasNext()) { MemoryReservationFuture future = iterator.next(); - synchronized (future) { - if (future.isCancelled() || future.isDone()) { - iterator.remove(); - continue; - } - long bytesToReserve = future.getBytesToReserve(); - String curQueryId = future.getQueryId(); - String curFragmentInstanceId = future.getFragmentInstanceId(); - String curPlanNodeId = future.getPlanNodeId(); - long maxBytesCanReserve = future.getMaxBytesCanReserve(); - if (tryReserve( - curQueryId, curFragmentInstanceId, curPlanNodeId, bytesToReserve, maxBytesCanReserve)) { - future.set(null); - iterator.remove(); - } else { - rollbackReserve(curQueryId, curFragmentInstanceId, curPlanNodeId, bytesToReserve); - } + if (tryCompleteReservation(future)) { + removePendingReservation(future); + } + } + } + + private void removePendingReservation(MemoryReservationFuture future) { + memoryReservationFutures.remove(future.reservationId, future); + Map fragments = + queryMemoryReservations.get(future.getQueryId()); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(future.getFragmentInstanceId()); + if (fragment != null) { + fragment.pendingReservations.remove(future.reservationId, future); + } + } + + private boolean tryCompleteReservation(MemoryReservationFuture future) { + synchronized (future) { + if (future.isDone()) { + return true; + } + long bytesToReserve = future.getBytesToReserve(); + String queryId = future.getQueryId(); + String fragmentInstanceId = future.getFragmentInstanceId(); + String planNodeId = future.getPlanNodeId(); + if (tryReserve( + queryId, + fragmentInstanceId, + planNodeId, + bytesToReserve, + future.getMaxBytesCanReserve())) { + future.set(null); + return true; } + return false; } } public long getQueryMemoryReservedBytes(String queryId) { - if (!queryMemoryReservations.containsKey(queryId)) { + Map fragments = queryMemoryReservations.get(queryId); + if (fragments == null) { return 0L; } long sum = 0; - for (Map map : queryMemoryReservations.get(queryId).values()) { - sum = sum + map.values().stream().reduce(0L, Long::sum); + for (FragmentMemoryReservation fragment : fragments.values()) { + sum += fragment.inputUsedBytes.get() + fragment.outputUsedBytes.get(); } return sum; } + public long getFragmentMemoryReservedBytes( + String queryId, String fragmentInstanceId, MemorySide side) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + return fragment == null ? 0L : fragment.getUsedBytes(side).get(); + } + + public long getPlanNodeMemoryReservedBytes( + String queryId, String fragmentInstanceId, String planNodeId) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + PlanNodeReservation reservation = fragment == null ? null : fragment.planNodes.get(planNodeId); + return reservation == null ? 0L : reservation.bytesReserved; + } + public long getReservedBytes() { return memoryBlock.getUsedMemoryInBytes(); } @@ -460,22 +579,61 @@ public boolean tryReserve( String planNodeId, long bytesToReserve, long maxBytesCanReserve) { - long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve); - long queryRemainingBytes = - maxBytesCanReserve - - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .merge(planNodeId, bytesToReserve, Long::sum); - return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() && queryRemainingBytes >= 0; + FragmentMemoryReservation fragment = + getFragmentReservation(queryId, fragmentInstanceId, planNodeId); + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + // Legacy quotas are independent per plan node. Only shared quotas require the FI lock. + Object accountingLock = fragment.sharedBudget ? fragment : reservation; + synchronized (accountingLock) { + long usedBytes = + fragment.sharedBudget + ? fragment.getUsedBytes(reservation.side).get() + : reservation.bytesReserved; + // Do not publish speculative charges. Two competing requests must not reject one another + // against different counters and then both wait after their rollbacks restore capacity. + if (bytesToReserve > maxBytesCanReserve - usedBytes + || !memoryBlock.allocate(bytesToReserve)) { + return false; + } + fragment.getUsedBytes(reservation.side).addAndGet(bytesToReserve); + reservation.bytesReserved += bytesToReserve; + return true; + } + } + + private FragmentMemoryReservation getFragmentReservation( + String queryId, String fragmentInstanceId, String planNodeId) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + if (fragment == null) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F, + queryId, + fragmentInstanceId, + planNodeId)); + } + return fragment; } - private void rollbackReserve( - String queryId, String fragmentInstanceId, String planNodeId, long bytesToReserve) { - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .merge(planNodeId, -bytesToReserve, Long::sum); - memoryBlock.release(bytesToReserve); + private PlanNodeReservation requireReservation( + PlanNodeReservation reservation, + String queryId, + String fragmentInstanceId, + String planNodeId) { + if (reservation == null) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F, + queryId, + fragmentInstanceId, + planNodeId)); + } + return reservation; } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java new file mode 100644 index 0000000000000..7149960538559 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java @@ -0,0 +1,26 @@ +/* + * 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.execution.memory; + +/** The exchange budget charged by a registered plan node. */ +public enum MemorySide { + INPUT, + OUTPUT +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java index cb2f6f02acb57..4b4ca52b8558f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java @@ -22,6 +22,7 @@ import org.apache.iotdb.calc.execution.filter.TopKRuntimeFilter; import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; import org.apache.iotdb.calc.execution.operator.process.FilterAndProjectOperator; import org.apache.iotdb.calc.execution.operator.process.LimitOperator; import org.apache.iotdb.calc.execution.operator.process.OffsetOperator; @@ -47,7 +48,9 @@ import org.apache.iotdb.commons.queryengine.plan.relational.metadata.QualifiedObjectName; import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.AggregationNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.FunctionCall; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.LongLiteral; @@ -301,6 +304,20 @@ protected OperatorContext addOperatorContext( .addOperatorContext(context.getNextOperatorId(), planNodeId, operatorType); } + @Override + public Operator visitCollect(CollectNode node, LocalExecutionPlanContext context) { + CollectOperator collect = (CollectOperator) super.visitCollect(node, context); + context.addCollectOperator(collect); + return collect; + } + + @Override + public Operator visitUnion(UnionNode node, LocalExecutionPlanContext context) { + CollectOperator collect = (CollectOperator) super.visitUnion(node, context); + context.addCollectOperator(collect); + return collect; + } + @Override public Operator visitCteScan(CteScanNode node, LocalExecutionPlanContext context) { OperatorContext operatorContext = @@ -343,8 +360,7 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte Operator child = node.getChildren().get(0).accept(this, context); List children = new ArrayList<>(1); children.add(child); - sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); - context.getDriverContext().setSink(sinkHandle); + context.setISink(sinkHandle); return new IdentitySinkOperator( operatorContext, children, downStreamChannelIndex, sinkHandle); } else { @@ -385,7 +401,6 @@ public Operator visitTableExchange(ExchangeNode node, LocalExecutionPlanContext if (!isSameNode) { context.addExchangeSumNum(1); } - sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); ExchangeOperator exchangeOperator = new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId()); context.addExchangeOperator(exchangeOperator); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index 7a5f52749aa4a..055bc4b19f98f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner; import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; import org.apache.iotdb.calc.plan.planner.ITableOperatorGeneratorContext; import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; import org.apache.iotdb.calc.plan.planner.memory.PipelineMemoryEstimator; @@ -32,8 +33,11 @@ import org.apache.iotdb.db.queryengine.execution.driver.DriverContext; import org.apache.iotdb.db.queryengine.execution.driver.SchemaDriverContext; import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink; +import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; import org.apache.iotdb.db.queryengine.plan.analyze.TemplatedInfo; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; @@ -41,6 +45,7 @@ import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion; import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.read.TimeValuePair; @@ -55,6 +60,7 @@ import java.time.ZoneId; import java.util.ArrayList; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; @@ -71,6 +77,7 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext { private static final Logger LOGGER = LoggerFactory.getLogger(LocalExecutionPlanContext.class); + // Save operators in this pipeline, a new one will be created when creating another pipeline private final DriverContext driverContext; private final AtomicInteger nextOperatorId; @@ -83,6 +90,9 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext private List pipelineDriverFactories; private List exchangeOperatorList = new ArrayList<>(); private int exchangeSumNum = 0; + private final long fragmentExchangeMaxBytes; + private Map legacySinkLimits = new IdentityHashMap<>(); + private List collectOperators = new ArrayList<>(); private List cachedDataTypes; @@ -110,6 +120,8 @@ public LocalExecutionPlanContext( TypeProvider typeProvider, FragmentInstanceContext instanceContext, DataNodeQueryContext dataNodeQueryContext) { + this.fragmentExchangeMaxBytes = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); this.typeProvider = typeProvider; this.allSensorsMap = new ConcurrentHashMap<>(); this.nextOperatorId = new AtomicInteger(0); @@ -121,6 +133,9 @@ public LocalExecutionPlanContext( // For creating subContext, differ from parent context mainly in driver context public LocalExecutionPlanContext(LocalExecutionPlanContext parentContext) { + this.fragmentExchangeMaxBytes = parentContext.fragmentExchangeMaxBytes; + this.legacySinkLimits = parentContext.legacySinkLimits; + this.collectOperators = parentContext.collectOperators; this.nextOperatorId = parentContext.nextOperatorId; this.typeProvider = parentContext.typeProvider; this.allSensorsMap = parentContext.allSensorsMap; @@ -142,6 +157,8 @@ public LocalExecutionPlanContext( final TypeProvider typeProvider, final FragmentInstanceContext instanceContext, final ISchemaRegion schemaRegion) { + this.fragmentExchangeMaxBytes = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); this.allSensorsMap = new ConcurrentHashMap<>(); this.typeProvider = typeProvider; this.nextOperatorId = new AtomicInteger(0); @@ -247,10 +264,11 @@ public void setExchangeSumNum(int exchangeSumNum) { this.exchangeSumNum = exchangeSumNum; } + /** Original per-handle quota for tree planning and independent table FI accounting. */ public long getMaxBytesOneHandleCanReserve() { - long maxBytesPerFI = - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); - return exchangeSumNum == 0 ? maxBytesPerFI : maxBytesPerFI / exchangeSumNum; + return exchangeSumNum == 0 + ? fragmentExchangeMaxBytes + : fragmentExchangeMaxBytes / exchangeSumNum; } public void addExchangeSumNum(int addValue) { @@ -261,6 +279,10 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) { this.exchangeOperatorList.add(exchangeOperator); } + public void addCollectOperator(CollectOperator collectOperator) { + collectOperators.add(collectOperator); + } + public void setMaxBytesOneHandleCanReserve() { long maxBytesOneHandleCanReserve = getMaxBytesOneHandleCanReserve(); LOGGER.debug( @@ -273,6 +295,60 @@ public void setMaxBytesOneHandleCanReserve() { exchangeOperator.getSourceHandle().setMaxBytesCanReserve(maxBytesOneHandleCanReserve)); } + /** Called after operator generation, before root memory estimation and factory registration. */ + public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { + // Only remote sources charge this FI's INPUT. Local inputs are owned by the producing FI; + // internal pipeline queues cannot enable sharing because of the pipeline check below. + long inputExchangeCount = + exchangeOperatorList.stream() + .filter(exchange -> exchange.getSourceHandle() instanceof SourceHandle) + .count(); + int outputSinkCount = legacySinkLimits.size(); + long inputMaxBytes = + inputExchangeCount == 0 + ? 0 + : (outputSinkCount == 0 ? fragmentExchangeMaxBytes : fragmentExchangeMaxBytes / 2); + long outputMaxBytes = outputSinkCount == 0 ? 0 : fragmentExchangeMaxBytes - inputMaxBytes; + long maxBlockBytes = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + boolean sharedBudget = + (inputMaxBytes > 0 || outputMaxBytes > 0) + && SharedExchangeMemoryEligibility.isEligible(plan) + // Only the root pipeline may exist. Child factories are registered during generation; + // the root factory is registered after its memory has been estimated below. + && pipelineDriverFactories.isEmpty() + && driverContext.getDependencyDriverIndex() == -1 + && (inputExchangeCount == 0 || inputMaxBytes >= maxBlockBytes) + && (outputSinkCount == 0 || outputMaxBytes >= maxBlockBytes); + FragmentInstanceContext instanceContext = driverContext.getFragmentInstanceContext(); + instanceContext.initializeExchangeMemoryBudget( + sharedBudget ? inputMaxBytes : 0, sharedBudget ? outputMaxBytes : 0); + if (sharedBudget) { + memoryPool.registerFragmentInstanceSharedBudget( + instanceContext.getId().getQueryId().getId(), + instanceContext.getId().getFragmentInstanceId()); + collectOperators.forEach(CollectOperator::enableReadyInputConsumption); + LOGGER.debug( + DataNodeQueryMessages.LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109, + instanceContext.getId(), + instanceContext.getInputMaxBytes(), + instanceContext.getOutputMaxBytes()); + } else { + LOGGER.debug( + DataNodeQueryMessages + .MAXBYTESONEHANDLECANRESERVE_FOR_EXCHANGEOPERATOR_IS_ARG_EXCHANGESUMNUM_IS_ARG, + getMaxBytesOneHandleCanReserve(), + exchangeSumNum); + } + legacySinkLimits.forEach( + (sink, legacyLimit) -> + sink.setMaxBytesCanReserve( + sharedBudget ? instanceContext.getOutputMaxBytes() : legacyLimit)); + long sourceLimit = + sharedBudget ? instanceContext.getInputMaxBytes() : getMaxBytesOneHandleCanReserve(); + exchangeOperatorList.forEach( + exchangeOperator -> exchangeOperator.getSourceHandle().setMaxBytesCanReserve(sourceLimit)); + } + public Set getAllSensors(IDeviceID deviceId, String sensorId) { Set allSensors = allSensorsMap.computeIfAbsent(deviceId, k -> new HashSet<>()); allSensors.add(sensorId); @@ -303,6 +379,9 @@ public void setISink(ISink sink) { checkArgument( driverContext.getSink() == null, DataNodeQueryMessages.EXCEPTION_THERE_MUST_BE_AT_MOST_ONE_SINKNODE_A965AFE7); + if (sink instanceof ISinkHandle) { + legacySinkLimits.put(sink, getMaxBytesOneHandleCanReserve()); + } driverContext.setSink(sink); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java index feb6b0cb74afd..c600b4eb9016d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java @@ -35,6 +35,7 @@ import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.DeviceContext; import org.apache.iotdb.db.queryengine.execution.driver.DataDriverContext; +import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeService; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; @@ -115,6 +116,9 @@ public List plan( Operator root = generateOperator(instanceContext, context, plan); + // Select Collect consumption before estimating its retained input memory. + setExchangeMemoryLimits(plan, context); + PipelineMemoryEstimator memoryEstimator = context.constructPipelineMemoryEstimator(root, null, plan, -1); // set the map to null for gc @@ -132,9 +136,6 @@ public List plan( context.getTimePartitions().ifPresent(instanceContext::setTimePartitions); - // set maxBytes one SourceHandle can reserve after visiting the whole tree - context.setMaxBytesOneHandleCanReserve(); - return context.getPipelineDriverFactories(); } @@ -155,6 +156,9 @@ public List plan( Operator root = generateOperator(instanceContext, context, plan); + // Select Collect consumption before estimating its retained input memory. + setExchangeMemoryLimits(plan, context); + PipelineMemoryEstimator memoryEstimator = context.constructPipelineMemoryEstimator(root, null, plan, -1); // set the map to null for gc @@ -165,9 +169,6 @@ public List plan( context.addPipelineDriverFactory(root, context.getDriverContext(), 0); - // set maxBytes one SourceHandle can reserve after visiting the whole tree - context.setMaxBytesOneHandleCanReserve(); - return context.getPipelineDriverFactories(); } @@ -196,6 +197,22 @@ private Operator generateOperator( return root; } + private void setExchangeMemoryLimits(PlanNode plan, LocalExecutionPlanContext context) { + FragmentInstanceContext instanceContext = context.getInstanceContext(); + if (instanceContext.getSessionInfo() == null + || instanceContext.getSessionInfo().getSqlDialect() == TREE) { + context.setMaxBytesOneHandleCanReserve(); + return; + } + // Choose the table FI policy only after all operators and pipelines are known. + context.setExchangeMemoryLimits( + plan, + MPPDataExchangeService.getInstance() + .getMPPDataExchangeManager() + .getLocalMemoryManager() + .getQueryPool()); + } + private long checkMemory( final PipelineMemoryEstimator memoryEstimator, FragmentInstanceContext instanceContext) throws MemoryNotEnoughException { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java new file mode 100644 index 0000000000000..47c31964a3ff9 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iotdb.db.queryengine.plan.planner; + +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExchangeNode; + +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.List; + +/** Select shared budgets only for FIs whose fan-in can consume any ready input. */ +final class SharedExchangeMemoryEligibility { + private SharedExchangeMemoryEligibility() {} + + static boolean isEligible(PlanNode root) { + boolean hasCollect = false; + Deque nodes = new ArrayDeque<>(); + nodes.push(root); + while (!nodes.isEmpty()) { + PlanNode node = nodes.pop(); + if (node instanceof ExchangeNode) { + // An upstream plan, if still attached, belongs to a different FI. + continue; + } + // TopK's input ordering assumptions remain unchanged, including its single-child form. + if (node instanceof TopKNode) { + return false; + } + // Core relational nodes need not implement the serialization-only PlanNode.getType(). + boolean anyReady = node instanceof CollectNode || node instanceof UnionNode; + hasCollect |= anyReady; + List children = node.getChildren(); + // Unknown multi-input operators are not implicitly safe. In particular, sequential sinks + // and joins can wait for an input while another input holds the entire shared budget. + if (!anyReady && children.size() > 1) { + return false; + } + nodes.addAll(children); + } + return hasCollect; + } +} 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 66f6da892f3e8..016ee576584bd 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 @@ -689,6 +689,12 @@ public List visitGroup(GroupNode node, PlanContext context) { for (PlanNode child : childrenNodes) { if (canSortEliminated(node.getOrderingScheme(), nodeOrderingMap.get(child.getPlanNodeId()))) { result.add(child); + } else if (groupingInvalidatedByCollect( + child, node.getOrderingScheme().getOrderBy().subList(0, node.getPartitionKeyCount()))) { + SortNode sort = + new SortNode(queryId.genPlanNodeId(), child, node.getOrderingScheme(), false, false); + nodeOrderingMap.put(sort.getPlanNodeId(), sort.getOrderingScheme()); + result.add(sort); } else { GroupNode subGroupNode = new GroupNode( @@ -1100,7 +1106,7 @@ public List visitExternalTsFileAggregationScan( if (needFinalAggregation) { OrderingScheme childOrdering = nodeOrderingMap.get(result.get(0).getPlanNodeId()); finalAggregation.setChild(mergeChildrenViaCollectOrMergeSort(childOrdering, result)); - return Collections.singletonList(finalAggregation); + return Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return result; } @@ -2045,7 +2051,7 @@ public List visitAggregation(AggregationNode node, PlanContext context } List childrenNodes = node.getChild().accept(this, context); - OrderingScheme childOrdering = nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()); + OrderingScheme childOrdering = commonOrdering(childrenNodes); if (node.isStreamable()) { // Child has Ordering, we need to check if it is the Ordering we expected if (childOrdering != null) { @@ -2118,7 +2124,7 @@ public List visitAggregation(AggregationNode node, PlanContext context if (physicalAggregation.isStreamable() && childOrdering != null) { nodeOrderingMap.put(planNodeId, expectedOrderingSchema); } - return aggregationNode; + return validateAggregationGrouping(aggregationNode); }) .collect(Collectors.toList()); @@ -2128,10 +2134,10 @@ public List visitAggregation(AggregationNode node, PlanContext context new CollectNode(queryId.genPlanNodeId(), aggregationNodes.get(0).getOutputSymbols()); collectNode.setChildren(aggregationNodes); splitResult.left.setChild(collectNode); - return Collections.singletonList(splitResult.left); + return Collections.singletonList(validateAggregationGrouping(splitResult.left)); } - return Collections.singletonList(physicalAggregation); + return Collections.singletonList(validateAggregationGrouping(physicalAggregation)); } // We cannot do multi-stage Aggregate if any aggregation-function is distinct. @@ -2144,7 +2150,9 @@ public List visitAggregation(AggregationNode node, PlanContext context nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes); node.setChild(physicalChild); return Collections.singletonList( - withRateFunctionInputOrdering(node, nodeOrderingMap.get(physicalChild.getPlanNodeId()))); + validateAggregationGrouping( + withRateFunctionInputOrdering( + node, nodeOrderingMap.get(physicalChild.getPlanNodeId())))); } Pair splitResult = split(node, symbolAllocator, queryId); AggregationNode intermediate = splitResult.right; @@ -2167,13 +2175,73 @@ public List visitAggregation(AggregationNode node, PlanContext context if (node.isStreamable() && childOrdering != null) { nodeOrderingMap.put(planNodeId, expectedOrderingSchema); } - return aggregationNode; + return validateAggregationGrouping(aggregationNode); }) .collect(Collectors.toList()); splitResult.left.setChild( mergeChildrenViaCollectOrMergeSort( nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes)); - return Collections.singletonList(splitResult.left); + return Collections.singletonList(validateAggregationGrouping(splitResult.left)); + } + + private AggregationNode validateAggregationGrouping(AggregationNode node) { + if (!node.getPreGroupedSymbols().isEmpty() + && groupingInvalidatedByCollect(node.getChild(), node.getPreGroupedSymbols())) { + node.setPreGroupedSymbols(ImmutableList.of()); + nodeOrderingMap.remove(node.getPlanNodeId()); + } + return node; + } + + /** Whether an unordered fan-in invalidates a previously established grouping property. */ + private boolean groupingInvalidatedByCollect(PlanNode node, List grouping) { + if (grouping.isEmpty()) { + return false; + } + if ((node instanceof CollectNode || node instanceof UnionNode) + && node.getChildren().size() > 1) { + return true; + } + OrderingScheme ordering = nodeOrderingMap.get(node.getPlanNodeId()); + if (ordering != null && prefixMatched(ordering, grouping)) { + return false; + } + // A completed aggregation emits each full grouping key once, even when it hashes its input. + if (node instanceof AggregationNode) { + AggregationNode aggregation = (AggregationNode) node; + if (!aggregation.getStep().isOutputPartial() + && aggregation.getGroupingSetCount() == 1 + && grouping.containsAll(aggregation.getGroupingKeys())) { + return false; + } + } + if (node instanceof ProjectNode) { + ProjectNode project = (ProjectNode) node; + List inputs = new ArrayList<>(); + for (Symbol symbol : grouping) { + Expression expression = project.getAssignments().get(symbol); + if (!(expression instanceof SymbolReference)) { + // A computed key cannot inherit a branch-local continuity proof. + return containsUnorderedCollect(project.getChild()); + } + inputs.add(Symbol.from(expression)); + } + return groupingInvalidatedByCollect(project.getChild(), inputs); + } + for (PlanNode child : node.getChildren()) { + if (groupingInvalidatedByCollect(child, grouping)) { + return true; + } + } + return false; + } + + private boolean containsUnorderedCollect(PlanNode node) { + if ((node instanceof CollectNode || node instanceof UnionNode) + && node.getChildren().size() > 1) { + return true; + } + return node.getChildren().stream().anyMatch(this::containsUnorderedCollect); } private static AggregationNode withRateFunctionInputOrdering( @@ -2301,7 +2369,8 @@ public List visitAggregationTableScan( throw new IllegalStateException( DataNodeQueryMessages.LIST_PLANNODE_SIZE_SHOULD_1_BUT_NOW_IS); } - resultTableScanNodeList = Collections.singletonList(finalAggregation); + resultTableScanNodeList = + Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return resultTableScanNodeList; @@ -2460,7 +2529,7 @@ private List constructSpilledAggregationTableScanByRegionReplicaSet( mergeChildrenViaCollectOrMergeSort( nodeOrderingMap.get(result.get(0).getPlanNodeId()), result)); } - return Collections.singletonList(finalAggregation); + return Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return result; } @@ -2677,7 +2746,8 @@ public List visitAggregationTreeDeviceViewScan( throw new IllegalStateException( DataNodeQueryMessages.LIST_PLANNODE_SIZE_SHOULD_1_BUT_NOW_IS); } - resultTableScanNodeList = Collections.singletonList(finalAggregation); + resultTableScanNodeList = + Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return resultTableScanNodeList; @@ -2863,7 +2933,7 @@ private List constructSpilledAggregationTreeDeviceViewScanByRegionRepl mergeChildrenViaCollectOrMergeSort( nodeOrderingMap.get(result.get(0).getPlanNodeId()), result)); } - return Collections.singletonList(finalAggregation); + return Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return result; } @@ -4007,6 +4077,9 @@ public NodeDistribution getNodeDistribution(PlanNodeId nodeId) { public void clearExpectedOrderingScheme() { expectedOrderingScheme = null; hasSortProperty = false; + // Grouping pushdown depends on this ordering. A boundary such as a global window + // must not pass an outer Group's partitioning request to its scans. + pushDownGrouping = false; } public void setExpectedOrderingScheme(OrderingScheme expectedOrderingScheme) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java index 66c4eff6e3e03..ac3fc65fe6683 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java @@ -34,6 +34,7 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TableFunctionProcessorNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis; import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode; @@ -101,6 +102,21 @@ public PlanNode visitPlan(PlanNode node, Context context) { return newNode; } + @Override + public PlanNode visitUnion(UnionNode node, Context context) { + if (node.getChildren().size() <= 1) { + return visitPlan(node, context); + } + if (!context.canSkip()) { + context.canOptimized = TO_SORT; + } + // Neither branch may establish a property for the interleaved Union output. + return node.replaceChildren( + node.getChildren().stream() + .map(child -> child.accept(this, new Context(null, 0))) + .collect(Collectors.toList())); + } + /** We need to make sure: context#partitionKey can match the prefix of childOrderSchema */ private void checkPrefixMatch(Context context, List childOrder) { if (context.canSkip()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java index 64ef1b8c2adc6..c35f66128052e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java @@ -22,6 +22,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.OrderingScheme; import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.FillNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GapFillNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.PatternRecognitionNode; @@ -31,6 +32,7 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TableFunctionProcessorNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKRankingNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PruneTableScanColumns; @@ -38,6 +40,7 @@ import com.google.common.collect.ImmutableSet; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Optional; @@ -75,6 +78,28 @@ public PlanNode visitPlan(PlanNode node, Context context) { return newNode; } + @Override + public PlanNode visitCollect(CollectNode node, Context context) { + return visitUnorderedFanIn(node, context); + } + + @Override + public PlanNode visitUnion(UnionNode node, Context context) { + return visitUnorderedFanIn(node, context); + } + + private PlanNode visitUnorderedFanIn(PlanNode node, Context context) { + if (node.getChildren().size() <= 1) { + return visitPlan(node, context); + } + List children = new ArrayList<>(); + for (PlanNode child : node.getChildren()) { + children.add(child.accept(this, new Context())); + } + context.addCannotEliminateSort(true); + return node.replaceChildren(children); + } + @Override public PlanNode visitProject(ProjectNode node, Context context) { Context newContext = new Context(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java index ce49104248123..a5573159a94e5 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java @@ -40,9 +40,11 @@ import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent; +import org.apache.tsfile.read.common.block.TsBlock; import org.junit.Assert; import org.junit.Test; import org.mockito.Mockito; @@ -180,6 +182,13 @@ public void testCreateLocalSinkHandle() { Assert.assertEquals( ((LocalSinkChannel) shuffleSinkHandle.getChannel(0)).getSharedTsBlockQueue(), ((LocalSourceHandle) localSourceHandle).getSharedTsBlockQueue()); + assertProducerOwnsQueue( + spyMemoryPool, + shuffleSinkHandle, + localSourceHandle, + localFragmentInstanceId, + remoteFragmentInstanceId, + localPlanNodeId); } @Test @@ -237,5 +246,49 @@ public void testCreateLocalSourceHandle() { Assert.assertEquals( ((LocalSinkChannel) shuffleSinkHandle.getChannel(0)).getSharedTsBlockQueue(), ((LocalSourceHandle) localSourceHandle).getSharedTsBlockQueue()); + assertProducerOwnsQueue( + spyMemoryPool, + shuffleSinkHandle, + localSourceHandle, + localFragmentInstanceId, + remoteFragmentInstanceId, + localPlanNodeId); + } + + private static void assertProducerOwnsQueue( + MemoryPool pool, + ISinkHandle sink, + ISourceHandle source, + TFragmentInstanceId producer, + TFragmentInstanceId consumer, + String planNodeId) { + SharedTsBlockQueue queue = ((LocalSourceHandle) source).getSharedTsBlockQueue(); + sink.setMaxBytesCanReserve(512); + source.setMaxBytesCanReserve(128); + Assert.assertEquals(512, queue.getMaxBytesCanReserve()); + String producerId = + FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(producer); + String consumerId = + FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(consumer); + TsBlock block = Utils.createMockTsBlock(256); + synchronized (queue) { + Assert.assertTrue(queue.add(block).isDone()); + Assert.assertEquals(256, pool.getReservedBytes()); + Assert.assertEquals( + 256, pool.getPlanNodeMemoryReservedBytes(producer.queryId, producerId, planNodeId)); + Assert.assertEquals( + 256, + pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.OUTPUT)); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.INPUT)); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(consumer.queryId, consumerId, MemorySide.INPUT)); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(consumer.queryId, consumerId, MemorySide.OUTPUT)); + Assert.assertSame(block, queue.remove()); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.OUTPUT)); + } } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java index a92eaa5d0e283..8b97de5f3a259 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java @@ -254,4 +254,34 @@ public void run() { } } } + + @Test + public void testDependentQueuesKeepIndependentBudgetsByDefault() { + MemoryPool pool = new MemoryPool("pipeline-chain", new MemoryManager(1024), 512); + LocalMemoryManager manager = Mockito.mock(LocalMemoryManager.class); + Mockito.when(manager.getQueryPool()).thenReturn(pool); + ExecutorService executor = newDirectExecutorService(); + TFragmentInstanceId id = new TFragmentInstanceId("pipeline-chain", 0, "0"); + SharedTsBlockQueue upstream = new SharedTsBlockQueue(id, "a", manager, executor); + SharedTsBlockQueue downstream = new SharedTsBlockQueue(id, "b", manager, executor); + upstream.setMaxBytesCanReserve(256); + downstream.setMaxBytesCanReserve(256); + TsBlock block = Utils.createMockTsBlock(64); + try { + for (int i = 0; i < 4; i++) Assert.assertTrue(upstream.add(block).isDone()); + ListenableFuture fifth = upstream.add(block); + Assert.assertFalse(fifth.isDone()); + TsBlock forwarded = upstream.remove(); + Assert.assertTrue(fifth.isDone()); + Assert.assertTrue(downstream.add(forwarded).isDone()); + Assert.assertEquals(4, upstream.getNumOfBufferedTsBlocks()); + Assert.assertEquals(1, downstream.getNumOfBufferedTsBlocks()); + Assert.assertEquals(320, pool.getReservedBytes()); + } finally { + downstream.close(); + upstream.close(); + executor.shutdownNow(); + } + Assert.assertEquals(0, pool.getReservedBytes()); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index 2d9e5683cfdfb..3cfebb1ba295c 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -66,9 +66,10 @@ public class SourceHandleTest { public static void beforeClass() { maxBytesPerFI = IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + // INPUT receives half of the fragment budget, allowing five buffered blocks. IoTDBDescriptor.getInstance() .getMemoryConfig() - .setMaxBytesPerFragmentInstance(5 * MOCK_TSBLOCK_SIZE); + .setMaxBytesPerFragmentInstance(10 * MOCK_TSBLOCK_SIZE); } @AfterClass diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 5f2285e4ca910..d2b901d7f2355 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -32,8 +32,23 @@ import java.lang.management.ManagementFactory; import java.lang.management.ThreadInfo; import java.lang.reflect.Field; +import java.util.AbstractCollection; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentSkipListMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.FutureTask; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import static org.awaitility.Awaitility.await; @@ -50,7 +65,8 @@ public class MemoryPoolTest { public void before() { MemoryManager memoryManager = Mockito.spy(new MemoryManager(1024L)); pool = new MemoryPool("test", memoryManager, 512L); - pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); } @Test @@ -87,7 +103,7 @@ public void testDeregisterDoesNotRemoveReregisteredQueryMemory() throws Exceptio // Replace the query map while the delayed cleanup still refers to the old one. pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); Assert.assertEquals(0, pool.getQueryMemoryReservationSize()); - pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "result", PLAN_NODE_ID); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "result", PLAN_NODE_ID, MemorySide.INPUT); Assert.assertTrue( pool.tryReserveForTest(QUERY_ID, "result", PLAN_NODE_ID, 256L, Long.MAX_VALUE)); } @@ -123,7 +139,7 @@ public void testDeregisterMissingFragmentPreservesReservedMemory() { @Test public void testDeregisterPreservesOtherFragments() { - pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID, MemorySide.INPUT); Assert.assertTrue( pool.tryReserveForTest(QUERY_ID, "other", PLAN_NODE_ID, 256L, Long.MAX_VALUE)); @@ -172,6 +188,282 @@ public void testForceDeregisterReportsMemoryLeak() { Assert.assertEquals(256L, pool.getReservedBytes()); } + @Test + public void testDeregisterCancelsPendingReservationsBeforeUnrelatedFree() { + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "b", MemorySide.INPUT); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "a", 512, 512)); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "b", 512, 512)); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + Assert.assertFalse(pending.getFuture().isDone()); + + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); + + Assert.assertTrue(pending.getFuture().isCancelled()); + pool.free("healthy", "fi", "a", 512); + pool.free("healthy", "fi", "b", 512); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testForceDeregisterCancelsPendingReservationsEvenWhenReportingLeak() { + Assert.assertTrue(pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512)); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "a", MemorySide.INPUT); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "a", 256, 512)); + + Assert.assertThrows( + MemoryLeakException.class, + () -> + pool.deRegisterFragmentInstanceFromQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, true)); + + Assert.assertTrue(pending.getFuture().isCancelled()); + pool.free("healthy", "fi", "a", 256); + // Forced cleanup still reports and retains the original leaked allocation. + Assert.assertEquals(512, pool.getReservedBytes()); + } + + @Test + public void testNonForcedDeregisterKeepsPendingReservationWhileMemoryIsLive() { + Assert.assertTrue(pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512)); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, false); + + Assert.assertFalse(pending.getFuture().isDone()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512); + Assert.assertTrue(pending.getFuture().isDone()); + Assert.assertFalse(pending.getFuture().isCancelled()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, false); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testDeregisterDoesNotCancelOtherFragmentPendingReservation() { + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "b", MemorySide.INPUT); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "a", 512, 512)); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "b", 512, 512)); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID, MemorySide.INPUT); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + MemoryReservationResult other = + pool.reserveWithPriority(QUERY_ID, "other", PLAN_NODE_ID, 256, 512, false); + + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); + + Assert.assertTrue(pending.getFuture().isCancelled()); + Assert.assertFalse(other.getFuture().isDone()); + pool.free("healthy", "fi", "a", 512); + Assert.assertTrue(other.getFuture().isDone()); + Assert.assertFalse(other.getFuture().isCancelled()); + pool.free(QUERY_ID, "other", PLAN_NODE_ID, 256); + pool.free("healthy", "fi", "b", 512); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testOwnerCancellationAndDeregistrationDoNotScanUnrelatedWaiters() throws Exception { + assertDeregistrationDoesNotScanUnrelatedWaiters(true); + } + + @Test + public void testDeregistrationCancelsOnlyIndexedWaiters() throws Exception { + assertDeregistrationDoesNotScanUnrelatedWaiters(false); + } + + private void assertDeregistrationDoesNotScanUnrelatedWaiters(boolean ownerCancels) + throws Exception { + AtomicInteger visits = countGlobalPendingVisits(); + pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512); + pool.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "p", MemorySide.INPUT); + pool.tryReserve("holder", "fi", "p", 512, 512); + int count = 1024; + List victims = new ArrayList<>(); + List unrelated = new ArrayList<>(); + for (int i = 0; i < count; i++) { + String fi = Integer.toString(i); + pool.registerPlanNodeIdToQueryMemoryMap("victim", fi, "p", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("unrelated", fi, "p", MemorySide.INPUT); + victims.add(pool.reserveWithPriority("victim", fi, "p", 8, 512, false)); + unrelated.add(pool.reserveWithPriority("unrelated", fi, "p", 8, 512, false)); + } + for (int i = 0; i < count; i++) { + if (ownerCancels) { + Assert.assertEquals(8, pool.tryCancel(victims.get(i).getFuture())); + } + pool.deRegisterFragmentInstanceFromQueryMemoryMap("victim", Integer.toString(i), true); + Assert.assertTrue(victims.get(i).getFuture().isCancelled()); + Assert.assertFalse(unrelated.get(i).getFuture().isDone()); + } + // An operation-count assertion avoids timing noise and detects a global scan for every FI. + Assert.assertEquals(0, visits.get()); + Assert.assertEquals(count, pool.getMemoryReservationSize()); + for (int i = 0; i < count; i++) { + Assert.assertEquals(8, pool.tryCancel(unrelated.get(i).getFuture())); + pool.deRegisterFragmentInstanceFromQueryMemoryMap("unrelated", Integer.toString(i), true); + } + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512); + pool.free("holder", "fi", "p", 512); + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); + pool.deRegisterFragmentInstanceFromQueryMemoryMap("holder", "fi", true); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + Assert.assertEquals(0, pool.getQueryMemoryReservationSize()); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @SuppressWarnings("unchecked") + private AtomicInteger countGlobalPendingVisits() throws Exception { + AtomicInteger visits = new AtomicInteger(); + Field field = MemoryPool.class.getDeclaredField("memoryReservationFutures"); + field.setAccessible(true); + Object pending = field.get(pool); + if (pending instanceof Queue) { + // Keep this branch so the regression also runs against the former queue implementation. + Queue> counted = + new ConcurrentLinkedQueue>() { + @Override + public Iterator> iterator() { + return countVisits(super.iterator(), visits); + } + }; + counted.addAll((Queue>) pending); + field.set(pool, counted); + } else { + Map> counted = + new ConcurrentSkipListMap>() { + @Override + public Collection> values() { + Collection> values = super.values(); + return new AbstractCollection>() { + @Override + public Iterator> iterator() { + return countVisits(values.iterator(), visits); + } + + @Override + public int size() { + return values.size(); + } + }; + } + }; + counted.putAll((Map>) pending); + field.set(pool, counted); + } + return visits; + } + + private static Iterator countVisits(Iterator delegate, AtomicInteger visits) { + return new Iterator() { + @Override + public boolean hasNext() { + return delegate.hasNext(); + } + + @Override + public T next() { + visits.incrementAndGet(); + return delegate.next(); + } + + @Override + public void remove() { + delegate.remove(); + } + }; + } + + @Test + public void testCancellationKeepsFifoWakeupOrder() { + pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512); + pool.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "p", MemorySide.INPUT); + pool.tryReserve("holder", "fi", "p", 512, 512); + List requests = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + String fi = Integer.toString(i); + pool.registerPlanNodeIdToQueryMemoryMap("waiting", fi, "p", MemorySide.INPUT); + requests.add(pool.reserveWithPriority("waiting", fi, "p", 64, 512, false)); + } + Assert.assertEquals(64, pool.tryCancel(requests.get(1).getFuture())); + Assert.assertEquals(0, pool.tryCancel(requests.get(1).getFuture())); + Assert.assertEquals(2, pool.getMemoryReservationSize()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 64); + Assert.assertTrue(requests.get(0).getFuture().isDone()); + Assert.assertFalse(requests.get(2).getFuture().isDone()); + pool.free("waiting", "0", "p", 64); + Assert.assertTrue(requests.get(2).getFuture().isDone()); + pool.free("waiting", "2", "p", 64); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 448); + pool.free("holder", "fi", "p", 512); + for (int i = 0; i < 3; i++) { + pool.deRegisterFragmentInstanceFromQueryMemoryMap("waiting", Integer.toString(i), true); + } + Assert.assertEquals(0, pool.getMemoryReservationSize()); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testConcurrentCancellationAndGrantRemoveAllPendingIndexes() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int round = 0; round < 32; round++) { + MemoryPool concurrent = new MemoryPool("race", new MemoryManager(1024), 512); + concurrent.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "a", MemorySide.INPUT); + concurrent.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "b", MemorySide.INPUT); + concurrent.tryReserve("holder", "fi", "a", 512, 512); + concurrent.tryReserve("holder", "fi", "b", 512, 512); + List requests = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + String fi = Integer.toString(i); + concurrent.registerPlanNodeIdToQueryMemoryMap("waiting", fi, "p", MemorySide.INPUT); + requests.add(concurrent.reserveWithPriority("waiting", fi, "p", 16, 512, false)); + } + CountDownLatch start = new CountDownLatch(1); + Future grant = + executor.submit( + () -> { + start.await(); + concurrent.free("holder", "fi", "a", 512); + return null; + }); + Future cancel = + executor.submit( + () -> { + start.await(); + for (MemoryReservationResult request : requests) { + concurrent.tryCancel(request.getFuture()); + } + return null; + }); + start.countDown(); + grant.get(10, TimeUnit.SECONDS); + cancel.get(10, TimeUnit.SECONDS); + for (int i = 0; i < requests.size(); i++) { + Assert.assertTrue(requests.get(i).getFuture().isDone()); + if (!requests.get(i).getFuture().isCancelled()) { + concurrent.free("waiting", Integer.toString(i), "p", 16); + } + concurrent.deRegisterFragmentInstanceFromQueryMemoryMap( + "waiting", Integer.toString(i), true); + } + concurrent.free("holder", "fi", "b", 512); + concurrent.deRegisterFragmentInstanceFromQueryMemoryMap("holder", "fi", true); + Assert.assertEquals(0, concurrent.getMemoryReservationSize()); + Assert.assertEquals(0, concurrent.getQueryMemoryReservationSize()); + Assert.assertEquals(0, concurrent.getReservedBytes()); + } + } finally { + executor.shutdownNow(); + Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + @Test public void testTryReserve() { @@ -518,4 +810,491 @@ public void testReserveWithPriorityShowQueriesWithinAvailable() { Assert.assertTrue(r.getFuture().isDone()); Assert.assertEquals(256L, pool.getReservedBytes()); } + + @Test + public void testSharedInputBudgetAndIndependentOutput() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200, 256, false) + .isReserveSuccess()); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, 256, false); + Assert.assertFalse(pending.isReserveSuccess()); + Assert.assertFalse(pending.getFuture().isDone()); + Assert.assertEquals( + 200, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2")); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256, 256, false) + .isReserveSuccess()); + Assert.assertEquals(456, pool.getReservedBytes()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200); + Assert.assertTrue(pending.getFuture().isDone()); + Assert.assertEquals( + 100, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 256, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.OUTPUT)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + } + + @Test + public void testSideRegistrationCannotChange() { + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); + try { + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.OUTPUT); + Assert.fail("conflicting side must be rejected"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, pool.getReservedBytes()); + } + } + + @Test + public void testRequestsUseCurrentSuppliedLimitInsteadOfPoolCreationBudget() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 768, 768, false) + .isReserveSuccess()); + Assert.assertEquals(768, pool.getReservedBytes()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 768); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testRequestLargerThanGlobalCapacityDoesNotWait() { + try { + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 1025, 2048, false); + Assert.fail("a request larger than the global pool cannot become ready"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + } + } + + @Test + public void testSharedRequestCannotExceedSuppliedInputBudget() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + try { + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, 256, false); + Assert.fail("request exceeds the supplied input budget"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + } + } + + @Test + public void testSharedReservationsAndWakeupsUseSuppliedSideLimits() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200, 256, false) + .isReserveSuccess()); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, 256, false); + Assert.assertFalse(pending.getFuture().isDone()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256, 256, false) + .isReserveSuccess()); + // An OUTPUT release cannot grant an INPUT request whose shared side remains full. + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256); + Assert.assertFalse(pending.getFuture().isDone()); + Assert.assertFalse(pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, 256)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200); + pending.getFuture().get(10, TimeUnit.SECONDS); + Assert.assertEquals( + 100, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + } + + @Test + public void testSingleDirectionKeepsFullBudgetAndSingleBlockLimit() { + MemoryPool singleDirectionPool = + new MemoryPool("single-direction", new MemoryManager(1024), 513); + for (MemorySide side : MemorySide.values()) { + String fi = side.name(); + singleDirectionPool.registerFragmentInstanceSharedBudget(QUERY_ID, fi); + singleDirectionPool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, "active", side); + singleDirectionPool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, fi, "unused", side == MemorySide.INPUT ? MemorySide.OUTPUT : MemorySide.INPUT); + Assert.assertTrue( + singleDirectionPool + .reserveWithPriority(QUERY_ID, fi, "active", 513, 513, false) + .isReserveSuccess()); + try { + singleDirectionPool.reserveWithPriority(QUERY_ID, fi, "unused", 1, 0, false); + Assert.fail("an unused direction supplies a zero limit"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, singleDirectionPool.getMemoryReservationSize()); + Assert.assertEquals(513, singleDirectionPool.getReservedBytes()); + } + singleDirectionPool.free(QUERY_ID, fi, "active", 513); + Assert.assertEquals(0, singleDirectionPool.getReservedBytes()); + } + } + + @Test + public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + ExecutorService executor = Executors.newFixedThreadPool(16); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch attempted = new CountDownLatch(16); + CountDownLatch release = new CountDownLatch(1); + List> tasks = new ArrayList<>(); + try { + for (int i = 0; i < 16; i++) { + String id = "pipeline-" + i; + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT); + tasks.add( + executor.submit( + () -> { + start.await(); + boolean reserved = + pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128); + attempted.countDown(); + release.await(); + if (reserved) pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32); + return null; + })); + } + start.countDown(); + Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS)); + long details = 0; + for (int i = 0; i < 16; i++) { + details += + pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "pipeline-" + i); + } + Assert.assertTrue(details > 0 && details <= 128); + Assert.assertEquals( + details, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals(details, pool.getReservedBytes()); + release.countDown(); + for (Future task : tasks) task.get(10, TimeUnit.SECONDS); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + } finally { + start.countDown(); + release.countDown(); + executor.shutdownNow(); + } + } + + @Test + public void testReleaseBeforePendingReservationIsPublished() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + pool = Mockito.spy(pool); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false) + .isReserveSuccess()); + CountDownLatch failedAttempt = new CountDownLatch(1); + CountDownLatch released = new CountDownLatch(1); + AtomicBoolean firstAttempt = new AtomicBoolean(true); + Mockito.doAnswer( + invocation -> { + boolean granted = (boolean) invocation.callRealMethod(); + if (!granted && firstAttempt.compareAndSet(true, false)) { + failedAttempt.countDown(); + Assert.assertTrue(released.await(10, TimeUnit.SECONDS)); + } + return granted; + }) + .when(pool) + .tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256, 256); + FutureTask reservation = + new FutureTask<>( + () -> + pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256, 256, false)); + Thread reserver = new Thread(reservation, "reserve-shared-input"); + reserver.start(); + try { + Assert.assertTrue(failedAttempt.await(10, TimeUnit.SECONDS)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + released.countDown(); + MemoryReservationResult result = reservation.get(10, TimeUnit.SECONDS); + Assert.assertTrue( + "Releasing the last owner must wake the pending request", result.getFuture().isDone()); + result.getFuture().get(10, TimeUnit.SECONDS); + Assert.assertEquals(256, pool.getReservedBytes()); + Assert.assertEquals( + 256, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 256, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2")); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256); + Assert.assertEquals(0, pool.getReservedBytes()); + } finally { + released.countDown(); + reserver.join(10000); + Assert.assertFalse(reserver.isAlive()); + } + } + + @Test + public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT); + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false); + MemoryReservationResult cancelled = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", 128, 256, false); + MemoryReservationResult waiting = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", 256, 256, false); + Assert.assertFalse(cancelled.getFuture().isDone()); + Assert.assertFalse(waiting.getFuture().isDone()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 128, 256, false) + .isReserveSuccess()); + Assert.assertEquals(128, pool.tryCancel(cancelled.getFuture())); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + waiting.getFuture().get(10, TimeUnit.SECONDS); + Assert.assertTrue(cancelled.getFuture().isCancelled()); + Assert.assertEquals( + 0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled")); + Assert.assertEquals( + 256, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 128, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.OUTPUT)); + Assert.assertEquals(384, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", 256); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 128); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testConcurrentPendingReservationsCompleteAndReleaseAllAccounting() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + int count = 16; + ExecutorService executor = Executors.newFixedThreadPool(count); + CountDownLatch attempted = new CountDownLatch(count); + List> tasks = new ArrayList<>(); + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 128, 128, false); + try { + for (int i = 0; i < count; i++) { + String id = "waiting-" + i; + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT); + tasks.add( + executor.submit( + () -> { + for (int round = 0; round < 50; round++) { + MemoryReservationResult result = + pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128, false); + if (round == 0) attempted.countDown(); + result.getFuture().get(10, TimeUnit.SECONDS); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, result.getReservedBytes()); + } + return null; + })); + } + Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS)); + Assert.assertEquals(count, pool.getMemoryReservationSize()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 128); + for (Future task : tasks) task.get(30, TimeUnit.SECONDS); + for (int i = 0; i < count; i++) { + Assert.assertEquals( + 0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting-" + i)); + } + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + } finally { + executor.shutdownNow(); + Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void testDefaultFragmentKeepsIndependentPlanNodeQuotas() { + MemoryPool legacy = new MemoryPool("legacy", new MemoryManager(1024), 512); + legacy.registerPlanNodeIdToQueryMemoryMap("q", "f", "a", MemorySide.INPUT); + legacy.registerPlanNodeIdToQueryMemoryMap("q", "f", "b", MemorySide.INPUT); + Assert.assertTrue( + legacy.reserveWithPriority("q", "f", "a", 256, 256, false).isReserveSuccess()); + Assert.assertTrue( + legacy.reserveWithPriority("q", "f", "b", 256, 256, false).isReserveSuccess()); + Assert.assertEquals(512, legacy.getFragmentMemoryReservedBytes("q", "f", MemorySide.INPUT)); + legacy.free("q", "f", "a", 256); + legacy.free("q", "f", "b", 256); + Assert.assertEquals(0, legacy.getReservedBytes()); + } + + @Test + public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "b", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", MemorySide.INPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false) + .isReserveSuccess()); + MemoryReservationResult blocked = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", 64, 256, false); + Assert.assertFalse(blocked.getFuture().isDone()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, "legacy", "a", 256, 256, false).isReserveSuccess()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, "legacy", "b", 256, 256, false).isReserveSuccess()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + Assert.assertTrue(blocked.getFuture().isDone()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", 64); + pool.free(QUERY_ID, "legacy", "a", 256); + pool.free(QUERY_ID, "legacy", "b", 256); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { + assertConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits(true); + } + + @Test + public void testLegacyConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { + assertConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits(false); + } + + private void assertConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits(boolean shared) + throws Exception { + if (shared) { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + } + for (int i = 0; i < 3; i++) { + String fi = "other" + i; + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, PLAN_NODE_ID, MemorySide.OUTPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, fi, PLAN_NODE_ID, 256, 256, false).isReserveSuccess()); + } + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", MemorySide.INPUT); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 100; i++) { + CountDownLatch start = new CountDownLatch(1); + Future first = + executor.submit( + () -> { + start.await(); + return pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false); + }); + Future second = + executor.submit( + () -> { + start.await(); + return pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, "b", 256, 256, false); + }); + start.countDown(); + MemoryReservationResult a = first.get(10, TimeUnit.SECONDS); + MemoryReservationResult b = second.get(10, TimeUnit.SECONDS); + Assert.assertTrue(a.getFuture().isDone() ^ b.getFuture().isDone()); + String granted = a.getFuture().isDone() ? PLAN_NODE_ID : "b"; + String waiting = a.getFuture().isDone() ? "b" : PLAN_NODE_ID; + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, granted, 256); + Assert.assertTrue(a.getFuture().isDone()); + Assert.assertTrue(b.getFuture().isDone()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, waiting, 256); + Assert.assertEquals(256, pool.getRemainingBytes()); + Assert.assertEquals( + 0, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + } + } finally { + executor.shutdownNow(); + } + for (int i = 0; i < 3; i++) pool.free(QUERY_ID, "other" + i, PLAN_NODE_ID, 256); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testConcurrentLegacyReservationsKeepIndependentQuotasAndAccurateSideTotals() + throws Exception { + for (boolean samePlanNode : new boolean[] {true, false}) { + MemoryPool legacyPool = new MemoryPool("legacy", new MemoryManager(1024), 512); + ExecutorService executor = Executors.newFixedThreadPool(16); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch attempted = new CountDownLatch(16); + CountDownLatch release = new CountDownLatch(1); + List> tasks = new ArrayList<>(); + try { + for (int i = 0; i < 16; i++) { + String id = samePlanNode ? "same" : "pipeline-" + i; + legacyPool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT); + tasks.add( + executor.submit( + () -> { + start.await(); + boolean reserved = + legacyPool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128); + attempted.countDown(); + release.await(); + if (reserved) { + legacyPool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32); + } + return null; + })); + } + start.countDown(); + Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS)); + long expected = samePlanNode ? 128 : 512; + Assert.assertEquals(expected, legacyPool.getReservedBytes()); + Assert.assertEquals( + expected, + legacyPool.getFragmentMemoryReservedBytes( + QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + if (samePlanNode) { + Assert.assertEquals( + 128, + legacyPool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "same")); + } else { + for (int i = 0; i < 16; i++) { + Assert.assertEquals( + 32, + legacyPool.getPlanNodeMemoryReservedBytes( + QUERY_ID, FRAGMENT_INSTANCE_ID, "pipeline-" + i)); + } + } + release.countDown(); + for (Future task : tasks) { + task.get(10, TimeUnit.SECONDS); + } + Assert.assertEquals(0, legacyPool.getReservedBytes()); + Assert.assertEquals(0, legacyPool.getQueryMemoryReservedBytes(QUERY_ID)); + } finally { + start.countDown(); + release.countDown(); + executor.shutdownNow(); + } + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java new file mode 100644 index 0000000000000..2b3c847e89e46 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -0,0 +1,629 @@ +/* + * 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.planner; + +import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; +import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; +import org.apache.iotdb.calc.execution.operator.process.TableMergeSortOperator; +import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.memory.MemoryManager; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SourceHandleListener; +import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; +import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; +import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; +import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; +import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; +import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata; +import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.apache.tsfile.read.common.block.column.TsBlockSerde; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.concurrent.ExecutorService; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class LocalExecutionPlanContextBudgetTest { + private long previousBudget; + private int previousBlockSize; + private MemoryPool pool; + private LocalExecutionPlanContext context; + + @Before + public void setUp() { + previousBudget = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + previousBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1024); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(64); + pool = new MemoryPool("budget-test", new MemoryManager(4096), 1024); + context = newContext(); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(previousBudget); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(previousBlockSize); + } + + private LocalExecutionPlanContext newContext() { + FragmentInstanceContext instance = + FragmentInstanceContext.createFragmentInstanceContext( + new FragmentInstanceId(new PlanFragmentId(new QueryId("budget"), 0), "0"), + mock(FragmentInstanceStateMachine.class), + mock(MemoryReservationManager.class)); + LocalExecutionPlanContext result = + new LocalExecutionPlanContext(new TypeProvider(), instance, (DataNodeQueryContext) null); + result.setExchangeSumNum(4); + return result; + } + + private static PlanNode node(PlanNodeType type, PlanNode... children) { + Class nodeClass; + switch (type) { + case TABLE_EXCHANGE_NODE: + nodeClass = org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExchangeNode.class; + break; + case EXCHANGE: + nodeClass = + org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ExchangeNode.class; + break; + case TABLE_COLLECT_NODE: + nodeClass = CollectNode.class; + break; + case TABLE_TOPK_NODE: + nodeClass = TopKNode.class; + break; + case TREE_COLLECT: + nodeClass = + org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.CollectNode.class; + break; + default: + nodeClass = PlanNode.class; + break; + } + PlanNode node = mock(nodeClass); + when(node.getType()).thenReturn(type); + when(node.getChildren()).thenReturn(Arrays.asList(children)); + return node; + } + + private static PlanNode collect() { + return new CollectNode( + new PlanNodeId("collect"), + Arrays.asList( + node(PlanNodeType.TABLE_EXCHANGE_NODE), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + Collections.emptyList()); + } + + private CollectOperator generateTableCollect(PlanNode plan) { + Operator first = mock(Operator.class); + Operator second = mock(Operator.class); + when(first.isBlocked()).thenReturn(SettableFuture.create()); + doReturn(Operator.NOT_BLOCKED).when(second).isBlocked(); + when(first.calculateMaxPeekMemoryWithCounter()).thenReturn(100L); + when(first.calculateRetainedSizeAfterCallingNext()).thenReturn(30L); + when(second.calculateMaxPeekMemoryWithCounter()).thenReturn(80L); + when(second.calculateRetainedSizeAfterCallingNext()).thenReturn(50L); + doReturn(first).when(plan.getChildren().get(0)).accept(any(), same(context)); + doReturn(second).when(plan.getChildren().get(1)).accept(any(), same(context)); + return (CollectOperator) + plan.accept(new DataNodeTableOperatorGenerator(mock(Metadata.class)), context); + } + + private ISourceHandle addSource(LocalExecutionPlanContext target) { + ISourceHandle handle = mock(SourceHandle.class); + ExchangeOperator operator = mock(ExchangeOperator.class); + when(operator.getSourceHandle()).thenReturn(handle); + target.addExchangeOperator(operator); + return handle; + } + + private void assertAdmissionIsShared(boolean shared) { + assertAdmissionIsShared(shared, MemorySide.INPUT); + } + + private void assertAdmissionIsShared(boolean shared, MemorySide side) { + assertAdmissionIsShared(shared, side, shared ? 1024 : 256); + } + + private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit) { + FragmentInstanceContext instance = context.getDriverContext().getFragmentInstanceContext(); + if (shared) { + assertEquals( + limit, + side == MemorySide.INPUT ? instance.getInputMaxBytes() : instance.getOutputMaxBytes()); + } else { + assertEquals(0, instance.getInputMaxBytes()); + assertEquals(0, instance.getOutputMaxBytes()); + } + String fi = context.getFragmentInstanceId().getFragmentInstanceId(); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a" + side, side); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b" + side, side); + assertTrue( + pool.reserveWithPriority("budget", fi, "a" + side, limit, limit, false).isReserveSuccess()); + MemoryPool.MemoryReservationResult second = + pool.reserveWithPriority("budget", fi, "b" + side, limit, limit, false); + if (shared) { + assertFalse(second.getFuture().isDone()); + pool.free("budget", fi, "a" + side, limit); + assertTrue(second.getFuture().isDone()); + } else { + assertTrue(second.isReserveSuccess()); + pool.free("budget", fi, "a" + side, limit); + } + pool.free("budget", fi, "b" + side, limit); + } + + @Test + public void testRealSourceHandlesShareCanonicalFragmentAccountAndCleanup() throws Exception { + FragmentInstanceId id = context.getFragmentInstanceId(); + String fi = FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(id.toThrift()); + LocalMemoryManager manager = mock(LocalMemoryManager.class); + when(manager.getQueryPool()).thenReturn(pool); + SourceHandle first = realSourceHandle(id, "first", manager); + SourceHandle second = realSourceHandle(id, "second", manager); + for (SourceHandle source : Arrays.asList(first, second)) { + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(source); + context.addExchangeOperator(exchange); + } + try { + context.setExchangeMemoryLimits(collect(), pool); + assertEquals(1024, context.getInstanceContext().getInputMaxBytes()); + first.isBlocked(); + second.isBlocked(); + first.updatePendingDataBlockInfo(0, List.of(1024L)); + second.updatePendingDataBlockInfo(0, List.of(1024L)); + assertEquals(1024, pool.getFragmentMemoryReservedBytes("budget", fi, MemorySide.INPUT)); + assertEquals(0, pool.getPlanNodeMemoryReservedBytes("budget", fi, "second")); + assertEquals(1, pool.getMemoryReservationSize()); + first.abort(); + assertEquals(1024, pool.getPlanNodeMemoryReservedBytes("budget", fi, "second")); + } finally { + first.abort(); + second.abort(); + pool.deRegisterFragmentInstanceFromQueryMemoryMap("budget", fi, true); + } + assertEquals(0, pool.getMemoryReservationSize()); + assertEquals(0, pool.getReservedBytes()); + assertEquals(0, pool.getQueryMemoryReservationSize()); + } + + private static SourceHandle realSourceHandle( + FragmentInstanceId id, String planNode, LocalMemoryManager manager) { + return new SourceHandle( + new TEndPoint("remote", 12345), + new TFragmentInstanceId(id.getQueryId().getId(), 1, "0"), + id.toThrift(), + planNode, + 0, + manager, + mock(ExecutorService.class), + new TsBlockSerde(), + mock(SourceHandleListener.class), + mock(IClientManager.class)); + } + + @Test + public void testPureCollectUsesSharedSideBudgets() { + ISourceHandle source = addSource(context); + ISourceHandle secondSource = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(source).setMaxBytesCanReserve(512); + verify(secondSource).setMaxBytesCanReserve(512); + verify(sink).setMaxBytesCanReserve(512); + assertAdmissionIsShared(true, MemorySide.INPUT, 512); + assertAdmissionIsShared(true, MemorySide.OUTPUT, 512); + } + + @Test + public void testOddBudgetIsInitializedAfterBothDirectionsAreKnown() { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1025); + pool = new MemoryPool("odd-budget", new MemoryManager(4096), 1025); + context = newContext(); + ISourceHandle source = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(source).setMaxBytesCanReserve(512); + verify(sink).setMaxBytesCanReserve(513); + assertAdmissionIsShared(true, MemorySide.INPUT, 512); + assertAdmissionIsShared(true, MemorySide.OUTPUT, 513); + } + + @Test + public void testInputOnlyCollectUsesFullBudgetEvenForOneFullBlock() { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + context.setExchangeSumNum(1); + ISourceHandle source = addSource(context); + context.setExchangeMemoryLimits( + node(PlanNodeType.TABLE_COLLECT_NODE, node(PlanNodeType.TABLE_EXCHANGE_NODE)), pool); + verify(source).setMaxBytesCanReserve(1024); + assertEquals(0, context.getDriverContext().getFragmentInstanceContext().getOutputMaxBytes()); + assertAdmissionIsShared(true); + } + + @Test + public void testOutputOnlyCollectUsesFullBudgetEvenForOneFullBlock() { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + context.setExchangeSumNum(1); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits( + node( + PlanNodeType.IDENTITY_SINK, + node( + PlanNodeType.TABLE_COLLECT_NODE, + node(PlanNodeType.SERIES_SCAN), + node(PlanNodeType.SERIES_SCAN))), + pool); + verify(sink).setMaxBytesCanReserve(1024); + assertEquals(0, context.getDriverContext().getFragmentInstanceContext().getInputMaxBytes()); + assertAdmissionIsShared(true, MemorySide.OUTPUT); + } + + private void addLocalInput() { + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(mock(LocalSourceHandle.class)); + context.addExchangeOperator(exchange); + } + + @Test + public void testLocalInputsDoNotSplitOutputBudget() { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + addLocalInput(); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(sink).setMaxBytesCanReserve(1024); + assertEquals(0, context.getDriverContext().getFragmentInstanceContext().getInputMaxBytes()); + assertAdmissionIsShared(true, MemorySide.OUTPUT); + } + + @Test + public void testMixedLocalAndRemoteInputsStillSplitBidirectionalBudget() { + addLocalInput(); + ISourceHandle source = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(source).setMaxBytesCanReserve(512); + verify(sink).setMaxBytesCanReserve(512); + assertAdmissionIsShared(true, MemorySide.INPUT, 512); + assertAdmissionIsShared(true, MemorySide.OUTPUT, 512); + } + + @Test + public void testLocalInputsWithoutSinkDoNotRegisterUnusedSharedBudget() { + addLocalInput(); + context.setExchangeMemoryLimits(collect(), pool); + assertAdmissionIsShared(false); + } + + @Test + public void testMappingCollectUsesSharedBudget() { + addSource(context); + PlanNode union = + new UnionNode( + new PlanNodeId("union"), + Arrays.asList( + node(PlanNodeType.TABLE_EXCHANGE_NODE), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + ArrayListMultimap.create(), + Collections.emptyList()); + context.setExchangeMemoryLimits(union, pool); + assertAdmissionIsShared(true); + } + + @Test + public void testTreePlanningRetainsLegacySourceQuota() { + ISourceHandle source = addSource(context); + context.setMaxBytesOneHandleCanReserve(); + verify(source).setMaxBytesCanReserve(256); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node( + PlanNodeType.TREE_COLLECT, + node(PlanNodeType.EXCHANGE), + node(PlanNodeType.EXCHANGE)))); + assertAdmissionIsShared(false); + } + + @Test + public void testTableCollectEnablesReadyConsumptionOnlyForEligibleFragment() { + PlanNode plan = collect(); + CollectOperator collect = generateTableCollect(plan); + addSource(context); + context.setExchangeMemoryLimits(plan, pool); + assertTrue(collect.isBlocked().isDone()); + assertEquals(80, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(150, collect.calculateMaxPeekMemory()); + assertAdmissionIsShared(true); + } + + @Test + public void testTableUnionEnablesReadyConsumptionOnlyForEligibleFragment() { + PlanNode union = + new UnionNode( + new PlanNodeId("union"), + Arrays.asList( + node(PlanNodeType.TABLE_EXCHANGE_NODE), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + ArrayListMultimap.create(), + Collections.emptyList()); + CollectOperator collect = generateTableCollect(union); + addSource(context); + context.setExchangeMemoryLimits(union, pool); + assertTrue(collect.isBlocked().isDone()); + assertEquals(80, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(150, collect.calculateMaxPeekMemory()); + assertAdmissionIsShared(true); + } + + @Test + public void testSingleChildTopKRetainsLegacyBudgetAndConsumption() { + PlanNode child = collect(); + CollectOperator collect = generateTableCollect(child); + ISourceHandle source = addSource(context); + PlanNode topK = node(PlanNodeType.TABLE_TOPK_NODE, child); + context.setExchangeMemoryLimits(topK, pool); + verify(source).setMaxBytesCanReserve(256); + assertFalse(collect.isBlocked().isDone()); + assertAdmissionIsShared(false); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.TABLE_COLLECT_NODE, topK, collect()))); + assertTrue( + SharedExchangeMemoryEligibility.isEligible( + node( + PlanNodeType.TABLE_COLLECT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE, topK), + node(PlanNodeType.TABLE_EXCHANGE_NODE)))); + } + + @Test + public void testMultiplePipelinesRetainSequentialCollectConsumption() { + PlanNode plan = collect(); + CollectOperator collect = generateTableCollect(plan); + LocalExecutionPlanContext child = context.createSubContext(); + child.addPipelineDriverFactory(mock(Operator.class), child.getDriverContext(), 0); + addSource(context); + context.setExchangeMemoryLimits(plan, pool); + assertFalse(collect.isBlocked().isDone()); + assertAdmissionIsShared(false); + } + + @Test + public void testNestedDependencyDisablesSharingForEntireFragment() { + for (PlanNodeType dependency : + new PlanNodeType[] { + PlanNodeType.TABLE_MERGESORT_NODE, PlanNodeType.MERGE_SORT, PlanNodeType.AGG_MERGE_SORT, + PlanNodeType.TABLE_JOIN_NODE, PlanNodeType.TABLE_SEMI_JOIN_NODE, + PlanNodeType.TABLE_TOPK_NODE, + PlanNodeType.COUNT_MERGE, PlanNodeType.DEVICE_VIEW + }) { + PlanNode dependent = + node( + dependency, + node(PlanNodeType.TABLE_EXCHANGE_NODE), + node(PlanNodeType.TABLE_EXCHANGE_NODE)); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.TABLE_COLLECT_NODE, dependent, collect()))); + assertFalse( + SharedExchangeMemoryEligibility.isEligible(node(dependency, collect(), collect()))); + } + ISourceHandle source = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setExchangeSumNum(2); + context.setISink(sink); + context.setExchangeSumNum(4); + context.setExchangeMemoryLimits( + node( + PlanNodeType.TABLE_COLLECT_NODE, + collect(), + node( + PlanNodeType.TABLE_MERGESORT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE), + node(PlanNodeType.TABLE_EXCHANGE_NODE))), + pool); + verify(source).setMaxBytesCanReserve(256); + // The legacy sink quota is captured at its original point in operator generation. + verify(sink).setMaxBytesCanReserve(512); + assertAdmissionIsShared(false); + } + + @Test + public void testSequentialAndUnknownMultiInputOperatorsDisableSharing() { + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.IDENTITY_SINK, collect(), collect()))); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.SCHEMA_QUERY_MERGE, collect(), collect()))); + } + + @Test + public void testExchangeBoundaryDoesNotInspectUpstreamFragment() { + PlanNode remote = node(PlanNodeType.TABLE_MERGESORT_NODE, collect(), collect()); + assertTrue( + SharedExchangeMemoryEligibility.isEligible( + node( + PlanNodeType.TABLE_COLLECT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE, remote), + node(PlanNodeType.TABLE_EXCHANGE_NODE)))); + } + + @Test + public void testNoCollectKeepsLegacyBudget() { + context.setExchangeMemoryLimits( + node( + PlanNodeType.TABLE_MERGESORT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE), + node(PlanNodeType.TABLE_EXCHANGE_NODE)), + pool); + assertAdmissionIsShared(false); + } + + @Test + public void testInternalPipelineDisablesSharingForEntireFragment() { + LocalExecutionPlanContext child = context.createSubContext(); + child.addPipelineDriverFactory(mock(Operator.class), child.getDriverContext(), 0); + ISourceHandle source = addSource(child); + PipelineSourceHandle pipelineSource = mock(PipelineSourceHandle.class); + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(pipelineSource); + context.addExchangeOperator(exchange); + context.setExchangeMemoryLimits(collect(), pool); + verify(source).setMaxBytesCanReserve(256); + verify(pipelineSource).setMaxBytesCanReserve(256); + assertAdmissionIsShared(false); + } + + @Test + public void testPipelineSchedulingDependencyDisablesSharing() { + context.getDriverContext().setDependencyDriverIndex(0); + context.setExchangeMemoryLimits(collect(), pool); + assertAdmissionIsShared(false); + } + + @Test + public void testSideSmallerThanOneBlockKeepsLegacyBudget() { + addSource(context); + context.setISink(mock(ISinkHandle.class)); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + context.setExchangeMemoryLimits(collect(), pool); + assertAdmissionIsShared(false); + } + + @Test + public void testSinkOnlyFragmentCanStillReserveOneFullBlock() { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(64); + context = newContext(); + context.setExchangeSumNum(1); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits( + node(PlanNodeType.IDENTITY_SINK, node(PlanNodeType.SERIES_SCAN)), pool); + verify(sink).setMaxBytesCanReserve(64); + } + + @Test + public void testMergeSortGetsBothHeadsDespitePrefetchOnOneInput() throws Exception { + context.setExchangeMemoryLimits( + node(PlanNodeType.TABLE_MERGESORT_NODE, collect(), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + pool); + String fi = context.getFragmentInstanceId().getFragmentInstanceId(); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", MemorySide.INPUT); + for (int i = 0; i < 4; i++) { + assertTrue(pool.reserveWithPriority("budget", fi, "a", 64, 256, false).isReserveSuccess()); + } + MemoryPool.MemoryReservationResult prefetch = + pool.reserveWithPriority("budget", fi, "a", 64, 256, false); + assertFalse(prefetch.getFuture().isDone()); + MemoryPool.MemoryReservationResult headB = + pool.reserveWithPriority("budget", fi, "b", 64, 256, false); + assertTrue(headB.isReserveSuccess()); + Operator a = mock(Operator.class); + Operator b = mock(Operator.class); + doReturn(Operator.NOT_BLOCKED).when(a).isBlocked(); + doReturn(headB.getFuture()).when(b).isBlocked(); + when(a.hasNextWithTimer()).thenReturn(true); + when(b.hasNextWithTimer()).thenReturn(true); + TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32)); + builder.getTimeColumnBuilder().writeLong(1); + builder.getColumnBuilder(0).writeInt(1); + builder.declarePosition(); + TsBlock first = builder.build(); + builder.reset(); + builder.getTimeColumnBuilder().writeLong(2); + builder.getColumnBuilder(0).writeInt(2); + builder.declarePosition(); + TsBlock second = builder.build(); + when(a.nextWithTimer()) + .thenAnswer( + call -> { + pool.free("budget", fi, "a", 64); + return first; + }); + when(b.nextWithTimer()) + .thenAnswer( + call -> { + pool.free("budget", fi, "b", 64); + return second; + }); + TableMergeSortOperator merge = + new TableMergeSortOperator( + mock(CommonOperatorContext.class), + new ArrayList<>(Arrays.asList(a, b)), + Collections.singletonList(TSDataType.INT32), + Comparator.comparingLong(key -> key.tsBlock.getTimeByIndex(key.rowIndex))); + assertTrue(merge.isBlocked().isDone()); + TsBlock result = merge.next(); + assertTrue(result != null && !result.isEmpty()); + assertTrue(prefetch.getFuture().isDone()); + pool.free("budget", fi, "a", 256); + merge.close(); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java new file mode 100644 index 0000000000000..31c562199fd63 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java @@ -0,0 +1,245 @@ +/* + * 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.planner; + +import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; +import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.common.SqlDialect; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeService; +import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; +import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; +import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; +import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; +import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; + +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class LocalExecutionPlannerBudgetTest { + private long previousBudget; + private int previousBlockSize; + private boolean previousEstimation; + private int previousDataNodeId; + + @Before + public void setUp() { + previousDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId(); + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); + previousBudget = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + previousBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + previousEstimation = + IoTDBDescriptor.getInstance().getMemoryConfig().isEnableQueryMemoryEstimation(); + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1024); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(64); + IoTDBDescriptor.getInstance().getMemoryConfig().setEnableQueryMemoryEstimation(true); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(previousDataNodeId); + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(previousBudget); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(previousBlockSize); + IoTDBDescriptor.getInstance() + .getMemoryConfig() + .setEnableQueryMemoryEstimation(previousEstimation); + } + + @Test + public void testTablePlanEnablesSharingBeforeMemoryEstimationAndRootFactory() throws Exception { + LocalExecutionPlanner planner = LocalExecutionPlanner.getInstance(); + MemoryPool pool = + MPPDataExchangeService.getInstance() + .getMPPDataExchangeManager() + .getLocalMemoryManager() + .getQueryPool(); + FragmentInstanceContext instance = instance("table_planner_budget", SqlDialect.TABLE); + PlanNode first = mock(PlanNode.class); + PlanNode second = mock(PlanNode.class); + SourceHandle firstHandle = mock(SourceHandle.class); + SourceHandle secondHandle = mock(SourceHandle.class); + addInput(first, firstHandle, false, 100, 30); + addInput(second, secondHandle, true, 80, 50); + CollectNode plan = + spy( + new CollectNode( + new PlanNodeId("collect"), Arrays.asList(first, second), Collections.emptyList())); + doReturn(PlanNodeType.TABLE_COLLECT_NODE).when(plan).getType(); + when(first.getChildren()).thenReturn(Collections.emptyList()); + when(second.getChildren()).thenReturn(Collections.emptyList()); + long freeBefore = planner.getFreeMemoryForOperators(); + try { + List factories = + planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext(1)); + assertEquals(1, factories.size()); + CollectOperator collect = (CollectOperator) factories.get(0).getOperation(); + // Check the actual reservation, not only the formula after planning has completed. + assertEquals(150 + collect.ramBytesUsed(), freeBefore - planner.getFreeMemoryForOperators()); + assertEquals(1024, instance.getInputMaxBytes()); + assertEquals(0, instance.getOutputMaxBytes()); + assertTrue(collect.isBlocked().isDone()); + verify(firstHandle).setMaxBytesCanReserve(1024); + verify(secondHandle).setMaxBytesCanReserve(1024); + + String query = instance.getId().getQueryId().getId(); + String fi = instance.getId().getFragmentInstanceId(); + pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "first", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "second", MemorySide.INPUT); + assertTrue( + pool.reserveWithPriority(query, fi, "first", 1024, 1024, false).isReserveSuccess()); + MemoryPool.MemoryReservationResult pending = + pool.reserveWithPriority(query, fi, "second", 64, 1024, false); + assertFalse(pending.getFuture().isDone()); + pool.free(query, fi, "first", 1024); + assertTrue(pending.getFuture().isDone()); + pool.free(query, fi, "second", 64); + collect.close(); + } finally { + pool.deRegisterFragmentInstanceFromQueryMemoryMap( + instance.getId().getQueryId().getId(), instance.getId().getFragmentInstanceId(), true); + long reserved = freeBefore - planner.getFreeMemoryForOperators(); + if (reserved > 0) { + planner.releaseToFreeMemoryForOperators(reserved); + } + } + } + + @Test + public void testNullSessionRetainsTreeGeneratorAndIndependentInputQuotas() throws Exception { + FragmentInstanceContext instance = instance("null_session_planner_budget", null); + PlanNode plan = mock(PlanNode.class); + when(plan.getType()).thenReturn(PlanNodeType.TREE_COLLECT); + when(plan.getPlanNodeId()).thenReturn(new PlanNodeId("collect")); + SourceHandle firstHandle = mock(SourceHandle.class); + SourceHandle secondHandle = mock(SourceHandle.class); + Operator first = input(false, 100, 30); + Operator second = input(true, 80, 50); + doAnswer( + invocation -> { + assertTrue(invocation.getArgument(0) instanceof OperatorTreeGenerator); + LocalExecutionPlanContext context = invocation.getArgument(1); + registerInput(context, firstHandle); + registerInput(context, secondHandle); + CollectOperator collect = + new CollectOperator( + context + .getDriverContext() + .addOperatorContext(0, new PlanNodeId("collect"), "CollectOperator"), + new java.util.ArrayList<>(Arrays.asList(first, second))); + return collect; + }) + .when(plan) + .accept(any(), any()); + LocalExecutionPlanner planner = LocalExecutionPlanner.getInstance(); + long freeBefore = planner.getFreeMemoryForOperators(); + try { + List factories = + planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext(1)); + assertEquals(1, factories.size()); + CollectOperator collect = (CollectOperator) factories.get(0).getOperation(); + assertEquals(100 + collect.ramBytesUsed(), freeBefore - planner.getFreeMemoryForOperators()); + assertFalse(collect.isBlocked().isDone()); + assertEquals(0, instance.getInputMaxBytes()); + verify(firstHandle).setMaxBytesCanReserve(512); + verify(secondHandle).setMaxBytesCanReserve(512); + collect.close(); + } finally { + long reserved = freeBefore - planner.getFreeMemoryForOperators(); + if (reserved > 0) { + planner.releaseToFreeMemoryForOperators(reserved); + } + } + } + + private static FragmentInstanceContext instance(String query, SqlDialect dialect) { + FragmentInstanceContext instance = + spy( + FragmentInstanceContext.createFragmentInstanceContext( + new FragmentInstanceId(new PlanFragmentId(new QueryId(query), 0), "0"), + mock(FragmentInstanceStateMachine.class), + mock(MemoryReservationManager.class))); + SessionInfo session = null; + if (dialect != null) { + session = mock(SessionInfo.class); + when(session.getSqlDialect()).thenReturn(dialect); + } + doReturn(session).when(instance).getSessionInfo(); + return instance; + } + + private static void addInput( + PlanNode node, SourceHandle handle, boolean ready, long peak, long retained) { + Operator operator = input(ready, peak, retained); + doAnswer( + invocation -> { + LocalExecutionPlanContext context = invocation.getArgument(1); + registerInput(context, handle); + return operator; + }) + .when(node) + .accept(any(), any()); + } + + private static void registerInput(LocalExecutionPlanContext context, SourceHandle handle) { + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(handle); + context.addExchangeOperator(exchange); + context.addExchangeSumNum(1); + } + + private static Operator input(boolean ready, long peak, long retained) { + Operator operator = mock(Operator.class); + doReturn(ready ? Operator.NOT_BLOCKED : SettableFuture.create()).when(operator).isBlocked(); + when(operator.calculateMaxPeekMemoryWithCounter()).thenReturn(peak); + when(operator.calculateRetainedSizeAfterCallingNext()).thenReturn(retained); + return operator; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java new file mode 100644 index 0000000000000..aa20dd0005bc2 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java @@ -0,0 +1,158 @@ +/* + * 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.planner; + +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.AggregationNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode; + +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class CollectInputPropertyTest { + private static PlanNode plan(String sql) { + PlanTester tester = new PlanTester(); + tester.createPlan(sql); + return tester.getFragmentPlan(0); + } + + private static List nodes(PlanNode node, Class type) { + List result = new ArrayList<>(); + if (type.isInstance(node)) result.add(type.cast(node)); + for (PlanNode child : node.getChildren()) result.addAll(nodes(child, type)); + return result; + } + + @Test + public void testAggregationAboveGlobalWindowHashesInterleavedGroups() { + for (String aggregate : new String[] {"SUM(s1)", "COUNT(DISTINCT s1)"}) { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, " + + aggregate + + " FROM (" + + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n " + + "FROM table1 WHERE tag1 != 'shenzhen') t WHERE n > 0 GROUP BY tag1, tag2, tag3"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + for (AggregationNode aggregation : aggregations) { + if (!nodes(aggregation.getChild(), CollectNode.class).isEmpty()) { + assertTrue(aggregation.getPreGroupedSymbols().isEmpty()); + assertFalse(aggregation.isStreamable()); + } + } + assertTrue( + "Hash aggregation should not add a grouping sort", nodes(root, SortNode.class).isEmpty()); + } + } + + @Test + public void testPartialPreGroupingAlsoFallsBackToHash() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, s1, COUNT(*) FROM (" + + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n " + + "FROM table1 WHERE tag1 != 'shenzhen') t WHERE n > 0 GROUP BY tag1, tag2, tag3, s1"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertTrue(aggregations.get(0).getPreGroupedSymbols().isEmpty()); + } + + @Test + public void testNestedWindowRestoresPartitionContiguityWithSort() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, SUM(s1) OVER (" + + "PARTITION BY tag1, tag2, tag3 ORDER BY s1) AS running_sum " + + "FROM (SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1) t WHERE n > 0"); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertTrue(nodes(root, GroupNode.class).isEmpty()); + List windows = nodes(root, WindowNode.class); + assertEquals(2, windows.size()); + WindowNode outer = + windows.stream() + .filter(w -> !w.getSpecification().getPartitionBy().isEmpty()) + .findFirst() + .get(); + assertFalse(nodes(outer.getChild(), SortNode.class).isEmpty()); + } + + @Test + public void testStreamSortAfterGlobalWindowUsesFullSort() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n " + + "FROM table1 ORDER BY tag1, tag2, tag3, s1"); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertFalse(nodes(root, SortNode.class).isEmpty()); + assertTrue(nodes(root, StreamSortNode.class).isEmpty()); + } + + @Test + public void testConstantFillDoesNotRemoveRestoredSort() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1 " + + "FILL METHOD CONSTANT 0 ORDER BY tag1, tag2, tag3"); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertFalse(nodes(root, SortNode.class).isEmpty()); + assertTrue(nodes(root, StreamSortNode.class).isEmpty()); + } + + @Test + public void testUnionFinalAggregationDoesNotInheritBranchGrouping() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, SUM(s1) FROM (" + + "SELECT tag1, tag2, tag3, s1 FROM table1 UNION ALL " + + "SELECT tag1, tag2, tag3, s1 FROM table2) u GROUP BY tag1, tag2, tag3"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + boolean foundFinalCollect = false; + for (AggregationNode aggregation : aggregations) { + if (aggregation.getChild() instanceof CollectNode) { + foundFinalCollect = true; + assertTrue(aggregation.getPreGroupedSymbols().isEmpty()); + } + } + assertTrue(foundFinalCollect); + } + + @Test + public void testOrderedScanBranchesKeepStreamingAggregation() { + PlanNode root = + plan("SELECT tag1, tag2, tag3, SUM(s1 + 1) FROM table1 GROUP BY tag1, tag2, tag3"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + assertTrue(aggregations.stream().anyMatch(AggregationNode::isStreamable)); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java index e2326688112b2..b5b69f7663f03 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java @@ -26,6 +26,9 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.OrderingScheme; import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.ValuesNode; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.execution.warnings.WarningCollector; @@ -36,6 +39,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExternalTsFileScanNode; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableListMultimap; import com.google.common.collect.ImmutableMap; import org.junit.Test; @@ -108,4 +112,54 @@ region, new ColumnSchema("region", STRING, false, TAG), resource.closeByQueryExecution(); } } + + @Test + public void testUnionDoesNotInheritFirstBranchGrouping() { + Symbol key = new Symbol("key"); + OrderingScheme ordering = + new OrderingScheme(ImmutableList.of(key), ImmutableMap.of(key, ASC_NULLS_LAST)); + PlanNode left = + new SortNode( + new PlanNodeId("leftSort"), + new ValuesNode(new PlanNodeId("left"), ImmutableList.of(key), ImmutableList.of()), + ordering, + false, + false); + PlanNode right = + new SortNode( + new PlanNodeId("rightSort"), + new ValuesNode(new PlanNodeId("right"), ImmutableList.of(key), ImmutableList.of()), + ordering, + false, + false); + UnionNode union = + new UnionNode( + new PlanNodeId("union"), + ImmutableList.of(left, right), + ImmutableListMultimap.of(key, key, key, key), + ImmutableList.of(key)); + assertTrue( + optimizeGroup(new GroupNode(new PlanNodeId("group"), union, ordering, 1)) + instanceof SortNode); + // An explicit sort above the fan-in restores the grouping property. + SortNode restored = new SortNode(new PlanNodeId("restored"), union, ordering, false, false); + assertTrue( + optimizeGroup(new GroupNode(new PlanNodeId("group2"), restored, ordering, 1)) + instanceof GroupNode); + } + + private PlanNode optimizeGroup(GroupNode group) { + return new ParallelizeGrouping() + .optimize( + group, + new PlanOptimizer.Context( + SESSION_INFO, + new Analysis(null, Collections.emptyMap()), + TEST_MATADATA, + new MPPQueryContext(new QueryId("group_union")), + new SymbolAllocator(), + new QueryId("group_union"), + WarningCollector.NOOP, + PlanOptimizersStatsCollector.createPlanOptimizersStatsCollector())); + } }