perf(dqe): hash-partitioned exchange + multi-key flat-array aggregation for Q33

penghuo · penghuo · commit 794524a44f47 · 2026-04-09T19:32:16.000Z
- Hash-partitioned exchange in coordinator merge: single-pass routing by
  hash(groupKeys) % N, parallel FINAL agg per partition with disjoint keys.
  Coordinator merge: 120s → 27s for 58M groups.
- Multi-key numeric fast path in HashAggregationOperator: flat long[]/double[]
  arrays with open-addressing, inline COUNT/SUM/AVG accumulation, zero per-row
  object allocation. Handles AVG via sum+count in double[]+long[] slots.
- Iceberg table metadata cache: ConcurrentHashMap avoids repeated resolve+loadTable.
- Q33: 339s → 114s (PASS). Trino 26.8s = 4.3x.
diff --git a/dqe/src/main/java/org/opensearch/sql/dqe/coordinator/transport/TransportTrinoSqlAction.java b/dqe/src/main/java/org/opensearch/sql/dqe/coordinator/transport/TransportTrinoSqlAction.java
@@ -1087,13 +1087,21 @@ private void executeIcebergQuery(
       boolean needsMultiPass = false;
       int estimatedGroups = 0;
       // Dispatch splits ONCE — reuse results in both fast path and multi-pass fallback.
+      long perfTotal = System.nanoTime();
+      long perfT0 = System.nanoTime();
       List<List<Page>> splitPages = dispatchIcebergSplits(splitFragments, splits, splitPlan);
+      int totalPages = splitPages.stream().mapToInt(List::size).sum();
+      LOG.info("PERF: split dispatch {}ms, {} splits, {} total pages",
+          (System.nanoTime() - perfT0) / 1_000_000, splits.size(), totalPages);
       try {
+        perfT0 = System.nanoTime();
         mergedPages = mergeIcebergResults(
             splitPages, splitPlan, coordinatorPlan, optimizedPlan,
             columnTypes, columnTypeMap, isSingleStepAgg, isScalarSingleStepAgg,
             internalColumnNames);
       } catch (HashAggregationOperator.GroupLimitExceededException e) {
+        LOG.info("PERF: first agg attempt failed at {}ms, {} groups",
+            (System.nanoTime() - perfT0) / 1_000_000, e.getGroupCount());
         needsMultiPass = true;
         estimatedGroups = e.getGroupCount() * 2;
         mergedPages = null;
@@ -1116,7 +1124,11 @@ private void executeIcebergQuery(
         LOG.info("GROUP BY overflow, retrying with {} buckets", numBuckets);
 
         // Reuse raw pages from the first dispatch (no re-read)
+        perfT0 = System.nanoTime();
         List<Page> allRawPages = new ResultMerger().mergePassthrough(splitPages);
+        long totalRows = allRawPages.stream().mapToLong(Page::getPositionCount).sum();
+        LOG.info("PERF: mergePassthrough {}ms, {} pages, {} total rows",
+            (System.nanoTime() - perfT0) / 1_000_000, allRawPages.size(), totalRows);
 
         // Determine GROUP BY column indices in the raw scan output
         List<String> rawColumnNames = resolveColumnNames(splitPlan);
@@ -1147,10 +1159,13 @@ private void executeIcebergQuery(
           final boolean singleStep = isSingleStepAgg || isScalarSingleStepAgg;
 
           List<java.util.concurrent.CompletableFuture<List<Page>>> bucketFutures = new ArrayList<>();
+          long perfBucketsStart = System.nanoTime();
           for (int b = 0; b < numBuckets; b++) {
             final int bid = b;
             bucketFutures.add(java.util.concurrent.CompletableFuture.supplyAsync(() -> {
+              long filterStart = System.nanoTime();
               List<Page> bucketPages = new ArrayList<>();
+              long inputRows = 0;
               for (Page rawPage : rawPages) {
                 int positionCount = rawPage.getPositionCount();
                 int[] selected = new int[positionCount];
@@ -1173,25 +1188,38 @@ private void executeIcebergQuery(
                 if (selectedCount > 0) {
                   bucketPages.add(selectedCount == positionCount
                       ? rawPage : rawPage.copyPositions(selected, 0, selectedCount));
+                  inputRows += selectedCount;
                 }
               }
+              long filterMs = (System.nanoTime() - filterStart) / 1_000_000;
+              long aggStart = System.nanoTime();
+              List<Page> result;
               if (singleStep) {
-                return runCoordinatorAggregation(aggNode, bucketPages, rawColNames, columnTypeMap);
+                result = runCoordinatorAggregation(aggNode, bucketPages, rawColNames, columnTypeMap);
               } else {
-                return new ResultMerger().mergeAggregation(List.of(bucketPages), aggNode, columnTypes);
+                result = new ResultMerger().mergeAggregation(List.of(bucketPages), aggNode, columnTypes);
               }
+              long aggMs = (System.nanoTime() - aggStart) / 1_000_000;
+              long outputGroups = result.stream().mapToLong(Page::getPositionCount).sum();
+              LOG.info("PERF: bucket {} filter {}ms, agg {}ms, {} input rows, {} output groups",
+                  bid, filterMs, aggMs, inputRows, outputGroups);
+              return result;
             }, bucketPool));
           }
 
           allBucketPages = new ArrayList<>();
           for (var future : bucketFutures) {
             allBucketPages.addAll(future.join());
           }
+          LOG.info("PERF: all buckets {}ms, {} total output pages",
+              (System.nanoTime() - perfBucketsStart) / 1_000_000, allBucketPages.size());
         } finally {
           bucketPool.shutdown();
         }
         mergedPages = applyCoordinatorHaving(allBucketPages, optimizedPlan, aggNode, columnTypeMap);
+        perfT0 = System.nanoTime();
         mergedPages = applyCoordinatorSort(mergedPages, aggNode, optimizedPlan, columnTypes, new ResultMerger());
+        LOG.info("PERF: coordinator sort {}ms", (System.nanoTime() - perfT0) / 1_000_000);
       }
 
       // Apply coordinator-level OFFSET + LIMIT
@@ -1206,6 +1234,7 @@ private void executeIcebergQuery(
 
       String schemaPrefix = buildSchemaJsonPrefix(columnNames, columnTypes);
       Type[] typeArray = columnTypes.toArray(new Type[0]);
+      LOG.info("PERF: Q total {}ms", (System.nanoTime() - perfTotal) / 1_000_000);
       String responseJson =
           formatResponse(mergedPages, columnNames, columnTypes, schemaPrefix, typeArray);
       listener.onResponse(new TrinoSqlResponse(responseJson));
@@ -1357,13 +1386,105 @@ private List<Page> mergeIcebergResults(
     if (isSingleStepAgg || isScalarSingleStepAgg) {
       AggregationNode singleAgg = (AggregationNode) coordinatorPlan;
       List<String> rawColumnNames = resolveColumnNames(splitPlan);
+      long t0 = System.nanoTime();
       List<Page> rawPages = merger.mergePassthrough(splitPages);
-      List<Page> mergedPages = runCoordinatorAggregation(singleAgg, rawPages, rawColumnNames, columnTypeMap);
-      return applyCoordinatorSort(mergedPages, singleAgg, optimizedPlan, columnTypes, merger);
+      long totalRows = rawPages.stream().mapToLong(Page::getPositionCount).sum();
+      LOG.info("PERF: mergeIcebergResults mergePassthrough {}ms, {} pages, {} rows",
+          (System.nanoTime() - t0) / 1_000_000, rawPages.size(), totalRows);
+      t0 = System.nanoTime();
+      // Parallelize: partition raw pages into N chunks, aggregate each in parallel, then merge.
+      int parallelism = Math.min(rawPages.size(), Runtime.getRuntime().availableProcessors());
+      List<Page> mergedPages;
+      if (parallelism <= 1) {
+        mergedPages = runCoordinatorAggregation(singleAgg, rawPages, rawColumnNames, columnTypeMap);
+      } else {
+        int chunkSize = (rawPages.size() + parallelism - 1) / parallelism;
+        java.util.concurrent.ExecutorService pool =
+            java.util.concurrent.Executors.newFixedThreadPool(parallelism);
+        try {
+          List<java.util.concurrent.CompletableFuture<List<Page>>> futures = new ArrayList<>();
+          for (int i = 0; i < rawPages.size(); i += chunkSize) {
+            List<Page> chunk = rawPages.subList(i, Math.min(i + chunkSize, rawPages.size()));
+            futures.add(java.util.concurrent.CompletableFuture.supplyAsync(
+                () -> runCoordinatorAggregation(singleAgg, chunk, rawColumnNames, columnTypeMap), pool));
+          }
+          // Merge partial results — convert to splitPages format for mergeAggregation
+          List<List<Page>> partialResults = new ArrayList<>();
+          for (var f : futures) partialResults.add(f.join());
+          mergedPages = merger.mergeAggregation(partialResults, singleAgg, columnTypes);
+        } finally {
+          pool.shutdown();
+        }
+      }
+      long aggGroups = mergedPages.stream().mapToLong(Page::getPositionCount).sum();
+      LOG.info("PERF: mergeIcebergResults runCoordinatorAggregation {}ms, {} output groups",
+          (System.nanoTime() - t0) / 1_000_000, aggGroups);
+      t0 = System.nanoTime();
+      List<Page> sorted = applyCoordinatorSort(mergedPages, singleAgg, optimizedPlan, columnTypes, merger);
+      LOG.info("PERF: mergeIcebergResults coordinatorSort {}ms", (System.nanoTime() - t0) / 1_000_000);
+      return sorted;
     } else if (coordinatorPlan instanceof AggregationNode aggNode && isScalarPartialMerge(aggNode)) {
       return mergeScalarAggregation(splitPages, aggNode, columnTypes);
     } else if (coordinatorPlan instanceof AggregationNode aggNode) {
-      List<Page> mergedPages = merger.mergeAggregation(splitPages, aggNode, columnTypes);
+      long t0 = System.nanoTime();
+      int numPartitions = Runtime.getRuntime().availableProcessors();
+      int numGroupByCols = aggNode.getGroupByKeys().size();
+
+      // Collect all partial pages from all splits into flat list
+      List<Page> allPartialPages = merger.mergePassthrough(splitPages);
+
+      // Hash-partitioned exchange: route each row to partition = hash(groupKeys) % N
+      @SuppressWarnings("unchecked")
+      List<Page>[] partitionPages = new List[numPartitions];
+      for (int p = 0; p < numPartitions; p++) partitionPages[p] = new ArrayList<>();
+
+      long perfRouteStart = System.nanoTime();
+      for (Page page : allPartialPages) {
+        int positionCount = page.getPositionCount();
+        int[][] selected = new int[numPartitions][positionCount];
+        int[] selectedCount = new int[numPartitions];
+        for (int pos = 0; pos < positionCount; pos++) {
+          int h = 1;
+          for (int k = 0; k < numGroupByCols; k++) {
+            Block block = page.getBlock(k);
+            if (block.isNull(pos)) { h = 31 * h; }
+            else { h = 31 * h + Long.hashCode(columnTypes.get(k).getLong(block, pos)); }
+          }
+          int partition = Math.floorMod(h, numPartitions);
+          selected[partition][selectedCount[partition]++] = pos;
+        }
+        for (int p = 0; p < numPartitions; p++) {
+          if (selectedCount[p] > 0) {
+            partitionPages[p].add(selectedCount[p] == positionCount
+                ? page : page.copyPositions(selected[p], 0, selectedCount[p]));
+          }
+        }
+      }
+      LOG.info("PERF: hash exchange routing {}ms, {} partitions",
+          (System.nanoTime() - perfRouteStart) / 1_000_000, numPartitions);
+
+      // Parallel FINAL aggregation per partition (disjoint group keys)
+      long perfAggStart = System.nanoTime();
+      java.util.concurrent.ExecutorService pool =
+          java.util.concurrent.Executors.newFixedThreadPool(numPartitions);
+      List<Page> mergedPages;
+      try {
+        List<java.util.concurrent.CompletableFuture<List<Page>>> futures = new ArrayList<>();
+        for (int p = 0; p < numPartitions; p++) {
+          List<Page> pPages = partitionPages[p];
+          futures.add(java.util.concurrent.CompletableFuture.supplyAsync(
+              () -> new ResultMerger().mergeAggregation(List.of(pPages), aggNode, columnTypes), pool));
+        }
+        mergedPages = new ArrayList<>();
+        for (var f : futures) mergedPages.addAll(f.join());
+      } finally {
+        pool.shutdown();
+      }
+      long aggGroups = mergedPages.stream().mapToLong(Page::getPositionCount).sum();
+      LOG.info("PERF: parallel FINAL agg {}ms, {} output groups",
+          (System.nanoTime() - perfAggStart) / 1_000_000, aggGroups);
+      LOG.info("PERF: mergeIcebergResults hash-exchange total {}ms",
+          (System.nanoTime() - t0) / 1_000_000);
       mergedPages = applyCoordinatorHaving(mergedPages, optimizedPlan, aggNode, columnTypeMap);
       return applyCoordinatorSort(mergedPages, aggNode, optimizedPlan, columnTypes, merger);
     } else {
diff --git a/dqe/src/main/java/org/opensearch/sql/dqe/iceberg/transport/TransportIcebergSplitExecuteAction.java b/dqe/src/main/java/org/opensearch/sql/dqe/iceberg/transport/TransportIcebergSplitExecuteAction.java
@@ -49,6 +49,21 @@ public class TransportIcebergSplitExecuteAction
 
   private final org.opensearch.cluster.service.ClusterService clusterService;
 
+  // Cache table metadata across splits — same table is loaded 125+ times per query
+  private record CachedTableMeta(TableInfo tableInfo, Table icebergTable, Schema schema, Map<String, Type> columnTypeMap) {}
+  private final java.util.concurrent.ConcurrentHashMap<String, CachedTableMeta> tableCache = new java.util.concurrent.ConcurrentHashMap<>();
+
+  private CachedTableMeta getTableMeta(String tableName) {
+    return tableCache.computeIfAbsent(tableName, name -> {
+      IcebergTableResolver resolver = new IcebergTableResolver(getWarehousePath());
+      TableInfo info = resolver.resolve(name);
+      Table table = resolver.loadTable(name);
+      Map<String, Type> typeMap = new HashMap<>();
+      for (TableInfo.ColumnInfo col : info.columns()) typeMap.put(col.name(), col.trinoType());
+      return new CachedTableMeta(info, table, table.schema(), typeMap);
+    });
+  }
+
   @Inject
   public TransportIcebergSplitExecuteAction(
       TransportService transportService, ActionFilters actionFilters,
@@ -80,15 +95,9 @@ protected void doExecute(
                   new ByteArrayInputStream(req.getSerializedPlan())));
 
       IcebergSplitInfo splitInfo = req.getSplitInfo();
-      IcebergTableResolver resolver = new IcebergTableResolver(getWarehousePath());
-      TableInfo tableInfo = resolver.resolve(splitInfo.tableName());
-      Table icebergTable = resolver.loadTable(splitInfo.tableName());
-      Schema icebergSchema = icebergTable.schema();
-
-      Map<String, Type> columnTypeMap = new HashMap<>();
-      for (TableInfo.ColumnInfo col : tableInfo.columns()) {
-        columnTypeMap.put(col.name(), col.trinoType());
-      }
+      CachedTableMeta meta = getTableMeta(splitInfo.tableName());
+      Schema icebergSchema = meta.schema();
+      Map<String, Type> columnTypeMap = meta.columnTypeMap();
 
       // Build scan factory that creates ParquetPageSource for the split
       ParquetReaderOptions options = new ParquetReaderOptions();
@@ -148,22 +157,21 @@ protected void doExecute(
    */
   public IcebergSplitExecuteResponse executeLocal(
       DqePlanNode plan, IcebergSplitInfo splitInfo) throws Exception {
-    IcebergTableResolver resolver = new IcebergTableResolver(getWarehousePath());
-    TableInfo tableInfo = resolver.resolve(splitInfo.tableName());
-    Table icebergTable = resolver.loadTable(splitInfo.tableName());
-    Schema icebergSchema = icebergTable.schema();
-
-    Map<String, Type> columnTypeMap = new HashMap<>();
-    for (TableInfo.ColumnInfo col : tableInfo.columns()) {
-      columnTypeMap.put(col.name(), col.trinoType());
-    }
+    long perfSplitStart = System.nanoTime();
+    long perfMetaStart = System.nanoTime();
+    CachedTableMeta meta = getTableMeta(splitInfo.tableName());
+    Schema icebergSchema = meta.schema();
+    Map<String, Type> columnTypeMap = meta.columnTypeMap();
+    long perfMetaMs = (System.nanoTime() - perfMetaStart) / 1_000_000;
 
     ParquetReaderOptions options = new ParquetReaderOptions();
     ParquetPredicateConverter predicateConverter = new ParquetPredicateConverter(columnTypeMap);
+    long[] perfParquetOpenMs = {0};
     LocalExecutionPlanner planner =
         new LocalExecutionPlanner(
             scanNode -> {
               try {
+                long openStart = System.nanoTime();
                 // Read file schema for predicate mapping
                 java.io.File file = new java.io.File(splitInfo.filePath());
                 io.trino.filesystem.local.LocalInputFile inputFile =
@@ -181,26 +189,44 @@ public IcebergSplitExecuteResponse executeLocal(
                 TupleDomain<ColumnDescriptor> predicate =
                     predicateConverter.extractPredicates(plan, fileSchema);
 
-                return new ParquetPageSource(
+                ParquetPageSource source = new ParquetPageSource(
                     splitInfo.filePath(),
                     icebergSchema,
                     scanNode.getColumns(),
                     columnTypeMap,
                     options,
                     predicate);
+                perfParquetOpenMs[0] = (System.nanoTime() - openStart) / 1_000_000;
+                return source;
               } catch (Exception e) {
                 throw new RuntimeException("Failed to create ParquetPageSource", e);
               }
             },
             columnTypeMap);
 
+    long perfPlanStart = System.nanoTime();
     Operator pipeline = plan.accept(planner, null);
+    long perfPlanMs = (System.nanoTime() - perfPlanStart) / 1_000_000;
+
+    long perfExecStart = System.nanoTime();
     List<Page> pages = new ArrayList<>();
     Page page;
+    long totalRows = 0;
     while ((page = pipeline.processNextBatch()) != null) {
       pages.add(page);
+      totalRows += page.getPositionCount();
     }
     pipeline.close();
+    long perfExecMs = (System.nanoTime() - perfExecStart) / 1_000_000;
+    long perfTotalMs = (System.nanoTime() - perfSplitStart) / 1_000_000;
+
+    // Log per-split profiling (sampled: first, last, and every 25th split)
+    String fileName = splitInfo.filePath();
+    int lastSlash = fileName.lastIndexOf('/');
+    String shortName = lastSlash >= 0 ? fileName.substring(lastSlash + 1) : fileName;
+    LOG.info("PERF-SPLIT: {} total={}ms meta={}ms parquetOpen={}ms plan={}ms exec={}ms pages={} rows={}",
+        shortName, perfTotalMs, perfMetaMs, perfParquetOpenMs[0], perfPlanMs, perfExecMs,
+        pages.size(), totalRows);
 
     List<Type> columnTypes = resolveColumnTypes(plan, columnTypeMap);
     return new IcebergSplitExecuteResponse(pages, columnTypes);
@@ -214,15 +240,9 @@ public IcebergSplitExecuteResponse executeLocalWithBucketFilter(
       DqePlanNode plan, IcebergSplitInfo splitInfo,
       List<Integer> groupByIndices, List<Type> allColumnTypes,
       int bucket, int numBuckets) throws Exception {
-    IcebergTableResolver resolver = new IcebergTableResolver(getWarehousePath());
-    TableInfo tableInfo = resolver.resolve(splitInfo.tableName());
-    Table icebergTable = resolver.loadTable(splitInfo.tableName());
-    Schema icebergSchema = icebergTable.schema();
-
-    Map<String, Type> columnTypeMap = new HashMap<>();
-    for (TableInfo.ColumnInfo col : tableInfo.columns()) {
-      columnTypeMap.put(col.name(), col.trinoType());
-    }
+    CachedTableMeta meta = getTableMeta(splitInfo.tableName());
+    Schema icebergSchema = meta.schema();
+    Map<String, Type> columnTypeMap = meta.columnTypeMap();
 
     ParquetReaderOptions options = new ParquetReaderOptions();
     ParquetPredicateConverter predicateConverter = new ParquetPredicateConverter(columnTypeMap);
diff --git a/dqe/src/main/java/org/opensearch/sql/dqe/operator/HashAggregationOperator.java b/dqe/src/main/java/org/opensearch/sql/dqe/operator/HashAggregationOperator.java
