DataTable V3 implementation and measure data table serialization cost on server

pinot-common/src/main/java/org/apache/pinot/common/utils/CommonConstants.java

@@ -339,6 +339,9 @@ public static class SegmentCompletionProtocol {
     public static final String CONFIG_OF_ENABLE_THREAD_CPU_TIME_MEASUREMENT =
         "pinot.server.instance.enableThreadCpuTimeMeasurement";
     public static final boolean DEFAULT_ENABLE_THREAD_CPU_TIME_MEASUREMENT = false;
+
+    public static final String CONFIG_OF_CURRENT_DATA_TABLE_VERSION = "pinot.server.instance.currentDataTableVersion";
+    public static final int DEFAULT_CURRENT_DATA_TABLE_VERSION = 3;
   }
 
   public static class Controller {

pinot-common/src/main/java/org/apache/pinot/common/utils/DataTable.java

@@ -19,7 +19,9 @@
 package org.apache.pinot.common.utils;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.Map;
+import javax.annotation.Nullable;
 import org.apache.pinot.common.response.ProcessingException;
 import org.apache.pinot.spi.utils.ByteArray;
 
@@ -44,10 +46,11 @@ public interface DataTable {
   String REQUEST_ID_METADATA_KEY = "requestId";
   String NUM_RESIZES_METADATA_KEY = "numResizes";
   String RESIZE_TIME_MS_METADATA_KEY = "resizeTimeMs";
-  String EXECUTION_THREAD_CPU_TIME_NS_METADATA_KEY = "executionThreadCpuTimeNs";
 
   void addException(ProcessingException processingException);
 
+  Map<Integer, String> getExceptions();
+
   byte[] toBytes()
       throws IOException;
 
@@ -80,4 +83,77 @@ byte[] toBytes()
   double[] getDoubleArray(int rowId, int colId);
 
   String[] getStringArray(int rowId, int colId);
+
+  enum MetadataValueType {
+    INT, LONG, STRING
+  }
+
+  /* The MetadataKey is used in V3, where we present metadata as Map<MetadataKey, String>
+   * ATTENTION:
+   *  - Don't change existing keys.
+   *  - Don't remove existing keys.
+   *  - Always add new keys to the end.
+   *  Otherwise, backward compatibility will be broken.
+   */
+  enum MetadataKey {
+    UNKNOWN("unknown", MetadataValueType.STRING),
+    TABLE("table", MetadataValueType.STRING), // NOTE: this key is only used in PrioritySchedulerTest
+    NUM_DOCS_SCANNED("numDocsScanned", MetadataValueType.LONG),
+    NUM_ENTRIES_SCANNED_IN_FILTER("numEntriesScannedInFilter", MetadataValueType.LONG),
+    NUM_ENTRIES_SCANNED_POST_FILTER("numEntriesScannedPostFilter", MetadataValueType.LONG),
+    NUM_SEGMENTS_QUERIED("numSegmentsQueried", MetadataValueType.INT),
+    NUM_SEGMENTS_PROCESSED("numSegmentsProcessed", MetadataValueType.INT),
+    NUM_SEGMENTS_MATCHED("numSegmentsMatched", MetadataValueType.INT),
+    NUM_CONSUMING_SEGMENTS_PROCESSED("numConsumingSegmentsProcessed", MetadataValueType.INT),
+    MIN_CONSUMING_FRESHNESS_TIME_MS("minConsumingFreshnessTimeMs", MetadataValueType.LONG),
+    TOTAL_DOCS("totalDocs", MetadataValueType.LONG),
+    NUM_GROUPS_LIMIT_REACHED("numGroupsLimitReached", MetadataValueType.STRING),
+    TIME_USED_MS("timeUsedMs", MetadataValueType.LONG),
+    TRACE_INFO("traceInfo", MetadataValueType.STRING),
+    REQUEST_ID("requestId", MetadataValueType.LONG),
+    NUM_RESIZES("numResizes", MetadataValueType.INT),
+    RESIZE_TIME_MS("resizeTimeMs", MetadataValueType.LONG),
+    THREAD_CPU_TIME_NS("threadCpuTimeNs", MetadataValueType.LONG);
+
+    private static final Map<String, MetadataKey> _nameToEnumKeyMap = new HashMap<>();
+    private final String _name;
+    private final MetadataValueType _valueType;
+
+    MetadataKey(String name, MetadataValueType valueType) {
+      _name = name;
+      _valueType = valueType;
+    }
+
+    // getByOrdinal returns an enum key for a given ordinal or null if the key does not exist.
+    @Nullable
+    public static MetadataKey getByOrdinal(int ordinal) {
+      if (ordinal >= MetadataKey.values().length) {
+        return null;
+      }
+      return MetadataKey.values()[ordinal];
+    }
+
+    // getByName returns an enum key for a given name or null if the key does not exist.
+    public static MetadataKey getByName(String name) {
+      return _nameToEnumKeyMap.getOrDefault(name, null);
+    }
+
+    // getName returns the associated name(string) of the enum key.
+    public String getName() {
+      return _name;
+    }
+
+    // getValueType returns the value type(int/long/String) of the enum key.
+    public MetadataValueType getValueType() {
+      return _valueType;
+    }
+
+    static {
+      for (MetadataKey key : MetadataKey.values()) {
+        if (_nameToEnumKeyMap.put(key.getName(), key) != null) {
+          throw new IllegalArgumentException("Duplicate name defined in the MetadataKey definition: " + key.getName());
+        }
+      }
+    }
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTable.java

@@ -0,0 +1,283 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.common.datatable;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataTable;
+import org.apache.pinot.common.utils.StringUtil;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.BytesUtils;
+
+import static org.apache.pinot.core.common.datatable.DataTableUtils.decodeString;
+
+
+/**
+ * Base implementation of the DataTable interface.
+ */
+public abstract class BaseDataTable implements DataTable {
+  protected int _numRows;
+  protected int _numColumns;
+  protected DataSchema _dataSchema;
+  protected int[] _columnOffsets;
+  protected int _rowSizeInBytes;
+  protected Map<String, Map<Integer, String>> _dictionaryMap;
+  protected byte[] _fixedSizeDataBytes;
+  protected ByteBuffer _fixedSizeData;
+  protected byte[] _variableSizeDataBytes;
+  protected ByteBuffer _variableSizeData;
+  protected Map<String, String> _metadata;
+
+  public BaseDataTable(int numRows, DataSchema dataSchema, Map<String, Map<Integer, String>> dictionaryMap,
+      byte[] fixedSizeDataBytes, byte[] variableSizeDataBytes) {
+    _numRows = numRows;
+    _numColumns = dataSchema.size();
+    _dataSchema = dataSchema;
+    _columnOffsets = new int[_numColumns];
+    _rowSizeInBytes = DataTableUtils.computeColumnOffsets(dataSchema, _columnOffsets);
+    _dictionaryMap = dictionaryMap;
+    _fixedSizeDataBytes = fixedSizeDataBytes;
+    _fixedSizeData = ByteBuffer.wrap(fixedSizeDataBytes);
+    _variableSizeDataBytes = variableSizeDataBytes;
+    _variableSizeData = ByteBuffer.wrap(variableSizeDataBytes);
+    _metadata = new HashMap<>();
+  }
+
+  /**
+   * Construct empty data table. (Server side)
+   */
+  public BaseDataTable() {
+    _numRows = 0;
+    _numColumns = 0;
+    _dataSchema = null;
+    _columnOffsets = null;
+    _rowSizeInBytes = 0;
+    _dictionaryMap = null;
+    _fixedSizeDataBytes = null;
+    _fixedSizeData = null;
+    _variableSizeDataBytes = null;
+    _variableSizeData = null;
+    _metadata = new HashMap<>();
+  }
+
+  /**
+   * Helper method to serialize dictionary map.
+   */
+  protected byte[] serializeDictionaryMap()
+      throws IOException {
+    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+    DataOutputStream dataOutputStream = new DataOutputStream(byteArrayOutputStream);
+
+    dataOutputStream.writeInt(_dictionaryMap.size());
+    for (Map.Entry<String, Map<Integer, String>> dictionaryMapEntry : _dictionaryMap.entrySet()) {
+      String columnName = dictionaryMapEntry.getKey();
+      Map<Integer, String> dictionary = dictionaryMapEntry.getValue();
+      byte[] bytes = StringUtil.encodeUtf8(columnName);
+      dataOutputStream.writeInt(bytes.length);
+      dataOutputStream.write(bytes);
+      dataOutputStream.writeInt(dictionary.size());
+
+      for (Map.Entry<Integer, String> dictionaryEntry : dictionary.entrySet()) {
+        dataOutputStream.writeInt(dictionaryEntry.getKey());
+        byte[] valueBytes = StringUtil.encodeUtf8(dictionaryEntry.getValue());
+        dataOutputStream.writeInt(valueBytes.length);
+        dataOutputStream.write(valueBytes);
+      }
+    }
+
+    return byteArrayOutputStream.toByteArray();
+  }
+
+  /**
+   * Helper method to deserialize dictionary map.
+   */
+  protected Map<String, Map<Integer, String>> deserializeDictionaryMap(byte[] bytes)
+      throws IOException {
+    try (ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes);
+        DataInputStream dataInputStream = new DataInputStream(byteArrayInputStream)) {
+      int numDictionaries = dataInputStream.readInt();
+      Map<String, Map<Integer, String>> dictionaryMap = new HashMap<>(numDictionaries);
+
+      for (int i = 0; i < numDictionaries; i++) {
+        String column = decodeString(dataInputStream);
+        int dictionarySize = dataInputStream.readInt();
+        Map<Integer, String> dictionary = new HashMap<>(dictionarySize);
+        for (int j = 0; j < dictionarySize; j++) {
+          int key = dataInputStream.readInt();
+          String value = decodeString(dataInputStream);
+          dictionary.put(key, value);
+        }
+        dictionaryMap.put(column, dictionary);
+      }
+
+      return dictionaryMap;
+    }
+  }
+
+  public Map<String, String> getMetadata() {
+    return _metadata;
+  }
+
+  public DataSchema getDataSchema() {
+    return _dataSchema;
+  }
+
+  public int getNumberOfRows() {
+    return _numRows;
+  }
+
+  public int getInt(int rowId, int colId) {
+    _fixedSizeData.position(rowId * _rowSizeInBytes + _columnOffsets[colId]);
+    return _fixedSizeData.getInt();
+  }
+
+  public long getLong(int rowId, int colId) {
+    _fixedSizeData.position(rowId * _rowSizeInBytes + _columnOffsets[colId]);
+    return _fixedSizeData.getLong();
+  }
+
+  public float getFloat(int rowId, int colId) {
+    _fixedSizeData.position(rowId * _rowSizeInBytes + _columnOffsets[colId]);
+    return _fixedSizeData.getFloat();
+  }
+
+  public double getDouble(int rowId, int colId) {
+    _fixedSizeData.position(rowId * _rowSizeInBytes + _columnOffsets[colId]);
+    return _fixedSizeData.getDouble();
+  }
+
+  public String getString(int rowId, int colId) {
+    _fixedSizeData.position(rowId * _rowSizeInBytes + _columnOffsets[colId]);
+    int dictId = _fixedSizeData.getInt();
+    return _dictionaryMap.get(_dataSchema.getColumnName(colId)).get(dictId);
+  }
+
+  public ByteArray getBytes(int rowId, int colId) {
+    // NOTE: DataTable V2/V3 uses String to store BYTES value
+    return BytesUtils.toByteArray(getString(rowId, colId));
+  }
+
+  public <T> T getObject(int rowId, int colId) {
+    int size = positionCursorInVariableBuffer(rowId, colId);
+    int objectTypeValue = _variableSizeData.getInt();
+    ByteBuffer byteBuffer = _variableSizeData.slice();
+    byteBuffer.limit(size);
+    return ObjectSerDeUtils.deserialize(byteBuffer, objectTypeValue);
+  }
+
+  public int[] getIntArray(int rowId, int colId) {
+    int length = positionCursorInVariableBuffer(rowId, colId);
+    int[] ints = new int[length];
+    for (int i = 0; i < length; i++) {
+      ints[i] = _variableSizeData.getInt();
+    }
+    return ints;
+  }
+
+  public long[] getLongArray(int rowId, int colId) {
+    int length = positionCursorInVariableBuffer(rowId, colId);
+    long[] longs = new long[length];
+    for (int i = 0; i < length; i++) {
+      longs[i] = _variableSizeData.getLong();
+    }
+    return longs;
+  }
+
+  public float[] getFloatArray(int rowId, int colId) {
+    int length = positionCursorInVariableBuffer(rowId, colId);
+    float[] floats = new float[length];
+    for (int i = 0; i < length; i++) {
+      floats[i] = _variableSizeData.getFloat();
+    }
+    return floats;
+  }
+
+  public double[] getDoubleArray(int rowId, int colId) {
+    int length = positionCursorInVariableBuffer(rowId, colId);
+    double[] doubles = new double[length];
+    for (int i = 0; i < length; i++) {
+      doubles[i] = _variableSizeData.getDouble();
+    }
+    return doubles;
+  }
+
+  public String[] getStringArray(int rowId, int colId) {
+    int length = positionCursorInVariableBuffer(rowId, colId);
+    String[] strings = new String[length];
+    Map<Integer, String> dictionary = _dictionaryMap.get(_dataSchema.getColumnName(colId));
+    for (int i = 0; i < length; i++) {
+      strings[i] = dictionary.get(_variableSizeData.getInt());
+    }
+    return strings;
+  }
+
+  private int positionCursorInVariableBuffer(int rowId, int colId) {
+    _fixedSizeData.position(rowId * _rowSizeInBytes + _columnOffsets[colId]);
+    _variableSizeData.position(_fixedSizeData.getInt());
+    return _fixedSizeData.getInt();
+  }
+
+  public String toString() {
+    if (_dataSchema == null) {
+      return _metadata.toString();
+    }
+
+    StringBuilder stringBuilder = new StringBuilder();
+    stringBuilder.append(_dataSchema.toString()).append('\n');
+    stringBuilder.append("numRows: ").append(_numRows).append('\n');
+
+    _fixedSizeData.position(0);
+    for (int rowId = 0; rowId < _numRows; rowId++) {
+      for (int colId = 0; colId < _numColumns; colId++) {
+        switch (_dataSchema.getColumnDataType(colId)) {
+          case INT:
+            stringBuilder.append(_fixedSizeData.getInt());
+            break;
+          case LONG:
+            stringBuilder.append(_fixedSizeData.getLong());
+            break;
+          case FLOAT:
+            stringBuilder.append(_fixedSizeData.getFloat());
+            break;
+          case DOUBLE:
+            stringBuilder.append(_fixedSizeData.getDouble());
+            break;
+          case STRING:
+            stringBuilder.append(_fixedSizeData.getInt());
+            break;
+          // Object and array.
+          default:
+            stringBuilder.append(String.format("(%s:%s)", _fixedSizeData.getInt(), _fixedSizeData.getInt()));
+            break;
+        }
+        stringBuilder.append("\t");
+      }
+      stringBuilder.append("\n");
+    }
+    return stringBuilder.toString();
+  }
+}