yashmayya commented on code in PR #19804:
URL: https://github.com/apache/pinot/pull/19804#discussion_r4275974261
##########
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:
Taking back my last point: this was a miss. The window flag came in #13180,
three days after `isPartialResult()`, and the later limit flags (early
termination, lite leaf limit) were all added to it. The flag is set only in
overflow mode BREAK, after the window operator cuts its input, so the result
really is partial, as with the join flag. The default mode is THROW, so only
queries that opt into BREAK and hit the limit change. Added it to
`isPartialResult()` in 947c303a27, and the check here is now just
`isPartialResult()`.
--
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]
