Skip to content

Commit a05c5af

Browse files
LantaoJinLantaoJin
committed
address comments
Signed-off-by: Lantao Jin <[email protected]>
1 parent 446c8c0 commit a05c5af

File tree

3 files changed

+68
-83
lines changed

3 files changed

+68
-83
lines changed

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

Lines changed: 50 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -234,10 +234,10 @@ public RelNode visitParse(Parse node, CalcitePlanContext context) {
234234
RexNode sourceField = rexVisitor.analyze(node.getSourceField(), context);
235235
ParseMethod parseMethod = node.getParseMethod();
236236
java.util.Map<String, Literal> arguments = node.getArguments();
237-
assert arguments.isEmpty();
238237
String pattern = (String) node.getPattern().getValue();
239238
List<String> groupCandidates =
240239
ParseUtils.getNamedGroupCandidates(parseMethod, pattern, arguments);
240+
List<RexNode> overrideFields = new ArrayList<>();
241241
List<RexNode> newFields =
242242
groupCandidates.stream()
243243
.map(
@@ -247,84 +247,69 @@ public RelNode visitParse(Parse node, CalcitePlanContext context) {
247247
SqlLibraryOperators.REGEXP_EXTRACT,
248248
sourceField,
249249
context.rexBuilder.makeLiteral(pattern));
250+
if (originalFieldNames.contains(group)) {
251+
overrideFields.add(context.relBuilder.field(group));
252+
}
250253
return context.relBuilder.alias(regexp, group);
251254
})
252255
.toList();
253256
context.relBuilder.projectPlus(newFields);
254-
255-
List<String> overriding = new ArrayList<>(groupCandidates);
256-
overriding.retainAll(originalFieldNames);
257-
renameForOverriding(overriding, context);
258-
257+
context.relBuilder.projectExcept(overrideFields);
258+
renameForOverriding(groupCandidates, context);
259259
return context.relBuilder.peek();
260260
}
261261

262-
private static void renameForOverriding(List<String> overriding, CalcitePlanContext context) {
263-
if (!overriding.isEmpty()) {
264-
List<RexNode> toDrop = context.relBuilder.fields(overriding);
265-
context.relBuilder.projectExcept(toDrop);
266-
// the overriding field in Calcite will add a numeric suffix, for example:
267-
// `| eval SAL = SAL + 1` creates a field SAL0 to replace SAL, so we rename it back to SAL,
268-
// or query `| eval SAL=SAL + 1 | where exists [ source=DEPT | where emp.SAL=HISAL ]` fails.
269-
List<String> newNames =
270-
context.relBuilder.peek().getRowType().getFieldNames().stream()
271-
.map(
272-
cur -> {
273-
String noNumericSuffix = cur.replaceAll("\\d", "");
274-
if (overriding.contains(noNumericSuffix)) {
275-
return noNumericSuffix;
276-
} else {
277-
return cur;
278-
}
279-
})
280-
.toList();
281-
context.relBuilder.rename(newNames);
282-
}
283-
}
284-
285262
@Override
286263
public RelNode visitEval(Eval node, CalcitePlanContext context) {
287264
visitChildren(node, context);
288265
List<String> originalFieldNames = context.relBuilder.peek().getRowType().getFieldNames();
289-
List<RexNode> evalList =
290-
node.getExpressionList().stream()
291-
.map(
292-
expr -> {
293-
boolean containsSubqueryExpression = containsSubqueryExpression(expr);
294-
final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
295-
if (containsSubqueryExpression) {
296-
context.relBuilder.variable(v::set);
297-
context.pushCorrelVar(v.get());
298-
}
299-
RexNode eval = rexVisitor.analyze(expr, context);
300-
if (containsSubqueryExpression) {
301-
// RelBuilder.projectPlus doesn't have a parameter with variablesSet:
302-
// projectPlus(Iterable<CorrelationId> variablesSet, RexNode... nodes)
303-
context.relBuilder.project(
304-
Iterables.concat(context.relBuilder.fields(), ImmutableList.of(eval)),
305-
ImmutableList.of(),
306-
false,
307-
ImmutableList.of(v.get().id));
308-
context.popCorrelVar();
309-
} else {
310-
context.relBuilder.projectPlus(eval);
311-
}
312-
return eval;
313-
})
314-
.toList();
315-
// Overriding the existing field if the alias has the same name with original field name.
316-
List<String> overriding =
317-
evalList.stream()
318-
.filter(expr -> expr.getKind() == AS)
319-
.map(
320-
expr ->
321-
((RexLiteral) ((RexCall) expr).getOperands().get(1)).getValueAs(String.class))
322-
.collect(Collectors.toList());
323-
overriding.retainAll(originalFieldNames);
324-
renameForOverriding(overriding, context);
266+
node.getExpressionList()
267+
.forEach(
268+
expr -> {
269+
boolean containsSubqueryExpression = containsSubqueryExpression(expr);
270+
final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
271+
if (containsSubqueryExpression) {
272+
context.relBuilder.variable(v::set);
273+
context.pushCorrelVar(v.get());
274+
}
275+
RexNode eval = rexVisitor.analyze(expr, context);
276+
if (containsSubqueryExpression) {
277+
// RelBuilder.projectPlus doesn't have a parameter with variablesSet:
278+
// projectPlus(Iterable<CorrelationId> variablesSet, RexNode... nodes)
279+
context.relBuilder.project(
280+
Iterables.concat(context.relBuilder.fields(), ImmutableList.of(eval)),
281+
ImmutableList.of(),
282+
false,
283+
ImmutableList.of(v.get().id));
284+
context.popCorrelVar();
285+
} else {
286+
// Overriding the existing field if the alias has the same name with original field
287+
// name.
288+
RexNode overrideField = null;
289+
String alias =
290+
((RexLiteral) ((RexCall) eval).getOperands().get(1)).getValueAs(String.class);
291+
if (originalFieldNames.contains(alias)) {
292+
overrideField = context.relBuilder.field(alias);
293+
context.relBuilder.projectPlus(eval);
294+
context.relBuilder.projectExcept(overrideField);
295+
renameForOverriding(List.of(alias), context);
296+
} else {
297+
context.relBuilder.projectPlus(eval);
298+
}
299+
}
300+
});
325301
return context.relBuilder.peek();
326302
}
327303

304+
private static void renameForOverriding(List<String> newNames, CalcitePlanContext context) {
305+
List<String> originalFieldNames = context.relBuilder.peek().getRowType().getFieldNames();
306+
int length = originalFieldNames.size();
307+
List<String> expectedRenameFields =
308+
new ArrayList<>(originalFieldNames.subList(0, length - newNames.size()));
309+
expectedRenameFields.addAll(newNames);
310+
context.relBuilder.rename(expectedRenameFields);
311+
}
312+
328313
@Override
329314
public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
330315
visitChildren(node, context);

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

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111

1212
import java.math.BigDecimal;
1313
import java.util.List;
14-
import java.util.Map;
1514
import java.util.stream.Collectors;
1615
import lombok.RequiredArgsConstructor;
1716
import org.apache.calcite.rel.RelNode;
@@ -248,16 +247,6 @@ public RexNode visitQualifiedName(QualifiedName node, CalcitePlanContext context
248247
.peekCorrelVar()
249248
.map(correlVar -> context.relBuilder.field(correlVar, qualifiedName))
250249
.orElseGet(() -> context.relBuilder.field(qualifiedName));
251-
}
252-
// 3. resolve overriding fields, for example, `eval SAL = SAL + 1` will delete the original SAL
253-
// and add a SAL0. SAL0 in currentFields, but qualifiedName is SAL.
254-
// TODO now we cannot handle the case using a overriding fields in subquery, for example
255-
// source = EMP | eval DEPTNO = DEPTNO + 1 | where exists [ source = DEPT | where emp.DEPTNO =
256-
// DEPTNO ]
257-
Map<String, String> fieldMap =
258-
currentFields.stream().collect(Collectors.toMap(s -> s.replaceAll("\\d", ""), s -> s));
259-
if (fieldMap.containsKey(qualifiedName)) {
260-
return context.relBuilder.field(fieldMap.get(qualifiedName));
261250
} else {
262251
throw new IllegalArgumentException(
263252
String.format(
@@ -313,13 +302,6 @@ private boolean isTimeBased(SpanUnit unit) {
313302
return !(unit == NONE || unit == UNKNOWN);
314303
}
315304

316-
// @Override
317-
// public RexNode visitAggregateFunction(AggregateFunction node, Context context) {
318-
// RexNode field = analyze(node.getField(), context);
319-
// AggregateCall aggregateCall = translateAggregateCall(node, field, relBuilder);
320-
// return new MyAggregateCall(aggregateCall);
321-
// }
322-
323305
@Override
324306
public RexNode visitLet(Let node, CalcitePlanContext context) {
325307
RexNode expr = analyze(node.getExpression(), context);

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.io.IOException;
1616
import org.json.JSONObject;
1717
import org.junit.Test;
18+
import org.opensearch.client.Request;
1819

1920
public class CalcitePPLParseIT extends CalcitePPLIntegTestCase {
2021
@Override
@@ -128,4 +129,21 @@ public void testParseMultipleGroups() {
128129
verifyErrorMessageContains(
129130
e, "Multiple capturing groups (count=2) not allowed in regex input for REGEXP_EXTRACT");
130131
}
132+
133+
@Test
134+
public void testParseOverriding2() throws IOException {
135+
Request request1 = new Request("PUT", "/test/_doc/1?refresh=true");
136+
request1.setJsonEntity(
137+
"{\"email\": \"[email protected]\", \"email0\": \"[email protected]\", \"email1\": \"[email protected]\"}");
138+
client().performRequest(request1);
139+
JSONObject result;
140+
result =
141+
executeQuery(
142+
"source = test | parse email '.+@(?<email0>.+)' | fields email, email0, email1");
143+
verifyDataRows(result, rows("[email protected]", "a.com", "[email protected]"));
144+
result =
145+
executeQuery(
146+
"source = test | parse email '.+@(?<email>.+)' | fields email, email0, email1");
147+
verifyDataRows(result, rows("a.com", "[email protected]", "[email protected]"));
148+
}
131149
}

0 commit comments

Comments
 (0)