diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeMap.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeMap.java index b976238130466..60cc5d3f36257 100644 --- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeMap.java +++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeMap.java @@ -12,6 +12,7 @@ import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.Map; +import java.util.Map.Entry; import java.util.Set; import java.util.function.BiConsumer; import java.util.stream.Stream; @@ -20,7 +21,7 @@ import static java.util.Collections.unmodifiableCollection; import static java.util.Collections.unmodifiableSet; -public class AttributeMap implements Map { +public class AttributeMap { static class AttributeWrapper { @@ -144,14 +145,13 @@ public String toString() { private static final AttributeMap EMPTY = new AttributeMap<>(); @SuppressWarnings("unchecked") - public static final AttributeMap emptyAttributeMap() { + public static AttributeMap emptyAttributeMap() { return EMPTY; } private final Map delegate; private Set keySet = null; private Collection values = null; - private Set> entrySet = null; public AttributeMap() { delegate = new LinkedHashMap<>(); @@ -161,16 +161,16 @@ public AttributeMap(Attribute key, E value) { delegate = singletonMap(new AttributeWrapper(key), value); } - void add(Attribute key, E value) { + protected void add(Attribute key, E value) { delegate.put(new AttributeWrapper(key), value); } // a set from a collection of sets without (too much) copying - void addAll(AttributeMap other) { + protected void addAll(AttributeMap other) { delegate.putAll(other.delegate); } - - public AttributeMap combine(AttributeMap other) { + + AttributeMap combine(AttributeMap other) { AttributeMap combine = new AttributeMap<>(); combine.addAll(this); combine.addAll(other); @@ -178,7 +178,7 @@ public AttributeMap combine(AttributeMap other) { return combine; } - public AttributeMap subtract(AttributeMap other) { + AttributeMap subtract(AttributeMap other) { AttributeMap diff = new AttributeMap<>(); for (Entry entry : this.delegate.entrySet()) { if (other.delegate.containsKey(entry.getKey()) == false) { @@ -189,7 +189,7 @@ public AttributeMap subtract(AttributeMap other) { return diff; } - public AttributeMap intersect(AttributeMap other) { + AttributeMap intersect(AttributeMap other) { AttributeMap smaller = (other.size() > size() ? this : other); AttributeMap larger = (smaller == this ? other : this); @@ -203,7 +203,7 @@ public AttributeMap intersect(AttributeMap other) { return intersect; } - public boolean subsetOf(AttributeMap other) { + boolean subsetOf(AttributeMap other) { if (this.size() > other.size()) { return false; } @@ -216,7 +216,7 @@ public boolean subsetOf(AttributeMap other) { return true; } - public Set attributeNames() { + Set attributeNames() { Set s = new LinkedHashSet<>(size()); for (AttributeWrapper aw : delegate.keySet()) { @@ -225,67 +225,22 @@ public Set attributeNames() { return s; } - @Override public int size() { return delegate.size(); } - @Override public boolean isEmpty() { return delegate.isEmpty(); } - - @Override - public boolean containsKey(Object key) { - if (key instanceof NamedExpression) { - return delegate.keySet().contains(new AttributeWrapper(((NamedExpression) key).toAttribute())); - } - return false; - } - - @Override - public boolean containsValue(Object value) { - return delegate.values().contains(value); - } - - @Override - public E get(Object key) { + + public E getOrDefault(Object key, E defaultValue) { if (key instanceof NamedExpression) { - return delegate.get(new AttributeWrapper(((NamedExpression) key).toAttribute())); + return delegate.getOrDefault(new AttributeWrapper(((NamedExpression) key).toAttribute()), defaultValue); } - return null; - } - - @Override - public E getOrDefault(Object key, E defaultValue) { - E e; - return (((e = get(key)) != null) || containsKey(key)) - ? e - : defaultValue; - } - - @Override - public E put(Attribute key, E value) { - throw new UnsupportedOperationException(); + return defaultValue; } - @Override - public E remove(Object key) { - throw new UnsupportedOperationException(); - } - - @Override - public void putAll(Map m) { - throw new UnsupportedOperationException(); - } - - @Override - public void clear() { - throw new UnsupportedOperationException(); - } - - @Override - public Set keySet() { + protected Set keySet() { if (keySet == null) { keySet = new UnwrappingSet<>(delegate.keySet()) { @Override @@ -296,44 +251,14 @@ protected Attribute unwrap(AttributeWrapper next) { } return keySet; } - - @Override - public Collection values() { + + protected Collection values() { if (values == null) { values = unmodifiableCollection(delegate.values()); } return values; } - @Override - public Set> entrySet() { - if (entrySet == null) { - entrySet = new UnwrappingSet<>(delegate.entrySet()) { - @Override - protected Entry unwrap(final Entry next) { - return new Entry<>() { - @Override - public Attribute getKey() { - return next.getKey().attr; - } - - @Override - public E getValue() { - return next.getValue(); - } - - @Override - public E setValue(E value) { - throw new UnsupportedOperationException(); - } - }; - } - }; - } - return entrySet; - } - - @Override public void forEach(BiConsumer action) { delegate.forEach((k, v) -> action.accept(k.attr, v)); } diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeSet.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeSet.java index cdaa48d949e4d..dea5ddc008cc0 100644 --- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeSet.java +++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeSet.java @@ -6,15 +6,14 @@ */ package org.elasticsearch.xpack.ql.expression; -import java.util.Collection; import java.util.Iterator; +import java.util.List; import java.util.Set; import java.util.Spliterator; import java.util.function.Consumer; -import java.util.function.Predicate; import java.util.stream.Stream; -public class AttributeSet implements Set { +public class AttributeSet implements Iterable { private static final AttributeMap EMPTY_DELEGATE = AttributeMap.emptyAttributeMap(); @@ -32,8 +31,8 @@ public AttributeSet() { public AttributeSet(Attribute attr) { delegate = new AttributeMap<>(attr, PRESENT); } - - public AttributeSet(Collection attr) { + + public AttributeSet(List attr) { if (attr.isEmpty()) { delegate = EMPTY_DELEGATE; } @@ -52,7 +51,7 @@ private AttributeSet(AttributeMap delegate) { // package protected - should be called through Expressions to cheaply create // a set from a collection of sets without too much copying - void addAll(AttributeSet other) { + protected void addAll(AttributeSet other) { delegate.addAll(other.delegate); } @@ -81,96 +80,28 @@ public void forEach(Consumer action) { delegate.forEach((k, v) -> action.accept(k)); } - @Override public int size() { return delegate.size(); } - @Override public boolean isEmpty() { return delegate.isEmpty(); } - @Override - public boolean contains(Object o) { - return delegate.containsKey(o); - } - - @Override - public boolean containsAll(Collection c) { - for (Object o : c) { - if (delegate.containsKey(o) == false) { - return false; - } - } - return true; - } - @Override public Iterator iterator() { return delegate.keySet().iterator(); } - @Override - public Object[] toArray() { - return delegate.keySet().toArray(); - } - - @Override - public T[] toArray(T[] a) { - return delegate.keySet().toArray(a); - } - - @Override - public boolean add(Attribute e) { - throw new UnsupportedOperationException(); - } - - @Override - public boolean remove(Object o) { - throw new UnsupportedOperationException(); - } - - @Override - public boolean addAll(Collection c) { - throw new UnsupportedOperationException(); - } - - @Override - public boolean retainAll(Collection c) { - throw new UnsupportedOperationException(); - } - - @Override - public boolean removeAll(Collection c) { - throw new UnsupportedOperationException(); - } - - @Override - public void clear() { - throw new UnsupportedOperationException(); - } - @Override public Spliterator spliterator() { throw new UnsupportedOperationException(); } - @Override - public boolean removeIf(Predicate filter) { - throw new UnsupportedOperationException(); - } - - @Override public Stream stream() { return delegate.keySet().stream(); } - @Override - public Stream parallelStream() { - return delegate.keySet().parallelStream(); - } - @Override public boolean equals(Object o) { return delegate.equals(o); diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/Expressions.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/Expressions.java index b8e867653528d..085c51f6f04c4 100644 --- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/Expressions.java +++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/Expressions.java @@ -63,18 +63,6 @@ public static List asAttributes(List named return list; } - public static AttributeMap asAttributeMap(List named) { - if (named.isEmpty()) { - return AttributeMap.emptyAttributeMap(); - } - - AttributeMap map = new AttributeMap<>(); - for (NamedExpression exp : named) { - map.add(exp.toAttribute(), exp); - } - return map; - } - public static boolean anyMatch(List exps, Predicate predicate) { for (Expression exp : exps) { if (exp.anyMatch(predicate)) { @@ -174,7 +162,7 @@ public static List> aliases(List output) { + public static boolean hasReferenceAttribute(Iterable output) { for (Attribute attribute : output) { if (attribute instanceof ReferenceAttribute) { return true; diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/CollectionUtils.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/CollectionUtils.java index e8a5ccdfb4c11..cc19b9dbb1a39 100644 --- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/CollectionUtils.java +++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/CollectionUtils.java @@ -48,7 +48,6 @@ public static List combine(Collection... collections) { List list = new ArrayList<>(); for (Collection col : collections) { - // typically AttributeSet which ends up iterating anyway plus creating a redundant array if (col instanceof Set) { for (T t : col) { list.add(t); @@ -61,6 +60,20 @@ public static List combine(Collection... collections) { return list; } + @SafeVarargs + @SuppressWarnings({ "varargs", "unchecked" }) + public static List combineIterables(Iterable... iterables) { + if (org.elasticsearch.common.util.CollectionUtils.isEmpty(iterables)) { + return emptyList(); + } + + List list = new ArrayList<>(); + for (Iterable col : iterables) { + addAll(list, col); + } + return list; + } + @SafeVarargs @SuppressWarnings("varargs") public static List combine(Collection left, T... entries) { @@ -73,6 +86,22 @@ public static List combine(Collection left, T... entries) { } return list; } + + @SafeVarargs + @SuppressWarnings("varargs") + public static List appendToIterable(Iterable left, T... entries) { + List list = null; + if (left instanceof Collection) { + list = new ArrayList<>(((Collection) left).size() + entries.length); + } else { + list = new ArrayList<>(); + } + addAll(list, left); + if (entries.length > 0) { + Collections.addAll(list, entries); + } + return list; + } public static int mapSize(int size) { if (size < 2) { @@ -80,4 +109,17 @@ public static int mapSize(int size) { } return (int) (size / 0.75f + 1f); } + + public static void addAll(List list, Iterable col) { + // typically AttributeSet which ends up iterating anyway plus creating a redundant array + if (col instanceof Set || col instanceof Collection == false) { + for (T t : col) { + list.add(t); + } + } else { + if (((Collection) col).isEmpty() == false) { + list.addAll((Collection) col); + } + } + } } diff --git a/x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/expression/AttributeMapTests.java b/x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/expression/AttributeMapTests.java index 1ec93d2436526..840ce5c4edb23 100644 --- a/x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/expression/AttributeMapTests.java +++ b/x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/expression/AttributeMapTests.java @@ -12,11 +12,8 @@ import java.util.Collection; import java.util.List; -import java.util.Map; -import java.util.Map.Entry; import java.util.Set; -import static java.util.stream.Collectors.toList; import static org.hamcrest.Matchers.arrayContaining; import static org.hamcrest.Matchers.arrayWithSize; import static org.hamcrest.Matchers.contains; @@ -54,11 +51,9 @@ public void testAttributeMapWithSameAliasesCanResolveAttributes() { mapBuilder.put(a.toAttribute(), a.child()); } AttributeMap newAttributeMap = mapBuilder.build(); - - assertTrue(newAttributeMap.containsKey(param1.toAttribute())); - assertTrue(newAttributeMap.get(param1.toAttribute()) == param1.child()); - assertTrue(newAttributeMap.containsKey(param2.toAttribute())); - assertTrue(newAttributeMap.get(param2.toAttribute()) == param2.child()); + + assertTrue(newAttributeMap.getOrDefault(param1.toAttribute(), null) == param1.child()); + assertTrue(newAttributeMap.getOrDefault(param2.toAttribute(), null) == param2.child()); } private Alias createIntParameterAlias(int index, int value) { @@ -83,13 +78,9 @@ public void testBuilder() { AttributeMap m = builder.build(); assertThat(m.size(), is(3)); assertThat(m.isEmpty(), is(false)); - - Attribute one = m.keySet().iterator().next(); - assertThat(m.containsKey(one), is(true)); - assertThat(m.containsKey(a("one")), is(false)); - assertThat(m.containsValue("one"), is(true)); - assertThat(m.containsValue("on"), is(false)); + assertThat(m.attributeNames(), contains("one", "two", "three")); + assertThat(m.values().size(), is(3)); assertThat(m.values(), contains("one", "two", "three")); } @@ -99,10 +90,8 @@ public void testSingleItemConstructor() { assertThat(m.size(), is(1)); assertThat(m.isEmpty(), is(false)); - assertThat(m.containsKey(one), is(true)); - assertThat(m.containsKey(a("one")), is(false)); - assertThat(m.containsValue("one"), is(true)); - assertThat(m.containsValue("on"), is(false)); + assertThat(m.values().size(), is(1)); + assertThat(m.values(), contains("one")); } public void testSubtract() { @@ -166,25 +155,6 @@ public void testValues() { assertThat(values, contains("one", "two", "three")); } - public void testEntrySet() { - Attribute one = a("one"); - Attribute two = a("two"); - Attribute three = a("three"); - - Set> set = threeMap().entrySet(); - - assertThat(set, hasSize(3)); - - List keys = set.stream().map(Map.Entry::getKey).collect(toList()); - List values = set.stream().map(Map.Entry::getValue).collect(toList()); - - assertThat(keys, hasSize(3)); - - - assertThat(values, hasSize(3)); - assertThat(values, contains("one", "two", "three")); - } - public void testCopy() { AttributeMap m = threeMap(); AttributeMap copy = AttributeMap.builder().putAll(m).build(); diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java index de480b2c4b261..b53e62aea7298 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java @@ -78,6 +78,7 @@ import static org.elasticsearch.xpack.ql.analyzer.AnalyzerRules.AnalyzerRule; import static org.elasticsearch.xpack.ql.analyzer.AnalyzerRules.BaseAnalyzerRule; import static org.elasticsearch.xpack.ql.util.CollectionUtils.combine; +import static org.elasticsearch.xpack.ql.util.CollectionUtils.combineIterables; public class Analyzer extends RuleExecutor { /** @@ -638,7 +639,7 @@ protected LogicalPlan doRule(LogicalPlan plan) { AttributeMap.Builder builder = AttributeMap.builder(); // collect aliases child.forEachUp(p -> p.forEachExpressionUp(Alias.class, a -> builder.put(a.toAttribute(), a.child()))); - final Map collectRefs = builder.build(); + final AttributeMap collectRefs = builder.build(); referencesStream = referencesStream.filter(r -> { for (Attribute attr : child.outputSet()) { @@ -756,7 +757,7 @@ private static LogicalPlan propagateMissing(LogicalPlan plan, AttributeSet missi if (plan instanceof Project) { Project p = (Project) plan; AttributeSet diff = missing.subtract(p.child().outputSet()); - return new Project(p.source(), propagateMissing(p.child(), diff, failed), combine(p.projections(), missing)); + return new Project(p.source(), propagateMissing(p.child(), diff, failed), combineIterables(p.projections(), missing)); } if (plan instanceof Aggregate) { @@ -777,7 +778,7 @@ private static LogicalPlan propagateMissing(LogicalPlan plan, AttributeSet missi if (failed.isEmpty() == false) { return plan; } - return new Aggregate(a.source(), a.child(), a.groupings(), combine(a.aggregates(), missing)); + return new Aggregate(a.source(), a.child(), a.groupings(), combineIterables(a.aggregates(), missing)); } // LeafPlans are tables and BinaryPlans are joins so pushing can only happen on unary @@ -785,7 +786,7 @@ private static LogicalPlan propagateMissing(LogicalPlan plan, AttributeSet missi return plan.replaceChildrenSameSize(singletonList(propagateMissing(((UnaryPlan) plan).child(), missing, failed))); } - failed.addAll(missing); + CollectionUtils.addAll(failed, missing); return plan; } diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java index ce2759deb7a36..abb986ffaac20 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java @@ -75,7 +75,7 @@ import static org.elasticsearch.xpack.ql.analyzer.VerifierChecks.checkFilterConditionType; import static org.elasticsearch.xpack.ql.common.Failure.fail; import static org.elasticsearch.xpack.ql.type.DataTypes.BINARY; -import static org.elasticsearch.xpack.ql.util.CollectionUtils.combine; +import static org.elasticsearch.xpack.ql.util.CollectionUtils.appendToIterable; import static org.elasticsearch.xpack.sql.stats.FeatureMetric.COMMAND; import static org.elasticsearch.xpack.sql.stats.FeatureMetric.GROUPBY; import static org.elasticsearch.xpack.sql.stats.FeatureMetric.HAVING; @@ -406,7 +406,7 @@ private static boolean checkGroupByHavingHasOnlyAggs(Expression e, Set source, List private static void checkPivot(LogicalPlan p, Set localFailures, AttributeMap attributeRefs) { p.forEachDown(Pivot.class, pv -> { // check only exact fields are used inside PIVOTing - if (onlyExactFields(combine(pv.groupingSet(), pv.column()), localFailures) == false + if (onlyExactFields(appendToIterable(pv.groupingSet(), pv.column()), localFailures) == false || onlyRawFields(pv.groupingSet(), localFailures, attributeRefs) == false) { // if that is not the case, no need to do further validation since the declaration is fundamentally wrong return; diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java index 9942d1df31cc9..87af9c7e48b77 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java @@ -222,7 +222,7 @@ public LogicalPlan apply(LogicalPlan plan) { AttributeMap.Builder builder = AttributeMap.builder(); // collect aliases plan.forEachExpressionUp(Alias.class, a -> builder.put(a.toAttribute(), a.child())); - final Map collectRefs = builder.build(); + final AttributeMap collectRefs = builder.build(); java.util.function.Function replaceReference = r -> collectRefs.getOrDefault(r, r); plan = plan.transformUp(p -> { @@ -300,7 +300,7 @@ static class PruneOrderByNestedFields extends OptimizerRule { private void findNested(Expression exp, AttributeMap functions, Consumer onFind) { exp.forEachUp(e -> { if (e instanceof ReferenceAttribute) { - Function f = functions.get(e); + Function f = functions.getOrDefault(e, null); if (f != null) { findNested(f, functions, onFind); } diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/Pivot.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/Pivot.java index 2a0561fbf9b51..eb5e01cbf84f5 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/Pivot.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plan/logical/Pivot.java @@ -7,6 +7,7 @@ package org.elasticsearch.xpack.sql.plan.logical; +import org.elasticsearch.common.util.CollectionUtils; import org.elasticsearch.xpack.ql.capabilities.Resolvables; import org.elasticsearch.xpack.ql.expression.Alias; import org.elasticsearch.xpack.ql.expression.Attribute; @@ -57,7 +58,7 @@ public Pivot(Source source, LogicalPlan child, Expression column, List(new AttributeSet(Expressions.onlyPrimitiveFieldAttributes(child().output())) + grouping = CollectionUtils.iterableAsArrayList(new AttributeSet(Expressions.onlyPrimitiveFieldAttributes(child().output())) // make sure to have the column as the last entry (helps with translation) so substract it .subtract(columnSet) .subtract(Expressions.references(aggregates)) @@ -157,7 +158,7 @@ public AttributeMap valuesToLiterals() { @Override public List output() { if (output == null) { - output = new ArrayList<>(groupingSet() + output = CollectionUtils.iterableAsArrayList(groupingSet() .subtract(Expressions.references(singletonList(column))) .combine(valuesOutput())); } diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java index 46bd7fe17a1c6..8b26e72dd6007 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java @@ -7,6 +7,7 @@ package org.elasticsearch.xpack.sql.planner; import org.elasticsearch.common.collect.Tuple; +import org.elasticsearch.common.util.CollectionUtils; import org.elasticsearch.xpack.ql.execution.search.AggRef; import org.elasticsearch.xpack.ql.execution.search.FieldExtraction; import org.elasticsearch.xpack.ql.expression.Alias; @@ -94,7 +95,7 @@ import java.util.Map.Entry; import java.util.concurrent.atomic.AtomicReference; -import static org.elasticsearch.xpack.ql.util.CollectionUtils.combine; +import static org.elasticsearch.xpack.ql.util.CollectionUtils.combineIterables; import static org.elasticsearch.xpack.sql.expression.function.grouping.Histogram.DAY_INTERVAL; import static org.elasticsearch.xpack.sql.expression.function.grouping.Histogram.MONTH_INTERVAL; import static org.elasticsearch.xpack.sql.expression.function.grouping.Histogram.YEAR_INTERVAL; @@ -707,7 +708,7 @@ protected PhysicalPlan rule(OrderExec plan) { // if it's a reference, get the target expression if (orderExpression instanceof ReferenceAttribute) { - orderExpression = qContainer.aliases().get(orderExpression); + orderExpression = qContainer.aliases().getOrDefault(orderExpression, null); } String lookup = Expressions.id(orderExpression); GroupByKey group = qContainer.findGroupForAgg(lookup); @@ -798,7 +799,7 @@ protected PhysicalPlan rule(PivotExec plan) { Pivot p = plan.pivot(); EsQueryExec fold = FoldAggregate .fold(new AggregateExec(plan.source(), exec, - new ArrayList<>(p.groupingSet()), combine(p.groupingSet(), p.aggregates())), exec); + CollectionUtils.iterableAsArrayList(p.groupingSet()), combineIterables(p.groupingSet(), p.aggregates())), exec); // replace the aggregate extractors with pivot specific extractors // these require a reference to the pivoting column in order to compare the value @@ -813,10 +814,8 @@ protected PhysicalPlan rule(PivotExec plan) { for (int i = startingIndex; i < fields.size(); i++) { Tuple tuple = fields.remove(i); - for (Map.Entry entry : values.entrySet()) { - fields.add(new Tuple<>( - new PivotColumnRef(groupTuple.v1(), tuple.v1(), entry.getValue().value()), Expressions.id(entry.getKey()))); - } + values.forEach((attr, literal) -> + fields.add(new Tuple<>(new PivotColumnRef(groupTuple.v1(), tuple.v1(), literal.value()), Expressions.id(attr)))); i += values.size(); } diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainer.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainer.java index e0b1d081ec322..51c4b8cf974b0 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainer.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainer.java @@ -77,7 +77,7 @@ public class QueryContainer { // scalar function processors - recorded as functions get folded; // at scrolling, their inputs (leaves) get updated - private final AttributeMap scalarFunctions; + private AttributeMap scalarFunctions; private final Map sort; private final int limit; @@ -89,8 +89,6 @@ public class QueryContainer { // computed private Boolean aggsOnly; private Boolean customSort; - // associate Attributes with aliased FieldAttributes (since they map directly to ES fields) - private Map fieldAlias; public QueryContainer() { @@ -312,16 +310,8 @@ public QueryContainer addSort(String expressionId, Sort sortable) { } private String aliasName(Attribute attr) { - if (fieldAlias == null) { - fieldAlias = new LinkedHashMap<>(); - for (Map.Entry entry : aliases.entrySet()) { - if (entry.getValue() instanceof FieldAttribute) { - fieldAlias.put(entry.getKey(), (FieldAttribute) entry.getValue()); - } - } - } - FieldAttribute fa = fieldAlias.get(attr); - return fa != null ? fa.name() : attr.name(); + Expression fa = aliases.getOrDefault(attr, null); + return fa instanceof FieldAttribute ? ((FieldAttribute)fa).name() : attr.name(); } // @@ -375,7 +365,11 @@ static Query rewriteToContainNestedField(@Nullable Query query, Source source, S // replace function/operators's input with references private Tuple resolvedTreeComputingRef(ScalarFunction function, Attribute attr) { - Pipe proc = scalarFunctions.computeIfAbsent(attr, v -> function.asPipe()); + Pipe proc = null; + if ((proc = scalarFunctions.getOrDefault(attr, null)) == null) { + proc = function.asPipe(); + scalarFunctions = AttributeMap.builder(scalarFunctions).put(attr, proc).build(); + } // find the processor inputs (Attributes) and convert them into references // no need to promote them to the top since the container doesn't have to be aware