-
Notifications
You must be signed in to change notification settings - Fork 244
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
HSEARCH-5133 Implement Lucene sum aggregations
- Loading branch information
Showing
10 changed files
with
428 additions
and
3 deletions.
There are no files selected for viewing
63 changes: 63 additions & 0 deletions
63
...org/hibernate/search/backend/lucene/lowlevel/aggregation/collector/impl/SumCollector.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,63 @@ | ||
/* | ||
* 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 org.hibernate.search.backend.lucene.lowlevel.docvalues.impl.LongMultiValues; | ||
import org.hibernate.search.backend.lucene.lowlevel.docvalues.impl.LongMultiValuesSource; | ||
|
||
import org.apache.lucene.index.LeafReaderContext; | ||
import org.apache.lucene.search.Collector; | ||
import org.apache.lucene.search.LeafCollector; | ||
import org.apache.lucene.search.Scorable; | ||
import org.apache.lucene.search.ScoreMode; | ||
|
||
public class SumCollector implements Collector { | ||
|
||
private final LongMultiValuesSource valueSource; | ||
private long sum = 0L; | ||
|
||
public SumCollector(LongMultiValuesSource valueSource) { | ||
this.valueSource = valueSource; | ||
} | ||
|
||
public long sum() { | ||
return sum; | ||
} | ||
|
||
@Override | ||
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { | ||
return new SumLeafCollector( valueSource.getValues( context ) ); | ||
} | ||
|
||
@Override | ||
public ScoreMode scoreMode() { | ||
return ScoreMode.COMPLETE_NO_SCORES; | ||
} | ||
|
||
public class SumLeafCollector implements LeafCollector { | ||
private final LongMultiValues values; | ||
|
||
public SumLeafCollector(LongMultiValues values) { | ||
this.values = values; | ||
} | ||
|
||
@Override | ||
public void collect(int doc) throws IOException { | ||
if ( values.advanceExact( doc ) ) { | ||
while ( values.hasNextValue() ) { | ||
long value = values.nextValue(); | ||
sum += value; | ||
} | ||
} | ||
} | ||
|
||
@Override | ||
public void setScorer(Scorable scorer) { | ||
// no-op by default | ||
} | ||
} | ||
} |
31 changes: 31 additions & 0 deletions
31
...ernate/search/backend/lucene/lowlevel/aggregation/collector/impl/SumCollectorFactory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
/* | ||
* SPDX-License-Identifier: Apache-2.0 | ||
* Copyright Red Hat Inc. and Hibernate Authors | ||
*/ | ||
package org.hibernate.search.backend.lucene.lowlevel.aggregation.collector.impl; | ||
|
||
import org.hibernate.search.backend.lucene.lowlevel.collector.impl.CollectorExecutionContext; | ||
import org.hibernate.search.backend.lucene.lowlevel.collector.impl.CollectorFactory; | ||
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> { | ||
|
||
private final JoiningLongMultiValuesSource source; | ||
private final CollectorKey<SumCollector, Long> key; | ||
|
||
public SumCollectorFactory(JoiningLongMultiValuesSource source, CollectorKey<SumCollector, Long> key) { | ||
this.source = source; | ||
this.key = key; | ||
} | ||
|
||
@Override | ||
public SumCollectorManager createCollectorManager(CollectorExecutionContext context) { | ||
return new SumCollectorManager( source ); | ||
} | ||
|
||
@Override | ||
public CollectorKey<SumCollector, Long> getCollectorKey() { | ||
return key; | ||
} | ||
} |
35 changes: 35 additions & 0 deletions
35
...ernate/search/backend/lucene/lowlevel/aggregation/collector/impl/SumCollectorManager.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,35 @@ | ||
/* | ||
* 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 org.hibernate.search.backend.lucene.lowlevel.docvalues.impl.JoiningLongMultiValuesSource; | ||
|
||
import org.apache.lucene.search.CollectorManager; | ||
|
||
public class SumCollectorManager implements CollectorManager<SumCollector, Long> { | ||
|
||
private final JoiningLongMultiValuesSource source; | ||
|
||
public SumCollectorManager(JoiningLongMultiValuesSource source) { | ||
this.source = source; | ||
} | ||
|
||
@Override | ||
public SumCollector newCollector() { | ||
return new SumCollector( source ); | ||
} | ||
|
||
@Override | ||
public Long reduce(Collection<SumCollector> collectors) throws IOException { | ||
long result = 0L; | ||
for ( SumCollector sumCollector : collectors ) { | ||
result += sumCollector.sum(); | ||
} | ||
return result; | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
142 changes: 142 additions & 0 deletions
142
...ate/search/backend/lucene/types/aggregation/impl/LuceneMetricNumericFieldAggregation.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,142 @@ | ||
/* | ||
* SPDX-License-Identifier: Apache-2.0 | ||
* Copyright Red Hat Inc. and Hibernate Authors | ||
*/ | ||
package org.hibernate.search.backend.lucene.types.aggregation.impl; | ||
|
||
import java.util.Set; | ||
|
||
import org.hibernate.search.backend.lucene.lowlevel.aggregation.collector.impl.SumCollector; | ||
import org.hibernate.search.backend.lucene.lowlevel.aggregation.collector.impl.SumCollectorFactory; | ||
import org.hibernate.search.backend.lucene.lowlevel.collector.impl.CollectorKey; | ||
import org.hibernate.search.backend.lucene.lowlevel.docvalues.impl.JoiningLongMultiValuesSource; | ||
import org.hibernate.search.backend.lucene.search.aggregation.impl.AggregationExtractContext; | ||
import org.hibernate.search.backend.lucene.search.aggregation.impl.AggregationRequestContext; | ||
import org.hibernate.search.backend.lucene.search.common.impl.AbstractLuceneCodecAwareSearchQueryElementFactory; | ||
import org.hibernate.search.backend.lucene.search.common.impl.LuceneSearchIndexScope; | ||
import org.hibernate.search.backend.lucene.search.common.impl.LuceneSearchIndexValueFieldContext; | ||
import org.hibernate.search.backend.lucene.types.codec.impl.AbstractLuceneNumericFieldCodec; | ||
import org.hibernate.search.backend.lucene.types.lowlevel.impl.LuceneNumericDomain; | ||
import org.hibernate.search.engine.backend.types.converter.spi.ProjectionConverter; | ||
import org.hibernate.search.engine.search.aggregation.spi.FieldMetricAggregationBuilder; | ||
import org.hibernate.search.engine.search.common.ValueConvert; | ||
|
||
/** | ||
* @param <F> The type of field values. | ||
* @param <K> The type of returned value. It can be {@code F}, {@link Double} | ||
* or a different type if value converters are used. | ||
*/ | ||
public class LuceneMetricNumericFieldAggregation<F, E extends Number, K> extends AbstractLuceneNestableAggregation<K> { | ||
|
||
private final Set<String> indexNames; | ||
private final String absoluteFieldPath; | ||
private final AbstractLuceneNumericFieldCodec<F, E> codec; | ||
private final LuceneNumericDomain<E> numericDomain; | ||
private final ProjectionConverter<F, ? extends K> fromFieldValueConverter; | ||
private final String operation; | ||
private final CollectorKey<SumCollector, Long> collectorKey; | ||
|
||
LuceneMetricNumericFieldAggregation(Builder<F, E, K> builder) { | ||
super( builder ); | ||
this.indexNames = builder.scope.hibernateSearchIndexNames(); | ||
this.absoluteFieldPath = builder.field.absolutePath(); | ||
this.codec = builder.codec; | ||
this.numericDomain = codec.getDomain(); | ||
this.fromFieldValueConverter = builder.fromFieldValueConverter; | ||
this.operation = builder.operation; | ||
this.collectorKey = CollectorKey.create(); | ||
} | ||
|
||
@Override | ||
public Extractor<K> request(AggregationRequestContext context) { | ||
JoiningLongMultiValuesSource source = JoiningLongMultiValuesSource.fromField( | ||
absoluteFieldPath, createNestedDocsProvider( context ) | ||
); | ||
if ( "sum".equals( operation ) ) { | ||
context.requireCollector( new SumCollectorFactory( source, collectorKey ) ); | ||
} | ||
return new LuceneNumericMetricFieldAggregationExtraction(); | ||
} | ||
|
||
@Override | ||
public Set<String> indexNames() { | ||
return indexNames; | ||
} | ||
|
||
private class LuceneNumericMetricFieldAggregationExtraction implements Extractor<K> { | ||
|
||
@Override | ||
public K extract(AggregationExtractContext context) { | ||
Long collector = context.getFacets( collectorKey ); | ||
E e = numericDomain.sortedDocValueToTerm( collector ); | ||
F decode = codec.decode( e ); | ||
return fromFieldValueConverter.fromDocumentValue( decode, context.fromDocumentValueConvertContext() ); | ||
} | ||
} | ||
|
||
public static class Factory<F> | ||
extends AbstractLuceneCodecAwareSearchQueryElementFactory<FieldMetricAggregationBuilder.TypeSelector, | ||
F, | ||
AbstractLuceneNumericFieldCodec<F, ?>> { | ||
|
||
private final String operation; | ||
|
||
public Factory(AbstractLuceneNumericFieldCodec<F, ?> codec, String operation) { | ||
super( codec ); | ||
this.operation = operation; | ||
} | ||
|
||
@Override | ||
public FieldMetricAggregationBuilder.TypeSelector create(LuceneSearchIndexScope<?> scope, | ||
LuceneSearchIndexValueFieldContext<F> field) { | ||
return new TypeSelector<>( codec, scope, field, operation ); | ||
} | ||
} | ||
|
||
private 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; | ||
|
||
private TypeSelector(AbstractLuceneNumericFieldCodec<F, ?> codec, | ||
LuceneSearchIndexScope<?> scope, LuceneSearchIndexValueFieldContext<F> field, | ||
String operation) { | ||
this.codec = codec; | ||
this.scope = scope; | ||
this.field = field; | ||
this.operation = operation; | ||
} | ||
|
||
@Override | ||
public <T> Builder<F, ?, T> type(Class<T> expectedType, ValueConvert convert) { | ||
return new Builder<>( codec, scope, field, | ||
field.type().projectionConverter( convert ).withConvertedType( expectedType, field ), | ||
operation | ||
); | ||
} | ||
} | ||
|
||
private static class Builder<F, E extends Number, K> extends AbstractBuilder<K> | ||
implements FieldMetricAggregationBuilder<K> { | ||
|
||
private final AbstractLuceneNumericFieldCodec<F, E> codec; | ||
private final ProjectionConverter<F, ? extends K> fromFieldValueConverter; | ||
private final String operation; | ||
|
||
public Builder(AbstractLuceneNumericFieldCodec<F, E> codec, LuceneSearchIndexScope<?> scope, | ||
LuceneSearchIndexValueFieldContext<F> field, | ||
ProjectionConverter<F, ? extends K> fromFieldValueConverter, | ||
String operation) { | ||
super( scope, field ); | ||
this.codec = codec; | ||
this.fromFieldValueConverter = fromFieldValueConverter; | ||
this.operation = operation; | ||
} | ||
|
||
@Override | ||
public LuceneMetricNumericFieldAggregation<F, E, K> build() { | ||
return new LuceneMetricNumericFieldAggregation<>( this ); | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.