Integrate aggregators to convert result from datafusion (#19441)

Signed-off-by: expani <[email protected]>

Author: expani

server/src/main/java/org/opensearch/search/aggregations/ShardResultConvertor.java

@@ -0,0 +1,18 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.aggregations;
+
+import java.util.List;
+import java.util.Map;
+
+public interface ShardResultConvertor {
+
+    List<InternalAggregation> convert(Map<String, Object[]> shardResult);
+
+}

server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java

@@ -51,6 +51,7 @@
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.LeafBucketCollector;
 import org.opensearch.search.aggregations.LeafBucketCollectorBase;
+import org.opensearch.search.aggregations.ShardResultConvertor;
 import org.opensearch.search.aggregations.StarTreeBucketCollector;
 import org.opensearch.search.aggregations.StarTreePreComputeCollector;
 import org.opensearch.search.aggregations.support.ValuesSource;
@@ -59,7 +60,9 @@
 import org.opensearch.search.startree.StarTreeQueryHelper;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
@@ -71,7 +74,7 @@
  *
  * @opensearch.internal
  */
-class MaxAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector {
+class MaxAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector, ShardResultConvertor {
 
     final ValuesSource.Numeric valuesSource;
     final DocValueFormat formatter;
@@ -280,4 +283,14 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
     public void doReset() {
         maxes.fill(0, maxes.size(), Double.NEGATIVE_INFINITY);
     }
+
+    @Override
+    public List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
+            Object[] values = shardResult.get(name);
+            List<InternalAggregation> results = new ArrayList<>(values.length);
+            for (Object value : values) {
+                results.add(new InternalMax(name, (Long) value, formatter, metadata()));
+            }
+            return results;
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java

@@ -51,6 +51,7 @@
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.LeafBucketCollector;
 import org.opensearch.search.aggregations.LeafBucketCollectorBase;
+import org.opensearch.search.aggregations.ShardResultConvertor;
 import org.opensearch.search.aggregations.StarTreeBucketCollector;
 import org.opensearch.search.aggregations.StarTreePreComputeCollector;
 import org.opensearch.search.aggregations.support.ValuesSource;
@@ -59,6 +60,8 @@
 import org.opensearch.search.startree.StarTreeQueryHelper;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
@@ -70,7 +73,7 @@
  *
  * @opensearch.internal
  */
-class MinAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector {
+class MinAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector, ShardResultConvertor {
     private static final int MAX_BKD_LOOKUPS = 1024;
 
     final ValuesSource.Numeric valuesSource;
@@ -271,4 +274,14 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
             (bucket, metricValue) -> mins.set(bucket, Math.min(mins.get(bucket), NumericUtils.sortableLongToDouble(metricValue)))
         );
     }
+
+    @Override
+    public List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
+            Object[] values = shardResult.get(name);
+            List<InternalAggregation> results = new ArrayList<>(values.length);
+            for (Object value : values) {
+                results.add(new InternalMin(name, (Long) value, format, metadata()));
+            }
+            return results;
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java

@@ -45,6 +45,7 @@
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.LeafBucketCollector;
 import org.opensearch.search.aggregations.LeafBucketCollectorBase;
+import org.opensearch.search.aggregations.ShardResultConvertor;
 import org.opensearch.search.aggregations.StarTreeBucketCollector;
 import org.opensearch.search.aggregations.StarTreePreComputeCollector;
 import org.opensearch.search.aggregations.support.ValuesSource;
@@ -53,6 +54,8 @@
 import org.opensearch.search.startree.StarTreeQueryHelper;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 
 import static org.opensearch.search.startree.StarTreeQueryHelper.getSupportedStarTree;
@@ -62,7 +65,7 @@
  *
  * @opensearch.internal
  */
-public class SumAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector {
+public class SumAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector, ShardResultConvertor {
 
     private final ValuesSource.Numeric valuesSource;
     private final DocValueFormat format;
@@ -215,4 +218,14 @@ public InternalAggregation buildEmptyAggregation() {
     public void doClose() {
         Releasables.close(sums, compensations);
     }
+
+    @Override
+    public List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
+            Object[] values = shardResult.get(name);
+            List<InternalAggregation> results = new ArrayList<>(values.length);
+            for (Object value : values) {
+                results.add(new InternalSum(name, (Long) value, format, metadata()));
+            }
+            return results;
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java

@@ -45,6 +45,7 @@
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.LeafBucketCollector;
 import org.opensearch.search.aggregations.LeafBucketCollectorBase;
+import org.opensearch.search.aggregations.ShardResultConvertor;
 import org.opensearch.search.aggregations.StarTreeBucketCollector;
 import org.opensearch.search.aggregations.StarTreePreComputeCollector;
 import org.opensearch.search.aggregations.support.ValuesSource;
@@ -53,6 +54,8 @@
 import org.opensearch.search.startree.StarTreeQueryHelper;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 
 import static org.opensearch.search.startree.StarTreeQueryHelper.getSupportedStarTree;
@@ -65,7 +68,7 @@
  *
  * @opensearch.internal
  */
-public class ValueCountAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector {
+public class ValueCountAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector, ShardResultConvertor {
 
     final ValuesSource valuesSource;
 
@@ -209,4 +212,14 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
             (bucket, metricValue) -> counts.increment(bucket, metricValue)
         );
     }
+
+    @Override
+    public List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
+            Object[] values = shardResult.get(name);
+            List<InternalAggregation> results = new ArrayList<>(values.length);
+            for (Object value : values) {
+                results.add(new InternalValueCount(name, (Long) value, metadata()));
+            }
+            return results;
+    }
 }

server/src/main/java/org/opensearch/search/query/QueryPhase.java

@@ -60,9 +60,7 @@
 import org.opensearch.search.aggregations.AggregationProcessor;
 import org.opensearch.search.aggregations.DefaultAggregationProcessor;
 import org.opensearch.search.aggregations.GlobalAggCollectorManager;
-import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.InternalAggregations;
-import org.opensearch.search.aggregations.metrics.InternalValueCount;
 import org.opensearch.search.internal.ContextIndexSearcher;
 import org.opensearch.search.internal.ScrollContext;
 import org.opensearch.search.internal.SearchContext;
@@ -75,7 +73,6 @@
 import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -168,28 +165,16 @@ public void execute(SearchContext searchContext) throws QueryPhaseExecutionExcep
 
         // boolean rescore = executeInternal(searchContext, queryPhaseSearcher);
 
+        // Post process
+        final InternalAggregations internalAggregations = SearchEngineResultConversionUtils.convertDFResultGeneric(searchContext);
+        LOGGER.info("InternalAggregation created is {}", internalAggregations.asList());
+        searchContext.queryResult().aggregations(internalAggregations);
+
         // if (rescore) { // only if we do a regular search
         // rescoreProcessor.process(searchContext);
         // }
         // suggestProcessor.process(searchContext);
-        // aggregationProcessor.postProcess(searchContext);
-
-        // Post process
-        // Create a list to store the InternalValueCount objects
-        // Can we map from the preprocess
-        List<InternalAggregation> internalAggList = new ArrayList<>();
-        Map<String, Object[]> map = searchContext.getDFResults();
-        for (Map.Entry<String, Object[]> entry : map.entrySet()) {
-            String key = entry.getKey();
-            Object[] value = entry.getValue();
-            // SUM, Count will work with integer casting, but (Integer) value casting may not work well for avg
-            InternalValueCount ivc = new InternalValueCount(key, (long) value[0], null);
-            internalAggList.add(ivc);
-        }
-
-        final InternalAggregations internalAggregations = InternalAggregations.from(internalAggList);
-        QuerySearchResult querySearchResult = searchContext.queryResult();
-        querySearchResult.aggregations(internalAggregations);
+         aggregationProcessor.postProcess(searchContext);
 
         if (searchContext.getProfilers() != null) {
             ProfileShardResult shardResults = SearchProfileShardResults.buildShardResults(

server/src/main/java/org/opensearch/search/query/SearchEngineResultConversionUtils.java

@@ -0,0 +1,56 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.query;
+
+import org.opensearch.search.aggregations.Aggregator;
+import org.opensearch.search.aggregations.InternalAggregations;
+import org.opensearch.search.aggregations.ShardResultConvertor;
+import org.opensearch.search.internal.SearchContext;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class SearchEngineResultConversionUtils {
+
+    public static InternalAggregations convertDFResultGeneric(SearchContext searchContext) {
+        Map<String, Object[]> dfResult = searchContext.getDFResults();
+
+        // Create aggregators which will process the result from DataFusion
+        try {
+
+            List<Aggregator> aggregators = new ArrayList<>();
+
+            if (searchContext.aggregations().factories().hasGlobalAggregator()) {
+                aggregators.addAll(searchContext.aggregations().factories().createTopLevelGlobalAggregators(searchContext));
+            }
+
+            if (searchContext.aggregations().factories().hasNonGlobalAggregator()) {
+                aggregators.addAll(searchContext.aggregations().factories().createTopLevelNonGlobalAggregators(searchContext));
+            }
+
+            List<ShardResultConvertor> shardResultConvertors = aggregators.stream().map(x -> {
+                if (x instanceof ShardResultConvertor) {
+                    return ((ShardResultConvertor) x);
+                } else {
+                    throw new UnsupportedOperationException("Aggregator doesn't support converting results from shard: " + x);
+                }
+            }).toList();
+
+            return InternalAggregations.from(
+                shardResultConvertors.stream().flatMap(x -> x.convert(dfResult).stream()).collect(Collectors.toList())
+            );
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+}