SQL: Switch is aggs only to OO (elastic/x-pack-elasticsearch#3432)

Switches the "is this query aggs only?" question from pattern matching
on the column tree to an OO-style "ask the columns if they can be aggs
only" method.

I expect this could have been less code if I was willing to define
`supportedByAggsOnlyQuery` as `true` at the top of the
`ProcessorDefinition` object tree and override it only on nodes
`ReferenceInput` and `ScoreProcessorDefinition` but this feels dirty to
me. I tend to think of a superclass as a list of questions that all the
subclasses have to answer rather than a list of behaviors to share.
Pulling the `return true` up saves a few lines of code but breaks my
ability to reason about subclasses.


Original commit: elastic/x-pack-elasticsearch@b1338543cb

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

@@ -19,4 +19,9 @@ public interface FieldExtraction {
      */
     void collectFields(SqlSourceBuilder sourceBuilder);
 
+    /**
+     * Is this aggregation supported in an "aggregation only" query
+     * ({@code true}) or should it force a scroll query ({@code false})?
+     */
+    boolean supportedByAggsOnlyQuery();
 }

sql/server/src/main/java/org/elasticsearch/xpack/sql/expression/Expression.java

@@ -24,7 +24,7 @@ public abstract class Expression extends Node<Expression> implements Resolvable
     public static class TypeResolution {
         private final boolean failed;
         private final String message;
-        
+
         public static final TypeResolution TYPE_RESOLVED = new TypeResolution(false, StringUtils.EMPTY);
 
         public TypeResolution(String message, Object... args) {
@@ -124,4 +124,4 @@ public abstract class Expression extends Node<Expression> implements Resolvable
     public String toString() {
         return nodeName() + "[" + NodeUtils.propertiesToString(this, false) + "]";
     }
-}
+}

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

@@ -12,4 +12,9 @@ public class AggNameInput extends NonExecutableInput<String> {
     public AggNameInput(Expression expression, String context) {
         super(expression, context);
     }
+
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return true;
+    }
 }

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

@@ -43,21 +43,26 @@ public class AggPathInput extends NonExecutableInput<String> {
         return true;
     }
 
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return true;
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(context(), innerKey);
     }
-    
+
     @Override
     public boolean equals(Object obj) {
         if (this == obj) {
             return true;
         }
-        
+
         if (obj == null || getClass() != obj.getClass()) {
             return false;
         }
-        
+
         AggPathInput other = (AggPathInput) obj;
         return Objects.equals(context(), other.context())
                 && Objects.equals(innerKey, other.innerKey)

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

@@ -33,21 +33,26 @@ public class AggValueInput extends LeafInput<Supplier<Object>> {
         return new SuppliedProcessor(() -> matrixProcessor != null ? matrixProcessor.process(context().get()) : context().get());
     }
 
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return true;
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(context(), innerKey);
     }
-    
+
     @Override
     public boolean equals(Object obj) {
         if (this == obj) {
             return true;
         }
-        
+
         if (obj == null || getClass() != obj.getClass()) {
             return false;
         }
-        
+
         AggValueInput other = (AggValueInput) obj;
         return Objects.equals(context(), other.context())
                 && Objects.equals(innerKey, other.innerKey);

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

@@ -13,4 +13,9 @@ public class AttributeInput extends NonExecutableInput<Attribute> {
     public AttributeInput(Expression expression, Attribute context) {
         super(expression, context);
     }
+
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return true;
+    }
 }

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

@@ -27,6 +27,11 @@ public abstract class BinaryProcessorDefinition extends ProcessorDefinition {
         return right;
     }
 
+    @Override
+    public boolean supportedByAggsOnlyQuery() {
+        return left.supportedByAggsOnlyQuery() && right.supportedByAggsOnlyQuery();
+    }
+
     @Override
     public boolean resolved() {
         return left().resolved() && right().resolved();

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

@@ -19,4 +19,9 @@ public class ConstantInput extends LeafInput<Object> {
     public Processor asProcessor() {
         return new ConstantProcessor(context());
     }
+
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return true;
+    }
 }

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

@@ -20,4 +20,9 @@ public class HitExtractorInput extends LeafInput<HitExtractor> {
     public Processor asProcessor() {
         return new HitExtractorProcessor(context());
     }
+
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return true;
+    }
 }

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

@@ -45,7 +45,7 @@ public abstract class LeafInput<T> extends ProcessorDefinition {
         }
 
         AggPathInput other = (AggPathInput) obj;
-        return Objects.equals(context(), other.context()) 
+        return Objects.equals(context(), other.context())
                 && Objects.equals(expression(), other.expression());
     }
 }

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

@@ -9,7 +9,7 @@ import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 import org.elasticsearch.xpack.sql.expression.Expression;
 import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.Processor;
 
-public class NonExecutableInput<T> extends LeafInput<T> {
+public abstract class NonExecutableInput<T> extends LeafInput<T> {
 
     NonExecutableInput(Expression expression, T context) {
         super(expression, context);

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

@@ -13,4 +13,9 @@ public class ReferenceInput extends NonExecutableInput<ColumnReference> {
     public ReferenceInput(Expression expression, ColumnReference context) {
         super(expression, context);
     }
+
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return false;
+    }
 }

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

@@ -34,4 +34,8 @@ public class ScoreProcessorDefinition extends ProcessorDefinition {
         sourceBuilder.trackScores();
     }
 
+    @Override
+    public boolean supportedByAggsOnlyQuery() {
+        return false;
+    }
 }

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

@@ -42,6 +42,11 @@ public class UnaryProcessorDefinition extends ProcessorDefinition {
         return new ChainingProcessor(child.asProcessor(), action);
     }
 
+    @Override
+    public boolean supportedByAggsOnlyQuery() {
+        return child.supportedByAggsOnlyQuery();
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(expression(), child, action);
@@ -52,11 +57,11 @@ public class UnaryProcessorDefinition extends ProcessorDefinition {
         if (this == obj) {
             return true;
         }
-        
+
         if (obj == null || getClass() != obj.getClass()) {
             return false;
         }
-        
+
         UnaryProcessorDefinition other = (UnaryProcessorDefinition) obj;
         return Objects.equals(action, other.action)
                 && Objects.equals(child, other.child)

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

@@ -35,4 +35,9 @@ public class AggRef implements ColumnReference {
     public void collectFields(SqlSourceBuilder sourceBuilder) {
         // Aggregations do not need any special fields
     }
+
+    @Override
+    public boolean supportedByAggsOnlyQuery() {
+        return true;
+    }
 }

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

@@ -42,6 +42,11 @@ public class ComputedRef implements ColumnReference {
         return depth;
     }
 
+    @Override
+    public boolean supportedByAggsOnlyQuery() {
+        return processor.supportedByAggsOnlyQuery();
+    }
+
     @Override
     public void collectFields(SqlSourceBuilder sourceBuilder) {
         processor.collectFields(sourceBuilder);

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

@@ -5,17 +5,21 @@
  */
 package org.elasticsearch.xpack.sql.querydsl.container;
 
-public interface FieldReference extends ColumnReference {
+public abstract class FieldReference implements ColumnReference {
+    /**
+     * Field name.
+     *
+     * @return field name.
+     */
+    public abstract String name();
 
     @Override
-    default int depth() {
+    public final int depth() {
         return 0;
     }
 
-    /**
-     * Field name.
-     * 
-     * @return field name.
-     */
-    String name();
+    @Override
+    public final boolean supportedByAggsOnlyQuery() {
+        return false;
+    }
 }

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

@@ -89,26 +89,8 @@ public class QueryContainer {
         this.columns = refs == null || refs.isEmpty() ? emptyList() : refs;
         this.sort = sort == null || sort.isEmpty() ? emptySet() : sort;
         this.limit = limit;
-
-        int aggLevel = 0;
-        boolean onlyAggs = true;
-
-        for (ColumnReference ref : this.columns) {
-            if (ref.depth() > aggLevel) {
-                aggLevel = ref.depth();
-            }
-            if (ref instanceof ComputedRef) {
-                // check field references
-                if (((ComputedRef) ref).processor().anyMatch(p -> p instanceof ReferenceInput && ((ReferenceInput) p).context() instanceof FieldReference)) {
-                    onlyAggs = false;
-                }
-            }
-            if (ref instanceof FieldReference) {
-                onlyAggs = false;
-            }
-        }
-        aggsOnly = onlyAggs;
-        aggDepth = aggLevel;
+        aggsOnly = columns.stream().allMatch(ColumnReference::supportedByAggsOnlyQuery);
+        aggDepth = columns.stream().mapToInt(ColumnReference::depth).max().orElse(0);
     }
 
     public Query query() {

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

@@ -8,7 +8,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 {
+public class ScriptFieldRef extends FieldReference {
 
     private final String name;
     private final ScriptTemplate script;

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

@@ -7,7 +7,7 @@ package org.elasticsearch.xpack.sql.querydsl.container;
 
 import org.elasticsearch.xpack.sql.execution.search.SqlSourceBuilder;
 
-public class SearchHitFieldRef implements FieldReference {
+public class SearchHitFieldRef extends FieldReference {
     private final String name;
     private final boolean docValue;
     private final String hitName;

sql/server/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/processor/definition/BinaryProcessorDefinitionTests.java

@@ -0,0 +1,56 @@
+/*
+ * 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.expression.function.scalar.processor.definition;
+
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.Processor;
+
+import static java.util.Collections.emptyList;
+
+public class BinaryProcessorDefinitionTests extends ESTestCase {
+    public void testSupportedByAggsOnlyQuery() {
+        ProcessorDefinition supported = new DummyProcessorDefinition(true);
+        ProcessorDefinition unsupported = new DummyProcessorDefinition(false);
+
+        assertFalse(newBinaryProcessor(unsupported, unsupported).supportedByAggsOnlyQuery());
+        assertFalse(newBinaryProcessor(unsupported, supported).supportedByAggsOnlyQuery());
+        assertFalse(newBinaryProcessor(supported, unsupported).supportedByAggsOnlyQuery());
+        assertTrue(newBinaryProcessor(supported, supported).supportedByAggsOnlyQuery());
+    }
+
+    private ProcessorDefinition newBinaryProcessor(ProcessorDefinition left, ProcessorDefinition right) {
+        return new BinaryProcessorDefinition(null, left, right) {
+            @Override
+            public Processor asProcessor() {
+                return null;
+            }
+        };
+    }
+
+    static class DummyProcessorDefinition extends ProcessorDefinition {
+        private final boolean supportedByAggsOnlyQuery;
+
+        DummyProcessorDefinition(boolean supportedByAggsOnlyQuery) {
+            super(null, emptyList());
+            this.supportedByAggsOnlyQuery = supportedByAggsOnlyQuery;
+        }
+
+        @Override
+        public boolean supportedByAggsOnlyQuery() {
+            return supportedByAggsOnlyQuery;
+        }
+
+        @Override
+        public boolean resolved() {
+            return true;
+        }
+
+        @Override
+        public Processor asProcessor() {
+            return null;
+        }
+    }
+}

sql/server/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/processor/definition/UnaryProcessorDefinitionTests.java

@@ -0,0 +1,26 @@
+/*
+ * 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.expression.function.scalar.processor.definition;
+
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.sql.expression.function.scalar.processor.definition.BinaryProcessorDefinitionTests.DummyProcessorDefinition;
+import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.Processor;
+
+import static java.util.Collections.emptyList;
+
+public class UnaryProcessorDefinitionTests extends ESTestCase {
+    public void testSupportedByAggsOnlyQuery() {
+        ProcessorDefinition supported = new DummyProcessorDefinition(true);
+        ProcessorDefinition unsupported = new DummyProcessorDefinition(false);
+
+        assertFalse(newUnaryProcessor(unsupported).supportedByAggsOnlyQuery());
+        assertTrue(newUnaryProcessor(supported).supportedByAggsOnlyQuery());
+    }
+
+    private ProcessorDefinition newUnaryProcessor(ProcessorDefinition child) {
+        return new UnaryProcessorDefinition(null, child, null);
+    }
+}