From 1f55bf1f62ea065a782f2878b2ac7364b137caaf Mon Sep 17 00:00:00 2001
From: Fabio Massimo Ercoli <fabiomassimo.ercoli@gmail.com>
Date: Wed, 28 Aug 2024 14:57:13 +0200
Subject: [PATCH] HSEARCH-5133 Push type-sel into each Lucene*Aggregation

---
 ...LuceneMetricCompensatedSumAggregation.java | 26 +++++++------------
 .../LuceneAvgCompensatedSumAggregation.java   | 17 +++++++++++-
 .../LuceneSumCompensatedSumAggregation.java   | 17 +++++++++++-
 3 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/AbstractLuceneMetricCompensatedSumAggregation.java b/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/AbstractLuceneMetricCompensatedSumAggregation.java
index feb1d25c2d9..099d81ea70d 100644
--- a/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/AbstractLuceneMetricCompensatedSumAggregation.java
+++ b/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/AbstractLuceneMetricCompensatedSumAggregation.java
@@ -76,19 +76,16 @@ public K extract(AggregationExtractContext context) {
 
 	abstract E extractEncoded(AggregationExtractContext context, LuceneNumericDomain<E> numericDomain);
 
-	protected static class TypeSelector<F> implements FieldMetricAggregationBuilder.TypeSelector {
-		private final AbstractLuceneNumericFieldCodec<F, ?> codec;
-		private final LuceneSearchIndexScope<?> scope;
-		private final LuceneSearchIndexValueFieldContext<F> field;
-		private final String operation;
+	protected abstract static class TypeSelector<F> implements FieldMetricAggregationBuilder.TypeSelector {
+		protected final AbstractLuceneNumericFieldCodec<F, ?> codec;
+		protected final LuceneSearchIndexScope<?> scope;
+		protected final LuceneSearchIndexValueFieldContext<F> field;
 
 		protected TypeSelector(AbstractLuceneNumericFieldCodec<F, ?> codec,
-				LuceneSearchIndexScope<?> scope, LuceneSearchIndexValueFieldContext<F> field,
-				String operation) {
+				LuceneSearchIndexScope<?> scope, LuceneSearchIndexValueFieldContext<F> field) {
 			this.codec = codec;
 			this.scope = scope;
 			this.field = field;
-			this.operation = operation;
 		}
 
 		@Override
@@ -101,16 +98,11 @@ protected TypeSelector(AbstractLuceneNumericFieldCodec<F, ?> codec,
 						.withConvertedType( expectedType, field );
 			}
 
-			if ( "sum".equals( operation ) ) {
-				return new LuceneSumCompensatedSumAggregation.Builder<>( codec, scope, field, projectionConverter );
-			}
-			else if ( "avg".equals( operation ) ) {
-				return new LuceneAvgCompensatedSumAggregation.Builder<>( codec, scope, field, projectionConverter );
-			}
-			else {
-				throw new AssertionFailure( "Aggregation operation not supported: " + operation );
-			}
+			return getFtBuilder( projectionConverter );
 		}
+
+		protected abstract <T> Builder<F, ? extends Number, T> getFtBuilder(
+				ProjectionConverter<F, ? extends T> projectionConverter);
 	}
 
 	protected abstract static class Builder<F, E extends Number, K> extends AbstractBuilder<K>
diff --git a/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneAvgCompensatedSumAggregation.java b/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneAvgCompensatedSumAggregation.java
index 38c24be6b8f..c8b677752c3 100644
--- a/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneAvgCompensatedSumAggregation.java
+++ b/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneAvgCompensatedSumAggregation.java
@@ -61,7 +61,22 @@ protected Factory(AbstractLuceneNumericFieldCodec<F, ?> codec) {
 		@Override
 		public FieldMetricAggregationBuilder.TypeSelector create(LuceneSearchIndexScope<?> scope,
 				LuceneSearchIndexValueFieldContext<F> field) {
-			return new TypeSelector<>( codec, scope, field, "avg" );
+			return new FunctionTypeSelector<>( codec, scope, field );
+		}
+	}
+
+	protected static class FunctionTypeSelector<F> extends TypeSelector<F>
+			implements FieldMetricAggregationBuilder.TypeSelector {
+
+		protected FunctionTypeSelector(AbstractLuceneNumericFieldCodec<F, ?> codec, LuceneSearchIndexScope<?> scope,
+				LuceneSearchIndexValueFieldContext<F> field) {
+			super( codec, scope, field );
+		}
+
+		@Override
+		protected <T> Builder<F, ? extends Number, T> getFtBuilder(
+				ProjectionConverter<F, ? extends T> projectionConverter) {
+			return new Builder<>( codec, scope, field, projectionConverter );
 		}
 	}
 
diff --git a/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneSumCompensatedSumAggregation.java b/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneSumCompensatedSumAggregation.java
index b5c49605389..0397eac198a 100644
--- a/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneSumCompensatedSumAggregation.java
+++ b/backend/lucene/src/main/java/org/hibernate/search/backend/lucene/types/aggregation/impl/LuceneSumCompensatedSumAggregation.java
@@ -54,7 +54,22 @@ protected Factory(AbstractLuceneNumericFieldCodec<F, ?> codec) {
 		@Override
 		public FieldMetricAggregationBuilder.TypeSelector create(LuceneSearchIndexScope<?> scope,
 				LuceneSearchIndexValueFieldContext<F> field) {
-			return new TypeSelector<>( codec, scope, field, "sum" );
+			return new FunctionTypeSelector<>( codec, scope, field );
+		}
+	}
+
+	protected static class FunctionTypeSelector<F> extends TypeSelector<F>
+			implements FieldMetricAggregationBuilder.TypeSelector {
+
+		protected FunctionTypeSelector(AbstractLuceneNumericFieldCodec<F, ?> codec, LuceneSearchIndexScope<?> scope,
+				LuceneSearchIndexValueFieldContext<F> field) {
+			super( codec, scope, field );
+		}
+
+		@Override
+		protected <T> Builder<F, ? extends Number, T> getFtBuilder(
+				ProjectionConverter<F, ? extends T> projectionConverter) {
+			return new Builder<>( codec, scope, field, projectionConverter );
 		}
 	}