Skip to content

Commit 579d44a

Browse files
forestmveyforestmvey
committed
Added indexing for each part of array column name.
Signed-off-by: forestmvey <[email protected]>
1 parent 95e2c82 commit 579d44a

File tree

8 files changed

+121
-68
lines changed

8 files changed

+121
-68
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,4 +105,8 @@ public static String format(final String format, Object... args) {
105105
private static boolean isQuoted(String text, String mark) {
106106
return !Strings.isNullOrEmpty(text) && text.startsWith(mark) && text.endsWith(mark);
107107
}
108+
109+
public static String removeParenthesis(String qualifier) {
110+
return qualifier.replaceAll("\\[.+\\]", "");
111+
}
108112
}

‎core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.util.OptionalInt;
2121
import java.util.stream.Collectors;
2222
import lombok.Getter;
23+
import org.apache.commons.lang3.tuple.Pair;
2324
import org.opensearch.sql.analysis.symbol.Namespace;
2425
import org.opensearch.sql.analysis.symbol.Symbol;
2526
import org.opensearch.sql.ast.AbstractNodeVisitor;
@@ -51,6 +52,7 @@
5152
import org.opensearch.sql.ast.expression.When;
5253
import org.opensearch.sql.ast.expression.WindowFunction;
5354
import org.opensearch.sql.ast.expression.Xor;
55+
import org.opensearch.sql.common.utils.StringUtils;
5456
import org.opensearch.sql.data.model.ExprValueUtils;
5557
import org.opensearch.sql.data.type.ExprCoreType;
5658
import org.opensearch.sql.data.type.ExprType;
@@ -387,7 +389,7 @@ public Expression visitArrayQualifiedName(ArrayQualifiedName node, AnalysisConte
387389
return reserved;
388390
}
389391

390-
return visitArrayIdentifier(qualifierAnalyzer.unqualified(node), context, node.getIndex());
392+
return visitArrayIdentifier(qualifierAnalyzer.unqualified(node), context, node.getPartsAndIndexes());
391393
}
392394

393395
private Expression checkForReservedIdentifier(List<String> parts, AnalysisContext context, QualifierAnalyzer qualifierAnalyzer, QualifiedName node) {
@@ -454,7 +456,8 @@ private Expression visitIdentifier(String ident, AnalysisContext context) {
454456
return ref;
455457
}
456458

457-
private Expression visitArrayIdentifier(String ident, AnalysisContext context, OptionalInt index) {
459+
private Expression visitArrayIdentifier(String ident, AnalysisContext context,
460+
List<Pair<String, OptionalInt>> partsAndIndexes) {
458461
// ParseExpression will always override ReferenceExpression when ident conflicts
459462
for (NamedExpression expr : context.getNamedParseExpressions()) {
460463
if (expr.getNameOrAlias().equals(ident) && expr.getDelegated() instanceof ParseExpression) {
@@ -463,13 +466,7 @@ private Expression visitArrayIdentifier(String ident, AnalysisContext context, O
463466
}
464467

465468
TypeEnvironment typeEnv = context.peek();
466-
ArrayReferenceExpression ref = index.isEmpty() ?
467-
DSL.indexedRef(ident,
468-
typeEnv.resolve(new Symbol(Namespace.FIELD_NAME, ident)))
469-
:
470-
DSL.indexedRef(ident,
471-
typeEnv.resolve(new Symbol(Namespace.FIELD_NAME, ident)), index.getAsInt());
472-
473-
return ref;
469+
return new ArrayReferenceExpression(ident,
470+
typeEnv.resolve(new Symbol(Namespace.FIELD_NAME, StringUtils.removeParenthesis(ident))), partsAndIndexes);
474471
}
475472
}

‎core/src/main/java/org/opensearch/sql/analysis/QualifierAnalyzer.java

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,10 @@
1111
import lombok.RequiredArgsConstructor;
1212
import org.opensearch.sql.analysis.symbol.Namespace;
1313
import org.opensearch.sql.analysis.symbol.Symbol;
14+
import org.opensearch.sql.ast.expression.ArrayQualifiedName;
1415
import org.opensearch.sql.ast.expression.QualifiedName;
1516
import org.opensearch.sql.common.antlr.SyntaxCheckException;
17+
import org.opensearch.sql.common.utils.StringUtils;
1618
import org.opensearch.sql.exception.SemanticCheckException;
1719

1820
/**
@@ -38,6 +40,10 @@ public String unqualified(QualifiedName fullName) {
3840
return isQualifierIndexOrAlias(fullName) ? fullName.rest().toString() : fullName.toString();
3941
}
4042

43+
public String unqualified(ArrayQualifiedName fullName) {
44+
return isQualifierIndexOrAlias(fullName) ? fullName.rest().toString() : fullName.toString();
45+
}
46+
4147
private boolean isQualifierIndexOrAlias(QualifiedName fullName) {
4248
Optional<String> qualifier = fullName.first();
4349
if (qualifier.isPresent()) {
@@ -50,10 +56,22 @@ private boolean isQualifierIndexOrAlias(QualifiedName fullName) {
5056
return false;
5157
}
5258

59+
private boolean isQualifierIndexOrAlias(ArrayQualifiedName fullName) {
60+
Optional<String> qualifier = fullName.first();
61+
if (qualifier.isPresent()) {
62+
if (isFieldName(qualifier.get())) {
63+
return false;
64+
}
65+
resolveQualifierSymbol(fullName, qualifier.get());
66+
return true;
67+
}
68+
return false;
69+
}
70+
5371
private boolean isFieldName(String qualifier) {
5472
try {
5573
// Resolve the qualifier in Namespace.FIELD_NAME
56-
context.peek().resolve(new Symbol(Namespace.FIELD_NAME, qualifier));
74+
context.peek().resolve(new Symbol(Namespace.FIELD_NAME, StringUtils.removeParenthesis(qualifier)));
5775
return true;
5876
} catch (SemanticCheckException e2) {
5977
return false;

‎core/src/main/java/org/opensearch/sql/ast/expression/ArrayQualifiedName.java

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,24 +9,22 @@
99

1010
import lombok.EqualsAndHashCode;
1111
import lombok.Getter;
12+
import org.apache.commons.lang3.tuple.Pair;
1213
import org.opensearch.sql.ast.AbstractNodeVisitor;
1314

15+
import java.util.List;
1416
import java.util.OptionalInt;
17+
import java.util.stream.Collectors;
1518

1619
@Getter
1720
@EqualsAndHashCode(callSuper = false)
1821
public class ArrayQualifiedName extends QualifiedName {
1922

20-
private final OptionalInt index;
23+
private final List<Pair<String, OptionalInt>> partsAndIndexes;
2124

22-
public ArrayQualifiedName(String name, int index) {
23-
super(name);
24-
this.index = OptionalInt.of(index);
25-
}
26-
27-
public ArrayQualifiedName(String name) {
28-
super(name);
29-
this.index = OptionalInt.empty();
25+
public ArrayQualifiedName(List<Pair<String, OptionalInt>> parts) {
26+
super(parts.stream().map(p -> p.getLeft()).collect(Collectors.toList()));
27+
this.partsAndIndexes = parts;
3028
}
3129

3230
@Override

‎core/src/main/java/org/opensearch/sql/expression/ArrayReferenceExpression.java

Lines changed: 19 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,12 @@
1111
import java.util.Arrays;
1212
import java.util.List;
1313
import java.util.OptionalInt;
14+
import java.util.stream.Collectors;
1415

1516
import lombok.EqualsAndHashCode;
1617
import lombok.Getter;
18+
import org.apache.commons.lang3.tuple.Pair;
19+
import org.opensearch.sql.common.utils.StringUtils;
1720
import org.opensearch.sql.data.model.ExprTupleValue;
1821
import org.opensearch.sql.data.model.ExprValue;
1922
import org.opensearch.sql.data.model.ExprValueUtils;
@@ -24,30 +27,20 @@
2427
@EqualsAndHashCode
2528
public class ArrayReferenceExpression extends ReferenceExpression {
2629
@Getter
27-
private final OptionalInt index; // Should be a list of indexes to support multiple nesting levels
30+
private final List<Pair<String, OptionalInt>> partsAndIndexes;
2831
@Getter
2932
private final ExprType type;
30-
@Getter
31-
private final List<String> paths;
32-
public ArrayReferenceExpression(String ref, ExprType type, int index) {
33-
super(ref, type);
34-
this.index = OptionalInt.of(index);
35-
this.type = type;
36-
this.paths = Arrays.asList(ref.split("\\."));
37-
}
38-
39-
public ArrayReferenceExpression(String ref, ExprType type) {
33+
public ArrayReferenceExpression(String ref, ExprType type, List<Pair<String, OptionalInt>> partsAndIndexes) {
4034
super(ref, type);
41-
this.index = OptionalInt.empty();
35+
this.partsAndIndexes = partsAndIndexes;
4236
this.type = type;
43-
this.paths = Arrays.asList(ref.split("\\."));
4437
}
4538

4639
public ArrayReferenceExpression(ReferenceExpression ref) {
4740
super(ref.toString(), ref.type());
48-
this.index = OptionalInt.empty();
41+
this.partsAndIndexes = Arrays.stream(ref.toString().split("\\.")).map(e -> Pair.of(e, OptionalInt.empty())).collect(
42+
Collectors.toList());
4943
this.type = ref.type();
50-
this.paths = Arrays.asList(ref.toString().split("\\."));
5144
}
5245

5346
@Override
@@ -66,13 +59,15 @@ public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
6659
}
6760

6861
public ExprValue resolve(ExprTupleValue value) {
69-
return resolve(value, paths);
62+
return resolve(value, partsAndIndexes);
7063
}
7164

72-
private ExprValue resolve(ExprValue value, List<String> paths) {
73-
ExprValue wholePathValue = value.keyValue(String.join(PATH_SEP, paths));
65+
private ExprValue resolve(ExprValue value, List<Pair<String, OptionalInt>> paths) {
66+
List<String> pathsWithoutParenthesis =
67+
paths.stream().map(p -> StringUtils.removeParenthesis(p.getLeft())).collect(Collectors.toList());
68+
ExprValue wholePathValue = value.keyValue(String.join(PATH_SEP, pathsWithoutParenthesis));
7469

75-
if (!index.isEmpty()) {
70+
if (!paths.get(0).getRight().isEmpty()) {
7671
wholePathValue = getIndexValueOfCollection(wholePathValue);
7772
} else if (paths.size() == 1 && !wholePathValue.type().equals(ExprCoreType.ARRAY)) {
7873
wholePathValue = ExprValueUtils.missingValue();
@@ -81,19 +76,20 @@ private ExprValue resolve(ExprValue value, List<String> paths) {
8176
if (!wholePathValue.isMissing() || paths.size() == 1) {
8277
return wholePathValue;
8378
} else {
84-
return resolve(value.keyValue(paths.get(0)), paths.subList(1, paths.size()));
79+
return resolve(value.keyValue(pathsWithoutParenthesis.get(0)
80+
), paths.subList(1, paths.size()));
8581
}
8682
}
8783

8884
private ExprValue getIndexValueOfCollection(ExprValue value) {
8985
ExprValue collectionVal = value;
90-
for (OptionalInt currentIndex : List.of(index)) {
86+
// for (OptionalInt currentIndex : List.of(index)) {
9187
if (collectionVal.type().equals(ExprCoreType.ARRAY)) {
92-
collectionVal = collectionVal.collectionValue().get(currentIndex.getAsInt());
88+
// collectionVal = collectionVal.collectionValue().get(currentIndex.getAsInt());
9389
} else {
9490
return ExprValueUtils.missingValue();
9591
}
96-
}
92+
// }
9793
return collectionVal;
9894
}
9995
}

‎core/src/main/java/org/opensearch/sql/expression/DSL.java

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -87,14 +87,6 @@ public static ReferenceExpression ref(String ref, ExprType type) {
8787
return new ReferenceExpression(ref, type);
8888
}
8989

90-
public static ArrayReferenceExpression indexedRef(String ref, ExprType type, int index) {
91-
return new ArrayReferenceExpression(ref, type, index);
92-
}
93-
94-
public static ArrayReferenceExpression indexedRef(String ref, ExprType type) {
95-
return new ArrayReferenceExpression(ref, type);
96-
}
97-
9890
/**
9991
* Wrap a named expression if not yet. The intent is that different languages may use
10092
* Alias or not when building AST. This caused either named or unnamed expression

‎sql/src/main/antlr/OpenSearchSQLParser.g4

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,7 @@ expressions
303303
expressionAtom
304304
: constant #constantExpressionAtom
305305
| columnName #fullColumnNameExpressionAtom
306-
| arrayColumnName #arrayColumnNameExpressionAtom
306+
// | arrayColumnName #arrayColumnNameExpressionAtom
307307
| functionCall #functionCallExpressionAtom
308308
| LR_BRACKET expression RR_BRACKET #nestedExpressionAtom
309309
| left=expressionAtom
@@ -815,10 +815,11 @@ columnName
815815
: qualifiedName
816816
;
817817

818-
arrayColumnName
819-
: qualifiedName LT_SQR_PRTHS COLON_SYMB RT_SQR_PRTHS
820-
| qualifiedName LT_SQR_PRTHS decimalLiteral RT_SQR_PRTHS
821-
;
818+
//arrayColumnName
819+
// : arrayQualifiedName
820+
// | qualifiedName LT_SQR_PRTHS COLON_SYMB RT_SQR_PRTHS (arrayColumnName | columnName) *
821+
// | qualifiedName LT_SQR_PRTHS decimalLiteral RT_SQR_PRTHS (arrayColumnName | columnName) *
822+
// ;
822823

823824
allTupleFields
824825
: path=qualifiedName DOT STAR
@@ -828,12 +829,21 @@ alias
828829
: ident
829830
;
830831

832+
//arrayQualifiedName
833+
// : (ident | indexedIdentifier) (DOT (ident | indexedIdentifier))*
834+
//// (( LT_SQR_PRTHS decimalLiteral RT_SQR_PRTHS ) | (DOT ident ( LT_SQR_PRTHS decimalLiteral RT_SQR_PRTHS )* ))+
835+
// ;
836+
837+
//indexedIdentifier
838+
// : ident LT_SQR_PRTHS decimalLiteral RT_SQR_PRTHS
839+
// ;
840+
831841
qualifiedName
832842
: ident (DOT ident)*
833843
;
834844

835845
ident
836-
: DOT? ID
846+
: DOT? ID (LT_SQR_PRTHS decimalLiteral RT_SQR_PRTHS)?
837847
| BACKTICK_QUOTE_ID
838848
| keywordsCanBeId
839849
| scalarFunctionName

‎sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java

Lines changed: 49 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,14 @@
6969

7070
import com.google.common.collect.ImmutableList;
7171
import com.google.common.collect.ImmutableMap;
72+
73+
import java.util.ArrayList;
7274
import java.util.Arrays;
7375
import java.util.Collections;
7476
import java.util.HashMap;
7577
import java.util.List;
7678
import java.util.Map;
79+
import java.util.OptionalInt;
7780
import java.util.stream.Collectors;
7881
import org.antlr.v4.runtime.RuleContext;
7982
import org.apache.commons.lang3.tuple.ImmutablePair;
@@ -106,7 +109,6 @@
106109
import org.opensearch.sql.expression.function.BuiltinFunctionName;
107110
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.AlternateMultiMatchQueryContext;
108111
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.AndExpressionContext;
109-
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ArrayColumnNameContext;
110112
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ColumnNameContext;
111113
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.IdentContext;
112114
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.IntervalLiteralContext;
@@ -130,16 +132,26 @@ public UnresolvedExpression visitColumnName(ColumnNameContext ctx) {
130132
return visit(ctx.qualifiedName());
131133
}
132134

133-
@Override
134-
public UnresolvedExpression visitArrayColumnName(ArrayColumnNameContext ctx) {
135-
UnresolvedExpression qualifiedName = visit(ctx.qualifiedName());
136-
if (ctx.decimalLiteral() == null) {
137-
return new ArrayQualifiedName(qualifiedName.toString());
138-
} else {
139-
return new ArrayQualifiedName(
140-
qualifiedName.toString(), Integer.parseInt(ctx.decimalLiteral().getText()));
141-
}
142-
}
135+
// @Override
136+
// public UnresolvedExpression visitArrayColumnName(ArrayColumnNameContext ctx) {
137+
//// return new QualifiedName(
138+
//// identifiers.stream()
139+
//// .map(RuleContext::getText)
140+
//// .map(StringUtils::unquoteIdentifier)
141+
//// .collect(Collectors.toList()));
142+
//
143+
// var blah = ctx.arrayQualifiedName().indexedIdentifier();
144+
// var hmm = ctx.arrayQualifiedName().ident();
145+
//
146+
//
147+
// UnresolvedExpression qualifiedName = visit(ctx.arrayQualifiedName());
148+
//// if (ctx.arrayQualifiedName().decimalLiteral() == null) {
149+
// return new ArrayQualifiedName(qualifiedName.toString());
150+
//// } else {
151+
// return new ArrayQualifiedName(
152+
// qualifiedName.toString(), Integer.parseInt(ctx.arrayQualifiedName().decimalLiteral().toString()));
153+
//// }
154+
// }
143155

144156
@Override
145157
public UnresolvedExpression visitIdent(IdentContext ctx) {
@@ -544,6 +556,32 @@ public UnresolvedExpression visitExtractFunctionCall(ExtractFunctionCallContext
544556

545557

546558
private QualifiedName visitIdentifiers(List<IdentContext> identifiers) {
559+
560+
List<Pair<String, OptionalInt>> parts = new ArrayList<>();
561+
boolean supportsArrays = false;
562+
for(var blah : identifiers) {
563+
if (blah.decimalLiteral() != null) {
564+
parts.add(
565+
Pair.of(
566+
StringUtils.unquoteIdentifier(blah.getText()),
567+
OptionalInt.of(Integer.parseInt(blah.decimalLiteral().getText()))
568+
)
569+
);
570+
supportsArrays = true;
571+
} else {
572+
parts.add(
573+
Pair.of(
574+
StringUtils.unquoteIdentifier(blah.getText()),
575+
OptionalInt.empty()
576+
)
577+
);
578+
}
579+
}
580+
581+
if (supportsArrays) {
582+
return new ArrayQualifiedName(parts);
583+
}
584+
547585
return new QualifiedName(
548586
identifiers.stream()
549587
.map(RuleContext::getText)

0 commit comments

Comments
 (0)