xiangfu0 commented on code in PR #20093:
URL: https://github.com/apache/pinot/pull/20093#discussion_r4416395513
##########
pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QueryScheduler.java:
##########
@@ -197,40 +205,58 @@ protected DataBuffer
processQueryAndSerializeResponse(ServerQueryRequest queryRe
responseMetadata.put(MetadataKey.REQUEST_ID.getName(),
Long.toString(requestId));
responseMetadata.put(MetadataKey.QUERY_ID.getName(), queryId);
responseMetadata.put(MetadataKey.WORKLOAD_NAME.getName(), workloadName);
- DataBuffer response = serializeResponse(queryRequest, instanceResponse);
+ DataBuffer response;
+ boolean responseSizeExceeded = false;
+ try {
+ response = serializeResponse(queryRequest, instanceResponse);
+ } catch (ResponseSizeLimitExceededException e) {
+ // A data table builder stopped because the rows already exceed the
limit. The full size is unknown, so the
+ // meter records the bytes written until the builder stopped.
+
_serverMetrics.addMeteredTableValue(queryRequest.getTableNameWithType(),
ServerMeter.QUERY_RESPONSE_SIZE,
+ e.getSizeBytes());
+ response = responseSizeExceeded(queryRequest, executionContext,
e.getMessage());
Review Comment:
On this path the query log shows the wrong serialization time.
`responseSizeExceeded()` calls `serializeResponse()` again, and
`startNewPhaseTimer` overwrites the `RESPONSE_SERIALIZATION` slot.
`_queryLogger.logQuery` then runs below, so it logs the time to serialize the
small error table, not the time spent building rows. The post-serialization
path logs before it replaces the response. Moving `logQuery` above the
size-limit handling would fix both paths.
##########
pinot-core/src/main/java/org/apache/pinot/core/common/datatable/DataTableBuilderV4.java:
##########
@@ -61,11 +65,26 @@ public void setColumn(int colId, String[] values)
_currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (String value : values) {
- int dictId = _dictionary.computeIntIfAbsent(value, k ->
_dictionary.size());
+ int dictId = getDictId(value);
_variableSizeDataOutputStream.writeInt(dictId);
}
}
+ private int getDictId(String value) {
+ int dictId = _dictionary.getInt(value);
+ if (dictId < 0) {
+ dictId = _dictionary.size();
+ _dictionary.put(value, dictId);
Review Comment:
`getInt` followed by `put` probes the hash table twice for every new string.
In a JMH run building 100k rows with unique strings, this went from 16.0 ms to
18.7 ms (+17%; a first run showed +19%). Since `defaultReturnValue` is -1, `int
dictId = _dictionary.putIfAbsent(value, _dictionary.size())` does one probe and
returns -1 when it inserted. That variant measured 15.6 ms, back at baseline.
Repeated strings didn't change.
##########
pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerResponseSizeTest.java:
##########
@@ -0,0 +1,186 @@
+/**
+ * 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.query.scheduler;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.LongAccumulator;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.DataTableFactory;
+import org.apache.pinot.common.metrics.ServerMeter;
+import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.request.InstanceRequest;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.common.datatable.DataTableBuilder;
+import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
+import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
+import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
+import org.apache.pinot.core.query.executor.QueryExecutor;
+import org.apache.pinot.core.query.request.ServerQueryRequest;
+import
org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
+import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Tests that [QueryScheduler] stops building a response that exceeds
`maxServerResponseSizeBytes`, and replies with
+/// an error instead.
+public class QuerySchedulerResponseSizeTest {
+ private static final PinotConfiguration CONFIG = new PinotConfiguration();
+ private static final String TABLE_NAME = "testTable_OFFLINE";
+ private static final DataSchema DATA_SCHEMA =
+ new DataSchema(new String[]{"value"}, new
ColumnDataType[]{ColumnDataType.LONG});
+ private static final int NUM_ROWS = 1000;
+
+ private ServerMetrics _previousServerMetrics;
+ private ServerMetrics _serverMetrics;
+ private ExecutorService _executorService;
+
+ /// [ServerMetrics#register] only swaps in against the NOOP default, so the
instance that an earlier test class
+ /// registered is cleared first, and restored in [#tearDown].
+ @BeforeClass
+ public void setUp() {
+ _previousServerMetrics = ServerMetrics.get();
+ ServerMetrics.deregister();
+ _serverMetrics = mock(ServerMetrics.class);
+ ServerMetrics.register(_serverMetrics);
+ _executorService = Executors.newSingleThreadExecutor();
+ }
+
+ @AfterClass(alwaysRun = true)
+ public void tearDown() {
+ _executorService.shutdownNow();
+ ServerMetrics.deregister();
+ ServerMetrics.register(_previousServerMetrics);
+ }
+
+ @Test
+ public void testResponseWithinLimit()
+ throws Exception {
+ AtomicInteger numRowsWritten = new AtomicInteger();
+ DataTable response =
processQuery("OPTION(maxServerResponseSizeBytes=100000)", numRowsWritten);
+ assertTrue(response.getExceptions().isEmpty(),
response.getExceptions().toString());
+ assertEquals(response.getNumberOfRows(), NUM_ROWS);
+ }
+
+ @Test
+ public void testStopsBuildingResponseOverLimit()
+ throws Exception {
+ AtomicInteger numRowsWritten = new AtomicInteger();
+ DataTable response =
processQuery("OPTION(maxServerResponseSizeBytes=100)", numRowsWritten);
+ String error =
response.getExceptions().get(QueryErrorCode.QUERY_CANCELLATION.getId());
+ assertNotNull(error, response.getExceptions().toString());
+ assertEquals(error,
+ "Serialized query response size of at least 104 exceeds threshold 100
for requestId 123 from broker unknown");
+ // Each row is 8 bytes, so the builder stops at the 13th row instead of
writing all of them.
+ assertEquals(numRowsWritten.get(), 13);
+ verify(_serverMetrics).addMeteredTableValue(TABLE_NAME,
ServerMeter.LARGE_QUERY_RESPONSE_SIZE_EXCEPTIONS, 1);
+ verify(_serverMetrics).addMeteredTableValue(TABLE_NAME,
ServerMeter.QUERY_RESPONSE_SIZE, 104);
Review Comment:
This only checks for one call with the value 104. A second
`QUERY_RESPONSE_SIZE` call with another value, such as the error table's size,
would still pass. Adding `verify(_serverMetrics,
times(1)).addMeteredTableValue(eq(TABLE_NAME),
eq(ServerMeter.QUERY_RESPONSE_SIZE), anyLong())` here and in
`testRejectsSerializedResponseOverLimit` would lock in the "recorded once"
contract. A builder test that sets a terminate exception and writes more than 1
MB would also cover the new termination check.
##########
pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryExecutionContext.java:
##########
@@ -294,6 +298,14 @@ public void setQueryId(@Nullable String queryId) {
_queryId = queryId;
}
+ public long getMaxResponseSizeBytes() {
Review Comment:
nit: this holds `maxServerResponseSizeBytes`, and the broker also has
`maxQueryResponseSizeBytes`. Since this is a public SPI class,
`get/setMaxServerResponseSizeBytes` would make it clear which limit this is.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
