fix DataSchema thread-safe issue

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

@@ -48,7 +48,7 @@
 public class DataSchema {
   private final String[] _columnNames;
   private final ColumnDataType[] _columnDataTypes;
-  private ColumnDataType[] _storedColumnDataTypes;
+  private final ColumnDataType[] _storedColumnDataTypes;
 
   /** Used by both Broker and Server to generate results for EXPLAIN PLAN queries. */
   public static final DataSchema EXPLAIN_RESULT_SCHEMA =
@@ -61,6 +61,7 @@ public DataSchema(@JsonProperty("columnNames") String[] columnNames,
       @JsonProperty("columnDataTypes") ColumnDataType[] columnDataTypes) {
     _columnNames = columnNames;
     _columnDataTypes = columnDataTypes;
+    _storedColumnDataTypes = computeStoredColumnDataType(columnDataTypes);
   }
 
   public int size() {
@@ -85,13 +86,6 @@ public ColumnDataType[] getColumnDataTypes() {
 
   @JsonIgnore
   public ColumnDataType[] getStoredColumnDataTypes() {
-    if (_storedColumnDataTypes == null) {
-      int numColumns = _columnDataTypes.length;
-      _storedColumnDataTypes = new ColumnDataType[numColumns];
-      for (int i = 0; i < numColumns; i++) {
-        _storedColumnDataTypes[i] = _columnDataTypes[i].getStoredType();
-      }
-    }
     return _storedColumnDataTypes;
   }
 
@@ -126,29 +120,34 @@ public boolean isTypeCompatibleWith(DataSchema anotherDataSchema) {
    * <code>LONG</code>.
    * <p>NOTE: The given data schema should be type compatible with this one.
    *
+   * @param originalSchema the original Data schema
    * @param anotherDataSchema Data schema to cover
    */
-  public void upgradeToCover(DataSchema anotherDataSchema) {
-    int numColumns = _columnDataTypes.length;
+  public static DataSchema upgradeToCover(DataSchema originalSchema, DataSchema anotherDataSchema) {
+    int numColumns = originalSchema._columnDataTypes.length;
+    ColumnDataType[] columnDataTypes = new ColumnDataType[numColumns];
     for (int i = 0; i < numColumns; i++) {
-      ColumnDataType thisColumnDataType = _columnDataTypes[i];
+      ColumnDataType thisColumnDataType = originalSchema._columnDataTypes[i];
       ColumnDataType thatColumnDataType = anotherDataSchema._columnDataTypes[i];
       if (thisColumnDataType != thatColumnDataType) {
         if (thisColumnDataType.isArray()) {
           if (thisColumnDataType.isWholeNumberArray() && thatColumnDataType.isWholeNumberArray()) {
-            _columnDataTypes[i] = ColumnDataType.LONG_ARRAY;
+            columnDataTypes[i] = ColumnDataType.LONG_ARRAY;
           } else {
-            _columnDataTypes[i] = ColumnDataType.DOUBLE_ARRAY;
+            columnDataTypes[i] = ColumnDataType.DOUBLE_ARRAY;
           }
         } else {
           if (thisColumnDataType.isWholeNumber() && thatColumnDataType.isWholeNumber()) {
-            _columnDataTypes[i] = ColumnDataType.LONG;
+            columnDataTypes[i] = ColumnDataType.LONG;
           } else {
-            _columnDataTypes[i] = ColumnDataType.DOUBLE;
+            columnDataTypes[i] = ColumnDataType.DOUBLE;
           }
         }
+      } else {
+        columnDataTypes[i] = originalSchema._columnDataTypes[i];
       }
     }
+    return new DataSchema(originalSchema._columnNames, columnDataTypes);
   }
 
   public byte[] toBytes()
@@ -240,6 +239,15 @@ public int hashCode() {
     return EqualityUtils.hashCodeOf(Arrays.hashCode(_columnNames), Arrays.hashCode(_columnDataTypes));
   }
 
+  private static ColumnDataType[] computeStoredColumnDataType(ColumnDataType[] columnDataTypes) {
+    int numColumns = columnDataTypes.length;
+    ColumnDataType[] storedColumnDataTypes = new ColumnDataType[numColumns];
+    for (int i = 0; i < numColumns; i++) {
+      storedColumnDataTypes[i] = columnDataTypes[i].getStoredType();
+    }
+    return storedColumnDataTypes;
+  }
+
   public enum ColumnDataType {
     INT(0),
     LONG(0L),

pinot-common/src/test/java/org/apache/pinot/common/utils/DataSchemaTest.java

@@ -82,7 +82,7 @@ public void testTypeCompatible() {
     DataSchema incompatibleDataSchema = new DataSchema(anotherColumnNames, COLUMN_DATA_TYPES);
     Assert.assertFalse(dataSchema.isTypeCompatibleWith(incompatibleDataSchema));
 
-    dataSchema.upgradeToCover(compatibleDataSchema);
+    dataSchema = DataSchema.upgradeToCover(dataSchema, compatibleDataSchema);
     DataSchema upgradedDataSchema = new DataSchema(COLUMN_NAMES, UPGRADED_COLUMN_DATA_TYPES);
     Assert.assertEquals(dataSchema, upgradedDataSchema);
   }

pinot-core/src/main/java/org/apache/pinot/core/query/reduce/SelectionDataTableReducer.java

@@ -118,7 +118,7 @@ private List<ServerRoutingInstance> removeConflictingResponses(DataSchema dataSc
         droppedServers.add(entry.getKey());
         iterator.remove();
       } else {
-        dataSchema.upgradeToCover(dataSchemaToCompare);
+        dataSchema = DataSchema.upgradeToCover(dataSchema, dataSchemaToCompare);
       }
     }
     return droppedServers;

pinot-core/src/test/java/org/apache/pinot/core/query/selection/SelectionOperatorServiceTest.java

@@ -208,7 +208,7 @@ public void testCompatibleRowsDataTableTransformation()
     rows.add(_compatibleRow1);
     DataSchema dataSchema = _dataSchema.clone();
     assertTrue(dataSchema.isTypeCompatibleWith(_compatibleDataSchema));
-    dataSchema.upgradeToCover(_compatibleDataSchema);
+    dataSchema = DataSchema.upgradeToCover(dataSchema, _compatibleDataSchema);
     assertEquals(dataSchema, _upgradedDataSchema);
     DataTable dataTable = SelectionOperatorUtils.getDataTableFromRows(rows, dataSchema, false);
     Object[] expectedRow1 = {