Support IdSet literals of millions of ids

[yashmayya]

PR flow

Enable efficient parsing and shared decoding of large IdSet literals across query segments.

flowchart TD
  N0["SQL Parsing with GrowingCharStream #40;F4#44; F5#41;"]:::stAdded
  N1["Transform Factory Init with QueryContext #40;F10#44; F7#41;"]:::stAdded
  N2["Get#47;Compute Shared IdSet from QueryContext #40;F7#44; F10#41;"]:::stAdded
  N3["Deserialize and Cache IdSet #40;F7#41;"]:::stAdded
  N4["Execute Filter Using Shared IdSet #40;F7#41;"]:::stAdded
  N0 --> N1
  N1 --> N2
  N2 --> N3
  N3 --> N4
  classDef stAdded fill:#dafbe1,stroke:#1a7f37,color:#1f2328,stroke-width:2px
  classDef stModified fill:#fff8c5,stroke:#9a6700,color:#1f2328,stroke-width:2px
  classDef stRemoved fill:#ffebe9,stroke:#cf222e,color:#1f2328,stroke-width:2px
  classDef stUnchanged fill:#f6f8fa,stroke:#656d76,color:#1f2328,stroke-width:1px

Loading

AI-generated · Green: added · Yellow: modified · Red: removed · Gray: existing

Diff evidence

• F4: pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java — before · after
• F5: pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/GrowingCharStream.java — after
• F7: pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunction.java — before · after
• F10: pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java — before · after

• Regenerate PR flow

Phase 1 (#19757) is merged, and this PR is based on master.

Problem

A query can filter on an IdSet of millions of ids, for example IN_ID_SET(col, '<16 MB literal>') = 1. With a literal that large, both engines have three problems:

• Parsing is quadratic. The generated JavaCC char stream grows its buffer by 2,048 chars at a time, and copies the buffer every time. To parse a literal of 1M chars, the broker allocates 2.6 GB. A 6 MB literal takes 2.5 s and 84 GiB.
• Each segment decodes its own copy of the IdSet.
• The MSE message size limits are fixed. MSE servers accept dispatched plans of up to 64 MB, and brokers accept messages from MSE servers of up to 64 MB. No config changes these limits.

Changes

• New GrowingCharStream: the SQL parser reads through this subclass of the generated SimpleCharStream. When a token outgrows the 4,096-char buffer, the buffer grows one time, to the length of the SQL.
  • A literal of 1M chars now parses in 13 ms with 12 MB.
  • Tokens of up to 2,048 chars never grow the buffer, so most queries allocate as before. For example, a 1M-value IN list allocates 526 MB, as before.
  • The cost is in a long SQL with one longer token. A token of more than 2,048 chars can grow the buffer, and one of more than 4,096 chars always does. For example, a 5,000-char literal in an 8M-char IN list query allocates an 80 MB buffer. Doubling the buffer would avoid that, but each long literal would then allocate 2 to 3 times as much.
• InIdSetTransformFunction decodes the IdSet once per query context, not once per segment. A query context covers one query on a server in SSE, and one leaf worker and table type (OFFLINE or REALTIME) in MSE. IdSet lookups only read, so many threads can share the IdSet.
  • TransformFunction gets an init overload that takes the QueryContext. The factory calls it, and its default calls the existing init.
  • QueryContext now keeps its shared values in its own map. Before, a ClassValue kept them: a value stayed in the cache of its class after the query, until about 21 later queries used the same class. For IdSets, that kept tens of MB for each old query. This also frees the shared literals and IN lists of other queries sooner. MemoizedClassAssociation had no other user, so it is deleted.
• New configs for the MSE gRPC message size limits. The defaults do not change. Each value must be positive, or the server or broker fails at startup.

Config	Default	Limits
pinot.query.multistage.query.server.max.inbound.message.size.bytes	64 MB	Plans that brokers dispatch to MSE servers. The plans carry the literals.
pinot.query.multistage.dispatch.channel.max.inbound.message.size.bytes	64 MB	Messages that MSE servers send to brokers, for example stats and EXPLAIN plans.

No change to the wire protocol or to serialized formats.

Results

• With a literal of about 16M chars (an IdSet for 10M ids at 1% false positives), the integration test query takes about 0.7 s in each engine. MSE planning of a query with this literal takes about 0.4 s and 352 MB.

Known limits

• HTTP queries: the broker parses the request body with Jackson's default limit of 20M chars for a string. So over HTTP, the SQL can have a literal of up to about 20M chars: an IdSet of about 12M ids at 1% false positives, or 16M at 3%. Larger sets need phase 3, below. The broker gRPC query endpoint does not help: it accepts requests of up to 4 MB, the gRPC default. The same Jackson limit applies to clients that read an IDSET result. This is not new, and SSE has the same limit.

Tests

• GrowingCharStreamTest:
  • The new stream reads the same tokens as the generated stream, with the same line and column numbers, in the lexical state of the SQL parser. The SQL has string literals, quoted identifiers, numbers and comments of up to 100K chars, at offsets that make the circular buffer wrap around. The test also makes sure that each SQL has its long token.
  • Only a long token grows the buffer, and only one time.
• CalciteSqlParserTest: a literal of 2M chars parses with less than 50 bytes per char. The generated stream allocates 10.8 GB for it. An empty query fails as before.
• InIdSetTransformFunctionTest: the segments of one query share one decoded IdSet. Another IdSet, or another query, gets its own.
• QueryServerTest and DispatchClientTest: each config sets the limit of a real gRPC server or channel. A larger message fails with RESOURCE_EXHAUSTED. Zero and negative values fail.
• MultiStageBrokerRequestHandlerTest: the broker gives the dispatch config to its dispatcher.
• IdSetIntegrationTest: both engines build an IdSet literal of about 16M chars and filter with it.

Each new test fails when its change is reverted. I did this check for each change. The QueryContext change has no automated test: checking that a value is freed needs a garbage collection, which a test cannot force reliably.

Next phases

• Phase 3, single-query forms in MSE: IN_SUBQUERY(expr, 'sql') and uncorrelated IN_ID_SET(expr, (SELECT IDSET(...) ...)). The broker runs the subquery first, so the set never goes through the parser, the client or the HTTP limit.
• Phase 4 (optional, measure first): one shared Bloom filter per query for IDSET, and shorter literals in EXPLAIN (Explain plan should elide function parameters #8009).

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.95775% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.54%. Comparing base (5b6c7a2) to head (90da69d).
⚠️ Report is 5 commits behind head on master.

Files with missing lines	Patch %	Lines
...r/transform/function/InIdSetTransformFunction.java	76.92%	2 Missing and 1 partial ⚠️
...e/pinot/query/service/dispatch/DispatchClient.java	50.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19759      +/-   ##
============================================
+ Coverage     68.52%   68.54%   +0.02%     
  Complexity     1470     1470              
============================================
  Files          3522     3522              
  Lines        229450   229612     +162     
  Branches      36436    36461      +25     
============================================
+ Hits         157222   157397     +175     
+ Misses        59906    59888      -18     
- Partials      12322    12327       +5     

Flag	Coverage Δ
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	68.54% <92.95%> (+0.02%)	⬆️
lane-a	100.00% <ø> (ø)
lane-b	0.00% <ø> (ø)
temurin	68.54% <92.95%> (+0.02%)	⬆️
unittests	68.54% <92.95%> (+0.02%)	⬆️
unittests1	58.25% <91.42%> (-0.01%)	⬇️
unittests2	40.14% <21.12%> (+0.01%)	⬆️

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

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

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[gortiz]

Thanks, this is a careful change and the tests are convincing. Approving.

I have one concern, and it does not block the PR. With this change, a query with a long token makes the parser allocate a buffer sized to the whole SQL, at 10 bytes per char (the char[] plus the line and column int[] arrays). At the 20M-char HTTP limit, that is about 200 MB per query on the broker. That comes on top of the SQL string itself and the 3 or 4 copies of the literal that Calcite and the Thrift conversion make. Nothing limits how many of these queries run at the same time, so a few concurrent large IdSet queries can put real pressure on the broker heap. It is still far better than the quadratic behavior we have today, but I think we should keep it in mind, and maybe mention broker heap sizing in the docs for these configs.

I think we can remove most of that cost by reading the SQL string directly instead of copying it into a buffer. The generated token manager only calls nine methods on the stream: readChar, backup, BeginToken, GetImage, GetSuffix and the begin/end line/column getters. All of them are public and not final. So a SimpleCharStream subclass can ignore the buffer completely:

readChar()     -> if (pos == len) throw new IOException(); return sql.charAt(pos++);  // EOF as in the generated stream
backup(n)      -> pos -= n
BeginToken()   -> tokenBegin = pos; save the line and column; return readChar()
GetImage()     -> sql.substring(tokenBegin, pos)   // the one copy the token needs anyway
GetSuffix(len) -> copy the last len chars from the string
getEnd*()      -> advance from the saved begin position

The extra memory of the stream then becomes O(1), and it no longer depends on the size of the whole SQL. That also removes the 80 MB case that the description mentions, for a 5,000-char literal in a long IN list query.

The tricky part is line and column tracking. Tabs, CR, LF and CRLF must behave exactly as in UpdateLineColumn. The line and column must also be computed by advancing from the position saved at BeginToken. Computing the column from the start of the line would become quadratic again for a single-line SQL. GrowingCharStreamTest already compares the tokens and their positions with the generated stream, so it would validate this version as it is. I'm happy to see this in a follow-up if you prefer to keep this PR as it is.

[yashmayya]

I'd prefer to leave the additional parser optimization as a follow-up. I've addressed the other comments.

[xiangfu0]

Docs follow-up merged: pinot-contrib/pinot-docs#1092. The IdSet guide now covers large literals and the MSE message-size settings are in the broker/server configuration references.