diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/IdSetSubqueryRewriter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/IdSetSubqueryRewriter.java new file mode 100644 index 000000000000..ea8980ae129b --- /dev/null +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/IdSetSubqueryRewriter.java @@ -0,0 +1,264 @@ +/** + * 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.pinot.broker.requesthandler; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.sql.SqlBasicCall; +import org.apache.calcite.sql.SqlCall; +import org.apache.calcite.sql.SqlCharStringLiteral; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlIdentifier; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlLiteral; +import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.sql.SqlOrderBy; +import org.apache.calcite.sql.SqlUnresolvedFunction; +import org.apache.calcite.sql.SqlWith; +import org.apache.calcite.sql.SqlWithItem; +import org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.calcite.sql.parser.SqlParserUtil; +import org.apache.calcite.sql.util.SqlBasicVisitor; +import org.apache.calcite.util.Litmus; +import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.core.query.utils.idset.IdSets; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.sql.parsers.CalciteSqlParser; + + +/// Runs the IdSet subqueries of a multi-stage query before the query is compiled, and puts the IdSets they return in +/// their place, so that the planner and the servers only see `IN_ID_SET` calls on literals: +/// - `IN_SUBQUERY(expr, 'subquery')` becomes `IN_ID_SET(expr, 'idSet')`, anywhere in the query. +/// - `IN_ID_SET(expr, (subquery))` becomes `IN_ID_SET(expr, 'idSet')` when the subquery validates on its own, i.e. +/// when it is not correlated, and does not name a WITH item of the query. Otherwise the subquery stays, and runs as +/// part of the query. +/// +/// A subquery that appears more than once runs once. To explain a query, no subquery runs: the empty IdSet takes the +/// place of each result. +/// +/// The rewrite changes the parsed query in place, so it is not thread-safe. +final class IdSetSubqueryRewriter { + private static final String IN_SUBQUERY = "insubquery"; + private static final String IN_ID_SET = "inidset"; + + static final String EMPTY_ID_SET; + + static { + try { + EMPTY_ID_SET = IdSets.emptyIdSet().toBase64String(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + private final String _sql; + private final boolean _explain; + private final SubqueryRunner _runner; + // The IdSets of the subqueries, by subquery, in the order the subqueries ran + private final Map _idSets = new LinkedHashMap<>(); + + private IdSetSubqueryRewriter(String sql, boolean explain, SubqueryRunner runner) { + _sql = sql; + _explain = explain; + _runner = runner; + } + + /// Validates and runs the subqueries for [IdSetSubqueryRewriter]. + interface SubqueryRunner { + /// Returns whether the subquery validates on its own, i.e. whether it does not refer to the query it is in. + boolean validates(String subquery); + + /// Runs the subquery and returns the serialized IdSet it returns, or `null` if it returns no rows or NULL. + @Nullable + String run(String subquery); + } + + /// Runs the IdSet subqueries of the query and replaces them with the IdSets they return, or with the empty IdSet to + /// explain the query. Returns the subqueries that were replaced, each once, in the order they ran. + /// + /// @param sql the text of the query, which scalar subqueries are read from + /// @param query the parsed query, without the EXPLAIN + static List rewrite(String sql, SqlNode query, boolean explain, SubqueryRunner runner) { + IdSetSubqueryRewriter rewriter = new IdSetSubqueryRewriter(sql, explain, runner); + rewriter.rewrite(query, Set.of()); + return new ArrayList<>(rewriter._idSets.keySet()); + } + + /// @param withNames the names of the WITH items in scope, in lower case + private void rewrite(SqlNode node, Set withNames) { + IdSetCallFinder finder = new IdSetCallFinder(withNames); + node.accept(finder); + for (IdSetCall idSetCall : finder._calls) { + SqlBasicCall call = idSetCall._call; + if (IN_SUBQUERY.equals(getFunctionName(call))) { + String subquery = getSubquery(call); + call.setOperator(new SqlUnresolvedFunction(new SqlIdentifier("IN_ID_SET", call.getParserPosition()), null, + null, null, null, SqlFunctionCategory.USER_DEFINED_FUNCTION)); + setIdSet(call, getIdSet(subquery)); + } else { + SqlNode subqueryNode = call.operand(1); + String subquery = getText(subqueryNode); + // A subquery that names a WITH item cannot run on its own: it would read the table that the item hides, if any + boolean runsFirst = subquery != null && !hasAnyName(subqueryNode, idSetCall._withNames) + && (_idSets.containsKey(subquery) || _runner.validates(subquery)); + if (runsFirst) { + setIdSet(call, getIdSet(subquery)); + } else { + // The planner runs the subquery as part of the query, but the IdSet subqueries in it still run first + rewrite(subqueryNode, idSetCall._withNames); + } + } + } + } + + private String getIdSet(String subquery) { + return _idSets.computeIfAbsent(subquery, key -> { + String idSet = _explain ? null : _runner.run(key); + return idSet != null ? idSet : EMPTY_ID_SET; + }); + } + + private static void setIdSet(SqlBasicCall call, String idSet) { + call.setOperand(1, SqlLiteral.createCharString(idSet, call.operand(1).getParserPosition())); + } + + @Nullable + private static String getFunctionName(SqlCall call) { + return call.getOperator() instanceof SqlFunction + ? FunctionRegistry.canonicalize(call.getOperator().getName()) + : null; + } + + private static String getSubquery(SqlBasicCall call) { + if (call.operandCount() == 2) { + SqlNode subquery = call.operand(1); + if (subquery.getKind() == SqlKind.LITERAL_CHAIN) { + subquery = SqlLiteral.unchain(subquery); + } + if (subquery instanceof SqlCharStringLiteral) { + return ((SqlCharStringLiteral) subquery).getValueAs(String.class); + } + } + throw QueryErrorCode.QUERY_VALIDATION.asException( + "IN_SUBQUERY takes an expression and a string literal with the subquery, got: " + call); + } + + /// Returns the text of the subquery in the query, with its parentheses, or `null` if the text does not parse to the + /// same subquery, e.g. when the parser read a different text than the given one. + @Nullable + private String getText(SqlNode subquery) { + SqlParserPos pos = subquery.getParserPosition(); + try { + String text = _sql.substring(SqlParserUtil.lineColToIndex(_sql, pos.getLineNum(), pos.getColumnNum()), + SqlParserUtil.lineColToIndex(_sql, pos.getEndLineNum(), pos.getEndColumnNum()) + 1); + SqlNode parsed = CalciteSqlParser.compileToSqlNodeAndOptions(text).getSqlNode(); + return parsed.equalsDeep(subquery, Litmus.IGNORE) ? text : null; + } catch (RuntimeException e) { + return null; + } + } + + /// Returns whether the node has an identifier with one of the given names, ignoring the case. + private static boolean hasAnyName(SqlNode node, Set names) { + if (names.isEmpty()) { + return false; + } + NameFinder nameFinder = new NameFinder(names); + node.accept(nameFinder); + return nameFinder._found; + } + + /// An IdSet subquery, and the names of the WITH items in scope where it appears. + private static class IdSetCall { + final SqlBasicCall _call; + final Set _withNames; + + IdSetCall(SqlBasicCall call, Set withNames) { + _call = call; + _withNames = withNames; + } + } + + /// Finds the IdSet subqueries of a query, except those in other IdSet subqueries. + private static class IdSetCallFinder extends SqlBasicVisitor { + final List _calls = new ArrayList<>(); + // The names of the WITH items in scope, in lower case. Each WITH gets a new set, so the calls can keep theirs. + Set _withNames; + + IdSetCallFinder(Set withNames) { + _withNames = withNames; + } + + @Override + public Void visit(SqlCall call) { + // The ORDER BY of a WITH query parses outside the WITH, but the validator puts it in the scope of the WITH + SqlNode with = call instanceof SqlOrderBy ? ((SqlOrderBy) call).query : call; + if (with instanceof SqlWith) { + Set withNames = _withNames; + _withNames = new HashSet<>(withNames); + for (SqlNode withItem : ((SqlWith) with).withList) { + _withNames.add(((SqlWithItem) withItem).name.getSimple().toLowerCase(Locale.ROOT)); + } + super.visit(call); + _withNames = withNames; + return null; + } + if (call instanceof SqlBasicCall && call.operandCount() > 0) { + String functionName = getFunctionName(call); + if (IN_SUBQUERY.equals(functionName) || (IN_ID_SET.equals(functionName) && call.operandCount() == 2 + && call.operand(1).getKind().belongsTo(SqlKind.QUERY))) { + _calls.add(new IdSetCall((SqlBasicCall) call, _withNames)); + // The tested expression can hold more of them + call.operand(0).accept(this); + return null; + } + } + return super.visit(call); + } + } + + /// Finds whether a node has an identifier with one of the given names, ignoring the case. + private static class NameFinder extends SqlBasicVisitor { + final Set _names; + boolean _found; + + NameFinder(Set names) { + _names = names; + } + + @Override + public Void visit(SqlIdentifier id) { + for (String name : id.names) { + if (_names.contains(name.toLowerCase(Locale.ROOT))) { + _found = true; + } + } + return null; + } + } +} diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index e2708b725e9a..adf76c39c2e4 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -28,6 +28,7 @@ import java.util.Arrays; import java.util.Collection; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Locale; import java.util.Map; @@ -46,7 +47,9 @@ import javax.ws.rs.WebApplicationException; import javax.ws.rs.core.HttpHeaders; import javax.ws.rs.core.Response; +import org.apache.calcite.sql.SqlExplain; import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlNode; import org.apache.commons.lang3.StringUtils; import org.apache.hc.client5.http.io.HttpClientConnectionManager; import org.apache.pinot.broker.api.AccessControl; @@ -67,12 +70,14 @@ import org.apache.pinot.common.response.broker.QueryProcessingException; import org.apache.pinot.common.response.broker.ResultTable; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.DatabaseUtils; import org.apache.pinot.common.utils.ExceptionUtils; import org.apache.pinot.common.utils.NamedThreadFactory; import org.apache.pinot.common.utils.Timer; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.common.utils.request.QueryFingerprintUtils; +import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.common.utils.tls.TlsUtils; import org.apache.pinot.core.routing.MultiClusterRoutingContext; import org.apache.pinot.core.routing.RoutingManager; @@ -101,9 +106,13 @@ import org.apache.pinot.spi.metrics.PinotMeter; import org.apache.pinot.spi.query.QueryExecutionContext; import org.apache.pinot.spi.query.QueryThreadContext; +import org.apache.pinot.spi.trace.DefaultRequestContext; import org.apache.pinot.spi.trace.QueryFingerprint; import org.apache.pinot.spi.trace.RequestContext; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.apache.pinot.sql.parsers.PinotSqlType; import org.apache.pinot.sql.parsers.SqlNodeAndOptions; import org.apache.pinot.sql.parsers.rewriter.RlsUtils; import org.slf4j.Logger; @@ -131,6 +140,9 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { private static final Marker MSE_STATS_MARKER = MarkerFactory.getMarker("MSE_STATS_MARKER"); private static final int NUM_UNAVAILABLE_SEGMENTS_TO_LOG = 10; + // Extra column of the EXPLAIN response with the IdSet subqueries as a JSON array. EXPLAIN does not run them: the plan + // uses the empty IdSet in their place. + private static final String EXPLAIN_ID_SET_SUBQUERIES_FIELD = "ID_SET_SUBQUERIES_NOT_RUN"; private final WorkerManager _workerManager; private final WorkerManager _multiClusterWorkerManager; @@ -449,6 +461,15 @@ protected BrokerResponse handleRequestThrowing(long requestId, String query, Sql long activeDeadlineMs = startTimeMs + timeoutMs; long passiveDeadlineMs = activeDeadlineMs + getExtraPassiveTimeoutMs(options); + // Runs before the query context opens, because each subquery opens its own. As in the single-stage engine, they + // run before the query is checked for access and quota, and each is checked like any other query. + List idSetSubqueries = + rewriteIdSetSubqueries(requestId, query, sqlNodeAndOptions, requesterIdentity, httpHeaders, activeDeadlineMs); + if (!idSetSubqueries.isEmpty()) { + // Keeps the time of the subqueries out of the compilation time of the query + queryTimer = new Timer(queryTimer.getRemainingTimeMs(), TimeUnit.MILLISECONDS); + } + QueryExecutionContext executionContext = new QueryExecutionContext(QueryExecutionContext.QueryType.MSE, requestId, cid, workloadName, startTimeMs, activeDeadlineMs, passiveDeadlineMs, _brokerId, _brokerId, queryHash); @@ -468,7 +489,7 @@ protected BrokerResponse handleRequestThrowing(long requestId, String query, Sql prepareCompiledQueryForPlanning(compiledQuery, requestId, requestContext, httpHeaders); if (sqlNodeAndOptions.getSqlNode().getKind() == SqlKind.EXPLAIN) { - return explain(compiledQuery, requestId, requestContext, queryTimer); + return explain(compiledQuery, requestId, requestContext, queryTimer, idSetSubqueries); } else { return query(compiledQuery, requestId, requesterIdentity, requestContext, httpHeaders, queryTimer, rlsFiltersApplied.get(), queryWasLogged, workloadName); @@ -673,8 +694,10 @@ private long getExtraPassiveTimeoutMs(Map queryOptions) { /// Explains the query and returns the broker response. /// /// Throws using the same conventions as handleRequestThrowing. + /// + /// @param idSetSubqueries the IdSet subqueries that did not run, see [IdSetSubqueryRewriter] private BrokerResponse explain(QueryEnvironment.CompiledQuery query, long requestId, RequestContext requestContext, - Timer timer) + Timer timer, List idSetSubqueries) throws WebApplicationException, QueryException { Map queryOptions = query.getOptions(); @@ -689,9 +712,158 @@ private BrokerResponse explain(QueryEnvironment.CompiledQuery query, long reques () -> query.explain(requestId, fragmentToPlanNode), timer); String plan = queryPlanResult.getExplainPlan(); Map extraFields = queryPlanResult.getExtraFields(); + if (!idSetSubqueries.isEmpty()) { + ArrayNode subqueries = JsonNodeFactory.instance.arrayNode(); + idSetSubqueries.forEach(subqueries::add); + extraFields = new LinkedHashMap<>(extraFields); + extraFields.put(EXPLAIN_ID_SET_SUBQUERIES_FIELD, subqueries.toString()); + } return constructMultistageExplainPlan(query.getTextQuery(), plan, extraFields); } + /// Runs the IdSet subqueries of the query, or not to explain it, and replaces them with their IdSets. See + /// [IdSetSubqueryRewriter]. Returns the subqueries that were replaced. + private List rewriteIdSetSubqueries(long requestId, String query, SqlNodeAndOptions sqlNodeAndOptions, + @Nullable RequesterIdentity requesterIdentity, @Nullable HttpHeaders httpHeaders, long deadlineMs) { + SqlNode sqlNode = sqlNodeAndOptions.getSqlNode(); + boolean explain = sqlNode.getKind() == SqlKind.EXPLAIN; + SqlNode queryNode = explain ? ((SqlExplain) sqlNode).getExplicandum() : sqlNode; + Map queryOptions = sqlNodeAndOptions.getOptions(); + return IdSetSubqueryRewriter.rewrite(query, queryNode, explain, new IdSetSubqueryRewriter.SubqueryRunner() { + @Override + public boolean validates(String subquery) { + return validatesOnItsOwn(requestId, subquery, queryOptions, httpHeaders, deadlineMs); + } + + @Nullable + @Override + public String run(String subquery) { + return runIdSetSubquery(requestId, subquery, queryOptions, requesterIdentity, httpHeaders, deadlineMs); + } + }); + } + + /// Returns whether the subquery validates on its own, i.e. whether it does not refer to the query it is in. + /// Validation does not optimize the subquery, so it is quick enough to run without a timeout. + /// + /// Any failure, e.g. from an invalid option the subquery inherits, means no: the subquery then stays in the query, + /// and compiling the query reports the problem. + private boolean validatesOnItsOwn(long requestId, String subquery, Map queryOptions, + @Nullable HttpHeaders httpHeaders, long deadlineMs) { + getTimeLeftMs(subquery, deadlineMs); + try { + SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(subquery); + RequestUtils.setOptions(sqlNodeAndOptions, queryOptions); + QueryEnvironment queryEnv = new QueryEnvironment( + getQueryEnvConf(httpHeaders, sqlNodeAndOptions.getOptions(), requestId), _multiClusterRoutingContext); + queryEnv.validate(sqlNodeAndOptions); + return true; + } catch (RuntimeException e) { + return false; + } + } + + /// Runs an IdSet subquery as a query of its own, and returns the serialized IdSet it returns, or `null` if it returns + /// no rows or NULL. + /// + /// The subquery inherits the options of the query, which the options set in the subquery override, and the time left + /// until the deadline of the query. It runs like any other query: it is checked for access and quota, and it can be + /// cancelled. As it inherits the client query id, cancelling the query by that id cancels the subquery, which fails + /// the query. + @Nullable + private String runIdSetSubquery(long requestId, String subquery, Map queryOptions, + @Nullable RequesterIdentity requesterIdentity, @Nullable HttpHeaders httpHeaders, long deadlineMs) { + SqlNodeAndOptions sqlNodeAndOptions; + try { + sqlNodeAndOptions = RequestUtils.parseQuery(subquery); + } catch (Exception e) { + throw new QueryException(QueryErrorCode.fromThrowable(e, QueryErrorCode.SQL_PARSING), + "Failed to parse subquery: " + subquery + ": " + e.getMessage(), e); + } + if (sqlNodeAndOptions.getSqlType() != PinotSqlType.DQL + || !sqlNodeAndOptions.getSqlNode().getKind().belongsTo(SqlKind.QUERY)) { + throw QueryErrorCode.QUERY_VALIDATION.asException("Subquery is not a query: " + subquery); + } + RequestUtils.setOptions(sqlNodeAndOptions, queryOptions); + Map options = sqlNodeAndOptions.getOptions(); + long timeLeftMs = getTimeLeftMs(subquery, deadlineMs); + Long timeoutMs = QueryOptionsUtils.getTimeoutMs(options); + options.put(QueryOptionKey.TIMEOUT_MS, + Long.toString(timeoutMs != null ? Math.min(timeoutMs, timeLeftMs) : timeLeftMs)); + // The rows of the subquery are its result + options.remove(QueryOptionKey.DROP_RESULTS); + + // Like BaseBrokerRequestHandler.handleRequest, except for the first access check and the application quota, which + // the query passed + long subqueryRequestId = _requestIdGenerator.get(); + RequestContext requestContext = new DefaultRequestContext(); + requestContext.setRequestArrivalTimeMillis(System.currentTimeMillis()); + requestContext.setBrokerId(_brokerId); + requestContext.setRequestId(subqueryRequestId); + requestContext.setQuery(subquery); + setTrackedHeadersInRequestContext(requestContext, httpHeaders, _trackedHeaders); + LOGGER.info("Request {} runs a subquery as request {}", requestId, subqueryRequestId); + BrokerResponse response = handleRequest(subqueryRequestId, subquery, sqlNodeAndOptions, + JsonNodeFactory.instance.objectNode().put(CommonConstants.Broker.Request.SQL, subquery), requesterIdentity, + requestContext, httpHeaders, _accessControlFactory.create()); + response.setBrokerId(_brokerId); + response.setRequestId(Long.toString(subqueryRequestId)); + onQueryCompletion(requestContext, response); + + if (!response.getExceptions().isEmpty()) { + QueryProcessingException exception = response.getExceptions().get(0); + throw new QueryException(QueryErrorCode.fromErrorCode(exception.getErrorCode()), + "Subquery failed: " + subquery + ": " + exception.getMessage()); + } + // An IdSet built from part of the rows would make the query miss rows without an error + if (response.isPartialResult()) { + List reasons = getPartialResultReasons(response); + throw QueryErrorCode.QUERY_EXECUTION.asException( + "Subquery returned a partial result" + (reasons.isEmpty() ? "" : " " + reasons) + ": " + subquery); + } + ResultTable resultTable = response.getResultTable(); + DataSchema dataSchema = resultTable != null ? resultTable.getDataSchema() : null; + if (dataSchema == null || dataSchema.size() != 1 || dataSchema.getColumnDataType(0) != ColumnDataType.STRING) { + throw QueryErrorCode.QUERY_VALIDATION.asException( + "Subquery must return one STRING column, the serialized IdSet, got: " + dataSchema + " for: " + subquery); + } + List rows = resultTable.getRows(); + if (rows.size() > 1) { + throw QueryErrorCode.QUERY_EXECUTION.asException( + "Subquery must return at most one row, got: " + rows.size() + " for: " + subquery); + } + return rows.isEmpty() ? null : (String) rows.get(0)[0]; + } + + /// Returns the known reasons why the response holds a partial result, for the error message. + private static List getPartialResultReasons(BrokerResponse response) { + List reasons = new ArrayList<>(); + if (response.isNumGroupsLimitReached()) { + reasons.add("numGroupsLimitReached"); + } + if (response.isMaxRowsInJoinReached()) { + reasons.add("maxRowsInJoinReached"); + } + if (response.isMaxRowsInWindowReached()) { + reasons.add("maxRowsInWindowReached"); + } + if (response.isMseLiteLeafStageLimitReached()) { + reasons.add("mseLiteLeafStageLimitReached"); + } + if (response instanceof BrokerResponseNativeV2) { + reasons.addAll(((BrokerResponseNativeV2) response).getEarlyTerminationReasons()); + } + return reasons; + } + + private static long getTimeLeftMs(String subquery, long deadlineMs) { + long timeLeftMs = deadlineMs - System.currentTimeMillis(); + if (timeLeftMs <= 0) { + throw QueryErrorCode.BROKER_TIMEOUT.asException("Timed out before running subquery: " + subquery); + } + return timeLeftMs; + } + private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestId, RequesterIdentity requesterIdentity, RequestContext requestContext, HttpHeaders httpHeaders, Timer timer, boolean rlsFiltersApplied, boolean queryWasLogged, String workloadName) diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/IdSetSubqueryRewriterTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/IdSetSubqueryRewriterTest.java new file mode 100644 index 000000000000..5ffc16e2291d --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/IdSetSubqueryRewriterTest.java @@ -0,0 +1,309 @@ +/** + * 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.pinot.broker.requesthandler; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.sql.SqlExplain; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.util.Litmus; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +/// Tests [IdSetSubqueryRewriter] on parsed queries, with a fake runner that returns a given IdSet per subquery. +public class IdSetSubqueryRewriterTest { + private static final String EMPTY = IdSetSubqueryRewriter.EMPTY_ID_SET; + + @Test + public void testRewritesInSubqueryInFilters() { + FakeRunner runner = new FakeRunner().returning("SELECT IDSET(col1) FROM b", "set1") + .returning("SELECT IDSET(col3) FROM b WHERE col2 = 'x'", "set2"); + List subqueries = assertRewrite( + "SELECT col1 FROM a WHERE IN_SUBQUERY(col1, 'SELECT IDSET(col1) FROM b') = 1 " + + "AND NOT IN_SUBQUERY(col3, 'SELECT IDSET(col3) FROM b WHERE col2 = ''x''')", + "SELECT col1 FROM a WHERE IN_ID_SET(col1, 'set1') = 1 AND NOT IN_ID_SET(col3, 'set2')", false, runner); + assertEquals(subqueries, List.of("SELECT IDSET(col1) FROM b", "SELECT IDSET(col3) FROM b WHERE col2 = 'x'")); + assertEquals(runner._ran, subqueries); + assertEquals(runner._validated, List.of()); + } + + @Test + public void testRewritesInSubqueryInEveryClause() { + FakeRunner runner = new FakeRunner().returning("q1", "set1").returning("q2", "set2").returning("q3", "set3") + .returning("q4", "set4").returning("q5", "set5"); + assertRewrite( + "SELECT IN_SUBQUERY(a.col1, 'q1'), COUNT(*) FROM a JOIN b ON a.col2 = b.col2 AND IN_SUBQUERY(b.col3, 'q2') = 1 " + + "GROUP BY 1 HAVING IN_SUBQUERY(MAX(a.col3), 'q3') = 1 " + + "ORDER BY CASE WHEN IN_SUBQUERY(MIN(b.col1), 'q4') = 1 THEN 1 ELSE IN_SUBQUERY(1, 'q5') END", + "SELECT IN_ID_SET(a.col1, 'set1'), COUNT(*) FROM a JOIN b ON a.col2 = b.col2 AND IN_ID_SET(b.col3, 'set2') = 1 " + + "GROUP BY 1 HAVING IN_ID_SET(MAX(a.col3), 'set3') = 1 " + + "ORDER BY CASE WHEN IN_ID_SET(MIN(b.col1), 'set4') = 1 THEN 1 ELSE IN_ID_SET(1, 'set5') END", + false, runner); + assertEquals(runner._ran.size(), 5); + } + + @Test + public void testRewritesInSubqueryInNestedQueries() { + FakeRunner runner = new FakeRunner().returning("q1", "set1").returning("q2", "set2").returning("q3", "set3") + .returning("q4", "set4"); + assertRewrite( + "WITH w AS (SELECT col1 FROM b WHERE IN_SUBQUERY(col1, 'q1') = 1) " + + "SELECT col1 FROM (SELECT col1 FROM a WHERE IN_SUBQUERY(col1, 'q2') = 1) " + + "WHERE col1 IN (SELECT col1 FROM w) " + + "AND col1 NOT IN (SELECT col1 FROM c WHERE IN_SUBQUERY(col1, 'q3') = 0) " + + "UNION ALL SELECT col1 FROM d WHERE IN_SUBQUERY(col1, 'q4') = 1", + "WITH w AS (SELECT col1 FROM b WHERE IN_ID_SET(col1, 'set1') = 1) " + + "SELECT col1 FROM (SELECT col1 FROM a WHERE IN_ID_SET(col1, 'set2') = 1) " + + "WHERE col1 IN (SELECT col1 FROM w) " + + "AND col1 NOT IN (SELECT col1 FROM c WHERE IN_ID_SET(col1, 'set3') = 0) " + + "UNION ALL SELECT col1 FROM d WHERE IN_ID_SET(col1, 'set4') = 1", + false, runner); + assertEquals(runner._ran.size(), 4); + } + + @Test + public void testMatchesEveryNameOfInSubquery() { + FakeRunner runner = new FakeRunner().returning("q1", "set1").returning("q2", "set2").returning("q3", "set3"); + assertRewrite("SELECT col1 FROM a WHERE INSUBQUERY(col1, 'q1') = 1 AND inSubquery(col2, 'q2') = 1 " + + "AND in_sub_query(col3, 'q3') = 1", + "SELECT col1 FROM a WHERE IN_ID_SET(col1, 'set1') = 1 AND IN_ID_SET(col2, 'set2') = 1 " + + "AND IN_ID_SET(col3, 'set3') = 1", false, runner); + } + + @Test + public void testUsesTheEmptyIdSetForNoResult() { + FakeRunner runner = new FakeRunner().returning("q1", null); + assertRewrite("SELECT col1 FROM a WHERE IN_SUBQUERY(col1, 'q1') = 0", + "SELECT col1 FROM a WHERE IN_ID_SET(col1, '" + EMPTY + "') = 0", false, runner); + } + + @Test + public void testRunsScalarSubqueries() { + String subquery1 = "SELECT IDSET(col1) FROM b WHERE col3 > 1"; + String subquery2 = "SELECT IDSET(col2, 'expectedInsertions=10') FROM (SELECT col2 FROM c UNION SELECT col2 FROM d)"; + String subquery3 = "WITH w AS (SELECT col3 FROM b) SELECT IDSET(col3) FROM w"; + String subquery4 = "SELECT IDSET(col1) FROM b\n -- a comment\n GROUP BY col2\n ORDER BY col2 LIMIT 1"; + // The text of a scalar subquery includes its parentheses + List expected = + List.of("(" + subquery1 + ")", "(" + subquery2 + ")", "(" + subquery3 + ")", "(\n " + subquery4 + "\n)"); + FakeRunner runner = new FakeRunner().returning(expected.get(0), "set1").returning(expected.get(1), "set2") + .returning(expected.get(2), "set3").returning(expected.get(3), "set4"); + List subqueries = assertRewrite( + "SET timeoutMs = 1000;\nSELECT col1 FROM a WHERE IN_ID_SET(col1, " + expected.get(0) + ") = 1\n" + + "AND IN_ID_SET(col2, " + expected.get(1) + ") = 1 AND inIdSet(col3, " + expected.get(2) + ") = 1\n" + + "AND IN_ID_SET(col1, " + expected.get(3) + ") = 0", + "SELECT col1 FROM a WHERE IN_ID_SET(col1, 'set1') = 1 AND IN_ID_SET(col2, 'set2') = 1 " + + "AND inIdSet(col3, 'set3') = 1 AND IN_ID_SET(col1, 'set4') = 0", false, runner); + assertEquals(subqueries, expected); + assertEquals(runner._validated, expected); + assertEquals(runner._ran, expected); + } + + @Test + public void testLeavesCorrelatedScalarSubqueries() { + String correlated = "(SELECT IDSET(b.col1) FROM b WHERE b.col2 = a.col2 AND IN_SUBQUERY(b.col3, 'q1') = 1)"; + FakeRunner runner = new FakeRunner().correlated(correlated).returning("q1", "set1"); + List subqueries = assertRewrite("SELECT col1 FROM a WHERE IN_ID_SET(a.col1, " + correlated + ") = 1", + "SELECT col1 FROM a WHERE IN_ID_SET(a.col1, (SELECT IDSET(b.col1) FROM b WHERE b.col2 = a.col2 " + + "AND IN_ID_SET(b.col3, 'set1') = 1)) = 1", false, runner); + // The IN_SUBQUERY in the correlated subquery still runs first + assertEquals(subqueries, List.of("q1")); + assertEquals(runner._validated, List.of(correlated)); + assertEquals(runner._ran, List.of("q1")); + } + + @Test + public void testLeavesScalarSubqueriesWhoseTextDiffers() { + // The parser read another text than the given one, e.g. after the broker sanitized it + String sql = "SELECT col1 FROM a WHERE IN_ID_SET(col1, (SELECT IDSET(col1) FROM b)) = 1"; + SqlNode query = parse(sql); + FakeRunner runner = new FakeRunner(); + List subqueries = + IdSetSubqueryRewriter.rewrite(sql.replace("FROM b", "FROM c"), query, false, runner); + assertEquals(subqueries, List.of()); + assertEquals(runner._validated, List.of()); + assertTrue(query.equalsDeep(parse(sql), Litmus.THROW)); + } + + @Test + public void testExplainRunsNoSubqueries() { + String scalarSubquery = "(SELECT IDSET(col1) FROM b)"; + String correlated = "(SELECT IDSET(b.col1) FROM b WHERE b.col2 = a.col2)"; + FakeRunner runner = new FakeRunner().correlated(correlated); + List subqueries = assertRewrite( + "EXPLAIN PLAN FOR SELECT col1 FROM a WHERE IN_SUBQUERY(col1, 'q1') = 1 AND IN_ID_SET(col2, " + + scalarSubquery + ") = 1 AND IN_ID_SET(col3, " + correlated + ") = 1", + "SELECT col1 FROM a WHERE IN_ID_SET(col1, '" + EMPTY + "') = 1 AND IN_ID_SET(col2, '" + EMPTY + "') = 1 " + + "AND IN_ID_SET(col3, " + correlated + ") = 1", true, runner); + assertEquals(subqueries, List.of("q1", scalarSubquery)); + assertEquals(runner._validated, List.of(scalarSubquery, correlated)); + assertEquals(runner._ran, List.of()); + } + + @Test + public void testLeavesQueriesWithoutIdSetSubqueries() { + FakeRunner runner = new FakeRunner(); + String sql = "SELECT col1 FROM a WHERE IN_ID_SET(col1, 'AwAAAA==') = 1 AND col2 IN (SELECT col2 FROM b) " + + "AND col3 = (SELECT MAX(col3) FROM c)"; + assertEquals(assertRewrite(sql, sql, false, runner), List.of()); + assertEquals(runner._validated, List.of()); + assertEquals(runner._ran, List.of()); + } + + @Test + public void testRejectsInvalidInSubqueryCalls() { + for (String sql : List.of("SELECT col1 FROM a WHERE IN_SUBQUERY(col1) = 1", + "SELECT col1 FROM a WHERE IN_SUBQUERY(col1, 'q1', 'q2') = 1", + "SELECT col1 FROM a WHERE IN_SUBQUERY(col1, col2) = 1", + "SELECT col1 FROM a WHERE IN_SUBQUERY(col1, (SELECT 'q1')) = 1")) { + QueryException exception = expectThrows(QueryException.class, + () -> IdSetSubqueryRewriter.rewrite(sql, parse(sql), false, new FakeRunner())); + assertEquals(exception.getErrorCode(), QueryErrorCode.QUERY_VALIDATION); + assertTrue(exception.getMessage().contains("IN_SUBQUERY takes an expression and a string literal"), + exception.getMessage()); + } + } + + @Test + public void testRunsEachSubqueryOnce() { + String scalarSubquery = "(SELECT IDSET(col1) FROM b)"; + FakeRunner runner = new FakeRunner().returning("q1", "set1").returning(scalarSubquery, "set2"); + List subqueries = assertRewrite( + "SELECT col1 FROM a WHERE IN_SUBQUERY(col1, 'q1') = 1 OR IN_SUBQUERY(col2, 'q1') = 1 " + + "OR IN_ID_SET(col1, " + scalarSubquery + ") = 1 OR IN_ID_SET(col2, " + scalarSubquery + ") = 1", + "SELECT col1 FROM a WHERE IN_ID_SET(col1, 'set1') = 1 OR IN_ID_SET(col2, 'set1') = 1 " + + "OR IN_ID_SET(col1, 'set2') = 1 OR IN_ID_SET(col2, 'set2') = 1", false, runner); + assertEquals(subqueries, List.of("q1", scalarSubquery)); + assertEquals(runner._ran, List.of("q1", scalarSubquery)); + assertEquals(runner._validated, List.of(scalarSubquery)); + } + + @Test + public void testLeavesScalarSubqueriesThatNameWithItems() { + // On its own, a subquery that names the WITH item b would read table b + String namesWithItem = "(SELECT IDSET(col1) FROM b)"; + String other = "(SELECT IDSET(col3) FROM c)"; + FakeRunner runner = new FakeRunner().returning(other, "set1"); + assertRewrite("WITH b AS (SELECT col1 FROM d) SELECT col1 FROM a WHERE IN_ID_SET(col1, " + namesWithItem + + ") = 1 AND IN_ID_SET(col3, " + other + ") = 1", + "WITH b AS (SELECT col1 FROM d) SELECT col1 FROM a WHERE IN_ID_SET(col1, " + namesWithItem + + ") = 1 AND IN_ID_SET(col3, 'set1') = 1", false, runner); + assertEquals(runner._validated, List.of(other)); + + // The WITH items stay in scope in a subquery that stays, so the subqueries in it stay too + String correlated = "(SELECT IDSET(c.col1) FROM c WHERE c.col2 = a.col2 AND IN_ID_SET(c.col3, " + namesWithItem + + ") = 1)"; + runner = new FakeRunner(); + String sql = "WITH b AS (SELECT col1 FROM d) SELECT col1 FROM a WHERE IN_ID_SET(a.col1, " + correlated + ") = 1"; + assertRewrite(sql, sql, false, runner); + assertEquals(runner._validated, List.of()); + assertEquals(runner._ran, List.of()); + + // The ORDER BY of a WITH query parses outside the WITH, but the WITH items are in its scope too + runner = new FakeRunner(); + sql = "WITH b AS (SELECT col1 FROM d) SELECT col1 FROM a ORDER BY CASE WHEN IN_ID_SET(col1, " + namesWithItem + + ") = 1 THEN 0 ELSE 1 END LIMIT 10"; + assertRewrite(sql, sql, false, runner); + assertEquals(runner._validated, List.of()); + + // But not outside the query that has the WITH + runner = new FakeRunner().returning(namesWithItem, "set1"); + assertRewrite("SELECT col1 FROM (WITH b AS (SELECT col1 FROM d) SELECT col1 FROM b) WHERE IN_ID_SET(col1, " + + namesWithItem + ") = 1", + "SELECT col1 FROM (WITH b AS (SELECT col1 FROM d) SELECT col1 FROM b) WHERE IN_ID_SET(col1, 'set1') = 1", + false, runner); + } + + @Test + public void testRewritesIdSetSubqueriesInTheTestedExpression() { + FakeRunner runner = new FakeRunner().returning("q1", "set1").returning("q2", "set2"); + assertRewrite( + "SELECT col1 FROM a WHERE IN_SUBQUERY(CASE WHEN IN_SUBQUERY(col2, 'q2') = 1 THEN col1 END, 'q1') = 1", + "SELECT col1 FROM a WHERE IN_ID_SET(CASE WHEN IN_ID_SET(col2, 'set2') = 1 THEN col1 END, 'set1') = 1", + false, runner); + assertEquals(runner._ran, List.of("q1", "q2")); + } + + @Test + public void testReadsSubqueriesSplitAcrossLiterals() { + FakeRunner runner = new FakeRunner().returning("SELECT IDSET(col1) FROM b", "set1"); + assertRewrite("SELECT col1 FROM a WHERE IN_SUBQUERY(col1, 'SELECT IDSET(col1) '\n'FROM b') = 1", + "SELECT col1 FROM a WHERE IN_ID_SET(col1, 'set1') = 1", false, runner); + } + + /// Rewrites the query, checks that it becomes the expected query, and returns the replaced subqueries. + private static List assertRewrite(String sql, String expectedSql, boolean explain, FakeRunner runner) { + SqlNode query = parse(sql); + assertEquals(query.getKind() == SqlKind.EXPLAIN, explain); + if (explain) { + query = ((SqlExplain) query).getExplicandum(); + } + List subqueries = IdSetSubqueryRewriter.rewrite(sql, query, explain, runner); + assertTrue(query.equalsDeep(parse(expectedSql), Litmus.THROW)); + return subqueries; + } + + private static SqlNode parse(String sql) { + return CalciteSqlParser.compileToSqlNodeAndOptions(sql).getSqlNode(); + } + + private static class FakeRunner implements IdSetSubqueryRewriter.SubqueryRunner { + final Map _results = new HashMap<>(); + final Set _correlated = new HashSet<>(); + final List _validated = new ArrayList<>(); + final List _ran = new ArrayList<>(); + + FakeRunner returning(String subquery, @Nullable String idSet) { + _results.put(subquery, idSet); + return this; + } + + FakeRunner correlated(String subquery) { + _correlated.add(subquery); + return this; + } + + @Override + public boolean validates(String subquery) { + _validated.add(subquery); + return !_correlated.contains(subquery); + } + + @Nullable + @Override + public String run(String subquery) { + _ran.add(subquery); + assertTrue(_results.containsKey(subquery), "Unexpected subquery: " + subquery); + return _results.get(subquery); + } + } +} diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageIdSetSubqueryTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageIdSetSubqueryTest.java new file mode 100644 index 000000000000..54d7eb8471d8 --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageIdSetSubqueryTest.java @@ -0,0 +1,491 @@ +/** + * 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.pinot.broker.requesthandler; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import javax.ws.rs.WebApplicationException; +import javax.ws.rs.core.HttpHeaders; +import org.apache.pinot.broker.api.AccessControl; +import org.apache.pinot.broker.broker.AccessControlFactory; +import org.apache.pinot.broker.broker.AllowAllAccessControlFactory; +import org.apache.pinot.broker.queryquota.QueryQuotaManager; +import org.apache.pinot.common.failuredetector.FailureDetector; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.response.BrokerResponse; +import org.apache.pinot.common.response.broker.BrokerResponseNativeV2; +import org.apache.pinot.common.response.broker.QueryProcessingException; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.request.RequestUtils; +import org.apache.pinot.core.query.utils.idset.IdSet; +import org.apache.pinot.core.query.utils.idset.IdSets; +import org.apache.pinot.core.routing.MockRoutingManagerFactory; +import org.apache.pinot.core.routing.RoutingManager; +import org.apache.pinot.query.QueryEnvironmentTestBase; +import org.apache.pinot.query.routing.WorkerManager; +import org.apache.pinot.spi.accounting.ThreadAccountantUtils; +import org.apache.pinot.spi.auth.TableAuthorizationResult; +import org.apache.pinot.spi.auth.broker.RequesterIdentity; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; +import org.apache.pinot.spi.trace.DefaultRequestContext; +import org.apache.pinot.spi.trace.RequestContext; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; +import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner; +import org.apache.pinot.spi.utils.NetUtils; +import org.apache.pinot.sql.parsers.SqlNodeAndOptions; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +/// Tests how [MultiStageBrokerRequestHandler] runs the IdSet subqueries of a query, see [IdSetSubqueryRewriter]. +/// +/// Queries plan against the tables of [QueryEnvironmentTestBase]. The handler answers the subqueries it is given +/// responses for, and no server runs, so a query that gets past planning stops at the mocked throttler. +public class MultiStageIdSetSubqueryTest { + private static final long REQUEST_ID = 1L; + private static final String SUBQUERY = "SELECT IDSET(col1) FROM b"; + // The text of a scalar subquery includes its parentheses + private static final String SCALAR_SUBQUERY = "(" + SUBQUERY + ")"; + private static final String ID_SET; + + static { + IdSet idSet = IdSets.create(DataType.STRING); + idSet.add("foo"); + try { + ID_SET = idSet.toBase64String(); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + private final RequesterIdentity _requesterIdentity = mock(RequesterIdentity.class); + private final HttpHeaders _httpHeaders = mock(HttpHeaders.class); + // Responses to the subqueries, and errors they throw, by subquery + private final Map _subqueryResponses = new HashMap<>(); + private final Map _subqueryErrors = new HashMap<>(); + // All the requests the handler runs, including the queries of the tests + private final List _requests = new ArrayList<>(); + private final List _completedRequests = new ArrayList<>(); + + @BeforeMethod + public void setUp() { + _subqueryResponses.clear(); + _subqueryErrors.clear(); + _requests.clear(); + _completedRequests.clear(); + } + + @Test + public void testSubqueryInheritsTheOptionsAndTheTimeLeft() + throws Exception { + String subquery = "SET numGroupsLimit = 8; " + SUBQUERY; + _subqueryResponses.put(subquery, idSetResponse(ID_SET)); + String sql = "SET timeoutMs = 60000; SET clientQueryId = 'myQuery'; SET dropResults = true; " + + "SET numGroupsLimit = 7; SET maxRowsInJoin = 9; " + + "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, '" + subquery + "') = 1"; + SqlNodeAndOptions sqlNodeAndOptions = RequestUtils.parseQuery(sql); + long startTimeMs = System.currentTimeMillis(); + handleRequest(newHandler(new AllowAllAccessControlFactory()), sql, sqlNodeAndOptions); + + Request request = getSubqueryRequest(); + assertNotEquals(request._requestId, REQUEST_ID); + assertEquals(request._query, subquery); + assertEquals(request._requestContext.getRequestId(), request._requestId); + assertEquals(request._requestContext.getQuery(), subquery); + assertTrue(request._requestContext.getRequestArrivalTimeMillis() >= startTimeMs); + assertSame(request._requesterIdentity, _requesterIdentity); + assertSame(request._httpHeaders, _httpHeaders); + Map options = request._options; + // The options set in the subquery override those of the query + assertEquals(options.get(QueryOptionKey.NUM_GROUPS_LIMIT), "8"); + assertEquals(options.get(QueryOptionKey.MAX_ROWS_IN_JOIN), "9"); + // The client query id lets the client cancel the subquery, which fails the query + assertEquals(options.get(QueryOptionKey.CLIENT_QUERY_ID), "myQuery"); + assertFalse(options.containsKey(QueryOptionKey.DROP_RESULTS)); + long timeoutMs = Long.parseLong(options.get(QueryOptionKey.TIMEOUT_MS)); + assertTrue(timeoutMs > 0 && timeoutMs <= 60_000, "Timeout: " + timeoutMs); + + // The query uses the IdSet, and its options do not change + String rewrittenQuery = sqlNodeAndOptions.getSqlNode().toString(); + assertTrue(rewrittenQuery.contains(ID_SET), rewrittenQuery); + assertFalse(rewrittenQuery.toUpperCase().contains("IN_SUBQUERY"), rewrittenQuery); + assertEquals(sqlNodeAndOptions.getOptions().get(QueryOptionKey.NUM_GROUPS_LIMIT), "7"); + assertEquals(sqlNodeAndOptions.getOptions().get(QueryOptionKey.DROP_RESULTS), "true"); + + // The subquery is reported to the query event listener like any other query + assertEquals(_completedRequests.size(), 1); + assertSame(_completedRequests.get(0), request._requestContext); + } + + @Test + public void testSubqueryGetsTheTimeLeftOfTheQuery() + throws Exception { + _subqueryResponses.put(SUBQUERY, idSetResponse(ID_SET)); + String sql = "SET timeoutMs = 10000; SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, '" + SUBQUERY + "') = 1"; + MultiStageBrokerRequestHandler handler = newHandler(new AllowAllAccessControlFactory()); + + // The query arrived 4 seconds ago, so the subquery gets at most the 6 seconds left + RequestContext requestContext = new DefaultRequestContext(); + requestContext.setRequestArrivalTimeMillis(System.currentTimeMillis() - 4_000); + handler.handleRequestThrowing(REQUEST_ID, sql, RequestUtils.parseQuery(sql), null, requestContext, null); + long timeoutMs = Long.parseLong(getSubqueryRequest()._options.get(QueryOptionKey.TIMEOUT_MS)); + assertTrue(timeoutMs > 0 && timeoutMs <= 6_000, "Timeout: " + timeoutMs); + + // A shorter timeout of the subquery wins + _requests.clear(); + String subquery = "SET timeoutMs = 100; " + SUBQUERY; + _subqueryResponses.put(subquery, idSetResponse(ID_SET)); + String shorterTimeoutSql = "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, '" + subquery + "') = 1"; + handleRequest(handler, shorterTimeoutSql, RequestUtils.parseQuery(shorterTimeoutSql)); + assertEquals(getSubqueryRequest()._options.get(QueryOptionKey.TIMEOUT_MS), "100"); + + // No time left + _requests.clear(); + requestContext.setRequestArrivalTimeMillis(System.currentTimeMillis() - 10_000); + QueryException exception = expectThrows(QueryException.class, + () -> handler.handleRequestThrowing(REQUEST_ID, sql, RequestUtils.parseQuery(sql), null, requestContext, + null)); + assertEquals(exception.getErrorCode(), QueryErrorCode.BROKER_TIMEOUT); + assertNoSubqueryRan(); + } + + @Test + public void testUsesTheEmptyIdSetForNoRowsOrNull() + throws Exception { + MultiStageBrokerRequestHandler handler = newHandler(new AllowAllAccessControlFactory()); + for (BrokerResponse response : List.of(response(stringColumn(), List.of()), + response(stringColumn(), List.of(new Object[]{null})))) { + _subqueryResponses.put(SUBQUERY, response); + String sql = "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, '" + SUBQUERY + "') = 0"; + SqlNodeAndOptions sqlNodeAndOptions = RequestUtils.parseQuery(sql); + handleRequest(handler, sql, sqlNodeAndOptions); + String rewrittenQuery = sqlNodeAndOptions.getSqlNode().toString(); + assertTrue(rewrittenQuery.contains("'" + IdSetSubqueryRewriter.EMPTY_ID_SET + "'"), rewrittenQuery); + } + } + + @Test + public void testFailsWhenTheSubqueryFails() + throws Exception { + MultiStageBrokerRequestHandler handler = newHandler(new AllowAllAccessControlFactory()); + String sql = "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, '" + SUBQUERY + "') = 1"; + + BrokerResponseNativeV2 failedResponse = new BrokerResponseNativeV2(); + failedResponse.addException(new QueryProcessingException(QueryErrorCode.EXECUTION_TIMEOUT, "Too slow")); + _subqueryResponses.put(SUBQUERY, failedResponse); + assertFails(handler, sql, QueryErrorCode.EXECUTION_TIMEOUT, "Subquery failed: " + SUBQUERY + ": Too slow"); + + _subqueryResponses.clear(); + _subqueryErrors.put(SUBQUERY, QueryErrorCode.QUERY_VALIDATION.asException("Unknown column")); + String message = assertFails(handler, sql, QueryErrorCode.QUERY_VALIDATION, "Subquery failed: " + SUBQUERY); + assertTrue(message.contains("Unknown column"), message); + _subqueryErrors.clear(); + + // An IdSet built from part of the rows would make the query miss rows + BrokerResponseNativeV2 partialResponse = (BrokerResponseNativeV2) idSetResponse(ID_SET); + partialResponse.mergeNumGroupsLimitReached(true); + _subqueryResponses.put(SUBQUERY, partialResponse); + assertFails(handler, sql, QueryErrorCode.QUERY_EXECUTION, + "Subquery returned a partial result [numGroupsLimitReached]: " + SUBQUERY); + partialResponse = (BrokerResponseNativeV2) idSetResponse(ID_SET); + partialResponse.mergeMaxRowsInWindowReached(true); + _subqueryResponses.put(SUBQUERY, partialResponse); + assertFails(handler, sql, QueryErrorCode.QUERY_EXECUTION, + "Subquery returned a partial result [maxRowsInWindowReached]: " + SUBQUERY); + // A partial result that no known flag explains, e.g. a flag added later + BrokerResponseNativeV2 unexplainedPartialResponse = new BrokerResponseNativeV2() { + @Override + public boolean isPartialResult() { + return true; + } + }; + unexplainedPartialResponse.setResultTable(new ResultTable(stringColumn(), List.of(new Object[]{ID_SET}))); + _subqueryResponses.put(SUBQUERY, unexplainedPartialResponse); + assertFails(handler, sql, QueryErrorCode.QUERY_EXECUTION, "Subquery returned a partial result: " + SUBQUERY); + + _subqueryResponses.put(SUBQUERY, response(new DataSchema(new String[]{"a", "b"}, + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING}), List.of())); + assertFails(handler, sql, QueryErrorCode.QUERY_VALIDATION, "Subquery must return one STRING column"); + + _subqueryResponses.put(SUBQUERY, + response(new DataSchema(new String[]{"a"}, new ColumnDataType[]{ColumnDataType.LONG}), List.of())); + assertFails(handler, sql, QueryErrorCode.QUERY_VALIDATION, "Subquery must return one STRING column"); + + _subqueryResponses.put(SUBQUERY, + response(stringColumn(), List.of(new Object[]{ID_SET}, new Object[]{ID_SET}))); + assertFails(handler, sql, QueryErrorCode.QUERY_EXECUTION, "Subquery must return at most one row, got: 2"); + } + + @Test + public void testFailsOnSubqueriesThatAreNotQueriesOrSetRejectedOptions() + throws Exception { + MultiStageBrokerRequestHandler handler = newHandler(new AllowAllAccessControlFactory()); + assertFails(handler, "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, 'EXPLAIN PLAN FOR " + SUBQUERY + "') = 1", + QueryErrorCode.QUERY_VALIDATION, "Subquery is not a query"); + assertFails(handler, "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, 'SELECT FROM') = 1", + QueryErrorCode.SQL_PARSING, "Failed to parse subquery"); + + // The request does not allow options in the SQL, so the subquery cannot set any either + String sql = "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, 'SET numGroupsLimit = 8; " + SUBQUERY + "') = 1"; + SqlNodeAndOptions sqlNodeAndOptions = RequestUtils.parseQuery(sql); + sqlNodeAndOptions.getOptions().put(QueryOptionKey.SQL_OPTIONS_MODE, "REJECT"); + QueryException exception = expectThrows(QueryException.class, + () -> handleRequest(handler, sql, sqlNodeAndOptions)); + assertEquals(exception.getErrorCode(), QueryErrorCode.QUERY_VALIDATION); + assertTrue(exception.getMessage().contains("Query options are not allowed in the SQL"), exception.getMessage()); + assertNoSubqueryRan(); + } + + @Test + public void testChecksAccessToTheTablesOfTheSubquery() + throws Exception { + // Denies table b, which only the subquery reads. The test does not answer the subquery, so it compiles and gets + // checked like any other query. + AccessControl denyB = new AccessControl() { + @Override + public TableAuthorizationResult authorize(RequesterIdentity requesterIdentity, Set tables) { + Set failedTables = tables.stream().filter(table -> table.startsWith("b")).collect(Collectors.toSet()); + return failedTables.isEmpty() ? TableAuthorizationResult.success() : new TableAuthorizationResult(failedTables); + } + }; + MultiStageBrokerRequestHandler handler = newHandler(new AccessControlFactory() { + @Override + public AccessControl create() { + return denyB; + } + }); + String sql = "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, '" + SUBQUERY + "') = 1"; + WebApplicationException exception = expectThrows(WebApplicationException.class, + () -> handleRequest(handler, sql, RequestUtils.parseQuery(sql))); + assertEquals(exception.getResponse().getStatus(), 403); + } + + @Test + public void testExplainRunsNoSubqueries() + throws Exception { + MultiStageBrokerRequestHandler handler = newHandler(new AllowAllAccessControlFactory()); + _subqueryResponses.put(SUBQUERY, idSetResponse(ID_SET)); + _subqueryResponses.put(SCALAR_SUBQUERY, idSetResponse(ID_SET)); + String sql = "EXPLAIN PLAN FOR SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, '" + SUBQUERY + "') = 1 " + + "AND IN_ID_SET(col2, " + SCALAR_SUBQUERY + ") = 1"; + BrokerResponse response = handleRequest(handler, sql, RequestUtils.parseQuery(sql)); + assertNoSubqueryRan(); + assertTrue(response.getExceptions().isEmpty(), response.getExceptions().toString()); + ResultTable resultTable = response.getResultTable(); + List columnNames = List.of(resultTable.getDataSchema().getColumnNames()); + assertEquals(columnNames.get(columnNames.size() - 1), "ID_SET_SUBQUERIES_NOT_RUN"); + Object[] row = resultTable.getRows().get(0); + String plan = (String) row[1]; + assertTrue(plan.contains("'" + IdSetSubqueryRewriter.EMPTY_ID_SET + "'"), plan); + assertFalse(plan.contains("Join"), plan); + assertEquals(row[columnNames.size() - 1], "[\"" + SUBQUERY + "\",\"" + SCALAR_SUBQUERY + "\"]"); + + // Without IdSet subqueries, the response has no extra column + String plainSql = "EXPLAIN PLAN FOR SELECT COUNT(*) FROM a"; + response = handleRequest(handler, plainSql, RequestUtils.parseQuery(plainSql)); + assertFalse(List.of(response.getResultTable().getDataSchema().getColumnNames()) + .contains("ID_SET_SUBQUERIES_NOT_RUN")); + } + + @Test + public void testRunsUncorrelatedScalarSubqueriesFirst() + throws Exception { + MultiStageBrokerRequestHandler handler = newHandler(new AllowAllAccessControlFactory()); + _subqueryResponses.put(SCALAR_SUBQUERY, idSetResponse(ID_SET)); + String sql = "SELECT COUNT(*) FROM a WHERE IN_ID_SET(col1, " + SCALAR_SUBQUERY + ") = 1"; + SqlNodeAndOptions sqlNodeAndOptions = RequestUtils.parseQuery(sql); + handleRequest(handler, sql, sqlNodeAndOptions); + assertEquals(getSubqueryRequest()._query, SCALAR_SUBQUERY); + assertTrue(sqlNodeAndOptions.getSqlNode().toString().contains(ID_SET)); + + // A correlated subquery does not validate on its own, so it runs as part of the query + _requests.clear(); + String correlatedSubquery = "(SELECT IDSET(b.col1) FROM b WHERE b.col2 = a.col2)"; + _subqueryResponses.put(correlatedSubquery, idSetResponse(ID_SET)); + String correlatedSql = "SELECT COUNT(*) FROM a WHERE IN_ID_SET(a.col1, " + correlatedSubquery + ") = 1"; + handleRequest(handler, correlatedSql, RequestUtils.parseQuery(correlatedSql)); + assertNoSubqueryRan(); + + // An option that fails the subquery's validation also fails the query, which compiles the subquery in its plan + _requests.clear(); + String invalidOptionSql = "SET inferPartitionHint = 'foo'; " + sql; + QueryException exception = expectThrows(QueryException.class, + () -> handleRequest(handler, invalidOptionSql, RequestUtils.parseQuery(invalidOptionSql))); + assertTrue(exception.getMessage().contains("inferPartitionHint"), exception.getMessage()); + assertNoSubqueryRan(); + + // On its own, the subquery would read table b instead of the WITH item b + _requests.clear(); + String withSql = "WITH b AS (SELECT col1 FROM c WHERE col2 = 'foo') " + + "SELECT COUNT(*) FROM a WHERE IN_ID_SET(col1, " + SCALAR_SUBQUERY + ") = 1"; + handleRequest(handler, withSql, RequestUtils.parseQuery(withSql)); + assertNoSubqueryRan(); + } + + /// Checks that the query fails with the error code and a message that contains the given one. Returns the message. + private String assertFails(MultiStageBrokerRequestHandler handler, String sql, QueryErrorCode errorCode, + String message) { + QueryException exception = + expectThrows(QueryException.class, () -> handleRequest(handler, sql, RequestUtils.parseQuery(sql))); + assertEquals(exception.getErrorCode(), errorCode); + assertTrue(exception.getMessage().contains(message), exception.getMessage()); + return exception.getMessage(); + } + + private BrokerResponse handleRequest(MultiStageBrokerRequestHandler handler, String sql, + SqlNodeAndOptions sqlNodeAndOptions) { + RequestContext requestContext = new DefaultRequestContext(); + requestContext.setRequestArrivalTimeMillis(System.currentTimeMillis()); + return handler.handleRequestThrowing(REQUEST_ID, sql, sqlNodeAndOptions, _requesterIdentity, requestContext, + _httpHeaders); + } + + /// Returns the only subquery request, i.e. the only request other than the query of the test. + private Request getSubqueryRequest() { + List subqueryRequests = getSubqueryRequests(); + assertEquals(subqueryRequests.size(), 1, subqueryRequests.toString()); + return subqueryRequests.get(0); + } + + private void assertNoSubqueryRan() { + assertEquals(getSubqueryRequests(), List.of()); + } + + private List getSubqueryRequests() { + return _requests.stream().filter(request -> request._requestId != REQUEST_ID).collect(Collectors.toList()); + } + + private static BrokerResponse idSetResponse(String idSet) { + return response(stringColumn(), List.of(new Object[]{idSet})); + } + + private static DataSchema stringColumn() { + return new DataSchema(new String[]{"idset(col1)"}, new ColumnDataType[]{ColumnDataType.STRING}); + } + + private static BrokerResponse response(DataSchema dataSchema, List rows) { + BrokerResponseNativeV2 response = new BrokerResponseNativeV2(); + response.setResultTable(new ResultTable(dataSchema, rows)); + return response; + } + + private MultiStageBrokerRequestHandler newHandler(AccessControlFactory accessControlFactory) + throws IOException { + int port1 = NetUtils.findOpenPort(); + int port2 = NetUtils.findOpenPort(); + MockRoutingManagerFactory factory = new MockRoutingManagerFactory(port1, port2); + QueryEnvironmentTestBase.TABLE_SCHEMAS.forEach((tableName, schema) -> factory.registerTable(schema, tableName)); + QueryEnvironmentTestBase.SERVER1_SEGMENTS.forEach( + (tableName, segments) -> segments.forEach(segment -> factory.registerSegment(port1, tableName, segment))); + QueryEnvironmentTestBase.SERVER2_SEGMENTS.forEach( + (tableName, segments) -> segments.forEach(segment -> factory.registerSegment(port2, tableName, segment))); + RoutingManager routingManager = factory.buildRoutingManager(null); + WorkerManager workerManager = + new WorkerManager("Broker_localhost", "localhost", NetUtils.findOpenPort(), routingManager); + + PinotConfiguration config = new PinotConfiguration(); + config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME, "localhost"); + config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, Integer.toString(NetUtils.findOpenPort())); + BrokerQueryEventListenerFactory.init(config); + BrokerMetrics.register(mock(BrokerMetrics.class)); + QueryQuotaManager queryQuotaManager = mock(QueryQuotaManager.class); + when(queryQuotaManager.acquire(anyString())).thenReturn(true); + when(queryQuotaManager.acquireDatabase(anyString())).thenReturn(true); + + return new MultiStageBrokerRequestHandler(config, "testBrokerId", new BrokerRequestIdGenerator(), routingManager, + accessControlFactory, queryQuotaManager, factory.buildTableCache(), mock(MultiStageQueryThrottler.class), + mock(FailureDetector.class), ThreadAccountantUtils.getNoOpAccountant(), null, workerManager, workerManager) { + @Override + public void start() { + } + + @Override + public void shutDown() { + } + + @Override + protected BrokerResponse handleRequestThrowing(long requestId, String query, + SqlNodeAndOptions sqlNodeAndOptions, @Nullable RequesterIdentity requesterIdentity, + RequestContext requestContext, @Nullable HttpHeaders httpHeaders) { + _requests.add(new Request(requestId, query, new HashMap<>(sqlNodeAndOptions.getOptions()), requestContext, + requesterIdentity, httpHeaders)); + QueryException error = _subqueryErrors.get(query); + if (error != null) { + throw error; + } + BrokerResponse response = _subqueryResponses.get(query); + return response != null + ? response + : super.handleRequestThrowing(requestId, query, sqlNodeAndOptions, requesterIdentity, requestContext, + httpHeaders); + } + + @Override + protected void onQueryCompletion(RequestContext requestContext, BrokerResponse brokerResponse) { + _completedRequests.add(requestContext); + } + }; + } + + private static class Request { + final long _requestId; + final String _query; + final Map _options; + final RequestContext _requestContext; + final RequesterIdentity _requesterIdentity; + final HttpHeaders _httpHeaders; + + Request(long requestId, String query, Map options, RequestContext requestContext, + @Nullable RequesterIdentity requesterIdentity, @Nullable HttpHeaders httpHeaders) { + _requestId = requestId; + _query = query; + _options = options; + _requestContext = requestContext; + _requesterIdentity = requesterIdentity; + _httpHeaders = httpHeaders; + } + + @Override + public String toString() { + return _requestId + ": " + _query; + } + } +} diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java index c9657f39ea6c..aadf20312f79 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java @@ -124,7 +124,7 @@ public int getNumRowsResultSet() { @Override public boolean isPartialResult() { return getExceptionsSize() > 0 || isNumGroupsLimitReached() || !getEarlyTerminationReasons().isEmpty() - || isMaxRowsInJoinReached() || isMseLiteLeafStageLimitReached(); + || isMaxRowsInJoinReached() || isMaxRowsInWindowReached() || isMseLiteLeafStageLimitReached(); } @Override diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/request/RequestUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/request/RequestUtils.java index a120893c8677..4a27337c0e70 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/request/RequestUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/request/RequestUtils.java @@ -116,7 +116,13 @@ public static void setOptions(SqlNodeAndOptions sqlNodeAndOptions, JsonNode json return; } LOGGER.debug("Query options are set to: {}", requestOptions); - requestOptions = QueryOptionsUtils.resolveCaseInsensitiveOptions(requestOptions); + setOptions(sqlNodeAndOptions, QueryOptionsUtils.resolveCaseInsensitiveOptions(requestOptions)); + } + + /// Merges the given options into the options parsed from the SQL, like [#setOptions(SqlNodeAndOptions, JsonNode)] + /// does with the options of the request payload. The keys of the given options must be resolved, see + /// [QueryOptionsUtils#resolveCaseInsensitiveOptions(Map)]. + public static void setOptions(SqlNodeAndOptions sqlNodeAndOptions, Map requestOptions) { SqlOptionsMode sqlOptionsMode = QueryOptionsUtils.getSqlOptionsMode(requestOptions); Map sqlOptions = sqlNodeAndOptions.getOptions(); if (sqlOptionsMode != SqlOptionsMode.ALLOW && !sqlOptions.isEmpty()) { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java b/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java index 22cf653d6e53..838b5300f29f 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java @@ -52,6 +52,21 @@ public void testEarlyTerminationReasonsMarkPartialResponse() { assertTrue(brokerResponse.isPartialResult()); } + @Test + public void testRowLimitsMarkPartialResponse() + throws Exception { + // A join or window stops at its row limit only in overflow mode BREAK, and then returns part of its result + BrokerResponseNativeV2 joinResponse = new BrokerResponseNativeV2(); + joinResponse.mergeMaxRowsInJoinReached(true); + assertTrue(joinResponse.isPartialResult()); + + BrokerResponseNativeV2 windowResponse = new BrokerResponseNativeV2(); + assertFalse(windowResponse.isPartialResult()); + windowResponse.mergeMaxRowsInWindowReached(true); + assertTrue(windowResponse.isPartialResult()); + assertTrue(JsonUtils.objectToJsonNode(windowResponse).get("partialResult").asBoolean()); + } + @Test public void testResponseMetadataSerialization() throws Exception { diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java index d13a97232146..9f4c0e4b31d1 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java @@ -107,10 +107,9 @@ protected int getNumQueriesToGenerate() { public void testHardcodedQueries() throws Exception { testHardcodedQueriesCommon(); + testHardcodedIdSetQueries(); if (useMultiStageQueryEngine()) { testHardcodedQueriesV2(); - } else { - testHardCodedQueriesV1(); } } @@ -361,9 +360,9 @@ private void testHardcodedQueriesV2() testQuery(query, h2Query); } - private void testHardCodedQueriesV1() + /// Tests IN_ID_SET and IN_SUBQUERY, which are not standard SQL. + private void testHardcodedIdSetQueries() throws Exception { - // Non-Standard SQL syntax: // IN_ID_SET { IdSet idSet = IdSets.create(FieldSpec.DataType.LONG); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/JmxMetricsIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/JmxMetricsIntegrationTest.java index f96e42f298da..53d0cc96f5fa 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/JmxMetricsIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/JmxMetricsIntegrationTest.java @@ -110,15 +110,15 @@ public void testMultiStageMigrationMetric() throws Exception { postQuery("SELECT COUNT(*) FROM mytable"); + // Works with the multi-stage query engine too. The subquery counts as a query of its own. + JsonNode response = postQuery("SELECT AirlineID, count(*) FROM mytable WHERE IN_SUBQUERY(airlineID, 'SELECT " + + "ID_SET(AirlineID) FROM mytable WHERE Carrier = ''AA''') = 1 GROUP BY AirlineID;"); + assertFalse(response.get("resultTable").get("rows").isEmpty()); + // Run some queries that are known to not work as is with the multi-stage query engine // MV column requires ARRAY_TO_MV wrapper to be used in filter predicates - JsonNode response = postQuery("SELECT COUNT(*) FROM mytable WHERE DivAirports = 'JFK'"); - assertFalse(response.get("resultTable").get("rows").isEmpty()); - - // Unsupported function - response = postQuery("SELECT AirlineID, count(*) FROM mytable WHERE IN_SUBQUERY(airlineID, 'SELECT " - + "ID_SET(AirlineID) FROM mytable WHERE Carrier = ''AA''') = 1 GROUP BY AirlineID;"); + response = postQuery("SELECT COUNT(*) FROM mytable WHERE DivAirports = 'JFK'"); assertFalse(response.get("resultTable").get("rows").isEmpty()); // Repeated columns in an ORDER BY query @@ -131,14 +131,14 @@ public void testMultiStageMigrationMetric() throws Exception { TestUtils.waitForCondition((aVoid) -> { try { multiStageMigrationMetricValue.set((Long) MBEAN_SERVER.getAttribute(multiStageMigrationMetric, "Count")); - return multiStageMigrationMetricValue.get() == 3L; + return multiStageMigrationMetricValue.get() == 2L; } catch (Exception e) { throw new RuntimeException(e); } }, 5000, "Expected value of MBean 'pinot.broker.singleStageQueriesInvalidMultiStage' to be: " - + 3L + "; actual value: " + multiStageMigrationMetricValue.get()); + + 2L + "; actual value: " + multiStageMigrationMetricValue.get()); - assertEquals((Long) MBEAN_SERVER.getAttribute(multiStageMigrationMetric, "Count"), 3L); + assertEquals((Long) MBEAN_SERVER.getAttribute(multiStageMigrationMetric, "Count"), 2L); } @Test(dataProvider = "useBothQueryEngines") diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java index e10ae92a223c..c4697d74a088 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java @@ -32,6 +32,7 @@ import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; @@ -55,6 +56,10 @@ public class IdSetIntegrationTest extends CustomDataQueryClusterIntegrationTest private static final String BLOOM_FILTER_PARAMS = "expectedInsertions=1000;fpp=0.0001"; // A Bloom filter sized for 10M ids, which serializes to a literal of about 16 MB private static final String LARGE_BLOOM_FILTER_PARAMS = "expectedInsertions=10000000;fpp=0.01"; + // A Bloom filter sized for 20M ids, which serializes to a literal of about 32 MB + private static final String HUGE_BLOOM_FILTER_PARAMS = "expectedInsertions=20000000;fpp=0.01"; + // The number of rows whose value is in the sets built from the rows with id < NUM_SET_VALUES, without false positives + private static final long NUM_ROWS_IN_SET = (long) NUM_SET_VALUES * (NUM_ROWS / NUM_VALUES); @Override protected long getCountStarResult() { @@ -99,15 +104,93 @@ public void testFilterBySetFromEitherEngine(boolean buildSetWithMultiStageEngine @Test public void testScalarSubquery() throws Exception { - // The set argument is a scalar subquery, so IN_ID_SET runs as a scalar function above the join that supplies it + // The subquery is not correlated, so it runs before the query, like the subquery of an IN_SUBQUERY setUseMultiStageQueryEngine(true); for (String column : List.of(INT_COL, LONG_COL, STR_COL, BYTES_COL)) { String idSetQuery = "SELECT " + idSetExpression(column) + " FROM " + getTableName() + " WHERE " + SET_FILTER; - String query = - "SELECT COUNT(*) FROM " + getTableName() + " WHERE IN_ID_SET(" + column + ", (" + idSetQuery + ")) = 1"; IdSet idSet = IdSets.fromBase64String(getSingleValue(postQuery(idSetQuery)).asText()); - assertEquals(getSingleValue(postQuery(query)).asLong(), getExpectedCount(column, idSet), column); + assertEquals(getCount("IN_ID_SET(" + column + ", (" + idSetQuery + ")) = 1"), getExpectedCount(column, idSet), + column); } + + // So the filter runs in the leaf stage, without a join. EXPLAIN does not run the subquery, and says so. + JsonNode response = postQuery("EXPLAIN PLAN FOR SELECT COUNT(*) FROM " + getTableName() + " WHERE IN_ID_SET(" + + INT_COL + ", (SELECT IDSET(" + INT_COL + ") FROM " + getTableName() + " WHERE " + SET_FILTER + ")) = 1"); + assertTrue(response.get("exceptions").isEmpty(), response.get("exceptions").toString()); + JsonNode resultTable = response.get("resultTable"); + String plan = resultTable.get("rows").get(0).get(1).asText(); + assertFalse(plan.contains("Join"), plan); + assertTrue(resultTable.get("dataSchema").get("columnNames").toString().contains("SUBQUERIES"), + resultTable.toString()); + + // A correlated subquery runs as part of the query + assertEquals(getSingleValue(postQuery("SELECT COUNT(*) FROM " + getTableName() + " t1 WHERE IN_ID_SET(t1." + + INT_COL + ", (SELECT IDSET(t2." + INT_COL + ") FROM " + getTableName() + " t2 WHERE t2." + SET_FILTER + + " AND t2." + LONG_COL + " = t1." + LONG_COL + ")) = 1")).asLong(), NUM_ROWS_IN_SET); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testInSubquery(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + for (String column : List.of(INT_COL, LONG_COL, STR_COL, BYTES_COL)) { + String idSetQuery = "SELECT " + idSetExpression(column) + " FROM " + getTableName() + " WHERE " + SET_FILTER; + long expectedCount = + getExpectedCount(column, IdSets.fromBase64String(getSingleValue(postQuery(idSetQuery)).asText())); + String inSubquery = "IN_SUBQUERY(" + column + ", " + quote(idSetQuery) + ")"; + assertEquals(getCount(inSubquery + " = 1"), expectedCount, column); + assertEquals(getCount(inSubquery + " = 0"), NUM_ROWS - expectedCount, column); + } + } + + @Test + public void testInSubqueryOutsideFilters() + throws Exception { + // Only the multi-stage engine runs IN_SUBQUERY outside the WHERE clause + setUseMultiStageQueryEngine(true); + String subquery = quote("SELECT IDSET(" + INT_COL + ") FROM " + getTableName() + " WHERE " + SET_FILTER); + String inSubquery = "IN_SUBQUERY(" + INT_COL + ", " + subquery + ") = 1"; + assertEquals(getSingleValue(postQuery( + "SELECT SUM(CASE WHEN " + inSubquery + " THEN 1 ELSE 0 END) FROM " + getTableName())).asLong(), + NUM_ROWS_IN_SET); + JsonNode response = postQuery("SELECT " + INT_COL + ", COUNT(*) FROM " + getTableName() + " GROUP BY " + INT_COL + + " HAVING " + inSubquery + " LIMIT " + NUM_VALUES); + assertTrue(response.get("exceptions").isEmpty(), response.get("exceptions").toString()); + assertEquals(response.get("resultTable").get("rows").size(), NUM_SET_VALUES); + assertEquals(getSingleValue(postQuery("SELECT COUNT(*) FROM " + getTableName() + " t1 JOIN " + getTableName() + + " t2 ON t1." + ID + " = t2." + ID + " WHERE IN_SUBQUERY(t2." + INT_COL + ", " + subquery + ") = 1")) + .asLong(), NUM_ROWS_IN_SET); + + // A subquery can have IN_SUBQUERY too + String nestedInSubquery = "IN_SUBQUERY(" + ID + ", " + quote( + "SELECT IDSET(" + ID + ") FROM " + getTableName() + " WHERE " + SET_FILTER) + ") = 1"; + assertEquals(getCount("IN_SUBQUERY(" + INT_COL + ", " + quote( + "SELECT IDSET(" + INT_COL + ") FROM " + getTableName() + " WHERE " + nestedInSubquery) + ") = 1"), + NUM_ROWS_IN_SET); + } + + @Test + public void testFailsOnPartialSubqueryResult() + throws Exception { + // The subquery stops at 10 groups, so its IdSet would miss ids + setUseMultiStageQueryEngine(true); + String subquery = "SET numGroupsLimit = 10; SELECT IDSET(" + INT_COL + ") FROM (SELECT " + INT_COL + " FROM " + + getTableName() + " GROUP BY " + INT_COL + ")"; + JsonNode response = postQuery("SELECT COUNT(*) FROM " + getTableName() + " WHERE IN_SUBQUERY(" + INT_COL + ", " + + quote(subquery) + ") = 1"); + String exceptions = response.get("exceptions").toString(); + assertTrue(exceptions.contains("Subquery returned a partial result [numGroupsLimitReached]"), exceptions); + } + + @Test + public void testHugeIdSetInSubquery() + throws Exception { + // The set never reaches the client or the parser, so it can be larger than the 20M chars a JSON request can hold + setUseMultiStageQueryEngine(true); + String idSetQuery = "SELECT IDSET(" + STR_COL + ", '" + HUGE_BLOOM_FILTER_PARAMS + "') FROM " + getTableName() + + " WHERE " + SET_FILTER; + // With 60 of the 20M ids it is sized for, the Bloom filter has practically no false positives + assertEquals(getCount("IN_SUBQUERY(" + STR_COL + ", " + quote(idSetQuery) + ") = 1"), NUM_ROWS_IN_SET); } @Test(dataProvider = "useBothQueryEngines") @@ -181,6 +264,17 @@ private static String stringValue(int value) { return "value_" + value; } + /// Returns the SQL string literal for the given string. + private static String quote(String string) { + return "'" + string.replace("'", "''") + "'"; + } + + /// Returns the number of rows that pass the filter. + private long getCount(String filter) + throws Exception { + return getSingleValue(postQuery("SELECT COUNT(*) FROM " + getTableName() + " WHERE " + filter)).asLong(); + } + private static JsonNode getSingleValue(JsonNode response) { assertTrue(response.get("exceptions").isEmpty(), response.get("exceptions").toString()); return response.get("resultTable").get("rows").get(0).get(0); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java index e05ee827547d..c6750013e74d 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java @@ -299,6 +299,12 @@ public static PinotOperatorTable instance(boolean nullHandlingEnabled) { OperandTypes.family( List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.INTEGER), i -> i == 2)), + // IN_SUBQUERY(expr, 'subquery') never runs: the broker runs the subquery first and replaces the call with an + // IN_ID_SET on the result. It is registered for the paths that only compile queries, such as the controller + // routing a query to a broker. Nullable like IN_ID_SET. + new PinotSqlFunction("IN_SUBQUERY", ReturnTypes.BOOLEAN_NULLABLE, + OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)), + // Placeholder for special functions to handle MV // NOTE: // ARRAY_TO_MV is not deterministic. diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index 34cc9e7c92d4..ce5d25c7558a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -347,6 +347,20 @@ public boolean canCompileQuery(String query) { } } + /// Validates the query without converting or optimizing it, which costs less than [#compile(String, + /// SqlNodeAndOptions)]. The validation changes the parsed query. + /// + /// @throws QueryException if the query does not validate, e.g. because of an invalid query option + public void validate(SqlNodeAndOptions sqlNodeAndOptions) { + try (PlannerContext plannerContext = getPlannerContext(sqlNodeAndOptions)) { + validate(sqlNodeAndOptions.getSqlNode(), plannerContext); + } catch (QueryException e) { + throw e; + } catch (RuntimeException e) { + throw QueryErrorCode.QUERY_VALIDATION.asException(e.getMessage(), e); + } + } + /// Results of planning a query public static class QueryPlannerResult { private final DispatchableSubPlan _dispatchableSubPlan; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java index 714d39839443..0d1fe39bd3c3 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java @@ -47,6 +47,11 @@ import org.apache.pinot.query.planner.plannode.SetOpNode; import org.apache.pinot.query.planner.plannode.WindowNode; import org.apache.pinot.query.routing.QueryServerInstance; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; +import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.apache.pinot.sql.parsers.SqlNodeAndOptions; import org.testng.annotations.DataProvider; import org.testng.annotations.Ignore; import org.testng.annotations.Test; @@ -1689,6 +1694,35 @@ private static boolean isOneOf(List> allowedNodeTy return false; } + @Test + public void testCompilesInSubquery() { + // The broker runs the subquery before it compiles the query, but the paths that only compile queries accept + // IN_SUBQUERY too. They see the tables outside the subquery only. + try (QueryEnvironment.CompiledQuery compiledQuery = _queryEnvironment.compile( + "SELECT COUNT(*) FROM a WHERE IN_SUBQUERY(col1, 'SELECT IDSET(col1) FROM b') = 1")) { + assertEquals(compiledQuery.getTableNames(), Set.of("a")); + } + assertTrue(_queryEnvironment.canCompileQuery( + "SELECT inSubquery(col3, 'SELECT IDSET(col3) FROM b') FROM a WHERE INSUBQUERY(col2, 'SELECT 1') = 0")); + } + + @Test + public void testValidate() { + _queryEnvironment.validate(CalciteSqlParser.compileToSqlNodeAndOptions("SELECT col1 FROM a WHERE col3 > 1")); + + // A correlated subquery does not validate on its own + expectThrows(QueryException.class, () -> _queryEnvironment.validate( + CalciteSqlParser.compileToSqlNodeAndOptions("SELECT IDSET(b.col1) FROM b WHERE b.col2 = a.col2"))); + + // An invalid option fails before the validation starts, but is reported the same way + SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions("SELECT col1 FROM a"); + sqlNodeAndOptions.getOptions().put(QueryOptionKey.INFER_PARTITION_HINT, "foo"); + QueryException exception = + expectThrows(QueryException.class, () -> _queryEnvironment.validate(sqlNodeAndOptions)); + assertEquals(exception.getErrorCode(), QueryErrorCode.QUERY_VALIDATION); + assertTrue(exception.getMessage().contains(QueryOptionKey.INFER_PARTITION_HINT), exception.getMessage()); + } + @DataProvider(name = "testQueryExceptionDataProvider") private Object[][] provideQueriesWithException() { return new Object[][]{