xiangfu0 commented on code in PR #19804:
URL: https://github.com/apache/pinot/pull/19804#discussion_r4278299322
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -689,9 +712,158 @@ 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.
+ ///
+ /// 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,
Review Comment:
A subquery that itself contains `IN_SUBQUERY` recurses through
`handleRequest` → `handleRequestThrowing` → `rewriteIdSetSubqueries` → here
again. Depth is effectively bounded by the shared deadline and the SQL text, so
it's not a real risk, but a small explicit max-depth guard would let a
pathological nested query fail fast with a clear message instead of leaning on
the timeout. Optional.
##########
pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java:
##########
@@ -124,7 +124,7 @@ public int getNumRowsResultSet() {
@Override
public boolean isPartialResult() {
return getExceptionsSize() > 0 || isNumGroupsLimitReached() ||
!getEarlyTerminationReasons().isEmpty()
- || isMaxRowsInJoinReached() || isMseLiteLeafStageLimitReached();
+ || isMaxRowsInJoinReached() || isMaxRowsInWindowReached() ||
isMseLiteLeafStageLimitReached();
Review Comment:
This now marks a response partial whenever the window row limit was hit.
Worth confirming it only triggers in overflow mode BREAK (the default THROW
throws instead and never sets the flag), so existing default-mode queries are
unaffected. It mirrors the existing `maxRowsInJoinReached` handling and is
required for the subquery partial-result guard, so this looks correct — just
noting the client-visible `partialResult` behavior change for the record.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -689,9 +712,158 @@ 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.
+ ///
+ /// 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
Review Comment:
Nice that this is called out explicitly. Just to confirm the intent: because
subqueries go through the 8-arg `handleRequest` rather than the public wrapper,
the application-level QPS quota (`acquireApplication`) is consumed only once
for the whole query, so a single accepted request with N `IN_SUBQUERY` calls
fans out to N+1 executions on one app-quota token. That matches the
single-stage engine, so this is probably intended — flagging only so it's a
conscious decision.
--
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]
