Skip to content

Commit 48553f3

Browse files
forestmveyforestmvey
committed
Adding nested function support in SELECT clause.
Signed-off-by: forestmvey <forestv@bitquilltech.com>
1 parent 23cc0f6 commit 48553f3

66 files changed

Lines changed: 2187 additions & 107 deletions

File tree

Some content is hidden

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

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,14 @@ public LogicalPlan visitProject(Project node, AnalysisContext context) {
359359
List<NamedExpression> namedExpressions =
360360
selectExpressionAnalyzer.analyze(node.getProjectList(), context,
361361
new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child));
362+
363+
for (UnresolvedExpression expr : node.getProjectList()) {
364+
NestedAnalyzer nestedAnalyzer = new NestedAnalyzer(
365+
namedExpressions, expressionAnalyzer, child
366+
);
367+
child = nestedAnalyzer.analyze(expr, context);
368+
}
369+
362370
// new context
363371
context.push();
364372
TypeEnvironment newEnv = context.peek();

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

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -325,17 +325,6 @@ private Expression visitIdentifier(String ident, AnalysisContext context) {
325325
ReferenceExpression ref = DSL.ref(ident,
326326
typeEnv.resolve(new Symbol(Namespace.FIELD_NAME, ident)));
327327

328-
// Fall back to old engine too if type is not supported semantically
329-
if (isTypeNotSupported(ref.type())) {
330-
throw new SyntaxCheckException(String.format(
331-
"Identifier [%s] of type [%s] is not supported yet", ident, ref.type()));
332-
}
333328
return ref;
334329
}
335-
336-
// Array type is not supporte yet.
337-
private boolean isTypeNotSupported(ExprType type) {
338-
return "array".equalsIgnoreCase(type.typeName());
339-
}
340-
341330
}

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.analysis;
7+
8+
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
9+
10+
import java.util.List;
11+
import java.util.Map;
12+
import lombok.RequiredArgsConstructor;
13+
import org.opensearch.sql.ast.AbstractNodeVisitor;
14+
import org.opensearch.sql.ast.expression.Alias;
15+
import org.opensearch.sql.ast.expression.Function;
16+
import org.opensearch.sql.ast.expression.UnresolvedExpression;
17+
import org.opensearch.sql.expression.NamedExpression;
18+
import org.opensearch.sql.expression.ReferenceExpression;
19+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
20+
import org.opensearch.sql.planner.logical.LogicalNested;
21+
import org.opensearch.sql.planner.logical.LogicalPlan;
22+
23+
/**
24+
* Analyze the Nested Function in the {@link AnalysisContext} to construct the {@link
25+
* LogicalPlan}.
26+
*/
27+
@RequiredArgsConstructor
28+
public class NestedAnalyzer extends AbstractNodeVisitor<LogicalPlan, AnalysisContext> {
29+
private final List<NamedExpression> namedExpressions;
30+
private final ExpressionAnalyzer expressionAnalyzer;
31+
private final LogicalPlan child;
32+
33+
public LogicalPlan analyze(UnresolvedExpression projectItem, AnalysisContext context) {
34+
LogicalPlan nested = projectItem.accept(this, context);
35+
return (nested == null) ? child : nested;
36+
}
37+
38+
@Override
39+
public LogicalPlan visitAlias(Alias node, AnalysisContext context) {
40+
return node.getDelegated().accept(this, context);
41+
}
42+
43+
@Override
44+
public LogicalPlan visitFunction(Function node, AnalysisContext context) {
45+
if (node.getFuncName().equalsIgnoreCase(BuiltinFunctionName.NESTED.name())) {
46+
47+
List<UnresolvedExpression> expressions = node.getFuncArgs();
48+
ReferenceExpression nestedField =
49+
(ReferenceExpression)expressionAnalyzer.analyze(expressions.get(0), context);
50+
Map<String, ReferenceExpression> args;
51+
if (expressions.size() == 2) {
52+
args = Map.of(
53+
"field", nestedField,
54+
"path", (ReferenceExpression)expressionAnalyzer.analyze(expressions.get(1), context)
55+
);
56+
} else {
57+
args = Map.of(
58+
"field", (ReferenceExpression)expressionAnalyzer.analyze(expressions.get(0), context),
59+
"path", generatePath(nestedField.toString())
60+
);
61+
}
62+
return new LogicalNested(child, List.of(args), namedExpressions);
63+
}
64+
return null;
65+
}
66+
67+
private ReferenceExpression generatePath(String field) {
68+
return new ReferenceExpression(field.substring(0, field.lastIndexOf(".")), STRING);
69+
}
70+
}

‎core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,6 @@
6060
import org.opensearch.sql.ast.tree.TableFunction;
6161
import org.opensearch.sql.ast.tree.UnresolvedPlan;
6262
import org.opensearch.sql.ast.tree.Values;
63-
import org.opensearch.sql.expression.function.BuiltinFunctionName;
6463

6564
/**
6665
* Class of static methods to create specific node instances.

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -624,6 +624,10 @@ public static FunctionExpression xor(Expression... expressions) {
624624
return compile(FunctionProperties.None, BuiltinFunctionName.XOR, expressions);
625625
}
626626

627+
public static FunctionExpression nested(Expression... expressions) {
628+
return compile(FunctionProperties.None, BuiltinFunctionName.NESTED, expressions);
629+
}
630+
627631
public static FunctionExpression not(Expression... expressions) {
628632
return compile(FunctionProperties.None, BuiltinFunctionName.NOT, expressions);
629633
}

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import lombok.RequiredArgsConstructor;
1616
import org.opensearch.sql.data.model.ExprTupleValue;
1717
import org.opensearch.sql.data.model.ExprValue;
18+
import org.opensearch.sql.data.type.ExprCoreType;
1819
import org.opensearch.sql.data.type.ExprType;
1920
import org.opensearch.sql.expression.env.Environment;
2021

@@ -100,7 +101,12 @@ public ExprValue resolve(ExprTupleValue value) {
100101
}
101102

102103
private ExprValue resolve(ExprValue value, List<String> paths) {
103-
final ExprValue wholePathValue = value.keyValue(String.join(PATH_SEP, paths));
104+
ExprValue wholePathValue = value.keyValue(String.join(PATH_SEP, paths));
105+
// For array types only first index currently supported.
106+
if (value.type().equals(ExprCoreType.ARRAY)) {
107+
wholePathValue = value.collectionValue().get(0).keyValue(paths.get(0));
108+
}
109+
104110
if (!wholePathValue.isMissing() || paths.size() == 1) {
105111
return wholePathValue;
106112
} else {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,8 @@ public enum BuiltinFunctionName {
187187
STDDEV_POP(FunctionName.of("stddev_pop")),
188188
// take top documents from aggregation bucket.
189189
TAKE(FunctionName.of("take")),
190+
// Not always an aggregation query
191+
NESTED(FunctionName.of("nested")),
190192

191193
/**
192194
* Text Functions.

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.expression.function;
7+
8+
import lombok.Getter;
9+
import lombok.RequiredArgsConstructor;
10+
import org.apache.commons.lang3.tuple.Pair;
11+
import org.opensearch.sql.exception.SemanticCheckException;
12+
import org.opensearch.sql.expression.nested.NestedFunction;
13+
14+
@RequiredArgsConstructor
15+
public class NestedFunctionResolver
16+
implements FunctionResolver {
17+
18+
@Getter
19+
private final FunctionName functionName;
20+
21+
@Override
22+
public Pair<FunctionSignature, FunctionBuilder> resolve(FunctionSignature unresolvedSignature) {
23+
if (!unresolvedSignature.getFunctionName().equals(functionName)) {
24+
throw new SemanticCheckException(String.format("Expected '%s' but got '%s'",
25+
functionName.getFunctionName(), unresolvedSignature.getFunctionName().getFunctionName()));
26+
}
27+
28+
FunctionBuilder buildFunction = (functionProperties, args)
29+
-> new NestedFunction(functionName, args);
30+
return Pair.of(unresolvedSignature, buildFunction);
31+
}
32+
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ public void register(BuiltinFunctionRepository repository) {
4040
repository.register(match_phrase_prefix());
4141
repository.register(wildcard_query(BuiltinFunctionName.WILDCARD_QUERY));
4242
repository.register(wildcard_query(BuiltinFunctionName.WILDCARDQUERY));
43+
// Functions supported in SELECT clause
44+
repository.register(nested(BuiltinFunctionName.NESTED));
4345
}
4446

4547
private static FunctionResolver match_bool_prefix() {
@@ -86,6 +88,11 @@ private static FunctionResolver wildcard_query(BuiltinFunctionName wildcardQuery
8688
return new RelevanceFunctionResolver(funcName);
8789
}
8890

91+
private static FunctionResolver nested(BuiltinFunctionName nested) {
92+
FunctionName funcName = nested.getName();
93+
return new NestedFunctionResolver(funcName);
94+
}
95+
8996
public static class OpenSearchFunction extends FunctionExpression {
9097
private final FunctionName functionName;
9198
private final List<Expression> arguments;

‎core/src/main/java/org/opensearch/sql/expression/nested/NestedFunction.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.expression.nested;
7+
8+
import java.util.List;
9+
import org.opensearch.sql.data.model.ExprValue;
10+
import org.opensearch.sql.data.type.ExprType;
11+
import org.opensearch.sql.expression.Expression;
12+
import org.opensearch.sql.expression.env.Environment;
13+
import org.opensearch.sql.expression.function.FunctionName;
14+
import org.opensearch.sql.expression.function.OpenSearchFunctions;
15+
16+
public class NestedFunction extends OpenSearchFunctions.OpenSearchFunction {
17+
private final List<Expression> arguments;
18+
19+
/**
20+
* Required argument constructor.
21+
* @param functionName name of the function
22+
* @param arguments a list of expressions
23+
*/
24+
public NestedFunction(FunctionName functionName, List<Expression> arguments) {
25+
super(functionName, arguments);
26+
this.arguments = arguments;
27+
}
28+
29+
@Override
30+
public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
31+
// Only need to resolve field argument which is always first index in member variable.
32+
return valueEnv.resolve(this.arguments.get(0));
33+
}
34+
35+
@Override
36+
public ExprType type() {
37+
return this.arguments.get(0).type();
38+
}
39+
}

0 commit comments

Comments
 (0)