yashmayya commented on code in PR #19804:
URL: https://github.com/apache/pinot/pull/19804#discussion_r4275860199
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -689,9 +712,155 @@ private BrokerResponse
explain(QueryEnvironment.CompiledQuery query, long reques
() -> 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.
+ private boolean validatesOnItsOwn(long requestId, String subquery,
Map<String, String> queryOptions,
+ @Nullable HttpHeaders httpHeaders, long deadlineMs) {
+ getTimeLeftMs(subquery, deadlineMs);
+ SqlNodeAndOptions sqlNodeAndOptions =
CalciteSqlParser.compileToSqlNodeAndOptions(subquery);
+ RequestUtils.setOptions(sqlNodeAndOptions, queryOptions);
+ QueryEnvironment queryEnv = new QueryEnvironment(
+ getQueryEnvConf(httpHeaders, sqlNodeAndOptions.getOptions(),
requestId), _multiClusterRoutingContext);
+ try {
+ queryEnv.validate(sqlNodeAndOptions);
Review Comment:
Good catch, fixed in 4db75a8e7d. The check now returns false on any failure,
so the subquery stays in the query and the main compile reports the real error.
`QueryEnvironment.validate` also wraps those failures in a `QueryException`
now, so its Javadoc holds. The time-left check stays outside the try, so a
timeout still fails fast.
##########
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) {
Review Comment:
Agreed. It is a bit wider than one stage: the broker sends the plans of all
the stages a server runs in one message, so copies in different stages add up
too. I added this to the PR description (Scale section), with the server key to
raise. I will add it to the docs PR for the IdSet changes, and I listed the
pre-dispatch error as a follow-up under Next.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -689,9 +712,155 @@ private BrokerResponse
explain(QueryEnvironment.CompiledQuery query, long reques
() -> 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.
+ private boolean validatesOnItsOwn(long requestId, String subquery,
Map<String, String> queryOptions,
+ @Nullable HttpHeaders httpHeaders, long deadlineMs) {
+ getTimeLeftMs(subquery, deadlineMs);
+ SqlNodeAndOptions sqlNodeAndOptions =
CalciteSqlParser.compileToSqlNodeAndOptions(subquery);
+ RequestUtils.setOptions(sqlNodeAndOptions, queryOptions);
+ QueryEnvironment queryEnv = new QueryEnvironment(
+ getQueryEnvConf(httpHeaders, sqlNodeAndOptions.getOptions(),
requestId), _multiClusterRoutingContext);
+ try {
+ queryEnv.validate(sqlNodeAndOptions);
+ return true;
+ } catch (QueryException 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
+ List<String> partialResultReasons = getPartialResultReasons(response);
+ if (!partialResultReasons.isEmpty()) {
+ throw QueryErrorCode.QUERY_EXECUTION.asException(
+ "Subquery returned a partial result " + partialResultReasons + ": "
+ 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 reasons why the response holds a partial result without an
exception, if any.
+ private static List<String> getPartialResultReasons(BrokerResponse response)
{
Review Comment:
Makes sense, done in 4db75a8e7d: it now gates on `isPartialResult()` plus
the window flag, and the method only builds the message. I left
`isPartialResult()` itself as is, because adding the window flag there changes
`partialResult` for every MSE window query. I would rather do that in a
separate PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
