This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 1ea817199aa Support IdSet literals of millions of ids (#19759)
1ea817199aa is described below
commit 1ea817199aa0b29e473c4d6bf49b660b0fbe219c
Author: Yash Mayya <[email protected]>
AuthorDate: Wed Oct 7 16:38:58 2026 -0700
Support IdSet literals of millions of ids (#19759)
---
.../MultiStageBrokerRequestHandler.java | 5 +-
.../MultiStageBrokerRequestHandlerTest.java | 19 ++-
.../apache/pinot/sql/parsers/CalciteSqlParser.java | 15 +--
.../sql/parsers/parser/GrowingCharStream.java | 86 +++++++++++++
.../pinot/sql/parsers/CalciteSqlParserTest.java | 57 +++++++--
.../sql/parsers/parser/GrowingCharStreamTest.java | 138 +++++++++++++++++++++
.../function/InIdSetTransformFunction.java | 36 +++++-
.../transform/function/TransformFunction.java | 12 ++
.../function/TransformFunctionFactory.java | 2 +-
.../core/query/request/context/QueryContext.java | 8 +-
.../pinot/core/util/MemoizedClassAssociation.java | 42 -------
.../function/InIdSetTransformFunctionTest.java | 98 +++++++++++++++
.../tests/custom/IdSetIntegrationTest.java | 17 +++
.../query/service/dispatch/DispatchClient.java | 18 ++-
.../query/service/dispatch/QueryDispatcher.java | 35 +++++-
.../pinot/query/service/server/QueryServer.java | 26 ++--
.../query/service/dispatch/DispatchClientTest.java | 138 +++++++++++++++++++++
.../query/service/server/QueryServerTest.java | 55 ++++++++
.../apache/pinot/spi/utils/CommonConstants.java | 14 +++
19 files changed, 738 insertions(+), 83 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
index f42c17adbb7..e2708b725e9 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
@@ -212,11 +212,14 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
CommonConstants.Broker.DEFAULT_STREAM_STATS_DRAIN_MS);
boolean enableProtoSegmentList =
_config.getProperty(CommonConstants.Broker.CONFIG_OF_MSE_ENABLE_PROTO_SEGMENT_LIST,
CommonConstants.Broker.DEFAULT_MSE_ENABLE_PROTO_SEGMENT_LIST);
+ int dispatchMaxInboundMessageSizeBytes = config.getProperty(
+
CommonConstants.MultiStageQueryRunner.KEY_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES,
+
CommonConstants.MultiStageQueryRunner.DEFAULT_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES);
_mailboxService = new MailboxService(hostname, port, InstanceType.BROKER,
config, tlsConfig);
_queryDispatcher =
new QueryDispatcher(_mailboxService, failureDetector, tlsConfig,
isQueryCancellationEnabled(), cancelTimeout,
dispatchKeepAliveTimeMs, dispatchKeepAliveTimeoutMs,
dispatchKeepAliveWithoutCalls, _streamStatsDefault,
- streamStatsDrainMs, enableProtoSegmentList);
+ streamStatsDrainMs, enableProtoSegmentList,
dispatchMaxInboundMessageSizeBytes);
LOGGER.info("Initialized MultiStageBrokerRequestHandler on host: {}, port:
{} with broker id: {}, timeout: {}ms, "
+ "query log max length: {}, query log max rate: {}, query
cancellation enabled: {}", hostname, port,
_brokerId, _brokerTimeoutMs, _queryLogger.getMaxQueryLengthToLog(),
_queryLogger.getLogRateLimit(),
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
index fc36e0cf2cb..d2d99c2ebbd 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
@@ -210,6 +210,16 @@ public class MultiStageBrokerRequestHandlerTest extends
QueryEnvironmentTestBase
"No option should be injected when the broker default is unset");
}
+ @Test
+ public void testDispatchMaxInboundMessageSizeFromConfig()
+ throws Exception {
+ assertEquals(newHandler(new
PinotConfiguration()).getQueryDispatcher().getDispatchMaxInboundMessageSizeBytes(),
+
MultiStageQueryRunner.DEFAULT_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES);
+ PinotConfiguration config = new PinotConfiguration();
+
config.setProperty(MultiStageQueryRunner.KEY_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES,
"12345678");
+
assertEquals(newHandler(config).getQueryDispatcher().getDispatchMaxInboundMessageSizeBytes(),
12_345_678);
+ }
+
@Test
public void testDefaultDisabledPlannerRulesFromClusterConfig() {
// ServiceStartableUtils.applyClusterConfig applies cluster configs with
setProperty, which doesn't split lists
@@ -247,8 +257,6 @@ public class MultiStageBrokerRequestHandlerTest extends
QueryEnvironmentTestBase
@Nullable String streamingGroupByFlushThreshold, @Nullable String
streamingDistinctFlushThreshold)
throws Exception {
PinotConfiguration config = new PinotConfiguration();
- config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME,
"localhost");
- config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT,
Integer.toString(NetUtils.findOpenPort()));
if (streamingGroupByFlushThreshold != null) {
config.setProperty(CommonConstants.Broker.CONFIG_OF_MSE_STREAMING_GROUP_BY_FLUSH_THRESHOLD,
streamingGroupByFlushThreshold);
@@ -257,6 +265,13 @@ public class MultiStageBrokerRequestHandlerTest extends
QueryEnvironmentTestBase
config.setProperty(CommonConstants.Broker.CONFIG_OF_MSE_STREAMING_DISTINCT_FLUSH_THRESHOLD,
streamingDistinctFlushThreshold);
}
+ return newHandler(config);
+ }
+
+ private static MultiStageBrokerRequestHandler newHandler(PinotConfiguration
config)
+ throws Exception {
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME,
"localhost");
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT,
Integer.toString(NetUtils.findOpenPort()));
BrokerQueryEventListenerFactory.init(config);
BrokerMetrics.register(mock(BrokerMetrics.class));
diff --git
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
index 0d85fb9b03d..d1f8f9bef86 100644
---
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
+++
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
@@ -21,7 +21,6 @@ package org.apache.pinot.sql.parsers;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableSet;
-import java.io.StringReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
@@ -72,6 +71,7 @@ import
org.apache.pinot.common.utils.config.QueryOptionsUtils.SqlOptionsMode;
import org.apache.pinot.common.utils.request.RequestUtils;
import org.apache.pinot.segment.spi.AggregationFunctionType;
import org.apache.pinot.sql.FilterKind;
+import org.apache.pinot.sql.parsers.parser.GrowingCharStream;
import org.apache.pinot.sql.parsers.parser.SqlInsertFromFile;
import org.apache.pinot.sql.parsers.parser.SqlParserImpl;
import org.apache.pinot.sql.parsers.parser.SqlPinotCreateMaterializedView;
@@ -137,8 +137,8 @@ public class CalciteSqlParser {
}
}
- try (StringReader inStream = new StringReader(sql)) {
- SqlParserImpl sqlParser = newSqlParser(inStream);
+ try {
+ SqlParserImpl sqlParser = newSqlParser(sql);
SqlNodeList sqlNodeList = sqlParser.parseSqlStmtList();
sqlNodeList = (SqlNodeList) PostgreSqlCastRewriter.rewrite(sqlNodeList);
// Extract OPTION statements from sql.
@@ -667,8 +667,8 @@ public class CalciteSqlParser {
/// @throws SqlCompilationException if String is not a valid expression.
public static Expression compileToExpression(String expression) {
SqlNode sqlNode;
- try (StringReader inStream = new StringReader(expression)) {
- SqlParserImpl sqlParser = newSqlParser(inStream);
+ try {
+ SqlParserImpl sqlParser = newSqlParser(expression);
sqlNode = sqlParser.parseSqlExpressionEof();
} catch (Throwable e) {
throw new SqlCompilationException("Caught exception while parsing
expression: " + expression, e);
@@ -679,8 +679,9 @@ public class CalciteSqlParser {
}
@VisibleForTesting
- static SqlParserImpl newSqlParser(StringReader inStream) {
- SqlParserImpl sqlParser = new SqlParserImpl(inStream);
+ static SqlParserImpl newSqlParser(String sql) {
+ // Reads long tokens, e.g. a multi-MB IdSet literal, in linear time
+ SqlParserImpl sqlParser = GrowingCharStream.newParser(sql);
sqlParser.switchTo(SqlAbstractParserImpl.LexicalState.DQID);
// TODO: convert to MySQL conformance once we retired most of the
un-tested BABEL tokens
sqlParser.setConformance(SqlConformanceEnum.BABEL);
diff --git
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/GrowingCharStream.java
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/GrowingCharStream.java
new file mode 100644
index 00000000000..2171a2a98c7
--- /dev/null
+++
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/GrowingCharStream.java
@@ -0,0 +1,86 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.sql.parsers.parser;
+
+import java.io.StringReader;
+
+
+/// A [SimpleCharStream] over a SQL string that grows its buffer to fit the
whole SQL when a token outgrows it.
+///
+/// The generated stream keeps the current token in a buffer of 4,096 chars,
which it grows by 2,048 chars at a time.
+/// Each growth copies the buffer and its line and column arrays (10 bytes per
char), so reading a token of n chars
+/// costs O(n^2) time and memory, e.g. 2.6 GB of allocation for a string
literal of 1M chars such as a serialized IdSet.
+/// Growing the buffer once, to a size that fits any token of the SQL, makes
it O(n). Tokens of up to 2,048 chars never
+/// grow the buffer, so most queries, including those with long IN lists, are
not affected.
+///
+/// The cost is in a long SQL with one longer token: the buffer grows to the
length of the whole SQL. A token of more
+/// than 2,048 chars can grow it, depending on where the token starts in the
buffer, and a token of more than 4,096
+/// chars always does. For example, a 5,000-char literal in an IN list query
of 8M chars allocates 80 MB. Doubling the
+/// buffer would avoid that, but a literal of n chars would then allocate 2 to
3 times as much, and such literals are
+/// the reason for this class.
+///
+/// Not thread-safe, like the generated stream.
+public class GrowingCharStream extends SimpleCharStream {
+ // Fits the whole SQL, so the buffer grows at most once
+ private final int _maxBufferSize;
+
+ private GrowingCharStream(String sql) {
+ super(new StringReader(sql), 1, 1);
+ _maxBufferSize = sql.length() + 1;
+ }
+
+ /// Returns a parser that reads the SQL through a [GrowingCharStream].
+ public static SqlParserImpl newParser(String sql) {
+ GrowingCharStream charStream = new GrowingCharStream(sql);
+ SqlParserImpl sqlParser = new SqlParserImpl(new
SqlParserImplTokenManager(charStream));
+ // The parser also uses the stream directly, e.g. to set the tab size
+ sqlParser.jj_input_stream = charStream;
+ return sqlParser;
+ }
+
+ /// Same as [SimpleCharStream#ExpandBuff], except for the size of the new
buffer.
+ @Override
+ protected void ExpandBuff(boolean wrapAround) {
+ int newBufsize = Math.max(bufsize + 2048, _maxBufferSize);
+ char[] newBuffer = new char[newBufsize];
+ int[] newBufline = new int[newBufsize];
+ int[] newBufcolumn = new int[newBufsize];
+ if (wrapAround) {
+ System.arraycopy(buffer, tokenBegin, newBuffer, 0, bufsize - tokenBegin);
+ System.arraycopy(buffer, 0, newBuffer, bufsize - tokenBegin, bufpos);
+ System.arraycopy(bufline, tokenBegin, newBufline, 0, bufsize -
tokenBegin);
+ System.arraycopy(bufline, 0, newBufline, bufsize - tokenBegin, bufpos);
+ System.arraycopy(bufcolumn, tokenBegin, newBufcolumn, 0, bufsize -
tokenBegin);
+ System.arraycopy(bufcolumn, 0, newBufcolumn, bufsize - tokenBegin,
bufpos);
+ bufpos += bufsize - tokenBegin;
+ } else {
+ System.arraycopy(buffer, tokenBegin, newBuffer, 0, bufsize - tokenBegin);
+ System.arraycopy(bufline, tokenBegin, newBufline, 0, bufsize -
tokenBegin);
+ System.arraycopy(bufcolumn, tokenBegin, newBufcolumn, 0, bufsize -
tokenBegin);
+ bufpos -= tokenBegin;
+ }
+ maxNextCharInd = bufpos;
+ buffer = newBuffer;
+ bufline = newBufline;
+ bufcolumn = newBufcolumn;
+ bufsize = newBufsize;
+ available = bufsize;
+ tokenBegin = 0;
+ }
+}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java
index 552aeea604c..2e3d93f30d5 100644
---
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java
@@ -18,7 +18,6 @@
*/
package org.apache.pinot.sql.parsers;
-import java.io.StringReader;
import java.util.List;
import org.apache.calcite.sql.SqlBinaryStringLiteral;
import org.apache.calcite.sql.SqlCall;
@@ -32,8 +31,10 @@ import org.apache.pinot.common.request.Expression;
import org.apache.pinot.common.request.Function;
import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
import org.apache.pinot.sql.parsers.parser.SqlPhysicalExplain;
import org.apache.pinot.sql.parsers.parser.SqlPinotCreateMaterializedView;
+import org.testng.SkipException;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
@@ -171,8 +172,8 @@ public class CalciteSqlParserTest {
@Test
public void testPostgreSqlByteaLiteralUnderNodeWithoutSetOperandIsRejected()
throws Exception {
- SqlNode query = CalciteSqlParser.newSqlParser(new StringReader("SELECT a
FROM t")).parseSqlStmtEof();
- SqlNode bytea = CalciteSqlParser.newSqlParser(new
StringReader("'\\x01'::bytea")).parseSqlExpressionEof();
+ SqlNode query = CalciteSqlParser.newSqlParser("SELECT a FROM
t").parseSqlStmtEof();
+ SqlNode bytea =
CalciteSqlParser.newSqlParser("'\\x01'::bytea").parseSqlExpressionEof();
SqlOrderBy orderBy = new SqlOrderBy(SqlParserPos.ZERO, query,
SqlNodeList.EMPTY, bytea, null);
SqlCompilationException e =
expectThrows(SqlCompilationException.class, () ->
PostgreSqlCastRewriter.rewrite(orderBy));
@@ -183,7 +184,7 @@ public class CalciteSqlParserTest {
@Test
public void testPostgreSqlByteaLiteralInImmutableNodeListIsRejected()
throws Exception {
- SqlNode bytea = CalciteSqlParser.newSqlParser(new
StringReader("'\\x01'::bytea")).parseSqlExpressionEof();
+ SqlNode bytea =
CalciteSqlParser.newSqlParser("'\\x01'::bytea").parseSqlExpressionEof();
SqlNodeList immutable = SqlNodeList.of(SqlParserPos.ZERO, List.of(bytea));
SqlCompilationException e =
expectThrows(SqlCompilationException.class, () ->
PostgreSqlCastRewriter.rewrite(immutable));
@@ -199,17 +200,17 @@ public class CalciteSqlParserTest {
throws Exception {
String sql = "SELECT a FROM t WHERE b = " + BYTEA_CONSTANT;
SqlSelect select = (SqlSelect) PostgreSqlCastRewriter.rewrite(
- CalciteSqlParser.newSqlParser(new
StringReader(sql)).parseSqlStmtEof());
+ CalciteSqlParser.newSqlParser(sql).parseSqlStmtEof());
assertSpansByteaConstant(((SqlCall) select.getWhere()).operand(1), sql);
sql = "SELECT a FROM t WHERE c LIKE 'x' AND b = " + BYTEA_CONSTANT;
select = (SqlSelect) PostgreSqlCastRewriter.rewrite(
- CalciteSqlParser.newSqlParser(new
StringReader(sql)).parseSqlStmtEof());
+ CalciteSqlParser.newSqlParser(sql).parseSqlStmtEof());
assertSpansByteaConstant(((SqlCall) ((SqlCall)
select.getWhere()).operand(1)).operand(1), sql);
sql = "a || " + BYTEA_CONSTANT;
SqlCall concat = (SqlCall) PostgreSqlCastRewriter.rewrite(
- CalciteSqlParser.newSqlParser(new
StringReader(sql)).parseSqlExpressionEof());
+ CalciteSqlParser.newSqlParser(sql).parseSqlExpressionEof());
assertSpansByteaConstant(concat.operand(1), sql);
}
@@ -519,6 +520,48 @@ public class CalciteSqlParserTest {
assertQualifyRejected("EXPLAIN PLAN FOR SELECT city FROM myTable QUALIFY
city > 'a'");
}
+ /// A string literal of millions of chars, e.g. a serialized IdSet, must
parse in linear time and memory. The
+ /// generated JavaCC char stream grows its buffer by 2,048 chars at a time
while it reads a token, so with it, parsing
+ /// a literal of 2M chars allocates over 10 GB.
+ @Test
+ public void testParsesLongStringLiteralInLinearMemory() {
+ boolean threadMemoryMeasurementEnabled =
ThreadResourceUsageProvider.isThreadMemoryMeasurementEnabled();
+ ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(true);
+ try {
+ if (!ThreadResourceUsageProvider.isThreadMemoryMeasurementEnabled()) {
+ throw new SkipException("Thread memory measurement is not supported");
+ }
+ // Loads the parser classes before the allocation is measured
+ CalciteSqlParser.compileToPinotQuery("SELECT COUNT(*) FROM myTable WHERE
IN_ID_SET(col, 'AA==') = 1");
+ String literal = "A".repeat(2 * 1024 * 1024);
+ String sql = "SELECT COUNT(*) FROM myTable WHERE IN_ID_SET(col, '" +
literal + "') = 1";
+ long allocatedBytesBefore =
ThreadResourceUsageProvider.getCurrentThreadAllocatedBytes();
+ PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+ long allocatedBytes =
ThreadResourceUsageProvider.getCurrentThreadAllocatedBytes() -
allocatedBytesBefore;
+ Function inIdSet =
pinotQuery.getFilterExpression().getFunctionCall().getOperands().get(0).getFunctionCall();
+ assertEquals(inIdSet.getOperands().get(1).getLiteral().getStringValue(),
literal);
+ // The char buffer takes 10 bytes per char, and the literal is copied a
few times
+ assertTrue(allocatedBytes < 50L * sql.length(),
+ "Allocated " + allocatedBytes + " bytes to parse " + sql.length() +
" chars");
+ } finally {
+
ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(threadMemoryMeasurementEnabled);
+ }
+ }
+
+ /// The parser reads the SQL through a custom char stream, which must also
handle an empty SQL.
+ @Test
+ public void testRejectsEmptyQuery() {
+ // The second query reaches the parser as a comment
+ for (String sql : List.of("", "/* comment */")) {
+ SqlCompilationException e =
+ expectThrows(SqlCompilationException.class, () ->
CalciteSqlParser.compileToPinotQuery(sql));
+ assertTrue(e.getMessage().contains("SqlNode with executable statement
not found"), e.getMessage());
+ }
+ SqlCompilationException e =
+ expectThrows(SqlCompilationException.class, () ->
CalciteSqlParser.compileToExpression(""));
+ assertTrue(e.getCause().getMessage().contains("Encountered \"<EOF>\""),
e.getCause().getMessage());
+ }
+
private void assertQualifyRejected(String query) {
SqlCompilationException e =
expectThrows(SqlCompilationException.class, () ->
CalciteSqlParser.compileToPinotQuery(query));
diff --git
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/parser/GrowingCharStreamTest.java
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/parser/GrowingCharStreamTest.java
new file mode 100644
index 00000000000..09f58966136
--- /dev/null
+++
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/parser/GrowingCharStreamTest.java
@@ -0,0 +1,138 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.sql.parsers.parser;
+
+import java.io.StringReader;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// Compares the tokens read through a [GrowingCharStream] with the tokens
read through the generated
+/// [SimpleCharStream], including their line and column positions.
+public class GrowingCharStreamTest {
+ private static final int[] TOKEN_LENGTHS = {1, 2047, 2048, 2049, 4095, 4096,
4097, 6000, 8193, 50_000};
+ // Moves the start of the long token across the 4,096-char initial buffer,
which then wraps around
+ private static final int[] PREFIX_LENGTHS = {0, 1, 2047, 2049, 3000, 4090,
4096, 5000, 10_000};
+
+ @Test
+ public void testReadsTheSameTokensAsTheGeneratedStream() {
+ Random random = new Random(42);
+ for (int tokenLength : TOKEN_LENGTHS) {
+ for (int prefixLength : PREFIX_LENGTHS) {
+ String prefix = "SELECT " + shortTokens(prefixLength);
+ String suffix = "\nFROM myTable\n WHERE b = 1";
+ assertSameTokens(prefix + "'" + text(random, tokenLength) + "'" +
suffix, tokenLength);
+ // A quoted identifier cannot hold a line break
+ assertSameTokens(prefix + "\"" + text(random,
tokenLength).replace('\n', ' ') + "\"" + suffix, tokenLength);
+ assertSameTokens(prefix + "1" + "0".repeat(tokenLength) + suffix,
tokenLength);
+ assertSameTokens(prefix + "/* " + text(random, tokenLength) + " */ a"
+ suffix, tokenLength);
+ // Two long tokens, the second one read into a grown buffer
+ assertSameTokens(prefix + "'" + text(random, tokenLength) + "', " +
shortTokens(prefixLength) + "'"
+ + text(random, 2 * tokenLength) + "'" + suffix, 2 * tokenLength);
+ }
+ }
+ }
+
+ @Test
+ public void testGrowsTheBufferOnlyForLongTokens() {
+ // Short tokens, e.g. of a long IN list, keep the initial buffer
+ String inListSql = "SELECT * FROM myTable WHERE a IN (" +
IntStream.range(0, 10_000)
+ .mapToObj(Integer::toString)
+ .collect(Collectors.joining(", ")) + ")";
+ SqlParserImpl sqlParser = GrowingCharStream.newParser(inListSql);
+ getTokens(sqlParser.token_source);
+ assertEquals(sqlParser.jj_input_stream.bufsize, 4096);
+
+ // A long token grows the buffer once, to fit the whole SQL
+ String longLiteralSql = "SELECT * FROM myTable WHERE a = '" +
"A".repeat(100_000) + "'";
+ sqlParser = GrowingCharStream.newParser(longLiteralSql);
+ getTokens(sqlParser.token_source);
+ assertEquals(sqlParser.jj_input_stream.bufsize, longLiteralSql.length() +
1);
+ }
+
+ /// Returns short tokens of about the given length in total.
+ private static String shortTokens(int length) {
+ StringBuilder stringBuilder = new StringBuilder();
+ for (int i = 0; stringBuilder.length() < length; i++) {
+ stringBuilder.append("c").append(i).append(i % 10 == 0 ? ",\n" : ", ");
+ }
+ return stringBuilder.toString();
+ }
+
+ /// Returns text without quotes or asterisks, but with line breaks, so that
a string literal or a comment of the text
+ /// spans several lines.
+ private static String text(Random random, int length) {
+ char[] chars = new char[length];
+ for (int i = 0; i < length; i++) {
+ int value = random.nextInt(64);
+ chars[i] = value == 0 ? '\n' : value == 1 ? '\t' : (char) ('A' + value %
26);
+ }
+ return new String(chars);
+ }
+
+ private static void assertSameTokens(String sql, int longTokenLength) {
+ List<Token> expectedTokens = getTokens(new SqlParserImpl(new
StringReader(sql)).token_source);
+ List<Token> actualTokens =
getTokens(GrowingCharStream.newParser(sql).token_source);
+ assertEquals(actualTokens.size(), expectedTokens.size());
+ int maxImageLength = 0;
+ for (int i = 0; i < expectedTokens.size(); i++) {
+ // Compare one token at a time, so that a failure does not print the
whole SQL
+ String actualToken = describe(actualTokens.get(i));
+ String expectedToken = describe(expectedTokens.get(i));
+ if (!actualToken.equals(expectedToken)) {
+ assertEquals(abbreviate(actualToken), abbreviate(expectedToken),
+ "Token " + i + " of SQL of length " + sql.length());
+ }
+ maxImageLength = Math.max(maxImageLength,
expectedTokens.get(i).image.length());
+ }
+ // Makes sure that the long text is read as one token
+ assertTrue(maxImageLength >= longTokenLength, "Longest token: " +
maxImageLength + " chars");
+ }
+
+ /// Returns the tokens and special tokens (e.g. comments) in the lexical
state that the SQL parser uses.
+ private static List<Token> getTokens(SqlParserImplTokenManager tokenManager)
{
+ tokenManager.SwitchTo(SqlParserImplConstants.DQID);
+ List<Token> tokens = new ArrayList<>();
+ Token token;
+ do {
+ token = tokenManager.getNextToken();
+ for (Token specialToken = token.specialToken; specialToken != null;
specialToken = specialToken.specialToken) {
+ tokens.add(specialToken);
+ }
+ tokens.add(token);
+ } while (token.kind != SqlParserImplConstants.EOF);
+ return tokens;
+ }
+
+ private static String describe(Token token) {
+ return token.kind + "@" + token.beginLine + ":" + token.beginColumn + "-"
+ token.endLine + ":" + token.endColumn
+ + ":" + token.image;
+ }
+
+ private static String abbreviate(String token) {
+ return token.length() <= 200 ? token : token.substring(0, 100) + "..." +
token.substring(token.length() - 100);
+ }
+}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunction.java
index ac7594f399a..8224909537a 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunction.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.core.operator.transform.function;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import java.io.IOException;
import java.util.List;
@@ -27,6 +28,7 @@ import org.apache.pinot.common.function.TransformFunctionType;
import org.apache.pinot.core.operator.ColumnContext;
import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.transform.TransformResultMetadata;
+import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.core.query.utils.idset.IdSet;
import org.apache.pinot.core.query.utils.idset.IdSets;
import org.apache.pinot.spi.data.FieldSpec.DataType;
@@ -50,8 +52,25 @@ public class InIdSetTransformFunction extends
BaseTransformFunction {
return TransformFunctionType.IN_ID_SET.getName();
}
+ @Override
+ public void init(List<TransformFunction> arguments, Map<String,
ColumnContext> columnContextMap,
+ QueryContext queryContext) {
+ doInit(arguments, columnContextMap, queryContext);
+ // Sets the flag like init(arguments, columnContextMap,
nullHandlingEnabled), which would call the 2-argument init
+ // and deserialize an IdSet that the segments do not share
+ _nullHandlingEnabled = queryContext.isNullHandlingEnabled();
+ }
+
@Override
public void init(List<TransformFunction> arguments, Map<String,
ColumnContext> columnContextMap) {
+ doInit(arguments, columnContextMap, null);
+ }
+
+ /// Initializes the function: checks the arguments and deserializes the
IdSet. Without a query context, every instance
+ /// deserializes its own IdSet. With one, the segments that the query
context covers share one deserialized IdSet,
+ /// which can hold millions of ids.
+ private void doInit(List<TransformFunction> arguments, Map<String,
ColumnContext> columnContextMap,
+ @Nullable QueryContext queryContext) {
super.init(arguments, columnContextMap);
Preconditions.checkArgument(arguments.size() == 2,
"2 arguments are required for IN_ID_SET transform function:
expression, base64 encoded IdSet");
@@ -61,12 +80,25 @@ public class InIdSetTransformFunction extends
BaseTransformFunction {
"Second argument for IN_ID_SET transform function must be a literal
string of the base64 encoded IdSet");
_transformFunction = arguments.get(0);
+ String serializedIdSet = ((LiteralTransformFunction)
arguments.get(1)).getStringLiteral();
+ // The IdSet is only read after this, which is safe from the threads of
several segments
+ _idSet = queryContext != null
+ ? queryContext.getOrComputeSharedValue(IdSet.class, serializedIdSet,
InIdSetTransformFunction::deserialize)
+ : deserialize(serializedIdSet);
+ IdSets.validateValueType(_idSet,
_transformFunction.getResultMetadata().getDataType().getStoredType());
+ }
+
+ private static IdSet deserialize(String serializedIdSet) {
try {
- _idSet = IdSets.fromBase64String(((LiteralTransformFunction)
arguments.get(1)).getStringLiteral());
+ return IdSets.fromBase64String(serializedIdSet);
} catch (IOException | RuntimeException e) {
throw new IllegalArgumentException("Caught exception while deserializing
IdSet: " + e, e);
}
- IdSets.validateValueType(_idSet,
_transformFunction.getResultMetadata().getDataType().getStoredType());
+ }
+
+ @VisibleForTesting
+ IdSet getIdSet() {
+ return _idSet;
}
@Override
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunction.java
index 908221b0d76..0e9daf956ff 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunction.java
@@ -25,6 +25,7 @@ import javax.annotation.Nullable;
import org.apache.pinot.core.operator.ColumnContext;
import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.transform.TransformResultMetadata;
+import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.segment.spi.index.reader.Dictionary;
import org.roaringbitmap.RoaringBitmap;
@@ -55,6 +56,17 @@ public interface TransformFunction {
init(arguments, columnContextMap);
}
+ /// Initializes the transform function for a query. A transform function
that has state to share across the segments
+ /// of the query can keep it in the query context.
+ ///
+ /// @param arguments Arguments for the transform function
+ /// @param columnContextMap Map from column name to context
+ /// @param queryContext Context of the query
+ default void init(List<TransformFunction> arguments, Map<String,
ColumnContext> columnContextMap,
+ QueryContext queryContext) {
+ init(arguments, columnContextMap, queryContext.isNullHandlingEnabled());
+ }
+
/// Returns the metadata for the result of the transform function.
///
/// @return Transform result metadata
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
index c4647f9a5bc..8b8842374d5 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
@@ -370,7 +370,7 @@ public class TransformFunctionFactory {
}
try {
- transformFunction.init(transformFunctionArguments, columnContextMap,
queryContext.isNullHandlingEnabled());
+ transformFunction.init(transformFunctionArguments, columnContextMap,
queryContext);
} catch (Exception e) {
throw new BadQueryRequestException("Caught exception while
initializing transform function: "
+ functionName + ": " + e.getMessage(), e);
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
index 1371fd3f3c7..41c12e6adce 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
@@ -40,7 +40,6 @@ import
org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
import
org.apache.pinot.core.query.aggregation.function.AggregationFunctionFactory;
import org.apache.pinot.core.util.GroupByUtils;
-import org.apache.pinot.core.util.MemoizedClassAssociation;
import org.apache.pinot.segment.spi.datasource.DataSource;
import org.apache.pinot.spi.config.table.FieldConfig;
import org.apache.pinot.spi.data.Schema;
@@ -80,7 +79,9 @@ public class QueryContext {
private final Map<ExpressionContext, ExpressionContext>
_expressionOverrideHints;
private final ExplainMode _explain;
- private final Function<Class<?>, Map<?, ?>> _sharedValues =
MemoizedClassAssociation.of(ConcurrentHashMap::new);
+ // Values that the segments of the query share, by value type. The query
context owns the map, so the values can be
+ // collected with it. A ClassValue would keep each value in its type's cache
until later queries clear it.
+ private final ConcurrentHashMap<Class<?>, ConcurrentHashMap<?, ?>>
_sharedValues = new ConcurrentHashMap<>();
// Pre-calculate the aggregation functions and columns for the query so that
it can be shared across all the segments
private AggregationFunction[] _aggregationFunctions;
@@ -560,7 +561,8 @@ public class QueryContext {
/// @param <V> the value type
/// @return the shared value
public <K, V> V getOrComputeSharedValue(Class<V> type, K key, Function<K, V>
mapper) {
- return ((ConcurrentHashMap<K, V>)
_sharedValues.apply(type)).computeIfAbsent(key, mapper);
+ return ((ConcurrentHashMap<K, V>) _sharedValues.computeIfAbsent(type, k ->
new ConcurrentHashMap<>()))
+ .computeIfAbsent(key, mapper);
}
public int getEffectiveSegmentGroupTrimSize() {
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/util/MemoizedClassAssociation.java
b/pinot-core/src/main/java/org/apache/pinot/core/util/MemoizedClassAssociation.java
deleted file mode 100644
index 6e066cbae74..00000000000
---
a/pinot-core/src/main/java/org/apache/pinot/core/util/MemoizedClassAssociation.java
+++ /dev/null
@@ -1,42 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-package org.apache.pinot.core.util;
-
-import java.util.function.Function;
-import java.util.function.Supplier;
-
-/// Creates a function which associates generic values with a class and
memoizes this association atomically.
-/// @param <T>
-public class MemoizedClassAssociation<T> extends ClassValue<T> {
-
- public static <T> Function<Class<?>, T> of(Supplier<T> supplier) {
- return new MemoizedClassAssociation<>(supplier)::get;
- }
-
- private final Supplier<T> _supplier;
-
- private MemoizedClassAssociation(Supplier<T> supplier) {
- _supplier = supplier;
- }
-
- @Override
- protected T computeValue(Class<?> type) {
- return _supplier.get();
- }
-}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunctionTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunctionTest.java
new file mode 100644
index 00000000000..d1ce67db23d
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunctionTest.java
@@ -0,0 +1,98 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.operator.transform.function;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.RequestContextUtils;
+import org.apache.pinot.core.operator.ColumnContext;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.core.query.utils.idset.IdSet;
+import org.apache.pinot.core.query.utils.idset.IdSets;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
+
+
+/// Tests the IN_ID_SET transform function on a segment; `IdSetQueriesTest`
runs it in queries.
+public class InIdSetTransformFunctionTest extends BaseTransformFunctionTest {
+
+ @Test
+ public void testSharesTheIdSetAcrossTheSegmentsOfAQuery()
+ throws IOException {
+ IdSet idSet = createIdSet(0, NUM_ROWS / 2);
+ ExpressionContext expression = createExpression(idSet);
+ QueryContext queryContext = createQueryContext();
+
+ // Each segment of a query builds its own transform function, but they
share one deserialized IdSet
+ InIdSetTransformFunction segment1Function =
createTransformFunction(expression, queryContext);
+ InIdSetTransformFunction segment2Function =
createTransformFunction(expression, queryContext);
+ assertNotSame(segment2Function, segment1Function);
+ assertSame(segment2Function.getIdSet(), segment1Function.getIdSet());
+ assertResults(segment2Function, idSet);
+
+ // Another IdSet of the same query is deserialized on its own
+ IdSet otherIdSet = createIdSet(NUM_ROWS / 2, NUM_ROWS);
+ InIdSetTransformFunction otherIdSetFunction =
createTransformFunction(createExpression(otherIdSet), queryContext);
+ assertNotSame(otherIdSetFunction.getIdSet(), segment1Function.getIdSet());
+ assertResults(otherIdSetFunction, otherIdSet);
+
+ // Another query deserializes its own IdSet
+ InIdSetTransformFunction otherQueryFunction =
createTransformFunction(expression, createQueryContext());
+ assertNotSame(otherQueryFunction.getIdSet(), segment1Function.getIdSet());
+ }
+
+ private IdSet createIdSet(int fromRow, int toRow) {
+ IdSet idSet = IdSets.create(DataType.INT);
+ for (int i = fromRow; i < toRow; i++) {
+ idSet.add(_intSVValues[i]);
+ }
+ return idSet;
+ }
+
+ private static ExpressionContext createExpression(IdSet idSet)
+ throws IOException {
+ return RequestContextUtils.getExpression(
+ String.format("inIdSet(%s, '%s')", INT_SV_COLUMN,
idSet.toBase64String()));
+ }
+
+ private static QueryContext createQueryContext() {
+ return QueryContextConverterUtils.getQueryContext("SELECT * FROM
testTable");
+ }
+
+ private InIdSetTransformFunction createTransformFunction(ExpressionContext
expression, QueryContext queryContext) {
+ Map<String, ColumnContext> columnContextMap = new HashMap<>();
+ _dataSourceMap.forEach(
+ (column, dataSource) -> columnContextMap.put(column,
ColumnContext.fromDataSource(dataSource)));
+ return (InIdSetTransformFunction) TransformFunctionFactory.get(expression,
columnContextMap, queryContext);
+ }
+
+ private void assertResults(InIdSetTransformFunction transformFunction, IdSet
idSet) {
+ int[] results = transformFunction.transformToIntValuesSV(_projectionBlock);
+ for (int i = 0; i < NUM_ROWS; i++) {
+ assertEquals(results[i], idSet.contains(_intSVValues[i]) ? 1 : 0);
+ }
+ }
+}
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java
index 1c4b2e43af3..e10ae92a223 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/IdSetIntegrationTest.java
@@ -53,6 +53,8 @@ public class IdSetIntegrationTest extends
CustomDataQueryClusterIntegrationTest
private static final String SET_FILTER = ID + " < " + NUM_SET_VALUES;
// STRING and BYTES sets are Bloom filters; these params keep their literals
short
private static final String BLOOM_FILTER_PARAMS =
"expectedInsertions=1000;fpp=0.0001";
+ // A Bloom filter sized for 10M ids, which serializes to a literal of about
16 MB
+ private static final String LARGE_BLOOM_FILTER_PARAMS =
"expectedInsertions=10000000;fpp=0.01";
@Override
protected long getCountStarResult() {
@@ -108,6 +110,21 @@ public class IdSetIntegrationTest extends
CustomDataQueryClusterIntegrationTest
}
}
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testLargeIdSet(boolean useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ String idSet = getSingleValue(postQuery(
+ "SELECT IDSET(" + STR_COL + ", '" + LARGE_BLOOM_FILTER_PARAMS + "')
FROM " + getTableName() + " WHERE "
+ + SET_FILTER)).asText();
+ assertTrue(idSet.length() > 15_000_000, "IdSet length: " + idSet.length());
+ String inQuery =
+ "SELECT COUNT(*) FROM " + getTableName() + " WHERE IN_ID_SET(" +
STR_COL + ", '" + idSet + "') = 1";
+ // Over HTTP, because the broker's gRPC endpoint accepts requests of up to
4 MB
+ assertEquals(getSingleValue(queryBrokerHttpEndpoint(inQuery)).asLong(),
+ getExpectedCount(STR_COL, IdSets.fromBase64String(idSet)));
+ }
+
@Test(dataProvider = "useBothQueryEngines")
public void testMismatchedValueType(boolean useMultiStageQueryEngine)
throws Exception {
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/DispatchClient.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/DispatchClient.java
index 489ab1a3595..05c5f0905b0 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/DispatchClient.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/DispatchClient.java
@@ -37,6 +37,7 @@ import org.apache.pinot.query.grpc.GrpcKeepAliveConfig;
import org.apache.pinot.query.routing.QueryServerInstance;
import
org.apache.pinot.query.service.dispatch.streaming.StreamingDispatchObserver;
import org.apache.pinot.query.service.dispatch.streaming.StreamingQuerySession;
+import org.apache.pinot.spi.utils.CommonConstants;
/// Dispatches a query plan to given a [QueryServerInstance]. Each
[DispatchClient] has its own gRPC Channel
@@ -48,11 +49,6 @@ class DispatchClient {
/// Shared buffer allocator configured to prefer direct (off-heap) buffers
for better performance.
/// Using a static allocator allows for better memory pooling across all
DispatchClient instances.
private static final PooledByteBufAllocator BUF_ALLOCATOR = new
PooledByteBufAllocator(true);
- /// Max size of an inbound message on the broker's dispatch channel. The
default gRPC limit is 4 MB, but a single
- /// `SubmitWithStream` `OpChainComplete` can carry a large
[Worker.MultiStageStatsTree] (wide/deep
- /// plans, large STRING stats), and exceeding the limit fails the whole
query with `RESOURCE_EXHAUSTED` even
- /// though the query results are already in hand. Matches the server's
inbound limit (`QueryServer`).
- private static final int MAX_INBOUND_MESSAGE_SIZE = 64 * 1024 * 1024;
private final ManagedChannel _channel;
private final PinotQueryWorkerGrpc.PinotQueryWorkerStub _dispatchStub;
@@ -67,11 +63,21 @@ class DispatchClient {
DispatchClient(String host, int port, @Nullable TlsConfig tlsConfig,
@Nullable SslContext sslContext,
GrpcKeepAliveConfig keepAliveConfig) {
+ this(host, port, tlsConfig, sslContext, keepAliveConfig,
+
CommonConstants.MultiStageQueryRunner.DEFAULT_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES);
+ }
+
+ /// @param maxInboundMessageSizeBytes Max size of an inbound message on the
broker's dispatch channel. The default
+ /// gRPC limit is 4 MB, but a single `SubmitWithStream`
`OpChainComplete` can carry a large
+ /// [Worker.MultiStageStatsTree] (wide/deep plans, large STRING stats),
and the plans servers return for
+ /// `EXPLAIN` carry the query's literals. Exceeding the limit fails the
whole query with `RESOURCE_EXHAUSTED`.
+ DispatchClient(String host, int port, @Nullable TlsConfig tlsConfig,
@Nullable SslContext sslContext,
+ GrpcKeepAliveConfig keepAliveConfig, int maxInboundMessageSizeBytes) {
// Always use NettyChannelBuilder to allow setting Netty-specific channel
options like the buffer allocator.
// This ensures we can explicitly configure direct (off-heap) buffers for
better performance.
NettyChannelBuilder channelBuilder = NettyChannelBuilder.forAddress(host,
port)
.withOption(ChannelOption.ALLOCATOR, BUF_ALLOCATOR)
- .maxInboundMessageSize(MAX_INBOUND_MESSAGE_SIZE);
+ .maxInboundMessageSize(maxInboundMessageSizeBytes);
if (tlsConfig == null) {
channelBuilder.usePlaintext();
} else {
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java
index c6524bde0a0..b3a1e605da6 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java
@@ -129,6 +129,7 @@ public class QueryDispatcher implements
PinotClusterConfigChangeListener {
@Nullable
private final SslContext _clientGrpcSslContext;
private final GrpcKeepAliveConfig _keepAliveConfig;
+ private final int _dispatchMaxInboundMessageSizeBytes;
// maps broker-generated query id to the set of servers that the query was
dispatched to
private final Map<Long, Set<QueryServerInstance>> _serversByQuery;
private final FailureDetector _failureDetector;
@@ -147,7 +148,8 @@ public class QueryDispatcher implements
PinotClusterConfigChangeListener {
boolean enableCancellation, Duration cancelTimeout) {
this(mailboxService, failureDetector, tlsConfig, enableCancellation,
cancelTimeout,
GrpcKeepAliveConfig.DISABLED, false,
CommonConstants.Broker.DEFAULT_STREAM_STATS_DRAIN_MS,
- CommonConstants.Broker.DEFAULT_MSE_ENABLE_PROTO_SEGMENT_LIST);
+ CommonConstants.Broker.DEFAULT_MSE_ENABLE_PROTO_SEGMENT_LIST,
+
CommonConstants.MultiStageQueryRunner.DEFAULT_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES);
}
/// Overload that accepts gRPC keep-alive settings for broker dispatch
channels. A non-positive `keepAliveTimeMs`
@@ -165,14 +167,29 @@ public class QueryDispatcher implements
PinotClusterConfigChangeListener {
boolean enableCancellation, Duration cancelTimeout, int keepAliveTimeMs,
int keepAliveTimeoutMs,
boolean keepAliveWithoutCalls, boolean streamStatsDefault, long
statsDrainMs,
boolean enableProtoSegmentList) {
+ this(mailboxService, failureDetector, tlsConfig, enableCancellation,
cancelTimeout, keepAliveTimeMs,
+ keepAliveTimeoutMs, keepAliveWithoutCalls, streamStatsDefault,
statsDrainMs, enableProtoSegmentList,
+
CommonConstants.MultiStageQueryRunner.DEFAULT_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES);
+ }
+
+ /// Overload that also takes the max size of a message the dispatch channels
accept from the servers.
+ public QueryDispatcher(MailboxService mailboxService, FailureDetector
failureDetector, @Nullable TlsConfig tlsConfig,
+ boolean enableCancellation, Duration cancelTimeout, int keepAliveTimeMs,
int keepAliveTimeoutMs,
+ boolean keepAliveWithoutCalls, boolean streamStatsDefault, long
statsDrainMs,
+ boolean enableProtoSegmentList, int dispatchMaxInboundMessageSizeBytes) {
this(mailboxService, failureDetector, tlsConfig, enableCancellation,
cancelTimeout,
new GrpcKeepAliveConfig(keepAliveTimeMs, keepAliveTimeoutMs,
keepAliveWithoutCalls),
- streamStatsDefault, statsDrainMs, enableProtoSegmentList);
+ streamStatsDefault, statsDrainMs, enableProtoSegmentList,
dispatchMaxInboundMessageSizeBytes);
}
private QueryDispatcher(MailboxService mailboxService, FailureDetector
failureDetector, @Nullable TlsConfig tlsConfig,
boolean enableCancellation, Duration cancelTimeout, GrpcKeepAliveConfig
keepAliveConfig,
- boolean streamStatsDefault, long statsDrainMs, boolean
enableProtoSegmentList) {
+ boolean streamStatsDefault, long statsDrainMs, boolean
enableProtoSegmentList,
+ int dispatchMaxInboundMessageSizeBytes) {
+ // Checked here, because the dispatch channels are created only when
queries are dispatched
+ Preconditions.checkArgument(dispatchMaxInboundMessageSizeBytes > 0, "%s
must be positive, got: %s",
+
CommonConstants.MultiStageQueryRunner.KEY_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES,
+ dispatchMaxInboundMessageSizeBytes);
_cancelTimeout = cancelTimeout;
_statsDrainMs = statsDrainMs;
_mailboxService = mailboxService;
@@ -181,6 +198,7 @@ public class QueryDispatcher implements
PinotClusterConfigChangeListener {
_tlsConfig = tlsConfig;
_clientGrpcSslContext = initClientSslContext(tlsConfig);
_keepAliveConfig = keepAliveConfig;
+ _dispatchMaxInboundMessageSizeBytes = dispatchMaxInboundMessageSizeBytes;
_failureDetector = failureDetector;
_streamStatsDefault = streamStatsDefault;
_enableProtoSegmentList = enableProtoSegmentList;
@@ -749,6 +767,11 @@ public class QueryDispatcher implements
PinotClusterConfigChangeListener {
return _enableProtoSegmentList;
}
+ @VisibleForTesting
+ public int getDispatchMaxInboundMessageSizeBytes() {
+ return _dispatchMaxInboundMessageSizeBytes;
+ }
+
/// Builds the request for one server: the plans of the stages it takes part
in, with only its own workers'
/// metadata. The leaf-stage segment lists are encoded here, once per
worker, rather than at plan time.
private static Worker.QueryRequest createRequest(QueryServerInstance
serverInstance,
@@ -864,11 +887,13 @@ public class QueryDispatcher implements
PinotClusterConfigChangeListener {
return true;
}
- private DispatchClient getOrCreateDispatchClient(QueryServerInstance
queryServerInstance) {
+ @VisibleForTesting
+ DispatchClient getOrCreateDispatchClient(QueryServerInstance
queryServerInstance) {
String hostname = queryServerInstance.getHostname();
int port = queryServerInstance.getQueryServicePort();
return _dispatchClientMap.computeIfAbsent(toHostnamePortKey(hostname,
port),
- k -> new DispatchClient(hostname, port, _tlsConfig,
_clientGrpcSslContext, _keepAliveConfig));
+ k -> new DispatchClient(hostname, port, _tlsConfig,
_clientGrpcSslContext, _keepAliveConfig,
+ _dispatchMaxInboundMessageSizeBytes));
}
/// Reset the connection backoff for a server. When the GRPC channel enters
a TRANSIENT_FAILURE state from
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
index 42dbb7344cb..bd5bde1ee96 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
@@ -19,6 +19,7 @@
package org.apache.pinot.query.service.server;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
import com.google.protobuf.ByteString;
import io.grpc.Server;
import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
@@ -86,10 +87,6 @@ import org.slf4j.LoggerFactory;
public class QueryServer extends PinotQueryWorkerGrpc.PinotQueryWorkerImplBase
{
private static final Logger LOGGER =
LoggerFactory.getLogger(QueryServer.class);
- // TODO: Inbound messages can get quite large because we send the entire
stage metadata map in each call.
- // See https://github.com/apache/pinot/issues/10331
- private static final int MAX_INBOUND_MESSAGE_SIZE = 64 * 1024 * 1024;
-
private final String _instanceId;
private final int _port;
private final QueryRunner _queryRunner;
@@ -100,6 +97,9 @@ public class QueryServer extends
PinotQueryWorkerGrpc.PinotQueryWorkerImplBase {
private final ThreadAccountant _threadAccountant;
private final int _permitKeepAliveTimeMs;
private final boolean _permitKeepAliveWithoutCalls;
+ // TODO: Inbound messages can get quite large because we send the entire
stage metadata map in each call.
+ // See https://github.com/apache/pinot/issues/10331
+ private final int _maxInboundMessageSizeBytes;
// query submission service is only used for plan submission for now.
// TODO: with complex query submission logic we should allow asynchronous
query submission return instead of
@@ -162,6 +162,12 @@ public class QueryServer extends
PinotQueryWorkerGrpc.PinotQueryWorkerImplBase {
_permitKeepAliveWithoutCalls = serverConf.getProperty(
CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_PERMIT_KEEP_ALIVE_WITHOUT_CALLS,
CommonConstants.MultiStageQueryRunner.DEFAULT_OF_QUERY_SERVER_PERMIT_KEEP_ALIVE_WITHOUT_CALLS);
+ _maxInboundMessageSizeBytes = serverConf.getProperty(
+
CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES,
+
CommonConstants.MultiStageQueryRunner.DEFAULT_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES);
+ Preconditions.checkArgument(_maxInboundMessageSizeBytes > 0, "%s must be
positive, got: %s",
+
CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES,
+ _maxInboundMessageSizeBytes);
ExecutorService baseExecutorService =
ExecutorServiceUtils.create(serverConf,
CommonConstants.Server.MULTISTAGE_SUBMISSION_EXEC_CONFIG_PREFIX,
@@ -193,8 +199,9 @@ public class QueryServer extends
PinotQueryWorkerGrpc.PinotQueryWorkerImplBase {
}
_server = buildGrpcServer(serverBuilder);
LOGGER.info(
- "Initialized QueryServer on port: {} with permitKeepAliveTimeMs:
{}, permitKeepAliveWithoutCalls: {}",
- _port, _permitKeepAliveTimeMs, _permitKeepAliveWithoutCalls);
+ "Initialized QueryServer on port: {} with permitKeepAliveTimeMs:
{}, permitKeepAliveWithoutCalls: {}, "
+ + "maxInboundMessageSizeBytes: {}", _port,
_permitKeepAliveTimeMs, _permitKeepAliveWithoutCalls,
+ _maxInboundMessageSizeBytes);
}
_queryRunner.start();
_server.start();
@@ -213,6 +220,11 @@ public class QueryServer extends
PinotQueryWorkerGrpc.PinotQueryWorkerImplBase {
return _permitKeepAliveWithoutCalls;
}
+ @VisibleForTesting
+ int getMaxInboundMessageSizeBytes() {
+ return _maxInboundMessageSizeBytes;
+ }
+
private Server buildGrpcServer(NettyServerBuilder builder) {
// By using directExecutor, GRPC doesn't need to manage its own thread pool
builder.directExecutor();
@@ -225,7 +237,7 @@ public class QueryServer extends
PinotQueryWorkerGrpc.PinotQueryWorkerImplBase {
builder.permitKeepAliveTime(_permitKeepAliveTimeMs,
TimeUnit.MILLISECONDS);
}
builder.permitKeepAliveWithoutCalls(_permitKeepAliveWithoutCalls);
- return
builder.addService(this).maxInboundMessageSize(MAX_INBOUND_MESSAGE_SIZE).build();
+ return
builder.addService(this).maxInboundMessageSize(_maxInboundMessageSizeBytes).build();
}
public void shutdown() {
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/dispatch/DispatchClientTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/dispatch/DispatchClientTest.java
new file mode 100644
index 00000000000..264eb8528a9
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/dispatch/DispatchClientTest.java
@@ -0,0 +1,138 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.service.dispatch;
+
+import io.grpc.Deadline;
+import io.grpc.Server;
+import io.grpc.Status;
+import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
+import io.grpc.stub.StreamObserver;
+import java.io.IOException;
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.failuredetector.FailureDetector;
+import org.apache.pinot.common.proto.PinotQueryWorkerGrpc;
+import org.apache.pinot.common.proto.Worker;
+import org.apache.pinot.query.mailbox.MailboxService;
+import org.apache.pinot.query.routing.QueryServerInstance;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Tests the max size of the messages the broker's dispatch channels accept
from the servers.
+public class DispatchClientTest {
+ // Larger than the 4 MB gRPC accepts by default
+ private static final String LARGE_VALUE = "x".repeat(5 * 1024 * 1024);
+
+ private Server _server;
+ private QueryServerInstance _serverInstance;
+
+ @BeforeClass
+ public void setUp()
+ throws IOException {
+ // Replies to every submit with a response larger than LARGE_VALUE
+ _server = NettyServerBuilder.forPort(0).addService(new
PinotQueryWorkerGrpc.PinotQueryWorkerImplBase() {
+ @Override
+ public void submit(Worker.QueryRequest request,
StreamObserver<Worker.QueryResponse> responseObserver) {
+
responseObserver.onNext(Worker.QueryResponse.newBuilder().putMetadata("value",
LARGE_VALUE).build());
+ responseObserver.onCompleted();
+ }
+ }).build().start();
+ _serverInstance = new QueryServerInstance("server", "localhost",
_server.getPort(), 0);
+ }
+
+ @AfterClass
+ public void tearDown() {
+ _server.shutdownNow();
+ }
+
+ @Test
+ public void testDefaultMaxInboundMessageSize()
+ throws Exception {
+ QueryDispatcher queryDispatcher =
+ new QueryDispatcher(mock(MailboxService.class),
mock(FailureDetector.class), null, false,
+ Duration.ofSeconds(1));
+ try {
+ AsyncResponse<Worker.QueryResponse> response = submit(queryDispatcher);
+ assertNull(response.getThrowable());
+ assertNotNull(response.getResponse());
+ assertEquals(response.getResponse().getMetadataMap().get("value"),
LARGE_VALUE);
+ } finally {
+ queryDispatcher.shutdown();
+ }
+ }
+
+ @Test
+ public void testConfiguredMaxInboundMessageSize()
+ throws Exception {
+ QueryDispatcher queryDispatcher =
createQueryDispatcher(LARGE_VALUE.length() / 2);
+ try {
+ AsyncResponse<Worker.QueryResponse> response = submit(queryDispatcher);
+ assertNull(response.getResponse());
+ assertEquals(Status.fromThrowable(response.getThrowable()).getCode(),
Status.Code.RESOURCE_EXHAUSTED);
+ } finally {
+ queryDispatcher.shutdown();
+ }
+
+ queryDispatcher = createQueryDispatcher(LARGE_VALUE.length() * 2);
+ try {
+ AsyncResponse<Worker.QueryResponse> response = submit(queryDispatcher);
+ assertNull(response.getThrowable());
+ assertNotNull(response.getResponse());
+ assertEquals(response.getResponse().getMetadataMap().get("value"),
LARGE_VALUE);
+ } finally {
+ queryDispatcher.shutdown();
+ }
+ }
+
+ @Test
+ public void testRejectsNonPositiveMaxInboundMessageSize() {
+ for (int dispatchMaxInboundMessageSizeBytes : new int[]{0, -1}) {
+ IllegalArgumentException exception =
+ expectThrows(IllegalArgumentException.class, () ->
createQueryDispatcher(dispatchMaxInboundMessageSizeBytes));
+ assertTrue(exception.getMessage()
+
.contains(CommonConstants.MultiStageQueryRunner.KEY_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES));
+ }
+ }
+
+ private static QueryDispatcher createQueryDispatcher(int
dispatchMaxInboundMessageSizeBytes) {
+ return new QueryDispatcher(mock(MailboxService.class),
mock(FailureDetector.class), null, false,
+ Duration.ofSeconds(1), 0, 0, false, false,
CommonConstants.Broker.DEFAULT_STREAM_STATS_DRAIN_MS, false,
+ dispatchMaxInboundMessageSizeBytes);
+ }
+
+ private AsyncResponse<Worker.QueryResponse> submit(QueryDispatcher
queryDispatcher)
+ throws Exception {
+ CompletableFuture<AsyncResponse<Worker.QueryResponse>> future = new
CompletableFuture<>();
+ queryDispatcher.getOrCreateDispatchClient(_serverInstance)
+ .submit(Worker.QueryRequest.getDefaultInstance(), _serverInstance,
Deadline.after(10, TimeUnit.SECONDS),
+ future::complete);
+ return future.get(10, TimeUnit.SECONDS);
+ }
+}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java
index c8c177fd51c..0abc060435e 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java
@@ -23,6 +23,8 @@ import com.google.protobuf.ByteString;
import io.grpc.Deadline;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import java.util.HashMap;
import java.util.List;
@@ -77,6 +79,7 @@ import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
public class QueryServerTest extends QueryTestSet {
@@ -148,6 +151,58 @@ public class QueryServerTest extends QueryTestSet {
assertTrue(server.isPermitKeepAliveWithoutCalls());
}
+ @Test
+ public void testMaxInboundMessageSizeDefault() {
+ QueryServer server = new QueryServer(QueryTestUtils.getAvailablePort(),
mock(QueryRunner.class));
+ assertEquals(server.getMaxInboundMessageSizeBytes(), 64 * 1024 * 1024);
+ }
+
+ @Test
+ public void testRejectsNonPositiveMaxInboundMessageSize() {
+ for (int maxInboundMessageSizeBytes : new int[]{0, -1}) {
+ PinotConfiguration config = new PinotConfiguration(
+
Map.of(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES,
+ maxInboundMessageSizeBytes));
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> new QueryServer(config, "serverId",
QueryTestUtils.getAvailablePort(), mock(QueryRunner.class), null,
+ ThreadAccountantUtils.getNoOpAccountant()));
+ assertTrue(exception.getMessage()
+
.contains(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES));
+ }
+ }
+
+ @Test
+ public void testMaxInboundMessageSizeFromConfig()
+ throws Exception {
+ // A request with large literals, e.g. an IdSet, must fit the configured
limit
+ Worker.QueryRequest queryRequest;
+ try (QueryEnvironment.CompiledQuery compiledQuery =
_queryEnvironment.compile("SELECT * FROM a")) {
+ queryRequest =
getQueryRequest(compiledQuery.planQuery(0).getQueryPlan(), 1);
+ }
+ int maxInboundMessageSizeBytes = queryRequest.getSerializedSize() - 1;
+ int port = QueryTestUtils.getAvailablePort();
+ QueryServer server = new QueryServer(new PinotConfiguration(
+
Map.of(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES,
+ maxInboundMessageSizeBytes)), "serverId", port,
mock(QueryRunner.class), null,
+ ThreadAccountantUtils.getNoOpAccountant());
+ assertEquals(server.getMaxInboundMessageSizeBytes(),
maxInboundMessageSizeBytes);
+ server.start();
+ ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost",
port).usePlaintext().build();
+ try {
+ PinotQueryWorkerGrpc.PinotQueryWorkerBlockingStub stub =
+
PinotQueryWorkerGrpc.newBlockingStub(channel).withDeadline(Deadline.after(10,
TimeUnit.SECONDS));
+ StatusRuntimeException exception =
expectThrows(StatusRuntimeException.class, () -> stub.submit(queryRequest));
+ assertEquals(exception.getStatus().getCode(),
Status.Code.RESOURCE_EXHAUSTED);
+ // A smaller request reaches the server
+ Worker.QueryResponse response =
+
stub.submit(Worker.QueryRequest.newBuilder().setMetadata(queryRequest.getMetadata()).build());
+
assertTrue(response.getMetadataMap().containsKey(CommonConstants.Query.Response.ServerResponseStatus.STATUS_OK));
+ } finally {
+ channel.shutdownNow();
+ server.shutdown();
+ }
+ }
+
@Test
public void testException()
throws Exception {
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 70b97dbc3f5..ee6fd925d58 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -2963,6 +2963,20 @@ public class CommonConstants {
public static final String
KEY_OF_QUERY_SERVER_PERMIT_KEEP_ALIVE_WITHOUT_CALLS =
"pinot.query.multistage.query.server.permit.keep.alive.without.calls";
public static final boolean
DEFAULT_OF_QUERY_SERVER_PERMIT_KEEP_ALIVE_WITHOUT_CALLS = false;
+
+ /// Maximum size, in bytes, of a gRPC message the MSE
[org.apache.pinot.query.service.server.QueryServer] accepts.
+ /// Must be positive. A message carries the stage plans a broker
dispatches to the server, literals included, so
+ /// queries with large literals (e.g. an `IN_ID_SET` IdSet of millions of
ids) can need more than the default.
+ public static final String
KEY_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES =
+ "pinot.query.multistage.query.server.max.inbound.message.size.bytes";
+ public static final int
DEFAULT_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES = 64 * 1024 * 1024;
+
+ /// Maximum size, in bytes, of a gRPC message the broker dispatch channels
accept from the MSE query servers, such
+ /// as query stats, or the plans the servers return for `EXPLAIN`, which
carry the query's literals. Must be
+ /// positive. See [#KEY_OF_QUERY_SERVER_MAX_INBOUND_MESSAGE_SIZE_BYTES]
for the messages the servers accept.
+ public static final String
KEY_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES =
+
"pinot.query.multistage.dispatch.channel.max.inbound.message.size.bytes";
+ public static final int
DEFAULT_OF_DISPATCH_CHANNEL_MAX_INBOUND_MESSAGE_SIZE_BYTES = 64 * 1024 * 1024;
}
public static class NullValuePlaceHolder {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
