Skip to content

Commit ef783f1

Browse files
ykmr1224ykmr1224
authored
Add MAP_CONCAT internal function (opensearch-project#4477)
* Add MAP_CONCAT internal function Signed-off-by: Tomoyuki Morita <moritato@amazon.com> * Minor fix Signed-off-by: Tomoyuki Morita <moritato@amazon.com> --------- Signed-off-by: Tomoyuki Morita <moritato@amazon.com> Signed-off-by: Tomoyuki MORITA <moritato@amazon.com>
1 parent a8f08ad commit ef783f1

3 files changed

Lines changed: 183 additions & 0 deletions

File tree

‎core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ public enum BuiltinFunctionName {
6262
/** Collection functions */
6363
ARRAY(FunctionName.of("array")),
6464
ARRAY_LENGTH(FunctionName.of("array_length")),
65+
MAP_CONCAT(FunctionName.of("map_concat"), true),
6566
MVAPPEND(FunctionName.of("mvappend")),
6667
MVJOIN(FunctionName.of("mvjoin")),
6768
FORALL(FunctionName.of("forall")),

‎core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@
123123
import static org.opensearch.sql.expression.function.BuiltinFunctionName.LTRIM;
124124
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAKEDATE;
125125
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAKETIME;
126+
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAP_CONCAT;
126127
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH;
127128
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_BOOL_PREFIX;
128129
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_PHRASE;
@@ -896,6 +897,7 @@ void populate() {
896897
registerOperator(ARRAY, PPLBuiltinOperators.ARRAY);
897898
registerOperator(MVAPPEND, PPLBuiltinOperators.MVAPPEND);
898899
registerOperator(ARRAY_LENGTH, SqlLibraryOperators.ARRAY_LENGTH);
900+
registerOperator(MAP_CONCAT, SqlLibraryOperators.MAP_CONCAT);
899901
registerOperator(FORALL, PPLBuiltinOperators.FORALL);
900902
registerOperator(EXISTS, PPLBuiltinOperators.EXISTS);
901903
registerOperator(FILTER, PPLBuiltinOperators.FILTER);

‎integ-test/src/test/java/org/opensearch/sql/calcite/standalone/MapConcatFunctionIT.java‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.standalone;
7+
8+
import java.io.IOException;
9+
import java.sql.PreparedStatement;
10+
import java.sql.ResultSet;
11+
import java.sql.SQLException;
12+
import java.util.List;
13+
import java.util.Map;
14+
import org.apache.calcite.plan.Contexts;
15+
import org.apache.calcite.plan.RelTraitDef;
16+
import org.apache.calcite.rel.RelNode;
17+
import org.apache.calcite.rel.type.RelDataType;
18+
import org.apache.calcite.rex.RexBuilder;
19+
import org.apache.calcite.rex.RexNode;
20+
import org.apache.calcite.schema.SchemaPlus;
21+
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
22+
import org.apache.calcite.sql.parser.SqlParser;
23+
import org.apache.calcite.sql.type.SqlTypeName;
24+
import org.apache.calcite.tools.Frameworks;
25+
import org.apache.calcite.tools.Programs;
26+
import org.apache.calcite.tools.RelBuilder;
27+
import org.junit.jupiter.api.Test;
28+
import org.opensearch.sql.calcite.CalcitePlanContext;
29+
import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners;
30+
import org.opensearch.sql.common.setting.Settings;
31+
import org.opensearch.sql.executor.QueryType;
32+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
33+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
34+
35+
public class MapConcatFunctionIT extends CalcitePPLIntegTestCase {
36+
37+
private static final String MAP_FIELD = "map";
38+
private static final String ID_FIELD = "id";
39+
TestContext context;
40+
41+
@Override
42+
public void init() throws IOException {
43+
super.init();
44+
context = createTestContext();
45+
enableCalcite();
46+
}
47+
48+
@Test
49+
public void testMapConcatWithNullValues() throws Exception {
50+
RelDataType mapType = createMapType(context.rexBuilder);
51+
RexNode map1 = context.rexBuilder.makeNullLiteral(mapType);
52+
RexNode map2 = context.rexBuilder.makeNullLiteral(mapType);
53+
54+
RexNode mapConcatCall =
55+
PPLFuncImpTable.INSTANCE.resolve(
56+
context.rexBuilder, BuiltinFunctionName.MAP_CONCAT, map1, map2);
57+
58+
RelNode relNode =
59+
context
60+
.relBuilder
61+
.values(new String[] {ID_FIELD}, 1)
62+
.project(context.relBuilder.alias(mapConcatCall, MAP_FIELD))
63+
.build();
64+
65+
executeRelNodeAndVerify(
66+
context.planContext,
67+
relNode,
68+
resultSet -> {
69+
assertTrue(resultSet.next());
70+
verifyColumns(resultSet, MAP_FIELD);
71+
assertNull(resultSet.getObject(1));
72+
});
73+
}
74+
75+
@Test
76+
public void testMapConcat() throws Exception {
77+
RexNode map1 =
78+
context.rexBuilder.makeCall(
79+
SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR,
80+
context.rexBuilder.makeLiteral("key1"),
81+
context.rexBuilder.makeLiteral("value1"),
82+
context.rexBuilder.makeLiteral("key2"),
83+
context.rexBuilder.makeLiteral("value2"));
84+
RexNode map2 =
85+
context.rexBuilder.makeCall(
86+
SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR,
87+
context.rexBuilder.makeLiteral("key2"),
88+
context.rexBuilder.makeLiteral("updated_value2"),
89+
context.rexBuilder.makeLiteral("key3"),
90+
context.rexBuilder.makeLiteral("value3"));
91+
92+
RexNode mapConcatCall =
93+
PPLFuncImpTable.INSTANCE.resolve(
94+
context.rexBuilder, BuiltinFunctionName.MAP_CONCAT, map1, map2);
95+
96+
RelNode relNode =
97+
context
98+
.relBuilder
99+
.values(new String[] {ID_FIELD}, 1)
100+
.project(context.relBuilder.alias(mapConcatCall, MAP_FIELD))
101+
.build();
102+
103+
executeRelNodeAndVerify(
104+
context.planContext,
105+
relNode,
106+
resultSet -> {
107+
assertTrue(resultSet.next());
108+
verifyColumns(resultSet, MAP_FIELD);
109+
Map<String, String> result = (Map<String, String>) resultSet.getObject(1);
110+
assertEquals("value1", result.get("key1"));
111+
assertEquals("updated_value2", result.get("key2"));
112+
assertEquals("value3", result.get("key3"));
113+
});
114+
}
115+
116+
private static class TestContext {
117+
final CalcitePlanContext planContext;
118+
final RelBuilder relBuilder;
119+
final RexBuilder rexBuilder;
120+
121+
TestContext(CalcitePlanContext planContext, RelBuilder relBuilder, RexBuilder rexBuilder) {
122+
this.planContext = planContext;
123+
this.relBuilder = relBuilder;
124+
this.rexBuilder = rexBuilder;
125+
}
126+
}
127+
128+
@FunctionalInterface
129+
private interface ResultVerifier {
130+
void verify(ResultSet resultSet) throws SQLException;
131+
}
132+
133+
private TestContext createTestContext() {
134+
CalcitePlanContext planContext = createCalcitePlanContext();
135+
return new TestContext(planContext, planContext.relBuilder, planContext.rexBuilder);
136+
}
137+
138+
private RelDataType createMapType(RexBuilder rexBuilder) {
139+
RelDataType stringType = rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR);
140+
RelDataType anyType = rexBuilder.getTypeFactory().createSqlType(SqlTypeName.ANY);
141+
return rexBuilder.getTypeFactory().createMapType(stringType, anyType);
142+
}
143+
144+
private void executeRelNodeAndVerify(
145+
CalcitePlanContext planContext, RelNode relNode, ResultVerifier verifier)
146+
throws SQLException {
147+
try (PreparedStatement statement = OpenSearchRelRunners.run(planContext, relNode)) {
148+
ResultSet resultSet = statement.executeQuery();
149+
verifier.verify(resultSet);
150+
}
151+
}
152+
153+
private void verifyColumns(ResultSet resultSet, String... expectedColumnNames)
154+
throws SQLException {
155+
assertEquals(expectedColumnNames.length, resultSet.getMetaData().getColumnCount());
156+
157+
for (int i = 0; i < expectedColumnNames.length; i++) {
158+
String expectedName = expectedColumnNames[i];
159+
String actualName = resultSet.getMetaData().getColumnName(i + 1);
160+
assertEquals(expectedName, actualName);
161+
}
162+
}
163+
164+
private CalcitePlanContext createCalcitePlanContext() {
165+
// Create a Frameworks.ConfigBuilder similar to CalcitePPLAbstractTest
166+
final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
167+
Frameworks.ConfigBuilder config =
168+
Frameworks.newConfigBuilder()
169+
.parserConfig(SqlParser.Config.DEFAULT)
170+
.defaultSchema(rootSchema)
171+
.traitDefs((List<RelTraitDef>) null)
172+
.programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2));
173+
174+
config.context(Contexts.of(RelBuilder.Config.DEFAULT));
175+
176+
Settings settings = getSettings();
177+
return CalcitePlanContext.create(
178+
config.build(), settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT), QueryType.PPL);
179+
}
180+
}

0 commit comments

Comments
 (0)