From 303de72a1d8bf80758bcd64178b278a389ea4afb Mon Sep 17 00:00:00 2001 From: Tim Owen Date: Wed, 4 Nov 2020 16:53:34 +0000 Subject: [PATCH 1/2] Expose some FacetContext and SlotAcc fields and classes to enable our custom aggregates to be moved into a plugin --- .../solr/search/facet/FacetContext.java | 36 ++++++++++++ .../org/apache/solr/search/facet/SlotAcc.java | 56 +++++++++---------- 2 files changed, 64 insertions(+), 28 deletions(-) diff --git a/solr/core/src/java/org/apache/solr/search/facet/FacetContext.java b/solr/core/src/java/org/apache/solr/search/facet/FacetContext.java index 86aa3add7fc2..cef9d018cc78 100644 --- a/solr/core/src/java/org/apache/solr/search/facet/FacetContext.java +++ b/solr/core/src/java/org/apache/solr/search/facet/FacetContext.java @@ -53,6 +53,42 @@ public boolean isShard() { return (flags & IS_SHARD) != 0; } + public FacetProcessor getFacetProcessor() { + return processor; + } + + public Map getFacetInfo() { + return facetInfo; + } + + public QueryContext getQueryContext() { + return qcontext; + } + + public SolrQueryRequest getRequest() { + return req; + } + + public SolrIndexSearcher getSearcher() { + return searcher; + } + + public Query getFilter() { + return filter; + } + + public DocSet getBase() { + return base; + } + + public FacetContext getParent() { + return parent; + } + + public int getFlags() { + return flags; + } + /** * @param filter The filter for the bucket that resulted in this context/domain. Can be null if this is the root context. * @param domain The resulting set of documents for this facet. diff --git a/solr/core/src/java/org/apache/solr/search/facet/SlotAcc.java b/solr/core/src/java/org/apache/solr/search/facet/SlotAcc.java index d6da5956a441..f00a51628f24 100644 --- a/solr/core/src/java/org/apache/solr/search/facet/SlotAcc.java +++ b/solr/core/src/java/org/apache/solr/search/facet/SlotAcc.java @@ -275,11 +275,11 @@ public boolean isAllBuckets() { // TODO: we should really have a decoupled value provider... // This would enhance reuse and also prevent multiple lookups of same value across diff stats - abstract static class FuncSlotAcc extends SlotAcc { + public abstract static class FuncSlotAcc extends SlotAcc { protected final ValueSource valueSource; protected FunctionValues values; - public FuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { + public FuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { super(fcontext); this.valueSource = values; } @@ -298,15 +298,15 @@ public void setNextReader(LeafReaderContext readerContext) throws IOException { // double-slot-func -> func-slot -> slot -> acc // double-slot-func -> double-slot -> slot -> acc - abstract static class DoubleFuncSlotAcc extends FuncSlotAcc { - double[] result; // TODO: use DoubleArray - double initialValue; + public abstract static class DoubleFuncSlotAcc extends FuncSlotAcc { + protected double[] result; // TODO: use DoubleArray + protected double initialValue; - public DoubleFuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { + public DoubleFuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { this(values, fcontext, numSlots, 0); } - public DoubleFuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots, double initialValue) { + public DoubleFuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots, double initialValue) { super(values, fcontext, numSlots); this.initialValue = initialValue; result = new double[numSlots]; @@ -336,11 +336,11 @@ public void resize(Resizer resizer) { } } - abstract static class LongFuncSlotAcc extends FuncSlotAcc { - long[] result; - long initialValue; + public abstract static class LongFuncSlotAcc extends FuncSlotAcc { + protected long[] result; + protected long initialValue; - public LongFuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots, long initialValue) { + public LongFuncSlotAcc(ValueSource values, FacetContext fcontext, int numSlots, long initialValue) { super(values, fcontext, numSlots); this.initialValue = initialValue; result = new long[numSlots]; @@ -370,11 +370,11 @@ public void resize(Resizer resizer) { } } - abstract class IntSlotAcc extends SlotAcc { - int[] result; // use LongArray32 - int initialValue; + public abstract static class IntSlotAcc extends SlotAcc { + protected int[] result; // use LongArray32 + protected int initialValue; - public IntSlotAcc(FacetContext fcontext, int numSlots, int initialValue) { + public IntSlotAcc(FacetContext fcontext, int numSlots, int initialValue) { super(fcontext); this.initialValue = initialValue; result = new int[numSlots]; @@ -432,7 +432,7 @@ public void collect(int doc, int slotNum, IntFunction slotContext) static class AvgSlotAcc extends DoubleFuncSlotAcc { int[] counts; - public AvgSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { + public AvgSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { super(values, fcontext, numSlots); counts = new int[numSlots]; } @@ -486,7 +486,7 @@ static class VarianceSlotAcc extends DoubleFuncSlotAcc { int[] counts; double[] sum; - public VarianceSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { + public VarianceSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { super(values, fcontext, numSlots); counts = new int[numSlots]; sum = new double[numSlots]; @@ -507,7 +507,7 @@ public void resize(Resizer resizer) { } private double variance(int slot) { - return AggUtil.uncorrectedVariance(result[slot], sum[slot], counts[slot]); // calc once and cache in result? + return AggUtil.uncorrectedVariance(result[slot], sum[slot], counts[slot]); // calc once and cache in result? } @Override @@ -543,7 +543,7 @@ static class StddevSlotAcc extends DoubleFuncSlotAcc { int[] counts; double[] sum; - public StddevSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { + public StddevSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { super(values, fcontext, numSlots); counts = new int[numSlots]; sum = new double[numSlots]; @@ -564,7 +564,7 @@ public void resize(Resizer resizer) { } private double stdDev(int slot) { - return AggUtil.uncorrectedStdDev(result[slot], sum[slot], counts[slot]); // calc once and cache in result? + return AggUtil.uncorrectedStdDev(result[slot], sum[slot], counts[slot]); // calc once and cache in result? } @Override @@ -602,9 +602,9 @@ public CountSlotAcc(FacetContext fcontext) { super(fcontext); } - public abstract void incrementCount(int slot, int count); + public abstract void incrementCount(int slot, int count); - public abstract int getCount(int slot); + public abstract int getCount(int slot); } /** @@ -667,9 +667,9 @@ public int getCount(int slot) { static class CountSlotArrAcc extends CountSlotAcc { int[] result; - public CountSlotArrAcc(FacetContext fcontext, int numSlots) { + public CountSlotArrAcc(FacetContext fcontext, int numSlots) { super(fcontext); - result = new int[numSlots]; + result = new int[numSlots]; } @Override @@ -681,7 +681,7 @@ public void collect(int doc, int slotNum, IntFunction slotContext) @Override public int compare(int slotA, int slotB) { - return Integer.compare(result[slotA], result[slotB]); + return Integer.compare(result[slotA], result[slotB]); } @Override @@ -689,16 +689,16 @@ public Object getValue(int slotNum) throws IOException { return result[slotNum]; } - public void incrementCount(int slot, int count) { + public void incrementCount(int slot, int count) { result[slot] += count; } - public int getCount(int slot) { + public int getCount(int slot) { return result[slot]; } // internal and expert - int[] getCountArray() { + int[] getCountArray() { return result; } From 49046a01935ff46f01553217342777bfc0dd8b6c Mon Sep 17 00:00:00 2001 From: Tim Owen Date: Fri, 6 Nov 2020 13:04:40 +0000 Subject: [PATCH 2/2] Add test to confirm custom aggregates can be written outside the solr.search.facet package --- .../search/function/AggValueSourceTest.java | 106 ++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 solr/core/src/test/org/apache/solr/search/function/AggValueSourceTest.java diff --git a/solr/core/src/test/org/apache/solr/search/function/AggValueSourceTest.java b/solr/core/src/test/org/apache/solr/search/function/AggValueSourceTest.java new file mode 100644 index 000000000000..287fe34463b4 --- /dev/null +++ b/solr/core/src/test/org/apache/solr/search/function/AggValueSourceTest.java @@ -0,0 +1,106 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.solr.search.function; + +import java.io.IOException; +import java.util.function.IntFunction; + +import org.apache.lucene.queries.function.ValueSource; +import org.apache.lucene.queries.function.valuesource.ConstValueSource; +import org.apache.solr.SolrTestCase; +import org.apache.solr.search.facet.FacetContext; +import org.apache.solr.search.facet.FacetMerger; +import org.apache.solr.search.facet.SimpleAggValueSource; +import org.apache.solr.search.facet.SlotAcc; +import org.junit.Test; + +/** Tests that AggValueSource can be extended outside its package */ +public class AggValueSourceTest extends SolrTestCase { + + @Test + public void testCustomAgg() { + // All we're really interested in testing here is that the custom agg compiles and can be created + final CustomAggregate customAggregate = new CustomAggregate(new ConstValueSource(123.0f)); + final FacetMerger facetMerger = customAggregate.createFacetMerger(0.0D); + } + + static class CustomAggregate extends SimpleAggValueSource { + + CustomAggregate(ValueSource vs) { + super("customagg", vs); + } + + @Override + public SlotAcc createSlotAcc(FacetContext fcontext, int numDocs, int numSlots) { + // check we can get access to the request and searcher, via the context + if (fcontext.getRequest().getCore() != fcontext.getSearcher().getCore()) { + throw new IllegalStateException("Searcher and request out of sync"); + } + return new CustomSlotAcc(getArg(), fcontext, numSlots); + } + + static class CustomSlotAcc extends SlotAcc.DoubleFuncSlotAcc { + + CustomSlotAcc(ValueSource values, FacetContext fcontext, int numSlots) { + super(values, fcontext, numSlots); + } + + @Override + public void collect(int doc, int slot, IntFunction slotContext) throws IOException { + result[slot] += values.doubleVal(doc); + } + + @Override + public Object getValue(int slot) { + if (fcontext.isShard()) { + // shard-specific logic here + } + return super.getValue(slot); + } + } + + @Override + public FacetMerger createFacetMerger(Object prototype) { + return new FacetMerger() { + double total = 0.0D; + + @Override + public void merge(Object facetResult, Context mcontext) { + total += (Double)facetResult; + } + + @Override + public void finish(Context mcontext) { } + + @Override + public Object getMergedResult() { + return total; + } + }; + } + + @Override + public int hashCode() { + return 0; + } + + @Override + public String description() { + return "customagg()"; + } + } +}