Expose some FacetContext and SlotAcc fields and classes to enable our…

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<String, Object> 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.

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> 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,24 +681,24 @@ public void collect(int doc, int slotNum, IntFunction<SlotContext> slotContext)
 
     @Override
     public int compare(int slotA, int slotB) {
-    return Integer.compare(result[slotA], result[slotB]);
+      return Integer.compare(result[slotA], result[slotB]);
     }
 
     @Override
     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;
     }
 

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> 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()";
+    }
+  }
+}