Skip to content

Commit 4bc64ab

Browse files
author
elasticsearchmachine
committed
Extract sort values as a new attribute
1 parent f74c87b commit 4bc64ab

File tree

5 files changed

+77
-36
lines changed

5 files changed

+77
-36
lines changed

‎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,6 +244,6 @@ public boolean isPushable() {
244244

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

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import org.elasticsearch.compute.operator.DriverContext;
3434
import org.elasticsearch.compute.operator.SourceOperator;
3535
import org.elasticsearch.core.Releasables;
36+
import org.elasticsearch.search.sort.ScriptSortBuilder;
3637
import org.elasticsearch.search.sort.SortAndFormats;
3738
import org.elasticsearch.search.sort.SortBuilder;
3839

@@ -271,11 +272,20 @@ private Page emit() {
271272
DocBlock docBlock = null;
272273
DoubleBlock scores = null;
273274
Page page = null;
275+
276+
int extractSortCount = (int) sorts.stream().filter(s -> s instanceof ScriptSortBuilder).count();
277+
DoubleBlock.Builder[] sortValuesBlockBuilders = new DoubleBlock.Builder[extractSortCount];
278+
for(int i = 0; i < sortValuesBlockBuilders.length; i++) {
279+
sortValuesBlockBuilders[i] = blockFactory.newDoubleBlockBuilder(size);
280+
}
281+
DoubleBlock[] sortValuesBlocks = null;
282+
274283
try (
275284
IntVector.Builder currentSegmentBuilder = blockFactory.newIntVectorFixedBuilder(size);
276285
IntVector.Builder currentDocsBuilder = blockFactory.newIntVectorFixedBuilder(size);
277286
DoubleVector.Builder currentScoresBuilder = scoreVectorOrNull(size);
278287
) {
288+
279289
int start = offset;
280290
offset += size;
281291
List<LeafReaderContext> leafContexts = perShardCollector.shardContext.searcher().getLeafContexts();
@@ -288,6 +298,27 @@ private Page emit() {
288298
float score = getScore(topDocs[i]);
289299
currentScoresBuilder.appendDouble(score);
290300
}
301+
// TODO Get the blocks according to the sort types for the sorts to extract - could be different types
302+
// Should we do the script field type materialization later? So it's part of loading the field via a script?
303+
// We don't reuse the order, but it's similar to how fetching source works.
304+
if ((extractSortCount > 0) && topDocs[i] instanceof FieldDoc fieldDoc) {
305+
int sortIndex = 0;
306+
int extractedSortIndex = 0;
307+
for (SortBuilder<?> sortBuilder : sorts) {
308+
if (sortBuilder instanceof ScriptSortBuilder) {
309+
Object sortValue = ((FieldDoc) topDocs[i]).fields[sortIndex];
310+
if (sortValue == null) {
311+
sortValuesBlockBuilders[extractedSortIndex].appendNull();
312+
} else if (sortValue instanceof Number numberSort) {
313+
sortValuesBlockBuilders[extractedSortIndex].appendDouble(numberSort.doubleValue());
314+
} else {
315+
throw new IllegalArgumentException("Unsupported sort value type: " + sortValue.getClass());
316+
}
317+
extractedSortIndex++;
318+
}
319+
sortIndex++;
320+
}
321+
}
291322
// Null the top doc so it can be GCed early, just in case.
292323
topDocs[i] = null;
293324
}
@@ -306,6 +337,11 @@ private Page emit() {
306337
scores = currentScoresBuilder.build().asBlock();
307338
page = new Page(size, docBlock, scores);
308339
}
340+
sortValuesBlocks = new DoubleBlock[extractSortCount];
341+
for(int i = 0; i < extractSortCount; i++) {
342+
sortValuesBlocks[i] = sortValuesBlockBuilders[i].build();
343+
}
344+
page = page.appendBlocks(sortValuesBlocks);
309345
} finally {
310346
if (page == null) {
311347
Releasables.closeExpectNoException(shard, segments, docs, docBlock, scores);

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

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -53,18 +53,18 @@ public static Iterable<Object[]> parameters() throws Exception {
5353
if (EsqlCapabilities.Cap.COSINE_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
5454
params.add(new Object[] { "v_cosine", (SimilarityEvaluatorFunction) CosineSimilarity::calculateSimilarity});
5555
}
56-
if (EsqlCapabilities.Cap.DOT_PRODUCT_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
57-
params.add(new Object[] { "v_dot_product", (SimilarityEvaluatorFunction) DotProduct::calculateSimilarity });
58-
}
59-
if (EsqlCapabilities.Cap.L1_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
60-
params.add(new Object[] { "v_l1_norm", (SimilarityEvaluatorFunction) L1Norm::calculateSimilarity });
61-
}
62-
if (EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
63-
params.add(new Object[] { "v_l2_norm", (SimilarityEvaluatorFunction) L2Norm::calculateSimilarity });
64-
}
65-
if (EsqlCapabilities.Cap.HAMMING_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
66-
params.add(new Object[] { "v_hamming", (SimilarityEvaluatorFunction) Hamming::calculateSimilarity });
67-
}
56+
// if (EsqlCapabilities.Cap.DOT_PRODUCT_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
57+
// params.add(new Object[] { "v_dot_product", (SimilarityEvaluatorFunction) DotProduct::calculateSimilarity });
58+
// }
59+
// if (EsqlCapabilities.Cap.L1_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
60+
// params.add(new Object[] { "v_l1_norm", (SimilarityEvaluatorFunction) L1Norm::calculateSimilarity });
61+
// }
62+
// if (EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
63+
// params.add(new Object[] { "v_l2_norm", (SimilarityEvaluatorFunction) L2Norm::calculateSimilarity });
64+
// }
65+
// if (EsqlCapabilities.Cap.HAMMING_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
66+
// params.add(new Object[] { "v_hamming", (SimilarityEvaluatorFunction) Hamming::calculateSimilarity });
67+
// }
6868

6969
return params;
7070
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushTopNToSource.java‎

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -132,28 +132,28 @@ private static PushableGeoDistance from(FoldContext ctx, Attribute attr, Express
132132

133133
record PushableCompoundExec(EvalExec evalExec, EsQueryExec queryExec, List<EsQueryExec.Sort> pushableSorts) implements Pushable {
134134
public PhysicalPlan rewrite(TopNExec topNExec) {
135-
// List<Alias> evalExecAlias = evalExec.fields();
136-
// List<FieldAttribute> additionalQueryAttrs = new ArrayList<>();
137-
// boolean aliasChanged = false;
138-
// for (EsQueryExec.Sort pushableSort : pushableSorts) {
139-
// if (pushableSort instanceof EsQueryExec.ScriptSort scriptSort) {
140-
// // Change eval alias to the script sort field
141-
// for(int i = 0; i < evalExecAlias.size(); i++) {
142-
// Alias alias = evalExecAlias.get(i);
143-
// if (alias.id().equals(scriptSort.alias().id())) {
144-
// evalExecAlias.set(i, alias.replaceChild(scriptSort.field()));
145-
// additionalQueryAttrs.add(scriptSort.field());
146-
// aliasChanged = true;
147-
// break;
148-
// }
149-
// }
150-
// }
151-
// }
152-
// if (aliasChanged) {
153-
// EsQueryExec newQueryExec = queryExec.withSorts(pushableSorts).withLimit(topNExec.limit());
154-
//// newQueryExec.attrs().addAll(additionalQueryAttrs);
155-
// EvalExec newEvalExec = new EvalExec(evalExec.source(), evalExec.child(), evalExecAlias);
156-
// }
135+
List<Alias> evalExecAlias = evalExec.fields();
136+
List<FieldAttribute> additionalQueryAttrs = new ArrayList<>();
137+
boolean aliasChanged = false;
138+
for (EsQueryExec.Sort pushableSort : pushableSorts) {
139+
if (pushableSort instanceof EsQueryExec.ScriptSort scriptSort) {
140+
// Change eval alias to the script sort field
141+
for(int i = 0; i < evalExecAlias.size(); i++) {
142+
Alias alias = evalExecAlias.get(i);
143+
if (alias.id().equals(scriptSort.alias().id())) {
144+
evalExecAlias.set(i, alias.replaceChild(scriptSort.field()));
145+
additionalQueryAttrs.add(scriptSort.field());
146+
aliasChanged = true;
147+
break;
148+
}
149+
}
150+
}
151+
}
152+
if (aliasChanged) {
153+
EsQueryExec newQueryExec = queryExec.withSorts(pushableSorts).withLimit(topNExec.limit());
154+
// newQueryExec.attrs().addAll(additionalQueryAttrs);
155+
EvalExec newEvalExec = new EvalExec(evalExec.source(), evalExec.child(), evalExecAlias);
156+
}
157157

158158
// We need to keep the EVAL in place because the coordinator will have its own TopNExec so we need to keep the distance
159159
return evalExec.replaceChild(queryExec.withSorts(pushableSorts).withLimit(topNExec.limit()));

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import java.util.List;
3434
import java.util.Map;
3535
import java.util.Objects;
36+
import java.util.stream.Stream;
3637

3738
public class EsQueryExec extends LeafExec implements EstimatesRowSize {
3839
public static final EsField DOC_ID_FIELD = new EsField(
@@ -311,7 +312,11 @@ public QueryBuilder query() {
311312

312313
@Override
313314
public List<Attribute> output() {
314-
return attrs;
315+
if (sorts == null || sorts.isEmpty()) {
316+
return attrs;
317+
}
318+
// Concat attributes with the fields used in script sorts, as they need to be extracted too
319+
return Stream.concat(attrs.stream(), sorts.stream().filter(s -> s instanceof ScriptSort).map(Sort::field)).toList();
315320
}
316321

317322
public Expression limit() {

0 commit comments

Comments
 (0)