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 3fc545d735d80..aec6102c6d906 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 @@ -15,7 +15,6 @@ import java.util.function.BiConsumer; import java.util.stream.Stream; -import static java.util.Collections.emptyMap; import static java.util.Collections.singletonMap; import static java.util.Collections.unmodifiableCollection; import static java.util.Collections.unmodifiableSet; @@ -141,8 +140,8 @@ public String toString() { } @SuppressWarnings("rawtypes") - public static final AttributeMap EMPTY = new AttributeMap<>(); - + private static final AttributeMap EMPTY = new AttributeMap<>(); + @SuppressWarnings("unchecked") public static final AttributeMap emptyAttributeMap() { return EMPTY; @@ -157,23 +156,6 @@ public AttributeMap() { delegate = new LinkedHashMap<>(); } - /** - * Please use the {@link AttributeMap#builder()} instead. - */ - @Deprecated - public AttributeMap(Map attr) { - if (attr.isEmpty()) { - delegate = emptyMap(); - } - else { - delegate = new LinkedHashMap<>(attr.size()); - - for (Entry entry : attr.entrySet()) { - delegate.put(new AttributeWrapper(entry.getKey()), entry.getValue()); - } - } - } - public AttributeMap(Attribute key, E value) { delegate = singletonMap(new AttributeWrapper(key), value); } @@ -377,23 +359,41 @@ public static Builder builder() { return new Builder<>(); } + public static Builder builder(AttributeMap map) { + return new Builder().putAll(map); + } + public static class Builder { - private final AttributeMap map = new AttributeMap<>(); + private AttributeMap map = null; + private AttributeMap previouslyBuiltMap = null; private Builder() {} + private AttributeMap map() { + if (map == null) { + map = new AttributeMap<>(); + if (previouslyBuiltMap != null) { + map.addAll(previouslyBuiltMap); + } + } + return map; + } + public Builder put(Attribute attr, E value) { - map.add(attr, value); + map().add(attr, value); return this; } public Builder putAll(AttributeMap m) { - map.addAll(m); + map().addAll(m); return this; } public AttributeMap build() { - return new AttributeMap<>(map); + AttributeMap m = map(); + previouslyBuiltMap = m; + map = null; + return m; } } } 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 585d3d5da10c9..3c46f42e80450 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 @@ -13,11 +13,9 @@ import java.util.function.Predicate; import java.util.stream.Stream; -import static java.util.Collections.emptyMap; - public class AttributeSet implements Set { - private static final AttributeMap EMPTY_DELEGATE = new AttributeMap<>(emptyMap()); + private static final AttributeMap EMPTY_DELEGATE = AttributeMap.emptyAttributeMap(); public static final AttributeSet EMPTY = new AttributeSet(EMPTY_DELEGATE); 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 3f716d760654b..c5ce1cea5f54b 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 @@ -21,7 +21,6 @@ import java.util.function.Predicate; import static java.util.Collections.emptyList; -import static java.util.Collections.emptyMap; public final class Expressions { @@ -65,7 +64,7 @@ public static List asAttributes(List named public static AttributeMap asAttributeMap(List named) { if (named.isEmpty()) { - return new AttributeMap<>(emptyMap()); + return AttributeMap.emptyAttributeMap(); } AttributeMap map = new AttributeMap<>(); 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 41ea7fd44d7ba..298ade4db11da 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 @@ -10,7 +10,6 @@ import org.elasticsearch.xpack.ql.type.DataTypes; import java.util.Collection; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; @@ -30,12 +29,12 @@ private static Attribute a(String name) { } private static AttributeMap threeMap() { - Map map = new LinkedHashMap<>(); - map.put(a("one"), "one"); - map.put(a("two"), "two"); - map.put(a("three"), "three"); + AttributeMap.Builder builder = AttributeMap.builder(); + builder.put(a("one"), "one"); + builder.put(a("two"), "two"); + builder.put(a("three"), "three"); - return new AttributeMap<>(map); + return builder.build(); } public void testAttributeMapWithSameAliasesCanResolveAttributes() { @@ -49,19 +48,6 @@ public void testAttributeMapWithSameAliasesCanResolveAttributes() { assertTrue(param1.toAttribute().equals(param2.toAttribute())); assertFalse(param1.toAttribute().semanticEquals(param2.toAttribute())); - Map collectRefs = new LinkedHashMap<>(); - for (Alias a : List.of(param1, param2)) { - collectRefs.put(a.toAttribute(), a.child()); - } - // we can look up the same item by both attributes - assertNotNull(collectRefs.get(param1.toAttribute())); - assertNotNull(collectRefs.get(param2.toAttribute())); - AttributeMap attributeMap = new AttributeMap<>(collectRefs); - - // validate that all Alias can be e - assertTrue(attributeMap.containsKey(param1.toAttribute())); - assertFalse(attributeMap.containsKey(param2.toAttribute())); // results in unknown attribute exception - AttributeMap.Builder mapBuilder = AttributeMap.builder(); for (Alias a : List.of(param1, param2)) { mapBuilder.put(a.toAttribute(), a.child()); @@ -69,7 +55,9 @@ public void testAttributeMapWithSameAliasesCanResolveAttributes() { AttributeMap newAttributeMap = mapBuilder.build(); assertTrue(newAttributeMap.containsKey(param1.toAttribute())); - assertTrue(newAttributeMap.containsKey(param2.toAttribute())); // no more unknown attribute exception + assertTrue(newAttributeMap.get(param1.toAttribute()) == param1.child()); + assertTrue(newAttributeMap.containsKey(param2.toAttribute())); + assertTrue(newAttributeMap.get(param2.toAttribute()) == param2.child()); } private Alias createIntParameterAlias(int index, int value) { @@ -85,13 +73,13 @@ public void testEmptyConstructor() { assertThat(m.isEmpty(), is(true)); } - public void testMapConstructor() { - Map map = new LinkedHashMap<>(); - map.put(a("one"), "one"); - map.put(a("two"), "two"); - map.put(a("three"), "three"); + public void testBuilder() { + AttributeMap.Builder builder = AttributeMap.builder(); + builder.put(a("one"), "one"); + builder.put(a("two"), "two"); + builder.put(a("three"), "three"); - AttributeMap m = new AttributeMap<>(map); + AttributeMap m = builder.build(); assertThat(m.size(), is(3)); assertThat(m.isEmpty(), is(false)); @@ -101,12 +89,16 @@ public void testMapConstructor() { assertThat(m.containsValue("one"), is(true)); assertThat(m.containsValue("on"), is(false)); assertThat(m.attributeNames(), contains("one", "two", "three")); - assertThat(m.values(), contains(map.values().toArray())); + assertThat(m.values(), contains("one", "two", "three")); // defensive copying - map.put(a("four"), "four"); + builder.put(a("four"), "four"); + AttributeMap m2 = builder.build(); assertThat(m.size(), is(3)); assertThat(m.isEmpty(), is(false)); + assertThat(m2.size(), is(4)); + assertThat(m.isEmpty(), is(false)); + assertThat(m2.attributeNames(), contains("one", "two", "three", "four")); } public void testSingleItemConstructor() { @@ -164,12 +156,7 @@ public void testKeySet() { Attribute two = a("two"); Attribute three = a("three"); - Map map = new LinkedHashMap<>(); - map.put(one, "one"); - map.put(two, "two"); - map.put(three, "three"); - - Set keySet = new AttributeMap<>(map).keySet(); + Set keySet = threeMap().keySet(); assertThat(keySet, contains(one, two, three)); // toObject @@ -192,12 +179,7 @@ public void testEntrySet() { Attribute two = a("two"); Attribute three = a("three"); - Map map = new LinkedHashMap<>(); - map.put(one, "one"); - map.put(two, "two"); - map.put(three, "three"); - - Set> set = new AttributeMap<>(map).entrySet(); + Set> set = threeMap().entrySet(); assertThat(set, hasSize(3)); @@ -211,12 +193,9 @@ public void testEntrySet() { assertThat(values, contains("one", "two", "three")); } - public void testForEach() { + public void testCopy() { AttributeMap m = threeMap(); - - Map collect = new LinkedHashMap<>(); - m.forEach(collect::put); - AttributeMap copy = new AttributeMap<>(collect); + AttributeMap copy = AttributeMap.builder().putAll(m).build(); assertThat(m, is(copy)); } 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 5744d6fbee43d..be4333d55fb3f 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 @@ -302,7 +302,7 @@ protected LogicalPlan rule(Project project) { OrderBy ob = (OrderBy) project.child(); // resolve function references (that maybe hiding the target) - final Map collectRefs = new LinkedHashMap<>(); + AttributeMap.Builder collectRefs = AttributeMap.builder(); // collect Attribute sources // only Aliases are interesting since these are the only ones that hide expressions @@ -316,7 +316,7 @@ protected LogicalPlan rule(Project project) { } })); - AttributeMap functions = new AttributeMap<>(collectRefs); + AttributeMap functions = collectRefs.build(); // track the direct parents Map nestedOrders = new LinkedHashMap<>(); @@ -541,14 +541,14 @@ private List combineProjections(List //TODO: this need rewriting when moving functions of NamedExpression // collect aliases in the lower list - Map map = new LinkedHashMap<>(); + AttributeMap.Builder aliasesBuilder = AttributeMap.builder(); for (NamedExpression ne : lower) { if ((ne instanceof Attribute) == false) { - map.put(ne.toAttribute(), ne); + aliasesBuilder.put(ne.toAttribute(), ne); } } - AttributeMap aliases = new AttributeMap<>(map); + AttributeMap aliases = aliasesBuilder.build(); List replaced = new ArrayList<>(); // replace any matching attribute with a lower alias (if there's a match) 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 55cec8d80a6fd..fb6cb91d03e3e 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 @@ -23,9 +23,7 @@ import org.elasticsearch.xpack.sql.SqlIllegalArgumentException; import java.util.ArrayList; -import java.util.LinkedHashMap; import java.util.List; -import java.util.Map; import java.util.Objects; import static java.util.Collections.singletonList; @@ -134,7 +132,7 @@ private AttributeSet valuesOutput() { public AttributeMap valuesToLiterals() { AttributeSet outValues = valuesOutput(); - Map valuesMap = new LinkedHashMap<>(); + AttributeMap.Builder valuesMap = AttributeMap.builder(); int index = 0; // for each attribute, associate its value @@ -152,7 +150,7 @@ public AttributeMap valuesToLiterals() { } } - return new AttributeMap<>(valuesMap); + return valuesMap.build(); } @Override 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 9165e1b5d035c..86da3948eaf03 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 @@ -421,20 +421,21 @@ static EsQueryExec fold(AggregateExec a, EsQueryExec exec) { // track aliases defined in the SELECT and used inside GROUP BY // SELECT x AS a ... GROUP BY a - Map aliasMap = new LinkedHashMap<>(); String id = null; + + AttributeMap.Builder aliases = AttributeMap.builder(); for (NamedExpression ne : a.aggregates()) { if (ne instanceof Alias) { - aliasMap.put(ne.toAttribute(), ((Alias) ne).child()); + aliases.put(ne.toAttribute(), ((Alias) ne).child()); } } - if (aliasMap.isEmpty() == false) { - Map newAliases = new LinkedHashMap<>(queryC.aliases()); - newAliases.putAll(aliasMap); - queryC = queryC.withAliases(new AttributeMap<>(newAliases)); + if (aliases.build().isEmpty() == false) { + aliases.putAll(queryC.aliases()); + queryC = queryC.withAliases(aliases.build()); } + // build the group aggregation // NB: any reference in grouping is already "optimized" by its source so there's no need to look for aliases GroupingContext groupingContext = groupBy(a.groupings()); 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 a2761db533f93..64e4245b557ce 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 @@ -432,9 +432,8 @@ public FieldExtraction resolve(Attribute attribute) { // update proc (if needed) if (qContainer.scalarFunctions().size() != scalarFunctions.size()) { - Map procs = new LinkedHashMap<>(qContainer.scalarFunctions()); - procs.put(attr, proc); - qContainer = qContainer.withScalarProcessors(new AttributeMap<>(procs)); + qContainer = qContainer.withScalarProcessors( + AttributeMap.builder(qContainer.scalarFunctions).put(attr, proc).build()); } return new Tuple<>(qContainer, new ComputedRef(proc)); diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/execution/search/SourceGeneratorTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/execution/search/SourceGeneratorTests.java index 6bee4b338659b..3ef98e6fff235 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/execution/search/SourceGeneratorTests.java +++ b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/execution/search/SourceGeneratorTests.java @@ -14,9 +14,7 @@ import org.elasticsearch.search.sort.FieldSortBuilder; import org.elasticsearch.search.sort.SortOrder; import org.elasticsearch.test.ESTestCase; -import org.elasticsearch.xpack.ql.expression.Attribute; import org.elasticsearch.xpack.ql.expression.AttributeMap; -import org.elasticsearch.xpack.ql.expression.Expression; import org.elasticsearch.xpack.ql.expression.FieldAttribute; import org.elasticsearch.xpack.ql.expression.ReferenceAttribute; import org.elasticsearch.xpack.ql.querydsl.container.AttributeSort; @@ -32,9 +30,6 @@ import org.elasticsearch.xpack.sql.querydsl.container.QueryContainer; import org.elasticsearch.xpack.sql.querydsl.container.ScoreSort; -import java.util.LinkedHashMap; -import java.util.Map; - import static java.util.Collections.singletonList; import static org.elasticsearch.index.query.QueryBuilders.boolQuery; import static org.elasticsearch.index.query.QueryBuilders.matchQuery; @@ -89,9 +84,7 @@ public void testSortNoneSpecified() { public void testSelectScoreForcesTrackingScore() { Score score = new Score(Source.EMPTY); ReferenceAttribute attr = new ReferenceAttribute(score.source(), "score", score.dataType()); - Map alias = new LinkedHashMap<>(); - alias.put(attr, score); - QueryContainer container = new QueryContainer().withAliases(new AttributeMap<>(alias)).addColumn(attr); + QueryContainer container = new QueryContainer().withAliases(new AttributeMap<>(attr, score)).addColumn(attr); SearchSourceBuilder sourceBuilder = SourceGenerator.sourceBuilder(container, null, randomIntBetween(1, 10)); assertTrue(sourceBuilder.trackScores()); } diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainerTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainerTests.java index 8886807ad5840..8b7cc542d93b9 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainerTests.java +++ b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainerTests.java @@ -9,7 +9,6 @@ import org.elasticsearch.xpack.ql.expression.Alias; import org.elasticsearch.xpack.ql.expression.Attribute; import org.elasticsearch.xpack.ql.expression.AttributeMap; -import org.elasticsearch.xpack.ql.expression.Expression; import org.elasticsearch.xpack.ql.expression.FieldAttribute; import org.elasticsearch.xpack.ql.querydsl.query.BoolQuery; import org.elasticsearch.xpack.ql.querydsl.query.MatchAll; @@ -24,8 +23,6 @@ import java.util.AbstractMap.SimpleImmutableEntry; import java.util.Arrays; import java.util.BitSet; -import java.util.LinkedHashMap; -import java.util.Map; import static java.util.Collections.emptyMap; import static java.util.Collections.singletonMap; @@ -86,11 +83,8 @@ public void testColumnMaskShouldDuplicateSameAttributes() { Attribute fourth = new FieldAttribute(Source.EMPTY, "fourth", esField); Alias firstAliased = new Alias(Source.EMPTY, "firstAliased", first); - Map aliasesMap = new LinkedHashMap<>(); - aliasesMap.put(firstAliased.toAttribute(), first); - QueryContainer queryContainer = new QueryContainer() - .withAliases(new AttributeMap<>(aliasesMap)) + .withAliases(new AttributeMap<>(firstAliased.toAttribute(), first)) .addColumn(third) .addColumn(first) .addColumn(fourth)