Repository navigation
Support IN_SUBQUERY in the multi-stage engine and run IdSet subqueries first - #19804
Conversation
yashmayya
added
multi-stage
Codecov Report✅ All modified and coverable lines are covered by tests.
Additional details and impacted files@@ Coverage Diff @@
## master #19804 +/- ##
============================================
- Coverage 68.55% 58.61% -9.94%
+ Complexity 1486 7 -1479
============================================
Files 3525 2722 -803
Lines 229852 169972 -59880
Branches 36507 27944 -8563
============================================
- Hits 157568 99629 -57939
- Misses 59965 61916 +1951
+ Partials 12319 8427 -3892
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| QueryEnvironment queryEnv = new QueryEnvironment( | ||
| getQueryEnvConf(httpHeaders, sqlNodeAndOptions.getOptions(), requestId), _multiClusterRoutingContext); | ||
| try { | ||
| queryEnv.validate(sqlNodeAndOptions); |
There was a problem hiding this 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.
There was a problem hiding this comment.
Good catch, fixed in 4db75a8. 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.
| } | ||
| } | ||
|
|
||
| private String getIdSet(String subquery) { |
There was a problem hiding this 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.
There was a problem hiding this 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.
xiangfu0
left a comment
xiangfu0
left a comment
There was a problem hiding this comment.
Reviewed the full diff with a focus on subquery semantics, planning, backward compatibility, and test coverage. This is a well-structured change: servers and the wire format never change (they only ever see IN_ID_SET(expr, '<literal>')), the EXPLAIN column is broker-local, and the test coverage is strong across the rewriter, the handler, compilation, and end-to-end.
I found no CRITICAL or MAJOR issues, so approving. A few minor observations are left inline for your confirmation — none are blocking.
Things I specifically verified and found sound:
- Subquery text extraction re-parses and
equalsDeep-compares, so any position/sanitization mismatch safely leaves the subquery in place rather than mis-rewriting. - WITH-item scoping, correlated detection (via
validate), run-once dedup, and the empty-IdSet null/no-row semantics all behave as documented and are tested. - Per-table access control and table QPS quota still apply to each subquery; deadline propagation caps the subquery timeout to the time left.
- The
isPartialResultwindow-flag change only affects BREAK overflow mode; the default THROW mode is unchanged.
| // 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 |
There was a problem hiding this 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.
| /// 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, |
There was a problem hiding this 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.
There was a problem hiding this comment.
I think I'll defer this one
| public boolean isPartialResult() { | ||
| return getExceptionsSize() > 0 || isNumGroupsLimitReached() || !getEarlyTerminationReasons().isEmpty() | ||
| || isMaxRowsInJoinReached() || isMseLiteLeafStageLimitReached(); | ||
| || isMaxRowsInJoinReached() || isMaxRowsInWindowReached() || isMseLiteLeafStageLimitReached(); |
There was a problem hiding this 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.
Follow-up to apache/pinot#19804. Updates engine compatibility and explains eligible scalar subqueries, independent checks/options/deadline, result shape, empty and partial results, cancellation, and EXPLAIN placeholders. --------- Co-authored-by: Xiang Fu <xiangfu@Xiang-mac-mtv-2.local>
|
Documentation follow-up merged: pinot-contrib/pinot-docs#1111 — MSE IN_SUBQUERY support and eligible IdSet pre-execution, options/deadline/checks, result/partial semantics, EXPLAIN, and plan-size caveats. |
4 participants
PR flow
The broker rewrites queries to run IdSet subqueries first, replacing them with IdSet literals before compilation.
AI-generated · Green: added · Yellow: modified · Red: removed · Gray: existing
Partial evidence: 0 file patches omitted; 1 truncated.
Diff evidence
Phase 3 of IdSet support in the multi-stage engine (MSE). Phases 1 (#19757) and 2 (#19759) are merged.
Problem
IN_SUBQUERY(col, 'subquery')with "No match found for function signature IN_SUBQUERY".IN_ID_SET(col, (SELECT IDSET(...) ...))has worked since Make IDSET and IN_ID_SET work in the multi-stage engine #19757, but slowly. The planner joins the subquery result to the rows, so every row leaves the leaf stage before the IdSet test runs.Changes
IN_ID_SET(col, '<IdSet>'), so the test runs in the leaf stage. Servers and the wire format do not change. Two forms run first:IN_SUBQUERY(expr, 'subquery'), anywhere in the query: SELECT, WHERE, HAVING, JOIN ON, nested queries, CTEs and set operations.IN_ID_SET(expr, (subquery)), when the subquery validates on its own (i.e. is not correlated) and names no WITH item of the query. Other subqueries run as part of the query, as before.SETstatements in anIN_SUBQUERYstring override them, andsqlOptionsModeapplies to them.dropResultsis not inherited.= 1is then false and= 0true for every row.numGroupsLimitReached, because the IdSet would miss ids.ID_SET_SUBQUERIES_NOT_RUNcolumn lists the subqueries as a JSON array.IN_SUBQUERYis registered in the MSE operator table. This lets paths that only compile queries accept it, such as the controller routing a query to a broker. As in the single-stage engine, those paths do not see the tables named in the subquery string.QueryEnvironment.validate(SqlNodeAndOptions): it validates without converting or optimizing, and reports every failure as aQueryException, including an invalid query option. The broker uses it to check whether a scalar subquery stands on its own. If that check fails for any reason, the subquery stays in the query, and the query's compilation reports the problem.RequestUtils.setOptions(SqlNodeAndOptions, Map)overload: it takes the request options as a map.Behavior changes
singleStageQueriesInvalidMultiStagemeter no longer counts single-stageIN_SUBQUERYqueries.IN_ID_SET(col, (subquery))now shows the IdSet test in the leaf stage instead of a join, plus the new column. Make IDSET and IN_ID_SET work in the multi-stage engine #19757 is in no release yet.partialResultis now true when a window stops at its row limit (overflow modeBREAK), as it already is for joins. The window flag was missing fromisPartialResult(). The query logger now logs these queries like other partial results. The default mode,THROW, does not change.Scale
In the integration test, a STRING IdSet sized for 20M ids (a 32M-char literal) filters the table in 0.5 s through
IN_SUBQUERY. The set never reaches the client or the parser, so the 20M-char HTTP limit does not apply.The limit is now the size of what the broker sends to each server:
pinot.query.multistage.query.server.max.inbound.message.size.bytes).So one use of a 32M-char set fits, but two uses on the same server need a higher limit. Otherwise, dispatch fails with
RESOURCE_EXHAUSTED.Tests
IdSetSubqueryRewriterTest: every SQL position, nested queries, name variants, the text of scalar subqueries, correlated subqueries, WITH names, repeated subqueries, EXPLAIN and invalid calls.MultiStageIdSetSubqueryTest: inherited options, identity and headers; time left; result checks, including partial results; errors; invalid inherited options; access checks on the subquery's tables; query events; and EXPLAIN.QueryCompilationTest: queries withIN_SUBQUERYcompile, andvalidatereports an invalid option as aQueryException.IdSetIntegrationTest:IN_SUBQUERYin both engines, outside filters, and nested; scalar subqueries, both uncorrelated and correlated; a partial subquery result; and a 20M-id set.BaseClusterIntegrationTestSet: theIN_ID_SETandIN_SUBQUERYhard-coded queries now run on both engines. The hybrid, realtime, logical table, cursor and MSE suites pass locally.JmxMetricsIntegrationTest: now expects 2 invalid queries instead of 3.Next
RESOURCE_EXHAUSTED.IN_SUBQUERYthrows an NPE when its subquery runs with null handling and matches no rows. A separate fix can also match the result checks above.