290 lines
11 KiB
Java
290 lines
11 KiB
Java
package it.cavallium.dbengine.lucene.collector;
|
|
|
|
import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
|
|
import java.io.IOException;
|
|
import java.util.ArrayList;
|
|
import java.util.Arrays;
|
|
import java.util.List;
|
|
import org.apache.commons.lang3.NotImplementedException;
|
|
import org.apache.lucene.facet.FacetResult;
|
|
import org.apache.lucene.facet.Facets;
|
|
import org.apache.lucene.facet.FacetsConfig;
|
|
import org.apache.lucene.facet.LabelAndValue;
|
|
import org.apache.lucene.facet.RandomSamplingFacetsCollector;
|
|
import org.apache.lucene.facet.range.DoubleRange;
|
|
import org.apache.lucene.facet.range.DoubleRangeFacetCounts;
|
|
import org.apache.lucene.facet.range.LongRange;
|
|
import org.apache.lucene.facet.range.LongRangeFacetCounts;
|
|
import org.apache.lucene.facet.range.Range;
|
|
import org.apache.lucene.search.BooleanClause.Occur;
|
|
import org.apache.lucene.search.BooleanQuery;
|
|
import org.apache.lucene.search.DoubleValuesSource;
|
|
import org.apache.lucene.search.IndexSearcher;
|
|
import org.apache.lucene.search.LongValuesSource;
|
|
import org.apache.lucene.search.MatchNoDocsQuery;
|
|
import org.apache.lucene.search.Query;
|
|
import org.apache.lucene.search.ScoreMode;
|
|
import org.jetbrains.annotations.Nullable;
|
|
|
|
public class DecimalBucketMultiCollectorManager implements CollectorMultiManager<Buckets, Buckets> {
|
|
|
|
private static final boolean USE_SINGLE_FACET_COLLECTOR = true;
|
|
private static final boolean AMORTIZE = true;
|
|
private final boolean randomSamplingEnabled;
|
|
private final FastFacetsCollectorManager facetsCollectorManager;
|
|
private final FastRandomSamplingFacetsCollector randomSamplingFacetsCollector;
|
|
private final Range[] bucketRanges;
|
|
|
|
private final List<Query> queries;
|
|
private final @Nullable Query normalizationQuery;
|
|
private final @Nullable Integer collectionRate;
|
|
private final @Nullable Integer sampleSize;
|
|
|
|
private final String bucketField;
|
|
private final BucketValueSource bucketValueSource;
|
|
|
|
private final double totalLength;
|
|
private final double bucketLength;
|
|
private final double minimum;
|
|
private final double maximum;
|
|
private final int buckets;
|
|
|
|
// todo: replace with an argument
|
|
private static final boolean USE_LONGS = true;
|
|
|
|
public DecimalBucketMultiCollectorManager(double minimum,
|
|
double maximum,
|
|
double buckets,
|
|
String bucketField,
|
|
BucketValueSource bucketValueSource,
|
|
List<Query> queries,
|
|
@Nullable Query normalizationQuery,
|
|
@Nullable Integer collectionRate,
|
|
@Nullable Integer sampleSize) {
|
|
this.queries = queries;
|
|
this.normalizationQuery = normalizationQuery;
|
|
var bucketsInt = (int) Math.ceil(buckets);
|
|
this.minimum = minimum;
|
|
this.maximum = maximum;
|
|
this.buckets = bucketsInt;
|
|
this.bucketLength = (maximum - minimum) / bucketsInt;
|
|
this.totalLength = bucketLength * bucketsInt;
|
|
this.bucketField = bucketField;
|
|
this.bucketValueSource = bucketValueSource;
|
|
this.collectionRate = collectionRate;
|
|
this.sampleSize = sampleSize;
|
|
|
|
if (USE_LONGS) {
|
|
this.bucketRanges = new LongRange[bucketsInt];
|
|
} else {
|
|
this.bucketRanges = new DoubleRange[bucketsInt];
|
|
}
|
|
for (int i = 0; i < bucketsInt; i++) {
|
|
double offsetMin = minimum + (bucketLength * i);
|
|
double offsetMax = minimum + (bucketLength * (i + 1));
|
|
if (USE_LONGS) {
|
|
this.bucketRanges[i] = new LongRange(Integer.toString(i),
|
|
(long) offsetMin,
|
|
true,
|
|
(long) offsetMax,
|
|
i == bucketsInt - 1
|
|
);
|
|
} else {
|
|
this.bucketRanges[i] = new DoubleRange(Integer.toString(i),
|
|
offsetMin,
|
|
true,
|
|
offsetMax,
|
|
i == bucketsInt - 1
|
|
);
|
|
}
|
|
}
|
|
|
|
this.randomSamplingEnabled = sampleSize != null;
|
|
int intCollectionRate = this.collectionRate == null ? 1 : this.collectionRate;
|
|
if (randomSamplingEnabled) {
|
|
randomSamplingFacetsCollector = new FastRandomSamplingFacetsCollector(intCollectionRate, sampleSize, 0);
|
|
this.facetsCollectorManager = null;
|
|
} else {
|
|
this.randomSamplingFacetsCollector = null;
|
|
this.facetsCollectorManager = new FastFacetsCollectorManager(intCollectionRate);
|
|
}
|
|
}
|
|
|
|
public double[] newBuckets() {
|
|
return new double[buckets];
|
|
}
|
|
|
|
public Buckets search(IndexSearcher indexSearcher) throws IOException {
|
|
Query query;
|
|
if (USE_SINGLE_FACET_COLLECTOR && normalizationQuery != null) {
|
|
query = normalizationQuery;
|
|
} else if (queries.size() == 0) {
|
|
query = new MatchNoDocsQuery();
|
|
} else if (queries.size() == 1) {
|
|
query = queries.get(0);
|
|
} else {
|
|
var booleanQueryBuilder = new BooleanQuery.Builder();
|
|
for (Query queryEntry : queries) {
|
|
booleanQueryBuilder.add(queryEntry, Occur.SHOULD);
|
|
}
|
|
booleanQueryBuilder.setMinimumNumberShouldMatch(1);
|
|
query = booleanQueryBuilder.build();
|
|
}
|
|
it.cavallium.dbengine.lucene.collector.FacetsCollector queryFacetsCollector;
|
|
if (randomSamplingEnabled) {
|
|
indexSearcher.search(query, randomSamplingFacetsCollector);
|
|
queryFacetsCollector = randomSamplingFacetsCollector;
|
|
} else {
|
|
queryFacetsCollector = indexSearcher.search(query, facetsCollectorManager);
|
|
}
|
|
double[] reducedNormalizationBuckets = newBuckets();
|
|
List<DoubleArrayList> seriesReducedBuckets = new ArrayList<>(queries.size());
|
|
for (int i = 0; i < queries.size(); i++) {
|
|
var buckets = newBuckets();
|
|
seriesReducedBuckets.add(DoubleArrayList.wrap(buckets));
|
|
}
|
|
int serieIndex = 0;
|
|
for (Query queryEntry : queries) {
|
|
var reducedBuckets = seriesReducedBuckets.get(serieIndex);
|
|
Facets facets;
|
|
if (USE_LONGS) {
|
|
LongValuesSource valuesSource;
|
|
if (bucketValueSource instanceof BucketValueSource.NullValueSource) {
|
|
|
|
valuesSource = null;
|
|
} else if (bucketValueSource instanceof BucketValueSource.ConstantValueSource constantValueSource) {
|
|
valuesSource = LongValuesSource.constant(constantValueSource.constant().longValue());
|
|
} else if (bucketValueSource instanceof BucketValueSource.LongBucketValueSource longBucketValueSource) {
|
|
valuesSource = longBucketValueSource.source();
|
|
} else {
|
|
throw new IllegalArgumentException("Wrong value source type: " + bucketValueSource);
|
|
}
|
|
facets = new LongRangeFacetCounts(bucketField,
|
|
valuesSource,
|
|
queryFacetsCollector.getLuceneFacetsCollector(),
|
|
USE_SINGLE_FACET_COLLECTOR && normalizationQuery != null || queries.size() > 1 ? queryEntry : null,
|
|
(LongRange[]) bucketRanges
|
|
);
|
|
} else {
|
|
DoubleValuesSource valuesSource;
|
|
if (bucketValueSource instanceof BucketValueSource.NullValueSource) {
|
|
valuesSource = null;
|
|
} else if (bucketValueSource instanceof BucketValueSource.ConstantValueSource constantValueSource) {
|
|
valuesSource = DoubleValuesSource.constant(constantValueSource.constant().longValue());
|
|
} else if (bucketValueSource instanceof BucketValueSource.DoubleBucketValueSource doubleBucketValueSource) {
|
|
valuesSource = doubleBucketValueSource.source();
|
|
} else {
|
|
throw new IllegalArgumentException("Wrong value source type: " + bucketValueSource);
|
|
}
|
|
facets = new DoubleRangeFacetCounts(bucketField,
|
|
valuesSource,
|
|
queryFacetsCollector.getLuceneFacetsCollector(),
|
|
USE_SINGLE_FACET_COLLECTOR && normalizationQuery != null || queries.size() > 1 ? queryEntry : null,
|
|
(DoubleRange[]) bucketRanges
|
|
);
|
|
}
|
|
FacetResult children = facets.getTopChildren(1, bucketField);
|
|
if (AMORTIZE && randomSamplingEnabled) {
|
|
var cfg = new FacetsConfig();
|
|
for (Range bucketRange : bucketRanges) {
|
|
cfg.setIndexFieldName(bucketRange.label, bucketField);
|
|
}
|
|
((RandomSamplingFacetsCollector) queryFacetsCollector.getLuceneFacetsCollector()).amortizeFacetCounts(children, cfg, indexSearcher);
|
|
}
|
|
for (LabelAndValue labelAndValue : children.labelValues) {
|
|
var index = Integer.parseInt(labelAndValue.label);
|
|
reducedBuckets.set(index, reducedBuckets.getDouble(index) + labelAndValue.value.doubleValue());
|
|
}
|
|
serieIndex++;
|
|
}
|
|
|
|
it.cavallium.dbengine.lucene.collector.FacetsCollector normalizationFacetsCollector;
|
|
Facets normalizationFacets;
|
|
if (normalizationQuery != null) {
|
|
if (USE_SINGLE_FACET_COLLECTOR) {
|
|
normalizationFacetsCollector = queryFacetsCollector;
|
|
} else if (randomSamplingEnabled) {
|
|
indexSearcher.search(normalizationQuery, randomSamplingFacetsCollector);
|
|
normalizationFacetsCollector = randomSamplingFacetsCollector;
|
|
} else {
|
|
normalizationFacetsCollector = indexSearcher.search(normalizationQuery, facetsCollectorManager);
|
|
}
|
|
if (USE_LONGS) {
|
|
LongValuesSource valuesSource;
|
|
if (bucketValueSource instanceof BucketValueSource.NullValueSource) {
|
|
valuesSource = null;
|
|
} else if (bucketValueSource instanceof BucketValueSource.ConstantValueSource constantValueSource) {
|
|
valuesSource = LongValuesSource.constant(constantValueSource.constant().longValue());
|
|
} else if (bucketValueSource instanceof BucketValueSource.LongBucketValueSource longBucketValueSource) {
|
|
valuesSource = longBucketValueSource.source();
|
|
} else {
|
|
throw new IllegalArgumentException("Wrong value source type: " + bucketValueSource);
|
|
}
|
|
normalizationFacets = new LongRangeFacetCounts(bucketField,
|
|
valuesSource,
|
|
normalizationFacetsCollector.getLuceneFacetsCollector(),
|
|
null,
|
|
(LongRange[]) bucketRanges
|
|
);
|
|
} else {
|
|
DoubleValuesSource valuesSource;
|
|
if (bucketValueSource instanceof BucketValueSource.NullValueSource) {
|
|
valuesSource = null;
|
|
} else if (bucketValueSource instanceof BucketValueSource.ConstantValueSource constantValueSource) {
|
|
valuesSource = DoubleValuesSource.constant(constantValueSource.constant().longValue());
|
|
} else if (bucketValueSource instanceof BucketValueSource.DoubleBucketValueSource doubleBucketValueSource) {
|
|
valuesSource = doubleBucketValueSource.source();
|
|
} else {
|
|
throw new IllegalArgumentException("Wrong value source type: " + bucketValueSource);
|
|
}
|
|
normalizationFacets = new DoubleRangeFacetCounts(bucketField,
|
|
valuesSource,
|
|
normalizationFacetsCollector.getLuceneFacetsCollector(),
|
|
null,
|
|
(DoubleRange[]) bucketRanges
|
|
);
|
|
}
|
|
var normalizationChildren = normalizationFacets.getTopChildren(0, bucketField);
|
|
if (AMORTIZE && randomSamplingEnabled) {
|
|
var cfg = new FacetsConfig();
|
|
for (Range bucketRange : bucketRanges) {
|
|
cfg.setIndexFieldName(bucketRange.label, bucketField);
|
|
}
|
|
((RandomSamplingFacetsCollector) normalizationFacetsCollector.getLuceneFacetsCollector()).amortizeFacetCounts(normalizationChildren, cfg, indexSearcher);
|
|
}
|
|
for (LabelAndValue labelAndValue : normalizationChildren.labelValues) {
|
|
var index = Integer.parseInt(labelAndValue.label);
|
|
reducedNormalizationBuckets[index] += labelAndValue.value.doubleValue();
|
|
}
|
|
} else {
|
|
Arrays.fill(reducedNormalizationBuckets, 1);
|
|
}
|
|
return new Buckets(seriesReducedBuckets, DoubleArrayList.wrap(reducedNormalizationBuckets));
|
|
}
|
|
|
|
@Override
|
|
public ScoreMode scoreMode() {
|
|
throw new NotImplementedException();
|
|
}
|
|
|
|
@Override
|
|
public Buckets reduce(List<Buckets> reducedBucketsList) {
|
|
List<DoubleArrayList> seriesReducedValues = new ArrayList<>();
|
|
double[] reducedTotals = newBuckets();
|
|
for (var seriesBuckets : reducedBucketsList) {
|
|
for (DoubleArrayList values : seriesBuckets.seriesValues()) {
|
|
double[] reducedValues = newBuckets();
|
|
for (int i = 0; i < values.size(); i++) {
|
|
reducedValues[i] += values.getDouble(i);
|
|
}
|
|
seriesReducedValues.add(DoubleArrayList.wrap(reducedValues));
|
|
}
|
|
var totals = seriesBuckets.totals();
|
|
for (int i = 0; i < totals.size(); i++) {
|
|
reducedTotals[i] += totals.getDouble(i);
|
|
}
|
|
}
|
|
return new Buckets(seriesReducedValues, DoubleArrayList.wrap(reducedTotals));
|
|
}
|
|
}
|