Cleanup controller and MSE broker entry points to keep error codes an… #15277

gortiz · 2025-03-14T14:53:48Z

This PR fixes some of the issues described in #14950.

It is mainly a clean up of MultiStageBrokerRequestHandler and PinotQueryResource to make parts reusable and query handling more consistent.

Specifically, it fixes:

On errors, Pinot controller query endpoints inconsistently return broker-like JSON error payloads (if the broker detects the error) or plain text (if the controller itself detected the error)
The Pinot controller detected errors are not logged.
MSE errors don't include the actual error code. Instead, error code 200 is always returned.

And it also improves (but not totally fixes) MSE errors including the stack trace

Notes for reviewers:

The main that have been changed are:
- pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
  - A slight refactor that makes code easier to extend. It also defines 3 level of errors (green, yellow and red) used to decide whether to log or not.
- pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java
  - Most methods return the streaming output instead of a string. This is more efficient (we don't need to allocate all the responses from brokers, which could be very expensive when returning large amounts of data!) and can also be used to return actual errors.
- pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
  I introduced the concept of CompiledQuery, which is the output of applying Calcite. This simplifies the code and makes it more efficient, given that we could be optimizing a query more than once previously.
The diff between some classes is larger than expected because I've had to add a new try-with-resources. I recommend to review the PR with hidden whitespaces

…d make it easier to add better ones

codecov-commenter · 2025-03-14T15:36:27Z

Codecov Report

Attention: Patch coverage is 27.56410% with 226 lines in your changes missing coverage. Please review.

Project coverage is 63.57%. Comparing base (59551e4) to head (34909ba).
Report is 1880 commits behind head on master.

Files with missing lines	Patch %	Lines
...requesthandler/MultiStageBrokerRequestHandler.java	0.00%	109 Missing ⚠️
...t/controller/api/resources/PinotQueryResource.java	14.51%	52 Missing and 1 partial ⚠️
.../java/org/apache/pinot/query/QueryEnvironment.java	60.37%	39 Missing and 3 partials ⚠️
...main/java/org/apache/pinot/common/utils/Timer.java	0.00%	10 Missing ⚠️
...g/apache/pinot/query/parser/utils/ParserUtils.java	0.00%	7 Missing ⚠️
...pinot/query/CalciteContextExceptionClassifier.java	44.44%	2 Missing and 3 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #15277      +/-   ##
============================================
+ Coverage     61.75%   63.57%   +1.82%     
- Complexity      207     1459    +1252     
============================================
  Files          2436     2783     +347     
  Lines        133233   156947   +23714     
  Branches      20636    24070    +3434     
============================================
+ Hits          82274    99781   +17507     
- Misses        44911    49653    +4742     
- Partials       6048     7513    +1465

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (+99.99%)`	⬆️
integration	`100.00% <ø> (+99.99%)`	⬆️
integration1	`100.00% <ø> (+99.99%)`	⬆️
integration2	`0.00% <ø> (ø)`
java-11	`63.55% <27.56%> (+1.84%)`	⬆️
java-21	`63.47% <27.56%> (+1.84%)`	⬆️
skip-bytebuffers-false	`63.57% <27.56%> (+1.82%)`	⬆️
skip-bytebuffers-true	`63.44% <27.56%> (+35.71%)`	⬆️
temurin	`63.57% <27.56%> (+1.82%)`	⬆️
unittests	`63.57% <27.56%> (+1.82%)`	⬆️
unittests1	`56.10% <47.51%> (+9.20%)`	⬆️
unittests2	`34.18% <18.58%> (+6.45%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…d make it easier to add better ones

…o improve-query-error-codes

…-codes

…ests

…-codes

gortiz · 2025-03-19T09:56:03Z

pinot-integration-test-base/pom.xml

+    <dependency>
+      <groupId>org.assertj</groupId>
+      <artifactId>assertj-core</artifactId>
+      <scope>test</scope>
+    </dependency>


I think it is time to start using better assertions.

gortiz · 2025-03-19T10:03:23Z

...on-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java

-  /**
-   * Test invalid queries which should cause query exceptions.
-   *
-   * @throws Exception
-   */
-  public void testQueryExceptions()
-      throws Exception {
-    testQueryException("POTATO", QueryErrorCode.SQL_PARSING);
-
-    // Ideally, we should attempt to unify the error codes returned by the two query engines if possible
-    testQueryException("SELECT COUNT(*) FROM potato",
-        useMultiStageQueryEngine()
-            ? QueryErrorCode.QUERY_PLANNING : QueryErrorCode.TABLE_DOES_NOT_EXIST);
-
-    testQueryException("SELECT POTATO(ArrTime) FROM mytable",
-        useMultiStageQueryEngine()
-            ? QueryErrorCode.QUERY_PLANNING : QueryErrorCode.QUERY_VALIDATION);
-
-    // ArrTime expects a numeric type
-    testQueryException("SELECT COUNT(*) FROM mytable where ArrTime = 'potato'",
-        useMultiStageQueryEngine()
-            ? QueryErrorCode.QUERY_EXECUTION : QueryErrorCode.QUERY_VALIDATION);
-
-    // Cannot use numeric aggregate function for string column
-    testQueryException("SELECT MAX(OriginState) FROM mytable where ArrTime > 5",
-        QueryErrorCode.QUERY_VALIDATION);
-  }
-
-  private void testQueryException(String query, QueryErrorCode errorCode)
-      throws Exception {
-    JsonNode jsonObject = postQuery(query);
-    assertEquals(jsonObject.get("exceptions").get(0).get("errorCode").asInt(), errorCode.getId());
-  }
-


For an unknown reason these tests were only executed in realtime integration tests. I don't think that makes sense. Instead now these tests are executed in ErrorCodesIntegrationTests and their subclasses. They are also split into different tests so we can have better coverage even if one of them fail. They are also executed against both controllers and brokers. So there are 2 configuration axes (MSE vs SSE and broker vs controller) and therefore 4 classes that actually run these tests. As a result, we need to startup the cluster 4 extra times. That is something we need to improve in the future (ie using JUnit 5 it would be trivial to do)

gortiz · 2025-03-19T11:01:02Z

pinot-query-planner/src/main/java/org/apache/pinot/query/catalog/PinotCatalog.java

-    if (tableName != null) {
-      if (_resolvedTables == null) {
-        _resolvedTables = new HashSet<>();
-      }
-      _resolvedTables.add(tableName);
-    }


@bziobrowski I don't think we need this anymore.

gortiz · 2025-03-19T11:02:42Z

pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java

@@ -439,51 +439,51 @@ public void testWindowFunctions() {
        "SELECT col1, col2, SUM(col3) OVER (PARTITION BY col1 ORDER BY col3 RANGE BETWEEN UNBOUNDED PRECEDING AND 1 "
            + "FOLLOWING) FROM a";
    e = expectThrows(RuntimeException.class, () -> _queryEnvironment.planQuery(sumQueryWithCustomRangeWindow));
-    assertTrue(e.getCause().getCause().getMessage()
+    assertTrue(e.getCause().getMessage()


Before the exception that was thrown was an ExecutionException whose cause was the actual exception. Now we get rid of the ExecutionException and therefore we should not call getCause

gortiz · 2025-03-19T11:03:32Z

pinot-query-runtime/src/test/resources/queries/Aggregates.json

@@ -644,13 +644,13 @@
        "description": "nested aggregation",
        "sql": "SELECT min(max(int_col)) FROM {tbl}",
        "comments": ".*Aggregate expressions cannot be nested.",
-        "expectedException": "Error composing query plan for.*"
+        "expectedException": ".*Aggregate expressions cannot be nested"


The error message now is better than before. In fact, we shouldn't have been asserting the original message.

yashmayya

Thanks @gortiz, this is a really nice cleanup!

yashmayya · 2025-03-19T10:38:54Z

pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/QueryAssert.java

+  /// assertions to be made before failing.
+  ///
+  /// See [Soft Assertions in AssertJ docs](https://assertj.github.io/doc/#assertj-core-soft-assertions)
+  public QueryErrorAssert.Soft softFirstException() {


When would we explicitly want non-soft assertions?

When only one condition is needed (then soft assertions are ok, but require more code to be used because they need to be closed) or when one condition only make sense if the next one is true (ie an assert verifies that something is not null and the next one uses it as a not null value).

There are some cases like that in this PR

yashmayya · 2025-03-19T10:39:38Z

pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/QueryAssert.java

+    return new QueryErrorAssert.Soft(actual.get("exceptions").get(0));
+  }
+
+  public static class Soft extends AbstractSoftAssertions implements AutoCloseable {


Yes, it is unused, but contrary to the other Soft class, this is the standard way to create soft assertions, so it is useful to keep it here in case we want to use soft assertions in the future for QueryAssert. I think it will be very useful once we start adding assertions on success scenarios.

yashmayya · 2025-03-19T10:47:05Z

...-tests/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java

+    assertThat(response.get("exceptions").get(0).get("message").asText())
+        .as("First exception message")
+        .contains("Illegal base64 character");


Why not use QueryAssert here?

I think I've modified this code before I created QueryAssert. Changing it.

yashmayya · 2025-03-19T10:52:02Z

pinot-common/src/main/java/org/apache/pinot/common/utils/Timer.java

-    return Math.max(remainingTime, 0);
+  public long getRemainingTimeMs() {
+    long remainingNs = _deadlineNs - _clock.nanos();
+    return Math.max(remainingNs / 1000, 0);


Shouldn't this be remainingNs / 1000000?

Yep. My bad. BTW, I've changed this class because nanos are recommended over millis and also because it was useful to add some extra methods

yashmayya · 2025-03-19T11:06:51Z

pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java

-    } catch (Throwable t) {
-      throw new RuntimeException("Error composing query plan for: " + sqlQuery, t);
+    try (CompiledQuery compiledQuery = compile(sqlQuery, sqlNodeAndOptions)) {
+      return compiledQuery.planQuery(requestId);


Neat refactor!

pinot-query-planner/src/main/java/org/apache/pinot/query/CalciteContextExceptionClassifier.java

...ration-tests/src/test/java/org/apache/pinot/integration/tests/ErrorCodesIntegrationTest.java

yashmayya · 2025-03-19T11:21:13Z

...t-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java

@@ -105,47 +104,41 @@ public class PinotQueryResource {
  @POST
  @Path("sql")
  @ManualAuthorization // performed by broker
-  public String handlePostSql(String requestJsonStr, @Context HttpHeaders httpHeaders) {
+  public StreamingOutput handlePostSql(String requestJsonStr, @Context HttpHeaders httpHeaders) {


What's the purpose of changing the return type to StreamingOutput?

The advantage of returning a StreamingOutput instead of a String is double:

We don't need to allocate the response on heap. Instead we can just pipeline the response from the broker. We use the same trick in LocalLogFileServer and TablesResource.

We can start to pipeline results to the customer as soon as the first bytes from broker arrive. Before we had to wait until the whole response was received in the controller.

...ker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

yashmayya · 2025-03-19T11:34:21Z

...ker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

+      case INTERNAL:
+      case UNKNOWN:


Why are these "yellow" errors and not "red"?

Red errors are errors we didn't catch and somehow bubbled up to the handleRequest. By definition, QueryExceptions are never red exceptions. Even if they are internal or unknown, we already caught them and explicitly set a user error message for them.

Colors are not about how problematic the error is. Instead, they are focused on how we want to log the error and how/what send to used.

Cleanup controller and MSE broker entry points to keep error codes an…

6164e18

…d make it easier to add better ones

gortiz requested review from vrajat, Jackie-Jiang and yashmayya March 14, 2025 14:53

gortiz mentioned this pull request Mar 14, 2025

Improve error handling and reporting in query side #14950

Open

9 tasks

yashmayya added cleanup multi-stage Related to the multi-stage query engine labels Mar 17, 2025

gortiz added 17 commits March 17, 2025 14:21

Cleanup controller and MSE broker entry points to keep error codes an…

e2f456b

…d make it easier to add better ones

Do not compile query twice

e5ebf3f

Merge remote-tracking branch 'personal/improve-query-error-codes' int…

2703ad7

…o improve-query-error-codes

Merge remote-tracking branch 'origin/master' into improve-query-error…

6948ec4

…-codes

Fix issue thread context issue using Calcite

745426d

Add javadoc

fa752cb

Fix tests now that exceptions are simpler

c23d681

Fix expected error message

e634b60

Add javadoc to some error codes^

34cec4a

Move assertNoError logic to QueryAssert^

f438c7d

Change some error codes^

f3119f7

Classify UNKNOWN_COLUMN and TABLE_DOES_NOT_EXIST in MSE

8f26fc4

Fix imports

bc6b760

Improve query exception test and move them in their own integration t…

f671fc5

…ests

Merge remote-tracking branch 'origin/master' into improve-query-error…

17d0f11

…-codes

Recover defaultEnableGroupTrim in MultiStageBrokerRequestHandler

d075093

Verify controller errors as well

ac8f1a3

gortiz commented Mar 19, 2025

View reviewed changes

Improve javadoc on QueryEnvironment and removed some unused functions

db91b39

gortiz commented Mar 19, 2025

View reviewed changes

gortiz requested a review from bziobrowski March 19, 2025 11:10

yashmayya reviewed Mar 19, 2025

View reviewed changes

This was referenced Mar 19, 2025

Improve error handling in controller #14951

Closed

Improve exceptions broker #14994

Closed

gortiz added 6 commits March 20, 2025 14:36

Rewrite an assert to use QueryAssert

165dbde

fix error in Timer.getRemainingTimeMs

f133195

Improve javadoc

3d46a27

Remove unused attributes

7e5e7c0

Fix incorrect link in javadoc

f695649

Fix typo in javadoc

34909ba

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Cleanup controller and MSE broker entry points to keep error codes an… #15277

Cleanup controller and MSE broker entry points to keep error codes an… #15277

gortiz commented Mar 14, 2025 •

edited

Loading

codecov-commenter commented Mar 14, 2025 •

edited

Loading

gortiz Mar 19, 2025

gortiz Mar 19, 2025

gortiz Mar 19, 2025

gortiz Mar 19, 2025

gortiz Mar 19, 2025

yashmayya left a comment

yashmayya Mar 19, 2025

gortiz Mar 20, 2025 •

edited

Loading

yashmayya Mar 19, 2025

gortiz Mar 20, 2025 •

edited

Loading

yashmayya Mar 19, 2025

gortiz Mar 20, 2025

yashmayya Mar 19, 2025

gortiz Mar 20, 2025

yashmayya Mar 19, 2025

yashmayya Mar 19, 2025

gortiz Mar 20, 2025

yashmayya Mar 19, 2025

gortiz Mar 20, 2025

Cleanup controller and MSE broker entry points to keep error codes an… #15277

Are you sure you want to change the base?

Cleanup controller and MSE broker entry points to keep error codes an… #15277

Conversation

gortiz commented Mar 14, 2025 • edited Loading

codecov-commenter commented Mar 14, 2025 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yashmayya left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gortiz Mar 20, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gortiz Mar 20, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gortiz commented Mar 14, 2025 •

edited

Loading

codecov-commenter commented Mar 14, 2025 •

edited

Loading

gortiz Mar 20, 2025 •

edited

Loading

gortiz Mar 20, 2025 •

edited

Loading