Skip to content

Commit a44c763

Browse files
ykmr1224ykmr1224
committed
Fix aggregation for dynamic fields
Signed-off-by: Tomoyuki Morita <[email protected]>
1 parent ac06da4 commit a44c763

File tree

9 files changed

+813
-88
lines changed

9 files changed

+813
-88
lines changed

‎common/src/main/java/org/opensearch/sql/common/utils/DebugUtils.java‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.common.utils;
7+
8+
import java.util.Collection;
9+
import java.util.Map;
10+
import java.util.stream.Collectors;
11+
12+
/** Utility class for debugging operations. Should not be committed. */
13+
public class DebugUtils {
14+
15+
private static void print(String format, Object... args) {
16+
System.out.println(String.format(format, args));
17+
}
18+
19+
public static <T> T debug(T obj, String message) {
20+
print("### %s: %s (at %s)", message, stringify(obj), getCalledFrom(1));
21+
return obj;
22+
}
23+
24+
public static <T> T debug(T obj) {
25+
print("### %s (at %s)", stringify(obj), getCalledFrom(1));
26+
return obj;
27+
}
28+
29+
private static String getCalledFrom(int pos) {
30+
RuntimeException e = new RuntimeException();
31+
StackTraceElement item = e.getStackTrace()[pos + 1];
32+
return item.getClassName() + "." + item.getMethodName() + ":" + item.getLineNumber();
33+
}
34+
35+
private static String stringify(Collection<?> items) {
36+
if (items == null) {
37+
return "null";
38+
}
39+
40+
if (items.isEmpty()) {
41+
return "()";
42+
}
43+
44+
String result = items.stream().map(i -> stringify(i)).collect(Collectors.joining(","));
45+
46+
return "(" + result + ")";
47+
}
48+
49+
private static String stringify(Map<?, ?> map) {
50+
if (map == null) {
51+
return "[[null]]";
52+
}
53+
54+
if (map.isEmpty()) {
55+
return "[[EMPTY]]";
56+
}
57+
58+
String result =
59+
map.entrySet().stream()
60+
.map(entry -> entry.getKey() + ": " + stringify(entry.getValue()))
61+
.collect(Collectors.joining(","));
62+
return "{" + result + "}";
63+
}
64+
65+
private static String stringify(Object obj) {
66+
if (obj instanceof Collection) {
67+
return stringify((Collection) obj);
68+
} else if (obj instanceof Map) {
69+
return stringify((Map) obj);
70+
}
71+
return String.valueOf(obj);
72+
}
73+
}

‎common/src/main/java/org/opensearch/sql/common/utils/JsonUtils.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.common.utils;
7+
8+
import lombok.experimental.UtilityClass;
9+
10+
@UtilityClass
11+
public class JsonUtils {
12+
/**
13+
* Utility method to build JSON string from multiple strings with single-quotes. This is just for
14+
* ease of read and maintain in tests. sjson("{", "'key': 'name'", "}") -> "{\n \"key\":
15+
* \"name\"\n}"
16+
*
17+
* @param lines lines using single-quote instead for double-quote
18+
* @return sting joined inputs and replaces single-quotes with double-quotes
19+
*/
20+
public static String sjson(String... lines) {
21+
StringBuilder builder = new StringBuilder();
22+
for (String line : lines) {
23+
builder.append(replaceQuote(line));
24+
builder.append("\n");
25+
}
26+
return builder.toString();
27+
}
28+
29+
private static String replaceQuote(String line) {
30+
return line.replace("'", "\"");
31+
}
32+
33+
/**
34+
* Utility method to build multiline string from list of strings. Last line will also have new
35+
* line at the end.
36+
*
37+
* @param lines input lines
38+
* @return string contains lines
39+
*/
40+
public static String lines(String... lines) {
41+
StringBuilder builder = new StringBuilder();
42+
for (String line : lines) {
43+
builder.append(line);
44+
builder.append("\n");
45+
}
46+
return builder.toString();
47+
}
48+
}

‎common/src/test/java/org/opensearch/sql/common/utils/JsonUtilsTest.java‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.common.utils;
7+
8+
import static org.junit.Assert.assertEquals;
9+
10+
import org.junit.Test;
11+
12+
public class JsonUtilsTest {
13+
14+
@Test
15+
public void testSjsonWithSingleLine() {
16+
String result = JsonUtils.sjson("{'key': 'value'}");
17+
assertEquals("{\"key\": \"value\"}\n", result);
18+
}
19+
20+
@Test
21+
public void testSjsonWithMultipleLines() {
22+
String result = JsonUtils.sjson("{", " 'name': 'John',", " 'age': 30", "}");
23+
assertEquals("{\n \"name\": \"John\",\n \"age\": 30\n}\n", result);
24+
}
25+
26+
@Test
27+
public void testSjsonWithEmptyString() {
28+
String result = JsonUtils.sjson("");
29+
assertEquals("\n", result);
30+
}
31+
32+
@Test
33+
public void testSjsonWithNoQuotes() {
34+
String result = JsonUtils.sjson("no quotes here");
35+
assertEquals("no quotes here\n", result);
36+
}
37+
38+
@Test
39+
public void testSjsonWithMixedQuotes() {
40+
String result = JsonUtils.sjson("'single' and \"double\" quotes");
41+
assertEquals("\"single\" and \"double\" quotes\n", result);
42+
}
43+
44+
@Test
45+
public void testSjsonWithMultipleSingleQuotes() {
46+
String result = JsonUtils.sjson("'key1': 'value1', 'key2': 'value2'");
47+
assertEquals("\"key1\": \"value1\", \"key2\": \"value2\"\n", result);
48+
}
49+
50+
@Test
51+
public void testSjsonWithNestedJson() {
52+
String result = JsonUtils.sjson("{", " 'outer': {", " 'inner': 'value'", " }", "}");
53+
assertEquals("{\n \"outer\": {\n \"inner\": \"value\"\n }\n}\n", result);
54+
}
55+
56+
@Test
57+
public void testSjsonWithArrays() {
58+
String result = JsonUtils.sjson("{", " 'items': ['item1', 'item2', 'item3']", "}");
59+
assertEquals("{\n \"items\": [\"item1\", \"item2\", \"item3\"]\n}\n", result);
60+
}
61+
62+
@Test
63+
public void testSjsonWithSpecialCharacters() {
64+
String result = JsonUtils.sjson("{'key': 'value with \\'escaped\\' quotes'}");
65+
assertEquals("{\"key\": \"value with \\\"escaped\\\" quotes\"}\n", result);
66+
}
67+
68+
@Test
69+
public void testSjsonWithEmptyArray() {
70+
String result = JsonUtils.sjson();
71+
assertEquals("", result);
72+
}
73+
74+
@Test
75+
public void testSjsonWithNullValues() {
76+
String result = JsonUtils.sjson("{'key': null}");
77+
assertEquals("{\"key\": null}\n", result);
78+
}
79+
80+
@Test
81+
public void testSjsonWithNumbers() {
82+
String result =
83+
JsonUtils.sjson("{", " 'integer': 42,", " 'float': 3.14,", " 'negative': -10", "}");
84+
assertEquals("{\n \"integer\": 42,\n \"float\": 3.14,\n \"negative\": -10\n}\n", result);
85+
}
86+
87+
@Test
88+
public void testSjsonWithBooleans() {
89+
String result = JsonUtils.sjson("{'active': true, 'deleted': false}");
90+
assertEquals("{\"active\": true, \"deleted\": false}\n", result);
91+
}
92+
93+
@Test
94+
public void testSjsonPreservesWhitespace() {
95+
String result = JsonUtils.sjson(" {'key': 'value'} ");
96+
assertEquals(" {\"key\": \"value\"} \n", result);
97+
}
98+
99+
@Test
100+
public void testSjsonWithComplexJson() {
101+
String result =
102+
JsonUtils.sjson(
103+
"{",
104+
" 'user': {",
105+
" 'name': 'Alice',",
106+
" 'email': '[email protected]',",
107+
" 'roles': ['admin', 'user'],",
108+
" 'active': true,",
109+
" 'loginCount': 42",
110+
" }",
111+
"}");
112+
String expected =
113+
"{\n"
114+
+ " \"user\": {\n"
115+
+ " \"name\": \"Alice\",\n"
116+
+ " \"email\": \"[email protected]\",\n"
117+
+ " \"roles\": [\"admin\", \"user\"],\n"
118+
+ " \"active\": true,\n"
119+
+ " \"loginCount\": 42\n"
120+
+ " }\n"
121+
+ "}\n";
122+
assertEquals(expected, result);
123+
}
124+
}

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1897,10 +1897,11 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
18971897
}
18981898

18991899
// 1. group the group-by list + field list and add a count() aggregation
1900-
List<UnresolvedExpression> groupExprList = new ArrayList<>(node.getGroupExprList());
1901-
List<UnresolvedExpression> fieldList =
1902-
node.getFields().stream().map(f -> (UnresolvedExpression) f).toList();
1903-
groupExprList.addAll(fieldList);
1900+
List<UnresolvedExpression> groupExprList = new ArrayList<>();
1901+
node.getGroupExprList().forEach(exp -> groupExprList.add(addAliasToDynamicFieldAccess(exp)));
1902+
// need alias for dynamic fields
1903+
node.getFields()
1904+
.forEach(field -> groupExprList.add(AstDSL.alias(field.getField().toString(), field)));
19041905
List<UnresolvedExpression> aggExprList =
19051906
List.of(AstDSL.alias(countFieldName, AstDSL.aggregate("count", null)));
19061907
aggregateWithTrimming(groupExprList, aggExprList, context);
@@ -2048,6 +2049,11 @@ public RelNode visitTimechart(
20482049
org.opensearch.sql.ast.tree.Timechart node, CalcitePlanContext context) {
20492050
visitChildren(node, context);
20502051

2052+
if (!context.fieldBuilder.isFieldSpecificType("@timestamp")) {
2053+
throw new IllegalArgumentException(
2054+
"`@timestamp` field needs to be specific type. Please cast explicitly.");
2055+
}
2056+
20512057
// Extract parameters
20522058
UnresolvedExpression spanExpr = node.getBinExpression();
20532059

@@ -2089,7 +2095,13 @@ public RelNode visitTimechart(
20892095

20902096
// Extract parameters for byField case
20912097
UnresolvedExpression byField = node.getByField();
2092-
String byFieldName = ((Field) byField).getField().toString();
2098+
2099+
String byFieldName = ((Field) node.getByField()).getField().toString();
2100+
if (!context.fieldBuilder.isFieldSpecificType(byFieldName)) {
2101+
throw new IllegalArgumentException(
2102+
String.format(
2103+
"By field `%s` needs to be specific type. Please cast explicitly.", byFieldName));
2104+
}
20932105
String valueFunctionName = getValueFunctionName(node.getAggregateFunction());
20942106

20952107
int limit = Optional.ofNullable(node.getLimit()).orElse(10);
@@ -2111,9 +2123,7 @@ public RelNode visitTimechart(
21112123
List<RexNode> outputFields = context.fieldBuilder.staticFields();
21122124
List<RexNode> reordered = new ArrayList<>();
21132125
reordered.add(context.fieldBuilder.staticField("@timestamp")); // timestamp first
2114-
reordered.add(
2115-
context.fieldBuilder.staticField(
2116-
byFieldName)); // byField second. TODO: allow dynamic fields
2126+
reordered.add(context.fieldBuilder.staticField(byFieldName)); // byField second.
21172127
reordered.add(outputFields.get(outputFields.size() - 1)); // value function last
21182128
context.relBuilder.project(reordered);
21192129

@@ -2145,6 +2155,14 @@ public RelNode visitTimechart(
21452155
}
21462156
}
21472157

2158+
private UnresolvedExpression addAliasToDynamicFieldAccess(UnresolvedExpression exp) {
2159+
if (exp instanceof Field f) {
2160+
return AstDSL.alias(f.getField().toString(), f);
2161+
} else {
2162+
return exp;
2163+
}
2164+
}
2165+
21482166
/** Build top categories query - simpler approach that works better with OTHER handling */
21492167
private RelNode buildTopCategoriesQuery(
21502168
RelNode completeResults, int limit, CalcitePlanContext context) {
@@ -2363,6 +2381,13 @@ public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
23632381
.forEach(
23642382
trendlineComputation -> {
23652383
RexNode field = rexVisitor.analyze(trendlineComputation.getDataField(), context);
2384+
String dataFieldName = trendlineComputation.getDataField().getField().toString();
2385+
if (!context.fieldBuilder.isFieldSpecificType(dataFieldName)) {
2386+
throw new IllegalArgumentException(
2387+
String.format(
2388+
"`%s` needs to be specific type. Please cast explicitly.", dataFieldName));
2389+
}
2390+
23662391
context.relBuilder.filter(context.relBuilder.isNotNull(field));
23672392

23682393
WindowFrame windowFrame =

‎core/src/main/java/org/opensearch/sql/calcite/rel/RelFieldBuilder.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import java.util.List;
1212
import java.util.stream.Collectors;
1313
import lombok.RequiredArgsConstructor;
14+
import org.apache.calcite.avatica.SqlType;
1415
import org.apache.calcite.rel.type.RelDataTypeField;
1516
import org.apache.calcite.rex.RexBuilder;
1617
import org.apache.calcite.rex.RexInputRef;
@@ -44,6 +45,16 @@ public List<String> getAllFieldNames() {
4445
return getAllFieldNames(0);
4546
}
4647

48+
public boolean isFieldSpecificType(String fieldName) {
49+
List<RelDataTypeField> fields = relBuilder.peek().getRowType().getFieldList();
50+
for (RelDataTypeField field : fields) {
51+
if (field.getName().equals(fieldName)) {
52+
return !field.getType().getSqlTypeName().equals(SqlType.ANY);
53+
}
54+
}
55+
return false;
56+
}
57+
4758
public List<String> getAllFieldNames(int inputCount, int inputOrdinal) {
4859
return getAllFieldNames(getStackPosition(inputCount, inputOrdinal));
4960
}

0 commit comments

Comments
 (0)