This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 49ce7be1798 Support IN_SUBQUERY in the multi-stage engine and run
IdSet subqueries first (#19804)
49ce7be1798 is described below
commit 49ce7be17983a16154fc639e45821256413e736d
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Oct 15 12:00:45 2026 -0700
Support IN_SUBQUERY in the multi-stage engine and run IdSet subqueries
first (#19804)
---
.../requesthandler/IdSetSubqueryRewriter.java | 264 +++++++++++
.../MultiStageBrokerRequestHandler.java | 176 +++++++-
.../requesthandler/IdSetSubqueryRewriterTest.java | 309 +++++++++++++
.../MultiStageIdSetSubqueryTest.java | 491 +++++++++++++++++++++
.../response/broker/BrokerResponseNativeV2.java | 2 +-
.../pinot/common/utils/request/RequestUtils.java | 8 +-
.../broker/BrokerResponseNativeV2Test.java | 15 +
.../tests/BaseClusterIntegrationTestSet.java | 7 +-
.../tests/JmxMetricsIntegrationTest.java | 18 +-
.../tests/custom/IdSetIntegrationTest.java | 102 ++++-
.../pinot/calcite/sql/fun/PinotOperatorTable.java | 6 +
.../org/apache/pinot/query/QueryEnvironment.java | 14 +
.../apache/pinot/query/QueryCompilationTest.java | 34 ++
13 files changed, 1425 insertions(+), 21 deletions(-)
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 00000000000..ea8980ae129
--- /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<String, String> _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<String> 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<String> 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<String> 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<String> _withNames;
+
+ IdSetCall(SqlBasicCall call, Set<String> withNames) {
+ _call = call;
+ _withNames = withNames;
+ }
+ }
+
+ /// Finds the IdSet subqueries of a query, except those in other IdSet
subqueries.
+ private static class IdSetCallFinder extends SqlBasicVisitor<Void> {
+ final List<IdSetCall> _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<String> _withNames;
+
+ IdSetCallFinder(Set<String> 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<String> 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<Void> {
+ final Set<String> _names;
+ boolean _found;
+
+ NameFinder(Set<String> 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 e2708b725e9..adf76c39c2e 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.ArrayList;
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.annotation.Nullable;
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.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.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.exception.QueryException;
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 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
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<String> 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 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
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 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
/// 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<String> idSetSubqueries)
throws WebApplicationException, QueryException {
Map<String, String> queryOptions = query.getOptions();
@@ -689,9 +712,158 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
() -> query.explain(requestId, fragmentToPlanNode), timer);
String plan = queryPlanResult.getExplainPlan();
Map<String, String> 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<String> 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<String, String> 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<String, String> 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<String,
String> 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<String, String> 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<String> 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<Object[]> 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<String> getPartialResultReasons(BrokerResponse response)
{
+ List<String> 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 00000000000..5ffc16e2291
--- /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<String> 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<String> 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<String> 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<String> 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<String> 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<String> 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<String> 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<String> 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<String> 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<String, String> _results = new HashMap<>();
+ final Set<String> _correlated = new HashSet<>();
+ final List<String> _validated = new ArrayList<>();
+ final List<String> _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 00000000000..54d7eb8471d
--- /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<String, BrokerResponse> _subqueryResponses = new
HashMap<>();
+ private final Map<String, QueryException> _subqueryErrors = new HashMap<>();
+ // All the requests the handler runs, including the queries of the tests
+ private final List<Request> _requests = new ArrayList<>();
+ private final List<RequestContext> _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<String, String> 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.<Object[]>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.<Object[]>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<String> tables) {
+ Set<String> 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<String> 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<Request> subqueryRequests = getSubqueryRequests();
+ assertEquals(subqueryRequests.size(), 1, subqueryRequests.toString());
+ return subqueryRequests.get(0);
+ }
+
+ private void assertNoSubqueryRan() {
+ assertEquals(getSubqueryRequests(), List.of());
+ }
+
+ private List<Request> getSubqueryRequests() {
+ return _requests.stream().filter(request -> request._requestId !=
REQUEST_ID).collect(Collectors.toList());
+ }
+
+ private static BrokerResponse idSetResponse(String idSet) {
+ return response(stringColumn(), List.<Object[]>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<Object[]>
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<String, String> _options;
+ final RequestContext _requestContext;
+ final RequesterIdentity _requesterIdentity;
+ final HttpHeaders _httpHeaders;
+
+ Request(long requestId, String query, Map<String, String> 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 d4b2152d4be..827c800d037 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 class BrokerResponseNativeV2 implements
BrokerResponse {
@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 a1bf57f25a1..0f913137795 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
@@ -118,7 +118,13 @@ public class RequestUtils {
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<String, String> requestOptions) {
SqlOptionsMode sqlOptionsMode =
QueryOptionsUtils.getSqlOptionsMode(requestOptions);
Map<String, String> 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 139fa96aa76..3a8d1aa4a88 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
@@ -82,6 +82,21 @@ public class BrokerResponseNativeV2Test {
assertEquals(json.get("directRemoteAccessBytes").asLong(), 20L);
}
+ @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 d13a9723214..9f4c0e4b31d 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 @@ public abstract class BaseClusterIntegrationTestSet
extends BaseClusterIntegrati
public void testHardcodedQueries()
throws Exception {
testHardcodedQueriesCommon();
+ testHardcodedIdSetQueries();
if (useMultiStageQueryEngine()) {
testHardcodedQueriesV2();
- } else {
- testHardCodedQueriesV1();
}
}
@@ -361,9 +360,9 @@ public abstract class BaseClusterIntegrationTestSet extends
BaseClusterIntegrati
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 f96e42f298d..53d0cc96f5f 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 class JmxMetricsIntegrationTest extends
BaseClusterIntegrationTestSet {
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 class JmxMetricsIntegrationTest extends
BaseClusterIntegrationTestSet {
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 e10ae92a223..c4697d74a08 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.apache.pinot.spi.data.Schema;
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 class IdSetIntegrationTest extends
CustomDataQueryClusterIntegrationTest
@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 @@ public class IdSetIntegrationTest extends
CustomDataQueryClusterIntegrationTest
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 e05ee827547..c6750013e74 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 class PinotOperatorTable implements
SqlOperatorTable {
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 34cc9e7c92d..ce5d25c7558 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 class QueryEnvironment {
}
}
+ /// 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 baa47e43635..c91508d5d41 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.ProjectNode;
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 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
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[][]{
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
