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

sql/server/src/main/java/org/elasticsearch/xpack/sql/execution/search/FieldExtraction.java

@@ -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);
+
+}

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

@@ -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 (container.query() == null) {
-            if (filter != null) {
-                source.query(new BoolQueryBuilder().must(container.query().asBuilder()).filter(filter));
-            } else {
-                source.query(container.query().asBuilder());
-            }
-        } else {
             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
+        SqlSourceBuilder sortBuilder = new SqlSourceBuilder();
-        Set<String> sourceFields = new LinkedHashSet<>();
+        // Iterate through all the columns requested, collecting the fields that
-        Set<String> docFields = new LinkedHashSet<>();
+        // need to be retrieved from the result documents
-        Map<String, Script> scriptFields = new LinkedHashMap<>();
+        container.columns().forEach(cr -> cr.collectFields(sortBuilder));
-
+        sortBuilder.build(source);
-        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());
-        }
 
         // 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.

sql/server/src/main/java/org/elasticsearch/xpack/sql/execution/search/SqlSourceBuilder.java

@@ -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));
+    }
+}

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

@@ -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
+    }
 }

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

@@ -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();
+    }
+
 }

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

@@ -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
+    }
 }

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

@@ -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

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

@@ -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 + ")";

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

@@ -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,6 +31,11 @@ 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;

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

@@ -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 + "}";

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

@@ -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,6 +25,15 @@ 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;

sql/server/src/test/java/org/elasticsearch/xpack/sql/execution/search/SqlSourceBuilderTests.java

@@ -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"));
+    }
+}