Skip to content

Commit e10997e

Browse files
ykmr1224ykmr1224
authored
Refactor JsonExtractAllFunctionIT and MapConcatFunctionIT (#4623)
1 parent 18ab4dc commit e10997e

4 files changed

Lines changed: 2 additions & 190 deletions

File tree

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

Lines changed: 1 addition & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -5,46 +5,23 @@
55

66
package org.opensearch.sql.calcite.standalone;
77

8-
import java.io.IOException;
9-
import java.sql.PreparedStatement;
108
import java.sql.ResultSet;
119
import java.sql.SQLException;
1210
import java.util.List;
1311
import java.util.Map;
1412
import java.util.stream.Collectors;
15-
import org.apache.calcite.plan.Contexts;
16-
import org.apache.calcite.plan.RelTraitDef;
1713
import org.apache.calcite.rel.RelNode;
1814
import org.apache.calcite.rel.type.RelDataType;
19-
import org.apache.calcite.rex.RexBuilder;
2015
import org.apache.calcite.rex.RexNode;
21-
import org.apache.calcite.schema.SchemaPlus;
22-
import org.apache.calcite.sql.parser.SqlParser;
2316
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;
2717
import org.junit.jupiter.api.Test;
28-
import org.opensearch.sql.calcite.CalcitePlanContext;
29-
import org.opensearch.sql.calcite.SysLimit;
30-
import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners;
31-
import org.opensearch.sql.executor.QueryType;
3218
import org.opensearch.sql.expression.function.BuiltinFunctionName;
3319
import org.opensearch.sql.expression.function.PPLFuncImpTable;
3420

35-
/** Integration test for internal function JSON_EXTRACT_ALL in Calcite PPL. */
36-
public class JsonExtractAllFunctionIT extends CalcitePPLIntegTestCase {
21+
public class JsonExtractAllFunctionIT extends CalcitePPLRelNodeIntegTestCase {
3722

3823
private static final String RESULT_FIELD = "result";
3924
private static final String ID_FIELD = "id";
40-
TestContext context;
41-
42-
@Override
43-
public void init() throws IOException {
44-
super.init();
45-
context = createTestContext();
46-
enableCalcite();
47-
}
4825

4926
@Test
5027
public void testJsonExtractAllWithNullInput() throws Exception {
@@ -310,61 +287,4 @@ public void testJsonExtractAllWithInvalidJson() throws Exception {
310287
assertEquals(1, map.size());
311288
});
312289
}
313-
314-
@FunctionalInterface
315-
private interface ResultVerifier {
316-
void verify(ResultSet resultSet) throws SQLException;
317-
}
318-
319-
private static class TestContext {
320-
final CalcitePlanContext planContext;
321-
final RelBuilder relBuilder;
322-
final RexBuilder rexBuilder;
323-
324-
TestContext(CalcitePlanContext planContext, RelBuilder relBuilder, RexBuilder rexBuilder) {
325-
this.planContext = planContext;
326-
this.relBuilder = relBuilder;
327-
this.rexBuilder = rexBuilder;
328-
}
329-
}
330-
331-
private TestContext createTestContext() {
332-
CalcitePlanContext planContext = createCalcitePlanContext();
333-
return new TestContext(planContext, planContext.relBuilder, planContext.rexBuilder);
334-
}
335-
336-
private void executeRelNodeAndVerify(
337-
CalcitePlanContext planContext, RelNode relNode, ResultVerifier verifier)
338-
throws SQLException {
339-
try (PreparedStatement statement = OpenSearchRelRunners.run(planContext, relNode)) {
340-
ResultSet resultSet = statement.executeQuery();
341-
verifier.verify(resultSet);
342-
}
343-
}
344-
345-
private void verifyColumns(ResultSet resultSet, String... expectedColumnNames)
346-
throws SQLException {
347-
assertEquals(expectedColumnNames.length, resultSet.getMetaData().getColumnCount());
348-
349-
for (int i = 0; i < expectedColumnNames.length; i++) {
350-
String expectedName = expectedColumnNames[i];
351-
String actualName = resultSet.getMetaData().getColumnName(i + 1);
352-
assertEquals(expectedName, actualName);
353-
}
354-
}
355-
356-
private CalcitePlanContext createCalcitePlanContext() {
357-
// Create a Frameworks.ConfigBuilder similar to CalcitePPLAbstractTest
358-
final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
359-
Frameworks.ConfigBuilder config =
360-
Frameworks.newConfigBuilder()
361-
.parserConfig(SqlParser.Config.DEFAULT)
362-
.defaultSchema(rootSchema)
363-
.traitDefs((List<RelTraitDef>) null)
364-
.programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2));
365-
366-
config.context(Contexts.of(RelBuilder.Config.DEFAULT));
367-
368-
return CalcitePlanContext.create(config.build(), SysLimit.DEFAULT, QueryType.PPL);
369-
}
370290
}

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
package org.opensearch.sql.calcite.standalone;
77

8-
import java.io.IOException;
98
import java.sql.ResultSet;
109
import java.sql.SQLException;
1110
import java.util.List;
@@ -23,13 +22,6 @@ public class MapAppendFunctionIT extends CalcitePPLRelNodeIntegTestCase {
2322
private static final String MAP_FIELD = "map";
2423
private static final String ID_FIELD = "id";
2524

26-
@Override
27-
public void init() throws IOException {
28-
super.init();
29-
context = createTestContext();
30-
enableCalcite();
31-
}
32-
3325
@Test
3426
public void testMapAppendWithNonOverlappingKeys() throws Exception {
3527
RexNode map1 = createMap("key1", "value1", "key2", "value2");

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

Lines changed: 1 addition & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -5,46 +5,19 @@
55

66
package org.opensearch.sql.calcite.standalone;
77

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;
138
import java.util.Map;
14-
import org.apache.calcite.plan.Contexts;
15-
import org.apache.calcite.plan.RelTraitDef;
169
import org.apache.calcite.rel.RelNode;
1710
import org.apache.calcite.rel.type.RelDataType;
18-
import org.apache.calcite.rex.RexBuilder;
1911
import org.apache.calcite.rex.RexNode;
20-
import org.apache.calcite.schema.SchemaPlus;
2112
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;
2713
import org.junit.jupiter.api.Test;
28-
import org.opensearch.sql.calcite.CalcitePlanContext;
29-
import org.opensearch.sql.calcite.SysLimit;
30-
import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners;
31-
import org.opensearch.sql.common.setting.Settings;
32-
import org.opensearch.sql.executor.QueryType;
3314
import org.opensearch.sql.expression.function.BuiltinFunctionName;
3415
import org.opensearch.sql.expression.function.PPLFuncImpTable;
3516

36-
public class MapConcatFunctionIT extends CalcitePPLIntegTestCase {
17+
public class MapConcatFunctionIT extends CalcitePPLRelNodeIntegTestCase {
3718

3819
private static final String MAP_FIELD = "map";
3920
private static final String ID_FIELD = "id";
40-
TestContext context;
41-
42-
@Override
43-
public void init() throws IOException {
44-
super.init();
45-
context = createTestContext();
46-
enableCalcite();
47-
}
4821

4922
@Test
5023
public void testMapConcatWithNullValues() throws Exception {
@@ -113,69 +86,4 @@ public void testMapConcat() throws Exception {
11386
assertEquals("value3", result.get("key3"));
11487
});
11588
}
116-
117-
private static class TestContext {
118-
final CalcitePlanContext planContext;
119-
final RelBuilder relBuilder;
120-
final RexBuilder rexBuilder;
121-
122-
TestContext(CalcitePlanContext planContext, RelBuilder relBuilder, RexBuilder rexBuilder) {
123-
this.planContext = planContext;
124-
this.relBuilder = relBuilder;
125-
this.rexBuilder = rexBuilder;
126-
}
127-
}
128-
129-
@FunctionalInterface
130-
private interface ResultVerifier {
131-
void verify(ResultSet resultSet) throws SQLException;
132-
}
133-
134-
private TestContext createTestContext() {
135-
CalcitePlanContext planContext = createCalcitePlanContext();
136-
return new TestContext(planContext, planContext.relBuilder, planContext.rexBuilder);
137-
}
138-
139-
private RelDataType createMapType(RexBuilder rexBuilder) {
140-
RelDataType stringType = rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR);
141-
RelDataType anyType = rexBuilder.getTypeFactory().createSqlType(SqlTypeName.ANY);
142-
return rexBuilder.getTypeFactory().createMapType(stringType, anyType);
143-
}
144-
145-
private void executeRelNodeAndVerify(
146-
CalcitePlanContext planContext, RelNode relNode, ResultVerifier verifier)
147-
throws SQLException {
148-
try (PreparedStatement statement = OpenSearchRelRunners.run(planContext, relNode)) {
149-
ResultSet resultSet = statement.executeQuery();
150-
verifier.verify(resultSet);
151-
}
152-
}
153-
154-
private void verifyColumns(ResultSet resultSet, String... expectedColumnNames)
155-
throws SQLException {
156-
assertEquals(expectedColumnNames.length, resultSet.getMetaData().getColumnCount());
157-
158-
for (int i = 0; i < expectedColumnNames.length; i++) {
159-
String expectedName = expectedColumnNames[i];
160-
String actualName = resultSet.getMetaData().getColumnName(i + 1);
161-
assertEquals(expectedName, actualName);
162-
}
163-
}
164-
165-
private CalcitePlanContext createCalcitePlanContext() {
166-
// Create a Frameworks.ConfigBuilder similar to CalcitePPLAbstractTest
167-
final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
168-
Frameworks.ConfigBuilder config =
169-
Frameworks.newConfigBuilder()
170-
.parserConfig(SqlParser.Config.DEFAULT)
171-
.defaultSchema(rootSchema)
172-
.traitDefs((List<RelTraitDef>) null)
173-
.programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2));
174-
175-
config.context(Contexts.of(RelBuilder.Config.DEFAULT));
176-
177-
Settings settings = getSettings();
178-
return CalcitePlanContext.create(
179-
config.build(), SysLimit.fromSettings(settings), QueryType.PPL);
180-
}
18189
}

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
package org.opensearch.sql.calcite.standalone;
77

8-
import java.io.IOException;
98
import java.sql.ResultSet;
109
import java.sql.SQLException;
1110
import java.util.Map;
@@ -22,13 +21,6 @@ public class MapRemoveFunctionIT extends CalcitePPLRelNodeIntegTestCase {
2221
private static final String MAP_FIELD = "map";
2322
private static final String ID_FIELD = "id";
2423

25-
@Override
26-
public void init() throws IOException {
27-
super.init();
28-
context = createTestContext();
29-
enableCalcite();
30-
}
31-
3224
@Test
3325
public void testMapRemoveWithNullMap() throws Exception {
3426
RelDataType mapType = createMapType(context.rexBuilder);

0 commit comments

Comments
 (0)