Implement distance_feature for runtime dates #60851
There are no files selected for viewing
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,218 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License; | ||
| * you may not use this file except in compliance with the Elastic License. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.runtimefields.query; | ||
|
|
||
| import org.apache.lucene.index.LeafReaderContext; | ||
| import org.apache.lucene.index.Term; | ||
| import org.apache.lucene.search.DocIdSetIterator; | ||
| import org.apache.lucene.search.Explanation; | ||
| import org.apache.lucene.search.IndexSearcher; | ||
| import org.apache.lucene.search.QueryVisitor; | ||
| import org.apache.lucene.search.ScoreMode; | ||
| import org.apache.lucene.search.Scorer; | ||
| import org.apache.lucene.search.TwoPhaseIterator; | ||
| import org.apache.lucene.search.Weight; | ||
| import org.elasticsearch.common.CheckedFunction; | ||
| import org.elasticsearch.script.Script; | ||
| import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Objects; | ||
| import java.util.Set; | ||
|
|
||
| public final class LongScriptFieldDistanceFeatureQuery extends AbstractScriptFieldQuery { | ||
| private final CheckedFunction<LeafReaderContext, AbstractLongScriptFieldScript, IOException> leafFactory; | ||
| private final long origin; | ||
| private final long pivot; | ||
| private final float boost; | ||
|
|
||
| public LongScriptFieldDistanceFeatureQuery( | ||
| Script script, | ||
| CheckedFunction<LeafReaderContext, AbstractLongScriptFieldScript, IOException> leafFactory, | ||
| String fieldName, | ||
| long origin, | ||
| long pivot, | ||
| float boost | ||
| ) { | ||
| super(script, fieldName); | ||
| this.leafFactory = leafFactory; | ||
| this.origin = origin; | ||
| this.pivot = pivot; | ||
| this.boost = boost; | ||
| } | ||
|
|
||
| @Override | ||
| public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { | ||
| return new Weight(this) { | ||
| @Override | ||
| public boolean isCacheable(LeafReaderContext ctx) { | ||
| return false; | ||
| } | ||
|
|
||
| @Override | ||
| public void extractTerms(Set<Term> terms) {} | ||
|
|
||
| @Override | ||
| public Scorer scorer(LeafReaderContext context) throws IOException { | ||
| return new DistanceScorer(this, leafFactory.apply(context), context.reader().maxDoc(), boost); | ||
| } | ||
|
|
||
| @Override | ||
| public Explanation explain(LeafReaderContext context, int doc) throws IOException { | ||
| AbstractLongScriptFieldScript script = leafFactory.apply(context); | ||
| script.runForDoc(doc); | ||
| long value = valueWithMinAbsoluteDistance(script); | ||
| float weight = LongScriptFieldDistanceFeatureQuery.this.boost * boost; | ||
| float score = score(weight, distanceFor(value)); | ||
| return Explanation.match( | ||
| score, | ||
| "Distance score, computed as weight * pivot / (pivot + abs(value - origin)) from:", | ||
| Explanation.match(weight, "weight"), | ||
| Explanation.match(pivot, "pivot"), | ||
| Explanation.match(origin, "origin"), | ||
| Explanation.match(value, "current value") | ||
| ); | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| private class DistanceScorer extends Scorer { | ||
| private final AbstractLongScriptFieldScript script; | ||
| private final TwoPhaseIterator twoPhase; | ||
| private final DocIdSetIterator disi; | ||
| private final float weight; | ||
|
|
||
| protected DistanceScorer(Weight weight, AbstractLongScriptFieldScript script, int maxDoc, float boost) { | ||
| super(weight); | ||
| this.script = script; | ||
| twoPhase = new TwoPhaseIterator(DocIdSetIterator.all(maxDoc)) { | ||
|
There was a problem hiding this comment. I am not familiar with runtime fields, but I am wondering if we intend always to create an iterator across all documents? Do we plan to add support to limit number of docs (e.g. only docs returned by a top filter)? There was a problem hiding this comment. I believe we're hoping for |
||
| @Override | ||
| public boolean matches() throws IOException { | ||
| script.runForDoc(approximation().docID()); | ||
| return script.count() > 0; | ||
| } | ||
|
|
||
| @Override | ||
| public float matchCost() { | ||
| return MATCH_COST; | ||
| } | ||
| }; | ||
| disi = TwoPhaseIterator.asDocIdSetIterator(twoPhase); | ||
| this.weight = LongScriptFieldDistanceFeatureQuery.this.boost * boost; | ||
| } | ||
|
|
||
| @Override | ||
| public int docID() { | ||
| return disi.docID(); | ||
| } | ||
|
|
||
| @Override | ||
| public DocIdSetIterator iterator() { | ||
| return disi; | ||
| } | ||
|
|
||
| @Override | ||
| public TwoPhaseIterator twoPhaseIterator() { | ||
| return twoPhase; | ||
| } | ||
|
|
||
| @Override | ||
| public float getMaxScore(int upTo) throws IOException { | ||
| return weight; | ||
| } | ||
|
|
||
| @Override | ||
| public float score() throws IOException { | ||
| if (script.count() == 0) { | ||
| return 0; | ||
| } | ||
| return LongScriptFieldDistanceFeatureQuery.this.score(weight, (double) minAbsoluteDistance(script)); | ||
| } | ||
| } | ||
|
|
||
| long minAbsoluteDistance(AbstractLongScriptFieldScript script) { | ||
| long minDistance = Long.MAX_VALUE; | ||
| for (int i = 0; i < script.count(); i++) { | ||
| minDistance = Math.min(minDistance, distanceFor(script.values()[i])); | ||
| } | ||
| return minDistance; | ||
| } | ||
|
|
||
| long valueWithMinAbsoluteDistance(AbstractLongScriptFieldScript script) { | ||
| long minDistance = Long.MAX_VALUE; | ||
| long minDistanceValue = Long.MAX_VALUE; | ||
| for (int i = 0; i < script.count(); i++) { | ||
| long distance = distanceFor(script.values()[i]); | ||
| if (distance < minDistance) { | ||
| minDistance = distance; | ||
| minDistanceValue = script.values()[i]; | ||
| } | ||
| } | ||
| return minDistanceValue; | ||
| } | ||
|
|
||
| long distanceFor(long value) { | ||
| long distance = Math.max(value, origin) - Math.min(value, origin); | ||
| if (distance < 0) { | ||
| // The distance doesn't fit into signed long so clamp it to MAX_VALUE | ||
| return Long.MAX_VALUE; | ||
| } | ||
| return distance; | ||
| } | ||
|
|
||
| float score(float weight, double distance) { | ||
| return (float) (weight * (pivot / (pivot + distance))); | ||
| } | ||
|
|
||
| @Override | ||
| public String toString(String field) { | ||
| StringBuilder b = new StringBuilder(); | ||
| if (false == fieldName().equals(field)) { | ||
| b.append(fieldName()).append(":"); | ||
| } | ||
| b.append(getClass().getSimpleName()); | ||
| b.append("(origin=").append(origin); | ||
| b.append(",pivot=").append(pivot); | ||
| b.append(",boost=").append(boost).append(")"); | ||
| return b.toString(); | ||
|
|
||
| } | ||
|
|
||
| @Override | ||
| public int hashCode() { | ||
| return Objects.hash(super.hashCode(), origin, pivot); | ||
|
There was a problem hiding this comment. should we incorporate the initial |
||
| } | ||
|
|
||
| @Override | ||
| public boolean equals(Object obj) { | ||
| if (false == super.equals(obj)) { | ||
| return false; | ||
| } | ||
| LongScriptFieldDistanceFeatureQuery other = (LongScriptFieldDistanceFeatureQuery) obj; | ||
| return origin == other.origin && pivot == other.pivot; | ||
| } | ||
|
|
||
| @Override | ||
| public void visit(QueryVisitor visitor) { | ||
| // No subclasses contain any Terms because those have to be strings. | ||
| if (visitor.acceptField(fieldName())) { | ||
| visitor.visitLeaf(this); | ||
| } | ||
| } | ||
|
|
||
| long origin() { | ||
| return origin; | ||
| } | ||
|
|
||
| long pivot() { | ||
| return pivot; | ||
| } | ||
|
|
||
| float boost() { | ||
| return boost; | ||
| } | ||
| } | ||
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
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I wonder what the plan is for the instanceof checks in DistanceFeatureQueryBuilder#doToQuery . Are we ok with keeping those?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
oh actually those are gone upstream, great! sorry for the noise then, you already did what I would asked you to do