diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridAggregationProcessor.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridAggregationProcessor.java index 20870bf26..f68579422 100644 --- 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, CollectorManager> 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, CollectorManager> 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 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); }