mirror of https://github.com/apache/lucene.git
SOLR-8673: Enable custom aggregate functions by opening up FacetContext
Prior to this change, it was not feasible to write a custom aggregate function in plugin code because the FacetContext class did not expose its fields for access outside its package. Also some of the useful abstract classes for SlotAcc were package-private too
This commit is contained in:
parent
478c6cb0f5
commit
510f96c69c
|
@ -216,6 +216,8 @@ Improvements
|
|||
|
||||
* SOLR-15015 : Add interleaving algorithm parameter support in Learning To Rank (Alessandro Benedetti)
|
||||
|
||||
* SOLR-8673: Make JSON Facets extendable (Tim Owen via Mikhail Khludnev)
|
||||
|
||||
* SOLR-14965: metrics: Adds two metrics to the SolrCloud Overseer: solr_metrics_overseer_stateUpdateQueueSize
|
||||
and solr_metrics_overseer_collectionWorkQueueSize with corresponding entries in the the Prometheus exporter's
|
||||
default/stock configuration. (Saatchi Bhalla, Megan Carey, Andrzej Białecki, David Smiley)
|
||||
|
|
|
@ -53,6 +53,42 @@ public class FacetContext {
|
|||
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.
|
||||
|
|
|
@ -278,7 +278,7 @@ public abstract class SlotAcc implements Closeable {
|
|||
|
||||
// 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;
|
||||
|
||||
|
@ -301,9 +301,9 @@ public abstract class SlotAcc implements Closeable {
|
|||
// 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) {
|
||||
this(values, fcontext, numSlots, 0);
|
||||
|
@ -339,9 +339,9 @@ public abstract class SlotAcc implements Closeable {
|
|||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
super(values, fcontext, numSlots);
|
||||
|
@ -373,9 +373,9 @@ public abstract class SlotAcc implements Closeable {
|
|||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
super(fcontext);
|
||||
|
|
|
@ -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, long 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()";
|
||||
}
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue