gortiz commented on code in PR #19804:
URL: https://github.com/apache/pinot/pull/19804#discussion_r4273762477
##########
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:
QueryEnvironment.validate(SqlNodeAndOptions) only wraps failures that happen
inside validate(SqlNode, PlannerContext). The QueryEnvironment constructor and
getPlannerContext() run before that try, so anything they throw because of an
inherited option escapes here as a plain RuntimeException. The probe then fails
the whole request as an UNKNOWN error, instead of returning false and keeping
the join plan. Could we catch RuntimeException here and return false? The main
compile still reports the real problem afterwards. The Javadoc of the new
QueryEnvironment.validate also says it throws QueryException, which is not true
for those paths, so either widen its try or adjust the doc.
##########
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:
This list repeats what BrokerResponse.isPartialResult() already checks, and
the two have already drifted: this one includes maxRowsInWindowReached and
isPartialResult() does not. What about gating on response.isPartialResult()
(and adding the window flag there, or here) and keeping this method only to
build the message? That way a partial-result flag added later is picked up here
too, instead of silently producing an incomplete IdSet.
##########
pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java:
##########
@@ -299,6 +299,12 @@ public static PinotOperatorTable instance(boolean
nullHandlingEnabled) {
OperandTypes.family(
List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
SqlTypeFamily.INTEGER), i -> i == 2)),
+ // IN_SUBQUERY(expr, 'subquery') never runs: the broker runs the
subquery first and replaces the call with an
+ // IN_ID_SET on the result. It is registered for the paths that only
compile queries, such as the controller
+ // routing a query to a broker. Nullable like IN_ID_SET.
+ new PinotSqlFunction("IN_SUBQUERY", ReturnTypes.BOOLEAN_NULLABLE,
Review Comment:
With IN_SUBQUERY registered here, any MSE path that compiles and executes a
query without IdSetSubqueryRewriter now plans fine and fails later on the
servers with an unknown function error. Today that would only be a subclass
that overrides handleRequestThrowing or a new entry point, so this is about the
failure mode, not a current bug. Would it be worth rejecting IN_SUBQUERY when
the query is planned for execution (physical planning / dispatch), which the
compile-only paths like the controller never reach? Then those paths keep
working and everything else gets a clear validation error.
##########
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:
The cache runs a repeated subquery once, but each call site still gets its
own copy of the literal in the plan. If the same 32 MB set appears twice in a
stage that goes to one server, that StagePlan is about 64 MB, which is the
default max inbound message size of the MSE QueryServer. So the practical
ceiling is now the gRPC dispatch limit per server message, not the 20M-char
HTTP limit. Could you mention that in the PR description and in the docs? A
clear error before dispatch would be nicer than a gRPC RESOURCE_EXHAUSTED, but
that can be a follow-up.
--
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]
