This is an automated email from the ASF dual-hosted git repository.
gortiz 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 3f6fef49e05 Send v1 DataTable responses to Netty and gRPC as pages
instead of one byte[] (#19943)
3f6fef49e05 is described below
commit 3f6fef49e05dec8414ebecca0373ce0cea14cef6
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Thu Oct 29 11:09:32 2026 +0100
Send v1 DataTable responses to Netty and gRPC as pages instead of one
byte[] (#19943)
DataTable.toDataBuffer() returns the serialized table as a DataBuffer over
the builder's existing pages, so the v1 server response path no longer copies
the response into one full-size byte[] (a G1 humongous object for large
responses). The wire format does not change.
- QueryScheduler gets submitResponse() and DataBuffer-returning helpers.
submit() and the byte[] helpers stay as deprecated defaults; a subclass must
override one of submit()/submitResponse(), checked in the constructor.
- FCFSQueryScheduler, PriorityScheduler, BinaryWorkloadScheduler and
WorkloadScheduler override submitResponse(); SchedulerQueryContext holds a
SettableFuture<DataBuffer>.
- InstanceRequestHandler sends the pages with
Unpooled.wrappedBuffer(ByteBuffer...).
- StreamingResponseUtils wraps the pages with
UnsafeByteOperations.unsafeWrap into a ByteString rope instead of
ByteString.copyFrom(byte[]).
- PageCacheWarmupServerQueryExecutor calls submitResponse().
---
.../apache/pinot/common/datatable/DataTable.java | 11 ++
.../pinot/common/datatable/DataTableImplV4.java | 89 ++++++++-------
.../operator/streaming/StreamingResponseUtils.java | 20 +++-
.../query/scheduler/BinaryWorkloadScheduler.java | 15 +--
.../core/query/scheduler/PriorityScheduler.java | 11 +-
.../pinot/core/query/scheduler/QueryScheduler.java | 119 ++++++++++++++++++---
.../query/scheduler/SchedulerQueryContext.java | 7 +-
.../core/query/scheduler/WorkloadScheduler.java | 9 +-
.../query/scheduler/fcfs/FCFSQueryScheduler.java | 7 +-
.../core/transport/InstanceRequestHandler.java | 34 +++---
.../core/common/datatable/DataTableSerDeTest.java | 23 +++-
.../streaming/StreamingResponseUtilsTest.java | 72 +++++++++++++
.../scheduler/QuerySchedulerResponseApiTest.java | 98 +++++++++++++++++
.../core/transport/InstanceRequestHandlerTest.java | 6 +-
.../pinot/core/transport/QueryRoutingTest.java | 62 ++++++++++-
.../warmup/PageCacheWarmupServerQueryExecutor.java | 5 +-
.../PageCacheWarmupServerQueryExecutorTest.java | 7 +-
17 files changed, 491 insertions(+), 104 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
index 14d916154c1..a329e7d2273 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
@@ -26,6 +26,8 @@ import java.util.Map;
import javax.annotation.Nullable;
import org.apache.pinot.common.CustomObject;
import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
+import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.utils.ByteArray;
import org.roaringbitmap.RoaringBitmap;
@@ -47,6 +49,15 @@ public interface DataTable {
byte[] toBytes()
throws IOException;
+ /// Serializes the table into the same bytes as [#toBytes()], but possibly
split into several buffers. An
+ /// implementation can return its sections without copying them, so a large
table does not need one large array.
+ ///
+ /// The default implementation wraps [#toBytes()].
+ default DataBuffer toDataBuffer()
+ throws IOException {
+ return PinotByteBuffer.wrap(toBytes());
+ }
+
Map<String, String> getMetadata();
DataSchema getDataSchema();
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
index 8b0b58cd4ba..a058d95addb 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
@@ -26,6 +26,7 @@ import java.io.DataOutputStream;
import java.io.IOException;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
import java.util.HashMap;
import java.util.Map;
import javax.annotation.Nullable;
@@ -34,6 +35,7 @@ import org.apache.pinot.common.datablock.DataBlockUtils;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.HashUtil;
import org.apache.pinot.common.utils.RoaringBitmapUtils;
+import org.apache.pinot.segment.spi.memory.CompoundDataBuffer;
import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
import org.apache.pinot.spi.accounting.ThreadResourceSnapshot;
@@ -433,43 +435,44 @@ public class DataTableImplV4 implements DataTable {
return _errCodeToExceptionMap;
}
- /// Serializes the table into one array of the exact final size, so the
fixed and variable size sections are copied
- /// once and no buffer grows by doubling.
+ /// Serializes the table into one array of the exact final size. Prefer
[#toDataBuffer()], which does not copy the
+ /// fixed and variable size sections.
@Override
public byte[] toBytes()
throws IOException {
- ThreadResourceSnapshot resourceSnapshot = new ThreadResourceSnapshot();
+ DataBuffer dataBuffer = toDataBuffer();
+ byte[] bytes = new byte[Math.toIntExact(dataBuffer.size())];
+ dataBuffer.copyTo(0, bytes);
+ return bytes;
+ }
- // Add table serialization time and memory metadata when the corresponding
measurement is enabled.
- // When CPU time/memory usage is not collectable, we omit these values
from the metadata.
- // The entries are reserved before the array is sized and filled in after
the sections are copied. Both are LONG
- // keys, whose encoded length does not depend on the value, so the
metadata length does not change.
- boolean measureCpuTime =
ThreadResourceUsageProvider.isThreadCpuTimeMeasurementEnabled();
- boolean measureAllocatedBytes =
ThreadResourceUsageProvider.isThreadMemoryMeasurementEnabled();
- if (measureCpuTime) {
- _metadata.put(MetadataKey.RESPONSE_SER_CPU_TIME_NS.getName(), "0");
- }
- if (measureAllocatedBytes) {
- _metadata.put(MetadataKey.RESPONSE_SER_MEM_ALLOCATED_BYTES.getName(),
"0");
- }
+ /// Serializes the table without copying the fixed and variable size
sections. The header, the small sections and
+ /// the metadata are new arrays, and the two data sections are the buffers
of this table. On a table built in this
+ /// process, they are the builder's pages, so no buffer is larger than one
page plus these small arrays.
+ @Override
+ public DataBuffer toDataBuffer()
+ throws IOException {
+ ThreadResourceSnapshot resourceSnapshot = new ThreadResourceSnapshot();
byte[] exceptionsBytes = serializeExceptions();
byte[] dictionaryBytes = _stringDictionary != null ?
serializeStringDictionary() : null;
byte[] dataSchemaBytes = _dataSchema != null ? _dataSchema.toBytes() :
null;
- int metadataLength = serializeMetadata().length;
- long totalLength = HEADER_SIZE + exceptionsBytes.length +
lengthOf(dictionaryBytes) + lengthOf(dataSchemaBytes)
- + sizeOf(_fixedSizeData) + sizeOf(_variableSizeData) + Integer.BYTES +
metadataLength;
- if (totalLength > Integer.MAX_VALUE) {
- throw new IllegalStateException("Serialized data table size " +
totalLength + " exceeds the 2 GB limit");
+ // The header stores every section offset as an int.
+ long sectionsLength = HEADER_SIZE + exceptionsBytes.length +
lengthOf(dictionaryBytes) + lengthOf(dataSchemaBytes)
+ + sizeOf(_fixedSizeData) + sizeOf(_variableSizeData);
+ if (sectionsLength > Integer.MAX_VALUE) {
+ throw new IllegalStateException("Serialized data table size " +
sectionsLength + " exceeds the 2 GB limit");
}
- byte[] bytes = new byte[(int) totalLength];
- ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
- writeLeadingSections(byteBuffer, exceptionsBytes, dictionaryBytes,
dataSchemaBytes);
+ ByteBuffer leadingSections = ByteBuffer.allocate(
+ HEADER_SIZE + exceptionsBytes.length + lengthOf(dictionaryBytes) +
lengthOf(dataSchemaBytes));
+ writeLeadingSections(leadingSections, exceptionsBytes, dictionaryBytes,
dataSchemaBytes);
- if (measureCpuTime) {
+ // Add table serialization time and memory metadata when the corresponding
measurement is enabled.
+ // When CPU time/memory usage is not collectable, we omit these values
from the metadata.
+ if (ThreadResourceUsageProvider.isThreadCpuTimeMeasurementEnabled()) {
_metadata.put(MetadataKey.RESPONSE_SER_CPU_TIME_NS.getName(),
String.valueOf(resourceSnapshot.getCpuTimeNs()));
}
- if (measureAllocatedBytes) {
+ if (ThreadResourceUsageProvider.isThreadMemoryMeasurementEnabled()) {
_metadata.put(MetadataKey.RESPONSE_SER_MEM_ALLOCATED_BYTES.getName(),
String.valueOf(resourceSnapshot.getAllocatedBytes()));
}
@@ -478,14 +481,25 @@ public class DataTableImplV4 implements DataTable {
// NOTE: We ignore metadata serialization time in
"responseSerializationCpuTimeNs" as it's negligible while
// considering it will bring a lot code complexity.
byte[] metadataBytes = serializeMetadata();
- if (metadataBytes.length != metadataLength) {
- throw new IllegalStateException(
- "Metadata length changed during serialization from " +
metadataLength + " to " + metadataBytes.length);
+ // Readers parse the whole table as one buffer, so the total must fit in
an int as well.
+ long totalLength = sectionsLength + Integer.BYTES + metadataBytes.length;
+ if (totalLength > Integer.MAX_VALUE) {
+ throw new IllegalStateException("Serialized data table size " +
totalLength + " exceeds the 2 GB limit");
}
- byteBuffer.putInt(metadataBytes.length);
- byteBuffer.put(metadataBytes);
- assert !byteBuffer.hasRemaining();
- return bytes;
+ ByteBuffer metadata = ByteBuffer.allocate(Integer.BYTES +
metadataBytes.length);
+ metadata.putInt(metadataBytes.length);
+ metadata.put(metadataBytes);
+
+ CompoundDataBuffer.Builder builder = new
CompoundDataBuffer.Builder(ByteOrder.BIG_ENDIAN, false);
+ builder.addBuffer(leadingSections.flip());
+ if (_fixedSizeData != null) {
+ builder.addBuffer(_fixedSizeData);
+ }
+ if (_variableSizeData != null) {
+ builder.addBuffer(_variableSizeData);
+ }
+ builder.addBuffer(metadata.flip());
+ return builder.build();
}
private static int lengthOf(@Nullable byte[] bytes) {
@@ -496,6 +510,8 @@ public class DataTableImplV4 implements DataTable {
return dataBuffer != null ? dataBuffer.size() : 0;
}
+ /// Writes the header and the exceptions, dictionary and data schema
sections. The fixed and variable size sections
+ /// follow them in the serialized table.
private void writeLeadingSections(ByteBuffer byteBuffer, byte[]
exceptionsBytes, @Nullable byte[] dictionaryBytes,
@Nullable byte[] dataSchemaBytes) {
int fixedSizeDataLength = (int) sizeOf(_fixedSizeData);
@@ -538,15 +554,6 @@ public class DataTableImplV4 implements DataTable {
if (dataSchemaBytes != null) {
byteBuffer.put(dataSchemaBytes);
}
- copySection(_fixedSizeData, byteBuffer, fixedSizeDataLength);
- copySection(_variableSizeData, byteBuffer, variableSizeDataLength);
- }
-
- private static void copySection(@Nullable DataBuffer section, ByteBuffer
byteBuffer, int length) {
- if (length > 0) {
- section.copyTo(0, byteBuffer.array(), byteBuffer.position(), length);
- byteBuffer.position(byteBuffer.position() + length);
- }
}
/// Serialize metadata section to bytes.
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingResponseUtils.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingResponseUtils.java
index bdf5464aaaf..c4f453144c9 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingResponseUtils.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingResponseUtils.java
@@ -19,9 +19,14 @@
package org.apache.pinot.core.operator.streaming;
import com.google.protobuf.ByteString;
+import com.google.protobuf.UnsafeByteOperations;
import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
import org.apache.pinot.common.datatable.DataTable;
import org.apache.pinot.common.proto.Server;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.spi.utils.CommonConstants.Query.Response;
@@ -47,6 +52,19 @@ public class StreamingResponseUtils {
private static Server.ServerResponse getResponse(DataTable dataTable, String
responseType)
throws IOException {
return
Server.ServerResponse.newBuilder().putMetadata(Response.MetadataKeys.RESPONSE_TYPE,
responseType)
- .setPayload(ByteString.copyFrom(dataTable.toBytes())).build();
+ .setPayload(toByteString(dataTable.toDataBuffer())).build();
+ }
+
+ /// Wraps the buffers of a serialized data table in one [ByteString] without
copying them. The buffers must not change
+ /// after this call, which holds for a serialized data table.
+ private static ByteString toByteString(DataBuffer dataBuffer) {
+ List<ByteBuffer> buffers = new ArrayList<>();
+ dataBuffer.appendAsByteBuffers(buffers);
+ List<ByteString> byteStrings = new ArrayList<>(buffers.size());
+ for (ByteBuffer buffer : buffers) {
+ byteStrings.add(UnsafeByteOperations.unsafeWrap(buffer));
+ }
+ // Despite its name, copyFrom(Iterable) concatenates the byte strings into
a balanced rope without copying them.
+ return ByteString.copyFrom(byteStrings);
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java
index fdae2557ab5..6d78893daa7 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java
@@ -33,6 +33,7 @@ import org.apache.pinot.core.query.executor.QueryExecutor;
import org.apache.pinot.core.query.request.ServerQueryRequest;
import
org.apache.pinot.core.query.scheduler.resources.BinaryWorkloadResourceManager;
import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.slf4j.Logger;
@@ -84,15 +85,15 @@ public class BinaryWorkloadScheduler extends QueryScheduler
{
}
@Override
- public ListenableFuture<byte[]> submit(ServerQueryRequest queryRequest) {
+ public ListenableFuture<DataBuffer> submitResponse(ServerQueryRequest
queryRequest) {
if (!_isRunning) {
- return shuttingDown(queryRequest);
+ return shuttingDownResponse(queryRequest);
}
queryRequest.getTimerContext().startNewPhaseTimer(ServerQueryPhase.SCHEDULER_WAIT);
if
(!QueryOptionsUtils.isSecondaryWorkload(queryRequest.getQueryContext().getQueryOptions()))
{
QueryExecutorService executorService =
_resourceManager.getExecutorService(queryRequest, null);
- ListenableFutureTask<byte[]> queryTask =
createQueryFutureTask(queryRequest, executorService);
+ ListenableFutureTask<DataBuffer> queryTask =
createResponseFutureTask(queryRequest, executorService);
_resourceManager.getQueryRunners().submit(queryTask);
return queryTask;
}
@@ -106,13 +107,13 @@ public class BinaryWorkloadScheduler extends
QueryScheduler {
} catch (OutOfCapacityException e) {
LOGGER.error("Out of capacity for query {} table {}, message: {}",
queryRequest.getRequestId(),
queryRequest.getTableNameWithType(), e.getMessage());
- return outOfCapacity(queryRequest);
+ return outOfCapacityResponse(queryRequest);
} catch (Exception e) {
// We should not throw any other exception other than
OutOfCapacityException. Signal that there's an issue with
// the scheduler if any other exception is thrown.
LOGGER.error("Internal error for query {} table {}, message {}",
queryRequest.getRequestId(),
queryRequest.getTableNameWithType(), e.getMessage());
- return shuttingDown(queryRequest);
+ return shuttingDownResponse(queryRequest);
}
return schedQueryContext.getResultFuture();
}
@@ -151,7 +152,7 @@ public class BinaryWorkloadScheduler extends QueryScheduler
{
ServerQueryRequest queryRequest = request.getQueryRequest();
SchedulerGroup schedulerGroup = request.getSchedulerGroup();
QueryExecutorService executorService =
_resourceManager.getExecutorService(queryRequest, schedulerGroup);
- ListenableFutureTask<byte[]> queryFutureTask =
createQueryFutureTask(queryRequest, executorService);
+ ListenableFutureTask<DataBuffer> queryFutureTask =
createResponseFutureTask(queryRequest, executorService);
queryFutureTask.addListener(new Runnable() {
@Override
public void run() {
@@ -209,7 +210,7 @@ public class BinaryWorkloadScheduler extends QueryScheduler
{
synchronized private void failAllPendingQueries() {
List<SchedulerQueryContext> pending = _secondaryQueryQ.drain();
for (SchedulerQueryContext queryContext : pending) {
- ListenableFuture<byte[]> serverShuttingDown =
shuttingDown(queryContext.getQueryRequest());
+ ListenableFuture<DataBuffer> serverShuttingDown =
shuttingDownResponse(queryContext.getQueryRequest());
queryContext.setResultFuture(serverShuttingDown);
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/PriorityScheduler.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/PriorityScheduler.java
index b1e0a8e5d70..fa155afe6eb 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/PriorityScheduler.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/PriorityScheduler.java
@@ -30,6 +30,7 @@ import org.apache.pinot.core.query.executor.QueryExecutor;
import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService;
import org.apache.pinot.core.query.scheduler.resources.ResourceManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.slf4j.Logger;
@@ -82,9 +83,9 @@ public abstract class PriorityScheduler extends
QueryScheduler {
}
@Override
- public ListenableFuture<byte[]> submit(ServerQueryRequest queryRequest) {
+ public ListenableFuture<DataBuffer> submitResponse(ServerQueryRequest
queryRequest) {
if (!_isRunning) {
- return shuttingDown(queryRequest);
+ return shuttingDownResponse(queryRequest);
}
queryRequest.getTimerContext().startNewPhaseTimer(ServerQueryPhase.SCHEDULER_WAIT);
final SchedulerQueryContext schedQueryContext = new
SchedulerQueryContext(queryRequest);
@@ -92,7 +93,7 @@ public abstract class PriorityScheduler extends
QueryScheduler {
_queryQueue.put(schedQueryContext);
} catch (OutOfCapacityException e) {
LOGGER.error("Out of capacity for table {}, message: {}",
queryRequest.getTableNameWithType(), e.getMessage());
- return outOfCapacity(queryRequest);
+ return outOfCapacityResponse(queryRequest);
}
return schedQueryContext.getResultFuture();
}
@@ -122,7 +123,7 @@ public abstract class PriorityScheduler extends
QueryScheduler {
ServerQueryRequest queryRequest = request.getQueryRequest();
SchedulerGroup schedulerGroup = request.getSchedulerGroup();
QueryExecutorService executorService =
_resourceManager.getExecutorService(queryRequest, schedulerGroup);
- ListenableFutureTask<byte[]> queryFutureTask =
createQueryFutureTask(queryRequest, executorService);
+ ListenableFutureTask<DataBuffer> queryFutureTask =
createResponseFutureTask(queryRequest, executorService);
queryFutureTask.addListener(new Runnable() {
@Override
public void run() {
@@ -170,7 +171,7 @@ public abstract class PriorityScheduler extends
QueryScheduler {
synchronized private void failAllPendingQueries() {
List<SchedulerQueryContext> pending = _queryQueue.drain();
for (SchedulerQueryContext queryContext : pending) {
-
queryContext.setResultFuture(shuttingDown(queryContext.getQueryRequest()));
+
queryContext.setResultFuture(shuttingDownResponse(queryContext.getQueryRequest()));
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QueryScheduler.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QueryScheduler.java
index 8b39f257c99..3e8bf2c8c87 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QueryScheduler.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QueryScheduler.java
@@ -18,9 +18,11 @@
*/
package org.apache.pinot.core.query.scheduler;
+import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListenableFutureTask;
+import com.google.common.util.concurrent.MoreExecutors;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.LongAccumulator;
@@ -37,6 +39,8 @@ import org.apache.pinot.core.query.logger.ServerQueryLogger;
import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.query.request.context.TimerContext;
import org.apache.pinot.core.query.scheduler.resources.ResourceManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
+import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.exception.QueryErrorCode;
@@ -74,6 +78,19 @@ public abstract class QueryScheduler {
_threadAccountant = threadAccountant;
_latestQueryTime = latestQueryTime;
_resourceManager = resourceManager;
+ // submit() and submitResponse() call each other by default. Fail here
instead of with a StackOverflowError on the
+ // first query.
+ Preconditions.checkState(overrides("submit") ||
overrides("submitResponse"),
+ "%s must override submitResponse(ServerQueryRequest) or the deprecated
submit(ServerQueryRequest)",
+ getClass().getName());
+ }
+
+ private boolean overrides(String methodName) {
+ try {
+ return getClass().getMethod(methodName,
ServerQueryRequest.class).getDeclaringClass() != QueryScheduler.class;
+ } catch (NoSuchMethodException e) {
+ throw new IllegalStateException(e);
+ }
}
public ResourceManager getResourceManager() {
@@ -84,7 +101,29 @@ public abstract class QueryScheduler {
/// @param queryRequest query to schedule for execution
/// @return Listenable future for query result representing serialized
response. It is possible that the
/// future may return immediately or be scheduled for execution at a
later time.
- public abstract ListenableFuture<byte[]> submit(ServerQueryRequest
queryRequest);
+ /// @deprecated Use [#submitResponse(ServerQueryRequest)], which does not
copy the response into one array. The
+ /// default implementation copies the result of
[#submitResponse(ServerQueryRequest)]. Subclasses must override
+ /// at least one of the two methods. The server calls
[#submitResponse(ServerQueryRequest)], so a subclass of a
+ /// scheduler that already overrides
[#submitResponse(ServerQueryRequest)] (all the schedulers in this package)
+ /// must override that method too: an override of this method alone is
not called.
+ @Deprecated
+ public ListenableFuture<byte[]> submit(ServerQueryRequest queryRequest) {
+ return Futures.transform(submitResponse(queryRequest),
QueryScheduler::toByteArray, MoreExecutors.directExecutor());
+ }
+
+ /// Submit a query for execution. The query will be scheduled for execution
as per the scheduling algorithm.
+ ///
+ /// The serialized response can be split into several buffers, so a large
response does not need one large array.
+ ///
+ /// The default implementation wraps the result of the deprecated
[#submit(ServerQueryRequest)], so that subclasses
+ /// that only override that method keep working. Subclasses must override at
least one of the two methods.
+ ///
+ /// @param queryRequest query to schedule for execution
+ /// @return Listenable future for the serialized response. It is possible
that the future may return immediately or
+ /// be scheduled for execution at a later time.
+ public ListenableFuture<DataBuffer> submitResponse(ServerQueryRequest
queryRequest) {
+ return Futures.transform(submit(queryRequest), QueryScheduler::wrap,
MoreExecutors.directExecutor());
+ }
/// Query scheduler name for logging
public abstract String name();
@@ -107,17 +146,34 @@ public abstract class QueryScheduler {
/// internally used to parallelize query processing.
/// @return Future task that can be scheduled for execution on an
ExecutorService. Ideally, this future
/// should be executed on a different executor service than
`executorService` to avoid deadlock.
+ protected ListenableFutureTask<DataBuffer>
createResponseFutureTask(ServerQueryRequest queryRequest,
+ ExecutorService executorService) {
+ return ListenableFutureTask.create(() ->
processQueryAndSerializeResponse(queryRequest, executorService));
+ }
+
+ /// @deprecated Use [#createResponseFutureTask(ServerQueryRequest,
ExecutorService)], which does not copy the
+ /// response into one array.
+ @Deprecated
protected ListenableFutureTask<byte[]>
createQueryFutureTask(ServerQueryRequest queryRequest,
ExecutorService executorService) {
return ListenableFutureTask.create(() ->
processQueryAndSerialize(queryRequest, executorService));
}
+ /// @deprecated Use [#processQueryAndSerializeResponse(ServerQueryRequest,
ExecutorService)], which does not copy
+ /// the response into one array.
+ @Deprecated
+ @Nullable
+ protected byte[] processQueryAndSerialize(ServerQueryRequest queryRequest,
ExecutorService executorService) {
+ return toByteArray(processQueryAndSerializeResponse(queryRequest,
executorService));
+ }
+
/// Process query and serialize response
/// @param queryRequest incoming query request
/// @param executorService Executor service to use for parallelizing query
processing
/// @return serialized query response
@Nullable
- protected byte[] processQueryAndSerialize(ServerQueryRequest queryRequest,
ExecutorService executorService) {
+ protected DataBuffer processQueryAndSerializeResponse(ServerQueryRequest
queryRequest,
+ ExecutorService executorService) {
QueryExecutionContext executionContext =
queryRequest.toExecutionContext(_instanceId);
_latestQueryTime.accumulate(executionContext.getStartTimeMs());
try (QueryThreadContext ignore = QueryThreadContext.open(executionContext,
_threadAccountant)) {
@@ -141,7 +197,7 @@ public abstract class QueryScheduler {
responseMetadata.put(MetadataKey.REQUEST_ID.getName(),
Long.toString(requestId));
responseMetadata.put(MetadataKey.QUERY_ID.getName(), queryId);
responseMetadata.put(MetadataKey.WORKLOAD_NAME.getName(), workloadName);
- byte[] responseBytes = serializeResponse(queryRequest, instanceResponse);
+ DataBuffer response = serializeResponse(queryRequest, instanceResponse);
// Log the statistics
if (_queryLogger != null) {
@@ -150,8 +206,8 @@ public abstract class QueryScheduler {
// TODO: Perform this check sooner during the serialization of DataTable.
Map<String, String> queryOptions =
queryRequest.getQueryContext().getQueryOptions();
- if (responseBytes != null) {
- int responseSizeBytes = responseBytes.length;
+ if (response != null) {
+ long responseSizeBytes = response.size();
String tableNameWithType = queryRequest.getTableNameWithType();
_serverMetrics.addMeteredTableValue(tableNameWithType,
ServerMeter.QUERY_RESPONSE_SIZE, responseSizeBytes);
Long maxResponseSizeBytes =
QueryOptionsUtils.getMaxServerResponseSizeBytes(queryOptions);
@@ -167,29 +223,29 @@ public abstract class QueryScheduler {
instanceResponse.addMetadata(MetadataKey.REQUEST_ID.getName(),
Long.toString(requestId));
instanceResponse.addMetadata(MetadataKey.QUERY_ID.getName(),
queryId);
instanceResponse.addMetadata(MetadataKey.WORKLOAD_NAME.getName(),
workloadName);
- responseBytes = serializeResponse(queryRequest, instanceResponse);
+ response = serializeResponse(queryRequest, instanceResponse);
}
}
- return responseBytes;
+ return response;
}
}
/// Serialize the instance response for query request
/// @param queryRequest Server query request for which response is serialized
/// @param instanceResponse instance response to serialize
- /// @return serialized response bytes
+ /// @return serialized response
@Nullable
- private byte[] serializeResponse(ServerQueryRequest queryRequest,
InstanceResponseBlock instanceResponse) {
+ private DataBuffer serializeResponse(ServerQueryRequest queryRequest,
InstanceResponseBlock instanceResponse) {
TimerContext timerContext = queryRequest.getTimerContext();
TimerContext.Timer responseSerializationTimer =
timerContext.startNewPhaseTimer(ServerQueryPhase.RESPONSE_SERIALIZATION);
long requestId = queryRequest.getRequestId();
String brokerId = queryRequest.getBrokerId();
- byte[] responseBytes = null;
+ DataBuffer response = null;
try {
- responseBytes = instanceResponse.toDataTable().toBytes();
+ response = instanceResponse.toDataTable().toDataBuffer();
} catch (Exception e) {
// First check terminate exception and use it as the response if exists.
We want to return the termination reason
// when query is explicitly terminated.
@@ -214,7 +270,7 @@ public abstract class QueryScheduler {
errorResponse.addMetadata(MetadataKey.QUERY_ID.getName(),
queryRequest.getCid());
String workloadName =
QueryOptionsUtils.getWorkloadName(queryRequest.getQueryContext().getQueryOptions());
errorResponse.addMetadata(MetadataKey.WORKLOAD_NAME.getName(),
workloadName);
- responseBytes = errorResponse.toDataTable().toBytes();
+ response = errorResponse.toDataTable().toDataBuffer();
} catch (Exception e1) {
_serverMetrics.addMeteredGlobalValue(ServerMeter.RESPONSE_SERIALIZATION_EXCEPTIONS,
1);
LOGGER.error("Caught exception while constructing error response for
requestId: {}, brokerId: {}",
@@ -226,12 +282,12 @@ public abstract class QueryScheduler {
timerContext.startNewPhaseTimer(ServerQueryPhase.TOTAL_QUERY_TIME,
timerContext.getQueryArrivalTimeMs())
.stopAndRecord();
- return responseBytes;
+ return response;
}
/// Error response future in case of internal error where query response is
not available. This can happen if the
/// query can not be executed.
- protected ListenableFuture<byte[]> immediateErrorResponse(ServerQueryRequest
queryRequest, QueryErrorCode errorCode) {
+ protected ListenableFuture<DataBuffer> errorResponse(ServerQueryRequest
queryRequest, QueryErrorCode errorCode) {
Map<String, String> queryOptions =
queryRequest.getQueryContext().getQueryOptions();
String workloadName = QueryOptionsUtils.getWorkloadName(queryOptions);
InstanceResponseBlock instanceResponse = new InstanceResponseBlock();
@@ -242,11 +298,46 @@ public abstract class QueryScheduler {
return Futures.immediateFuture(serializeResponse(queryRequest,
instanceResponse));
}
+ protected ListenableFuture<DataBuffer>
shuttingDownResponse(ServerQueryRequest queryRequest) {
+ return errorResponse(queryRequest, QueryErrorCode.SERVER_SHUTTING_DOWN);
+ }
+
+ protected ListenableFuture<DataBuffer>
outOfCapacityResponse(ServerQueryRequest queryRequest) {
+ return errorResponse(queryRequest, QueryErrorCode.SERVER_OUT_OF_CAPACITY);
+ }
+
+ /// @deprecated Use [#errorResponse(ServerQueryRequest, QueryErrorCode)].
+ @Deprecated
+ protected ListenableFuture<byte[]> immediateErrorResponse(ServerQueryRequest
queryRequest, QueryErrorCode errorCode) {
+ return Futures.transform(errorResponse(queryRequest, errorCode),
QueryScheduler::toByteArray,
+ MoreExecutors.directExecutor());
+ }
+
+ /// @deprecated Use [#shuttingDownResponse(ServerQueryRequest)].
+ @Deprecated
protected ListenableFuture<byte[]> shuttingDown(ServerQueryRequest
queryRequest) {
return immediateErrorResponse(queryRequest,
QueryErrorCode.SERVER_SHUTTING_DOWN);
}
+ /// @deprecated Use [#outOfCapacityResponse(ServerQueryRequest)].
+ @Deprecated
protected ListenableFuture<byte[]> outOfCapacity(ServerQueryRequest
queryRequest) {
return immediateErrorResponse(queryRequest,
QueryErrorCode.SERVER_OUT_OF_CAPACITY);
}
+
+ /// Copies a serialized response into one array, for the deprecated methods
that return `byte[]`.
+ @Nullable
+ private static byte[] toByteArray(@Nullable DataBuffer response) {
+ if (response == null) {
+ return null;
+ }
+ byte[] bytes = new byte[Math.toIntExact(response.size())];
+ response.copyTo(0, bytes);
+ return bytes;
+ }
+
+ @Nullable
+ private static DataBuffer wrap(@Nullable byte[] responseBytes) {
+ return responseBytes != null ? PinotByteBuffer.wrap(responseBytes) : null;
+ }
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/SchedulerQueryContext.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/SchedulerQueryContext.java
index c67bd784cb8..104527dabdc 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/SchedulerQueryContext.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/SchedulerQueryContext.java
@@ -23,13 +23,14 @@ import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import javax.annotation.Nullable;
import org.apache.pinot.core.query.request.ServerQueryRequest;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
/// Query specific scheduling context
public class SchedulerQueryContext {
private final ServerQueryRequest _queryRequest;
- private final SettableFuture<byte[]> _resultFuture;
+ private final SettableFuture<DataBuffer> _resultFuture;
private SchedulerGroup _schedulerGroup;
public SchedulerQueryContext(ServerQueryRequest queryRequest) {
@@ -42,11 +43,11 @@ public class SchedulerQueryContext {
return _queryRequest;
}
- public SettableFuture<byte[]> getResultFuture() {
+ public SettableFuture<DataBuffer> getResultFuture() {
return _resultFuture;
}
- public void setResultFuture(ListenableFuture<byte[]> f) {
+ public void setResultFuture(ListenableFuture<DataBuffer> f) {
_resultFuture.setFuture(f);
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java
index e84d6037d24..ea29289f80e 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java
@@ -28,6 +28,7 @@ import org.apache.pinot.core.query.executor.QueryExecutor;
import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService;
import
org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.accounting.WorkloadBudgetManager;
import org.apache.pinot.spi.accounting.WorkloadBudgetManagerFactory;
@@ -65,9 +66,9 @@ public class WorkloadScheduler extends QueryScheduler {
}
@Override
- public ListenableFuture<byte[]> submit(ServerQueryRequest queryRequest) {
+ public ListenableFuture<DataBuffer> submitResponse(ServerQueryRequest
queryRequest) {
if (!_isRunning) {
- return shuttingDown(queryRequest);
+ return shuttingDownResponse(queryRequest);
}
boolean isSecondary =
QueryOptionsUtils.isSecondaryWorkload(queryRequest.getQueryContext().getQueryOptions());
@@ -85,11 +86,11 @@ public class WorkloadScheduler extends QueryScheduler {
_serverMetrics.addMeteredValue(workloadName,
ServerMeter.WORKLOAD_BUDGET_EXCEEDED, 1L);
_serverMetrics.addMeteredTableValue(tableName,
ServerMeter.WORKLOAD_BUDGET_EXCEEDED, 1L);
_serverMetrics.addMeteredGlobalValue(ServerMeter.WORKLOAD_BUDGET_EXCEEDED, 1L);
- return immediateErrorResponse(queryRequest,
QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED);
+ return errorResponse(queryRequest,
QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED);
}
queryRequest.getTimerContext().startNewPhaseTimer(ServerQueryPhase.SCHEDULER_WAIT);
QueryExecutorService executorService =
_resourceManager.getExecutorService(queryRequest, null);
- ListenableFutureTask<byte[]> queryTask =
createQueryFutureTask(queryRequest, executorService);
+ ListenableFutureTask<DataBuffer> queryTask =
createResponseFutureTask(queryRequest, executorService);
_resourceManager.getQueryRunners().submit(queryTask);
return queryTask;
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java
index a34887d3bbf..83f615bde06 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java
@@ -27,6 +27,7 @@ import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.query.scheduler.QueryScheduler;
import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService;
import
org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.env.PinotConfiguration;
@@ -42,13 +43,13 @@ public class FCFSQueryScheduler extends QueryScheduler {
}
@Override
- public ListenableFuture<byte[]> submit(ServerQueryRequest queryRequest) {
+ public ListenableFuture<DataBuffer> submitResponse(ServerQueryRequest
queryRequest) {
if (!_isRunning) {
- return shuttingDown(queryRequest);
+ return shuttingDownResponse(queryRequest);
}
queryRequest.getTimerContext().startNewPhaseTimer(ServerQueryPhase.SCHEDULER_WAIT);
QueryExecutorService executorService =
_resourceManager.getExecutorService(queryRequest, null);
- ListenableFutureTask<byte[]> queryTask =
createQueryFutureTask(queryRequest, executorService);
+ ListenableFutureTask<DataBuffer> queryTask =
createResponseFutureTask(queryRequest, executorService);
_resourceManager.getQueryRunners().submit(queryTask);
return queryTask;
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/transport/InstanceRequestHandler.java
b/pinot-core/src/main/java/org/apache/pinot/core/transport/InstanceRequestHandler.java
index 1e6a1177374..2e1789607b0 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/transport/InstanceRequestHandler.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/transport/InstanceRequestHandler.java
@@ -30,7 +30,10 @@ import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.ssl.SslHandshakeCompletionEvent;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
import java.util.HashSet;
+import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CancellationException;
@@ -48,6 +51,7 @@ import org.apache.pinot.common.request.InstanceRequest;
import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.query.scheduler.QueryScheduler;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.server.access.AccessControl;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.env.PinotConfiguration;
@@ -185,7 +189,7 @@ public class InstanceRequestHandler extends
SimpleChannelInboundHandler<ByteBuf>
void submitQuery(ServerQueryRequest queryRequest, ChannelHandlerContext ctx,
long queryArrivalTimeMs) {
QueryExecutionContext executionContext =
queryRequest.toExecutionContext(_instanceName);
try (QueryThreadContext ignore = QueryThreadContext.open(executionContext,
_threadAccountant)) {
- ListenableFuture<byte[]> future = _queryScheduler.submit(queryRequest);
+ ListenableFuture<DataBuffer> future =
_queryScheduler.submitResponse(queryRequest);
if (_executionContexts != null) {
String queryId = queryRequest.getQueryId();
// Track the running query for cancellation.
@@ -199,11 +203,11 @@ public class InstanceRequestHandler extends
SimpleChannelInboundHandler<ByteBuf>
}
}
- private FutureCallback<byte[]> createCallback(ServerQueryRequest
queryRequest, ChannelHandlerContext ctx,
+ private FutureCallback<DataBuffer> createCallback(ServerQueryRequest
queryRequest, ChannelHandlerContext ctx,
long queryArrivalTimeMs) {
return new FutureCallback<>() {
@Override
- public void onSuccess(@Nullable byte[] responseBytes) {
+ public void onSuccess(@Nullable DataBuffer response) {
if (_executionContexts != null) {
String queryId = queryRequest.getQueryId();
if (LOGGER.isDebugEnabled()) {
@@ -213,9 +217,9 @@ public class InstanceRequestHandler extends
SimpleChannelInboundHandler<ByteBuf>
}
long requestId = queryRequest.getRequestId();
String tableNameWithType = queryRequest.getTableNameWithType();
- if (responseBytes != null) {
- // responseBytes contains either query results or exception.
- sendResponse(ctx, requestId, tableNameWithType, queryArrivalTimeMs,
responseBytes);
+ if (response != null) {
+ // response contains either query results or exception.
+ sendResponse(ctx, requestId, tableNameWithType, queryArrivalTimeMs,
response);
} else {
// Send exception response.
sendErrorResponse(ctx, requestId, tableNameWithType,
queryArrivalTimeMs,
@@ -315,8 +319,7 @@ public class InstanceRequestHandler extends
SimpleChannelInboundHandler<ByteBuf>
dataTable.addException(QueryErrorCode.QUERY_EXECUTION.getId(),
"Query execution error on: " + _instanceName + " " +
e.getMessage());
}
- byte[] serializedDataTable = dataTable.toBytes();
- sendResponse(ctx, requestId, tableNameWithType, queryArrivalTimeMs,
serializedDataTable);
+ sendResponse(ctx, requestId, tableNameWithType, queryArrivalTimeMs,
dataTable.toDataBuffer());
} catch (Exception exception) {
LOGGER.error("Exception while sending query processing error to
Broker.", exception);
} finally {
@@ -329,16 +332,21 @@ public class InstanceRequestHandler extends
SimpleChannelInboundHandler<ByteBuf>
}
/// Send a response (either query results or exception) back to broker as
response to the query request.
+ ///
+ /// The response buffers are wrapped in one composite Netty buffer without
copying them.
private void sendResponse(ChannelHandlerContext ctx, long requestId, String
tableNameWithType,
- long queryArrivalTimeMs, byte[] serializedDataTable) {
+ long queryArrivalTimeMs, DataBuffer response) {
long sendResponseStartTimeMs = System.currentTimeMillis();
int queryProcessingTimeMs = (int) (sendResponseStartTimeMs -
queryArrivalTimeMs);
-
ctx.writeAndFlush(Unpooled.wrappedBuffer(serializedDataTable)).addListener(f ->
{
+ long responseSizeBytes = response.size();
+ List<ByteBuffer> responseBuffers = new ArrayList<>();
+ response.appendAsByteBuffers(responseBuffers);
+ ctx.writeAndFlush(Unpooled.wrappedBuffer(responseBuffers.toArray(new
ByteBuffer[0]))).addListener(f -> {
if (f.isSuccess()) {
long sendResponseEndTimeMs = System.currentTimeMillis();
int sendResponseLatencyMs = (int) (sendResponseEndTimeMs -
sendResponseStartTimeMs);
_serverMetrics.addMeteredGlobalValue(ServerMeter.NETTY_CONNECTION_RESPONSES_SENT,
1);
-
_serverMetrics.addMeteredGlobalValue(ServerMeter.NETTY_CONNECTION_BYTES_SENT,
serializedDataTable.length);
+
_serverMetrics.addMeteredGlobalValue(ServerMeter.NETTY_CONNECTION_BYTES_SENT,
responseSizeBytes);
_serverMetrics.addTimedTableValue(tableNameWithType,
ServerTimer.NETTY_CONNECTION_SEND_RESPONSE_LATENCY,
sendResponseLatencyMs, TimeUnit.MILLISECONDS);
@@ -348,9 +356,9 @@ public class InstanceRequestHandler extends
SimpleChannelInboundHandler<ByteBuf>
"Slow query ({}): request handler processing time: {}, send
response latency: {}, total time to handle "
+ "request: {}", requestId, queryProcessingTimeMs,
sendResponseLatencyMs, totalQueryTimeMs);
}
- if (serializedDataTable.length > LARGE_RESPONSE_SIZE_THRESHOLD_BYTES) {
+ if (responseSizeBytes > LARGE_RESPONSE_SIZE_THRESHOLD_BYTES) {
LOGGER.warn("Large query ({}): response size in bytes: {}, table
name {}", requestId,
- serializedDataTable.length, tableNameWithType);
+ responseSizeBytes, tableNameWithType);
_serverMetrics.addMeteredTableValue(tableNameWithType,
ServerMeter.LARGE_QUERY_RESPONSES_SENT, 1);
}
} else {
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
index 41b2ed68b15..fd2dc9bf158 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
@@ -22,8 +22,10 @@ import java.io.IOException;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.UUID;
@@ -36,6 +38,8 @@ import org.apache.pinot.common.datatable.DataTableFactory;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
+import org.apache.pinot.segment.spi.memory.PagedPinotOutputStream;
import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.utils.ByteArray;
@@ -372,7 +376,24 @@ public class DataTableSerDeTest {
DataTable dataTable = dataTableBuilder.build();
DataTable newDataTable =
DataTableFactory.getDataTable(dataTable.toBytes());
- for (DataTable table : new DataTable[]{dataTable,
dataTable.toDataOnlyDataTable(), newDataTable}) {
+
+ // toDataBuffer() returns the builder's pages without copying them. The
first buffer holds the header and the
+ // small sections, and the last one the metadata. Every buffer between
them is one page at most.
+ DataBuffer dataBuffer = dataTable.toDataBuffer();
+ List<ByteBuffer> buffers = new ArrayList<>();
+ dataBuffer.appendAsByteBuffers(buffers);
+ assertTrue(buffers.size() > 2, "Expected the sections in pages: " +
buffers.size());
+ for (ByteBuffer buffer : buffers.subList(1, buffers.size() - 1)) {
+ assertTrue(buffer.remaining() <=
PagedPinotOutputStream.PageAllocator.MIN_RECOMMENDED_PAGE_SIZE,
+ "Buffer of " + buffer.remaining() + " bytes");
+ }
+ byte[] dataBufferBytes = new byte[Math.toIntExact(dataBuffer.size())];
+ dataBuffer.copyTo(0, dataBufferBytes);
+ DataTable dataBufferTable = DataTableFactory.getDataTable(dataBufferBytes);
+
+ for (DataTable table : new DataTable[]{
+ dataTable, dataTable.toDataOnlyDataTable(), newDataTable,
dataBufferTable
+ }) {
assertEquals(table.getNumberOfRows(), numRows, ERROR_MESSAGE);
RoaringBitmap nullRowIds = table.getNullRowIds(6);
assertNotNull(nullRowIds, ERROR_MESSAGE);
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingResponseUtilsTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingResponseUtilsTest.java
new file mode 100644
index 00000000000..51386128c13
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingResponseUtilsTest.java
@@ -0,0 +1,72 @@
+/**
+ * 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.streaming;
+
+import java.util.Arrays;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.DataTableFactory;
+import org.apache.pinot.common.proto.Server;
+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.spi.utils.ByteArray;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+/// Tests that [StreamingResponseUtils] puts a paged data table into the gRPC
payload with the same bytes as
+/// `DataTable.toBytes()`.
+public class StreamingResponseUtilsTest {
+ private static final int NUM_ROWS = 200;
+ private static final int VALUE_SIZE = 10_000;
+
+ @Test
+ public void testPayloadOfPagedDataTable()
+ throws Exception {
+ DataSchema dataSchema = new DataSchema(new String[]{"bytes"}, new
ColumnDataType[]{ColumnDataType.BYTES});
+ DataTableBuilder dataTableBuilder =
DataTableBuilderFactory.getDataTableBuilder(dataSchema);
+ for (int rowId = 0; rowId < NUM_ROWS; rowId++) {
+ dataTableBuilder.startRow();
+ dataTableBuilder.setColumn(0, new ByteArray(valueOf(rowId)));
+ dataTableBuilder.finishRow();
+ }
+ DataTable dataTable = dataTableBuilder.build();
+
+ for (Server.ServerResponse response : new Server.ServerResponse[]{
+ StreamingResponseUtils.getDataResponse(dataTable),
StreamingResponseUtils.getNonStreamingResponse(dataTable)
+ }) {
+ DataTable receivedDataTable =
DataTableFactory.getDataTable(response.getPayload().toByteArray());
+ assertEquals(receivedDataTable.getNumberOfRows(), NUM_ROWS);
+ for (int rowId = 0; rowId < NUM_ROWS; rowId++) {
+ assertEquals(receivedDataTable.getBytes(rowId, 0).getBytes(),
valueOf(rowId));
+ }
+ // The payload must also survive the protobuf wire encoding, which
writes each part of the rope.
+ Server.ServerResponse parsed =
Server.ServerResponse.parseFrom(response.toByteArray());
+ assertEquals(parsed.getPayload(), response.getPayload());
+ }
+ }
+
+ private static byte[] valueOf(int rowId) {
+ byte[] value = new byte[VALUE_SIZE];
+ Arrays.fill(value, (byte) rowId);
+ return value;
+ }
+}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerResponseApiTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerResponseApiTest.java
new file mode 100644
index 00000000000..23abaa22e1f
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerResponseApiTest.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.query.scheduler;
+
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import java.nio.ByteBuffer;
+import java.util.concurrent.atomic.LongAccumulator;
+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.CompoundDataBuffer;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
+import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Tests the bridge between the deprecated `byte[]` API of [QueryScheduler]
and the [DataBuffer] API, in both
+/// directions.
+public class QuerySchedulerResponseApiTest {
+ private static final byte[] RESPONSE_BYTES = {1, 2, 3, 4, 5, 6, 7};
+ private static final PinotConfiguration CONFIG = new PinotConfiguration();
+
+ /// A scheduler written against the old API, like the ones outside this
repository.
+ @Test
+ public void testSubmitResponseWrapsDeprecatedSubmit()
+ throws Exception {
+ QueryScheduler scheduler = new TestScheduler() {
+ @Override
+ @SuppressWarnings("deprecation")
+ public ListenableFuture<byte[]> submit(ServerQueryRequest queryRequest) {
+ return Futures.immediateFuture(RESPONSE_BYTES);
+ }
+ };
+ DataBuffer response =
scheduler.submitResponse(mock(ServerQueryRequest.class)).get();
+ byte[] bytes = new byte[Math.toIntExact(response.size())];
+ response.copyTo(0, bytes);
+ assertEquals(bytes, RESPONSE_BYTES);
+ }
+
+ /// A scheduler written against the new API, called through the deprecated
method.
+ @Test
+ @SuppressWarnings("deprecation")
+ public void testDeprecatedSubmitCopiesSubmitResponse()
+ throws Exception {
+ QueryScheduler scheduler = new TestScheduler() {
+ @Override
+ public ListenableFuture<DataBuffer> submitResponse(ServerQueryRequest
queryRequest) {
+ // Split the response into three buffers.
+ CompoundDataBuffer.Builder builder = new CompoundDataBuffer.Builder();
+ builder.addBuffer(ByteBuffer.wrap(RESPONSE_BYTES, 0, 2).slice());
+ builder.addBuffer(ByteBuffer.wrap(RESPONSE_BYTES, 2, 3).slice());
+ builder.addBuffer(ByteBuffer.wrap(RESPONSE_BYTES, 5, 2).slice());
+ return Futures.immediateFuture(builder.build());
+ }
+ };
+ assertEquals(scheduler.submit(mock(ServerQueryRequest.class)).get(),
RESPONSE_BYTES);
+ }
+
+ @Test
+ public void testSchedulerMustOverrideOneSubmitMethod() {
+ assertThrows(IllegalStateException.class, TestScheduler::new);
+ }
+
+ /// Overrides neither submit method, so subclasses choose which one to
override.
+ private static class TestScheduler extends QueryScheduler {
+ TestScheduler() {
+ super(CONFIG, "serverId", mock(QueryExecutor.class),
ThreadAccountantUtils.getNoOpAccountant(),
+ new LongAccumulator(Long::max, 0), new
UnboundedResourceManager(CONFIG));
+ }
+
+ @Override
+ public String name() {
+ return "test";
+ }
+ }
+}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/transport/InstanceRequestHandlerTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/transport/InstanceRequestHandlerTest.java
index ecf17e252f2..81eb6985029 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/transport/InstanceRequestHandlerTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/transport/InstanceRequestHandlerTest.java
@@ -98,7 +98,7 @@ public class InstanceRequestHandlerTest {
DataTable dataTable = DataTableFactory.getDataTable(replyBytes);
assertEquals(dataTable.getMetadata().get(DataTable.MetadataKey.REQUEST_ID.getName()),
Long.toString(requestId),
"The broker matches the reply to its ping by request id");
- verify(queryScheduler, never()).submit(any());
+ verify(queryScheduler, never()).submitResponse(any());
channel.finishAndReleaseAll();
}
@@ -117,12 +117,12 @@ public class InstanceRequestHandlerTest {
assertEquals(queryRequest.getRequestId(), -42L);
QueryScheduler queryScheduler = mock(QueryScheduler.class);
- when(queryScheduler.submit(any())).thenReturn(SettableFuture.create());
+
when(queryScheduler.submitResponse(any())).thenReturn(SettableFuture.create());
EmbeddedChannel channel = new EmbeddedChannel(
new InstanceRequestHandler("server01", new PinotConfiguration(),
queryScheduler, mock(AccessControl.class),
ThreadAccountantUtils.getNoOpAccountant()));
channel.writeInbound(Unpooled.wrappedBuffer(serialize(asSeenByOldServer)));
- verify(queryScheduler).submit(any());
+ verify(queryScheduler).submitResponse(any());
channel.finishAndReleaseAll();
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
index 379326ca7fa..b8516d1dc5d 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
@@ -35,7 +35,9 @@ import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketTimeoutException;
+import java.nio.ByteBuffer;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -53,6 +55,8 @@ import
org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
import org.apache.pinot.core.query.scheduler.QueryScheduler;
import org.apache.pinot.core.routing.SegmentsToQuery;
import
org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
+import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
import org.apache.pinot.server.access.AccessControl;
import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
import org.apache.pinot.spi.config.table.TableType;
@@ -60,6 +64,7 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.metrics.PinotMetricUtils;
import org.apache.pinot.spi.metrics.PinotMetricsRegistry;
+import org.apache.pinot.spi.utils.ByteArray;
import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
import org.apache.pinot.util.TestUtils;
@@ -140,8 +145,12 @@ public class QueryRoutingTest {
}
private QueryServer getQueryServer(int responseDelayMs, byte[]
responseBytes, int port) {
+ return getQueryServer(responseDelayMs,
PinotByteBuffer.wrap(responseBytes), port);
+ }
+
+ private QueryServer getQueryServer(int responseDelayMs, DataBuffer response,
int port) {
InstanceRequestHandler handler = new InstanceRequestHandler("server01",
new PinotConfiguration(),
- mockQueryScheduler(responseDelayMs, responseBytes),
mock(AccessControl.class),
+ mockQueryScheduler(responseDelayMs, response),
mock(AccessControl.class),
ThreadAccountantUtils.getNoOpAccountant());
return new QueryServer(port, null, handler);
}
@@ -154,11 +163,11 @@ public class QueryRoutingTest {
return ((InetSocketAddress) server.getChannel().localAddress()).getPort();
}
- private QueryScheduler mockQueryScheduler(int responseDelayMs, byte[]
responseBytes) {
+ private QueryScheduler mockQueryScheduler(int responseDelayMs, DataBuffer
response) {
QueryScheduler queryScheduler = mock(QueryScheduler.class);
- when(queryScheduler.submit(any())).thenAnswer(invocation -> {
+ when(queryScheduler.submitResponse(any())).thenAnswer(invocation -> {
Thread.sleep(responseDelayMs);
- return Futures.immediateFuture(responseBytes);
+ return Futures.immediateFuture(response);
});
return queryScheduler;
}
@@ -231,6 +240,51 @@ public class QueryRoutingTest {
assertEquals(_serverRoutingStatsManager.fetchNumInFlightRequestsForServer(serverId).intValue(),
0);
}
+ /// The server sends a response that is split into many pages. The broker
must receive the same table.
+ @Test
+ public void testPagedResponse()
+ throws Exception {
+ long requestId = 123;
+ DataSchema dataSchema =
+ new DataSchema(new String[]{"bytes"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BYTES});
+ DataTableBuilder dataTableBuilder =
DataTableBuilderFactory.getDataTableBuilder(dataSchema);
+ int numRows = 200;
+ for (int rowId = 0; rowId < numRows; rowId++) {
+ dataTableBuilder.startRow();
+ byte[] value = new byte[10_000];
+ Arrays.fill(value, (byte) rowId);
+ dataTableBuilder.setColumn(0, new ByteArray(value));
+ dataTableBuilder.finishRow();
+ }
+ DataTable dataTable = dataTableBuilder.build();
+ dataTable.getMetadata().put(MetadataKey.REQUEST_ID.getName(),
Long.toString(requestId));
+ DataBuffer response = dataTable.toDataBuffer();
+ List<ByteBuffer> responseBuffers = new ArrayList<>();
+ response.appendAsByteBuffers(responseBuffers);
+ assertTrue(responseBuffers.size() > 100, "Expected the response in many
pages: " + responseBuffers.size());
+
+ _queryServer = getQueryServer(0, response, 0);
+ initializeTestFixtures(startAndGetPort(_queryServer));
+
+ AsyncQueryResponse asyncQueryResponse =
+ _queryRouter.submitQuery(requestId, "testTable", BROKER_REQUEST,
_routingTable, null, null, 600_000L);
+ Map<ServerRoutingInstance, ServerResponse> responses =
asyncQueryResponse.getFinalResponses();
+ assertEquals(responses.size(), 1);
+ ServerResponse serverResponse =
responses.get(_offlineServerRoutingInstance);
+ assertNotNull(serverResponse);
+ assertEquals(serverResponse.getResponseSize(), response.size());
+ DataTable receivedDataTable = serverResponse.getDataTable();
+ assertNotNull(receivedDataTable);
+ assertEquals(receivedDataTable.getNumberOfRows(), numRows);
+ for (int rowId = 0; rowId < numRows; rowId++) {
+ byte[] expected = new byte[10_000];
+ Arrays.fill(expected, (byte) rowId);
+ assertEquals(receivedDataTable.getBytes(rowId, 0).getBytes(), expected);
+ }
+ _requestCount += 2;
+ waitForStatsUpdate(_requestCount);
+ }
+
@Test
public void testInvalidResponse()
throws Exception {
diff --git
a/pinot-server/src/main/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutor.java
b/pinot-server/src/main/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutor.java
index 84a2706f254..3ea61e403ba 100644
---
a/pinot-server/src/main/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutor.java
+++
b/pinot-server/src/main/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutor.java
@@ -44,6 +44,7 @@ import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUt
import org.apache.pinot.core.query.scheduler.QueryScheduler;
import org.apache.pinot.segment.local.data.manager.SegmentDataManager;
import org.apache.pinot.segment.local.data.manager.TableDataManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.spi.config.table.PageCacheWarmupConfig;
import org.apache.pinot.spi.config.table.QuotaConfig;
import org.apache.pinot.spi.config.table.TableConfig;
@@ -273,8 +274,8 @@ public class PageCacheWarmupServerQueryExecutor {
_serverMetrics.addMeteredGlobalValue(ServerMeter.PAGE_CACHE_WARMUP_QUERIES, 1L);
// Execute query with timeout
- ListenableFuture<byte[]> listenableFuture =
_queryScheduler.submit(serverQueryRequest);
- CompletableFuture<byte[]> queryFuture =
toCompletableFuture(listenableFuture);
+ ListenableFuture<DataBuffer> listenableFuture =
_queryScheduler.submitResponse(serverQueryRequest);
+ CompletableFuture<DataBuffer> queryFuture =
toCompletableFuture(listenableFuture);
CompletableFuture<Void> trackedFuture =
queryFuture.orTimeout(remainingWarmupTimeMs, TimeUnit.MILLISECONDS)
.thenAccept(result -> LOGGER.info("Successfully executed warmup
query={}", query))
.exceptionally(ex -> {
diff --git
a/pinot-server/src/test/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutorTest.java
b/pinot-server/src/test/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutorTest.java
index 9045e6ec347..f4c3d623562 100644
---
a/pinot-server/src/test/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutorTest.java
+++
b/pinot-server/src/test/java/org/apache/pinot/server/warmup/PageCacheWarmupServerQueryExecutorTest.java
@@ -30,6 +30,7 @@ import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUt
import org.apache.pinot.core.query.scheduler.QueryScheduler;
import org.apache.pinot.segment.local.data.manager.SegmentDataManager;
import org.apache.pinot.segment.local.data.manager.TableDataManager;
+import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
import org.apache.pinot.spi.config.table.PageCacheWarmupConfig;
import org.apache.pinot.spi.config.table.QuotaConfig;
import org.apache.pinot.spi.config.table.TableConfig;
@@ -123,8 +124,8 @@ public class PageCacheWarmupServerQueryExecutorTest {
String table = "testTable";
// Mock QueryScheduler to return an already‑completed future
- when(_queryScheduler.submit(any())).thenAnswer(inv ->
- com.google.common.util.concurrent.Futures.immediateFuture(new
byte[0]));
+ when(_queryScheduler.submitResponse(any())).thenAnswer(inv ->
+
com.google.common.util.concurrent.Futures.immediateFuture(PinotByteBuffer.wrap(new
byte[0])));
PageCacheWarmupConfig.Spec warmSpec = new PageCacheWarmupConfig.Spec(true,
2, null, null);
@@ -157,7 +158,7 @@ public class PageCacheWarmupServerQueryExecutorTest {
_warmupServerQueryExecutor.warmupTable(table, warmSpec, queries, null,
warmupQps);
// Submit should be called once per query
- verify(_queryScheduler, times(queries.size())).submit(any());
+ verify(_queryScheduler, times(queries.size())).submitResponse(any());
// No error metric should be recorded
verify(ServerMetrics.get(), never())
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
