Skip to content

Commit 11fbc8c

Browse files
fang-xing-esqlfang-xing-esql
authored
[ES|QL] Optional named arguments for function in map (#118619)
* MapExpression for functions
1 parent 377d893 commit 11fbc8c

File tree

35 files changed

+4053
-2015
lines changed

35 files changed

+4053
-2015
lines changed

‎docs/changelog/118619.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 118619
2+
summary: Optional named arguments for function in map
3+
area: EQL
4+
type: enhancement
5+
issues: []

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/EntryExpression.java‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
package org.elasticsearch.xpack.esql.core.expression;
8+
9+
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
10+
import org.elasticsearch.common.io.stream.StreamInput;
11+
import org.elasticsearch.common.io.stream.StreamOutput;
12+
import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
13+
import org.elasticsearch.xpack.esql.core.tree.Source;
14+
import org.elasticsearch.xpack.esql.core.type.DataType;
15+
import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
16+
17+
import java.io.IOException;
18+
import java.util.List;
19+
import java.util.Objects;
20+
21+
/**
22+
* Represent a key-value pair.
23+
*/
24+
public class EntryExpression extends Expression {
25+
public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
26+
Expression.class,
27+
"EntryExpression",
28+
EntryExpression::readFrom
29+
);
30+
31+
private final Expression key;
32+
33+
private final Expression value;
34+
35+
public EntryExpression(Source source, Expression key, Expression value) {
36+
super(source, List.of(key, value));
37+
this.key = key;
38+
this.value = value;
39+
}
40+
41+
private static EntryExpression readFrom(StreamInput in) throws IOException {
42+
return new EntryExpression(
43+
Source.readFrom((StreamInput & PlanStreamInput) in),
44+
in.readNamedWriteable(Expression.class),
45+
in.readNamedWriteable(Expression.class)
46+
);
47+
}
48+
49+
@Override
50+
public void writeTo(StreamOutput out) throws IOException {
51+
source().writeTo(out);
52+
out.writeNamedWriteable(key);
53+
out.writeNamedWriteable(value);
54+
}
55+
56+
@Override
57+
public String getWriteableName() {
58+
return ENTRY.name;
59+
}
60+
61+
@Override
62+
public Expression replaceChildren(List<Expression> newChildren) {
63+
return new EntryExpression(source(), newChildren.get(0), newChildren.get(1));
64+
}
65+
66+
@Override
67+
protected NodeInfo<? extends Expression> info() {
68+
return NodeInfo.create(this, EntryExpression::new, key, value);
69+
}
70+
71+
public Expression key() {
72+
return key;
73+
}
74+
75+
public Expression value() {
76+
return value;
77+
}
78+
79+
@Override
80+
public DataType dataType() {
81+
return value.dataType();
82+
}
83+
84+
@Override
85+
public Nullability nullable() {
86+
return Nullability.FALSE;
87+
}
88+
89+
@Override
90+
public int hashCode() {
91+
return Objects.hash(key, value);
92+
}
93+
94+
@Override
95+
public boolean equals(Object obj) {
96+
if (this == obj) {
97+
return true;
98+
}
99+
if (obj == null || getClass() != obj.getClass()) {
100+
return false;
101+
}
102+
103+
EntryExpression other = (EntryExpression) obj;
104+
return Objects.equals(key, other.key) && Objects.equals(value, other.value);
105+
}
106+
107+
@Override
108+
public String toString() {
109+
return key.toString() + ":" + value.toString();
110+
}
111+
}

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/ExpressionCoreWritables.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ public static List<NamedWriteableRegistry.Entry> expressions() {
2929
entries.add(new NamedWriteableRegistry.Entry(Expression.class, e.name, in -> (Expression) e.reader.read(in)));
3030
}
3131
entries.add(Literal.ENTRY);
32+
entries.addAll(mapExpressions());
3233
return entries;
3334
}
3435

@@ -45,4 +46,8 @@ public static List<NamedWriteableRegistry.Entry> namedExpressions() {
4546
public static List<NamedWriteableRegistry.Entry> attributes() {
4647
return List.of(FieldAttribute.ENTRY, MetadataAttribute.ENTRY, ReferenceAttribute.ENTRY);
4748
}
49+
50+
public static List<NamedWriteableRegistry.Entry> mapExpressions() {
51+
return List.of(EntryExpression.ENTRY, MapExpression.ENTRY);
52+
}
4853
}

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/MapExpression.java‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
package org.elasticsearch.xpack.esql.core.expression;
8+
9+
import org.apache.lucene.util.BytesRef;
10+
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
11+
import org.elasticsearch.common.io.stream.StreamInput;
12+
import org.elasticsearch.common.io.stream.StreamOutput;
13+
import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
14+
import org.elasticsearch.xpack.esql.core.tree.Source;
15+
import org.elasticsearch.xpack.esql.core.type.DataType;
16+
import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
17+
18+
import java.io.IOException;
19+
import java.util.ArrayList;
20+
import java.util.LinkedHashMap;
21+
import java.util.List;
22+
import java.util.Map;
23+
import java.util.Objects;
24+
import java.util.stream.Collectors;
25+
26+
import static org.elasticsearch.xpack.esql.core.type.DataType.UNSUPPORTED;
27+
28+
/**
29+
* Represent a collect of key-value pairs.
30+
*/
31+
public class MapExpression extends Expression {
32+
public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
33+
Expression.class,
34+
"MapExpression",
35+
MapExpression::readFrom
36+
);
37+
38+
private final List<EntryExpression> entryExpressions;
39+
40+
private final Map<Expression, Expression> map;
41+
42+
private final Map<Object, Expression> keyFoldedMap;
43+
44+
public MapExpression(Source source, List<Expression> entries) {
45+
super(source, entries);
46+
int entryCount = entries.size() / 2;
47+
this.entryExpressions = new ArrayList<>(entryCount);
48+
this.map = new LinkedHashMap<>(entryCount);
49+
// create a map with key folded and source removed to make the retrieval of value easier
50+
this.keyFoldedMap = new LinkedHashMap<>(entryCount);
51+
for (int i = 0; i < entryCount; i++) {
52+
Expression key = entries.get(i * 2);
53+
Expression value = entries.get(i * 2 + 1);
54+
entryExpressions.add(new EntryExpression(key.source(), key, value));
55+
map.put(key, value);
56+
if (key instanceof Literal l) {
57+
this.keyFoldedMap.put(l.value(), value);
58+
}
59+
}
60+
}
61+
62+
private static MapExpression readFrom(StreamInput in) throws IOException {
63+
return new MapExpression(
64+
Source.readFrom((StreamInput & PlanStreamInput) in),
65+
in.readNamedWriteableCollectionAsList(Expression.class)
66+
);
67+
}
68+
69+
@Override
70+
public void writeTo(StreamOutput out) throws IOException {
71+
source().writeTo(out);
72+
out.writeNamedWriteableCollection(children());
73+
}
74+
75+
@Override
76+
public String getWriteableName() {
77+
return ENTRY.name;
78+
}
79+
80+
@Override
81+
public MapExpression replaceChildren(List<Expression> newChildren) {
82+
return new MapExpression(source(), newChildren);
83+
}
84+
85+
@Override
86+
protected NodeInfo<MapExpression> info() {
87+
return NodeInfo.create(this, MapExpression::new, children());
88+
}
89+
90+
public List<EntryExpression> entryExpressions() {
91+
return entryExpressions;
92+
}
93+
94+
public Map<Expression, Expression> map() {
95+
return map;
96+
}
97+
98+
public Map<Object, Expression> keyFoldedMap() {
99+
return keyFoldedMap;
100+
}
101+
102+
@Override
103+
public Nullability nullable() {
104+
return Nullability.FALSE;
105+
}
106+
107+
@Override
108+
public DataType dataType() {
109+
return UNSUPPORTED;
110+
}
111+
112+
@Override
113+
public int hashCode() {
114+
return Objects.hash(entryExpressions);
115+
}
116+
117+
public Expression get(Object key) {
118+
if (key instanceof Expression) {
119+
return map.get(key);
120+
} else {
121+
// the key(literal) could be converted to BytesRef by ConvertStringToByteRef
122+
return keyFoldedMap.containsKey(key) ? keyFoldedMap.get(key) : keyFoldedMap.get(new BytesRef(key.toString()));
123+
}
124+
}
125+
126+
@Override
127+
public boolean equals(Object obj) {
128+
if (this == obj) {
129+
return true;
130+
}
131+
if (obj == null || getClass() != obj.getClass()) {
132+
return false;
133+
}
134+
135+
MapExpression other = (MapExpression) obj;
136+
return Objects.equals(entryExpressions, other.entryExpressions);
137+
}
138+
139+
@Override
140+
public String toString() {
141+
String str = entryExpressions.stream().map(String::valueOf).collect(Collectors.joining(", "));
142+
return "{ " + str + " }";
143+
}
144+
}

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/TypeResolutions.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,4 +253,19 @@ private static String acceptedTypesForErrorMsg(String... acceptedTypes) {
253253
return acceptedTypes[0];
254254
}
255255
}
256+
257+
public static TypeResolution isMapExpression(Expression e, String operationName, ParamOrdinal paramOrd) {
258+
if (e instanceof MapExpression == false) {
259+
return new TypeResolution(
260+
format(
261+
null,
262+
"{}argument of [{}] must be a map expression, received [{}]",
263+
paramOrd == null || paramOrd == DEFAULT ? "" : paramOrd.name().toLowerCase(Locale.ROOT) + " ",
264+
operationName,
265+
Expressions.name(e)
266+
)
267+
);
268+
}
269+
return TypeResolution.TYPE_RESOLVED;
270+
}
256271
}

‎x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/ConsumeProcessor.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ public Set<String> getSupportedAnnotationTypes() {
3939
"org.elasticsearch.injection.guice.Inject",
4040
"org.elasticsearch.xpack.esql.expression.function.FunctionInfo",
4141
"org.elasticsearch.xpack.esql.expression.function.Param",
42+
"org.elasticsearch.xpack.esql.expression.function.MapParam",
4243
"org.elasticsearch.rest.ServerlessScope",
4344
"org.elasticsearch.xcontent.ParserConstructor",
4445
"org.elasticsearch.core.UpdateForV9",

0 commit comments

Comments
 (0)