deserializeMetadata(ByteBuffer buffer)
+ /**
+ * Deserializes the metadata section from the given byte buffer.
+ *
+ * This is the counterpart of {@link #serializeMetadata(DataOutputStream)} and it is guaranteed that the buffer will
+ * be positioned at the start of the metadata section when this method is called.
+ *
+ * Important: It is mandatory for implementations to leave the cursor at the end of the metadata, in
+ * the exact same position as it was when {@link #serializeMetadata(DataOutputStream)} was called.
+ *
+ * Important: This method will be called at the end of the BaseDataConstructor constructor to read
+ * the metadata section. This means that it will be called before the subclass have been constructor
+ * have been called. Therefore it is not possible to use any subclass fields in this method.
+ */
+ protected void deserializeMetadata(ByteBuffer buffer)
throws IOException {
- return Collections.emptyMap();
+ buffer.getInt();
}
private byte[] serializeExceptions()
@@ -572,14 +585,9 @@ private Map deserializeExceptions(ByteBuffer buffer)
@Override
public String toString() {
if (_dataSchema == null) {
- return _metadata.toString();
+ return "{}";
} else {
- StringBuilder stringBuilder = new StringBuilder();
- stringBuilder.append("resultSchema:").append('\n');
- stringBuilder.append(_dataSchema).append('\n');
- stringBuilder.append("numRows: ").append(_numRows).append('\n');
- stringBuilder.append("metadata: ").append(_metadata.toString()).append('\n');
- return stringBuilder.toString();
+ return "resultSchema:" + '\n' + _dataSchema + '\n' + "numRows: " + _numRows + '\n';
}
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java
index 216f4d9d9131..51b01e6b69bf 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java
@@ -20,6 +20,8 @@
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Objects;
import org.apache.pinot.common.utils.DataSchema;
@@ -27,7 +29,7 @@
* Column-wise data table. It stores data in columnar-major format.
*/
public class ColumnarDataBlock extends BaseDataBlock {
- private static final int VERSION = 1;
+ private static final int VERSION = 2;
protected int[] _cumulativeColumnOffsetSizeInBytes;
protected int[] _columnSizeInBytes;
@@ -80,17 +82,21 @@ protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
}
@Override
- public ColumnarDataBlock toMetadataOnlyDataTable() {
- ColumnarDataBlock metadataOnlyDataTable = new ColumnarDataBlock();
- metadataOnlyDataTable._metadata.putAll(_metadata);
- metadataOnlyDataTable._errCodeToExceptionMap.putAll(_errCodeToExceptionMap);
- return metadataOnlyDataTable;
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof ColumnarDataBlock)) {
+ return false;
+ }
+ ColumnarDataBlock that = (ColumnarDataBlock) o;
+ return Objects.deepEquals(_cumulativeColumnOffsetSizeInBytes, that._cumulativeColumnOffsetSizeInBytes)
+ && Objects.deepEquals(_columnSizeInBytes, that._columnSizeInBytes);
}
@Override
- public ColumnarDataBlock toDataOnlyDataTable() {
- return new ColumnarDataBlock(_numRows, _dataSchema, _stringDictionary, _fixedSizeDataBytes, _variableSizeDataBytes);
+ public int hashCode() {
+ return Objects.hash(Arrays.hashCode(_cumulativeColumnOffsetSizeInBytes), Arrays.hashCode(_columnSizeInBytes));
}
-
- // TODO: add whole-column access methods.
+// TODO: add whole-column access methods.
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java
index 418426b4ac6b..768d6cdb885d 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java
@@ -80,10 +80,6 @@ byte[] toBytes()
@Nullable
RoaringBitmap getNullRowIds(int colId);
- DataBlock toMetadataOnlyDataTable();
-
- DataBlock toDataOnlyDataTable();
-
enum Type {
ROW(0),
COLUMNAR(1),
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java
index 27f114032849..1a67f980b115 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java
@@ -26,15 +26,20 @@
import org.apache.pinot.common.response.ProcessingException;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
public final class DataBlockUtils {
+ private static final Logger LOGGER = LoggerFactory.getLogger(DataBlockUtils.class);
+
private DataBlockUtils() {
}
static final int VERSION_TYPE_SHIFT = 5;
public static MetadataBlock getErrorDataBlock(Exception e) {
+ LOGGER.info("Caught exception while processing query", e);
if (e instanceof ProcessingException) {
return getErrorDataBlock(Collections.singletonMap(((ProcessingException) e).getErrorCode(), extractErrorMsg(e)));
} else {
@@ -51,33 +56,40 @@ private static String extractErrorMsg(Throwable t) {
}
public static MetadataBlock getErrorDataBlock(Map exceptions) {
- MetadataBlock errorBlock = new MetadataBlock(MetadataBlock.MetadataBlockType.ERROR);
- for (Map.Entry exception : exceptions.entrySet()) {
- errorBlock.addException(exception.getKey(), exception.getValue());
- }
- return errorBlock;
+ return MetadataBlock.newError(exceptions);
+ }
+
+ /**
+ * Reads an integer from the given byte buffer.
+ *
+ * The returned integer contains both the version and the type of the data block.
+ * {@link #getVersion(int)} and {@link #getType(int)} can be used to extract the version and the type.
+ * @param byteBuffer byte buffer to read from. A single int will be read
+ */
+ public static int readVersionType(ByteBuffer byteBuffer) {
+ return byteBuffer.getInt();
}
- public static MetadataBlock getEndOfStreamDataBlock() {
- return new MetadataBlock(MetadataBlock.MetadataBlockType.EOS);
+ public static int getVersion(int versionType) {
+ return versionType & ((1 << VERSION_TYPE_SHIFT) - 1);
}
- public static MetadataBlock getEndOfStreamDataBlock(Map stats) {
- return new MetadataBlock(MetadataBlock.MetadataBlockType.EOS, stats);
+ public static DataBlock.Type getType(int versionType) {
+ return DataBlock.Type.fromOrdinal(versionType >> VERSION_TYPE_SHIFT);
}
public static DataBlock getDataBlock(ByteBuffer byteBuffer)
throws IOException {
- int versionType = byteBuffer.getInt();
- int version = versionType & ((1 << VERSION_TYPE_SHIFT) - 1);
- DataBlock.Type type = DataBlock.Type.fromOrdinal(versionType >> VERSION_TYPE_SHIFT);
+ int versionType = readVersionType(byteBuffer);
+ int version = getVersion(versionType);
+ DataBlock.Type type = getType(versionType);
switch (type) {
case COLUMNAR:
return new ColumnarDataBlock(byteBuffer);
case ROW:
return new RowDataBlock(byteBuffer);
case METADATA:
- return new MetadataBlock(byteBuffer);
+ return MetadataBlock.deserialize(byteBuffer, version);
default:
throw new UnsupportedOperationException("Unsupported data table version: " + version + " with type: " + type);
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
index 41c2e9d0ea8c..164ca5157701 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
@@ -18,16 +18,19 @@
*/
package org.apache.pinot.common.datablock;
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.annotations.VisibleForTesting;
+import java.io.DataOutputStream;
import java.io.IOException;
+import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
-import java.util.HashMap;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
import java.util.Map;
+import java.util.Objects;
+import javax.annotation.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -36,108 +39,116 @@
*/
public class MetadataBlock extends BaseDataBlock {
- private static final ObjectMapper JSON = new ObjectMapper();
-
+ private static final Logger LOGGER = LoggerFactory.getLogger(MetadataBlock.class);
@VisibleForTesting
- static final int VERSION = 1;
-
- public enum MetadataBlockType {
- /**
- * Indicates that this block is the final block to be sent
- * (End Of Stream) as part of an operator chain computation.
- */
- EOS,
+ static final int VERSION = 2;
+ @Nullable
+ private List _statsByStage;
- /**
- * An {@code ERROR} metadata block indicates that there was
- * some error during computation. To retrieve the error that
- * occurred, use {@link MetadataBlock#getExceptions()}
- */
- ERROR
+ private MetadataBlock() {
+ this(Collections.emptyList());
}
- /**
- * Used to serialize the contents of the metadata block conveniently and in
- * a backwards compatible way. Use JSON because the performance of metadata block
- * SerDe should not be a bottleneck.
- */
- @JsonIgnoreProperties(ignoreUnknown = true)
- @VisibleForTesting
- static class Contents {
-
- private String _type;
- private Map _stats;
-
- @JsonCreator
- public Contents(@JsonProperty("type") String type, @JsonProperty("stats") Map stats) {
- _type = type;
- _stats = stats;
- }
-
- @JsonCreator
- public Contents() {
- this(null, new HashMap<>());
- }
-
- public String getType() {
- return _type;
- }
-
- public void setType(String type) {
- _type = type;
- }
-
- public Map getStats() {
- return _stats;
- }
+ public static MetadataBlock newEos() {
+ return new MetadataBlock();
+ }
- public void setStats(Map stats) {
- _stats = stats;
+ public static MetadataBlock newError(Map exceptions) {
+ MetadataBlock errorBlock = new MetadataBlock();
+ for (Map.Entry exception : exceptions.entrySet()) {
+ errorBlock.addException(exception.getKey(), exception.getValue());
}
+ return errorBlock;
}
- private final Contents _contents;
+ public MetadataBlock(List statsByStage) {
+ super(0, null, new String[0], new byte[0], new byte[0]);
+ _statsByStage = statsByStage;
+ }
- public MetadataBlock(MetadataBlockType type) {
- this(type, new HashMap<>());
+ MetadataBlock(ByteBuffer byteBuffer)
+ throws IOException {
+ super(byteBuffer);
}
- public MetadataBlock(MetadataBlockType type, Map stats) {
- super(0, null, new String[0], new byte[]{0}, toContents(new Contents(type.name(), stats)));
- _contents = new Contents(type.name(), stats);
+ @Override
+ protected void serializeMetadata(DataOutputStream output)
+ throws IOException {
+ if (_statsByStage == null) {
+ output.writeInt(0);
+ return;
+ }
+ int size = _statsByStage.size();
+ output.writeInt(size);
+ if (size > 0) {
+ byte[] bytes = new byte[4096];
+ for (ByteBuffer stat : _statsByStage) {
+ if (stat == null) {
+ output.writeBoolean(false);
+ } else {
+ output.writeBoolean(true);
+ output.writeInt(stat.remaining());
+ ByteBuffer duplicate = stat.duplicate();
+ while (duplicate.hasRemaining()) {
+ int length = Math.min(duplicate.remaining(), bytes.length);
+ duplicate.get(bytes, 0, length);
+ output.write(bytes, 0, length);
+ }
+ }
+ }
+ }
}
- private static byte[] toContents(Contents type) {
- try {
- return JSON.writeValueAsBytes(type);
- } catch (JsonProcessingException e) {
- throw new RuntimeException(e);
+ public static MetadataBlock deserialize(ByteBuffer byteBuffer, int version)
+ throws IOException {
+ switch (version) {
+ case 1:
+ case 2:
+ return new MetadataBlock(byteBuffer);
+ default:
+ throw new IOException("Unsupported metadata block version: " + version);
}
}
- public MetadataBlock(ByteBuffer byteBuffer)
+ @Override
+ protected void deserializeMetadata(ByteBuffer buffer)
throws IOException {
- super(byteBuffer);
- if (_variableSizeDataBytes != null && _variableSizeDataBytes.length > 0) {
- _contents = JSON.readValue(_variableSizeDataBytes, Contents.class);
- } else {
- _contents = new Contents();
+ try {
+ int statsSize = buffer.getInt();
+
+ List stats = new ArrayList<>(statsSize);
+
+ for (int i = 0; i < statsSize; i++) {
+ if (buffer.get() != 0) {
+ int length = buffer.getInt();
+ buffer.limit(buffer.position() + length);
+ stats.add(buffer.slice());
+ buffer.position(buffer.limit());
+ buffer.limit(buffer.capacity());
+ } else {
+ stats.add(null);
+ }
+ }
+ _statsByStage = stats;
+ } catch (BufferUnderflowException e) {
+ LOGGER.info("Failed to read stats from metadata block. Considering it empty", e);;
+ } catch (RuntimeException e) {
+ LOGGER.warn("Failed to read stats from metadata block. Considering it empty", e);;
}
}
public MetadataBlockType getType() {
- String type = _contents.getType();
-
- // if type is null, then we're reading a legacy block where we didn't encode any
- // data. assume that it is an EOS block if there's no exceptions and an ERROR block
- // otherwise
- return type == null
- ? (getExceptions().isEmpty() ? MetadataBlockType.EOS : MetadataBlockType.ERROR)
- : MetadataBlockType.valueOf(type);
+ return _errCodeToExceptionMap.isEmpty() ? MetadataBlockType.EOS : MetadataBlockType.ERROR;
}
- public Map getStats() {
- return _contents.getStats() != null ? _contents.getStats() : new HashMap<>();
+ /**
+ * Returns the list of serialized stats.
+ *
+ * The returned list may contain nulls, which would mean that no stats were available for that stage.
+ */
+ @Nullable
+ public List getStatsByStage() {
+ return _statsByStage;
}
@Override
@@ -156,12 +167,35 @@ protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
}
@Override
- public MetadataBlock toMetadataOnlyDataTable() {
- return this;
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof MetadataBlock)) {
+ return false;
+ }
+ MetadataBlock that = (MetadataBlock) o;
+ return Objects.equals(_statsByStage, that._statsByStage)
+ && _errCodeToExceptionMap.equals(that._errCodeToExceptionMap);
}
@Override
- public MetadataBlock toDataOnlyDataTable() {
- throw new UnsupportedOperationException();
+ public int hashCode() {
+ return Objects.hash(_statsByStage, _errCodeToExceptionMap);
+ }
+
+ public enum MetadataBlockType {
+ /**
+ * Indicates that this block is the final block to be sent
+ * (End Of Stream) as part of an operator chain computation.
+ */
+ EOS,
+
+ /**
+ * An {@code ERROR} metadata block indicates that there was
+ * some error during computation. To retrieve the error that
+ * occurred, use {@link MetadataBlock#getExceptions()}
+ */
+ ERROR
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java
index f5aa80648936..33a5e77cf764 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java
@@ -20,6 +20,8 @@
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Objects;
import org.apache.pinot.common.utils.DataSchema;
@@ -27,7 +29,7 @@
* Wrapper for row-wise data table. It stores data in row-major format.
*/
public class RowDataBlock extends BaseDataBlock {
- private static final int VERSION = 1;
+ private static final int VERSION = 2;
protected int[] _columnOffsets;
protected int _rowSizeInBytes;
@@ -72,22 +74,25 @@ protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
return _fixedSizeData.getInt(offset + 4);
}
- @Override
- public RowDataBlock toMetadataOnlyDataTable() {
- RowDataBlock metadataOnlyDataTable = new RowDataBlock();
- metadataOnlyDataTable._metadata.putAll(_metadata);
- metadataOnlyDataTable._errCodeToExceptionMap.putAll(_errCodeToExceptionMap);
- return metadataOnlyDataTable;
+ public int getRowSizeInBytes() {
+ return _rowSizeInBytes;
}
@Override
- public RowDataBlock toDataOnlyDataTable() {
- return new RowDataBlock(_numRows, _dataSchema, _stringDictionary, _fixedSizeDataBytes, _variableSizeDataBytes);
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof RowDataBlock)) {
+ return false;
+ }
+ RowDataBlock that = (RowDataBlock) o;
+ return _rowSizeInBytes == that._rowSizeInBytes && Objects.deepEquals(_columnOffsets, that._columnOffsets);
}
- public int getRowSizeInBytes() {
- return _rowSizeInBytes;
+ @Override
+ public int hashCode() {
+ return Objects.hash(Arrays.hashCode(_columnOffsets), _rowSizeInBytes);
}
-
- // TODO: add whole-row access methods.
+// TODO: add whole-row access methods.
}
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 72121611eb04..5fb1018dceae 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
@@ -110,6 +110,8 @@ enum MetadataKey {
NUM_SEGMENTS_PROCESSED(6, "numSegmentsProcessed", MetadataValueType.INT),
NUM_SEGMENTS_MATCHED(7, "numSegmentsMatched", MetadataValueType.INT),
NUM_CONSUMING_SEGMENTS_QUERIED(8, "numConsumingSegmentsQueried", MetadataValueType.INT),
+ // the timestamp indicating the freshness of the data queried in consuming segments.
+ // This can be ingestion timestamp if provided by the stream, or the last index time
MIN_CONSUMING_FRESHNESS_TIME_MS(9, "minConsumingFreshnessTimeMs", MetadataValueType.LONG),
TOTAL_DOCS(10, "totalDocs", MetadataValueType.LONG),
NUM_GROUPS_LIMIT_REACHED(11, "numGroupsLimitReached", MetadataValueType.STRING),
@@ -135,6 +137,7 @@ enum MetadataKey {
OPERATOR_ID(31, "operatorId", MetadataValueType.STRING),
OPERATOR_EXEC_START_TIME_MS(32, "operatorExecStartTimeMs", MetadataValueType.LONG),
OPERATOR_EXEC_END_TIME_MS(33, "operatorExecEndTimeMs", MetadataValueType.LONG),
+ // Not actually used
MAX_ROWS_IN_JOIN_REACHED(34, "maxRowsInJoinReached", MetadataValueType.STRING);
// We keep this constant to track the max id added so far for backward compatibility.
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datatable/StatMap.java b/pinot-common/src/main/java/org/apache/pinot/common/datatable/StatMap.java
new file mode 100644
index 000000000000..9d2818958bb7
--- /dev/null
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datatable/StatMap.java
@@ -0,0 +1,498 @@
+/**
+ * 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.common.datatable;
+
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.base.Preconditions;
+import java.io.DataInput;
+import java.io.DataOutput;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.EnumMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.annotation.Nullable;
+import org.apache.pinot.spi.utils.JsonUtils;
+
+
+/**
+ * A map that stores statistics.
+ *
+ * Statistics must be keyed by an enum that implements {@link StatMap.Key}.
+ *
+ * A stat map efficiently store, serialize and deserialize these statistics.
+ *
+ * Serialization and deserialization is backward and forward compatible as long as the only change in the keys are:
+ *
+ * - Adding new keys
+ * - Change the name of the keys
+ *
+ *
+ * Any other change (like changing the type of key, changing their literal order are not supported or removing keys)
+ * are backward incompatible changes.
+ * @param
+ */
+public class StatMap & StatMap.Key> {
+ private final Class _keyClass;
+ private final Map _map;
+
+ private static final ConcurrentHashMap, Object[]> KEYS_BY_CLASS = new ConcurrentHashMap<>();
+
+ public StatMap(Class keyClass) {
+ _keyClass = keyClass;
+ // TODO: Study whether this is fine or we should impose a single thread policy in StatMaps
+ _map = Collections.synchronizedMap(new EnumMap<>(keyClass));
+ }
+
+ public int getInt(K key) {
+ Preconditions.checkArgument(key.getType() == Type.INT, "Key %s is of type %s, not INT", key, key.getType());
+ Object o = _map.get(key);
+ return o == null ? 0 : (Integer) o;
+ }
+
+ public StatMap merge(K key, int value) {
+ if (key.getType() == Type.LONG) {
+ merge(key, (long) value);
+ return this;
+ }
+ int oldValue = getInt(key);
+ int newValue = key.merge(oldValue, value);
+ if (newValue == 0) {
+ _map.remove(key);
+ } else {
+ _map.put(key, newValue);
+ }
+ return this;
+ }
+
+ public long getLong(K key) {
+ if (key.getType() == Type.INT) {
+ return getInt(key);
+ }
+ Preconditions.checkArgument(key.getType() == Type.LONG, "Key %s is of type %s, not LONG", key, key.getType());
+ Object o = _map.get(key);
+ return o == null ? 0L : (Long) o;
+ }
+
+ public StatMap merge(K key, long value) {
+ Preconditions.checkArgument(key.getType() == Type.LONG, "Key %s is of type %s, not LONG", key, key.getType());
+ long oldValue = getLong(key);
+ long newValue = key.merge(oldValue, value);
+ if (newValue == 0) {
+ _map.remove(key);
+ } else {
+ _map.put(key, newValue);
+ }
+ return this;
+ }
+
+ public boolean getBoolean(K key) {
+ Preconditions.checkArgument(key.getType() == Type.BOOLEAN, "Key %s is of type %s, not BOOLEAN",
+ key, key.getType());
+ Object o = _map.get(key);
+ return o != null && (Boolean) o;
+ }
+
+ public StatMap merge(K key, boolean value) {
+ boolean oldValue = getBoolean(key);
+ boolean newValue = key.merge(oldValue, value);
+ if (!newValue) {
+ _map.remove(key);
+ } else {
+ _map.put(key, Boolean.TRUE);
+ }
+ return this;
+ }
+
+ public String getString(K key) {
+ Preconditions.checkArgument(key.getType() == Type.STRING, "Key %s is of type %s, not STRING", key, key.getType());
+ Object o = _map.get(key);
+ return o == null ? null : (String) o;
+ }
+
+ public StatMap merge(K key, String value) {
+ String oldValue = getString(key);
+ String newValue = key.merge(oldValue, value);
+ if (newValue == null) {
+ _map.remove(key);
+ } else {
+ _map.put(key, newValue);
+ }
+ return this;
+ }
+
+ /**
+ * Returns the value associated with the key.
+ *
+ * Primitives will be boxed, so it is recommended to use the specific methods for each type.
+ */
+ public Object getAny(K key) {
+ switch (key.getType()) {
+ case BOOLEAN:
+ return getBoolean(key);
+ case INT:
+ return getInt(key);
+ case LONG:
+ return getLong(key);
+ case STRING:
+ return getString(key);
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+
+ /**
+ * Modifies this object to merge the values of the other object.
+ *
+ * @param other The object to merge with. This argument will not be modified.
+ * @return this object once it is modified.
+ */
+ public StatMap merge(StatMap other) {
+ Preconditions.checkState(_keyClass.equals(other._keyClass),
+ "Different key classes %s and %s", _keyClass, other._keyClass);
+ for (Map.Entry entry : other._map.entrySet()) {
+ K key = entry.getKey();
+ Object value = entry.getValue();
+ if (value == null) {
+ continue;
+ }
+ switch (key.getType()) {
+ case BOOLEAN:
+ merge(key, (boolean) value);
+ break;
+ case INT:
+ merge(key, (int) value);
+ break;
+ case LONG:
+ merge(key, (long) value);
+ break;
+ case STRING:
+ merge(key, (String) value);
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+ return this;
+ }
+
+ public StatMap merge(DataInput input)
+ throws IOException {
+ byte serializedKeys = input.readByte();
+
+ K[] keys = (K[]) KEYS_BY_CLASS.computeIfAbsent(_keyClass, k -> k.getEnumConstants());
+ for (byte i = 0; i < serializedKeys; i++) {
+ int ordinal = input.readByte();
+ K key = keys[ordinal];
+ switch (key.getType()) {
+ case BOOLEAN:
+ merge(key, true);
+ break;
+ case INT:
+ merge(key, input.readInt());
+ break;
+ case LONG:
+ merge(key, input.readLong());
+ break;
+ case STRING:
+ merge(key, input.readUTF());
+ break;
+ default:
+ throw new IllegalStateException("Unknown type " + key.getType());
+ }
+ }
+ return this;
+ }
+
+ public ObjectNode asJson() {
+ ObjectNode node = JsonUtils.newObjectNode();
+
+ for (Map.Entry entry : _map.entrySet()) {
+ K key = entry.getKey();
+ Object value = entry.getValue();
+ switch (key.getType()) {
+ case BOOLEAN:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), false);
+ }
+ } else {
+ node.put(key.getStatName(), (boolean) value);
+ }
+ break;
+ case INT:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), 0);
+ }
+ } else {
+ node.put(key.getStatName(), (int) value);
+ }
+ break;
+ case LONG:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), 0L);
+ }
+ } else {
+ node.put(key.getStatName(), (long) value);
+ }
+ break;
+ case STRING:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), "");
+ }
+ } else {
+ node.put(key.getStatName(), (String) value);
+ }
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+
+ return node;
+ }
+
+ public void serialize(DataOutput output)
+ throws IOException {
+
+ assert checkContainsNoDefault() : "No default value should be stored in the map";
+ output.writeByte(_map.size());
+
+ // We use written keys just to fail fast in tests if the number of keys written
+ // is not the same as the number of keys
+ int writtenKeys = 0;
+ K[] keys = (K[]) KEYS_BY_CLASS.computeIfAbsent(_keyClass, k -> k.getEnumConstants());
+ for (int ordinal = 0; ordinal < keys.length; ordinal++) {
+ K key = keys[ordinal];
+ switch (key.getType()) {
+ case BOOLEAN: {
+ if (getBoolean(key)) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ }
+ break;
+ }
+ case INT: {
+ int value = getInt(key);
+ if (value != 0) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ output.writeInt(value);
+ }
+ break;
+ }
+ case LONG: {
+ long value = getLong(key);
+ if (value != 0) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ output.writeLong(value);
+ }
+ break;
+ }
+ case STRING: {
+ String value = getString(key);
+ if (value != null) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ output.writeUTF(value);
+ }
+ break;
+ }
+ default:
+ throw new IllegalStateException("Unknown type " + key.getType());
+ }
+ }
+ assert writtenKeys == _map.size() : "Written keys " + writtenKeys + " but map size " + _map.size();
+ }
+
+ private boolean checkContainsNoDefault() {
+ for (Map.Entry entry : _map.entrySet()) {
+ K key = entry.getKey();
+ Object value = entry.getValue();
+ switch (key.getType()) {
+ case BOOLEAN:
+ if (value == null || !(boolean) value) {
+ throw new IllegalStateException("Boolean value must be true but " + value + " is stored for key " + key);
+ }
+ break;
+ case INT:
+ if (value == null || (int) value == 0) {
+ throw new IllegalStateException("Int value must be non-zero but " + value + " is stored for key " + key);
+ }
+ break;
+ case LONG:
+ if (value == null || (long) value == 0) {
+ throw new IllegalStateException("Long value must be non-zero but " + value + " is stored for key " + key);
+ }
+ break;
+ case STRING:
+ if (value == null) {
+ throw new IllegalStateException("String value must be non-null but null is stored for key " + key);
+ }
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+ return true;
+ }
+
+ public static String getDefaultStatName(Key key) {
+ String name = key.name();
+ StringBuilder result = new StringBuilder();
+ boolean capitalizeNext = false;
+
+ for (char c : name.toCharArray()) {
+ if (c == '_') {
+ capitalizeNext = true;
+ } else {
+ if (capitalizeNext) {
+ result.append(c);
+ capitalizeNext = false;
+ } else {
+ result.append(Character.toLowerCase(c));
+ }
+ }
+ }
+
+ return result.toString();
+ }
+
+ public static & Key> StatMap deserialize(DataInput input, Class keyClass)
+ throws IOException {
+ StatMap result = new StatMap<>(keyClass);
+ result.merge(input);
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ StatMap statMap = (StatMap) o;
+ return Objects.equals(_map, statMap._map);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(_map);
+ }
+
+ @Override
+ public String toString() {
+ return asJson().toString();
+ }
+
+ public Class getKeyClass() {
+ return _keyClass;
+ }
+
+ public boolean isEmpty() {
+ return _map.isEmpty();
+ }
+
+ public Iterable keySet() {
+ return _map.keySet();
+ }
+
+ public interface Key {
+ String name();
+
+ /**
+ * The name of the stat used to report it. Names must be unique on the same key family.
+ */
+ default String getStatName() {
+ return getDefaultStatName(this);
+ }
+
+ default int merge(int value1, int value2) {
+ return value1 + value2;
+ }
+
+ default long merge(long value1, long value2) {
+ return value1 + value2;
+ }
+
+ default boolean merge(boolean value1, boolean value2) {
+ return value1 || value2;
+ }
+
+ default String merge(@Nullable String value1, @Nullable String value2) {
+ return value2 != null ? value2 : value1;
+ }
+
+ /**
+ * The type of the values associated to this key.
+ */
+ Type getType();
+
+ default boolean includeDefaultInJson() {
+ return false;
+ }
+
+ static int minPositive(int value1, int value2) {
+ if (value1 == 0 && value2 >= 0) {
+ return value2;
+ }
+ if (value2 == 0 && value1 >= 0) {
+ return value1;
+ }
+ return Math.min(value1, value2);
+ }
+
+ static long minPositive(long value1, long value2) {
+ if (value1 == 0 && value2 >= 0) {
+ return value2;
+ }
+ if (value2 == 0 && value1 >= 0) {
+ return value1;
+ }
+ return Math.min(value1, value2);
+ }
+
+ static int eqNotZero(int value1, int value2) {
+ if (value1 != value2) {
+ if (value1 == 0) {
+ return value2;
+ } else if (value2 == 0) {
+ return value1;
+ } else {
+ throw new IllegalStateException("Cannot merge non-zero values: " + value1 + " and " + value2);
+ }
+ }
+ return value1;
+ }
+ }
+
+ public enum Type {
+ BOOLEAN,
+ INT,
+ LONG,
+ STRING
+ }
+}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java b/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java
index 06f65fc2aa25..0d7ec268a749 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java
@@ -18,9 +18,12 @@
*/
package org.apache.pinot.common.response;
+import java.io.IOException;
import java.util.List;
+import javax.annotation.Nullable;
import org.apache.pinot.common.response.broker.QueryProcessingException;
import org.apache.pinot.common.response.broker.ResultTable;
+import org.apache.pinot.spi.utils.JsonUtils;
/**
@@ -33,6 +36,8 @@ public interface BrokerResponse {
*/
void setExceptions(List exceptions);
+ void addToExceptions(QueryProcessingException processingException);
+
/**
* Set the number of servers got queried by the broker.
*
@@ -47,21 +52,20 @@ public interface BrokerResponse {
*/
void setNumServersResponded(int numServersResponded);
+ long getTimeUsedMs();
+
/**
* Set the total time used in request handling, into the broker response.
*/
void setTimeUsedMs(long timeUsedMs);
- /**
- * Set the total number of rows in result set
- */
- void setNumRowsResultSet(int numRowsResultSet);
-
/**
* Convert the broker response to JSON String.
*/
- String toJsonString()
- throws Exception;
+ default String toJsonString()
+ throws IOException {
+ return JsonUtils.objectToString(this);
+ }
/**
* Returns the number of servers queried.
@@ -147,12 +151,13 @@ String toJsonString()
* set the result table.
* @param resultTable result table to be set.
*/
- void setResultTable(ResultTable resultTable);
+ void setResultTable(@Nullable ResultTable resultTable);
/**
* Get the result table.
* @return result table.
*/
+ @Nullable
ResultTable getResultTable();
/**
@@ -163,12 +168,10 @@ String toJsonString()
/**
* Get the total number of rows in result set
*/
- int getNumRowsResultSet();
-
- /**
- * Set the total thread cpu time used against offline table in request handling, into the broker response.
- */
- void setOfflineThreadCpuTimeNs(long offlineThreadCpuTimeNs);
+ default int getNumRowsResultSet() {
+ ResultTable resultTable = getResultTable();
+ return resultTable == null ? 0 : resultTable.getRows().size();
+ }
/**
* Get the thread cpu time used against offline table in request handling, from the broker response.
@@ -180,74 +183,43 @@ String toJsonString()
*/
long getRealtimeThreadCpuTimeNs();
- /**
- * Set the total thread cpu time used against realtime table in request handling, into the broker response.
- */
- void setRealtimeThreadCpuTimeNs(long realtimeThreadCpuTimeNs);
-
/**
* Get the system activities cpu time used against offline table in request handling, from the broker response.
*/
long getOfflineSystemActivitiesCpuTimeNs();
- /**
- * Set the system activities cpu time used against offline table in request handling, into the broker response.
- */
- void setOfflineSystemActivitiesCpuTimeNs(long offlineSystemActivitiesCpuTimeNs);
-
/**
* Get the system activities cpu time used against realtime table in request handling, from the broker response.
*/
long getRealtimeSystemActivitiesCpuTimeNs();
- /**
- * Set the system activities cpu time used against realtime table in request handling, into the broker response.
- */
- void setRealtimeSystemActivitiesCpuTimeNs(long realtimeSystemActivitiesCpuTimeNs);
-
/**
* Get the response serialization cpu time used against offline table in request handling, from the broker response.
*/
long getOfflineResponseSerializationCpuTimeNs();
- /**
- * Set the response serialization cpu time used against offline table in request handling, into the broker response.
- */
- void setOfflineResponseSerializationCpuTimeNs(long offlineResponseSerializationCpuTimeNs);
-
/**
* Get the response serialization cpu time used against realtime table in request handling, from the broker response.
*/
long getRealtimeResponseSerializationCpuTimeNs();
- /**
- * Set the response serialization cpu time used against realtime table in request handling, into the broker response.
- */
- void setRealtimeResponseSerializationCpuTimeNs(long realtimeResponseSerializationCpuTimeNs);
-
/**
* Get the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
* against offline table in request handling, from the broker response.
*/
- long getOfflineTotalCpuTimeNs();
-
- /**
- * Set the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
- * against offline table in request handling, into the broker response.
- */
- void setOfflineTotalCpuTimeNs(long offlineTotalCpuTimeNs);
+ default long getOfflineTotalCpuTimeNs() {
+ return getOfflineThreadCpuTimeNs() + getOfflineSystemActivitiesCpuTimeNs()
+ + getOfflineResponseSerializationCpuTimeNs();
+ }
/**
* Get the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
* against realtime table in request handling, from the broker response.
*/
- long getRealtimeTotalCpuTimeNs();
-
- /**
- * Set the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
- * against realtime table in request handling, into the broker response.
- */
- void setRealtimeTotalCpuTimeNs(long realtimeTotalCpuTimeNs);
+ default long getRealtimeTotalCpuTimeNs() {
+ return getRealtimeThreadCpuTimeNs() + getRealtimeSystemActivitiesCpuTimeNs()
+ + getRealtimeResponseSerializationCpuTimeNs();
+ }
/**
* Get the total number of segments pruned on the Broker side
@@ -264,11 +236,6 @@ String toJsonString()
*/
long getNumSegmentsPrunedByServer();
- /**
- * Set the total number of segments pruned on the Server side
- */
- void setNumSegmentsPrunedByServer(long numSegmentsPrunedByServer);
-
/**
* Get the total number of segments pruned due to invalid data or schema.
*
@@ -276,13 +243,6 @@ String toJsonString()
*/
long getNumSegmentsPrunedInvalid();
- /**
- * Set the total number of segments pruned due to invalid data or schema.
- *
- * This value is always lower or equal than {@link #getNumSegmentsPrunedByServer()}
- */
- void setNumSegmentsPrunedInvalid(long numSegmentsPrunedInvalid);
-
/**
* Get the total number of segments pruned by applying the limit optimization.
*
@@ -290,13 +250,6 @@ String toJsonString()
*/
long getNumSegmentsPrunedByLimit();
- /**
- * Set the total number of segments pruned by applying the limit optimization.
- *
- * This value is always lower or equal than {@link #getNumSegmentsPrunedByServer()}
- */
- void setNumSegmentsPrunedByLimit(long numSegmentsPrunedByLimit);
-
/**
* Get the total number of segments pruned applying value optimizations, like bloom filters.
*
@@ -304,33 +257,16 @@ String toJsonString()
*/
long getNumSegmentsPrunedByValue();
- /**
- * Set the total number of segments pruned applying value optimizations, like bloom filters.
- *
- * This value is always lower or equal than {@link #getNumSegmentsPrunedByServer()}
- */
- void setNumSegmentsPrunedByValue(long numSegmentsPrunedByValue);
-
/**
* Get the total number of segments with an EmptyFilterOperator when Explain Plan is called
*/
long getExplainPlanNumEmptyFilterSegments();
- /**
- * Set the total number of segments with an EmptyFilterOperator when Explain Plan is called
- */
- void setExplainPlanNumEmptyFilterSegments(long explainPlanNumEmptyFilterSegments);
-
/**
* Get the total number of segments with a MatchAllFilterOperator when Explain Plan is called
*/
long getExplainPlanNumMatchAllFilterSegments();
- /**
- * Set the total number of segments with a MatchAllFilterOperator when Explain Plan is called
- */
- void setExplainPlanNumMatchAllFilterSegments(long explainPlanNumMatchAllFilterSegments);
-
/**
* get request ID for the query
*/
@@ -354,4 +290,6 @@ String toJsonString()
long getBrokerReduceTimeMs();
void setBrokerReduceTimeMs(long brokerReduceTimeMs);
+
+ boolean isPartialResult();
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java
index 9fe098e26d51..d22ee51c89e1 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java
@@ -22,6 +22,7 @@
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
+import com.google.common.annotations.VisibleForTesting;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
@@ -37,7 +38,7 @@
/**
* This class implements pinot-broker's response format for any given query.
* All fields either primitive data types, or native objects (as opposed to JSONObjects).
- *
+ *
* Supports serialization via JSON.
*/
@JsonPropertyOrder({
@@ -78,8 +79,6 @@ public class BrokerResponseNative implements BrokerResponse {
private long _totalDocs = 0L;
private boolean _numGroupsLimitReached = false;
- private boolean _maxRowsInJoinReached = false;
- private boolean _partialResult = false;
private long _timeUsedMs = 0L;
private long _offlineThreadCpuTimeNs = 0L;
private long _realtimeThreadCpuTimeNs = 0L;
@@ -131,6 +130,7 @@ public static BrokerResponseNative empty() {
return new BrokerResponseNative();
}
+ @VisibleForTesting
public static BrokerResponseNative fromJsonString(String jsonString)
throws IOException {
return JsonUtils.stringToObject(jsonString, BrokerResponseNative.class);
@@ -143,7 +143,6 @@ public long getOfflineSystemActivitiesCpuTimeNs() {
}
@JsonProperty("offlineSystemActivitiesCpuTimeNs")
- @Override
public void setOfflineSystemActivitiesCpuTimeNs(long offlineSystemActivitiesCpuTimeNs) {
_offlineSystemActivitiesCpuTimeNs = offlineSystemActivitiesCpuTimeNs;
}
@@ -155,7 +154,6 @@ public long getRealtimeSystemActivitiesCpuTimeNs() {
}
@JsonProperty("realtimeSystemActivitiesCpuTimeNs")
- @Override
public void setRealtimeSystemActivitiesCpuTimeNs(long realtimeSystemActivitiesCpuTimeNs) {
_realtimeSystemActivitiesCpuTimeNs = realtimeSystemActivitiesCpuTimeNs;
}
@@ -167,7 +165,6 @@ public long getOfflineThreadCpuTimeNs() {
}
@JsonProperty("offlineThreadCpuTimeNs")
- @Override
public void setOfflineThreadCpuTimeNs(long timeUsedMs) {
_offlineThreadCpuTimeNs = timeUsedMs;
}
@@ -179,7 +176,6 @@ public long getRealtimeThreadCpuTimeNs() {
}
@JsonProperty("realtimeThreadCpuTimeNs")
- @Override
public void setRealtimeThreadCpuTimeNs(long timeUsedMs) {
_realtimeThreadCpuTimeNs = timeUsedMs;
}
@@ -191,7 +187,6 @@ public long getOfflineResponseSerializationCpuTimeNs() {
}
@JsonProperty("offlineResponseSerializationCpuTimeNs")
- @Override
public void setOfflineResponseSerializationCpuTimeNs(long offlineResponseSerializationCpuTimeNs) {
_offlineResponseSerializationCpuTimeNs = offlineResponseSerializationCpuTimeNs;
}
@@ -203,7 +198,6 @@ public long getRealtimeResponseSerializationCpuTimeNs() {
}
@JsonProperty("realtimeResponseSerializationCpuTimeNs")
- @Override
public void setRealtimeResponseSerializationCpuTimeNs(long realtimeResponseSerializationCpuTimeNs) {
_realtimeResponseSerializationCpuTimeNs = realtimeResponseSerializationCpuTimeNs;
}
@@ -215,7 +209,6 @@ public long getOfflineTotalCpuTimeNs() {
}
@JsonProperty("offlineTotalCpuTimeNs")
- @Override
public void setOfflineTotalCpuTimeNs(long offlineTotalCpuTimeNs) {
_offlineTotalCpuTimeNs = offlineTotalCpuTimeNs;
}
@@ -227,7 +220,6 @@ public long getRealtimeTotalCpuTimeNs() {
}
@JsonProperty("realtimeTotalCpuTimeNs")
- @Override
public void setRealtimeTotalCpuTimeNs(long realtimeTotalCpuTimeNs) {
_realtimeTotalCpuTimeNs = realtimeTotalCpuTimeNs;
}
@@ -251,7 +243,6 @@ public long getNumSegmentsPrunedByServer() {
}
@JsonProperty("numSegmentsPrunedByServer")
- @Override
public void setNumSegmentsPrunedByServer(long numSegmentsPrunedByServer) {
_numSegmentsPrunedByServer = numSegmentsPrunedByServer;
}
@@ -263,7 +254,6 @@ public long getNumSegmentsPrunedInvalid() {
}
@JsonProperty("numSegmentsPrunedInvalid")
- @Override
public void setNumSegmentsPrunedInvalid(long numSegmentsPrunedInvalid) {
_numSegmentsPrunedInvalid = numSegmentsPrunedInvalid;
}
@@ -275,7 +265,6 @@ public long getNumSegmentsPrunedByLimit() {
}
@JsonProperty("numSegmentsPrunedByLimit")
- @Override
public void setNumSegmentsPrunedByLimit(long numSegmentsPrunedByLimit) {
_numSegmentsPrunedByLimit = numSegmentsPrunedByLimit;
}
@@ -287,7 +276,6 @@ public long getNumSegmentsPrunedByValue() {
}
@JsonProperty("numSegmentsPrunedByValue")
- @Override
public void setNumSegmentsPrunedByValue(long numSegmentsPrunedByValue) {
_numSegmentsPrunedByValue = numSegmentsPrunedByValue;
}
@@ -299,7 +287,6 @@ public long getExplainPlanNumEmptyFilterSegments() {
}
@JsonProperty("explainPlanNumEmptyFilterSegments")
- @Override
public void setExplainPlanNumEmptyFilterSegments(long explainPlanNumEmptyFilterSegments) {
_explainPlanNumEmptyFilterSegments = explainPlanNumEmptyFilterSegments;
}
@@ -311,7 +298,6 @@ public long getExplainPlanNumMatchAllFilterSegments() {
}
@JsonProperty("explainPlanNumMatchAllFilterSegments")
- @Override
public void setExplainPlanNumMatchAllFilterSegments(long explainPlanNumMatchAllFilterSegments) {
_explainPlanNumMatchAllFilterSegments = explainPlanNumMatchAllFilterSegments;
}
@@ -350,7 +336,6 @@ public int getNumServersQueried() {
}
@JsonProperty("numServersQueried")
- @Override
public void setNumServersQueried(int numServersQueried) {
_numServersQueried = numServersQueried;
}
@@ -448,6 +433,7 @@ public void setNumConsumingSegmentsQueried(long numConsumingSegmentsQueried) {
public long getNumConsumingSegmentsProcessed() {
return _numConsumingSegmentsProcessed;
}
+
@JsonProperty("numConsumingSegmentsProcessed")
public void setNumConsumingSegmentsProcessed(long numConsumingSegmentsProcessed) {
_numConsumingSegmentsProcessed = numConsumingSegmentsProcessed;
@@ -497,27 +483,18 @@ public void setNumGroupsLimitReached(boolean numGroupsLimitReached) {
_numGroupsLimitReached = numGroupsLimitReached;
}
- @JsonProperty("maxRowsInJoinReached")
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
public boolean isMaxRowsInJoinReached() {
- return _maxRowsInJoinReached;
- }
-
- @JsonProperty("maxRowsInJoinReached")
- public void setMaxRowsInJoinReached(boolean maxRowsInJoinReached) {
- _maxRowsInJoinReached = maxRowsInJoinReached;
+ return false;
}
- @JsonProperty("partialResult")
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
public boolean isPartialResult() {
- return _partialResult;
- }
-
- @JsonProperty("partialResult")
- public void setPartialResult(boolean partialResult) {
- _partialResult = partialResult;
+ return isNumGroupsLimitReached() || getExceptionsSize() > 0 || isMaxRowsInJoinReached();
}
@JsonProperty("timeUsedMs")
+ @Override
public long getTimeUsedMs() {
return _timeUsedMs;
}
@@ -528,16 +505,10 @@ public void setTimeUsedMs(long timeUsedMs) {
_timeUsedMs = timeUsedMs;
}
- @JsonProperty("numRowsResultSet")
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
@Override
public int getNumRowsResultSet() {
- return _numRowsResultSet;
- }
-
- @JsonProperty("numRowsResultSet")
- @Override
- public void setNumRowsResultSet(int numRowsResultSet) {
- _numRowsResultSet = numRowsResultSet;
+ return BrokerResponse.super.getNumRowsResultSet();
}
@JsonProperty("segmentStatistics")
@@ -560,15 +531,11 @@ public void setTraceInfo(Map traceInfo) {
_traceInfo = traceInfo;
}
- @Override
- public String toJsonString()
- throws IOException {
- return JsonUtils.objectToString(this);
- }
-
@JsonIgnore
@Override
public void setExceptions(List exceptions) {
+ // TODO: This is incorrect. It is adding and not setting the exceptions
+ // But there is some code that seems to depend on this.
for (ProcessingException exception : exceptions) {
_processingExceptions.add(new QueryProcessingException(exception.getErrorCode(), exception.getMessage()));
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
index 60a22460f39a..40f521e25176 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
@@ -18,16 +18,17 @@
*/
package org.apache.pinot.common.response.broker;
+import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
-import java.io.IOException;
+import com.fasterxml.jackson.databind.node.ObjectNode;
import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Collections;
import java.util.List;
-import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.datatable.StatMap;
+import org.apache.pinot.common.response.BrokerResponse;
import org.apache.pinot.common.response.ProcessingException;
-import org.apache.pinot.common.utils.DataSchema;
-import org.apache.pinot.spi.utils.JsonUtils;
/**
@@ -46,28 +47,36 @@
"realtimeResponseSerializationCpuTimeNs", "offlineTotalCpuTimeNs", "realtimeTotalCpuTimeNs", "brokerReduceTimeMs",
"segmentStatistics", "traceInfo", "partialResult"
})
-public class BrokerResponseNativeV2 extends BrokerResponseNative {
+public class BrokerResponseNativeV2 implements BrokerResponse {
- private final Map _stageIdStats = new HashMap<>();
+ private ObjectNode _stageStats = null;
+ /**
+ * The max number of rows seen at runtime.
+ *
+ * In single-stage this doesn't make sense given it is the max number of rows read from the table. But in multi-stage
+ * virtual rows can be generated. For example, in a join query, the number of rows can be more than the number of rows
+ * in the table.
+ */
+ private long _maxRowsInOperator = 0;
+ private final StatMap _brokerStats = new StatMap<>(StatKey.class);
+ private final List _processingExceptions;
+ private ResultTable _resultTable;
+ private String _requestId;
+ private String _brokerId;
public BrokerResponseNativeV2() {
+ _processingExceptions = new ArrayList<>();
}
public BrokerResponseNativeV2(ProcessingException exception) {
- super(exception);
+ this(Collections.singletonList(exception));
}
public BrokerResponseNativeV2(List exceptions) {
- super(exceptions);
- }
-
- /** Generate EXPLAIN PLAN output when queries are evaluated by Broker without going to the Server. */
- private static BrokerResponseNativeV2 getBrokerResponseExplainPlanOutput() {
- BrokerResponseNativeV2 brokerResponse = BrokerResponseNativeV2.empty();
- List rows = new ArrayList<>();
- rows.add(new Object[]{"BROKER_EVALUATE", 0, -1});
- brokerResponse.setResultTable(new ResultTable(DataSchema.EXPLAIN_RESULT_SCHEMA, rows));
- return brokerResponse;
+ _processingExceptions = new ArrayList<>(exceptions.size());
+ for (ProcessingException exception : exceptions) {
+ _processingExceptions.add(new QueryProcessingException(exception.getErrorCode(), exception.getMessage()));
+ }
}
/**
@@ -77,20 +86,338 @@ public static BrokerResponseNativeV2 empty() {
return new BrokerResponseNativeV2();
}
- public static BrokerResponseNativeV2 fromJsonString(String jsonString)
- throws IOException {
- return JsonUtils.stringToObject(jsonString, BrokerResponseNativeV2.class);
+ public void setStageStats(ObjectNode stageStats) {
+ _stageStats = stageStats;
+ }
+
+ @JsonProperty
+ public ObjectNode getStageStats() {
+ return _stageStats;
+ }
+
+ /**
+ * Get the max number of rows seen by a single operator in the query processing chain.
+ */
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
+ public long getMaxRowsInOperator() {
+ return _maxRowsInOperator;
+ }
+
+ public void mergeMaxRowsInOperator(long maxRows) {
+ _maxRowsInOperator = Math.max(_maxRowsInOperator, maxRows);
+ }
+
+ @Override
+ public long getTimeUsedMs() {
+ return _brokerStats.getLong(StatKey.TIME_USED_MS);
+ }
+
+ @Override
+ public void setTimeUsedMs(long timeUsedMs) {
+ _brokerStats.merge(StatKey.TIME_USED_MS, timeUsedMs);
+ }
+
+ @Override
+ public long getNumDocsScanned() {
+ return _brokerStats.getLong(StatKey.NUM_DOCS_SCANNED);
+ }
+
+ @Override
+ public long getNumEntriesScannedInFilter() {
+ return _brokerStats.getLong(StatKey.NUM_ENTRIES_SCANNED_IN_FILTER);
+ }
+
+ @Override
+ public long getNumEntriesScannedPostFilter() {
+ return _brokerStats.getLong(StatKey.NUM_ENTRIES_SCANNED_POST_FILTER);
+ }
+
+ @Override
+ public long getNumSegmentsQueried() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_QUERIED);
+ }
+
+ @Override
+ public long getNumSegmentsProcessed() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PROCESSED);
+ }
+
+ @Override
+ public long getNumSegmentsMatched() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_MATCHED);
+ }
+
+ @Override
+ public long getNumConsumingSegmentsQueried() {
+ return _brokerStats.getLong(StatKey.NUM_CONSUMING_SEGMENTS_QUERIED);
+ }
+
+ @Override
+ public long getNumConsumingSegmentsProcessed() {
+ return _brokerStats.getLong(StatKey.NUM_CONSUMING_SEGMENTS_PROCESSED);
+ }
+
+ @Override
+ public long getNumConsumingSegmentsMatched() {
+ return _brokerStats.getLong(StatKey.NUM_CONSUMING_SEGMENTS_MATCHED);
+ }
+
+ @Override
+ public long getMinConsumingFreshnessTimeMs() {
+ return _brokerStats.getLong(StatKey.MIN_CONSUMING_FRESHNESS_TIME_MS);
+ }
+
+ @Override
+ public long getTotalDocs() {
+ return _brokerStats.getLong(StatKey.TOTAL_DOCS);
+ }
+
+ @Override
+ public boolean isNumGroupsLimitReached() {
+ return _brokerStats.getBoolean(StatKey.NUM_GROUPS_LIMIT_REACHED);
+ }
+
+ public void mergeNumGroupsLimitReached(boolean numGroupsLimitReached) {
+ _brokerStats.merge(StatKey.NUM_GROUPS_LIMIT_REACHED, numGroupsLimitReached);
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByServer() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_BY_SERVER);
}
- public void addStageStat(Integer stageId, BrokerResponseStats brokerResponseStats) {
- // StageExecutionWallTime will always be there, other stats are optional such as OperatorStats
- if (brokerResponseStats.getStageExecWallTimeMs() != -1) {
- _stageIdStats.put(stageId, brokerResponseStats);
+ @Override
+ public long getNumSegmentsPrunedInvalid() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_INVALID);
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByLimit() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_BY_LIMIT);
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByValue() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_BY_VALUE);
+ }
+
+ @Override
+ public long getExplainPlanNumEmptyFilterSegments() {
+ return _brokerStats.getLong(StatKey.EXPLAIN_PLAN_NUM_EMPTY_FILTER_SEGMENTS);
+ }
+
+ @Override
+ public long getExplainPlanNumMatchAllFilterSegments() {
+ return _brokerStats.getLong(StatKey.EXPLAIN_PLAN_NUM_MATCH_ALL_FILTER_SEGMENTS);
+ }
+
+ @Override
+ public long getOfflineTotalCpuTimeNs() {
+ return getOfflineThreadCpuTimeNs() + getOfflineSystemActivitiesCpuTimeNs()
+ + getOfflineResponseSerializationCpuTimeNs();
+ }
+
+ @Override
+ public long getRealtimeTotalCpuTimeNs() {
+ return getRealtimeThreadCpuTimeNs() + getRealtimeSystemActivitiesCpuTimeNs()
+ + getRealtimeResponseSerializationCpuTimeNs();
+ }
+
+ @Override
+ public void setExceptions(List exceptions) {
+ for (ProcessingException exception : exceptions) {
+ _processingExceptions.add(new QueryProcessingException(exception.getErrorCode(), exception.getMessage()));
}
}
- @JsonProperty("stageStats")
- public Map getStageIdStats() {
- return _stageIdStats;
+ public void addToExceptions(QueryProcessingException processingException) {
+ _processingExceptions.add(processingException);
+ }
+
+ @Override
+ public int getNumServersQueried() {
+ return _brokerStats.getInt(StatKey.NUM_SERVERS_QUERIED);
+ }
+
+ @Override
+ public void setNumServersQueried(int numServersQueried) {
+ _brokerStats.merge(StatKey.NUM_SERVERS_QUERIED, numServersQueried);
+ }
+
+ @Override
+ public int getNumServersResponded() {
+ return _brokerStats.getInt(StatKey.NUM_SERVERS_RESPONDED);
+ }
+
+ @Override
+ public void setNumServersResponded(int numServersResponded) {
+ _brokerStats.merge(StatKey.NUM_SERVERS_RESPONDED, numServersResponded);
+ }
+
+ @JsonProperty("maxRowsInJoinReached")
+ public boolean isMaxRowsInJoinReached() {
+ return _brokerStats.getBoolean(StatKey.MAX_ROWS_IN_JOIN_REACHED);
+ }
+
+ @JsonProperty("maxRowsInJoinReached")
+ public void mergeMaxRowsInJoinReached(boolean maxRowsInJoinReached) {
+ _brokerStats.merge(StatKey.MAX_ROWS_IN_JOIN_REACHED, maxRowsInJoinReached);
+ }
+
+ @Override
+ public int getExceptionsSize() {
+ return _processingExceptions.size();
+ }
+
+ @Override
+ public void setResultTable(@Nullable ResultTable resultTable) {
+ _resultTable = resultTable;
+ }
+
+ @Nullable
+ @Override
+ @JsonInclude(JsonInclude.Include.NON_NULL)
+ public ResultTable getResultTable() {
+ return _resultTable;
+ }
+
+ @JsonProperty("exceptions")
+ @Override
+ public List getProcessingExceptions() {
+ return List.of();
+ }
+
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
+ @Override
+ public int getNumRowsResultSet() {
+ return BrokerResponse.super.getNumRowsResultSet();
+ }
+
+ @Override
+ public long getOfflineThreadCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getRealtimeThreadCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getOfflineSystemActivitiesCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getRealtimeSystemActivitiesCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getOfflineResponseSerializationCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getRealtimeResponseSerializationCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByBroker() {
+ return _brokerStats.getInt(StatKey.NUM_SEGMENTS_PRUNED_BY_BROKER);
+ }
+
+ @Override
+ public void setNumSegmentsPrunedByBroker(long numSegmentsPrunedByBroker) {
+ _brokerStats.merge(StatKey.NUM_SEGMENTS_PRUNED_BY_BROKER, (int) numSegmentsPrunedByBroker);
+ }
+
+ @Override
+ public String getRequestId() {
+ return _requestId;
+ }
+
+ @Override
+ public void setRequestId(String requestId) {
+ _requestId = requestId;
+ }
+
+ @Override
+ public String getBrokerId() {
+ return _brokerId;
+ }
+
+ @Override
+ public void setBrokerId(String requestId) {
+ _brokerId = requestId;
+ }
+
+ @Override
+ public long getBrokerReduceTimeMs() {
+ return _brokerStats.getLong(StatKey.BROKER_REDUCE_TIME_MS);
+ }
+
+ @Override
+ public void setBrokerReduceTimeMs(long brokerReduceTimeMs) {
+ _brokerStats.merge(StatKey.BROKER_REDUCE_TIME_MS, brokerReduceTimeMs);
+ }
+
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
+ @Override
+ public boolean isPartialResult() {
+ return isNumGroupsLimitReached() || getExceptionsSize() > 0 || isMaxRowsInJoinReached();
+ }
+
+ public void addServerStats(StatMap serverStats) {
+ // Set execution statistics.
+ _brokerStats.merge(serverStats);
+ }
+
+ public enum StatKey implements StatMap.Key {
+ TIME_USED_MS(StatMap.Type.LONG),
+ NUM_DOCS_SCANNED(StatMap.Type.LONG),
+ NUM_ENTRIES_SCANNED_IN_FILTER(StatMap.Type.LONG),
+ NUM_ENTRIES_SCANNED_POST_FILTER(StatMap.Type.LONG),
+ NUM_SEGMENTS_QUERIED(StatMap.Type.INT),
+ NUM_SEGMENTS_PROCESSED(StatMap.Type.INT),
+ NUM_SEGMENTS_MATCHED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_QUERIED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_PROCESSED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_MATCHED(StatMap.Type.INT),
+ MIN_CONSUMING_FRESHNESS_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public long merge(long value1, long value2) {
+ return StatMap.Key.minPositive(value1, value2);
+ }
+ },
+ TOTAL_DOCS(StatMap.Type.LONG),
+ NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN),
+ NUM_SEGMENTS_PRUNED_BY_SERVER(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_INVALID(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_BY_LIMIT(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_BY_VALUE(StatMap.Type.INT),
+ NUM_SERGMENTS_PRUNED_BY_BROKER(StatMap.Type.INT),
+ EXPLAIN_PLAN_NUM_EMPTY_FILTER_SEGMENTS(StatMap.Type.INT),
+ EXPLAIN_PLAN_NUM_MATCH_ALL_FILTER_SEGMENTS(StatMap.Type.INT),
+ THREAD_CPU_TIME_NS(StatMap.Type.LONG),
+ SYSTEM_ACTIVITIES_CPU_TIME_NS(StatMap.Type.LONG),
+ RESPONSE_SER_CPU_TIME_NS(StatMap.Type.LONG),
+ NUM_SEGMENTS_PRUNED_BY_BROKER(StatMap.Type.INT),
+ MAX_ROWS_IN_JOIN_REACHED(StatMap.Type.BOOLEAN),
+ NUM_SERVERS_RESPONDED(StatMap.Type.INT),
+ NUM_SERVERS_QUERIED(StatMap.Type.INT),
+ BROKER_REDUCE_TIME_MS(StatMap.Type.LONG);
+
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java
index 3acc9f334932..97631876009a 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java
@@ -21,12 +21,10 @@
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
-import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import org.apache.pinot.spi.utils.JsonUtils;
// TODO: Decouple the execution stats aggregator logic and make it into a util that can aggregate 2 values with the
@@ -109,11 +107,6 @@ public void setStageExecutionUnit(int stageExecutionUnit) {
_stageExecutionUnit = stageExecutionUnit;
}
- public String toJsonString()
- throws IOException {
- return JsonUtils.objectToString(this);
- }
-
@JsonProperty("operatorStats")
public Map> getOperatorStats() {
return _operatorStats;
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datablock/BaseDataBlockContract.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/BaseDataBlockContract.java
new file mode 100644
index 000000000000..763fd7344cae
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/BaseDataBlockContract.java
@@ -0,0 +1,42 @@
+/**
+ * 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.common.datablock;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import static org.testng.Assert.*;
+
+
+public abstract class BaseDataBlockContract {
+
+ protected abstract BaseDataBlock deserialize(ByteBuffer byteBuffer, int versionType)
+ throws IOException;
+
+ public void testSerdeCorrectness(BaseDataBlock dataBlock)
+ throws IOException {
+ byte[] bytes = dataBlock.toBytes();
+ ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
+ int versionType = DataBlockUtils.readVersionType(byteBuffer);
+ BaseDataBlock deserialize = deserialize(byteBuffer, versionType);
+
+ assertEquals(byteBuffer.position(), bytes.length, "Buffer position should be at the end of the buffer");
+ assertEquals(deserialize, dataBlock, "Deserialized data block should be the same as the original data block");
+ }
+}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java
index 0f85f4a69ee2..cfceddc65d7c 100644
--- a/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java
@@ -18,43 +18,185 @@
*/
package org.apache.pinot.common.datablock;
-import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Lists;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Collections;
import java.util.HashMap;
+import java.util.List;
import org.testng.annotations.Test;
import static org.testng.Assert.*;
-public class MetadataBlockTest {
+public class MetadataBlockTest extends BaseDataBlockContract {
+ @Override
+ protected BaseDataBlock deserialize(ByteBuffer byteBuffer, int versionType)
+ throws IOException {
+ return new MetadataBlock(byteBuffer);
+ }
@Test
- public void shouldEncodeContentsAsJSON()
+ public void emptyMetadataBlock()
throws Exception {
// Given:
- MetadataBlock.MetadataBlockType type = MetadataBlock.MetadataBlockType.EOS;
-
// When:
- MetadataBlock metadataBlock = new MetadataBlock(type);
+ MetadataBlock metadataBlock = MetadataBlock.newEos();
// Then:
- byte[] expected = new ObjectMapper().writeValueAsBytes(new MetadataBlock.Contents("EOS", new HashMap<>()));
- assertEquals(metadataBlock._variableSizeDataBytes, expected);
+ byte[] expectedFixed = new byte[0];
+ assertEquals(metadataBlock._fixedSizeDataBytes, expectedFixed);
+ byte[] expectedVariable = new byte[0];
+ assertEquals(metadataBlock._variableSizeDataBytes, expectedVariable);
+
+ List statsByStage = metadataBlock.getStatsByStage();
+ assertNotNull(statsByStage, "Expected stats by stage to be non-null");
+ assertEquals(statsByStage.size(), 0, "Expected no stats by stage");
+ }
+
+ @Test
+ public void emptyDataBlockCorrectness()
+ throws IOException {
+ testSerdeCorrectness(MetadataBlock.newEos());
+ }
+
+ @Test
+ public void v1ErrorWithExceptionsIsDecodedAsV2ErrorWithSameExceptions()
+ throws IOException {
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(MetadataBlock.MetadataBlockType.ERROR);
+ v1MetadataBlock.addException(250, "timeout");
+ v1MetadataBlock.addException(500, "server error");
+
+ byte[] bytes = v1MetadataBlock.toBytes();
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(buff);
+
+ assertTrue(dataBlock instanceof MetadataBlock, "V1MetadataBlock should be always decoded as MetadataBlock");
+ MetadataBlock metadataBlock = (MetadataBlock) dataBlock;
+ assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.ERROR, "Expected error type");
+ assertEquals(metadataBlock.getStatsByStage(), Collections.emptyList(), "Expected no stats by stage");
+ assertEquals(metadataBlock.getExceptions(), v1MetadataBlock.getExceptions(), "Expected exceptions");
+ }
+
+ /**
+ * Verifies that a V2 EOS with empty stats is read in V1 as EOS without stats
+ */
+ @Test
+ public void v2EosWithoutStatsIsReadInV1AsEosWithoutStats()
+ throws IOException {
+ ByteBuffer stats = ByteBuffer.wrap(new byte[]{0, 0, 0, 0});
+ MetadataBlock metadataBlock = new MetadataBlock(Lists.newArrayList(stats));
+
+ byte[] bytes = metadataBlock.toBytes();
+
+ // This is how V1 blocks were deserialized
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(buff);
+
+ assertEquals(v1MetadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(v1MetadataBlock.getStats(), Collections.emptyMap(), "Expected no stats by stage");
+ assertEquals(v1MetadataBlock.getExceptions(), metadataBlock.getExceptions(), "Expected no exceptions");
}
+ /**
+ * Verifies that a V2 EOS with stats is read in V1 as EOS without stats
+ */
@Test
- public void shouldDefaultToEosWithNoErrorsOnLegacyMetadataBlock()
+ public void v2EosWithStatsIsReadInV1AsEosWithoutStats()
+ throws IOException {
+ MetadataBlock metadataBlock = MetadataBlock.newEos();
+
+ byte[] bytes = metadataBlock.toBytes();
+
+ // This is how V1 blocks were deserialized
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(buff);
+
+ assertEquals(v1MetadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(v1MetadataBlock.getStats(), Collections.emptyMap(), "Expected no stats by stage");
+ assertEquals(v1MetadataBlock.getExceptions(), metadataBlock.getExceptions(), "Expected no exceptions");
+ }
+
+ /**
+ * Verifies that a V2 error code is read in V1 as error with the same exceptions
+ */
+ @Test
+ public void v2ErrorIsReadInV1AsErrorWithSameExceptions()
+ throws IOException {
+ HashMap errorMap = new HashMap<>();
+ errorMap.put(250, "timeout");
+ MetadataBlock metadataBlock = MetadataBlock.newError(errorMap);
+
+ byte[] bytes = metadataBlock.toBytes();
+
+ // This is how V1 blocks were deserialized
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(buff);
+
+ assertEquals(v1MetadataBlock.getType(), MetadataBlock.MetadataBlockType.ERROR, "Expected error type");
+ assertEquals(v1MetadataBlock.getStats(), Collections.emptyMap(), "Expected no stats by stage");
+ assertEquals(v1MetadataBlock.getExceptions(), metadataBlock.getExceptions(), "Expected exceptions");
+ }
+
+ /**
+ * Verifies that a V1 error code is decoded as a V2 error code with the same exceptions
+ */
+ @Test
+ public void v1EosWithoutStatsIsDecodedAsV2EosWithoutStats()
+ throws IOException {
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(MetadataBlock.MetadataBlockType.EOS, new HashMap<>());
+
+ byte[] bytes = v1MetadataBlock.toBytes();
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(buff);
+
+ assertTrue(dataBlock instanceof MetadataBlock, "V1MetadataBlock should be always decoded as MetadataBlock");
+ MetadataBlock metadataBlock = (MetadataBlock) dataBlock;
+ assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(metadataBlock.getStatsByStage(), Collections.emptyList(), "Expected no stats by stage");
+ assertEquals(metadataBlock.getExceptions(), v1MetadataBlock.getExceptions(), "Expected exceptions");
+ }
+
+ /**
+ * Verifies that a V1 EOS with stats is decoded as a V2 EOS without stats
+ */
+ @Test
+ public void v1EosWithStatsIsDecodedAsV2EosWithoutStats()
+ throws IOException {
+ HashMap stats = new HashMap<>();
+ stats.put("foo", "bar");
+ stats.put("baz", "qux");
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(MetadataBlock.MetadataBlockType.EOS, stats);
+
+ byte[] bytes = v1MetadataBlock.toBytes();
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(buff);
+
+ assertTrue(dataBlock instanceof MetadataBlock, "V1MetadataBlock should be always decoded as MetadataBlock");
+ MetadataBlock metadataBlock = (MetadataBlock) dataBlock;
+ assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(metadataBlock.getStatsByStage(), Collections.emptyList(), "Expected no stats by stage");
+ assertEquals(metadataBlock.getExceptions(), Collections.emptyMap(), "Expected no exceptions");
+ }
+
+ /**
+ * Verifies that a V0 EOS without exceptions is decoded as a V2 EOS without stats
+ */
+ @Test
+ public void v0EosWithoutExceptionsIsDecodedAsV2EosWithoutStats()
throws IOException {
// Given:
// MetadataBlock used to be encoded without any data, we should make sure that
// during rollout or if server versions are mismatched that we can still handle
// the old format
- OldMetadataBlock legacyBlock = new OldMetadataBlock();
+ V0MetadataBlock legacyBlock = new V0MetadataBlock();
byte[] bytes = legacyBlock.toBytes();
// When:
ByteBuffer buff = ByteBuffer.wrap(bytes);
- buff.getInt(); // consume the version information before decoding
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
MetadataBlock metadataBlock = new MetadataBlock(buff);
// Then:
@@ -62,29 +204,30 @@ public void shouldDefaultToEosWithNoErrorsOnLegacyMetadataBlock()
}
@Test
- public void shouldDefaultToErrorOnLegacyMetadataBlockWithErrors()
+ public void v0EosWithExceptionsIsDecodedAsV2ErrorWithSameExceptions()
throws IOException {
// Given:
// MetadataBlock used to be encoded without any data, we should make sure that
// during rollout or if server versions are mismatched that we can still handle
// the old format
- OldMetadataBlock legacyBlock = new OldMetadataBlock();
+ V0MetadataBlock legacyBlock = new V0MetadataBlock();
legacyBlock.addException(250, "timeout");
byte[] bytes = legacyBlock.toBytes();
// When:
ByteBuffer buff = ByteBuffer.wrap(bytes);
- buff.getInt(); // consume the version information before decoding
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
MetadataBlock metadataBlock = new MetadataBlock(buff);
// Then:
assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.ERROR);
+ assertEquals(metadataBlock.getExceptions(), legacyBlock.getExceptions(), "Expected exceptions");
}
@Test(expectedExceptions = UnsupportedOperationException.class)
public void shouldThrowExceptionWhenUsingReadMethods() {
// Given:
- MetadataBlock block = new MetadataBlock(MetadataBlock.MetadataBlockType.EOS);
+ MetadataBlock block = MetadataBlock.newEos();
// When:
// (should through exception)
@@ -94,9 +237,9 @@ public void shouldThrowExceptionWhenUsingReadMethods() {
/**
* This is mostly just used as an internal serialization tool
*/
- private static class OldMetadataBlock extends BaseDataBlock {
+ private static class V0MetadataBlock extends BaseDataBlock {
- public OldMetadataBlock() {
+ public V0MetadataBlock() {
super(0, null, new String[0], new byte[0], new byte[0]);
}
@@ -114,15 +257,5 @@ protected int getOffsetInFixedBuffer(int rowId, int colId) {
protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
return 0;
}
-
- @Override
- public DataBlock toMetadataOnlyDataTable() {
- return null;
- }
-
- @Override
- public DataBlock toDataOnlyDataTable() {
- return null;
- }
}
}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datablock/RowDataBlockTest.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/RowDataBlockTest.java
new file mode 100644
index 000000000000..b3bbe26ff613
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/RowDataBlockTest.java
@@ -0,0 +1,40 @@
+/**
+ * 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.common.datablock;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.apache.pinot.common.utils.DataSchema;
+import org.testng.annotations.Test;
+
+
+public class RowDataBlockTest extends BaseDataBlockContract {
+ @Override
+ protected BaseDataBlock deserialize(ByteBuffer byteBuffer, int versionType)
+ throws IOException {
+ return new RowDataBlock(byteBuffer);
+ }
+
+ @Test
+ public void emptyDataBlockCorrectness()
+ throws IOException {
+ DataSchema dataSchema = new DataSchema(new String[0], new DataSchema.ColumnDataType[0]);
+ testSerdeCorrectness(new RowDataBlock(0, dataSchema, new String[0], new byte[0], new byte[0]));
+ }
+}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datablock/V1MetadataBlock.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/V1MetadataBlock.java
new file mode 100644
index 000000000000..990d9489ccc8
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/V1MetadataBlock.java
@@ -0,0 +1,150 @@
+/**
+ * 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.common.datablock;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+
+/**
+ * The datablock used before the introduction of {@link org.apache.pinot.common.datatable.StatMap}.
+ *
+ * This version stored the metadata in a {@code Map} which was encoded as JSON and stored in the
+ * variable size data buffer.
+ *
+ * Instances of this class are not actually seen by the operators. Instead, they are converted to {@link MetadataBlock}
+ * in {@link MetadataBlock#deserialize(ByteBuffer, int)}.
+ *
+ * The reason to keep it here is mostly for backwards compatibility and testing. In order to simplify the code, the
+ * stats engine just ignores the metadata of these objects, but we need to be able to deserialize them anyway.
+ */
+public class V1MetadataBlock extends BaseDataBlock {
+
+ private static final ObjectMapper JSON = new ObjectMapper();
+
+ @VisibleForTesting
+ static final int VERSION = 1;
+
+ /**
+ * Used to serialize the contents of the metadata block conveniently and in
+ * a backwards compatible way. Use JSON because the performance of metadata block
+ * SerDe should not be a bottleneck.
+ */
+ @JsonIgnoreProperties(ignoreUnknown = true)
+ @VisibleForTesting
+ static class Contents {
+
+ private String _type;
+ private Map _stats;
+
+ @JsonCreator
+ public Contents(@JsonProperty("type") String type, @JsonProperty("stats") Map stats) {
+ _type = type;
+ _stats = stats;
+ }
+
+ @JsonCreator
+ public Contents() {
+ this(null, new HashMap<>());
+ }
+
+ public String getType() {
+ return _type;
+ }
+
+ public void setType(String type) {
+ _type = type;
+ }
+
+ public Map getStats() {
+ return _stats;
+ }
+
+ public void setStats(Map stats) {
+ _stats = stats;
+ }
+ }
+
+ private final Contents _contents;
+
+ public V1MetadataBlock(MetadataBlock.MetadataBlockType type) {
+ this(type, new HashMap<>());
+ }
+
+ public V1MetadataBlock(MetadataBlock.MetadataBlockType type, Map stats) {
+ super(0, null, new String[0], new byte[]{0}, toContents(new Contents(type.name(), stats)));
+ _contents = new Contents(type.name(), stats);
+ }
+
+ private static byte[] toContents(Contents type) {
+ try {
+ return JSON.writeValueAsBytes(type);
+ } catch (JsonProcessingException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ public V1MetadataBlock(ByteBuffer byteBuffer)
+ throws IOException {
+ super(byteBuffer);
+ if (_variableSizeDataBytes != null && _variableSizeDataBytes.length > 0) {
+ _contents = JSON.readValue(_variableSizeDataBytes, Contents.class);
+ } else {
+ _contents = new Contents();
+ }
+ }
+
+ public MetadataBlock.MetadataBlockType getType() {
+ String type = _contents.getType();
+
+ // if type is null, then we're reading a legacy block where we didn't encode any
+ // data. assume that it is an EOS block if there's no exceptions and an ERROR block
+ // otherwise
+ return type == null
+ ? (getExceptions().isEmpty() ? MetadataBlock.MetadataBlockType.EOS : MetadataBlock.MetadataBlockType.ERROR)
+ : MetadataBlock.MetadataBlockType.valueOf(type);
+ }
+
+ public Map getStats() {
+ return _contents.getStats() != null ? _contents.getStats() : new HashMap<>();
+ }
+
+ @Override
+ public int getDataBlockVersionType() {
+ return VERSION + (Type.METADATA.ordinal() << DataBlockUtils.VERSION_TYPE_SHIFT);
+ }
+
+ @Override
+ protected int getOffsetInFixedBuffer(int rowId, int colId) {
+ throw new UnsupportedOperationException("Metadata block uses JSON encoding for field access");
+ }
+
+ @Override
+ protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
+ throw new UnsupportedOperationException("Metadata block uses JSON encoding for field access");
+ }
+}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datatable/StatMapTest.java b/pinot-common/src/test/java/org/apache/pinot/common/datatable/StatMapTest.java
new file mode 100644
index 000000000000..25f5ef8a07fc
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datatable/StatMapTest.java
@@ -0,0 +1,209 @@
+/**
+ * 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.common.datatable;
+
+import com.google.common.io.ByteArrayDataInput;
+import com.google.common.io.ByteArrayDataOutput;
+import com.google.common.io.ByteStreams;
+import java.io.IOException;
+import org.apache.pinot.common.response.broker.BrokerResponseNativeV2;
+import org.testng.Assert;
+import org.testng.SkipException;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+
+public class StatMapTest {
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckWhenAddInt(MyStats stat) {
+ if (stat.getType() == StatMap.Type.INT) {
+ throw new SkipException("Skipping INT test");
+ }
+ if (stat.getType() == StatMap.Type.LONG) {
+ throw new SkipException("Skipping LONG test");
+ }
+ StatMap statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, 1);
+ }
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckWhenAddLong(MyStats stat) {
+ if (stat.getType() == StatMap.Type.LONG) {
+ throw new SkipException("Skipping LONG test");
+ }
+ StatMap statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, 1L);
+ }
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckPutBoolean(MyStats stat) {
+ if (stat.getType() == StatMap.Type.BOOLEAN) {
+ throw new SkipException("Skipping BOOLEAN test");
+ }
+ StatMap statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, true);
+ }
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckPutString(MyStats stat) {
+ if (stat.getType() == StatMap.Type.STRING) {
+ throw new SkipException("Skipping STRING test");
+ }
+ StatMap statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, "foo");
+ }
+
+ @Test(dataProvider = "allTypeStats")
+ public void singleEncodeDecode(MyStats stat)
+ throws IOException {
+ StatMap statMap = new StatMap<>(MyStats.class);
+ switch (stat.getType()) {
+ case BOOLEAN:
+ statMap.merge(stat, true);
+ break;
+ case INT:
+ statMap.merge(stat, 1);
+ break;
+ case LONG:
+ statMap.merge(stat, 1L);
+ break;
+ case STRING:
+ statMap.merge(stat, "foo");
+ break;
+ default:
+ throw new IllegalStateException();
+ }
+ testSerializeDeserialize(statMap);
+ }
+
+ @Test
+ public void encodeDecodeAll()
+ throws IOException {
+ StatMap statMap = new StatMap<>(MyStats.class);
+ for (MyStats stat : MyStats.values()) {
+ switch (stat.getType()) {
+ case BOOLEAN:
+ statMap.merge(stat, true);
+ break;
+ case INT:
+ statMap.merge(stat, 1);
+ break;
+ case LONG:
+ statMap.merge(stat, 1L);
+ break;
+ case STRING:
+ statMap.merge(stat, "foo");
+ break;
+ default:
+ throw new IllegalStateException();
+ }
+ }
+ testSerializeDeserialize(statMap);
+ }
+
+ private & StatMap.Key> void testSerializeDeserialize(StatMap statMap)
+ throws IOException {
+ ByteArrayDataOutput output = ByteStreams.newDataOutput();
+ statMap.serialize(output);
+
+ ByteArrayDataInput input = ByteStreams.newDataInput(output.toByteArray());
+ StatMap deserializedStatMap = StatMap.deserialize(input, MyStats.class);
+
+ Assert.assertEquals(statMap, deserializedStatMap);
+ }
+
+ @Test(dataProvider = "complexStats")
+ public void mergeEquivalence(StatMap statMap)
+ throws IOException {
+ StatMap mergedOnHeap = new StatMap<>(statMap.getKeyClass());
+ mergedOnHeap.merge(statMap);
+
+ ByteArrayDataOutput output = ByteStreams.newDataOutput();
+ statMap.serialize(output);
+
+ ByteArrayDataInput input = ByteStreams.newDataInput(output.toByteArray());
+ StatMap mergedSerialized = new StatMap<>(statMap.getKeyClass());
+ mergedSerialized.merge(input);
+
+ Assert.assertEquals(mergedOnHeap, mergedSerialized,
+ "Merging objects should be equal to merging serialized buffers");
+ }
+
+ @DataProvider(name = "complexStats")
+ static StatMap[] complexStats() {
+ return new StatMap[] {
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, true)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 1)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, false)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 1)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, true)
+ .merge(MyStats.LONG_KEY, 0L)
+ .merge(MyStats.INT_KEY, 1)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, false)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 0)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, false)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 1),
+ new StatMap<>(BrokerResponseNativeV2.StatKey.class)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_SEGMENTS_QUERIED, 1)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_SEGMENTS_PROCESSED, 1)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_SEGMENTS_MATCHED, 1)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_DOCS_SCANNED, 10)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_ENTRIES_SCANNED_POST_FILTER, 5)
+ .merge(BrokerResponseNativeV2.StatKey.TOTAL_DOCS, 5)
+ .merge(BrokerResponseNativeV2.StatKey.TIME_USED_MS, 95)
+ };
+ }
+
+ @DataProvider(name = "allTypeStats")
+ static MyStats[] allTypeStats() {
+ return MyStats.values();
+ }
+
+ public enum MyStats implements StatMap.Key {
+ BOOL_KEY(StatMap.Type.BOOLEAN),
+ LONG_KEY(StatMap.Type.LONG),
+ INT_KEY(StatMap.Type.INT),
+ STR_KEY(StatMap.Type.STRING);
+
+ private final StatMap.Type _type;
+
+ MyStats(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
+}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java
index 0ef9bb22071d..5ee75e5bf177 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java
@@ -45,6 +45,9 @@ public class InstanceResponseOperator extends BaseOperator combineOperator, List segmentContexts,
List fetchContexts, QueryContext queryContext) {
_combineOperator = combineOperator;
@@ -81,12 +84,24 @@ public static long calSystemActivitiesCpuTimeNs(long totalWallClockTimeNs, long
@Override
protected InstanceResponseBlock getNextBlock() {
+ BaseResultsBlock baseResultsBlock = getBaseBlock();
+ return buildInstanceResponseBlock(baseResultsBlock);
+ }
+
+ protected InstanceResponseBlock buildInstanceResponseBlock(BaseResultsBlock baseResultsBlock) {
+ InstanceResponseBlock instanceResponseBlock = new InstanceResponseBlock(baseResultsBlock);
+ instanceResponseBlock.addMetadata(MetadataKey.THREAD_CPU_TIME_NS.getName(), String.valueOf(_threadCpuTimeNs));
+ instanceResponseBlock.addMetadata(MetadataKey.SYSTEM_ACTIVITIES_CPU_TIME_NS.getName(),
+ String.valueOf(_systemActivitiesCpuTimeNs));
+ return instanceResponseBlock;
+ }
+
+ protected BaseResultsBlock getBaseBlock() {
if (ThreadResourceUsageProvider.isThreadCpuTimeMeasurementEnabled()) {
long startWallClockTimeNs = System.nanoTime();
ThreadResourceUsageProvider mainThreadResourceUsageProvider = new ThreadResourceUsageProvider();
BaseResultsBlock resultsBlock = getCombinedResults();
- InstanceResponseBlock instanceResponseBlock = new InstanceResponseBlock(resultsBlock);
long mainThreadCpuTimeNs = mainThreadResourceUsageProvider.getThreadTimeNs();
long totalWallClockTimeNs = System.nanoTime() - startWallClockTimeNs;
@@ -97,22 +112,17 @@ protected InstanceResponseBlock getNextBlock() {
*/
long multipleThreadCpuTimeNs = resultsBlock.getExecutionThreadCpuTimeNs();
int numServerThreads = resultsBlock.getNumServerThreads();
- long systemActivitiesCpuTimeNs =
- calSystemActivitiesCpuTimeNs(totalWallClockTimeNs, multipleThreadCpuTimeNs, mainThreadCpuTimeNs,
- numServerThreads);
-
- long threadCpuTimeNs = mainThreadCpuTimeNs + multipleThreadCpuTimeNs;
- instanceResponseBlock.addMetadata(MetadataKey.THREAD_CPU_TIME_NS.getName(), String.valueOf(threadCpuTimeNs));
- instanceResponseBlock.addMetadata(MetadataKey.SYSTEM_ACTIVITIES_CPU_TIME_NS.getName(),
- String.valueOf(systemActivitiesCpuTimeNs));
+ _systemActivitiesCpuTimeNs = calSystemActivitiesCpuTimeNs(totalWallClockTimeNs, multipleThreadCpuTimeNs,
+ mainThreadCpuTimeNs, numServerThreads);
+ _threadCpuTimeNs = mainThreadCpuTimeNs + multipleThreadCpuTimeNs;
- return instanceResponseBlock;
+ return resultsBlock;
} else {
- return new InstanceResponseBlock(getCombinedResults());
+ return getCombinedResults();
}
}
- private BaseResultsBlock getCombinedResults() {
+ protected BaseResultsBlock getCombinedResults() {
try {
prefetchAll();
return _combineOperator.nextBlock();
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java
index 61a5b8cc83d9..1130efb4d716 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java
@@ -36,6 +36,12 @@
import org.apache.pinot.spi.trace.Tracing;
+/**
+ * Like {@link InstanceResponseOperator}, but instead of sending all the data to the broker at once, it streams the data
+ * to a given {@link ResultsBlockStreamer}.
+ *
+ * This is used in multi-stage to stream data to the receiving mailboxes.
+ */
public class StreamingInstanceResponseOperator extends InstanceResponseOperator {
private static final String EXPLAIN_NAME = "STREAMING_INSTANCE_RESPONSE";
@@ -58,7 +64,7 @@ protected InstanceResponseBlock getNextBlock() {
prefetchAll();
if (_streamingCombineOperator != null) {
_streamingCombineOperator.start();
- BaseResultsBlock resultsBlock = _streamingCombineOperator.nextBlock();
+ BaseResultsBlock resultsBlock = getBaseBlock();
while (!(resultsBlock instanceof MetadataResultsBlock)) {
if (resultsBlock instanceof ExceptionResultsBlock) {
return new InstanceResponseBlock(resultsBlock);
@@ -66,20 +72,20 @@ protected InstanceResponseBlock getNextBlock() {
if (resultsBlock.getNumRows() > 0) {
_streamer.send(resultsBlock);
}
- resultsBlock = _streamingCombineOperator.nextBlock();
+ resultsBlock = getBaseBlock();
}
// Return a metadata-only block in the end
- return new InstanceResponseBlock(resultsBlock);
+ return buildInstanceResponseBlock(resultsBlock);
} else {
// Handle single block combine operator in streaming fashion
- BaseResultsBlock resultsBlock = _combineOperator.nextBlock();
+ BaseResultsBlock resultsBlock = getBaseBlock();
if (resultsBlock instanceof ExceptionResultsBlock) {
return new InstanceResponseBlock(resultsBlock);
}
if (resultsBlock.getNumRows() > 0) {
_streamer.send(resultsBlock);
}
- return new InstanceResponseBlock(resultsBlock).toMetadataOnlyResponseBlock();
+ return buildInstanceResponseBlock(resultsBlock).toMetadataOnlyResponseBlock();
}
} catch (EarlyTerminationException e) {
Exception killedErrorMsg = Tracing.getThreadAccountant().getErrorStatus();
@@ -96,6 +102,10 @@ protected InstanceResponseBlock getNextBlock() {
}
}
+ protected BaseResultsBlock getCombinedResults() {
+ return _combineOperator.nextBlock();
+ }
+
@Override
public String toExplainString() {
return EXPLAIN_NAME;
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java
index c05644af5fbc..c40133baab8f 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java
@@ -34,7 +34,6 @@
import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.metrics.BrokerTimer;
import org.apache.pinot.common.response.broker.BrokerResponseNative;
-import org.apache.pinot.common.response.broker.BrokerResponseStats;
import org.apache.pinot.common.response.broker.QueryProcessingException;
import org.apache.pinot.core.transport.ServerRoutingInstance;
import org.apache.pinot.spi.config.table.TableType;
@@ -74,7 +73,6 @@ public class ExecutionStatsAggregator {
private long _explainPlanNumEmptyFilterSegments = 0L;
private long _explainPlanNumMatchAllFilterSegments = 0L;
private boolean _numGroupsLimitReached = false;
- private boolean _maxRowsInJoinReached = false;
private int _numBlocks = 0;
private int _numRows = 0;
private long _stageExecutionTimeMs = 0;
@@ -248,12 +246,10 @@ public synchronized void aggregate(@Nullable ServerRoutingInstance routingInstan
if (numTotalDocsString != null) {
_numTotalDocs += Long.parseLong(numTotalDocsString);
}
+
_numGroupsLimitReached |=
Boolean.parseBoolean(metadata.get(DataTable.MetadataKey.NUM_GROUPS_LIMIT_REACHED.getName()));
- _maxRowsInJoinReached |=
- Boolean.parseBoolean(metadata.get(DataTable.MetadataKey.MAX_ROWS_IN_JOIN_REACHED.getName()));
-
String numBlocksString = metadata.get(DataTable.MetadataKey.NUM_BLOCKS.getName());
if (numBlocksString != null) {
_numBlocks += Long.parseLong(numBlocksString);
@@ -309,15 +305,12 @@ public void setStats(@Nullable String rawTableName, BrokerResponseNative brokerR
brokerResponseNative.setNumSegmentsMatched(_numSegmentsMatched);
brokerResponseNative.setTotalDocs(_numTotalDocs);
brokerResponseNative.setNumGroupsLimitReached(_numGroupsLimitReached);
- brokerResponseNative.setMaxRowsInJoinReached(_maxRowsInJoinReached);
brokerResponseNative.setOfflineThreadCpuTimeNs(_offlineThreadCpuTimeNs);
brokerResponseNative.setRealtimeThreadCpuTimeNs(_realtimeThreadCpuTimeNs);
brokerResponseNative.setOfflineSystemActivitiesCpuTimeNs(_offlineSystemActivitiesCpuTimeNs);
brokerResponseNative.setRealtimeSystemActivitiesCpuTimeNs(_realtimeSystemActivitiesCpuTimeNs);
brokerResponseNative.setOfflineResponseSerializationCpuTimeNs(_offlineResponseSerializationCpuTimeNs);
brokerResponseNative.setRealtimeResponseSerializationCpuTimeNs(_realtimeResponseSerializationCpuTimeNs);
- brokerResponseNative.setOfflineTotalCpuTimeNs(_offlineTotalCpuTimeNs);
- brokerResponseNative.setRealtimeTotalCpuTimeNs(_realtimeTotalCpuTimeNs);
brokerResponseNative.setNumSegmentsPrunedByServer(_numSegmentsPrunedByServer);
brokerResponseNative.setNumSegmentsPrunedInvalid(_numSegmentsPrunedInvalid);
brokerResponseNative.setNumSegmentsPrunedByLimit(_numSegmentsPrunedByLimit);
@@ -364,25 +357,6 @@ public void setStats(@Nullable String rawTableName, BrokerResponseNative brokerR
}
}
- public void setStageLevelStats(@Nullable String rawTableName, BrokerResponseStats brokerResponseStats,
- @Nullable BrokerMetrics brokerMetrics) {
- if (_enableTrace) {
- setStats(rawTableName, brokerResponseStats, brokerMetrics);
- brokerResponseStats.setOperatorStats(_operatorStats);
- }
-
- brokerResponseStats.setNumBlocks(_numBlocks);
- brokerResponseStats.setNumRows(_numRows);
- brokerResponseStats.setMaxRowsInJoinReached(_maxRowsInJoinReached);
- brokerResponseStats.setNumGroupsLimitReached(_numGroupsLimitReached);
- brokerResponseStats.setStageExecutionTimeMs(_stageExecutionTimeMs);
- brokerResponseStats.setStageExecutionUnit(_stageExecutionUnit);
- brokerResponseStats.setTableNames(new ArrayList<>(_tableNames));
- if (_stageExecStartTimeMs >= 0 && _stageExecEndTimeMs >= 0) {
- brokerResponseStats.setStageExecWallTimeMs(_stageExecEndTimeMs - _stageExecStartTimeMs);
- }
- }
-
private void withNotNullLongMetadata(Map metadata, DataTable.MetadataKey key, LongConsumer consumer) {
String strValue = metadata.get(key.getName());
if (strValue != null) {
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
index e7c5af0dab9b..6bc1f0cbdae4 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
@@ -406,7 +406,7 @@ public void testBrokerResponseMetadata()
for (String query : queries) {
JsonNode response = postQuery(query);
for (String statName : statNames) {
- assertTrue(response.has(statName));
+ assertTrue(response.has(statName), "Response does not contain stat: " + statName);
}
}
}
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java
index f66c194d8538..a54e5314f7ff 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java
@@ -317,7 +317,8 @@ public void testDropResults(boolean useMultiStageQueryEngine)
final String resultTag = "resultTable";
// dropResults=true - resultTable must not be in the response
- Assert.assertFalse(postQueryWithOptions(query, "dropResults=true").has(resultTag));
+ JsonNode jsonNode = postQueryWithOptions(query, "dropResults=true");
+ Assert.assertFalse(jsonNode.has(resultTag));
// dropResults=TrUE (case insensitive match) - resultTable must not be in the response
Assert.assertFalse(postQueryWithOptions(query, "dropResults=TrUE").has(resultTag));
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java
index 5398e58b77f7..0dc2dc588434 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java
@@ -24,6 +24,7 @@
import org.apache.calcite.plan.RelOptRule;
import org.apache.calcite.plan.RelOptRuleCall;
import org.apache.calcite.plan.hep.HepRelVertex;
+import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.RelDistributions;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Exchange;
@@ -164,11 +165,15 @@ public void onMatch(RelOptRuleCall call) {
// 2. when hash key are the same but hash functions are different, it can be done via normal hash shuffle.
boolean isColocatedJoin = PinotHintStrategyTable.isHintOptionTrue(join.getHints(),
PinotHintOptions.JOIN_HINT_OPTIONS, PinotHintOptions.JoinHintOptions.IS_COLOCATED_BY_JOIN_KEYS);
- PinotLogicalExchange dynamicBroadcastExchange = isColocatedJoin
- ? PinotLogicalExchange.create(right.getInput(), RelDistributions.hash(join.analyzeCondition().rightKeys),
- PinotRelExchangeType.PIPELINE_BREAKER)
- : PinotLogicalExchange.create(right.getInput(), RelDistributions.BROADCAST_DISTRIBUTED,
- PinotRelExchangeType.PIPELINE_BREAKER);
+ PinotLogicalExchange dynamicBroadcastExchange;
+ RelNode rightInput = right.getInput();
+ if (isColocatedJoin) {
+ RelDistribution dist = RelDistributions.hash(join.analyzeCondition().rightKeys);
+ dynamicBroadcastExchange = PinotLogicalExchange.create(rightInput, dist, PinotRelExchangeType.PIPELINE_BREAKER);
+ } else {
+ RelDistribution dist = RelDistributions.BROADCAST_DISTRIBUTED;
+ dynamicBroadcastExchange = PinotLogicalExchange.create(rightInput, dist, PinotRelExchangeType.PIPELINE_BREAKER);
+ }
Join dynamicFilterJoin =
new LogicalJoin(join.getCluster(), join.getTraitSet(), left.getInput(), dynamicBroadcastExchange,
join.getCondition(), join.getVariablesSet(), join.getJoinType(), join.isSemiJoinDone(),
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
index 87a8bda603ea..e2bb53faf595 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
@@ -24,12 +24,14 @@
import java.io.IOException;
import java.util.concurrent.TimeUnit;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.proto.Mailbox.MailboxContent;
import org.apache.pinot.common.proto.PinotMailboxGrpc;
import org.apache.pinot.query.mailbox.channel.ChannelManager;
import org.apache.pinot.query.mailbox.channel.MailboxStatusObserver;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -45,16 +47,19 @@ public class GrpcSendingMailbox implements SendingMailbox {
private final String _hostname;
private final int _port;
private final long _deadlineMs;
+ private final StatMap _statMap;
private final MailboxStatusObserver _statusObserver = new MailboxStatusObserver();
private StreamObserver _contentObserver;
- public GrpcSendingMailbox(String id, ChannelManager channelManager, String hostname, int port, long deadlineMs) {
+ public GrpcSendingMailbox(String id, ChannelManager channelManager, String hostname, int port, long deadlineMs,
+ StatMap statMap) {
_id = id;
_channelManager = channelManager;
_hostname = hostname;
_port = port;
_deadlineMs = deadlineMs;
+ _statMap = statMap;
}
@Override
@@ -64,25 +69,31 @@ public void send(TransferableBlock block)
return;
}
if (LOGGER.isDebugEnabled()) {
- LOGGER.debug("==[GRPC SEND]== sending data to: " + _id);
+ LOGGER.debug("==[GRPC SEND]== sending message " + block + " to: " + _id);
}
if (_contentObserver == null) {
_contentObserver = getContentObserver();
}
_contentObserver.onNext(toMailboxContent(block));
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("==[GRPC SEND]== message " + block + " sent to: " + _id);
+ }
}
@Override
public void complete() {
if (isTerminated()) {
+ LOGGER.debug("Already terminated mailbox: {}", _id);
return;
}
+ LOGGER.debug("Completing mailbox: {}", _id);
_contentObserver.onCompleted();
}
@Override
public void cancel(Throwable t) {
if (isTerminated()) {
+ LOGGER.debug("Already terminated mailbox: {}", _id);
return;
}
LOGGER.debug("Cancelling mailbox: {}", _id);
@@ -118,9 +129,22 @@ private StreamObserver getContentObserver() {
private MailboxContent toMailboxContent(TransferableBlock block)
throws IOException {
- DataBlock dataBlock = block.getDataBlock();
- byte[] bytes = dataBlock.toBytes();
- ByteString byteString = UnsafeByteOperations.unsafeWrap(bytes);
- return MailboxContent.newBuilder().setMailboxId(_id).setPayload(byteString).build();
+ _statMap.merge(MailboxSendOperator.StatKey.RAW_MESSAGES, 1);
+ long start = System.currentTimeMillis();
+ try {
+ DataBlock dataBlock = block.getDataBlock();
+ byte[] bytes = dataBlock.toBytes();
+ ByteString byteString = UnsafeByteOperations.unsafeWrap(bytes);
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("Serialized block: {} to {} bytes", block, bytes.length);
+ }
+ _statMap.merge(MailboxSendOperator.StatKey.SERIALIZED_BYTES, bytes.length);
+ return MailboxContent.newBuilder().setMailboxId(_id).setPayload(byteString).build();
+ } catch (Throwable t) {
+ LOGGER.warn("Caught exception while serializing block: {}", block, t);
+ throw t;
+ } finally {
+ _statMap.merge(MailboxSendOperator.StatKey.SERIALIZATION_TIME_MS, System.currentTimeMillis() - start);
+ }
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
index d1452712bc33..8adf8db073b3 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
@@ -19,8 +19,10 @@
package org.apache.pinot.query.mailbox;
import java.util.concurrent.TimeoutException;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.spi.exception.QueryCancelledException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -36,11 +38,14 @@ public class InMemorySendingMailbox implements SendingMailbox {
private ReceivingMailbox _receivingMailbox;
private volatile boolean _isTerminated;
private volatile boolean _isEarlyTerminated;
+ private final StatMap _statMap;
- public InMemorySendingMailbox(String id, MailboxService mailboxService, long deadlineMs) {
+ public InMemorySendingMailbox(String id, MailboxService mailboxService, long deadlineMs,
+ StatMap statMap) {
_id = id;
_mailboxService = mailboxService;
_deadlineMs = deadlineMs;
+ _statMap = statMap;
}
@Override
@@ -52,8 +57,10 @@ public void send(TransferableBlock block)
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
+ _statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, timeoutMs);
+
switch (status) {
case SUCCESS:
break;
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
index 7c171f1bf43b..c71af0386d13 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
@@ -23,8 +23,10 @@
import com.google.common.cache.RemovalListener;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.mailbox.channel.ChannelManager;
import org.apache.pinot.query.mailbox.channel.GrpcMailboxServer;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -99,11 +101,12 @@ public int getPort() {
* not open the underlying channel or acquire any additional resources. Instead, it will initialize lazily when the
* data is sent for the first time.
*/
- public SendingMailbox getSendingMailbox(String hostname, int port, String mailboxId, long deadlineMs) {
+ public SendingMailbox getSendingMailbox(String hostname, int port, String mailboxId, long deadlineMs,
+ StatMap statMap) {
if (_hostname.equals(hostname) && _port == port) {
- return new InMemorySendingMailbox(mailboxId, this, deadlineMs);
+ return new InMemorySendingMailbox(mailboxId, this, deadlineMs, statMap);
} else {
- return new GrpcSendingMailbox(mailboxId, _channelManager, hostname, port, deadlineMs);
+ return new GrpcSendingMailbox(mailboxId, _channelManager, hostname, port, deadlineMs, statMap);
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
index 97c8731f95a8..e58b5f7dd58a 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
@@ -19,12 +19,19 @@
package org.apache.pinot.query.mailbox;
import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.Nullable;
+import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datablock.DataBlockUtils;
+import org.apache.pinot.common.datablock.MetadataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.slf4j.Logger;
@@ -50,9 +57,11 @@ public class ReceivingMailbox {
private final BlockingQueue _blocks = new ArrayBlockingQueue<>(DEFAULT_MAX_PENDING_BLOCKS);
private final AtomicReference _errorBlock = new AtomicReference<>();
private volatile boolean _isEarlyTerminated = false;
+ private long _lastArriveTime = System.currentTimeMillis();
@Nullable
private volatile Reader _reader;
+ private final StatMap _stats = new StatMap<>(StatKey.class);
public ReceivingMailbox(String id) {
_id = id;
@@ -72,11 +81,53 @@ public String getId() {
return _id;
}
+ /**
+ * Offers a raw block into the mailbox within the timeout specified, returns whether the block is successfully added.
+ * If the block is not added, an error block is added to the mailbox.
+ *
+ * Contrary to {@link #offer(TransferableBlock, long)}, the block may be an
+ * {@link TransferableBlock#isErrorBlock() error block}.
+ */
+ public ReceivingMailboxStatus offerRaw(ByteBuffer byteBuffer, long timeoutMs)
+ throws IOException {
+ TransferableBlock block;
+ long now = System.currentTimeMillis();
+ _stats.merge(StatKey.WAIT_CPU_TIME_MS, now - _lastArriveTime);
+ _lastArriveTime = now;
+ _stats.merge(StatKey.DESERIALIZED_BYTES, byteBuffer.remaining());
+ _stats.merge(StatKey.DESERIALIZED_MESSAGES, 1);
+
+ now = System.currentTimeMillis();
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(byteBuffer);
+ _stats.merge(StatKey.DESERIALIZATION_TIME_MS, System.currentTimeMillis() - now);
+
+ if (dataBlock instanceof MetadataBlock) {
+ Map exceptions = dataBlock.getExceptions();
+ if (exceptions.isEmpty()) {
+ block = TransferableBlockUtils.wrap(dataBlock);
+ } else {
+ setErrorBlock(TransferableBlockUtils.getErrorTransferableBlock(exceptions));
+ return ReceivingMailboxStatus.FIRST_ERROR;
+ }
+ } else {
+ block = TransferableBlockUtils.wrap(dataBlock);
+ }
+ return offerPrivate(block, timeoutMs);
+ }
+
+ public ReceivingMailboxStatus offer(TransferableBlock block, long timeoutMs) {
+ long now = System.currentTimeMillis();
+ _stats.merge(StatKey.WAIT_CPU_TIME_MS, now - _lastArriveTime);
+ _lastArriveTime = now;
+ _stats.merge(StatKey.IN_MEMORY_MESSAGES, 1);
+ return offerPrivate(block, timeoutMs);
+ }
+
/**
* Offers a non-error block into the mailbox within the timeout specified, returns whether the block is successfully
* added. If the block is not added, an error block is added to the mailbox.
*/
- public ReceivingMailboxStatus offer(TransferableBlock block, long timeoutMs) {
+ private ReceivingMailboxStatus offerPrivate(TransferableBlock block, long timeoutMs) {
TransferableBlock errorBlock = _errorBlock.get();
if (errorBlock != null) {
LOGGER.debug("Mailbox: {} is already cancelled or errored out, ignoring the late block", _id);
@@ -90,7 +141,10 @@ public ReceivingMailboxStatus offer(TransferableBlock block, long timeoutMs) {
return ReceivingMailboxStatus.TIMEOUT;
}
try {
- if (_blocks.offer(block, timeoutMs, TimeUnit.MILLISECONDS)) {
+ long now = System.currentTimeMillis();
+ boolean accepted = _blocks.offer(block, timeoutMs, TimeUnit.MILLISECONDS);
+ _stats.merge(StatKey.OFFER_CPU_TIME_MS, System.currentTimeMillis() - now);
+ if (accepted) {
errorBlock = _errorBlock.get();
if (errorBlock == null) {
if (LOGGER.isDebugEnabled()) {
@@ -170,11 +224,40 @@ private void notifyReader() {
}
}
+ public StatMap getStatMap() {
+ return _stats;
+ }
+
public interface Reader {
void blockReadyToRead();
}
public enum ReceivingMailboxStatus {
- SUCCESS, ERROR, TIMEOUT, CANCELLED, EARLY_TERMINATED
+ SUCCESS, FIRST_ERROR, ERROR, TIMEOUT, CANCELLED, EARLY_TERMINATED
+ }
+
+ public enum StatKey implements StatMap.Key {
+ DESERIALIZED_MESSAGES(StatMap.Type.INT),
+ DESERIALIZED_BYTES(StatMap.Type.LONG),
+ DESERIALIZATION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ IN_MEMORY_MESSAGES(StatMap.Type.INT),
+ OFFER_CPU_TIME_MS(StatMap.Type.LONG),
+ WAIT_CPU_TIME_MS(StatMap.Type.LONG);
+
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java
index ce91701c3ade..1d817f91aaed 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java
@@ -20,16 +20,12 @@
import io.grpc.Context;
import io.grpc.stub.StreamObserver;
-import java.util.Map;
+import java.nio.ByteBuffer;
import java.util.concurrent.TimeUnit;
-import org.apache.pinot.common.datablock.DataBlock;
-import org.apache.pinot.common.datablock.DataBlockUtils;
-import org.apache.pinot.common.datablock.MetadataBlock;
import org.apache.pinot.common.proto.Mailbox.MailboxContent;
import org.apache.pinot.common.proto.Mailbox.MailboxStatus;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.mailbox.ReceivingMailbox;
-import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -62,22 +58,9 @@ public void onNext(MailboxContent mailboxContent) {
_mailbox = _mailboxService.getReceivingMailbox(mailboxId);
}
try {
- TransferableBlock block;
- DataBlock dataBlock = DataBlockUtils.getDataBlock(mailboxContent.getPayload().asReadOnlyByteBuffer());
- if (dataBlock instanceof MetadataBlock) {
- Map exceptions = dataBlock.getExceptions();
- if (exceptions.isEmpty()) {
- block = TransferableBlockUtils.getEndOfStreamTransferableBlock(((MetadataBlock) dataBlock).getStats());
- } else {
- _mailbox.setErrorBlock(TransferableBlockUtils.getErrorTransferableBlock(exceptions));
- return;
- }
- } else {
- block = new TransferableBlock(dataBlock);
- }
-
long timeoutMs = Context.current().getDeadline().timeRemaining(TimeUnit.MILLISECONDS);
- ReceivingMailbox.ReceivingMailboxStatus status = _mailbox.offer(block, timeoutMs);
+ ByteBuffer buffer = mailboxContent.getPayload().asReadOnlyByteBuffer();
+ ReceivingMailbox.ReceivingMailboxStatus status = _mailbox.offerRaw(buffer, timeoutMs);
switch (status) {
case SUCCESS:
_responseObserver.onNext(MailboxStatus.newBuilder().setMailboxId(mailboxId)
@@ -88,6 +71,8 @@ public void onNext(MailboxContent mailboxContent) {
LOGGER.warn("Mailbox: {} already cancelled from upstream", mailboxId);
cancelStream();
break;
+ case FIRST_ERROR:
+ return;
case ERROR:
LOGGER.warn("Mailbox: {} already errored out (received error block before)", mailboxId);
cancelStream();
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java
new file mode 100644
index 000000000000..16cdda026e90
--- /dev/null
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java
@@ -0,0 +1,207 @@
+/**
+ * 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.runtime;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.function.IntFunction;
+import org.apache.pinot.query.planner.plannode.AbstractPlanNode;
+import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.ExchangeNode;
+import org.apache.pinot.query.planner.plannode.FilterNode;
+import org.apache.pinot.query.planner.plannode.JoinNode;
+import org.apache.pinot.query.planner.plannode.MailboxReceiveNode;
+import org.apache.pinot.query.planner.plannode.MailboxSendNode;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.planner.plannode.PlanNodeVisitor;
+import org.apache.pinot.query.planner.plannode.ProjectNode;
+import org.apache.pinot.query.planner.plannode.SetOpNode;
+import org.apache.pinot.query.planner.plannode.SortNode;
+import org.apache.pinot.query.planner.plannode.TableScanNode;
+import org.apache.pinot.query.planner.plannode.ValueNode;
+import org.apache.pinot.query.planner.plannode.WindowNode;
+import org.apache.pinot.query.runtime.operator.MultiStageOperator;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class InStageStatsTreeBuilder implements PlanNodeVisitor