HSEARCH-5133 Push type-sel into each Lucene*Aggregation

fax4ever · Aug 28, 2024 · 1f55bf1 · 1f55bf1
1 parent 3830727
commit 1f55bf1
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 19 deletions.
diff --git a/.../backend/lucene/types/aggregation/impl/AbstractLuceneMetricCompensatedSumAggregation.java b/.../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/...nate/search/backend/lucene/types/aggregation/impl/LuceneAvgCompensatedSumAggregation.java b/...nate/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/...nate/search/backend/lucene/types/aggregation/impl/LuceneSumCompensatedSumAggregation.java b/...nate/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 );
 		}
 	}