Refactor code in aggregation processor to make it simpler

Signed-off-by: Martin Gaievski <[email protected]>
martin-gaievski · Feb 20, 2024 · ec3f7dc · ec3f7dc
1 parent 02e0b09
commit ec3f7dc
Showing 1 changed file with 29 additions and 35 deletions.
diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridAggregationProcessor.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridAggregationProcessor.java
@@ -29,53 +29,47 @@ public class HybridAggregationProcessor implements AggregationProcessor {
 
     @Override
     public void preProcess(SearchContext context) {
-        if (shouldUseMultiCollectorManager(context)) {
-            if (context.shouldUseConcurrentSearch()) {
-                concurrentAggregationProcessor.preProcess(context);
-            } else {
-                defaultAggregationProcessor.preProcess(context);
-            }
-            if (HybridQueryUtil.isHybridQuery(context.query(), context)) {
-                HybridCollectorManager collectorManager;
-                try {
-                    collectorManager = HybridCollectorManager.createHybridCollectorManager(context);
-                } catch (IOException e) {
-                    throw new RuntimeException(e);
-                }
-                Map<Class<?>, CollectorManager<? extends Collector, ReduceableSearchResult>> collectorManagersByManagerClass = context
-                    .queryCollectorManagers();
-                collectorManagersByManagerClass.put(HybridCollectorManager.class, collectorManager);
-            }
+        if (context.shouldUseConcurrentSearch()) {
+            concurrentAggregationProcessor.preProcess(context);
         } else {
             defaultAggregationProcessor.preProcess(context);
         }
+
+        if (HybridQueryUtil.isHybridQuery(context.query(), context)) {
+            HybridCollectorManager collectorManager;
+            try {
+                collectorManager = HybridCollectorManager.createHybridCollectorManager(context);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            Map<Class<?>, CollectorManager<? extends Collector, ReduceableSearchResult>> collectorManagersByManagerClass = context
+                .queryCollectorManagers();
+            collectorManagersByManagerClass.put(HybridCollectorManager.class, collectorManager);
+        }
     }
 
     @Override
     public void postProcess(SearchContext context) {
-        if (shouldUseMultiCollectorManager(context)) {
-            if (HybridQueryUtil.isHybridQuery(context.query(), context)) {
+        if (HybridQueryUtil.isHybridQuery(context.query(), context)) {
+            if (!context.shouldUseConcurrentSearch()) {
                 CollectorManager<?, ReduceableSearchResult> collectorManager = context.queryCollectorManagers()
                     .get(HybridCollectorManager.class);
-                if (!context.shouldUseConcurrentSearch()) {
-                    try {
-                        final Collection collectors = List.of(collectorManager.newCollector());
-                        collectorManager.reduce(collectors).reduce(context.queryResult());
-                    } catch (IOException e) {
-                        throw new QueryPhaseExecutionException(
-                            context.shardTarget(),
-                            "failed to execute hybrid query aggregation processor",
-                            e
-                        );
-                    }
+                try {
+                    final Collection collectors = List.of(collectorManager.newCollector());
+                    collectorManager.reduce(collectors).reduce(context.queryResult());
+                } catch (IOException e) {
+                    throw new QueryPhaseExecutionException(
+                        context.shardTarget(),
+                        "failed to execute hybrid query aggregation processor",
+                        e
+                    );
                 }
             }
             updateQueryResult(context.queryResult(), context);
-            if (context.shouldUseConcurrentSearch()) {
-                concurrentAggregationProcessor.postProcess(context);
-            } else {
-                defaultAggregationProcessor.postProcess(context);
-            }
+        }
+
+        if (context.shouldUseConcurrentSearch()) {
+            concurrentAggregationProcessor.postProcess(context);
         } else {
             defaultAggregationProcessor.postProcess(context);
         }