honeymoose
/
OpenSearch
mirror of https://github.com/honeymoose/OpenSearch.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

SQL: Remove instanceof checks for field retrieval (elastic/x-pack-elasticsearch#3402) 

* SQL: Remove instanceof checks for field retrieval

This removes the `instanceof` checks iterating through the columns determining
what fields need to be retrieved from the ES document. It adds an interface
`FieldExtraction` that collects the fields needed in a builder. The builder can
then build the fields necessary from a `SearchSourceBuilder`.

* Remove default implementation in favor of pushing down exception throwing

Original commit: elastic/x-pack-elasticsearch@11d3d69eb1
This commit is contained in:
Lee Hinman 2017-12-21 17:11:31 -07:00 committed by GitHub
parent a288dde22f
commit 97eef004de
12 changed files with 217 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
sql/server/src/main/java/org/elasticsearch/xpack/sql/execution/search/FieldExtraction.java Normal file
View File

@ -0,0 +1,22 @@
/*
* 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.sql.execution.search;
import org.elasticsearch.search.builder.SearchSourceBuilder;
/**
* An interface for something that needs to extract field(s) from a result.
*/
public interface FieldExtraction {
/**
* Add whatever is necessary to the {@link SearchSourceBuilder}
* in order to fetch the field. This can include tracking the score,
* {@code _source} fields, doc values fields, and script fields.
*/
void collectFields(SqlSourceBuilder sourceBuilder);
}

67
sql/server/src/main/java/org/elasticsearch/xpack/sql/execution/search/SourceGenerator.java
View File

@ -61,40 +61,25 @@ public abstract class SourceGenerator {
private static final List<String> NO_STORED_FIELD = singletonList(StoredFieldsContext._NONE_);
public static SearchSourceBuilder sourceBuilder(QueryContainer container, QueryBuilder filter, Integer size) {
SearchSourceBuilder source = new SearchSourceBuilder();
final SearchSourceBuilder source = new SearchSourceBuilder();
// add the source
if (container.query() != null) {
if (filter != null) {
source.query(new BoolQueryBuilder().must(container.query().asBuilder()).filter(filter));
} else {
source.query(container.query().asBuilder());
}
} else {
if (container.query() == null) {
if (filter != null) {
source.query(new ConstantScoreQueryBuilder(filter));
}
} else {
if (filter == null) {
source.query(container.query().asBuilder());
} else {
source.query(new BoolQueryBuilder().must(container.query().asBuilder()).filter(filter));
}
}
// translate fields to source-fields or script fields
Set<String> sourceFields = new LinkedHashSet<>();
Set<String> docFields = new LinkedHashSet<>();
Map<String, Script> scriptFields = new LinkedHashMap<>();
for (ColumnReference ref : container.columns()) {
collectFields(source, ref, sourceFields, docFields, scriptFields);
}
if (!sourceFields.isEmpty()) {
source.fetchSource(sourceFields.toArray(new String[sourceFields.size()]), null);
}
for (String field : docFields) {
source.docValueField(field);
}
for (Entry<String, Script> entry : scriptFields.entrySet()) {
source.scriptField(entry.getKey(), entry.getValue());
}
SqlSourceBuilder sortBuilder = new SqlSourceBuilder();
// Iterate through all the columns requested, collecting the fields that
// need to be retrieved from the result documents
container.columns().forEach(cr -> cr.collectFields(sortBuilder));
sortBuilder.build(source);
// add the aggs
Aggs aggs = container.aggs();
@ -137,32 +122,6 @@ public abstract class SourceGenerator {
return source;
}
private static void collectFields(SearchSourceBuilder source, ColumnReference ref,
Set<String> sourceFields, Set<String> docFields, Map<String, Script> scriptFields) {
if (ref instanceof ComputedRef) {
ProcessorDefinition proc = ((ComputedRef) ref).processor();
if (proc instanceof ScoreProcessorDefinition) {
/*
* If we're SELECTing SCORE then force tracking scores just in case
* we're not sorting on them.
*/
source.trackScores(true);
}
proc.forEachUp(l -> collectFields(source, l.context(), sourceFields, docFields, scriptFields), ReferenceInput.class);
} else if (ref instanceof SearchHitFieldRef) {
SearchHitFieldRef sh = (SearchHitFieldRef) ref;
Set<String> collection = sh.useDocValue() ? docFields : sourceFields;
collection.add(sh.name());
} else if (ref instanceof ScriptFieldRef) {
ScriptFieldRef sfr = (ScriptFieldRef) ref;
scriptFields.put(sfr.name(), sfr.script().toPainless());
} else if (ref instanceof AggRef) {
// Nothing to do
} else {
throw new IllegalStateException("unhandled field in collectFields [" + ref.getClass() + "][" + ref + "]");
}
}
private static void sorting(QueryContainer container, SearchSourceBuilder source) {
if (source.aggregations() != null && source.aggregations().count() > 0) {
// Aggs can't be sorted using search sorting. That sorting is handled elsewhere.

71
sql/server/src/main/java/org/elasticsearch/xpack/sql/execution/search/SqlSourceBuilder.java Normal file
View File

@ -0,0 +1,71 @@
/*
* 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.sql.execution.search;
import org.elasticsearch.common.Strings;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.xpack.sql.execution.search.FieldExtraction;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
/**
* A {@code SqlSourceBuilder} is a builder object passed to objects implementing
* {@link FieldExtraction} that can "build" whatever needs to be extracted from
* the resulting ES document as a field.
*/
public class SqlSourceBuilder {
final Set<String> sourceFields = new HashSet<>();
final Set<String> docFields = new HashSet<>();
final Map<String, Script> scriptFields = new HashMap<>();
boolean trackScores = false;
public SqlSourceBuilder() {
}
/**
* Turns on returning the {@code _score} for documents.
*/
public void trackScores() {
this.trackScores = true;
}
/**
* Retrieve the requested field from the {@code _source} of the document
*/
public void addSourceField(String field) {
sourceFields.add(field);
}
/**
* Retrieve the requested field from doc values (or fielddata) of the document
*/
public void addDocField(String field) {
docFields.add(field);
}
/**
* Return the given field as a script field with the supplied script
*/
public void addScriptField(String name, Script script) {
scriptFields.put(name, script);
}
/**
* Collect the necessary fields, modifying the {@code SearchSourceBuilder}
* to retrieve them from the document.
*/
public void build(SearchSourceBuilder sourceBuilder) {
sourceBuilder.trackScores(this.trackScores);
sourceBuilder.fetchSource(sourceFields.toArray(Strings.EMPTY_ARRAY), null);
docFields.forEach(dvf -> sourceBuilder.docValueField(dvf));
scriptFields.forEach((k, v) -> sourceBuilder.scriptField(k, v));
}
}

9
sql/server/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/processor/definition/ProcessorDefinition.java
View File

@ -5,13 +5,15 @@
*/
package org.elasticsearch.xpack.sql.expression.function.scalar.processor.definition;
import org.elasticsearch.xpack.sql.execution.search.FieldExtraction;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
import org.elasticsearch.xpack.sql.expression.Expression;
import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.Processor;
import org.elasticsearch.xpack.sql.tree.Node;
import java.util.List;
public abstract class ProcessorDefinition extends Node<ProcessorDefinition> {
public abstract class ProcessorDefinition extends Node<ProcessorDefinition> implements FieldExtraction {
private final Expression expression;
@ -27,4 +29,9 @@ public abstract class ProcessorDefinition extends Node<ProcessorDefinition> {
public abstract boolean resolved();
public abstract Processor asProcessor();
@Override
public void collectFields(SqlSourceBuilder sourceBuilder) {
// No fields needed
}
}

8
sql/server/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/processor/definition/ScoreProcessorDefinition.java
View File

@ -5,6 +5,7 @@
*/
package org.elasticsearch.xpack.sql.expression.function.scalar.processor.definition;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
import org.elasticsearch.xpack.sql.execution.search.extractor.ScoreExtractor;
import org.elasticsearch.xpack.sql.expression.Expression;
import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.Processor;
@ -13,6 +14,7 @@ import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.
import static java.util.Collections.emptyList;
public class ScoreProcessorDefinition extends ProcessorDefinition {
public ScoreProcessorDefinition(Expression expression) {
super(expression, emptyList());
}
@ -26,4 +28,10 @@ public class ScoreProcessorDefinition extends ProcessorDefinition {
public Processor asProcessor() {
return new HitExtractorProcessor(ScoreExtractor.INSTANCE);
}
@Override
public void collectFields(SqlSourceBuilder sourceBuilder) {
sourceBuilder.trackScores();
}
}

6
sql/server/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/AggRef.java
View File

@ -5,6 +5,7 @@
*/
package org.elasticsearch.xpack.sql.querydsl.container;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
import org.elasticsearch.xpack.sql.querydsl.agg.AggPath;
public class AggRef implements ColumnReference {
@ -29,4 +30,9 @@ public class AggRef implements ColumnReference {
public String path() {
return path;
}
@Override
public void collectFields(SqlSourceBuilder sourceBuilder) {
// Aggregations do not need any special fields
}
}

9
sql/server/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/ColumnReference.java
View File

@ -5,10 +5,15 @@
*/
package org.elasticsearch.xpack.sql.querydsl.container;
import org.elasticsearch.xpack.sql.execution.search.FieldExtraction;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
/**
* Entity representing a 'column' backed by one or multiple results from ES.
* Entity representing a 'column' backed by one or multiple results from ES. A
* column reference can also extract a field (meta or otherwise) from a result
* set, so extends {@link FieldExtraction}.
*/
public interface ColumnReference {
public interface ColumnReference extends FieldExtraction {
/**
* Indicates the depth of the result. Used for counting the actual size of a

9
sql/server/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/ComputedRef.java
View File

@ -5,8 +5,11 @@
*/
package org.elasticsearch.xpack.sql.querydsl.container;
import org.elasticsearch.xpack.sql.execution.search.FieldExtraction;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
import org.elasticsearch.xpack.sql.expression.function.scalar.processor.definition.ProcessorDefinition;
import org.elasticsearch.xpack.sql.expression.function.scalar.processor.definition.ReferenceInput;
import org.elasticsearch.xpack.sql.expression.function.scalar.processor.definition.ScoreProcessorDefinition;
import java.util.concurrent.atomic.AtomicInteger;
@ -39,6 +42,12 @@ public class ComputedRef implements ColumnReference {
return depth;
}
@Override
public void collectFields(SqlSourceBuilder sourceBuilder) {
processor.collectFields(sourceBuilder);
processor.forEachUp(ri -> ri.context().collectFields(sourceBuilder), ReferenceInput.class);
}
@Override
public String toString() {
return processor + "(" + processor + ")";

10
sql/server/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/NestedFieldRef.java
View File

@ -5,6 +5,9 @@
*/
package org.elasticsearch.xpack.sql.querydsl.container;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
public class NestedFieldRef implements FieldReference {
private final String parent, name;
private final boolean docValue;
@ -28,8 +31,13 @@ public class NestedFieldRef implements FieldReference {
return docValue;
}
@Override
public void collectFields(SqlSourceBuilder sourceBuilder) {
throw new IllegalStateException("unhandled nested field while collecting source fields [" + getClass() + "]");
}
@Override
public String toString() {
return name;
}
}
}

6
sql/server/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/ScriptFieldRef.java
View File

@ -5,6 +5,7 @@
*/
package org.elasticsearch.xpack.sql.querydsl.container;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
import org.elasticsearch.xpack.sql.expression.function.scalar.script.ScriptTemplate;
public class ScriptFieldRef implements FieldReference {
@ -26,6 +27,11 @@ public class ScriptFieldRef implements FieldReference {
return script;
}
@Override
public void collectFields(SqlSourceBuilder sourceBuilder) {
sourceBuilder.addScriptField(name, script.toPainless());
}
@Override
public String toString() {
return "{" + name + "}";

13
sql/server/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/SearchHitFieldRef.java
View File

@ -5,6 +5,8 @@
*/
package org.elasticsearch.xpack.sql.querydsl.container;
import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
public class SearchHitFieldRef implements FieldReference {
private final String name;
private final boolean docValue;
@ -23,8 +25,17 @@ public class SearchHitFieldRef implements FieldReference {
return docValue;
}
@Override
public void collectFields(SqlSourceBuilder sourceBuilder) {
if (docValue) {
sourceBuilder.addDocField(name);
} else {
sourceBuilder.addSourceField(name);
}
}
@Override
public String toString() {
return name;
}
}
}

46
sql/server/src/test/java/org/elasticsearch/xpack/sql/execution/search/SqlSourceBuilderTests.java Normal file
View File

@ -0,0 +1,46 @@
/*
* 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.sql.execution.search;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.test.ESTestCase;
import java.util.Arrays;
import java.util.Map;
import java.util.stream.Collectors;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
public class SqlSourceBuilderTests extends ESTestCase {
public void testSqlSourceBuilder() {
final SqlSourceBuilder ssb = new SqlSourceBuilder();
final SearchSourceBuilder source = new SearchSourceBuilder();
ssb.trackScores();
ssb.addSourceField("foo");
ssb.addSourceField("foo2");
ssb.addDocField("bar");
ssb.addDocField("bar2");
final Script s = new Script("eggplant");
ssb.addScriptField("baz", s);
final Script s2 = new Script("potato");
ssb.addScriptField("baz2", s2);
ssb.build(source);
assertTrue(source.trackScores());
FetchSourceContext fsc = source.fetchSource();
assertThat(Arrays.asList(fsc.includes()), containsInAnyOrder("foo", "foo2"));
assertThat(source.docValueFields(), containsInAnyOrder("bar", "bar2"));
Map<String, Script> scriptFields = source.scriptFields()
.stream()
.collect(Collectors.toMap(SearchSourceBuilder.ScriptField::fieldName, SearchSourceBuilder.ScriptField::script));
assertThat(scriptFields.get("baz").getIdOrCode(), equalTo("eggplant"));
assertThat(scriptFields.get("baz2").getIdOrCode(), equalTo("potato"));
}
}