Performance (Both Memory and CPU) Improvements for Indexed Table

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -408,6 +408,11 @@ public static boolean isAccurateGroupByWithoutOrderBy(Map<String, String> queryO
         queryOptions.getOrDefault(QueryOptionKey.ACCURATE_GROUP_BY_WITHOUT_ORDER_BY, "false"));
   }
 
+  public static boolean isEnableOptimizedIndexedTable(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(
+        queryOptions.getOrDefault(QueryOptionKey.ENABLE_OPTIMIZED_INDEXED_TABLE, "false"));
+  }
+
   public static Boolean isUseMSEToFillEmptySchema(Map<String, String> queryOptions, boolean defaultValue) {
     String useMSEToFillEmptySchema = queryOptions.get(QueryOptionKey.USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA);
     return useMSEToFillEmptySchema != null ? Boolean.parseBoolean(useMSEToFillEmptySchema) : defaultValue;

pinot-core/src/main/java/org/apache/pinot/core/data/table/ConcurrentIndexedTableOptimized.java

@@ -0,0 +1,140 @@
+/**
+ * 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.data.table;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.StampedLock;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.query.request.context.QueryContext;
+
+
+/**
+ * OPTIMIZED Thread safe {@link Table} implementation for aggregating Records based on combination of keys
+ *
+ * Key optimizations:
+ * 1. Replaced ReentrantReadWriteLock with StampedLock for better read performance
+ * 2. Uses optimistic reads for the common case (no resize)
+ * 3. Added volatile _resizing flag to reduce double-check overhead
+ * 4. Optimized upsertWithOrderBy to use tryOptimisticRead first
+ */
+public class ConcurrentIndexedTableOptimized extends IndexedTableOptimized {
+  private final AtomicBoolean _noMoreNewRecords = new AtomicBoolean();
+  private final StampedLock _stampedLock = new StampedLock();
+  private volatile boolean _resizing = false;
+
+  public ConcurrentIndexedTableOptimized(DataSchema dataSchema, boolean hasFinalInput, QueryContext queryContext,
+      int resultSize, int trimSize, int trimThreshold, int initialCapacity, ExecutorService executorService) {
+    super(dataSchema, hasFinalInput, queryContext, resultSize, trimSize, trimThreshold,
+        new ConcurrentHashMap<>(initialCapacity), executorService);
+  }
+
+  /**
+   * Thread safe implementation of upsert for inserting {@link Record} into {@link Table}
+   */
+  @Override
+  public boolean upsert(Key key, Record record) {
+    if (_hasOrderBy) {
+      upsertWithOrderBy(key, record);
+    } else {
+      upsertWithoutOrderBy(key, record);
+    }
+    return true;
+  }
+
+  /**
+   * OPTIMIZATION: Uses StampedLock with optimistic reads for better performance
+   * Optimistic read attempts first, falls back to read lock only if validation fails
+   */
+  protected void upsertWithOrderBy(Key key, Record record) {
+    // Try optimistic read first - this is lock-free and very fast
+    long stamp = _stampedLock.tryOptimisticRead();
+    if (!_resizing && stamp != 0) {
+      addOrUpdateRecord(key, record);
+      // Validate that no write occurred during our optimistic read
+      if (_stampedLock.validate(stamp)) {
+        // Success - check if we need to resize
+        if (_lookupMap.size() >= _trimThreshold) {
+          tryResize();
+        }
+        return;
+      }
+    }
+
+    // Optimistic read failed, fall back to read lock
+    stamp = _stampedLock.readLock();
+    try {
+      addOrUpdateRecord(key, record);
+    } finally {
+      _stampedLock.unlockRead(stamp);
+    }
+
+    // Check resize after releasing read lock
+    if (_lookupMap.size() >= _trimThreshold) {
+      tryResize();
+    }
+  }
+
+  /**
+   * OPTIMIZATION: Extracted resize logic to reduce code duplication
+   * Uses volatile flag to prevent thundering herd on resize
+   */
+  private void tryResize() {
+    // Quick volatile check to avoid lock contention
+    if (_resizing) {
+      return;
+    }
+
+    long stamp = _stampedLock.tryWriteLock();
+    if (stamp == 0) {
+      // Another thread is resizing, skip
+      return;
+    }
+
+    try {
+      // Double-check with lock held
+      if (_lookupMap.size() >= _trimThreshold && !_resizing) {
+        _resizing = true;
+        try {
+          resize();
+        } finally {
+          _resizing = false;
+        }
+      }
+    } finally {
+      _stampedLock.unlockWrite(stamp);
+    }
+  }
+
+  /**
+   * OPTIMIZATION: No changes needed here - already optimal
+   * The noMoreNewRecords flag prevents unnecessary map lookups
+   */
+  protected void upsertWithoutOrderBy(Key key, Record record) {
+    if (_noMoreNewRecords.get()) {
+      updateExistingRecord(key, record);
+    } else {
+      addOrUpdateRecord(key, record);
+      if (_lookupMap.size() >= _resultSize) {
+        _noMoreNewRecords.set(true);
+      }
+    }
+  }
+}