Skip to content

Commit 61a479b

Browse files
carlosdelestcarlosdelest
committed
First version - adding toQuery() to ShardContext
1 parent 0a98bf8 commit 61a479b

File tree

17 files changed

+183
-45
lines changed

17 files changed

+183
-45
lines changed

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

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,15 @@
2525
import org.elasticsearch.compute.data.DocVector;
2626
import org.elasticsearch.compute.data.IntVector;
2727
import org.elasticsearch.compute.data.Page;
28+
import org.elasticsearch.compute.operator.DriverContext;
2829
import org.elasticsearch.compute.operator.EvalOperator;
2930
import org.elasticsearch.core.Releasable;
3031
import org.elasticsearch.core.Releasables;
32+
import org.elasticsearch.index.query.QueryBuilder;
3133

3234
import java.io.IOException;
3335
import java.io.UncheckedIOException;
36+
import java.util.List;
3437

3538
/**
3639
* {@link EvalOperator.ExpressionEvaluator} to run a Lucene {@link Query} during
@@ -44,19 +47,35 @@ public record ShardConfig(Query query, IndexSearcher searcher) {}
4447

4548
private final BlockFactory blockFactory;
4649
private final ShardConfig[] shards;
47-
private final int docChannel;
4850

4951
private ShardState[] perShardState = EMPTY_SHARD_STATES;
5052

51-
public LuceneQueryExpressionEvaluator(BlockFactory blockFactory, ShardConfig[] shards, int docChannel) {
53+
public LuceneQueryExpressionEvaluator(
54+
BlockFactory blockFactory,
55+
List<? extends ShardContext> shardContexts,
56+
QueryBuilder queryBuilder
57+
) {
58+
assert shardContexts.isEmpty() == false;
59+
assert queryBuilder != null;
60+
61+
this.blockFactory = blockFactory;
62+
this.shards = new ShardConfig[shardContexts.size()];
63+
64+
int i = 0;
65+
for (ShardContext shardContext : shardContexts) {
66+
this.shards[i++] = new ShardConfig(shardContext.toQuery(queryBuilder), shardContext.searcher());
67+
}
68+
}
69+
70+
public LuceneQueryExpressionEvaluator(BlockFactory blockFactory, ShardConfig[] shards) {
5271
this.blockFactory = blockFactory;
5372
this.shards = shards;
54-
this.docChannel = docChannel;
5573
}
5674

5775
@Override
5876
public Block eval(Page page) {
59-
DocVector docs = page.<DocBlock>getBlock(docChannel).asVector();
77+
// Lucene based operators retrieve DocVectors as first block
78+
DocVector docs = page.<DocBlock>getBlock(0).asVector();
6079
try {
6180
if (docs.singleSegmentNonDecreasing()) {
6281
return evalSingleSegmentNonDecreasing(docs).asBlock();
@@ -341,4 +360,19 @@ public void close() {
341360
Releasables.closeExpectNoException(builder);
342361
}
343362
}
363+
364+
public static class Factory implements EvalOperator.ExpressionEvaluator.Factory {
365+
private final List<? extends ShardContext> shardContexts;
366+
private final QueryBuilder queryBuilder;
367+
368+
public Factory(List<? extends ShardContext> shardContexts, QueryBuilder queryBuilder) {
369+
this.shardContexts = shardContexts;
370+
this.queryBuilder = queryBuilder;
371+
}
372+
373+
@Override
374+
public EvalOperator.ExpressionEvaluator get(DriverContext context) {
375+
return new LuceneQueryExpressionEvaluator(context.blockFactory(), shardContexts, queryBuilder);
376+
}
377+
}
344378
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
package org.elasticsearch.compute.lucene;
99

1010
import org.apache.lucene.search.IndexSearcher;
11+
import org.apache.lucene.search.Query;
12+
import org.elasticsearch.index.query.QueryBuilder;
1113
import org.elasticsearch.search.sort.SortAndFormats;
1214
import org.elasticsearch.search.sort.SortBuilder;
1315

@@ -39,4 +41,9 @@ public interface ShardContext {
3941
* {@code _cat/shards}.
4042
*/
4143
String shardIdentifier();
44+
45+
/**
46+
* Converts a {@link QueryBuilder} to a {@link Query} in the shard
47+
*/
48+
Query toQuery(QueryBuilder queryBuilder);
4249
}

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneQueryExpressionEvaluatorTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -183,8 +183,8 @@ private List<Page> runQuery(Set<String> values, Query query, boolean shuffleDocs
183183
);
184184
LuceneQueryExpressionEvaluator luceneQueryEvaluator = new LuceneQueryExpressionEvaluator(
185185
blockFactory,
186-
new LuceneQueryExpressionEvaluator.ShardConfig[] { shard },
187-
0
186+
new LuceneQueryExpressionEvaluator.ShardConfig[] { shard }
187+
188188
);
189189

190190
List<Operator> operators = new ArrayList<>();

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,7 @@ public MockShardContext(IndexReader reader, int index) {
211211
TrivialQueryCachingPolicy.NEVER,
212212
true
213213
);
214+
214215
} catch (IOException e) {
215216
throw new AssertionError(e);
216217
}

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,21 @@ public void testMatchWithinEval() {
269269
assertThat(error.getMessage(), containsString("[MATCH] function is only supported in WHERE commands"));
270270
}
271271

272+
public void testMatchNonPushedDown() {
273+
var query = """
274+
FROM test
275+
| WHERE match(content, "fox") OR length(content) < 20
276+
| KEEP id
277+
| SORT id
278+
""";
279+
280+
try (var resp = run(query)) {
281+
assertColumnNames(resp.columns(), List.of("id"));
282+
assertColumnTypes(resp.columns(), List.of("integer"));
283+
assertValues(resp.values(), List.of(List.of(1), List.of(2), List.of(6)));
284+
}
285+
}
286+
272287
private void createAndPopulateIndex() {
273288
var indexName = "test";
274289
var client = client().admin().indices();

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/EvalMapper.java‎

Lines changed: 50 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.compute.data.ElementType;
1717
import org.elasticsearch.compute.data.Page;
1818
import org.elasticsearch.compute.data.Vector;
19+
import org.elasticsearch.compute.lucene.ShardContext;
1920
import org.elasticsearch.compute.operator.DriverContext;
2021
import org.elasticsearch.compute.operator.EvalOperator;
2122
import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
@@ -31,6 +32,7 @@
3132
import org.elasticsearch.xpack.esql.core.expression.predicate.nulls.IsNull;
3233
import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;
3334
import org.elasticsearch.xpack.esql.evaluator.mapper.ExpressionMapper;
35+
import org.elasticsearch.xpack.esql.expression.function.fulltext.Match;
3436
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.InsensitiveEqualsMapper;
3537
import org.elasticsearch.xpack.esql.planner.Layout;
3638

@@ -50,13 +52,24 @@ public final class EvalMapper {
5052

5153
private EvalMapper() {}
5254

53-
@SuppressWarnings({ "rawtypes", "unchecked" })
5455
public static ExpressionEvaluator.Factory toEvaluator(FoldContext foldCtx, Expression exp, Layout layout) {
55-
if (exp instanceof EvaluatorMapper m) {
56+
return toEvaluator(foldCtx, exp, layout, List.of());
57+
}
58+
59+
@SuppressWarnings({ "rawtypes", "unchecked" })
60+
public static ExpressionEvaluator.Factory toEvaluator(
61+
FoldContext foldCtx,
62+
Expression exp,
63+
Layout layout,
64+
List<? extends ShardContext> shardContexts
65+
) {
66+
if (exp instanceof Match m) {
67+
return m.toEvaluator(shardContexts, m.queryBuilder());
68+
} else if (exp instanceof EvaluatorMapper m) {
5669
return m.toEvaluator(new EvaluatorMapper.ToEvaluator() {
5770
@Override
5871
public ExpressionEvaluator.Factory apply(Expression expression) {
59-
return toEvaluator(foldCtx, expression, layout);
72+
return toEvaluator(foldCtx, expression, layout, shardContexts);
6073
}
6174

6275
@Override
@@ -67,17 +80,22 @@ public FoldContext foldCtx() {
6780
}
6881
for (ExpressionMapper em : MAPPERS) {
6982
if (em.typeToken.isInstance(exp)) {
70-
return em.map(foldCtx, exp, layout);
83+
return em.map(foldCtx, exp, layout, shardContexts);
7184
}
7285
}
7386
throw new QlIllegalArgumentException("Unsupported expression [{}]", exp);
7487
}
7588

7689
static class BooleanLogic extends ExpressionMapper<BinaryLogic> {
7790
@Override
78-
public ExpressionEvaluator.Factory map(FoldContext foldCtx, BinaryLogic bc, Layout layout) {
79-
var leftEval = toEvaluator(foldCtx, bc.left(), layout);
80-
var rightEval = toEvaluator(foldCtx, bc.right(), layout);
91+
public ExpressionEvaluator.Factory map(
92+
FoldContext foldCtx,
93+
BinaryLogic bc,
94+
Layout layout,
95+
List<? extends ShardContext> shardContexts
96+
) {
97+
var leftEval = toEvaluator(foldCtx, bc.left(), layout, shardContexts);
98+
var rightEval = toEvaluator(foldCtx, bc.right(), layout, shardContexts);
8199
/**
82100
* Evaluator for the <href a="https://en.wikipedia.org/wiki/Three-valued_logic">three-valued boolean expressions</href>.
83101
* We can't generate these with the {@link Evaluator} annotation because that
@@ -153,7 +171,7 @@ public void close() {
153171

154172
static class Nots extends ExpressionMapper<Not> {
155173
@Override
156-
public ExpressionEvaluator.Factory map(FoldContext foldCtx, Not not, Layout layout) {
174+
public ExpressionEvaluator.Factory map(FoldContext foldCtx, Not not, Layout layout, List<? extends ShardContext> shardContexts) {
157175
var expEval = toEvaluator(foldCtx, not.field(), layout);
158176
return dvrCtx -> new org.elasticsearch.xpack.esql.evaluator.predicate.operator.logical.NotEvaluator(
159177
not.source(),
@@ -165,7 +183,12 @@ public ExpressionEvaluator.Factory map(FoldContext foldCtx, Not not, Layout layo
165183

166184
static class Attributes extends ExpressionMapper<Attribute> {
167185
@Override
168-
public ExpressionEvaluator.Factory map(FoldContext foldCtx, Attribute attr, Layout layout) {
186+
public ExpressionEvaluator.Factory map(
187+
FoldContext foldCtx,
188+
Attribute attr,
189+
Layout layout,
190+
List<? extends ShardContext> shardContexts
191+
) {
169192
record Attribute(int channel) implements ExpressionEvaluator {
170193
@Override
171194
public Block eval(Page page) {
@@ -200,7 +223,12 @@ public boolean eagerEvalSafeInLazy() {
200223
static class Literals extends ExpressionMapper<Literal> {
201224

202225
@Override
203-
public ExpressionEvaluator.Factory map(FoldContext foldCtx, Literal lit, Layout layout) {
226+
public ExpressionEvaluator.Factory map(
227+
FoldContext foldCtx,
228+
Literal lit,
229+
Layout layout,
230+
List<? extends ShardContext> shardContexts
231+
) {
204232
record LiteralsEvaluator(DriverContext context, Literal lit) implements ExpressionEvaluator {
205233
@Override
206234
public Block eval(Page page) {
@@ -257,7 +285,12 @@ private static Block block(Literal lit, BlockFactory blockFactory, int positions
257285
static class IsNulls extends ExpressionMapper<IsNull> {
258286

259287
@Override
260-
public ExpressionEvaluator.Factory map(FoldContext foldCtx, IsNull isNull, Layout layout) {
288+
public ExpressionEvaluator.Factory map(
289+
FoldContext foldCtx,
290+
IsNull isNull,
291+
Layout layout,
292+
List<? extends ShardContext> shardContexts
293+
) {
261294
var field = toEvaluator(foldCtx, isNull.field(), layout);
262295
return new IsNullEvaluatorFactory(field);
263296
}
@@ -305,7 +338,12 @@ public String toString() {
305338
static class IsNotNulls extends ExpressionMapper<IsNotNull> {
306339

307340
@Override
308-
public ExpressionEvaluator.Factory map(FoldContext foldCtx, IsNotNull isNotNull, Layout layout) {
341+
public ExpressionEvaluator.Factory map(
342+
FoldContext foldCtx,
343+
IsNotNull isNotNull,
344+
Layout layout,
345+
List<? extends ShardContext> shardContexts
346+
) {
309347
return new IsNotNullEvaluatorFactory(toEvaluator(foldCtx, isNotNull.field(), layout));
310348
}
311349

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/mapper/ExpressionMapper.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,26 @@
77

88
package org.elasticsearch.xpack.esql.evaluator.mapper;
99

10+
import org.elasticsearch.compute.lucene.ShardContext;
1011
import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
1112
import org.elasticsearch.xpack.esql.core.expression.Expression;
1213
import org.elasticsearch.xpack.esql.core.expression.FoldContext;
1314
import org.elasticsearch.xpack.esql.core.util.ReflectionUtils;
1415
import org.elasticsearch.xpack.esql.planner.Layout;
1516

17+
import java.util.List;
18+
1619
public abstract class ExpressionMapper<E extends Expression> {
1720
public final Class<E> typeToken;
1821

1922
public ExpressionMapper() {
2023
typeToken = ReflectionUtils.detectSuperTypeForRuleLike(getClass());
2124
}
2225

23-
public abstract ExpressionEvaluator.Factory map(FoldContext foldCtx, E expression, Layout layout);
26+
public abstract ExpressionEvaluator.Factory map(
27+
FoldContext foldCtx,
28+
E expression,
29+
Layout layout,
30+
List<? extends ShardContext> shardContexts
31+
);
2432
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,7 @@ private static void checkFullTextQueryFunctions(LogicalPlan plan, Failures failu
219219
m -> "[" + m.functionName() + "] " + m.functionType(),
220220
failures
221221
);
222-
checkFullTextSearchDisjunctions(condition, ftf -> "[" + ftf.functionName() + "] " + ftf.functionType(), failures);
222+
// checkFullTextSearchDisjunctions(condition, ftf -> "[" + ftf.functionName() + "] " + ftf.functionType(), failures);
223223
checkFullTextFunctionsParents(condition, failures);
224224
} else {
225225
plan.forEachExpression(FullTextFunction.class, ftf -> {

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/Match.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
1313
import org.elasticsearch.common.io.stream.StreamInput;
1414
import org.elasticsearch.common.io.stream.StreamOutput;
15+
import org.elasticsearch.compute.lucene.LuceneQueryExpressionEvaluator;
16+
import org.elasticsearch.compute.lucene.ShardContext;
1517
import org.elasticsearch.index.query.QueryBuilder;
1618
import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
1719
import org.elasticsearch.xpack.esql.common.Failure;
@@ -32,6 +34,7 @@
3234
import org.elasticsearch.xpack.esql.expression.function.scalar.convert.AbstractConvertFunction;
3335
import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
3436
import org.elasticsearch.xpack.esql.planner.EsqlExpressionTranslators;
37+
import org.elasticsearch.xpack.esql.planner.Layout;
3538
import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter;
3639

3740
import java.io.IOException;
@@ -292,4 +295,11 @@ private boolean isOperator() {
292295
}
293296
return isOperator;
294297
}
298+
299+
public LuceneQueryExpressionEvaluator.Factory toEvaluator(
300+
List<? extends ShardContext> shardContexts,
301+
QueryBuilder queryBuilder
302+
) {
303+
return new LuceneQueryExpressionEvaluator.Factory(shardContexts, queryBuilder);
304+
}
295305
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/operator/comparison/InsensitiveEqualsMapper.java‎

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,17 @@
1212
import org.apache.lucene.util.automaton.ByteRunAutomaton;
1313
import org.elasticsearch.common.TriFunction;
1414
import org.elasticsearch.common.lucene.BytesRefs;
15+
import org.elasticsearch.compute.lucene.ShardContext;
1516
import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
1617
import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
1718
import org.elasticsearch.xpack.esql.core.expression.FoldContext;
1819
import org.elasticsearch.xpack.esql.core.tree.Source;
1920
import org.elasticsearch.xpack.esql.core.type.DataType;
2021
import org.elasticsearch.xpack.esql.evaluator.mapper.ExpressionMapper;
21-
import org.elasticsearch.xpack.esql.expression.function.scalar.math.Cast;
2222
import org.elasticsearch.xpack.esql.planner.Layout;
2323

24+
import java.util.List;
25+
2426
import static org.elasticsearch.xpack.esql.evaluator.EvalMapper.toEvaluator;
2527

2628
public class InsensitiveEqualsMapper extends ExpressionMapper<InsensitiveEquals> {
@@ -29,12 +31,17 @@ public class InsensitiveEqualsMapper extends ExpressionMapper<InsensitiveEquals>
2931
InsensitiveEqualsEvaluator.Factory::new;
3032

3133
@Override
32-
public final ExpressionEvaluator.Factory map(FoldContext foldCtx, InsensitiveEquals bc, Layout layout) {
34+
public final ExpressionEvaluator.Factory map(
35+
FoldContext foldCtx,
36+
InsensitiveEquals bc,
37+
Layout layout,
38+
List<? extends ShardContext> shardContexts
39+
) {
3340
DataType leftType = bc.left().dataType();
3441
DataType rightType = bc.right().dataType();
3542

36-
var leftEval = toEvaluator(foldCtx, bc.left(), layout);
37-
var rightEval = toEvaluator(foldCtx, bc.right(), layout);
43+
var leftEval = toEvaluator(foldCtx, bc.left(), layout, shardContexts);
44+
var rightEval = toEvaluator(foldCtx, bc.right(), layout, shardContexts);
3845
if (DataType.isString(leftType)) {
3946
if (bc.right().foldable() && DataType.isString(rightType)) {
4047
BytesRef rightVal = BytesRefs.toBytesRef(bc.right().fold(FoldContext.small() /* TODO remove me */));
@@ -50,16 +57,4 @@ public final ExpressionEvaluator.Factory map(FoldContext foldCtx, InsensitiveEqu
5057
}
5158
throw new EsqlIllegalArgumentException("resolved type for [" + bc + "] but didn't implement mapping");
5259
}
53-
54-
public static ExpressionEvaluator.Factory castToEvaluator(
55-
FoldContext foldCtx,
56-
InsensitiveEquals op,
57-
Layout layout,
58-
DataType required,
59-
TriFunction<Source, ExpressionEvaluator.Factory, ExpressionEvaluator.Factory, ExpressionEvaluator.Factory> factory
60-
) {
61-
var lhs = Cast.cast(op.source(), op.left().dataType(), required, toEvaluator(foldCtx, op.left(), layout));
62-
var rhs = Cast.cast(op.source(), op.right().dataType(), required, toEvaluator(foldCtx, op.right(), layout));
63-
return factory.apply(op.source(), lhs, rhs);
64-
}
6560
}

0 commit comments

Comments
 (0)