Skip to content

Commit 2b8c412

Browse files
author
elasticsearchmachine
committed
Refactoring, added docs
1 parent d0bebb4 commit 2b8c412

File tree

13 files changed

+62
-75
lines changed

13 files changed

+62
-75
lines changed

server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -328,6 +328,7 @@ public BucketedSort newBucketedSort(
328328
};
329329
}
330330
case NUMBER -> {
331+
// TODO Temporary hack to allow vector similarity functions access from sorting scripts. Needs to be done properly.
331332
final ScoreScript.Factory numberSortFactory = context.compile(script, ScoreScript.CONTEXT);
332333
// searchLookup is unnecessary here, as it's just used for expressions
333334
final ScoreScript.LeafFactory numberSortScriptFactory = numberSortFactory.newFactory(script.getParams(), searchLookup);

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/Expression.java

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -212,14 +212,33 @@ public String propertiesToString(boolean skipIfChild) {
212212
return super.propertiesToString(false);
213213
}
214214

215+
/**
216+
* Returns whether the expression is pushable to Lucene:
217+
* <ul>
218+
* <li>PREFERRED: We should push whenever possible. An expression that has a PREFERRED expression and does not have a NOT_SUPPORTED one can be pushed down to Lucene.</li>
219+
* <li>SUPPORTED: Expression can be pushed down but does not offer a significant advantage in terms of performance, or depends on other expressions that are part of it.</li>
220+
* <li>NOT_SUPPORTED: Expression can't be pushed down.</li>
221+
* </ul>
222+
*
223+
* @return the pushable options for this expression
224+
*/
215225
public PushableOptions pushableOptions() {
216226
return PushableOptions.NOT_SUPPORTED;
217227
}
218228

219-
public String asScript() {
220-
throw new UnsupportedOperationException("asScript not implemented for " + getWriteableName());
229+
/**
230+
* Returns the expression as a painless script to be pushed down to Lucene
231+
*
232+
* @return the painless script representation of this expression
233+
* @throws UnsupportedOperationException if the expression does not support being represented as a script
234+
*/
235+
public String asPushableScript() {
236+
throw new UnsupportedOperationException("asPushableScript not implemented for " + getWriteableName());
221237
}
222238

239+
/**
240+
* See {@link #pushableOptions()}}
241+
*/
223242
public enum PushableOptions {
224243
SUPPORTED,
225244
NOT_SUPPORTED,

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/FieldAttribute.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,7 @@ public PushableOptions pushableOptions() {
266266
}
267267

268268
@Override
269-
public String asScript() {
269+
public String asPushableScript() {
270270
return "doc['" + fieldName().string() + "'].value";
271271
}
272272
}

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/Literal.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ public PushableOptions pushableOptions() {
244244
}
245245

246246
@Override
247-
public String asScript() {
247+
public String asPushableScript() {
248248
return Objects.toString(value());
249249
}
250250
}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperator.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,7 @@ private Page emit() {
288288
DoubleBlock scores = null;
289289
Page page = null;
290290

291+
// Prepares builders for sort values extraction
291292
int extractSortCount = (int) sorts.stream().filter(s -> s instanceof ScriptSortBuilder).count();
292293
DoubleBlock.Builder[] sortValuesBlockBuilders = new DoubleBlock.Builder[extractSortCount];
293294
for (int i = 0; i < sortValuesBlockBuilders.length; i++) {
@@ -310,6 +311,7 @@ private Page emit() {
310311
float score = getScore(topDocs[i]);
311312
currentScoresBuilder.appendDouble(score);
312313
}
314+
// Extracts sort values for script sorts, so they can be retrieved as values in the compute engine
313315
if ((extractSortCount > 0) && topDocs[i] instanceof FieldDoc fieldDoc) {
314316
int sortIndex = 0;
315317
int extractedSortIndex = 0;

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
import static org.elasticsearch.index.IndexSettings.INDEX_MAPPER_SOURCE_MODE_SETTING;
3535
import static org.elasticsearch.index.mapper.SourceFieldMapper.Mode.SYNTHETIC;
3636
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
37-
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION;
37+
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION_PUSHABLE;
3838

3939
public class DenseVectorFieldTypeIT extends AbstractEsqlIntegTestCase {
4040

@@ -91,7 +91,7 @@ public DenseVectorFieldTypeIT(
9191
private final Map<Integer, List<Number>> indexedVectors = new HashMap<>();
9292

9393
public void testRetrieveFieldType() {
94-
assumeTrue("Need L2_NORM available for dense_vector retrieval", L2_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled());
94+
assumeTrue("Need L2_NORM available for dense_vector retrieval", L2_NORM_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled());
9595

9696
var query = """
9797
FROM test
@@ -107,7 +107,7 @@ public void testRetrieveFieldType() {
107107

108108
@SuppressWarnings("unchecked")
109109
public void testRetrieveTopNDenseVectorFieldData() {
110-
assumeTrue("Need L2_NORM available for dense_vector retrieval", L2_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled());
110+
assumeTrue("Need L2_NORM available for dense_vector retrieval", L2_NORM_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled());
111111

112112
var query = """
113113
FROM test
@@ -137,7 +137,7 @@ public void testRetrieveTopNDenseVectorFieldData() {
137137

138138
@SuppressWarnings("unchecked")
139139
public void testRetrieveDenseVectorFieldData() {
140-
assumeTrue("Need L2_NORM available for dense_vector retrieval", L2_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled());
140+
assumeTrue("Need L2_NORM available for dense_vector retrieval", L2_NORM_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled());
141141

142142
var query = """
143143
FROM test

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/vector/VectorSimilarityFunctionsIT.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,19 +50,19 @@ public class VectorSimilarityFunctionsIT extends AbstractEsqlIntegTestCase {
5050
public static Iterable<Object[]> parameters() throws Exception {
5151
List<Object[]> params = new ArrayList<>();
5252

53-
if (EsqlCapabilities.Cap.COSINE_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
53+
if (EsqlCapabilities.Cap.COSINE_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled()) {
5454
params.add(new Object[] { "v_cosine", (SimilarityEvaluatorFunction) CosineSimilarity::calculateSimilarity });
5555
}
56-
if (EsqlCapabilities.Cap.DOT_PRODUCT_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
56+
if (EsqlCapabilities.Cap.DOT_PRODUCT_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled()) {
5757
params.add(new Object[] { "v_dot_product", (SimilarityEvaluatorFunction) DotProduct::calculateSimilarity });
5858
}
59-
if (EsqlCapabilities.Cap.L1_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
59+
if (EsqlCapabilities.Cap.L1_NORM_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled()) {
6060
params.add(new Object[] { "v_l1_norm", (SimilarityEvaluatorFunction) L1Norm::calculateSimilarity });
6161
}
62-
if (EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
62+
if (EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled()) {
6363
params.add(new Object[] { "v_l2_norm", (SimilarityEvaluatorFunction) L2Norm::calculateSimilarity });
6464
}
65-
if (EsqlCapabilities.Cap.HAMMING_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
65+
if (EsqlCapabilities.Cap.HAMMING_VECTOR_SIMILARITY_FUNCTION_PUSHABLE.isEnabled()) {
6666
params.add(new Object[] { "v_hamming", (SimilarityEvaluatorFunction) Hamming::calculateSimilarity });
6767
}
6868

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Magnitude.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ public PushableOptions pushableOptions() {
139139
}
140140

141141
@Override
142-
public String asScript() {
142+
public String asPushableScript() {
143143
return "doc['" + ((FieldAttribute) field()).name() + "'].magnitude";
144144
}
145145

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/VectorSimilarityFunction.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,14 +94,14 @@ public PushableOptions pushableOptions() {
9494
}
9595

9696
@Override
97-
public final String asScript() {
97+
public final String asPushableScript() {
9898
String queryVector = left().foldable() ? queryVectorAsScript(left()) : queryVectorAsScript(right());
9999
String fieldName = ((FieldAttribute) (left().foldable() ? right() : left())).name();
100100
return "doc['" + fieldName + "'].size() == 0 ? 0 : " + scriptFunctionName() + "(" + queryVector + ", '" + fieldName + "')";
101101
}
102102

103103
protected String queryVectorAsScript(Expression queryVector) {
104-
return queryVector.asScript();
104+
return queryVector.asPushableScript();
105105
}
106106

107107
protected abstract String scriptFunctionName();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/operator/arithmetic/EsqlArithmeticOperation.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -170,8 +170,8 @@ public PushableOptions pushableOptions() {
170170
}
171171

172172
@Override
173-
public String asScript() {
174-
return left().asScript() + function().symbol() + right().asScript();
173+
public String asPushableScript() {
174+
return left().asPushableScript() + function().symbol() + right().asPushableScript();
175175
}
176176

177177
@Override

0 commit comments

Comments
 (0)