HSEARCH-5133 Refactor extracting AggregationFunctionCollector

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/AggregationFunction.java

@@ -0,0 +1,17 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.search.backend.lucene.lowlevel.aggregation.collector.impl;
+
+public interface AggregationFunction<R extends AggregationFunction<?>> {
+
+	void apply(long value);
+
+	void merge(AggregationFunction<R> sibling);
+
+	long result();
+
+	R implementation();
+
+}

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/CountDistinctCollector.java → backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/AggregationFunctionCollector.java

@@ -15,37 +15,38 @@
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
 
-public class CountDistinctCollector implements Collector {
+public class AggregationFunctionCollector<R extends AggregationFunction<?>> implements Collector {
 
 	private final LongMultiValuesSource valueSource;
-	private final CountDistinct counter = new CountDistinct();
+	private final AggregationFunction<R> aggregationFunction;
 
-	public CountDistinctCollector(LongMultiValuesSource valueSource) {
+	public AggregationFunctionCollector(LongMultiValuesSource valueSource, AggregationFunction<R> aggregationFunction) {
 		this.valueSource = valueSource;
+		this.aggregationFunction = aggregationFunction;
 	}
 
-	public long count() {
-		return counter.result();
+	public void merge(AggregationFunctionCollector<R> sibling) {
+		aggregationFunction.merge( sibling.aggregationFunction );
 	}
 
-	public void merge(CountDistinctCollector collector) {
-		counter.merge( collector.counter );
+	public long result() {
+		return aggregationFunction.result();
 	}
 
 	@Override
 	public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
-		return new CountDistinctLeafCollector( valueSource.getValues( context ) );
+		return new AggregationFunctionLeafCollector( valueSource.getValues( context ) );
 	}
 
 	@Override
 	public ScoreMode scoreMode() {
 		return ScoreMode.COMPLETE_NO_SCORES;
 	}
 
-	public class CountDistinctLeafCollector implements LeafCollector {
+	public class AggregationFunctionLeafCollector implements LeafCollector {
 		private final LongMultiValues values;
 
-		public CountDistinctLeafCollector(LongMultiValues values) {
+		public AggregationFunctionLeafCollector(LongMultiValues values) {
 			this.values = values;
 		}
 
@@ -54,7 +55,7 @@ public void collect(int doc) throws IOException {
 			if ( values.advanceExact( doc ) ) {
 				while ( values.hasNextValue() ) {
 					long value = values.nextValue();
-					counter.increment( value );
+					aggregationFunction.apply( value );
 				}
 			}
 		}

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/AggregationFunctionCollectorManager.java

@@ -0,0 +1,46 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.search.backend.lucene.lowlevel.aggregation.collector.impl;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.LinkedList;
+import java.util.function.Supplier;
+
+import org.hibernate.search.backend.lucene.lowlevel.docvalues.impl.JoiningLongMultiValuesSource;
+
+import org.apache.lucene.search.CollectorManager;
+
+public class AggregationFunctionCollectorManager<T extends AggregationFunction<?>>
+		implements CollectorManager<AggregationFunctionCollector<T>, Long> {
+
+	private final JoiningLongMultiValuesSource source;
+	private final Supplier<AggregationFunction<T>> functionSupplier;
+
+	public AggregationFunctionCollectorManager(JoiningLongMultiValuesSource source,
+			Supplier<AggregationFunction<T>> functionSupplier) {
+		this.source = source;
+		this.functionSupplier = functionSupplier;
+	}
+
+	@Override
+	public AggregationFunctionCollector<T> newCollector() {
+		return new AggregationFunctionCollector<>( source, functionSupplier.get() );
+	}
+
+	@Override
+	public Long reduce(Collection<AggregationFunctionCollector<T>> collectors) throws IOException {
+		if ( collectors.isEmpty() ) {
+			return null;
+		}
+
+		LinkedList<AggregationFunctionCollector<T>> distinctCollectors = new LinkedList<>( collectors );
+		AggregationFunctionCollector<T> collector = distinctCollectors.removeLast();
+		for ( AggregationFunctionCollector<T> other : distinctCollectors ) {
+			collector.merge( other );
+		}
+		return collector.result();
+	}
+}

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/CountDistinct.java

@@ -12,12 +12,13 @@
  * The algorithm to collect distinct elements is inspired by {@code org.apache.lucene.facet.LongValueFacetCounts}
  * of <a href="https://lucene.apache.org/">Apache Lucene project</a>.
  */
-public class CountDistinct {
+public class CountDistinct implements AggregationFunction<CountDistinct> {
 
 	private final BitSet counts = new BitSet( 1024 );
 	private final HashSet<Long> hashCounts = new HashSet<>();
 
-	public void increment(long value) {
+	@Override
+	public void apply(long value) {
 		if ( value >= 0 && value < counts.size() ) {
 			counts.set( (int) value );
 		}
@@ -26,12 +27,20 @@ public void increment(long value) {
 		}
 	}
 
+	@Override
+	public void merge(AggregationFunction<CountDistinct> sibling) {
+		CountDistinct other = sibling.implementation();
+		counts.or( other.counts );
+		hashCounts.addAll( other.hashCounts );
+	}
+
+	@Override
 	public long result() {
 		return counts.cardinality() + hashCounts.size();
 	}
 
-	public void merge(CountDistinct other) {
-		counts.or( other.counts );
-		hashCounts.addAll( other.hashCounts );
+	@Override
+	public CountDistinct implementation() {
+		return this;
 	}
 }

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/CountDistinctCollectorFactory.java

@@ -10,23 +10,25 @@
 import org.hibernate.search.backend.lucene.lowlevel.docvalues.impl.JoiningLongMultiValuesSource;
 
 public class CountDistinctCollectorFactory
-		implements CollectorFactory<CountDistinctCollector, Long, CountDistinctCollectorManager> {
+		implements
+		CollectorFactory<AggregationFunctionCollector<CountDistinct>,
+				Long,
+				AggregationFunctionCollectorManager<CountDistinct>> {
 
 	private final JoiningLongMultiValuesSource source;
-	private final CollectorKey<CountDistinctCollector, Long> key;
+	private final CollectorKey<AggregationFunctionCollector<CountDistinct>, Long> key = CollectorKey.create();
 
-	public CountDistinctCollectorFactory(JoiningLongMultiValuesSource source, CollectorKey<CountDistinctCollector, Long> key) {
+	public CountDistinctCollectorFactory(JoiningLongMultiValuesSource source) {
 		this.source = source;
-		this.key = key;
 	}
 
 	@Override
-	public CountDistinctCollectorManager createCollectorManager(CollectorExecutionContext context) {
-		return new CountDistinctCollectorManager( source );
+	public AggregationFunctionCollectorManager<CountDistinct> createCollectorManager(CollectorExecutionContext context) {
+		return new AggregationFunctionCollectorManager<>( source, CountDistinct::new );
 	}
 
 	@Override
-	public CollectorKey<CountDistinctCollector, Long> getCollectorKey() {
+	public CollectorKey<AggregationFunctionCollector<CountDistinct>, Long> getCollectorKey() {
 		return key;
 	}
 }

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/Sum.java

@@ -0,0 +1,30 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.search.backend.lucene.lowlevel.aggregation.collector.impl;
+
+public class Sum implements AggregationFunction<Sum> {
+
+	private long sum = 0L;
+
+	@Override
+	public void apply(long value) {
+		sum += value;
+	}
+
+	@Override
+	public void merge(AggregationFunction<Sum> sibling) {
+		apply( sibling.result() );
+	}
+
+	@Override
+	public long result() {
+		return sum;
+	}
+
+	@Override
+	public Sum implementation() {
+		return this;
+	}
+}

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/SumCollectorFactory.java

@@ -9,23 +9,23 @@
 import org.hibernate.search.backend.lucene.lowlevel.collector.impl.CollectorKey;
 import org.hibernate.search.backend.lucene.lowlevel.docvalues.impl.JoiningLongMultiValuesSource;
 
-public class SumCollectorFactory implements CollectorFactory<SumCollector, Long, SumCollectorManager> {
+public class SumCollectorFactory
+		implements CollectorFactory<AggregationFunctionCollector<Sum>, Long, AggregationFunctionCollectorManager<Sum>> {
 
 	private final JoiningLongMultiValuesSource source;
-	private final CollectorKey<SumCollector, Long> key;
+	private final CollectorKey<AggregationFunctionCollector<Sum>, Long> key = CollectorKey.create();
 
-	public SumCollectorFactory(JoiningLongMultiValuesSource source, CollectorKey<SumCollector, Long> key) {
+	public SumCollectorFactory(JoiningLongMultiValuesSource source) {
 		this.source = source;
-		this.key = key;
 	}
 
 	@Override
-	public SumCollectorManager createCollectorManager(CollectorExecutionContext context) {
-		return new SumCollectorManager( source );
+	public AggregationFunctionCollectorManager<Sum> createCollectorManager(CollectorExecutionContext context) {
+		return new AggregationFunctionCollectorManager<>( source, Sum::new );
 	}
 
 	@Override
-	public CollectorKey<SumCollector, Long> getCollectorKey() {
+	public CollectorKey<AggregationFunctionCollector<Sum>, Long> getCollectorKey() {
 		return key;
 	}
 }