Skip to content

Commit

Permalink
Use a NestedScope in the KqlParsingContext to track nested queries.
Browse files Browse the repository at this point in the history
  • Loading branch information
@afoucret
afoucret committed Nov 14, 2024
1 parent cfde9fe commit cf1023f
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 29 deletions.
33 changes: 10 additions & 23 deletions x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/parser/KqlAstBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Supplier;

import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
import static org.elasticsearch.xpack.kql.parser.KqlParsingContext.isDateField;
Expand All @@ -40,8 +39,6 @@
class KqlAstBuilder extends KqlBaseBaseVisitor<QueryBuilder> {
private final KqlParsingContext kqlParsingContext;

private volatile String currentNestedPath;

KqlAstBuilder(KqlParsingContext kqlParsingContext) {
this.kqlParsingContext = kqlParsingContext;
}
Expand Down Expand Up @@ -107,10 +104,6 @@ public QueryBuilder visitParenthesizedQuery(KqlBaseParser.ParenthesizedQueryCont
public QueryBuilder visitNestedQuery(KqlBaseParser.NestedQueryContext ctx) {
String nestedFieldName = extractText(ctx.fieldName());

if (currentNestedPath != null) {
nestedFieldName = currentNestedPath + "." + nestedFieldName;
}

if (kqlParsingContext.isNestedField(nestedFieldName) == false) {
throw new KqlParsingException(
"[{}] is not a valid nested field name.",
Expand All @@ -119,13 +112,19 @@ public QueryBuilder visitNestedQuery(KqlBaseParser.NestedQueryContext ctx) {
nestedFieldName
);
}
QueryBuilder subQuery = withNestedPath(nestedFieldName, () -> typedParsing(this, ctx.nestedSubQuery(), QueryBuilder.class));
QueryBuilder subQuery = kqlParsingContext.withNestedPath(
nestedFieldName,
() -> typedParsing(this, ctx.nestedSubQuery(), QueryBuilder.class)
);

if (subQuery instanceof MatchNoneQueryBuilder) {
return subQuery;
}

return wrapWithNestedQuery(nestedFieldName, QueryBuilders.nestedQuery(nestedFieldName, subQuery, ScoreMode.None));
return wrapWithNestedQuery(
nestedFieldName,
QueryBuilders.nestedQuery(kqlParsingContext.fullFieldName(nestedFieldName), subQuery, ScoreMode.None)
);
}

@Override
Expand Down Expand Up @@ -258,22 +257,10 @@ private static boolean isOrQuery(ParserRuleContext ctx) {
};
}

private QueryBuilder withNestedPath(String nestedPath, Supplier<QueryBuilder> queryBuilderSupplier) {
String previousNestedPath = this.currentNestedPath;
this.currentNestedPath = nestedPath;
QueryBuilder queryBuilder = queryBuilderSupplier.get();
this.currentNestedPath = previousNestedPath;
return queryBuilder;
}

private void withFields(KqlBaseParser.FieldNameContext ctx, BiConsumer<String, MappedFieldType> fieldConsummer) {
assert ctx != null : "Field ctx cannot be null";
String fieldNamePattern = extractText(ctx);

if (currentNestedPath != null) {
fieldNamePattern = currentNestedPath + "." + fieldNamePattern;
}

Set<String> fieldNames = kqlParsingContext.resolveFieldNames(fieldNamePattern);

if (ctx.value.getType() == KqlBaseParser.QUOTED_STRING && Regex.isSimpleMatchPattern(fieldNamePattern)) {
Expand Down Expand Up @@ -325,9 +312,9 @@ private BiFunction<RangeQueryBuilder, String, RangeQueryBuilder> rangeOperation(
}

private QueryBuilder wrapWithNestedQuery(String fieldName, QueryBuilder query) {
String nestedPath = kqlParsingContext.nestedParent(fieldName);
String nestedPath = kqlParsingContext.nestedPath(fieldName);

if (nestedPath == null || nestedPath.equals(currentNestedPath)) {
if (nestedPath == null || nestedPath.equals(kqlParsingContext.currentNestedPath())) {
return query;
}

Expand Down
49 changes: 43 additions & 6 deletions x-pack/plugin/kql/src/main/java/org/elasticsearch/xpack/kql/parser/KqlParsingContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,18 @@
import org.elasticsearch.index.mapper.DateFieldMapper;
import org.elasticsearch.index.mapper.KeywordFieldMapper;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.index.mapper.ObjectMapper;
import org.elasticsearch.index.mapper.NestedLookup;
import org.elasticsearch.index.mapper.NestedObjectMapper;
import org.elasticsearch.index.query.QueryRewriteContext;
import org.elasticsearch.index.query.support.NestedScope;

import java.time.ZoneId;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;

import static org.elasticsearch.common.Strings.format;

public class KqlParsingContext {

Expand All @@ -38,6 +43,7 @@ public static Builder builder(QueryRewriteContext queryRewriteContext) {
private final boolean caseInsensitive;
private final ZoneId timeZone;
private final String defaultField;
private final NestedScope nestedScope = new NestedScope();

public KqlParsingContext(QueryRewriteContext queryRewriteContext, boolean caseInsensitive, ZoneId timeZone, String defaultField) {
this.queryRewriteContext = queryRewriteContext;
Expand All @@ -58,18 +64,17 @@ public String defaultField() {
return defaultField;
}

public String nestedParent(String fieldName) {
return queryRewriteContext.getMappingLookup().nestedLookup().getNestedParent(fieldName);
public String nestedPath(String fieldName) {
return nestedLookup().getNestedParent(fieldName);
}

public boolean isNestedField(String fieldName) {
Map<String, ObjectMapper> objectMappers = this.queryRewriteContext.getMappingLookup().objectMappers();
return objectMappers.containsKey(fieldName) && objectMappers.get(fieldName).isNested();
return nestedMappers().containsKey(fullFieldName(fieldName));
}

public Set<String> resolveFieldNames(String fieldNamePattern) {
assert fieldNamePattern != null && fieldNamePattern.isEmpty() == false : "fieldNamePattern cannot be null or empty";
return queryRewriteContext.getMatchingFieldNames(fieldNamePattern);
return queryRewriteContext.getMatchingFieldNames(fullFieldName(fieldNamePattern));
}

public Set<String> resolveDefaultFieldNames() {
Expand Down Expand Up @@ -100,6 +105,38 @@ public boolean isSearchableField(String fieldName) {
return isSearchableField(fieldName, fieldType(fieldName));
}

public NestedScope nestedScope() {
return nestedScope;
}

public <T> T withNestedPath(String nestedFieldName, Supplier<T> supplier) {
assert isNestedField(nestedFieldName);
nestedScope.nextLevel(nestedMappers().get(fullFieldName(nestedFieldName)));
T result = supplier.get();
nestedScope.previousLevel();
return result;
}

public String currentNestedPath() {
return nestedScope().getObjectMapper() != null ? nestedScope().getObjectMapper().fullPath() : null;
}

public String fullFieldName(String fieldName) {
if (nestedScope.getObjectMapper() == null) {
return fieldName;
}

return format("%s.%s", nestedScope.getObjectMapper().fullPath(), fieldName);
}

private NestedLookup nestedLookup() {
return queryRewriteContext.getMappingLookup().nestedLookup();
}

private Map<String, NestedObjectMapper> nestedMappers() {
return nestedLookup().getNestedMappers();
}

public static class Builder {
private final QueryRewriteContext queryRewriteContext;
private boolean caseInsensitive = true;
Expand Down

0 comments on commit cf1023f

Please sign in to comment.