Skip to content

Commit 13b64c2

Browse files
xinyualxinyual
committed
merge from main
Signed-off-by: xinyual <[email protected]>
2 parents adb8997 + 1841139 commit 13b64c2

File tree

60 files changed

+3004
-384
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

60 files changed

+3004
-384
lines changed

‎async-query-core/src/main/antlr/OpenSearchPPLParser.g4

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -868,7 +868,6 @@ jsonFunctionName
868868
| JSON_OBJECT
869869
| JSON_ARRAY
870870
| JSON_ARRAY_LENGTH
871-
| TO_JSON_STRING
872871
| JSON_EXTRACT
873872
| JSON_KEYS
874873
| JSON_VALID

‎benchmarks/build.gradle

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ repositories {
1414

1515
dependencies {
1616
implementation project(':core')
17+
implementation project(':opensearch')
1718

1819
// Dependencies required by JMH micro benchmark
1920
api group: 'org.openjdk.jmh', name: 'jmh-core', version: '1.36'

‎benchmarks/src/jmh/java/org/opensearch/sql/expression/operator/predicate/MergeArrayAndObjectMapBenchmark.java

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.expression.operator.predicate;
7+
8+
import java.util.ArrayList;
9+
import java.util.HashMap;
10+
import java.util.LinkedHashMap;
11+
import java.util.List;
12+
import java.util.Map;
13+
import org.openjdk.jmh.annotations.Benchmark;
14+
import org.opensearch.sql.opensearch.data.type.OpenSearchDataType;
15+
import org.opensearch.sql.opensearch.request.system.OpenSearchDescribeIndexRequest;
16+
17+
public class MergeArrayAndObjectMapBenchmark {
18+
private static final List<Map<String, OpenSearchDataType>> candidateMaps = prepareListOfMaps(120);
19+
20+
@Benchmark
21+
public void testMerge() {
22+
Map<String, OpenSearchDataType> finalResult = new HashMap<>();
23+
for (Map<String, OpenSearchDataType> map : candidateMaps) {
24+
OpenSearchDescribeIndexRequest.mergeObjectAndArrayInsideMap(finalResult, map);
25+
}
26+
}
27+
28+
private static Map<String, OpenSearchDataType> prepareMap(int recursive, String prefix) {
29+
Map<String, OpenSearchDataType> map = new HashMap<>();
30+
Map<String, Object> innerMap = prepareRecursiveMap(recursive, prefix);
31+
map.put("name", OpenSearchDataType.of(OpenSearchDataType.MappingType.Object, innerMap));
32+
return map;
33+
}
34+
35+
public static Map<String, Object> prepareRecursiveMap(int recursive, String prefix) {
36+
Map<String, Object> innerMap = new LinkedHashMap<>();
37+
if (recursive == 0) {
38+
innerMap.put("type", "string");
39+
} else {
40+
innerMap.put("type", "object");
41+
innerMap.put(
42+
"properties",
43+
Map.of(
44+
prefix + "_" + String.valueOf(recursive),
45+
Map.of("type", "text"),
46+
"recursive",
47+
prepareRecursiveMap(recursive - 1, prefix)));
48+
}
49+
return innerMap;
50+
}
51+
52+
private static List<Map<String, OpenSearchDataType>> prepareListOfMaps(int listNumber) {
53+
List<Map<String, OpenSearchDataType>> list = new ArrayList<>();
54+
for (int i = 0; i < listNumber; i++) {
55+
list.add(prepareMap(15, "prefix" + i));
56+
}
57+
return list;
58+
}
59+
}

‎benchmarks/src/jmh/java/org/opensearch/sql/expression/operator/predicate/PatternsWindowFunctionBenchmark.java

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
import org.opensearch.sql.expression.ReferenceExpression;
3232
import org.opensearch.sql.expression.window.WindowDefinition;
3333
import org.opensearch.sql.expression.window.frame.BufferPatternRowsWindowFrame;
34-
import org.opensearch.sql.expression.window.frame.CurrentRowWindowFrame;
3534
import org.opensearch.sql.expression.window.frame.WindowFrame;
3635

3736
@Warmup(iterations = 1)
@@ -62,14 +61,6 @@ public class PatternsWindowFunctionBenchmark {
6261
new BrainLogParser(),
6362
new NamedArgumentExpression("message", new ReferenceExpression("message", STRING)));
6463

65-
@Benchmark
66-
public void testSimplePattern() {
67-
CurrentRowWindowFrame windowFrame =
68-
new CurrentRowWindowFrame(new WindowDefinition(ImmutableList.of(), ImmutableList.of()));
69-
70-
run(windowFrame, DSL.simple_pattern(DSL.ref("message", STRING)));
71-
}
72-
7364
@Benchmark
7465
public void testBrain() {
7566
BufferPatternRowsWindowFrame windowFrame =

‎build.gradle

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ allprojects {
129129
}
130130
configurations.all {
131131
resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib:1.9.10"
132+
resolutionStrategy.force "net.minidev:json-smart:${versions.json_smart}"
132133
resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib-common:1.9.10"
133134
resolutionStrategy.force "net.bytebuddy:byte-buddy:1.14.19"
134135
resolutionStrategy.force "org.apache.httpcomponents.client5:httpclient5:${versions.httpclient5}"

‎core/build.gradle

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,8 @@ dependencies {
6363
api 'org.apache.calcite:calcite-linq4j:1.38.0'
6464
api project(':common')
6565
implementation "com.github.seancfoley:ipaddress:5.4.2"
66+
implementation "com.jayway.jsonpath:json-path:2.9.0"
67+
implementation "com.googlecode.aviator:aviator:5.4.3"
6668

6769
annotationProcessor('org.immutables:value:2.8.8')
6870
compileOnly('org.immutables:value-annotations:2.8.8')

‎core/src/main/java/org/opensearch/sql/ast/tree/Trendline.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
6666
}
6767

6868
public enum TrendlineType {
69-
SMA
69+
SMA,
70+
WMA
7071
}
7172
}

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

Lines changed: 124 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
import org.opensearch.sql.ast.expression.ParseMethod;
6060
import org.opensearch.sql.ast.expression.UnresolvedExpression;
6161
import org.opensearch.sql.ast.expression.WindowFrame;
62+
import org.opensearch.sql.ast.expression.WindowFrame.FrameType;
6263
import org.opensearch.sql.ast.expression.subquery.SubqueryExpression;
6364
import org.opensearch.sql.ast.tree.AD;
6465
import org.opensearch.sql.ast.tree.Aggregation;
@@ -86,11 +87,13 @@
8687
import org.opensearch.sql.ast.tree.SubqueryAlias;
8788
import org.opensearch.sql.ast.tree.TableFunction;
8889
import org.opensearch.sql.ast.tree.Trendline;
90+
import org.opensearch.sql.ast.tree.Trendline.TrendlineType;
8991
import org.opensearch.sql.ast.tree.UnresolvedPlan;
9092
import org.opensearch.sql.ast.tree.Window;
9193
import org.opensearch.sql.calcite.plan.OpenSearchConstants;
9294
import org.opensearch.sql.calcite.utils.JoinAndLookupUtils;
9395
import org.opensearch.sql.calcite.utils.PlanUtils;
96+
import org.opensearch.sql.common.utils.StringUtils;
9497
import org.opensearch.sql.exception.CalciteUnsupportedException;
9598
import org.opensearch.sql.exception.SemanticCheckException;
9699
import org.opensearch.sql.expression.function.BuiltinFunctionName;
@@ -975,6 +978,126 @@ public RelNode visitTableFunction(TableFunction node, CalcitePlanContext context
975978

976979
@Override
977980
public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
978-
throw new CalciteUnsupportedException("Trendline command is unsupported in Calcite");
981+
visitChildren(node, context);
982+
983+
node.getSortByField()
984+
.ifPresent(
985+
sortField -> {
986+
SortOption sortOption = analyzeSortOption(sortField.getFieldArgs());
987+
RexNode field = rexVisitor.analyze(sortField, context);
988+
if (sortOption == DEFAULT_DESC) {
989+
context.relBuilder.sort(context.relBuilder.desc(field));
990+
} else {
991+
context.relBuilder.sort(field);
992+
}
993+
});
994+
995+
List<RexNode> trendlineNodes = new ArrayList<>();
996+
List<String> aliases = new ArrayList<>();
997+
node.getComputations()
998+
.forEach(
999+
trendlineComputation -> {
1000+
RexNode field = rexVisitor.analyze(trendlineComputation.getDataField(), context);
1001+
context.relBuilder.filter(context.relBuilder.isNotNull(field));
1002+
1003+
WindowFrame windowFrame =
1004+
WindowFrame.of(
1005+
FrameType.ROWS,
1006+
StringUtils.format(
1007+
"%d PRECEDING", trendlineComputation.getNumberOfDataPoints() - 1),
1008+
"CURRENT ROW");
1009+
RexNode countExpr =
1010+
PlanUtils.makeOver(
1011+
context,
1012+
BuiltinFunctionName.COUNT,
1013+
null,
1014+
List.of(),
1015+
List.of(),
1016+
List.of(),
1017+
windowFrame);
1018+
// CASE WHEN count() over (ROWS (windowSize-1) PRECEDING) > windowSize - 1
1019+
RexNode whenConditionExpr =
1020+
PPLFuncImpTable.INSTANCE.resolve(
1021+
context.rexBuilder,
1022+
">",
1023+
countExpr,
1024+
context.relBuilder.literal(trendlineComputation.getNumberOfDataPoints() - 1));
1025+
1026+
RexNode thenExpr;
1027+
switch (trendlineComputation.getComputationType()) {
1028+
case TrendlineType.SMA:
1029+
// THEN avg(field) over (ROWS (windowSize-1) PRECEDING)
1030+
thenExpr =
1031+
PlanUtils.makeOver(
1032+
context,
1033+
BuiltinFunctionName.AVG,
1034+
field,
1035+
List.of(),
1036+
List.of(),
1037+
List.of(),
1038+
windowFrame);
1039+
break;
1040+
case TrendlineType.WMA:
1041+
// THEN wma expression
1042+
thenExpr =
1043+
buildWmaRexNode(
1044+
field,
1045+
trendlineComputation.getNumberOfDataPoints(),
1046+
windowFrame,
1047+
context);
1048+
break;
1049+
default:
1050+
throw new IllegalStateException("Unsupported trendline type");
1051+
}
1052+
1053+
// ELSE NULL
1054+
RexNode elseExpr = context.relBuilder.literal(null);
1055+
1056+
List<RexNode> caseOperands = new ArrayList<>();
1057+
caseOperands.add(whenConditionExpr);
1058+
caseOperands.add(thenExpr);
1059+
caseOperands.add(elseExpr);
1060+
RexNode trendlineNode =
1061+
context.rexBuilder.makeCall(SqlStdOperatorTable.CASE, caseOperands);
1062+
trendlineNodes.add(trendlineNode);
1063+
aliases.add(trendlineComputation.getAlias());
1064+
});
1065+
1066+
projectPlusOverriding(trendlineNodes, aliases, context);
1067+
return context.relBuilder.peek();
1068+
}
1069+
1070+
private RexNode buildWmaRexNode(
1071+
RexNode field,
1072+
Integer numberOfDataPoints,
1073+
WindowFrame windowFrame,
1074+
CalcitePlanContext context) {
1075+
1076+
// Divisor: 1 + 2 + 3 + ... + windowSize, aka (windowSize * (windowSize + 1) / 2)
1077+
RexNode divisor = context.relBuilder.literal(numberOfDataPoints * (numberOfDataPoints + 1) / 2);
1078+
1079+
// Divider: 1 * NTH_VALUE(field, 1) + 2 * NTH_VALUE(field, 2) + ... + windowSize *
1080+
// NTH_VALUE(field, windowSize)
1081+
RexNode divider = context.relBuilder.literal(0);
1082+
for (int i = 1; i <= numberOfDataPoints; i++) {
1083+
RexNode nthValueExpr =
1084+
PlanUtils.makeOver(
1085+
context,
1086+
BuiltinFunctionName.NTH_VALUE,
1087+
field,
1088+
List.of(context.relBuilder.literal(i)),
1089+
List.of(),
1090+
List.of(),
1091+
windowFrame);
1092+
divider =
1093+
context.relBuilder.call(
1094+
SqlStdOperatorTable.PLUS,
1095+
divider,
1096+
context.relBuilder.call(
1097+
SqlStdOperatorTable.MULTIPLY, nthValueExpr, context.relBuilder.literal(i)));
1098+
}
1099+
// Divider / CAST(Divisor, DOUBLE)
1100+
return context.relBuilder.call(
1101+
SqlStdOperatorTable.DIVIDE, divider, context.relBuilder.cast(divisor, SqlTypeName.DOUBLE));
9791102
}
9801103
}

‎core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java

Lines changed: 10 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,6 @@
1010
import static org.apache.calcite.rex.RexWindowBounds.UNBOUNDED_PRECEDING;
1111
import static org.apache.calcite.rex.RexWindowBounds.following;
1212
import static org.apache.calcite.rex.RexWindowBounds.preceding;
13-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.STDDEV_POP_NULLABLE;
14-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.STDDEV_SAMP_NULLABLE;
15-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.VAR_POP_NULLABLE;
16-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.VAR_SAMP_NULLABLE;
17-
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.TransferUserDefinedAggFunction;
1813

1914
import com.google.common.collect.ImmutableList;
2015
import java.util.ArrayList;
@@ -25,7 +20,6 @@
2520
import org.apache.calcite.rex.RexVisitorImpl;
2621
import org.apache.calcite.rex.RexWindowBound;
2722
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
28-
import org.apache.calcite.sql.type.ReturnTypes;
2923
import org.apache.calcite.sql.type.SqlTypeName;
3024
import org.apache.calcite.tools.RelBuilder;
3125
import org.opensearch.sql.ast.AbstractNodeVisitor;
@@ -37,9 +31,8 @@
3731
import org.opensearch.sql.ast.tree.Relation;
3832
import org.opensearch.sql.ast.tree.UnresolvedPlan;
3933
import org.opensearch.sql.calcite.CalcitePlanContext;
40-
import org.opensearch.sql.calcite.udf.udaf.PercentileApproxFunction;
41-
import org.opensearch.sql.calcite.udf.udaf.TakeAggFunction;
4234
import org.opensearch.sql.expression.function.BuiltinFunctionName;
35+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
4336

4437
public interface PlanUtils {
4538

@@ -116,6 +109,14 @@ static RexNode makeOver(
116109
true,
117110
lowerBound,
118111
upperBound);
112+
case NTH_VALUE:
113+
return withOver(
114+
context.relBuilder.aggregateCall(SqlStdOperatorTable.NTH_VALUE, field, argList.get(0)),
115+
partitions,
116+
orderKeys,
117+
true,
118+
lowerBound,
119+
upperBound);
119120
default:
120121
return withOver(
121122
makeAggCall(context, functionName, false, field, argList),
@@ -232,56 +233,7 @@ static RelBuilder.AggCall makeAggCall(
232233
boolean distinct,
233234
RexNode field,
234235
List<RexNode> argList) {
235-
switch (functionName) {
236-
case MAX:
237-
return context.relBuilder.max(field);
238-
case MIN:
239-
return context.relBuilder.min(field);
240-
case AVG:
241-
return context.relBuilder.avg(distinct, null, field);
242-
case COUNT:
243-
return context.relBuilder.count(
244-
distinct, null, field == null ? ImmutableList.of() : ImmutableList.of(field));
245-
case SUM:
246-
return context.relBuilder.sum(distinct, null, field);
247-
// case MEAN:
248-
// throw new UnsupportedOperationException("MEAN is not supported in PPL");
249-
// case STDDEV:
250-
// return context.relBuilder.aggregateCall(SqlStdOperatorTable.STDDEV,
251-
// field);
252-
case VARSAMP:
253-
return context.relBuilder.aggregateCall(VAR_SAMP_NULLABLE, field);
254-
case VARPOP:
255-
return context.relBuilder.aggregateCall(VAR_POP_NULLABLE, field);
256-
case STDDEV_POP:
257-
return context.relBuilder.aggregateCall(STDDEV_POP_NULLABLE, field);
258-
case STDDEV_SAMP:
259-
return context.relBuilder.aggregateCall(STDDEV_SAMP_NULLABLE, field);
260-
// case PERCENTILE_APPROX:
261-
// return
262-
// context.relBuilder.aggregateCall(SqlStdOperatorTable.PERCENTILE_CONT, field);
263-
case TAKE:
264-
return TransferUserDefinedAggFunction(
265-
TakeAggFunction.class,
266-
"TAKE",
267-
UserDefinedFunctionUtils.getReturnTypeInferenceForArray(),
268-
List.of(field),
269-
argList,
270-
context.relBuilder);
271-
case PERCENTILE_APPROX:
272-
List<RexNode> newArgList = new ArrayList<>(argList);
273-
newArgList.add(context.rexBuilder.makeFlag(field.getType().getSqlTypeName()));
274-
return TransferUserDefinedAggFunction(
275-
PercentileApproxFunction.class,
276-
"percentile_approx",
277-
ReturnTypes.ARG0_FORCE_NULLABLE,
278-
List.of(field),
279-
newArgList,
280-
context.relBuilder);
281-
default:
282-
throw new UnsupportedOperationException(
283-
"Unexpected aggregation: " + functionName.getName().getFunctionName());
284-
}
236+
return PPLFuncImpTable.INSTANCE.resolveAgg(functionName, distinct, field, argList, context);
285237
}
286238

287239
/** Get all uniq input references from a RexNode. */

‎core/src/main/java/org/opensearch/sql/calcite/utils/UserDefinedFunctionUtils.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ public static RelBuilder.AggCall TransferUserDefinedAggFunction(
7575
return relBuilder.aggregateCall(sqlUDAF, addArgList);
7676
}
7777

78-
static SqlReturnTypeInference getReturnTypeInferenceForArray() {
78+
public static SqlReturnTypeInference getReturnTypeInferenceForArray() {
7979
return opBinding -> {
8080
RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
8181

0 commit comments

Comments
 (0)