From 74f35f5892e36c9104b60bd84178de35fa54253e Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Mon, 29 Jun 2026 11:29:23 +0800 Subject: [PATCH 01/12] Implement PPL foreach command Signed-off-by: Songkan Tang --- .../sql/ast/AbstractNodeVisitor.java | 10 +++ .../ast/expression/ForeachPlaceholder.java | 33 ++++++++ .../org/opensearch/sql/ast/tree/Foreach.java | 52 ++++++++++++ .../sql/calcite/CalcitePlanContext.java | 9 +++ .../sql/calcite/CalciteRelNodeVisitor.java | 79 +++++++++++++++++++ .../sql/calcite/CalciteRexNodeVisitor.java | 14 ++++ .../sql/calcite/CalciteNoPushdownIT.java | 1 + .../sql/calcite/remote/CalciteExplainIT.java | 16 ++++ .../remote/CalciteForeachCommandIT.java | 65 +++++++++++++++ ppl/src/main/antlr/OpenSearchPPLLexer.g4 | 1 + ppl/src/main/antlr/OpenSearchPPLParser.g4 | 38 +++++++++ .../opensearch/sql/ppl/parser/AstBuilder.java | 23 ++++++ .../sql/ppl/parser/AstExpressionBuilder.java | 12 +++ .../sql/ppl/antlr/PPLSyntaxParserTest.java | 18 +++++ .../ppl/calcite/CalcitePPLForeachTest.java | 47 +++++++++++ 15 files changed, 418 insertions(+) create mode 100644 core/src/main/java/org/opensearch/sql/ast/expression/ForeachPlaceholder.java create mode 100644 core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java create mode 100644 ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java index be02547a2da..89c6c105bbe 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -18,6 +18,7 @@ import org.opensearch.sql.ast.expression.Compare; import org.opensearch.sql.ast.expression.EqualTo; import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.ForeachPlaceholder; import org.opensearch.sql.ast.expression.Function; import org.opensearch.sql.ast.expression.HighlightFunction; import org.opensearch.sql.ast.expression.In; @@ -62,6 +63,7 @@ import org.opensearch.sql.ast.tree.FillNull; import org.opensearch.sql.ast.tree.Filter; import org.opensearch.sql.ast.tree.Flatten; +import org.opensearch.sql.ast.tree.Foreach; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.Head; import org.opensearch.sql.ast.tree.Join; @@ -173,6 +175,10 @@ public T visitProject(Project node, C context) { return visitChildren(node, context); } + public T visitForeach(Foreach node, C context) { + return visitChildren(node, context); + } + public T visitAggregation(Aggregation node, C context) { return visitChildren(node, context); } @@ -253,6 +259,10 @@ public T visitField(Field node, C context) { return visitChildren(node, context); } + public T visitForeachPlaceholder(ForeachPlaceholder node, C context) { + return visitChildren(node, context); + } + public T visitQualifiedName(QualifiedName node, C context) { return visitChildren(node, context); } diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/ForeachPlaceholder.java b/core/src/main/java/org/opensearch/sql/ast/expression/ForeachPlaceholder.java new file mode 100644 index 00000000000..9362a70d22f --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/expression/ForeachPlaceholder.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ast.expression; + +import com.google.common.collect.ImmutableList; +import java.util.List; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import lombok.ToString; +import org.opensearch.sql.ast.AbstractNodeVisitor; + +/** Placeholder used by the PPL foreach command, such as {@code <>}. */ +@Getter +@ToString +@EqualsAndHashCode(callSuper = false) +@RequiredArgsConstructor +public class ForeachPlaceholder extends UnresolvedExpression { + private final String name; + + @Override + public List getChild() { + return ImmutableList.of(); + } + + @Override + public R accept(AbstractNodeVisitor nodeVisitor, C context) { + return nodeVisitor.visitForeachPlaceholder(this, context); + } +} diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java b/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java new file mode 100644 index 00000000000..dadfac87196 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java @@ -0,0 +1,52 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ast.tree; + +import com.google.common.collect.ImmutableList; +import java.util.List; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import lombok.ToString; +import org.opensearch.sql.ast.AbstractNodeVisitor; +import org.opensearch.sql.ast.expression.UnresolvedExpression; + +/** AST node representing the PPL foreach command. */ +@Getter +@ToString +@EqualsAndHashCode(callSuper = false) +@RequiredArgsConstructor +public class Foreach extends UnresolvedPlan { + private final String mode; + private final List fieldPatterns; + private final List evalClauses; + private UnresolvedPlan child; + + @Override + public Foreach attach(UnresolvedPlan child) { + this.child = child; + return this; + } + + @Override + public List getChild() { + return child == null ? ImmutableList.of() : ImmutableList.of(child); + } + + @Override + public T accept(AbstractNodeVisitor nodeVisitor, C context) { + return nodeVisitor.visitForeach(this, context); + } + + @Getter + @ToString + @EqualsAndHashCode + @RequiredArgsConstructor + public static class ForeachEvalClause { + private final String targetTemplate; + private final UnresolvedExpression expression; + } +} diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index 7ca7ab09304..6dedee9771f 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -65,6 +65,7 @@ public class CalcitePlanContext { private final Stack> windowPartitions = new Stack<>(); @Getter public Map rexLambdaRefMap; + @Getter private Map foreachBindings = new HashMap<>(); /** * Maps AggregateFunction AST nodes to their output field index for HAVING/post-aggregate @@ -179,6 +180,14 @@ public static boolean isLegacyPreferred() { return legacyPreferredFlag.get(); } + public void pushForeachBindings(Map bindings) { + foreachBindings = new HashMap<>(bindings); + } + + public void clearForeachBindings() { + foreachBindings.clear(); + } + public void putRexLambdaRefMap(Map candidateMap) { this.rexLambdaRefMap.putAll(candidateMap); } diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index c4bb8bcd910..a40fa6007be 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -129,6 +129,7 @@ import org.opensearch.sql.ast.tree.FillNull; import org.opensearch.sql.ast.tree.Filter; import org.opensearch.sql.ast.tree.Flatten; +import org.opensearch.sql.ast.tree.Foreach; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.GraphLookup.Direction; import org.opensearch.sql.ast.tree.Head; @@ -1239,6 +1240,84 @@ public RelNode visitEval(Eval node, CalcitePlanContext context) { return context.relBuilder.peek(); } + @Override + public RelNode visitForeach(Foreach node, CalcitePlanContext context) { + visitChildren(node, context); + if (!"multifield".equalsIgnoreCase(node.getMode())) { + throw new CalciteUnsupportedException( + "foreach mode [" + node.getMode() + "] is not supported"); + } + + List currentFields = context.relBuilder.peek().getRowType().getFieldNames(); + LinkedHashSet matchingFields = new LinkedHashSet<>(); + for (String pattern : node.getFieldPatterns()) { + matchingFields.addAll(WildcardUtils.expandWildcardPattern(pattern, currentFields)); + } + + for (String fieldName : matchingFields) { + Map bindings = buildForeachBindings(node.getFieldPatterns(), fieldName); + context.pushForeachBindings(bindings); + try { + for (Foreach.ForeachEvalClause clause : node.getEvalClauses()) { + RexNode expr = rexVisitor.analyze(clause.getExpression(), context); + String alias = substituteForeachTemplate(clause.getTargetTemplate(), bindings); + projectPlusOverriding( + List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); + } + } finally { + context.clearForeachBindings(); + } + } + return context.relBuilder.peek(); + } + + private Map buildForeachBindings(List patterns, String fieldName) { + Map bindings = new HashMap<>(); + bindings.put("FIELD", fieldName); + for (String pattern : patterns) { + List segments = wildcardSegments(pattern, fieldName); + if (segments != null) { + bindings.put("MATCHSTR", String.join("", segments)); + for (int i = 0; i < segments.size(); i++) { + bindings.put("MATCHSEG" + (i + 1), segments.get(i)); + } + break; + } + } + return bindings; + } + + private List wildcardSegments(String pattern, String fieldName) { + if (!WildcardUtils.matchesWildcardPattern(pattern, fieldName)) { + return null; + } + if (!WildcardUtils.containsWildcard(pattern)) { + return List.of(); + } + java.util.regex.Matcher matcher = + java.util.regex.Pattern.compile(WildcardUtils.convertWildcardPatternToRegex(pattern)) + .matcher(fieldName); + if (!matcher.matches()) { + return null; + } + List segments = new ArrayList<>(); + for (int i = 1; i <= matcher.groupCount(); i++) { + segments.add(matcher.group(i)); + } + return segments; + } + + private String substituteForeachTemplate(String template, Map bindings) { + String result = template; + for (Map.Entry entry : bindings.entrySet()) { + result = + result.replaceAll( + "(?i)<<" + java.util.regex.Pattern.quote(entry.getKey()) + ">>", + java.util.regex.Matcher.quoteReplacement(entry.getValue())); + } + return result; + } + @Override public RelNode visitConvert(Convert node, CalcitePlanContext context) { visitChildren(node, context); diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java index 849b615f970..b84afb2d1a5 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -53,6 +53,7 @@ import org.opensearch.sql.ast.expression.Cast; import org.opensearch.sql.ast.expression.Compare; import org.opensearch.sql.ast.expression.EqualTo; +import org.opensearch.sql.ast.expression.ForeachPlaceholder; import org.opensearch.sql.ast.expression.Function; import org.opensearch.sql.ast.expression.HighlightFunction; import org.opensearch.sql.ast.expression.In; @@ -402,6 +403,19 @@ public RexNode visitQualifiedName(QualifiedName node, CalcitePlanContext context return QualifiedNameResolver.resolve(node, context); } + @Override + public RexNode visitForeachPlaceholder(ForeachPlaceholder node, CalcitePlanContext context) { + String name = node.getName().toUpperCase(Locale.ROOT); + String value = context.getForeachBindings().get(name); + if (value == null) { + throw new SemanticCheckException("Unresolved foreach placeholder <<" + node.getName() + ">>"); + } + if ("FIELD".equals(name)) { + return context.relBuilder.field(value); + } + return context.rexBuilder.makeLiteral(value); + } + @Override public RexNode visitAlias(Alias node, CalcitePlanContext context) { RexNode expr = analyze(node.getDelegated(), context); diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/CalciteNoPushdownIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/CalciteNoPushdownIT.java index d46649d56d9..907f91fa98e 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/CalciteNoPushdownIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/CalciteNoPushdownIT.java @@ -37,6 +37,7 @@ CalciteDescribeCommandIT.class, CalciteExpandCommandIT.class, CalciteFieldFormatCommandIT.class, + CalciteForeachCommandIT.class, CalciteFieldsCommandIT.class, CalciteFillNullCommandIT.class, CalciteFlattenCommandIT.class, diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java index 95c47b9b0b7..b8c10799612 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java @@ -2910,6 +2910,22 @@ public void testNoMvWithEval() throws IOException { logical.contains("concat") && logical.contains("array_join")); } + @Test + public void testForeachExplain() throws IOException { + String query = + StringUtils.format( + "source=%s | foreach age balance [ eval <>_double = <> * 2 ] | fields" + + " age_double, balance_double", + TEST_INDEX_BANK); + String logical = logicalPlan(explainQueryYaml(query)); + Assert.assertTrue( + "Expected logical plan to contain foreach-expanded eval fields", + logical.contains("age_double") + && logical.contains("balance_double") + && logical.contains("*($") + && logical.contains(", 2)")); + } + /** * Return just the {@code logical:} section of a YAML explain result (everything before the {@code * physical:} key). The logical plan is deterministic across pushdown on/off, whereas the physical diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java new file mode 100644 index 00000000000..c2fe110b4b1 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java @@ -0,0 +1,65 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.remote; + +import static org.opensearch.sql.util.MatcherUtils.rows; +import static org.opensearch.sql.util.MatcherUtils.schema; +import static org.opensearch.sql.util.MatcherUtils.verifyDataRows; +import static org.opensearch.sql.util.MatcherUtils.verifySchema; + +import java.io.IOException; +import org.json.JSONObject; +import org.junit.jupiter.api.Test; +import org.opensearch.client.Request; +import org.opensearch.sql.legacy.TestUtils; +import org.opensearch.sql.ppl.PPLIntegTestCase; + +public class CalciteForeachCommandIT extends PPLIntegTestCase { + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); + + if (!TestUtils.isIndexExist(client(), "test_foreach")) { + String mapping = + "{\"mappings\":{\"properties\":{" + + "\"a\":{\"type\":\"long\"}," + + "\"b\":{\"type\":\"long\"}," + + "\"value_cpu\":{\"type\":\"long\"}," + + "\"value_mem\":{\"type\":\"long\"}}}}"; + TestUtils.createIndexByRestClient(client(), "test_foreach", mapping); + + Request request1 = new Request("PUT", "/test_foreach/_doc/1?refresh=true"); + request1.setJsonEntity("{\"a\": 1, \"b\": 2, \"value_cpu\": 10, \"value_mem\": 20}"); + client().performRequest(request1); + + Request request2 = new Request("PUT", "/test_foreach/_doc/2?refresh=true"); + request2.setJsonEntity("{\"a\": 3, \"b\": 4, \"value_cpu\": 30, \"value_mem\": 40}"); + client().performRequest(request2); + } + } + + @Test + public void testForeachExplicitFields() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | foreach a b [ eval <>_double = <> * 2 ] |" + + " fields a_double, b_double"); + verifySchema(result, schema("a_double", "bigint"), schema("b_double", "bigint")); + verifyDataRows(result, rows(2, 4), rows(6, 8)); + } + + @Test + public void testForeachWildcardMatchstr() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | foreach value_* [ eval copy_<> = <> ] |" + + " fields copy_cpu, copy_mem"); + verifySchema(result, schema("copy_cpu", "bigint"), schema("copy_mem", "bigint")); + verifyDataRows(result, rows(10, 20), rows(30, 40)); + } +} diff --git a/ppl/src/main/antlr/OpenSearchPPLLexer.g4 b/ppl/src/main/antlr/OpenSearchPPLLexer.g4 index 4bc69a8f295..2e3b918cb6a 100644 --- a/ppl/src/main/antlr/OpenSearchPPLLexer.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLLexer.g4 @@ -26,6 +26,7 @@ STREAMSTATS: 'STREAMSTATS'; DEDUP: 'DEDUP'; SORT: 'SORT'; EVAL: 'EVAL'; +FOREACH: 'FOREACH'; FIELDFORMAT: 'FIELDFORMAT'; HEAD: 'HEAD'; BIN: 'BIN'; diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4 index 98f85b08282..6635e52a7a0 100644 --- a/ppl/src/main/antlr/OpenSearchPPLParser.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLParser.g4 @@ -64,6 +64,7 @@ commands | dedupCommand | sortCommand | evalCommand + | foreachCommand | headCommand | binCommand | rareTopCommand @@ -115,6 +116,7 @@ commandName | DEDUP | SORT | EVAL + | FOREACH | FIELDFORMAT | HEAD | BIN @@ -375,6 +377,31 @@ evalCommand : EVAL evalClause (COMMA evalClause)* ; +foreachCommand + : FOREACH foreachOption* foreachFieldPattern (COMMA? foreachFieldPattern)* LT_SQR_PRTHS foreachEvalCommand RT_SQR_PRTHS + ; + +foreachFieldPattern + : wcFieldExpression + | STAR + ; + +foreachOption + : ident EQUAL ident + ; + +foreachEvalCommand + : EVAL foreachEvalClause (COMMA foreachEvalClause)* + ; + +foreachEvalClause + : target = foreachEvalTarget EQUAL logicalExpression + ; + +foreachEvalTarget + : (foreachPlaceholder | ident | DOT | MODULE)+ + ; + fieldformatCommand : FIELDFORMAT fieldFormatEvalClause (COMMA fieldFormatEvalClause)* ; @@ -942,6 +969,7 @@ valueExpression | literalValue # literalValueExpr | functionCall # functionCallExpr | lambda # lambdaExpr + | foreachPlaceholder # foreachPlaceholderExpr | LT_SQR_PRTHS subSearch RT_SQR_PRTHS # scalarSubqueryExpr | valueExpression NOT? IN LT_SQR_PRTHS subSearch RT_SQR_PRTHS # inSubqueryExpr | LT_PRTHS valueExpression (COMMA valueExpression)* RT_PRTHS NOT? IN LT_SQR_PRTHS subSearch RT_SQR_PRTHS # inSubqueryExpr @@ -1044,6 +1072,16 @@ wcFieldExpression : wcQualifiedName ; +foreachPlaceholder + : LESS LESS foreachPlaceholderName GREATER GREATER + ; + +foreachPlaceholderName + : FIELD + | ID + | NUMERIC_ID + ; + selectFieldExpression : wcQualifiedName | STAR diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java index 3741137f5a9..94482b21bb2 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java @@ -86,6 +86,8 @@ import org.opensearch.sql.ast.tree.FillNull; import org.opensearch.sql.ast.tree.Filter; import org.opensearch.sql.ast.tree.Flatten; +import org.opensearch.sql.ast.tree.Foreach; +import org.opensearch.sql.ast.tree.Foreach.ForeachEvalClause; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.GraphLookup.Direction; import org.opensearch.sql.ast.tree.Head; @@ -831,6 +833,27 @@ public UnresolvedPlan visitEvalCommand(EvalCommandContext ctx) { .collect(Collectors.toList())); } + @Override + public UnresolvedPlan visitForeachCommand(OpenSearchPPLParser.ForeachCommandContext ctx) { + String mode = + ctx.foreachOption().stream() + .filter(option -> "mode".equalsIgnoreCase(option.ident(0).getText())) + .map(option -> option.ident(1).getText()) + .findFirst() + .orElse("multifield"); + List patterns = + ctx.foreachFieldPattern().stream().map(this::getTextInQuery).collect(Collectors.toList()); + List evalClauses = + ctx.foreachEvalCommand().foreachEvalClause().stream() + .map( + clause -> + new ForeachEvalClause( + getTextInQuery(clause.target), + expressionBuilder.visit(clause.logicalExpression()))) + .collect(Collectors.toList()); + return new Foreach(mode, patterns, evalClauses); + } + @Override public UnresolvedPlan visitFieldformatCommand(FieldformatCommandContext ctx) { // Use the new fieldFormatEvalClause instead of evalClause diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java index 77d5c77a635..d4bfc83ede0 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java @@ -105,6 +105,18 @@ public UnresolvedExpression visitEvalClause(EvalClauseContext ctx) { return new Let((Field) visit(ctx.fieldExpression()), visit(ctx.logicalExpression())); } + @Override + public UnresolvedExpression visitForeachPlaceholderExpr( + OpenSearchPPLParser.ForeachPlaceholderExprContext ctx) { + return visit(ctx.foreachPlaceholder()); + } + + @Override + public UnresolvedExpression visitForeachPlaceholder( + OpenSearchPPLParser.ForeachPlaceholderContext ctx) { + return new ForeachPlaceholder(ctx.foreachPlaceholderName().getText()); + } + /** Field format eval clause - similar to evalClause but for fieldformat command. */ @Override public UnresolvedExpression visitFieldFormatEvalClause( diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java index 1238f1e7e6c..e249bad1b90 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java @@ -89,6 +89,24 @@ public void testSearchFieldsCommandShouldPass() { assertNotEquals(null, tree); } + @Test + public void testForeachCommandShouldPass() { + ParseTree tree = + new PPLSyntaxParser() + .parse("source=t | foreach a* b [ eval <>_double = <> * 2 ]"); + assertNotEquals(null, tree); + } + + @Test + public void testForeachCommandWithModeShouldPass() { + ParseTree tree = + new PPLSyntaxParser() + .parse( + "source=t | foreach mode=multifield value_* [ eval" + + " total_<> = <> + 1 ]"); + assertNotEquals(null, tree); + } + @Test public void testSearchFieldsCommandCrossClusterShouldPass() { ParseTree tree = new PPLSyntaxParser().parse("search source=c:t a=1 b=2 | fields a,b"); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java new file mode 100644 index 00000000000..7379d094e73 --- /dev/null +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java @@ -0,0 +1,47 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.calcite; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.test.CalciteAssert; +import org.junit.Test; + +public class CalcitePPLForeachTest extends CalcitePPLAbstractTest { + + public CalcitePPLForeachTest() { + super(CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL); + } + + @Test + public void testForeachExpandsExplicitFields() { + String ppl = + "source=EMP | foreach SAL COMM [ eval <>_double = <> * 2 ] | fields" + + " EMPNO, SAL_double, COMM_double"; + RelNode root = getRelNode(ppl); + + String expectedLogical = + "LogicalProject(EMPNO=[$0], SAL_double=[*($5, 2)], COMM_double=[*($6, 2)])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + + String expectedSparkSql = + "SELECT `EMPNO`, `SAL` * 2 `SAL_double`, `COMM` * 2 `COMM_double`\n" + "FROM `scott`.`EMP`"; + verifyPPLToSparkSQL(root, expectedSparkSql); + } + + @Test + public void testForeachWildcardAndMatchstr() { + String ppl = + "source=EMP | foreach *NO [ eval copy_<> = <> ] | fields copy_EMP," + + " copy_DEPT"; + RelNode root = getRelNode(ppl); + + String expectedLogical = + "LogicalProject(copy_EMP=[$0], copy_DEPT=[$7])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + } +} From 287febb5706db821f3800cab715b3eb407b689f2 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Mon, 6 Jul 2026 14:00:11 +0800 Subject: [PATCH 02/12] Complete PPL foreach collection modes Signed-off-by: Songkan Tang --- .../org/opensearch/sql/ast/tree/Foreach.java | 3 + .../sql/calcite/CalcitePlanContext.java | 15 +- .../sql/calcite/CalciteRelNodeVisitor.java | 136 ++++++++++++++++-- .../sql/calcite/CalciteRexNodeVisitor.java | 51 ++++++- .../function/BuiltinFunctionName.java | 1 + .../function/PPLBuiltinOperators.java | 3 + .../expression/function/PPLFuncImpTable.java | 2 + .../jsonUDF/ForeachJsonArrayFunctionImpl.java | 76 ++++++++++ .../remote/CalciteForeachCommandIT.java | 50 +++++++ ppl/src/main/antlr/OpenSearchPPLParser.g4 | 7 +- .../opensearch/sql/ppl/parser/AstBuilder.java | 42 ++++-- .../sql/ppl/antlr/PPLSyntaxParserTest.java | 25 ++++ .../ppl/calcite/CalcitePPLForeachTest.java | 55 +++++++ 13 files changed, 435 insertions(+), 31 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java b/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java index dadfac87196..486870e6efb 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java @@ -7,6 +7,7 @@ import com.google.common.collect.ImmutableList; import java.util.List; +import java.util.Map; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; @@ -21,7 +22,9 @@ @RequiredArgsConstructor public class Foreach extends UnresolvedPlan { private final String mode; + private final Map options; private final List fieldPatterns; + private final UnresolvedExpression collectionExpression; private final List evalClauses; private UnresolvedPlan child; diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index 6dedee9771f..6c3fac7ff0a 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -65,7 +65,7 @@ public class CalcitePlanContext { private final Stack> windowPartitions = new Stack<>(); @Getter public Map rexLambdaRefMap; - @Getter private Map foreachBindings = new HashMap<>(); + @Getter private Map foreachBindings = new HashMap<>(); /** * Maps AggregateFunction AST nodes to their output field index for HAVING/post-aggregate @@ -114,6 +114,7 @@ private CalcitePlanContext(CalcitePlanContext parent) { this.rexLambdaRefMap = new HashMap<>(); // New map for lambda variables this.capturedVariables = new ArrayList<>(); // New list for captured variables this.inLambdaContext = true; // Mark that we're inside a lambda + this.foreachBindings = new HashMap<>(parent.foreachBindings); } public RexNode resolveJoinCondition( @@ -180,7 +181,7 @@ public static boolean isLegacyPreferred() { return legacyPreferredFlag.get(); } - public void pushForeachBindings(Map bindings) { + public void pushForeachBindings(Map bindings) { foreachBindings = new HashMap<>(bindings); } @@ -188,6 +189,16 @@ public void clearForeachBindings() { foreachBindings.clear(); } + public enum ForeachBindingType { + FIELD, + LITERAL, + LAMBDA, + PAIR_ITEM, + PAIR_ITER + } + + public record ForeachBinding(String value, ForeachBindingType type) {} + public void putRexLambdaRefMap(Map candidateMap) { this.rexLambdaRefMap.putAll(candidateMap); } diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index a40fa6007be..3919ecffd49 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -98,6 +98,7 @@ import org.opensearch.sql.ast.expression.Argument.ArgumentMap; import org.opensearch.sql.ast.expression.Field; import org.opensearch.sql.ast.expression.Function; +import org.opensearch.sql.ast.expression.LambdaFunction; import org.opensearch.sql.ast.expression.Let; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.expression.ParseMethod; @@ -167,6 +168,8 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.ast.tree.Values; import org.opensearch.sql.ast.tree.Window; +import org.opensearch.sql.calcite.CalcitePlanContext.ForeachBinding; +import org.opensearch.sql.calcite.CalcitePlanContext.ForeachBindingType; import org.opensearch.sql.calcite.plan.AliasFieldsWrappable; import org.opensearch.sql.calcite.plan.HighlightPushDown; import org.opensearch.sql.calcite.plan.OpenSearchConstants; @@ -1243,11 +1246,21 @@ public RelNode visitEval(Eval node, CalcitePlanContext context) { @Override public RelNode visitForeach(Foreach node, CalcitePlanContext context) { visitChildren(node, context); - if (!"multifield".equalsIgnoreCase(node.getMode())) { - throw new CalciteUnsupportedException( - "foreach mode [" + node.getMode() + "] is not supported"); + String mode = node.getMode().toLowerCase(java.util.Locale.ROOT); + switch (mode) { + case "multifield": + return visitForeachMultifield(node, context); + case "multivalue": + case "json_array": + case "auto_collections": + return visitForeachCollection(node, context, mode); + default: + throw new CalciteUnsupportedException( + "foreach mode [" + node.getMode() + "] is not supported"); } + } + private RelNode visitForeachMultifield(Foreach node, CalcitePlanContext context) { List currentFields = context.relBuilder.peek().getRowType().getFieldNames(); LinkedHashSet matchingFields = new LinkedHashSet<>(); for (String pattern : node.getFieldPatterns()) { @@ -1255,7 +1268,8 @@ public RelNode visitForeach(Foreach node, CalcitePlanContext context) { } for (String fieldName : matchingFields) { - Map bindings = buildForeachBindings(node.getFieldPatterns(), fieldName); + Map bindings = + buildForeachMultifieldBindings(node.getFieldPatterns(), fieldName, node.getOptions()); context.pushForeachBindings(bindings); try { for (Foreach.ForeachEvalClause clause : node.getEvalClauses()) { @@ -1271,15 +1285,104 @@ public RelNode visitForeach(Foreach node, CalcitePlanContext context) { return context.relBuilder.peek(); } - private Map buildForeachBindings(List patterns, String fieldName) { - Map bindings = new HashMap<>(); - bindings.put("FIELD", fieldName); + private RelNode visitForeachCollection( + Foreach node, CalcitePlanContext context, String normalizedMode) { + if (node.getCollectionExpression() == null) { + throw new SemanticCheckException("foreach " + normalizedMode + " mode requires a field"); + } + if (node.getFieldPatterns().size() != 1) { + throw new SemanticCheckException( + "foreach " + normalizedMode + " mode accepts exactly one field"); + } + UnresolvedExpression collectionExpression = + collectionExpressionForMode(node.getCollectionExpression(), context, normalizedMode); + String itemName = option(node.getOptions(), "itemstr", "ITEM"); + String iterName = option(node.getOptions(), "iterstr", "ITER"); + String transformItemName = "__foreach_item"; + String transformIterName = "__foreach_iter"; + String pairName = "__foreach_pair"; + Function pairedCollection = + new Function( + BuiltinFunctionName.TRANSFORM.getName().getFunctionName(), + List.of( + collectionExpression, + new LambdaFunction( + new Function( + BuiltinFunctionName.ARRAY.getName().getFunctionName(), + List.of( + new QualifiedName(transformItemName), + new QualifiedName(transformIterName))), + List.of( + new QualifiedName(transformItemName), + new QualifiedName(transformIterName))))); + Map bindings = new HashMap<>(); + putBinding(bindings, itemName, pairName, ForeachBindingType.PAIR_ITEM); + putBinding(bindings, "ITEM", pairName, ForeachBindingType.PAIR_ITEM); + putBinding(bindings, iterName, pairName, ForeachBindingType.PAIR_ITER); + putBinding(bindings, "ITER", pairName, ForeachBindingType.PAIR_ITER); + + context.pushForeachBindings(bindings); + try { + for (Foreach.ForeachEvalClause clause : node.getEvalClauses()) { + String alias = substituteForeachTemplate(clause.getTargetTemplate(), bindings); + Function reduce = + new Function( + BuiltinFunctionName.REDUCE.getName().getFunctionName(), + List.of( + pairedCollection, + new Field(new QualifiedName(alias)), + new LambdaFunction( + clause.getExpression(), + List.of(new QualifiedName(alias), new QualifiedName(pairName))))); + RexNode expr = rexVisitor.analyze(reduce, context); + projectPlusOverriding( + List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); + } + } finally { + context.clearForeachBindings(); + } + return context.relBuilder.peek(); + } + + private UnresolvedExpression collectionExpressionForMode( + UnresolvedExpression collectionExpression, + CalcitePlanContext context, + String normalizedMode) { + if ("multivalue".equals(normalizedMode)) { + return collectionExpression; + } + if ("auto_collections".equals(normalizedMode)) { + RexNode rexNode = rexVisitor.analyze(collectionExpression, context); + if (rexNode.getType() instanceof ArraySqlType) { + return collectionExpression; + } + } + return new Function("foreach_json_array", List.of(collectionExpression)); + } + + private Map buildForeachMultifieldBindings( + List patterns, String fieldName, Map options) { + Map bindings = new HashMap<>(); + putBinding(bindings, "FIELD", fieldName, ForeachBindingType.FIELD); + putBinding(bindings, option(options, "fieldstr", "FIELD"), fieldName, ForeachBindingType.FIELD); for (String pattern : patterns) { List segments = wildcardSegments(pattern, fieldName); if (segments != null) { - bindings.put("MATCHSTR", String.join("", segments)); + String matchStr = String.join("", segments); + putBinding(bindings, "MATCHSTR", matchStr, ForeachBindingType.LITERAL); + putBinding( + bindings, + option(options, "matchstr", "MATCHSTR"), + matchStr, + ForeachBindingType.LITERAL); for (int i = 0; i < segments.size(); i++) { - bindings.put("MATCHSEG" + (i + 1), segments.get(i)); + String key = "MATCHSEG" + (i + 1); + putBinding(bindings, key, segments.get(i), ForeachBindingType.LITERAL); + putBinding( + bindings, + option(options, "matchseg" + (i + 1), key), + segments.get(i), + ForeachBindingType.LITERAL); } break; } @@ -1307,13 +1410,22 @@ private List wildcardSegments(String pattern, String fieldName) { return segments; } - private String substituteForeachTemplate(String template, Map bindings) { + private String option(Map options, String name, String defaultValue) { + return options.getOrDefault(name, defaultValue); + } + + private void putBinding( + Map bindings, String name, String value, ForeachBindingType type) { + bindings.put(name.toUpperCase(java.util.Locale.ROOT), new ForeachBinding(value, type)); + } + + private String substituteForeachTemplate(String template, Map bindings) { String result = template; - for (Map.Entry entry : bindings.entrySet()) { + for (Map.Entry entry : bindings.entrySet()) { result = result.replaceAll( "(?i)<<" + java.util.regex.Pattern.quote(entry.getKey()) + ">>", - java.util.regex.Matcher.quoteReplacement(entry.getValue())); + java.util.regex.Matcher.quoteReplacement(entry.getValue().value())); } return result; } diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java index b84afb2d1a5..5d228723c94 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -80,6 +80,7 @@ import org.opensearch.sql.ast.tree.Sort.SortOption; import org.opensearch.sql.ast.tree.Sort.SortOrder; import org.opensearch.sql.ast.tree.UnresolvedPlan; +import org.opensearch.sql.calcite.CalcitePlanContext.ForeachBinding; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType; import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; @@ -400,20 +401,58 @@ private boolean isBooleanLiteral(RexNode node) { /** Resolve qualified name. Note, the name should be case-sensitive. */ @Override public RexNode visitQualifiedName(QualifiedName node, CalcitePlanContext context) { + ForeachBinding binding = foreachBinding(node.toString(), context); + if (binding != null) { + return foreachBindingToRexNode(node.toString(), binding, context); + } return QualifiedNameResolver.resolve(node, context); } @Override public RexNode visitForeachPlaceholder(ForeachPlaceholder node, CalcitePlanContext context) { - String name = node.getName().toUpperCase(Locale.ROOT); - String value = context.getForeachBindings().get(name); - if (value == null) { + ForeachBinding binding = foreachBinding(node.getName(), context); + if (binding == null) { throw new SemanticCheckException("Unresolved foreach placeholder <<" + node.getName() + ">>"); } - if ("FIELD".equals(name)) { - return context.relBuilder.field(value); + return foreachBindingToRexNode(node.getName(), binding, context); + } + + private ForeachBinding foreachBinding(String name, CalcitePlanContext context) { + return context.getForeachBindings().get(name.toUpperCase(Locale.ROOT)); + } + + private RexNode foreachBindingToRexNode( + String name, ForeachBinding binding, CalcitePlanContext context) { + switch (binding.type()) { + case FIELD: + return context.relBuilder.field(binding.value()); + case LAMBDA: + RexLambdaRef lambdaRef = context.getRexLambdaRefMap().get(binding.value()); + if (lambdaRef == null) { + throw new SemanticCheckException("Unresolved foreach lambda placeholder " + name); + } + return lambdaRef; + case PAIR_ITEM: + return foreachPairItem(binding.value(), 0, name, context); + case PAIR_ITER: + return foreachPairItem(binding.value(), 1, name, context); + case LITERAL: + default: + return context.rexBuilder.makeLiteral(binding.value()); + } + } + + private RexNode foreachPairItem( + String pairName, int index, String placeholderName, CalcitePlanContext context) { + RexLambdaRef pair = context.getRexLambdaRefMap().get(pairName); + if (pair == null) { + throw new SemanticCheckException("Unresolved foreach lambda placeholder " + placeholderName); } - return context.rexBuilder.makeLiteral(value); + return PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, + BuiltinFunctionName.MVINDEX, + pair, + context.rexBuilder.makeExactLiteral(BigDecimal.valueOf(index))); } @Override diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java index 5ce02e6e0d3..5743f688417 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java @@ -69,6 +69,7 @@ public enum BuiltinFunctionName { /** Collection functions */ ARRAY(FunctionName.of("array")), + FOREACH_JSON_ARRAY(FunctionName.of("foreach_json_array"), true), ARRAY_LENGTH(FunctionName.of("array_length")), ARRAY_SLICE(FunctionName.of("array_slice"), true), ARRAY_COMPACT(FunctionName.of("array_compact")), diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java index a24015de992..519d6179503 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java @@ -54,6 +54,7 @@ import org.opensearch.sql.expression.function.CollectionUDF.MapRemoveFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.ReduceFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.TransformFunctionImpl; +import org.opensearch.sql.expression.function.jsonUDF.ForeachJsonArrayFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonAppendFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonArrayLengthFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonDeleteFunctionImpl; @@ -405,6 +406,8 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable { public static final SqlOperator FORALL = new ForallFunctionImpl().toUDF("forall"); public static final SqlOperator EXISTS = new ExistsFunctionImpl().toUDF("exists"); public static final SqlOperator ARRAY = new ArrayFunctionImpl().toUDF("array"); + public static final SqlOperator FOREACH_JSON_ARRAY = + new ForeachJsonArrayFunctionImpl().toUDF("foreach_json_array"); public static final SqlOperator MAP_APPEND = new MapAppendFunctionImpl().toUDF("map_append"); public static final SqlOperator MAP_REMOVE = new MapRemoveFunctionImpl().toUDF("MAP_REMOVE"); public static final SqlOperator MVAPPEND = new MVAppendFunctionImpl().toUDF("mvappend"); diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java index 5750bf6cae8..db47f3d558f 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java @@ -75,6 +75,7 @@ import static org.opensearch.sql.expression.function.BuiltinFunctionName.FIRST; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FLOOR; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FORALL; +import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_JSON_ARRAY; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FROM_DAYS; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FROM_UNIXTIME; import static org.opensearch.sql.expression.function.BuiltinFunctionName.GET_FORMAT; @@ -1087,6 +1088,7 @@ void populate() { registerOperator(FILTER, PPLBuiltinOperators.FILTER); registerOperator(TRANSFORM, PPLBuiltinOperators.TRANSFORM); registerOperator(REDUCE, PPLBuiltinOperators.REDUCE); + registerOperator(FOREACH_JSON_ARRAY, PPLBuiltinOperators.FOREACH_JSON_ARRAY); // Register Json function register( diff --git a/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java b/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java new file mode 100644 index 00000000000..387e2b312ec --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java @@ -0,0 +1,76 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.expression.function.jsonUDF; + +import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.gson; + +import com.google.gson.JsonSyntaxException; +import java.util.List; +import org.apache.calcite.adapter.enumerable.NotNullImplementor; +import org.apache.calcite.adapter.enumerable.NullPolicy; +import org.apache.calcite.adapter.enumerable.RexImpTable; +import org.apache.calcite.adapter.enumerable.RexToLixTranslator; +import org.apache.calcite.linq4j.tree.Expression; +import org.apache.calcite.linq4j.tree.Types; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.schema.impl.ScalarFunctionImpl; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.SqlReturnTypeInference; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.sql.type.SqlTypeUtil; +import org.opensearch.sql.expression.function.ImplementorUDF; +import org.opensearch.sql.expression.function.UDFOperandMetadata; + +/** Converts a JSON array string to a Calcite enumerable array for foreach collection modes. */ +public class ForeachJsonArrayFunctionImpl extends ImplementorUDF { + public ForeachJsonArrayFunctionImpl() { + super(new ForeachJsonArrayImplementor(), NullPolicy.ANY); + } + + @Override + public SqlReturnTypeInference getReturnTypeInference() { + return opBinding -> + SqlTypeUtil.createArrayType( + opBinding.getTypeFactory(), + opBinding + .getTypeFactory() + .createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(SqlTypeName.DOUBLE), true), + true); + } + + @Override + public UDFOperandMetadata getOperandMetadata() { + return UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.ANY)); + } + + public static class ForeachJsonArrayImplementor implements NotNullImplementor { + @Override + public Expression implement( + RexToLixTranslator translator, RexCall call, List translatedOperands) { + ScalarFunctionImpl function = + (ScalarFunctionImpl) + ScalarFunctionImpl.create( + Types.lookupMethod(ForeachJsonArrayFunctionImpl.class, "eval", Object[].class)); + return function.getImplementor().implement(translator, call, RexImpTable.NullAs.NULL); + } + } + + public static Object eval(Object... args) { + if (args.length != 1 || args[0] == null) { + return null; + } + if (args[0] instanceof List) { + return args[0]; + } + try { + return gson.fromJson(String.valueOf(args[0]), List.class); + } catch (JsonSyntaxException e) { + return null; + } + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java index c2fe110b4b1..d6cb5e61663 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java @@ -62,4 +62,54 @@ public void testForeachWildcardMatchstr() throws IOException { verifySchema(result, schema("copy_cpu", "bigint"), schema("copy_mem", "bigint")); verifyDataRows(result, rows(10, 20), rows(30, 40)); } + + @Test + public void testForeachCustomPlaceholders() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | foreach fieldstr=F matchstr=M value_* [ eval copy_<> = F ] |" + + " fields copy_cpu, copy_mem"); + verifySchema(result, schema("copy_cpu", "bigint"), schema("copy_mem", "bigint")); + verifyDataRows(result, rows(10, 20), rows(30, 40)); + } + + @Test + public void testForeachMultivalueMode() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval nums = array(1, 2, 3), total = 0 | foreach mode=multivalue" + + " itemstr=NUMBER nums [ eval total = total + NUMBER ] | fields total"); + verifySchema(result, schema("total", "int")); + verifyDataRows(result, rows(6), rows(6)); + } + + @Test + public void testForeachMultivalueIterPlaceholder() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval nums = array(10, 20, 30), total = 0 | foreach" + + " mode=multivalue iterstr=IDX nums [ eval total = total + IDX ] | fields total"); + verifySchema(result, schema("total", "int")); + verifyDataRows(result, rows(3), rows(3)); + } + + @Test + public void testForeachJsonArrayMode() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval total = 0 | foreach mode=json_array '[1,2,3]' [ eval total" + + " = total + <> ] | fields total"); + verifySchema(result, schema("total", "double")); + verifyDataRows(result, rows(6.0), rows(6.0)); + } + + @Test + public void testForeachAutoCollectionsMode() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval nums = array(1, 2, 3), total = 0 | foreach" + + " mode=auto_collections nums [ eval total = total + <> ] | fields total"); + verifySchema(result, schema("total", "int")); + verifyDataRows(result, rows(6), rows(6)); + } } diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4 index 6635e52a7a0..ce1bd19765b 100644 --- a/ppl/src/main/antlr/OpenSearchPPLParser.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLParser.g4 @@ -378,16 +378,17 @@ evalCommand ; foreachCommand - : FOREACH foreachOption* foreachFieldPattern (COMMA? foreachFieldPattern)* LT_SQR_PRTHS foreachEvalCommand RT_SQR_PRTHS + : FOREACH foreachOption* foreachTarget (COMMA? foreachTarget)* LT_SQR_PRTHS foreachEvalCommand RT_SQR_PRTHS ; -foreachFieldPattern +foreachTarget : wcFieldExpression | STAR + | logicalExpression ; foreachOption - : ident EQUAL ident + : ident EQUAL (ident | foreachPlaceholder) ; foreachEvalCommand diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java index 94482b21bb2..be609ffc363 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java @@ -835,14 +835,20 @@ public UnresolvedPlan visitEvalCommand(EvalCommandContext ctx) { @Override public UnresolvedPlan visitForeachCommand(OpenSearchPPLParser.ForeachCommandContext ctx) { - String mode = - ctx.foreachOption().stream() - .filter(option -> "mode".equalsIgnoreCase(option.ident(0).getText())) - .map(option -> option.ident(1).getText()) - .findFirst() - .orElse("multifield"); + java.util.Map options = new LinkedHashMap<>(); + ctx.foreachOption() + .forEach( + option -> + options.put( + option.ident(0).getText().toLowerCase(Locale.ROOT), + normalizeForeachOptionValue(option.getChild(2).getText()))); + String mode = options.getOrDefault("mode", "multifield"); List patterns = - ctx.foreachFieldPattern().stream().map(this::getTextInQuery).collect(Collectors.toList()); + ctx.foreachTarget().stream().map(this::getTextInQuery).collect(Collectors.toList()); + UnresolvedExpression collectionExpression = + "multifield".equalsIgnoreCase(mode) + ? null + : buildForeachCollectionExpression(ctx.foreachTarget()); List evalClauses = ctx.foreachEvalCommand().foreachEvalClause().stream() .map( @@ -851,7 +857,27 @@ public UnresolvedPlan visitForeachCommand(OpenSearchPPLParser.ForeachCommandCont getTextInQuery(clause.target), expressionBuilder.visit(clause.logicalExpression()))) .collect(Collectors.toList()); - return new Foreach(mode, patterns, evalClauses); + return new Foreach(mode, options, patterns, collectionExpression, evalClauses); + } + + private String normalizeForeachOptionValue(String value) { + if (value.startsWith("<<") && value.endsWith(">>")) { + return value.substring(2, value.length() - 2); + } + return value; + } + + private UnresolvedExpression buildForeachCollectionExpression( + List targets) { + if (targets.size() != 1) { + throw new IllegalArgumentException("foreach collection modes accept exactly one field"); + } + OpenSearchPPLParser.ForeachTargetContext target = targets.get(0); + if (target.logicalExpression() != null) { + return expressionBuilder.visit(target.logicalExpression()); + } + String fieldName = getTextInQuery(target); + return new Field(new QualifiedName(fieldName)); } @Override diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java index e249bad1b90..cfb24ae0725 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java @@ -107,6 +107,31 @@ public void testForeachCommandWithModeShouldPass() { assertNotEquals(null, tree); } + @Test + public void testForeachCollectionModesShouldPass() { + assertNotEquals( + null, + new PPLSyntaxParser() + .parse( + "source=t | foreach mode=multivalue itemstr=<> nums [ eval total = total +" + + " <> ]")); + assertNotEquals( + null, + new PPLSyntaxParser() + .parse( + "source=t | foreach mode=json_array '[1,2,3]' [ eval total = total + <> ]")); + assertNotEquals( + null, + new PPLSyntaxParser() + .parse("source=t | foreach mode=auto_collections nums [ eval total = total + ITEM ]")); + assertNotEquals( + null, + new PPLSyntaxParser() + .parse( + "source=t | foreach mode=multivalue iterstr=IDX nums [ eval total = total + IDX" + + " ]")); + } + @Test public void testSearchFieldsCommandCrossClusterShouldPass() { ParseTree tree = new PPLSyntaxParser().parse("search source=c:t a=1 b=2 | fields a,b"); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java index 7379d094e73..a1017f61880 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java @@ -5,6 +5,8 @@ package org.opensearch.sql.ppl.calcite; +import static org.junit.Assert.assertNotNull; + import org.apache.calcite.rel.RelNode; import org.apache.calcite.test.CalciteAssert; import org.junit.Test; @@ -44,4 +46,57 @@ public void testForeachWildcardAndMatchstr() { + " LogicalTableScan(table=[[scott, EMP]])\n"; verifyLogical(root, expectedLogical); } + + @Test + public void testForeachCustomMultifieldPlaceholders() { + String ppl = + "source=EMP | foreach fieldstr=F matchstr=M *NO [ eval copy_<> = F ] | fields" + + " copy_EMP, copy_DEPT"; + RelNode root = getRelNode(ppl); + + String expectedLogical = + "LogicalProject(copy_EMP=[$0], copy_DEPT=[$7])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + } + + @Test + public void testForeachMultivaluePlansReduce() { + String ppl = + "source=EMP | eval nums = array(1, 2, 3), total = 0 | foreach mode=multivalue" + + " itemstr=NUMBER nums [ eval total = total + NUMBER ] | fields total"; + RelNode root = getRelNode(ppl); + + assertNotNull(root); + } + + @Test + public void testForeachJsonArrayPlansReduce() { + String ppl = + "source=EMP | eval total = 0 | foreach mode=json_array '[1,2,3]' [ eval total = total +" + + " <> ] | fields total"; + RelNode root = getRelNode(ppl); + + assertNotNull(root); + } + + @Test + public void testForeachAutoCollectionsPlansReduce() { + String ppl = + "source=EMP | eval nums = array(1, 2, 3), total = 0 | foreach mode=auto_collections nums [" + + " eval total = total + <> ] | fields total"; + RelNode root = getRelNode(ppl); + + assertNotNull(root); + } + + @Test + public void testForeachIterPlaceholderPlansReduce() { + String ppl = + "source=EMP | eval nums = array(10, 20, 30), total = 0 | foreach mode=multivalue" + + " iterstr=IDX nums [ eval total = total + IDX ] | fields total"; + RelNode root = getRelNode(ppl); + + assertNotNull(root); + } } From bbd04c2f81858043eeaba6a463b2bbfdf6260871 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Wed, 8 Jul 2026 18:45:20 +0800 Subject: [PATCH 03/12] Extract foreach planner constants Signed-off-by: Songkan Tang --- .../sql/calcite/CalciteRelNodeVisitor.java | 79 ++++++++++++------- 1 file changed, 50 insertions(+), 29 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index 3919ecffd49..45c0dd2826d 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -39,6 +39,7 @@ import java.util.HashSet; import java.util.LinkedHashSet; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -210,6 +211,25 @@ public class CalciteRelNodeVisitor extends AbstractNodeVisitor bindings = new HashMap<>(); - putBinding(bindings, itemName, pairName, ForeachBindingType.PAIR_ITEM); - putBinding(bindings, "ITEM", pairName, ForeachBindingType.PAIR_ITEM); - putBinding(bindings, iterName, pairName, ForeachBindingType.PAIR_ITER); - putBinding(bindings, "ITER", pairName, ForeachBindingType.PAIR_ITER); + putBinding(bindings, itemName, FOREACH_PAIR, ForeachBindingType.PAIR_ITEM); + putBinding(bindings, FOREACH_PLACEHOLDER_ITEM, FOREACH_PAIR, ForeachBindingType.PAIR_ITEM); + putBinding(bindings, iterName, FOREACH_PAIR, ForeachBindingType.PAIR_ITER); + putBinding(bindings, FOREACH_PLACEHOLDER_ITER, FOREACH_PAIR, ForeachBindingType.PAIR_ITER); context.pushForeachBindings(bindings); try { @@ -1333,7 +1350,7 @@ private RelNode visitForeachCollection( new Field(new QualifiedName(alias)), new LambdaFunction( clause.getExpression(), - List.of(new QualifiedName(alias), new QualifiedName(pairName))))); + List.of(new QualifiedName(alias), new QualifiedName(FOREACH_PAIR))))); RexNode expr = rexVisitor.analyze(reduce, context); projectPlusOverriding( List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); @@ -1348,39 +1365,43 @@ private UnresolvedExpression collectionExpressionForMode( UnresolvedExpression collectionExpression, CalcitePlanContext context, String normalizedMode) { - if ("multivalue".equals(normalizedMode)) { + if (FOREACH_MODE_MULTIVALUE.equals(normalizedMode)) { return collectionExpression; } - if ("auto_collections".equals(normalizedMode)) { + if (FOREACH_MODE_AUTO_COLLECTIONS.equals(normalizedMode)) { RexNode rexNode = rexVisitor.analyze(collectionExpression, context); if (rexNode.getType() instanceof ArraySqlType) { return collectionExpression; } } - return new Function("foreach_json_array", List.of(collectionExpression)); + return new Function(FOREACH_JSON_ARRAY_FUNCTION, List.of(collectionExpression)); } private Map buildForeachMultifieldBindings( List patterns, String fieldName, Map options) { Map bindings = new HashMap<>(); - putBinding(bindings, "FIELD", fieldName, ForeachBindingType.FIELD); - putBinding(bindings, option(options, "fieldstr", "FIELD"), fieldName, ForeachBindingType.FIELD); + putBinding(bindings, FOREACH_PLACEHOLDER_FIELD, fieldName, ForeachBindingType.FIELD); + putBinding( + bindings, + option(options, FOREACH_OPTION_FIELDSTR, FOREACH_PLACEHOLDER_FIELD), + fieldName, + ForeachBindingType.FIELD); for (String pattern : patterns) { List segments = wildcardSegments(pattern, fieldName); if (segments != null) { String matchStr = String.join("", segments); - putBinding(bindings, "MATCHSTR", matchStr, ForeachBindingType.LITERAL); + putBinding(bindings, FOREACH_PLACEHOLDER_MATCHSTR, matchStr, ForeachBindingType.LITERAL); putBinding( bindings, - option(options, "matchstr", "MATCHSTR"), + option(options, FOREACH_OPTION_MATCHSTR, FOREACH_PLACEHOLDER_MATCHSTR), matchStr, ForeachBindingType.LITERAL); for (int i = 0; i < segments.size(); i++) { - String key = "MATCHSEG" + (i + 1); + String key = FOREACH_PLACEHOLDER_MATCHSEG + (i + 1); putBinding(bindings, key, segments.get(i), ForeachBindingType.LITERAL); putBinding( bindings, - option(options, "matchseg" + (i + 1), key), + option(options, FOREACH_OPTION_MATCHSEG + (i + 1), key), segments.get(i), ForeachBindingType.LITERAL); } @@ -1416,7 +1437,7 @@ private String option(Map options, String name, String defaultVa private void putBinding( Map bindings, String name, String value, ForeachBindingType type) { - bindings.put(name.toUpperCase(java.util.Locale.ROOT), new ForeachBinding(value, type)); + bindings.put(name.toUpperCase(Locale.ROOT), new ForeachBinding(value, type)); } private String substituteForeachTemplate(String template, Map bindings) { From c23aa22bca4fd82c0d21b15c11fc729aa8abddb1 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Wed, 8 Jul 2026 20:22:31 +0800 Subject: [PATCH 04/12] Complete foreach collection type inference Signed-off-by: Songkan Tang --- .../sql/calcite/CalcitePlanContext.java | 15 +- .../sql/calcite/CalciteRelNodeVisitor.java | 323 +++++++++++++++--- .../sql/calcite/CalciteRexNodeVisitor.java | 29 +- .../function/BuiltinFunctionName.java | 2 + .../ForeachPairCollectionFunctionImpl.java | 84 +++++ .../ForeachPairItemFunctionImpl.java | 81 +++++ .../function/CollectionUDF/LambdaUtils.java | 5 + .../function/PPLBuiltinOperators.java | 6 + .../expression/function/PPLFuncImpTable.java | 4 + .../jsonUDF/ForeachJsonArrayFunctionImpl.java | 63 ++-- .../remote/CalciteForeachCommandIT.java | 46 +++ ppl/src/main/antlr/OpenSearchPPLParser.g4 | 7 +- .../opensearch/sql/ppl/parser/AstBuilder.java | 43 ++- .../sql/ppl/antlr/PPLSyntaxParserTest.java | 8 + .../ppl/calcite/CalcitePPLForeachTest.java | 43 +++ 15 files changed, 660 insertions(+), 99 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairCollectionFunctionImpl.java create mode 100644 core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index 6c3fac7ff0a..e6c84eb9fb7 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -21,6 +21,7 @@ import org.apache.calcite.rex.RexLambdaRef; import org.apache.calcite.rex.RexNode; import org.apache.calcite.tools.FrameworkConfig; +import org.checkerframework.checker.nullness.qual.Nullable; import org.opensearch.sql.ast.expression.AggregateFunction; import org.opensearch.sql.ast.expression.Function; import org.opensearch.sql.ast.expression.UnresolvedExpression; @@ -194,10 +195,20 @@ public enum ForeachBindingType { LITERAL, LAMBDA, PAIR_ITEM, - PAIR_ITER + PAIR_ITER, + PAIR_EXTRA } - public record ForeachBinding(String value, ForeachBindingType type) {} + public record ForeachBinding( + String value, + ForeachBindingType type, + int pairIndex, + String typeName, + @Nullable String componentTypeName) { + public ForeachBinding(String value, ForeachBindingType type) { + this(value, type, -1, null, null); + } + } public void putRexLambdaRefMap(Map candidateMap) { this.rexLambdaRefMap.putAll(candidateMap); diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index 45c0dd2826d..27bef3d827d 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -37,6 +37,7 @@ import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Locale; @@ -44,6 +45,8 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -97,7 +100,9 @@ import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta; import org.opensearch.sql.ast.expression.Argument; import org.opensearch.sql.ast.expression.Argument.ArgumentMap; +import org.opensearch.sql.ast.expression.DataType; import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.ForeachPlaceholder; import org.opensearch.sql.ast.expression.Function; import org.opensearch.sql.ast.expression.LambdaFunction; import org.opensearch.sql.ast.expression.Let; @@ -229,6 +234,9 @@ public class CalciteRelNodeVisitor extends AbstractNodeVisitor 1) { throw new SemanticCheckException( "foreach " + normalizedMode + " mode accepts exactly one field"); } + String itemName = + node.getOptions().getOrDefault(FOREACH_OPTION_ITEMSTR, FOREACH_PLACEHOLDER_ITEM); + String iterName = + node.getOptions().getOrDefault(FOREACH_OPTION_ITERSTR, FOREACH_PLACEHOLDER_ITER); UnresolvedExpression collectionExpression = - collectionExpressionForMode(node.getCollectionExpression(), context, normalizedMode); - String itemName = option(node.getOptions(), FOREACH_OPTION_ITEMSTR, FOREACH_PLACEHOLDER_ITEM); - String iterName = option(node.getOptions(), FOREACH_OPTION_ITERSTR, FOREACH_PLACEHOLDER_ITER); - Function pairedCollection = - new Function( - BuiltinFunctionName.TRANSFORM.getName().getFunctionName(), - List.of( - collectionExpression, - new LambdaFunction( - new Function( - BuiltinFunctionName.ARRAY.getName().getFunctionName(), - List.of( - new QualifiedName(FOREACH_TRANSFORM_ITEM), - new QualifiedName(FOREACH_TRANSFORM_ITER))), - List.of( - new QualifiedName(FOREACH_TRANSFORM_ITEM), - new QualifiedName(FOREACH_TRANSFORM_ITER))))); + collectionExpressionForMode( + sourceCollection, context, normalizedMode, node.getEvalClauses(), itemName); + RexNode collectionRex = rexVisitor.analyze(collectionExpression, context); + RelDataType itemType = ((ArraySqlType) collectionRex.getType()).getComponentType(); + List targetAliases = + node.getEvalClauses().stream().map(Foreach.ForeachEvalClause::getTargetTemplate).toList(); + List capturedFields = + collectForeachCapturedFields( + node.getEvalClauses(), context, itemName, iterName, targetAliases); + List pairArguments = new ArrayList<>(); + pairArguments.add(collectionExpression); + capturedFields.stream() + .map(field -> new Field(new QualifiedName(field.name()))) + .forEach(pairArguments::add); + Function pairedCollection = new Function(FOREACH_PAIR_COLLECTION_FUNCTION, pairArguments); Map bindings = new HashMap<>(); - putBinding(bindings, itemName, FOREACH_PAIR, ForeachBindingType.PAIR_ITEM); - putBinding(bindings, FOREACH_PLACEHOLDER_ITEM, FOREACH_PAIR, ForeachBindingType.PAIR_ITEM); - putBinding(bindings, iterName, FOREACH_PAIR, ForeachBindingType.PAIR_ITER); - putBinding(bindings, FOREACH_PLACEHOLDER_ITER, FOREACH_PAIR, ForeachBindingType.PAIR_ITER); + bindings.put( + itemName.toUpperCase(Locale.ROOT), + newForeachPairBinding(FOREACH_PAIR, ForeachBindingType.PAIR_ITEM, 0, itemType)); + bindings.put( + FOREACH_PLACEHOLDER_ITEM, + newForeachPairBinding(FOREACH_PAIR, ForeachBindingType.PAIR_ITEM, 0, itemType)); + bindings.put( + iterName.toUpperCase(Locale.ROOT), + newForeachPairBinding( + FOREACH_PAIR, + ForeachBindingType.PAIR_ITER, + 1, + context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER))); + bindings.put( + FOREACH_PLACEHOLDER_ITER, + newForeachPairBinding( + FOREACH_PAIR, + ForeachBindingType.PAIR_ITER, + 1, + context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER))); + for (int i = 0; i < capturedFields.size(); i++) { + ForeachCapturedField field = capturedFields.get(i); + bindings.put( + field.name().toUpperCase(Locale.ROOT), + newForeachPairBinding(FOREACH_PAIR, ForeachBindingType.PAIR_EXTRA, i + 2, field.type())); + } context.pushForeachBindings(bindings); try { @@ -1361,10 +1397,75 @@ private RelNode visitForeachCollection( return context.relBuilder.peek(); } + private List collectForeachCapturedFields( + List evalClauses, + CalcitePlanContext context, + String itemName, + String iterName, + List targetAliases) { + Set excluded = + Stream.concat( + Stream.of( + itemName, + iterName, + FOREACH_PLACEHOLDER_ITEM, + FOREACH_PLACEHOLDER_ITER, + FOREACH_PAIR), + targetAliases.stream()) + .map(name -> name.toUpperCase(Locale.ROOT)) + .collect(Collectors.toSet()); + Map rowTypes = + context.relBuilder.peek().getRowType().getFieldList().stream() + .collect( + Collectors.toMap( + field -> field.getName().toUpperCase(Locale.ROOT), + RelDataTypeField::getType, + (left, right) -> left, + LinkedHashMap::new)); + Map fields = new LinkedHashMap<>(); + evalClauses.forEach( + clause -> collectForeachCapturedFields(clause.getExpression(), excluded, rowTypes, fields)); + return new ArrayList<>(fields.values()); + } + + private void collectForeachCapturedFields( + Node node, + Set excluded, + Map rowTypes, + Map fields) { + if (node instanceof QualifiedName qualifiedName) { + String key = qualifiedName.toString().toUpperCase(Locale.ROOT); + RelDataType type = rowTypes.get(key); + if (type != null && !excluded.contains(key)) { + fields.putIfAbsent(key, new ForeachCapturedField(qualifiedName.toString(), type)); + } + } + node.getChild() + .forEach(child -> collectForeachCapturedFields(child, excluded, rowTypes, fields)); + } + + private ForeachBinding newForeachPairBinding( + String pairName, ForeachBindingType type, int index, RelDataType relDataType) { + return new ForeachBinding( + pairName, type, index, relDataType.getSqlTypeName().name(), componentTypeName(relDataType)); + } + + @Nullable + private String componentTypeName(RelDataType relDataType) { + if (relDataType instanceof ArraySqlType arrayType) { + return arrayType.getComponentType().getSqlTypeName().name(); + } + return null; + } + + private record ForeachCapturedField(String name, RelDataType type) {} + private UnresolvedExpression collectionExpressionForMode( UnresolvedExpression collectionExpression, CalcitePlanContext context, - String normalizedMode) { + String normalizedMode, + List evalClauses, + String itemName) { if (FOREACH_MODE_MULTIVALUE.equals(normalizedMode)) { return collectionExpression; } @@ -1374,36 +1475,158 @@ private UnresolvedExpression collectionExpressionForMode( return collectionExpression; } } - return new Function(FOREACH_JSON_ARRAY_FUNCTION, List.of(collectionExpression)); + if (isJsonArrayFunction(collectionExpression)) { + Function jsonArray = (Function) collectionExpression; + validateHomogeneousJsonArrayArguments(jsonArray); + return new Function( + FOREACH_JSON_ARRAY_FUNCTION, + List.of( + collectionExpression, + new Literal( + inferForeachJsonArrayType(collectionExpression, evalClauses, itemName), + DataType.STRING))); + } + return new Function( + FOREACH_JSON_ARRAY_FUNCTION, + List.of( + collectionExpression, + new Literal( + inferForeachJsonArrayType(collectionExpression, evalClauses, itemName), + DataType.STRING))); + } + + private UnresolvedExpression inferAutoCollectionExpression(CalcitePlanContext context) { + return context.relBuilder.peek().getRowType().getFieldList().stream() + .filter(field -> field.getType() instanceof ArraySqlType) + .findFirst() + .map(field -> (UnresolvedExpression) new Field(new QualifiedName(field.getName()))) + .orElseThrow( + () -> + new SemanticCheckException( + "foreach auto_collections mode requires a multivalue field or JSON array")); + } + + private boolean isJsonArrayFunction(UnresolvedExpression expression) { + return expression instanceof Function function + && BuiltinFunctionName.JSON_ARRAY + .getName() + .getFunctionName() + .equalsIgnoreCase(function.getFuncName()); + } + + private String inferForeachJsonArrayType( + UnresolvedExpression collectionExpression, + List evalClauses, + String itemName) { + if (isJsonArrayFunction(collectionExpression)) { + return jsonArrayTypeFromArguments((Function) collectionExpression) + .orElse(FOREACH_JSON_ARRAY_STRING_TYPE); + } + if (foreachItemUsedInArithmetic(evalClauses, itemName)) { + return FOREACH_JSON_ARRAY_NUMERIC_TYPE; + } + return FOREACH_JSON_ARRAY_STRING_TYPE; + } + + private Optional jsonArrayTypeFromArguments(Function jsonArray) { + boolean hasNumeric = false; + boolean hasString = false; + for (UnresolvedExpression argument : jsonArray.getFuncArgs()) { + if (argument instanceof Literal literal) { + switch (literal.getType()) { + case INTEGER, LONG, SHORT, FLOAT, DOUBLE, DECIMAL: + hasNumeric = true; + break; + case STRING: + hasString = true; + break; + case NULL: + break; + default: + return Optional.empty(); + } + } else { + return Optional.empty(); + } + } + if (hasNumeric && hasString) { + throw new SemanticCheckException( + "foreach json_array elements must be consistently strings or numbers"); + } + if (hasNumeric) { + return Optional.of(FOREACH_JSON_ARRAY_NUMERIC_TYPE); + } + if (hasString) { + return Optional.of(FOREACH_JSON_ARRAY_STRING_TYPE); + } + return Optional.empty(); + } + + private void validateHomogeneousJsonArrayArguments(Function jsonArray) { + jsonArrayTypeFromArguments(jsonArray); + } + + private boolean foreachItemUsedInArithmetic( + List evalClauses, String itemName) { + return evalClauses.stream() + .anyMatch(clause -> foreachItemUsedInArithmetic(clause.getExpression(), itemName, false)); + } + + private boolean foreachItemUsedInArithmetic(Node node, String itemName, boolean inArithmetic) { + if (node instanceof UnresolvedExpression expression + && isForeachItemReference(expression, itemName)) { + return inArithmetic; + } + boolean arithmetic = + inArithmetic + || (node instanceof UnresolvedExpression expression + && isArithmeticExpression(expression)); + return node.getChild().stream() + .anyMatch(child -> foreachItemUsedInArithmetic(child, itemName, arithmetic)); + } + + private boolean isForeachItemReference(UnresolvedExpression expression, String itemName) { + if (expression instanceof ForeachPlaceholder placeholder) { + return FOREACH_PLACEHOLDER_ITEM.equalsIgnoreCase(placeholder.getName()) + || itemName.equalsIgnoreCase(placeholder.getName()); + } + return expression instanceof QualifiedName qualifiedName + && (FOREACH_PLACEHOLDER_ITEM.equalsIgnoreCase(qualifiedName.toString()) + || itemName.equalsIgnoreCase(qualifiedName.toString())); + } + + private boolean isArithmeticExpression(UnresolvedExpression expression) { + return expression instanceof Function function + && Set.of("+", "-", "*", "/", "%").contains(function.getFuncName()); } private Map buildForeachMultifieldBindings( List patterns, String fieldName, Map options) { Map bindings = new HashMap<>(); - putBinding(bindings, FOREACH_PLACEHOLDER_FIELD, fieldName, ForeachBindingType.FIELD); - putBinding( - bindings, - option(options, FOREACH_OPTION_FIELDSTR, FOREACH_PLACEHOLDER_FIELD), - fieldName, - ForeachBindingType.FIELD); + bindings.put( + FOREACH_PLACEHOLDER_FIELD, new ForeachBinding(fieldName, ForeachBindingType.FIELD)); + bindings.put( + options + .getOrDefault(FOREACH_OPTION_FIELDSTR, FOREACH_PLACEHOLDER_FIELD) + .toUpperCase(Locale.ROOT), + new ForeachBinding(fieldName, ForeachBindingType.FIELD)); for (String pattern : patterns) { List segments = wildcardSegments(pattern, fieldName); if (segments != null) { String matchStr = String.join("", segments); - putBinding(bindings, FOREACH_PLACEHOLDER_MATCHSTR, matchStr, ForeachBindingType.LITERAL); - putBinding( - bindings, - option(options, FOREACH_OPTION_MATCHSTR, FOREACH_PLACEHOLDER_MATCHSTR), - matchStr, - ForeachBindingType.LITERAL); + bindings.put( + FOREACH_PLACEHOLDER_MATCHSTR, new ForeachBinding(matchStr, ForeachBindingType.LITERAL)); + bindings.put( + options + .getOrDefault(FOREACH_OPTION_MATCHSTR, FOREACH_PLACEHOLDER_MATCHSTR) + .toUpperCase(Locale.ROOT), + new ForeachBinding(matchStr, ForeachBindingType.LITERAL)); for (int i = 0; i < segments.size(); i++) { String key = FOREACH_PLACEHOLDER_MATCHSEG + (i + 1); - putBinding(bindings, key, segments.get(i), ForeachBindingType.LITERAL); - putBinding( - bindings, - option(options, FOREACH_OPTION_MATCHSEG + (i + 1), key), - segments.get(i), - ForeachBindingType.LITERAL); + bindings.put(key, new ForeachBinding(segments.get(i), ForeachBindingType.LITERAL)); + bindings.put( + options.getOrDefault(FOREACH_OPTION_MATCHSEG + (i + 1), key).toUpperCase(Locale.ROOT), + new ForeachBinding(segments.get(i), ForeachBindingType.LITERAL)); } break; } @@ -1418,9 +1641,8 @@ private List wildcardSegments(String pattern, String fieldName) { if (!WildcardUtils.containsWildcard(pattern)) { return List.of(); } - java.util.regex.Matcher matcher = - java.util.regex.Pattern.compile(WildcardUtils.convertWildcardPatternToRegex(pattern)) - .matcher(fieldName); + Matcher matcher = + Pattern.compile(WildcardUtils.convertWildcardPatternToRegex(pattern)).matcher(fieldName); if (!matcher.matches()) { return null; } @@ -1431,15 +1653,6 @@ private List wildcardSegments(String pattern, String fieldName) { return segments; } - private String option(Map options, String name, String defaultValue) { - return options.getOrDefault(name, defaultValue); - } - - private void putBinding( - Map bindings, String name, String value, ForeachBindingType type) { - bindings.put(name.toUpperCase(Locale.ROOT), new ForeachBinding(value, type)); - } - private String substituteForeachTemplate(String template, Map bindings) { String result = template; for (Map.Entry entry : bindings.entrySet()) { diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java index 5d228723c94..f7c3aaf76d7 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -433,9 +433,10 @@ private RexNode foreachBindingToRexNode( } return lambdaRef; case PAIR_ITEM: - return foreachPairItem(binding.value(), 0, name, context); + case PAIR_EXTRA: + return foreachPairItem(binding, name, context); case PAIR_ITER: - return foreachPairItem(binding.value(), 1, name, context); + return foreachPairItem(binding, name, context); case LITERAL: default: return context.rexBuilder.makeLiteral(binding.value()); @@ -443,16 +444,21 @@ private RexNode foreachBindingToRexNode( } private RexNode foreachPairItem( - String pairName, int index, String placeholderName, CalcitePlanContext context) { - RexLambdaRef pair = context.getRexLambdaRefMap().get(pairName); + ForeachBinding binding, String placeholderName, CalcitePlanContext context) { + RexLambdaRef pair = context.getRexLambdaRefMap().get(binding.value()); if (pair == null) { throw new SemanticCheckException("Unresolved foreach lambda placeholder " + placeholderName); } return PPLFuncImpTable.INSTANCE.resolve( context.rexBuilder, - BuiltinFunctionName.MVINDEX, + BuiltinFunctionName.FOREACH_PAIR_ITEM, pair, - context.rexBuilder.makeExactLiteral(BigDecimal.valueOf(index))); + context.rexBuilder.makeExactLiteral(BigDecimal.valueOf(binding.pairIndex())), + context.rexBuilder.makeLiteral(binding.typeName()), + binding.componentTypeName() == null + ? context.rexBuilder.makeNullLiteral( + context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.NULL)) + : context.rexBuilder.makeLiteral(binding.componentTypeName())); } @Override @@ -631,7 +637,8 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) { List capturedVars = null; try { - for (UnresolvedExpression arg : args) { + for (int argIndex = 0; argIndex < args.size(); argIndex++) { + UnresolvedExpression arg = args.get(argIndex); if (arg instanceof LambdaFunction) { CalcitePlanContext lambdaContext = prepareLambdaContext( @@ -650,6 +657,14 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) { arguments.add(lambdaNode); // Capture any external variables that were referenced in the lambda capturedVars = lambdaContext.getCapturedVariables(); + } else if (node.getFuncName().equalsIgnoreCase("reduce") && argIndex == 0) { + Map foreachBindings = new HashMap<>(context.getForeachBindings()); + context.clearForeachBindings(); + try { + arguments.add(analyze(arg, context)); + } finally { + context.pushForeachBindings(foreachBindings); + } } else { arguments.add(analyze(arg, context)); } diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java index 5743f688417..e1a4abbf89a 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java @@ -70,6 +70,8 @@ public enum BuiltinFunctionName { /** Collection functions */ ARRAY(FunctionName.of("array")), FOREACH_JSON_ARRAY(FunctionName.of("foreach_json_array"), true), + FOREACH_PAIR_COLLECTION(FunctionName.of("foreach_pair_collection"), true), + FOREACH_PAIR_ITEM(FunctionName.of("foreach_pair_item"), true), ARRAY_LENGTH(FunctionName.of("array_length")), ARRAY_SLICE(FunctionName.of("array_slice"), true), ARRAY_COMPACT(FunctionName.of("array_compact")), diff --git a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairCollectionFunctionImpl.java b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairCollectionFunctionImpl.java new file mode 100644 index 00000000000..9de5d6b3452 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairCollectionFunctionImpl.java @@ -0,0 +1,84 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.expression.function.CollectionUDF; + +import java.util.ArrayList; +import java.util.List; +import org.apache.calcite.adapter.enumerable.NotNullImplementor; +import org.apache.calcite.adapter.enumerable.NullPolicy; +import org.apache.calcite.adapter.enumerable.RexToLixTranslator; +import org.apache.calcite.linq4j.tree.Expression; +import org.apache.calcite.linq4j.tree.Expressions; +import org.apache.calcite.linq4j.tree.Types; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.sql.type.SqlReturnTypeInference; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.sql.type.SqlTypeUtil; +import org.opensearch.sql.expression.function.ImplementorUDF; +import org.opensearch.sql.expression.function.UDFOperandMetadata; + +/** Builds the internal foreach pair array: [item, iter, captured-field...]. */ +public class ForeachPairCollectionFunctionImpl extends ImplementorUDF { + public ForeachPairCollectionFunctionImpl() { + super(new ForeachPairCollectionImplementor(), NullPolicy.NONE); + } + + @Override + public SqlReturnTypeInference getReturnTypeInference() { + return opBinding -> { + RelDataType pair = + opBinding + .getTypeFactory() + .createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(SqlTypeName.OTHER), true); + return SqlTypeUtil.createArrayType(opBinding.getTypeFactory(), pair, true); + }; + } + + @Override + public UDFOperandMetadata getOperandMetadata() { + return null; + } + + public static class ForeachPairCollectionImplementor implements NotNullImplementor { + @Override + public Expression implement( + RexToLixTranslator translator, RexCall call, List translatedOperands) { + return Expressions.call( + Types.lookupMethod(ForeachPairCollectionFunctionImpl.class, "eval", Object[].class), + translatedOperands); + } + } + + public static Object eval(Object... args) { + if (args.length == 0 || args[0] == null) { + return null; + } + List source = toList(args[0]); + List pairs = new ArrayList<>(); + for (int i = 0; i < source.size(); i++) { + Object[] pair = new Object[args.length + 1]; + pair[0] = source.get(i); + pair[1] = i; + for (int j = 1; j < args.length; j++) { + pair[j + 1] = args[j]; + } + pairs.add(pair); + } + return pairs; + } + + private static List toList(Object value) { + if (value instanceof List list) { + return list; + } + if (value instanceof Object[] array) { + return List.of(array); + } + throw new IllegalArgumentException("foreach pair collection requires an array input"); + } +} diff --git a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java new file mode 100644 index 00000000000..b93eefda881 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java @@ -0,0 +1,81 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.expression.function.CollectionUDF; + +import java.util.List; +import org.apache.calcite.adapter.enumerable.NotNullImplementor; +import org.apache.calcite.adapter.enumerable.NullPolicy; +import org.apache.calcite.adapter.enumerable.RexToLixTranslator; +import org.apache.calcite.linq4j.tree.Expression; +import org.apache.calcite.linq4j.tree.Expressions; +import org.apache.calcite.linq4j.tree.Types; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.sql.type.SqlReturnTypeInference; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.sql.type.SqlTypeUtil; +import org.opensearch.sql.expression.function.ImplementorUDF; +import org.opensearch.sql.expression.function.UDFOperandMetadata; + +/** Extracts an internal foreach pair slot while preserving the slot's inferred Calcite type. */ +public class ForeachPairItemFunctionImpl extends ImplementorUDF { + public ForeachPairItemFunctionImpl() { + super(new ForeachPairItemImplementor(), NullPolicy.NONE); + } + + @Override + public SqlReturnTypeInference getReturnTypeInference() { + return opBinding -> { + SqlTypeName typeName = SqlTypeName.valueOf(opBinding.getOperandLiteralValue(2, String.class)); + RelDataType type; + if (typeName == SqlTypeName.ARRAY) { + SqlTypeName componentTypeName = + SqlTypeName.valueOf(opBinding.getOperandLiteralValue(3, String.class)); + type = + SqlTypeUtil.createArrayType( + opBinding.getTypeFactory(), + opBinding + .getTypeFactory() + .createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(componentTypeName), true), + true); + } else { + type = opBinding.getTypeFactory().createSqlType(typeName); + } + return opBinding.getTypeFactory().createTypeWithNullability(type, true); + }; + } + + @Override + public UDFOperandMetadata getOperandMetadata() { + return null; + } + + public static class ForeachPairItemImplementor implements NotNullImplementor { + @Override + public Expression implement( + RexToLixTranslator translator, RexCall call, List translatedOperands) { + return Expressions.call( + Types.lookupMethod(ForeachPairItemFunctionImpl.class, "eval", Object[].class), + translatedOperands); + } + } + + public static Object eval(Object... args) { + if (args.length < 2 || args[0] == null || args[1] == null) { + return null; + } + int index = ((Number) args[1]).intValue(); + if (args[0] instanceof Object[] pair) { + return index < 0 || index >= pair.length ? null : pair[index]; + } + List pair = (List) args[0]; + if (index < 0 || index >= pair.size()) { + return null; + } + return pair.get(index); + } +} diff --git a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java index 314ac3ad945..57bace94116 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java @@ -16,6 +16,7 @@ import org.apache.calcite.rex.RexLambda; import org.apache.calcite.rex.RexLambdaRef; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.type.SqlReturnTypeInference; import org.apache.calcite.sql.type.SqlTypeName; @@ -68,6 +69,10 @@ public static RelDataType inferReturnTypeFromLambda( public static RexCall reInferReturnTypeForRexCallInsideLambda( RexCall rexCall, Map argTypes, RelDataTypeFactory typeFactory) { + if (rexCall.getKind() == SqlKind.CAST + || "CAST".equalsIgnoreCase(rexCall.getOperator().getName())) { + return rexCall; + } List filledOperands = new ArrayList<>(); List rexCallOperands = rexCall.getOperands(); for (RexNode rexNode : rexCallOperands) { diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java index 519d6179503..0b93c2d73d1 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java @@ -47,6 +47,8 @@ import org.opensearch.sql.expression.function.CollectionUDF.ExistsFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.FilterFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.ForallFunctionImpl; +import org.opensearch.sql.expression.function.CollectionUDF.ForeachPairCollectionFunctionImpl; +import org.opensearch.sql.expression.function.CollectionUDF.ForeachPairItemFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.MVAppendFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.MVFindFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.MVZipFunctionImpl; @@ -408,6 +410,10 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable { public static final SqlOperator ARRAY = new ArrayFunctionImpl().toUDF("array"); public static final SqlOperator FOREACH_JSON_ARRAY = new ForeachJsonArrayFunctionImpl().toUDF("foreach_json_array"); + public static final SqlOperator FOREACH_PAIR_COLLECTION = + new ForeachPairCollectionFunctionImpl().toUDF("foreach_pair_collection"); + public static final SqlOperator FOREACH_PAIR_ITEM = + new ForeachPairItemFunctionImpl().toUDF("foreach_pair_item"); public static final SqlOperator MAP_APPEND = new MapAppendFunctionImpl().toUDF("map_append"); public static final SqlOperator MAP_REMOVE = new MapRemoveFunctionImpl().toUDF("MAP_REMOVE"); public static final SqlOperator MVAPPEND = new MVAppendFunctionImpl().toUDF("mvappend"); diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java index db47f3d558f..76080b5abeb 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java @@ -76,6 +76,8 @@ import static org.opensearch.sql.expression.function.BuiltinFunctionName.FLOOR; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FORALL; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_JSON_ARRAY; +import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_PAIR_COLLECTION; +import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_PAIR_ITEM; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FROM_DAYS; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FROM_UNIXTIME; import static org.opensearch.sql.expression.function.BuiltinFunctionName.GET_FORMAT; @@ -1089,6 +1091,8 @@ void populate() { registerOperator(TRANSFORM, PPLBuiltinOperators.TRANSFORM); registerOperator(REDUCE, PPLBuiltinOperators.REDUCE); registerOperator(FOREACH_JSON_ARRAY, PPLBuiltinOperators.FOREACH_JSON_ARRAY); + registerOperator(FOREACH_PAIR_COLLECTION, PPLBuiltinOperators.FOREACH_PAIR_COLLECTION); + registerOperator(FOREACH_PAIR_ITEM, PPLBuiltinOperators.FOREACH_PAIR_ITEM); // Register Json function register( diff --git a/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java b/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java index 387e2b312ec..df4ea7b3da9 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java @@ -8,18 +8,16 @@ import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.gson; import com.google.gson.JsonSyntaxException; +import java.math.BigDecimal; import java.util.List; import org.apache.calcite.adapter.enumerable.NotNullImplementor; import org.apache.calcite.adapter.enumerable.NullPolicy; -import org.apache.calcite.adapter.enumerable.RexImpTable; import org.apache.calcite.adapter.enumerable.RexToLixTranslator; import org.apache.calcite.linq4j.tree.Expression; +import org.apache.calcite.linq4j.tree.Expressions; import org.apache.calcite.linq4j.tree.Types; import org.apache.calcite.rex.RexCall; -import org.apache.calcite.schema.impl.ScalarFunctionImpl; -import org.apache.calcite.sql.type.OperandTypes; import org.apache.calcite.sql.type.SqlReturnTypeInference; -import org.apache.calcite.sql.type.SqlTypeFamily; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.sql.type.SqlTypeUtil; import org.opensearch.sql.expression.function.ImplementorUDF; @@ -28,49 +26,66 @@ /** Converts a JSON array string to a Calcite enumerable array for foreach collection modes. */ public class ForeachJsonArrayFunctionImpl extends ImplementorUDF { public ForeachJsonArrayFunctionImpl() { - super(new ForeachJsonArrayImplementor(), NullPolicy.ANY); + super(new ForeachJsonArrayImplementor(), NullPolicy.NONE); } @Override public SqlReturnTypeInference getReturnTypeInference() { - return opBinding -> - SqlTypeUtil.createArrayType( - opBinding.getTypeFactory(), - opBinding - .getTypeFactory() - .createTypeWithNullability( - opBinding.getTypeFactory().createSqlType(SqlTypeName.DOUBLE), true), - true); + return opBinding -> { + SqlTypeName elementTypeName = + SqlTypeName.valueOf(opBinding.getOperandLiteralValue(1, String.class)); + return SqlTypeUtil.createArrayType( + opBinding.getTypeFactory(), + opBinding + .getTypeFactory() + .createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(elementTypeName), true), + true); + }; } @Override public UDFOperandMetadata getOperandMetadata() { - return UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.ANY)); + return null; } public static class ForeachJsonArrayImplementor implements NotNullImplementor { @Override public Expression implement( RexToLixTranslator translator, RexCall call, List translatedOperands) { - ScalarFunctionImpl function = - (ScalarFunctionImpl) - ScalarFunctionImpl.create( - Types.lookupMethod(ForeachJsonArrayFunctionImpl.class, "eval", Object[].class)); - return function.getImplementor().implement(translator, call, RexImpTable.NullAs.NULL); + return Expressions.call( + Types.lookupMethod(ForeachJsonArrayFunctionImpl.class, "eval", Object[].class), + translatedOperands); } } public static Object eval(Object... args) { - if (args.length != 1 || args[0] == null) { + if (args.length != 2 || args[0] == null) { return null; } - if (args[0] instanceof List) { - return args[0]; - } + SqlTypeName elementType = SqlTypeName.valueOf(String.valueOf(args[1])); try { - return gson.fromJson(String.valueOf(args[0]), List.class); + List values = + args[0] instanceof List list + ? list + : gson.fromJson(String.valueOf(args[0]), List.class); + return values == null + ? null + : values.stream().map(value -> cast(value, elementType)).toList(); } catch (JsonSyntaxException e) { return null; } } + + private static Object cast(Object value, SqlTypeName elementType) { + if (value == null) { + return null; + } + return switch (elementType) { + case DOUBLE -> ((Number) value).doubleValue(); + case VARCHAR -> String.valueOf(value); + case DECIMAL -> BigDecimal.valueOf(((Number) value).doubleValue()); + default -> value; + }; + } } diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java index d6cb5e61663..4cc5e25f389 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteForeachCommandIT.java @@ -11,6 +11,7 @@ import static org.opensearch.sql.util.MatcherUtils.verifySchema; import java.io.IOException; +import java.util.List; import org.json.JSONObject; import org.junit.jupiter.api.Test; import org.opensearch.client.Request; @@ -93,6 +94,21 @@ public void testForeachMultivalueIterPlaceholder() throws IOException { verifyDataRows(result, rows(3), rows(3)); } + @Test + public void testForeachStringItemWithNumericIter() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval word = split('ABCDE', ''), nums = array('1', '2', '3'," + + " '4', '5'), word_and_num = array() | foreach word mode=multivalue [ eval" + + " word_and_num = mvappend(word_and_num, concat(<>, mvindex(nums," + + " <>))) ] | fields word_and_num"); + verifySchema(result, schema("word_and_num", "array")); + verifyDataRows( + result, + rows(List.of("A1", "B2", "C3", "D4", "E5")), + rows(List.of("A1", "B2", "C3", "D4", "E5"))); + } + @Test public void testForeachJsonArrayMode() throws IOException { JSONObject result = @@ -103,6 +119,26 @@ public void testForeachJsonArrayMode() throws IOException { verifyDataRows(result, rows(6.0), rows(6.0)); } + @Test + public void testForeachJsonArrayFunctionWithoutMode() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval total = 0 | foreach json_array(1, 2, 3) [ eval total =" + + " total + <> ] | fields total"); + verifySchema(result, schema("total", "double")); + verifyDataRows(result, rows(6.0), rows(6.0)); + } + + @Test + public void testForeachJsonArrayStringElements() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval result = '' | foreach json_array('a', 'b', 'c') [ eval" + + " result = concat(result, <>) ] | fields result"); + verifySchema(result, schema("result", "string")); + verifyDataRows(result, rows("abc"), rows("abc")); + } + @Test public void testForeachAutoCollectionsMode() throws IOException { JSONObject result = @@ -112,4 +148,14 @@ public void testForeachAutoCollectionsMode() throws IOException { verifySchema(result, schema("total", "int")); verifyDataRows(result, rows(6), rows(6)); } + + @Test + public void testForeachAutoCollectionsWithoutTarget() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach | eval nums = array(1, 2, 3), total = 0 | foreach" + + " mode=auto_collections [ eval total = total + <> ] | fields total"); + verifySchema(result, schema("total", "int")); + verifyDataRows(result, rows(6), rows(6)); + } } diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4 index ce1bd19765b..0309167f344 100644 --- a/ppl/src/main/antlr/OpenSearchPPLParser.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLParser.g4 @@ -378,7 +378,12 @@ evalCommand ; foreachCommand - : FOREACH foreachOption* foreachTarget (COMMA? foreachTarget)* LT_SQR_PRTHS foreachEvalCommand RT_SQR_PRTHS + : FOREACH foreachArgument* LT_SQR_PRTHS foreachEvalCommand RT_SQR_PRTHS + ; + +foreachArgument + : foreachOption + | foreachTarget ; foreachTarget diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java index be609ffc363..8d326623fea 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java @@ -54,6 +54,7 @@ import org.opensearch.sql.ast.expression.DataType; import org.opensearch.sql.ast.expression.EqualTo; import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.Function; import org.opensearch.sql.ast.expression.Let; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.expression.Map; @@ -836,19 +837,28 @@ public UnresolvedPlan visitEvalCommand(EvalCommandContext ctx) { @Override public UnresolvedPlan visitForeachCommand(OpenSearchPPLParser.ForeachCommandContext ctx) { java.util.Map options = new LinkedHashMap<>(); - ctx.foreachOption() - .forEach( - option -> - options.put( - option.ident(0).getText().toLowerCase(Locale.ROOT), - normalizeForeachOptionValue(option.getChild(2).getText()))); - String mode = options.getOrDefault("mode", "multifield"); + List optionContexts = + ctx.foreachArgument().stream() + .map(OpenSearchPPLParser.ForeachArgumentContext::foreachOption) + .filter(java.util.Objects::nonNull) + .collect(Collectors.toList()); + List targetContexts = + ctx.foreachArgument().stream() + .map(OpenSearchPPLParser.ForeachArgumentContext::foreachTarget) + .filter(java.util.Objects::nonNull) + .collect(Collectors.toList()); + optionContexts.forEach( + option -> + options.put( + option.ident(0).getText().toLowerCase(Locale.ROOT), + normalizeForeachOptionValue(option.getChild(2).getText()))); + String mode = options.getOrDefault("mode", inferForeachMode(targetContexts)); List patterns = - ctx.foreachTarget().stream().map(this::getTextInQuery).collect(Collectors.toList()); + targetContexts.stream().map(this::getTextInQuery).collect(Collectors.toList()); UnresolvedExpression collectionExpression = "multifield".equalsIgnoreCase(mode) ? null - : buildForeachCollectionExpression(ctx.foreachTarget()); + : buildForeachCollectionExpression(targetContexts, mode); List evalClauses = ctx.foreachEvalCommand().foreachEvalClause().stream() .map( @@ -867,8 +877,21 @@ private String normalizeForeachOptionValue(String value) { return value; } + private String inferForeachMode(List targets) { + if (targets.size() == 1 + && targets.get(0).logicalExpression() != null + && expressionBuilder.visit(targets.get(0).logicalExpression()) instanceof Function function + && "json_array".equalsIgnoreCase(function.getFuncName())) { + return "json_array"; + } + return "multifield"; + } + private UnresolvedExpression buildForeachCollectionExpression( - List targets) { + List targets, String mode) { + if (targets.isEmpty() && "auto_collections".equalsIgnoreCase(mode)) { + return null; + } if (targets.size() != 1) { throw new IllegalArgumentException("foreach collection modes accept exactly one field"); } diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java index cfb24ae0725..0c5ad9dd4fe 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java @@ -124,6 +124,14 @@ public void testForeachCollectionModesShouldPass() { null, new PPLSyntaxParser() .parse("source=t | foreach mode=auto_collections nums [ eval total = total + ITEM ]")); + assertNotEquals( + null, + new PPLSyntaxParser() + .parse("source=t | foreach mode=auto_collections [ eval total = total + <> ]")); + assertNotEquals( + null, + new PPLSyntaxParser() + .parse("source=t | foreach json_array(1, 2, 3) [ eval total = total + <> ]")); assertNotEquals( null, new PPLSyntaxParser() diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java index a1017f61880..ae22f62e4a7 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java @@ -6,10 +6,12 @@ package org.opensearch.sql.ppl.calcite; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; import org.apache.calcite.rel.RelNode; import org.apache.calcite.test.CalciteAssert; import org.junit.Test; +import org.opensearch.sql.exception.SemanticCheckException; public class CalcitePPLForeachTest extends CalcitePPLAbstractTest { @@ -80,6 +82,25 @@ public void testForeachJsonArrayPlansReduce() { assertNotNull(root); } + @Test + public void testForeachJsonArrayFunctionWithoutModePlansReduce() { + String ppl = + "source=EMP | eval total = 0 | foreach json_array(1, 2, 3) [ eval total = total +" + + " <> ] | fields total"; + RelNode root = getRelNode(ppl); + + assertNotNull(root); + } + + @Test + public void testForeachJsonArrayMixedStringAndNumberFails() { + String ppl = + "source=EMP | eval total = 0 | foreach json_array(1, 2, 'hello') [ eval total = total +" + + " <> ] | fields total"; + + assertThrows(SemanticCheckException.class, () -> getRelNode(ppl)); + } + @Test public void testForeachAutoCollectionsPlansReduce() { String ppl = @@ -90,6 +111,28 @@ public void testForeachAutoCollectionsPlansReduce() { assertNotNull(root); } + @Test + public void testForeachAutoCollectionsWithoutTargetPlansReduce() { + String ppl = + "source=EMP | eval nums = array(1, 2, 3), total = 0 | foreach mode=auto_collections [" + + " eval total = total + <> ] | fields total"; + RelNode root = getRelNode(ppl); + + assertNotNull(root); + } + + @Test + public void testForeachStringItemWithNumericIterPlansReduce() { + String ppl = + "source=EMP | eval word = split('ABCDE', ''), nums = array('1', '2', '3', '4', '5')," + + " word_and_num = array() | foreach word mode=multivalue [ eval word_and_num =" + + " mvappend(word_and_num, concat(<>, mvindex(nums, <>))) ] | fields" + + " word_and_num"; + RelNode root = getRelNode(ppl); + + assertNotNull(root); + } + @Test public void testForeachIterPlaceholderPlansReduce() { String ppl = From 1976c0d745cf6c984310b6be7f6897244b37c6fa Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 9 Jul 2026 03:49:35 +0000 Subject: [PATCH 05/12] Refactor foreach command: reuse type system, extract planner, remove special cases - Extract ~400 lines of foreach planning from CalciteRelNodeVisitor into a dedicated ForeachPlanner class. - Replace string-based type plumbing (SqlTypeName names smuggled through ForeachBinding and reparsed at runtime) with RelDataType carried on the binding; foreach_pair_item now takes (pair, index) only and the call is built with the plan-time type directly. - Collapse ForeachBindingType {PAIR_ITEM, PAIR_ITER, PAIR_EXTRA, LAMBDA} into a single PAIR_SLOT; drop the unused LAMBDA variant. - Replace the AST arithmetic-scan heuristic for json_array element types with operand type inspection (json_array call) / plan-time JSON parse (string literal), using SqlTypeFamily. - Remove the reduce-arg0 special case in CalciteRexNodeVisitor: collection bindings are now staged on CalcitePlanContext and only activate inside cloned lambda contexts, so non-lambda reduce args resolve against the row naturally. - Foreach.Mode enum replaces stringly-typed mode comparisons; AstBuilder parses options/targets in a single pass without double-visiting expressions. - Delete dead code: FOREACH_TRANSFORM_* constants, FOREACH_PAIR_ITEM registry entry, validateHomogeneousJsonArrayArguments wrapper, duplicate collectionExpressionForMode branches. Signed-off-by: Songkan Tang --- .../org/opensearch/sql/ast/tree/Foreach.java | 24 +- .../sql/calcite/CalcitePlanContext.java | 46 +- .../sql/calcite/CalciteRelNodeVisitor.java | 424 +----------------- .../sql/calcite/CalciteRexNodeVisitor.java | 51 +-- .../sql/calcite/ForeachPlanner.java | 383 ++++++++++++++++ .../function/BuiltinFunctionName.java | 1 - .../ForeachPairItemFunctionImpl.java | 42 +- .../function/CollectionUDF/LambdaUtils.java | 11 +- .../expression/function/PPLFuncImpTable.java | 2 - .../opensearch/sql/ppl/parser/AstBuilder.java | 90 ++-- 10 files changed, 515 insertions(+), 559 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java b/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java index 486870e6efb..e11ecf541cd 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java @@ -7,6 +7,7 @@ import com.google.common.collect.ImmutableList; import java.util.List; +import java.util.Locale; import java.util.Map; import lombok.EqualsAndHashCode; import lombok.Getter; @@ -21,7 +22,7 @@ @EqualsAndHashCode(callSuper = false) @RequiredArgsConstructor public class Foreach extends UnresolvedPlan { - private final String mode; + private final Mode mode; private final Map options; private final List fieldPatterns; private final UnresolvedExpression collectionExpression; @@ -44,6 +45,27 @@ public T accept(AbstractNodeVisitor nodeVisitor, C context) { return nodeVisitor.visitForeach(this, context); } + /** Iteration mode of the foreach command. */ + public enum Mode { + MULTIFIELD, + MULTIVALUE, + JSON_ARRAY, + AUTO_COLLECTIONS; + + public static Mode of(String name) { + try { + return valueOf(name.toUpperCase(Locale.ROOT)); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("foreach mode [" + name + "] is not supported"); + } + } + + @Override + public String toString() { + return name().toLowerCase(Locale.ROOT); + } + } + @Getter @ToString @EqualsAndHashCode diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index e6c84eb9fb7..338d1ee6832 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -17,6 +17,7 @@ import java.util.function.BiFunction; import lombok.Getter; import lombok.Setter; +import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rex.RexCorrelVariable; import org.apache.calcite.rex.RexLambdaRef; import org.apache.calcite.rex.RexNode; @@ -66,8 +67,18 @@ public class CalcitePlanContext { private final Stack> windowPartitions = new Stack<>(); @Getter public Map rexLambdaRefMap; + + /** + * Foreach placeholder bindings active in this context, keyed by upper-cased placeholder name. + * Multifield mode activates bindings directly (placeholders resolve against the current row); + * collection modes stage bindings via {@link #stageForeachLambdaBindings} instead, so they only + * become active inside the lambda context cloned for the generated {@code reduce} call. + */ @Getter private Map foreachBindings = new HashMap<>(); + /** Bindings that become active in lambda contexts cloned from this one. */ + private Map stagedForeachLambdaBindings = new HashMap<>(); + /** * Maps AggregateFunction AST nodes to their output field index for HAVING/post-aggregate * resolution. @@ -115,7 +126,9 @@ private CalcitePlanContext(CalcitePlanContext parent) { this.rexLambdaRefMap = new HashMap<>(); // New map for lambda variables this.capturedVariables = new ArrayList<>(); // New list for captured variables this.inLambdaContext = true; // Mark that we're inside a lambda + // Active bindings carry over; staged bindings become active inside the lambda. this.foreachBindings = new HashMap<>(parent.foreachBindings); + this.foreachBindings.putAll(parent.stagedForeachLambdaBindings); } public RexNode resolveJoinCondition( @@ -186,30 +199,33 @@ public void pushForeachBindings(Map bindings) { foreachBindings = new HashMap<>(bindings); } - public void clearForeachBindings() { - foreachBindings.clear(); + public void stageForeachLambdaBindings(Map bindings) { + stagedForeachLambdaBindings = new HashMap<>(bindings); } - public enum ForeachBindingType { - FIELD, - LITERAL, - LAMBDA, - PAIR_ITEM, - PAIR_ITER, - PAIR_EXTRA + public void clearForeachBindings() { + foreachBindings.clear(); + stagedForeachLambdaBindings.clear(); } + /** + * A foreach placeholder binding. {@code FIELD} resolves to the named row field, {@code LITERAL} + * to a string literal, and {@code PAIR_SLOT} to slot {@code pairIndex} (typed {@code pairType}) + * of the named lambda pair variable. + */ public record ForeachBinding( - String value, - ForeachBindingType type, - int pairIndex, - String typeName, - @Nullable String componentTypeName) { + String value, ForeachBindingType type, int pairIndex, @Nullable RelDataType pairType) { public ForeachBinding(String value, ForeachBindingType type) { - this(value, type, -1, null, null); + this(value, type, -1, null); } } + public enum ForeachBindingType { + FIELD, + LITERAL, + PAIR_SLOT + } + public void putRexLambdaRefMap(Map candidateMap) { this.rexLambdaRefMap.putAll(candidateMap); } diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index 27bef3d827d..e29f127d8e0 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -37,16 +37,12 @@ import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; -import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; -import java.util.Locale; import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.regex.Matcher; -import java.util.regex.Pattern; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -100,11 +96,8 @@ import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta; import org.opensearch.sql.ast.expression.Argument; import org.opensearch.sql.ast.expression.Argument.ArgumentMap; -import org.opensearch.sql.ast.expression.DataType; import org.opensearch.sql.ast.expression.Field; -import org.opensearch.sql.ast.expression.ForeachPlaceholder; import org.opensearch.sql.ast.expression.Function; -import org.opensearch.sql.ast.expression.LambdaFunction; import org.opensearch.sql.ast.expression.Let; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.expression.ParseMethod; @@ -174,8 +167,6 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.ast.tree.Values; import org.opensearch.sql.ast.tree.Window; -import org.opensearch.sql.calcite.CalcitePlanContext.ForeachBinding; -import org.opensearch.sql.calcite.CalcitePlanContext.ForeachBindingType; import org.opensearch.sql.calcite.plan.AliasFieldsWrappable; import org.opensearch.sql.calcite.plan.HighlightPushDown; import org.opensearch.sql.calcite.plan.OpenSearchConstants; @@ -216,38 +207,18 @@ public class CalciteRelNodeVisitor extends AbstractNodeVisitor currentFields = context.relBuilder.peek().getRowType().getFieldNames(); - LinkedHashSet matchingFields = new LinkedHashSet<>(); - for (String pattern : node.getFieldPatterns()) { - matchingFields.addAll(WildcardUtils.expandWildcardPattern(pattern, currentFields)); - } - - for (String fieldName : matchingFields) { - Map bindings = - buildForeachMultifieldBindings(node.getFieldPatterns(), fieldName, node.getOptions()); - context.pushForeachBindings(bindings); - try { - for (Foreach.ForeachEvalClause clause : node.getEvalClauses()) { - RexNode expr = rexVisitor.analyze(clause.getExpression(), context); - String alias = substituteForeachTemplate(clause.getTargetTemplate(), bindings); - projectPlusOverriding( - List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); - } - } finally { - context.clearForeachBindings(); - } - } - return context.relBuilder.peek(); - } - - private RelNode visitForeachCollection( - Foreach node, CalcitePlanContext context, String normalizedMode) { - UnresolvedExpression sourceCollection = node.getCollectionExpression(); - if (sourceCollection == null && FOREACH_MODE_AUTO_COLLECTIONS.equals(normalizedMode)) { - sourceCollection = inferAutoCollectionExpression(context); - } - if (sourceCollection == null) { - throw new SemanticCheckException("foreach " + normalizedMode + " mode requires a field"); - } - if (node.getFieldPatterns().size() > 1) { - throw new SemanticCheckException( - "foreach " + normalizedMode + " mode accepts exactly one field"); - } - String itemName = - node.getOptions().getOrDefault(FOREACH_OPTION_ITEMSTR, FOREACH_PLACEHOLDER_ITEM); - String iterName = - node.getOptions().getOrDefault(FOREACH_OPTION_ITERSTR, FOREACH_PLACEHOLDER_ITER); - UnresolvedExpression collectionExpression = - collectionExpressionForMode( - sourceCollection, context, normalizedMode, node.getEvalClauses(), itemName); - RexNode collectionRex = rexVisitor.analyze(collectionExpression, context); - RelDataType itemType = ((ArraySqlType) collectionRex.getType()).getComponentType(); - List targetAliases = - node.getEvalClauses().stream().map(Foreach.ForeachEvalClause::getTargetTemplate).toList(); - List capturedFields = - collectForeachCapturedFields( - node.getEvalClauses(), context, itemName, iterName, targetAliases); - List pairArguments = new ArrayList<>(); - pairArguments.add(collectionExpression); - capturedFields.stream() - .map(field -> new Field(new QualifiedName(field.name()))) - .forEach(pairArguments::add); - Function pairedCollection = new Function(FOREACH_PAIR_COLLECTION_FUNCTION, pairArguments); - Map bindings = new HashMap<>(); - bindings.put( - itemName.toUpperCase(Locale.ROOT), - newForeachPairBinding(FOREACH_PAIR, ForeachBindingType.PAIR_ITEM, 0, itemType)); - bindings.put( - FOREACH_PLACEHOLDER_ITEM, - newForeachPairBinding(FOREACH_PAIR, ForeachBindingType.PAIR_ITEM, 0, itemType)); - bindings.put( - iterName.toUpperCase(Locale.ROOT), - newForeachPairBinding( - FOREACH_PAIR, - ForeachBindingType.PAIR_ITER, - 1, - context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER))); - bindings.put( - FOREACH_PLACEHOLDER_ITER, - newForeachPairBinding( - FOREACH_PAIR, - ForeachBindingType.PAIR_ITER, - 1, - context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER))); - for (int i = 0; i < capturedFields.size(); i++) { - ForeachCapturedField field = capturedFields.get(i); - bindings.put( - field.name().toUpperCase(Locale.ROOT), - newForeachPairBinding(FOREACH_PAIR, ForeachBindingType.PAIR_EXTRA, i + 2, field.type())); - } - - context.pushForeachBindings(bindings); - try { - for (Foreach.ForeachEvalClause clause : node.getEvalClauses()) { - String alias = substituteForeachTemplate(clause.getTargetTemplate(), bindings); - Function reduce = - new Function( - BuiltinFunctionName.REDUCE.getName().getFunctionName(), - List.of( - pairedCollection, - new Field(new QualifiedName(alias)), - new LambdaFunction( - clause.getExpression(), - List.of(new QualifiedName(alias), new QualifiedName(FOREACH_PAIR))))); - RexNode expr = rexVisitor.analyze(reduce, context); - projectPlusOverriding( - List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); - } - } finally { - context.clearForeachBindings(); - } - return context.relBuilder.peek(); - } - - private List collectForeachCapturedFields( - List evalClauses, - CalcitePlanContext context, - String itemName, - String iterName, - List targetAliases) { - Set excluded = - Stream.concat( - Stream.of( - itemName, - iterName, - FOREACH_PLACEHOLDER_ITEM, - FOREACH_PLACEHOLDER_ITER, - FOREACH_PAIR), - targetAliases.stream()) - .map(name -> name.toUpperCase(Locale.ROOT)) - .collect(Collectors.toSet()); - Map rowTypes = - context.relBuilder.peek().getRowType().getFieldList().stream() - .collect( - Collectors.toMap( - field -> field.getName().toUpperCase(Locale.ROOT), - RelDataTypeField::getType, - (left, right) -> left, - LinkedHashMap::new)); - Map fields = new LinkedHashMap<>(); - evalClauses.forEach( - clause -> collectForeachCapturedFields(clause.getExpression(), excluded, rowTypes, fields)); - return new ArrayList<>(fields.values()); - } - - private void collectForeachCapturedFields( - Node node, - Set excluded, - Map rowTypes, - Map fields) { - if (node instanceof QualifiedName qualifiedName) { - String key = qualifiedName.toString().toUpperCase(Locale.ROOT); - RelDataType type = rowTypes.get(key); - if (type != null && !excluded.contains(key)) { - fields.putIfAbsent(key, new ForeachCapturedField(qualifiedName.toString(), type)); - } - } - node.getChild() - .forEach(child -> collectForeachCapturedFields(child, excluded, rowTypes, fields)); - } - - private ForeachBinding newForeachPairBinding( - String pairName, ForeachBindingType type, int index, RelDataType relDataType) { - return new ForeachBinding( - pairName, type, index, relDataType.getSqlTypeName().name(), componentTypeName(relDataType)); - } - - @Nullable - private String componentTypeName(RelDataType relDataType) { - if (relDataType instanceof ArraySqlType arrayType) { - return arrayType.getComponentType().getSqlTypeName().name(); - } - return null; - } - - private record ForeachCapturedField(String name, RelDataType type) {} - - private UnresolvedExpression collectionExpressionForMode( - UnresolvedExpression collectionExpression, - CalcitePlanContext context, - String normalizedMode, - List evalClauses, - String itemName) { - if (FOREACH_MODE_MULTIVALUE.equals(normalizedMode)) { - return collectionExpression; - } - if (FOREACH_MODE_AUTO_COLLECTIONS.equals(normalizedMode)) { - RexNode rexNode = rexVisitor.analyze(collectionExpression, context); - if (rexNode.getType() instanceof ArraySqlType) { - return collectionExpression; - } - } - if (isJsonArrayFunction(collectionExpression)) { - Function jsonArray = (Function) collectionExpression; - validateHomogeneousJsonArrayArguments(jsonArray); - return new Function( - FOREACH_JSON_ARRAY_FUNCTION, - List.of( - collectionExpression, - new Literal( - inferForeachJsonArrayType(collectionExpression, evalClauses, itemName), - DataType.STRING))); - } - return new Function( - FOREACH_JSON_ARRAY_FUNCTION, - List.of( - collectionExpression, - new Literal( - inferForeachJsonArrayType(collectionExpression, evalClauses, itemName), - DataType.STRING))); - } - - private UnresolvedExpression inferAutoCollectionExpression(CalcitePlanContext context) { - return context.relBuilder.peek().getRowType().getFieldList().stream() - .filter(field -> field.getType() instanceof ArraySqlType) - .findFirst() - .map(field -> (UnresolvedExpression) new Field(new QualifiedName(field.getName()))) - .orElseThrow( - () -> - new SemanticCheckException( - "foreach auto_collections mode requires a multivalue field or JSON array")); - } - - private boolean isJsonArrayFunction(UnresolvedExpression expression) { - return expression instanceof Function function - && BuiltinFunctionName.JSON_ARRAY - .getName() - .getFunctionName() - .equalsIgnoreCase(function.getFuncName()); - } - - private String inferForeachJsonArrayType( - UnresolvedExpression collectionExpression, - List evalClauses, - String itemName) { - if (isJsonArrayFunction(collectionExpression)) { - return jsonArrayTypeFromArguments((Function) collectionExpression) - .orElse(FOREACH_JSON_ARRAY_STRING_TYPE); - } - if (foreachItemUsedInArithmetic(evalClauses, itemName)) { - return FOREACH_JSON_ARRAY_NUMERIC_TYPE; - } - return FOREACH_JSON_ARRAY_STRING_TYPE; - } - - private Optional jsonArrayTypeFromArguments(Function jsonArray) { - boolean hasNumeric = false; - boolean hasString = false; - for (UnresolvedExpression argument : jsonArray.getFuncArgs()) { - if (argument instanceof Literal literal) { - switch (literal.getType()) { - case INTEGER, LONG, SHORT, FLOAT, DOUBLE, DECIMAL: - hasNumeric = true; - break; - case STRING: - hasString = true; - break; - case NULL: - break; - default: - return Optional.empty(); - } - } else { - return Optional.empty(); - } - } - if (hasNumeric && hasString) { - throw new SemanticCheckException( - "foreach json_array elements must be consistently strings or numbers"); - } - if (hasNumeric) { - return Optional.of(FOREACH_JSON_ARRAY_NUMERIC_TYPE); - } - if (hasString) { - return Optional.of(FOREACH_JSON_ARRAY_STRING_TYPE); - } - return Optional.empty(); - } - - private void validateHomogeneousJsonArrayArguments(Function jsonArray) { - jsonArrayTypeFromArguments(jsonArray); - } - - private boolean foreachItemUsedInArithmetic( - List evalClauses, String itemName) { - return evalClauses.stream() - .anyMatch(clause -> foreachItemUsedInArithmetic(clause.getExpression(), itemName, false)); - } - - private boolean foreachItemUsedInArithmetic(Node node, String itemName, boolean inArithmetic) { - if (node instanceof UnresolvedExpression expression - && isForeachItemReference(expression, itemName)) { - return inArithmetic; - } - boolean arithmetic = - inArithmetic - || (node instanceof UnresolvedExpression expression - && isArithmeticExpression(expression)); - return node.getChild().stream() - .anyMatch(child -> foreachItemUsedInArithmetic(child, itemName, arithmetic)); - } - - private boolean isForeachItemReference(UnresolvedExpression expression, String itemName) { - if (expression instanceof ForeachPlaceholder placeholder) { - return FOREACH_PLACEHOLDER_ITEM.equalsIgnoreCase(placeholder.getName()) - || itemName.equalsIgnoreCase(placeholder.getName()); - } - return expression instanceof QualifiedName qualifiedName - && (FOREACH_PLACEHOLDER_ITEM.equalsIgnoreCase(qualifiedName.toString()) - || itemName.equalsIgnoreCase(qualifiedName.toString())); - } - - private boolean isArithmeticExpression(UnresolvedExpression expression) { - return expression instanceof Function function - && Set.of("+", "-", "*", "/", "%").contains(function.getFuncName()); - } - - private Map buildForeachMultifieldBindings( - List patterns, String fieldName, Map options) { - Map bindings = new HashMap<>(); - bindings.put( - FOREACH_PLACEHOLDER_FIELD, new ForeachBinding(fieldName, ForeachBindingType.FIELD)); - bindings.put( - options - .getOrDefault(FOREACH_OPTION_FIELDSTR, FOREACH_PLACEHOLDER_FIELD) - .toUpperCase(Locale.ROOT), - new ForeachBinding(fieldName, ForeachBindingType.FIELD)); - for (String pattern : patterns) { - List segments = wildcardSegments(pattern, fieldName); - if (segments != null) { - String matchStr = String.join("", segments); - bindings.put( - FOREACH_PLACEHOLDER_MATCHSTR, new ForeachBinding(matchStr, ForeachBindingType.LITERAL)); - bindings.put( - options - .getOrDefault(FOREACH_OPTION_MATCHSTR, FOREACH_PLACEHOLDER_MATCHSTR) - .toUpperCase(Locale.ROOT), - new ForeachBinding(matchStr, ForeachBindingType.LITERAL)); - for (int i = 0; i < segments.size(); i++) { - String key = FOREACH_PLACEHOLDER_MATCHSEG + (i + 1); - bindings.put(key, new ForeachBinding(segments.get(i), ForeachBindingType.LITERAL)); - bindings.put( - options.getOrDefault(FOREACH_OPTION_MATCHSEG + (i + 1), key).toUpperCase(Locale.ROOT), - new ForeachBinding(segments.get(i), ForeachBindingType.LITERAL)); - } - break; - } - } - return bindings; - } - - private List wildcardSegments(String pattern, String fieldName) { - if (!WildcardUtils.matchesWildcardPattern(pattern, fieldName)) { - return null; - } - if (!WildcardUtils.containsWildcard(pattern)) { - return List.of(); - } - Matcher matcher = - Pattern.compile(WildcardUtils.convertWildcardPatternToRegex(pattern)).matcher(fieldName); - if (!matcher.matches()) { - return null; - } - List segments = new ArrayList<>(); - for (int i = 1; i <= matcher.groupCount(); i++) { - segments.add(matcher.group(i)); - } - return segments; - } - - private String substituteForeachTemplate(String template, Map bindings) { - String result = template; - for (Map.Entry entry : bindings.entrySet()) { - result = - result.replaceAll( - "(?i)<<" + java.util.regex.Pattern.quote(entry.getKey()) + ">>", - java.util.regex.Matcher.quoteReplacement(entry.getValue().value())); - } - return result; + return foreachPlanner.plan(node, context); } @Override @@ -1795,7 +1379,7 @@ private RelNode buildConversionProjection(ConversionState state, CalcitePlanCont return context.relBuilder.peek(); } - private void projectPlusOverriding( + void projectPlusOverriding( List newFields, List newNames, CalcitePlanContext context) { Set originalFieldNameSet = new HashSet<>(context.relBuilder.peek().getRowType().getFieldNames()); diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java index f7c3aaf76d7..ee8d72a7b96 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -94,6 +94,7 @@ import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.expression.function.BuiltinFunctionName; import org.opensearch.sql.expression.function.CoercionUtils; +import org.opensearch.sql.expression.function.PPLBuiltinOperators; import org.opensearch.sql.expression.function.PPLFuncImpTable; @RequiredArgsConstructor @@ -426,41 +427,26 @@ private RexNode foreachBindingToRexNode( switch (binding.type()) { case FIELD: return context.relBuilder.field(binding.value()); - case LAMBDA: - RexLambdaRef lambdaRef = context.getRexLambdaRefMap().get(binding.value()); - if (lambdaRef == null) { + case PAIR_SLOT: + RexLambdaRef pair = context.getRexLambdaRefMap().get(binding.value()); + if (pair == null) { throw new SemanticCheckException("Unresolved foreach lambda placeholder " + name); } - return lambdaRef; - case PAIR_ITEM: - case PAIR_EXTRA: - return foreachPairItem(binding, name, context); - case PAIR_ITER: - return foreachPairItem(binding, name, context); + // The slot's type is known at plan time but the extraction UDF infers ANY, so build the + // call with the type set explicitly. LambdaUtils' re-inference preserves it (a CAST would + // not survive the enumerable backend for complex types like arrays). + return context.rexBuilder.makeCall( + binding.pairType(), + PPLBuiltinOperators.FOREACH_PAIR_ITEM, + List.of( + pair, + context.rexBuilder.makeExactLiteral(BigDecimal.valueOf(binding.pairIndex())))); case LITERAL: default: return context.rexBuilder.makeLiteral(binding.value()); } } - private RexNode foreachPairItem( - ForeachBinding binding, String placeholderName, CalcitePlanContext context) { - RexLambdaRef pair = context.getRexLambdaRefMap().get(binding.value()); - if (pair == null) { - throw new SemanticCheckException("Unresolved foreach lambda placeholder " + placeholderName); - } - return PPLFuncImpTable.INSTANCE.resolve( - context.rexBuilder, - BuiltinFunctionName.FOREACH_PAIR_ITEM, - pair, - context.rexBuilder.makeExactLiteral(BigDecimal.valueOf(binding.pairIndex())), - context.rexBuilder.makeLiteral(binding.typeName()), - binding.componentTypeName() == null - ? context.rexBuilder.makeNullLiteral( - context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.NULL)) - : context.rexBuilder.makeLiteral(binding.componentTypeName())); - } - @Override public RexNode visitAlias(Alias node, CalcitePlanContext context) { RexNode expr = analyze(node.getDelegated(), context); @@ -637,8 +623,7 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) { List capturedVars = null; try { - for (int argIndex = 0; argIndex < args.size(); argIndex++) { - UnresolvedExpression arg = args.get(argIndex); + for (UnresolvedExpression arg : args) { if (arg instanceof LambdaFunction) { CalcitePlanContext lambdaContext = prepareLambdaContext( @@ -657,14 +642,6 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) { arguments.add(lambdaNode); // Capture any external variables that were referenced in the lambda capturedVars = lambdaContext.getCapturedVariables(); - } else if (node.getFuncName().equalsIgnoreCase("reduce") && argIndex == 0) { - Map foreachBindings = new HashMap<>(context.getForeachBindings()); - context.clearForeachBindings(); - try { - arguments.add(analyze(arg, context)); - } finally { - context.pushForeachBindings(foreachBindings); - } } else { arguments.add(analyze(arg, context)); } diff --git a/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java b/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java new file mode 100644 index 00000000000..e2e7475caf2 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java @@ -0,0 +1,383 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite; + +import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.gson; + +import com.google.gson.JsonSyntaxException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.type.ArraySqlType; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.sql.ast.Node; +import org.opensearch.sql.ast.expression.DataType; +import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.Function; +import org.opensearch.sql.ast.expression.LambdaFunction; +import org.opensearch.sql.ast.expression.Literal; +import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.UnresolvedExpression; +import org.opensearch.sql.ast.tree.Foreach; +import org.opensearch.sql.ast.tree.Foreach.ForeachEvalClause; +import org.opensearch.sql.calcite.CalcitePlanContext.ForeachBinding; +import org.opensearch.sql.calcite.CalcitePlanContext.ForeachBindingType; +import org.opensearch.sql.calcite.utils.WildcardUtils; +import org.opensearch.sql.exception.SemanticCheckException; +import org.opensearch.sql.expression.function.BuiltinFunctionName; + +/** + * Plans the PPL {@code foreach} command. + * + *

Multifield mode expands each eval clause once per matching field, binding {@code <>} + * (and match placeholders) so the clause expression resolves against the current field. + * + *

Collection modes (multivalue / json_array / auto_collections) rewrite each eval clause into a + * {@code reduce} over the collection: elements are packed into internal pairs {@code [item, iter, + * captured-field...]} so the lambda can reference the loop item, the loop index, and any row fields + * the clause mentions. Placeholder bindings are staged on the context and only become active inside + * the lambda; the reduce call's other arguments resolve against the row as usual. + */ +class ForeachPlanner { + + private static final String OPTION_FIELDSTR = "fieldstr"; + private static final String OPTION_MATCHSTR = "matchstr"; + private static final String OPTION_MATCHSEG = "matchseg"; + private static final String OPTION_ITEMSTR = "itemstr"; + private static final String OPTION_ITERSTR = "iterstr"; + private static final String PLACEHOLDER_FIELD = "FIELD"; + private static final String PLACEHOLDER_MATCHSTR = "MATCHSTR"; + private static final String PLACEHOLDER_MATCHSEG = "MATCHSEG"; + private static final String PLACEHOLDER_ITEM = "ITEM"; + private static final String PLACEHOLDER_ITER = "ITER"; + + /** Name of the lambda variable holding the internal pair inside generated reduce calls. */ + private static final String PAIR_VAR = "__foreach_pair"; + + private final CalciteRelNodeVisitor relVisitor; + private final CalciteRexNodeVisitor rexVisitor; + + ForeachPlanner(CalciteRelNodeVisitor relVisitor, CalciteRexNodeVisitor rexVisitor) { + this.relVisitor = relVisitor; + this.rexVisitor = rexVisitor; + } + + RelNode plan(Foreach node, CalcitePlanContext context) { + return node.getMode() == Foreach.Mode.MULTIFIELD + ? planMultifield(node, context) + : planCollection(node, context); + } + + // ---------------------------------------------------------------------- multifield + + private RelNode planMultifield(Foreach node, CalcitePlanContext context) { + List currentFields = context.relBuilder.peek().getRowType().getFieldNames(); + Set matchingFields = new LinkedHashSet<>(); + for (String pattern : node.getFieldPatterns()) { + matchingFields.addAll(WildcardUtils.expandWildcardPattern(pattern, currentFields)); + } + + for (String fieldName : matchingFields) { + Map bindings = + multifieldBindings(node.getFieldPatterns(), fieldName, node.getOptions()); + context.pushForeachBindings(bindings); + try { + for (ForeachEvalClause clause : node.getEvalClauses()) { + RexNode expr = rexVisitor.analyze(clause.getExpression(), context); + String alias = substituteTemplate(clause.getTargetTemplate(), bindings); + relVisitor.projectPlusOverriding( + List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); + } + } finally { + context.clearForeachBindings(); + } + } + return context.relBuilder.peek(); + } + + private Map multifieldBindings( + List patterns, String fieldName, Map options) { + Map bindings = new LinkedHashMap<>(); + bind( + bindings, + new ForeachBinding(fieldName, ForeachBindingType.FIELD), + PLACEHOLDER_FIELD, + options.getOrDefault(OPTION_FIELDSTR, PLACEHOLDER_FIELD)); + for (String pattern : patterns) { + List segments = wildcardSegments(pattern, fieldName); + if (segments == null) { + continue; + } + bind( + bindings, + new ForeachBinding(String.join("", segments), ForeachBindingType.LITERAL), + PLACEHOLDER_MATCHSTR, + options.getOrDefault(OPTION_MATCHSTR, PLACEHOLDER_MATCHSTR)); + for (int i = 0; i < segments.size(); i++) { + String defaultName = PLACEHOLDER_MATCHSEG + (i + 1); + bind( + bindings, + new ForeachBinding(segments.get(i), ForeachBindingType.LITERAL), + defaultName, + options.getOrDefault(OPTION_MATCHSEG + (i + 1), defaultName)); + } + break; + } + return bindings; + } + + /** + * Returns the substrings of {@code fieldName} captured by the wildcards in {@code pattern}, an + * empty list if the pattern matches without wildcards, or null if it does not match. + */ + private List wildcardSegments(String pattern, String fieldName) { + if (!WildcardUtils.matchesWildcardPattern(pattern, fieldName)) { + return null; + } + if (!WildcardUtils.containsWildcard(pattern)) { + return List.of(); + } + Matcher matcher = + Pattern.compile(WildcardUtils.convertWildcardPatternToRegex(pattern)).matcher(fieldName); + if (!matcher.matches()) { + return null; + } + List segments = new ArrayList<>(); + for (int i = 1; i <= matcher.groupCount(); i++) { + segments.add(matcher.group(i)); + } + return segments; + } + + // ---------------------------------------------------------------------- collection modes + + private RelNode planCollection(Foreach node, CalcitePlanContext context) { + Foreach.Mode mode = node.getMode(); + UnresolvedExpression collection = node.getCollectionExpression(); + if (collection == null && mode == Foreach.Mode.AUTO_COLLECTIONS) { + collection = firstArrayField(context); + } + if (collection == null) { + throw new SemanticCheckException("foreach " + mode + " mode requires a field"); + } + collection = asArrayExpression(collection, mode, context); + RexNode collectionRex = rexVisitor.analyze(collection, context); + if (!(collectionRex.getType() instanceof ArraySqlType arrayType)) { + throw new SemanticCheckException("foreach " + mode + " mode requires a multivalue field"); + } + RelDataType itemType = arrayType.getComponentType(); + RelDataType iterType = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER); + + String itemName = node.getOptions().getOrDefault(OPTION_ITEMSTR, PLACEHOLDER_ITEM); + String iterName = node.getOptions().getOrDefault(OPTION_ITERSTR, PLACEHOLDER_ITER); + List targetAliases = + node.getEvalClauses().stream().map(ForeachEvalClause::getTargetTemplate).toList(); + List capturedFields = + capturedFields(node.getEvalClauses(), context, itemName, iterName, targetAliases); + + // Pack [item, iter, captured...] so the reduce lambda can reach everything through one var. + List pairArgs = new ArrayList<>(); + pairArgs.add(collection); + capturedFields.stream() + .map(field -> new Field(new QualifiedName(field.getName()))) + .forEach(pairArgs::add); + Function pairedCollection = + new Function( + BuiltinFunctionName.FOREACH_PAIR_COLLECTION.getName().getFunctionName(), pairArgs); + + Map bindings = new LinkedHashMap<>(); + bindPairSlot(bindings, 0, itemType, itemName, PLACEHOLDER_ITEM); + bindPairSlot(bindings, 1, iterType, iterName, PLACEHOLDER_ITER); + for (int i = 0; i < capturedFields.size(); i++) { + RelDataTypeField field = capturedFields.get(i); + bindPairSlot(bindings, i + 2, field.getType(), field.getName()); + } + + // Stage rather than activate: the pair collection and the accumulator's initial value are + // resolved against the row; only the lambda body sees the placeholder bindings. + context.stageForeachLambdaBindings(bindings); + try { + for (ForeachEvalClause clause : node.getEvalClauses()) { + String alias = substituteTemplate(clause.getTargetTemplate(), bindings); + Function reduce = + new Function( + BuiltinFunctionName.REDUCE.getName().getFunctionName(), + List.of( + pairedCollection, + new Field(new QualifiedName(alias)), + new LambdaFunction( + clause.getExpression(), + List.of(new QualifiedName(alias), new QualifiedName(PAIR_VAR))))); + RexNode expr = rexVisitor.analyze(reduce, context); + relVisitor.projectPlusOverriding( + List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); + } + } finally { + context.clearForeachBindings(); + } + return context.relBuilder.peek(); + } + + private void bindPairSlot( + Map bindings, int slot, RelDataType type, String... names) { + ForeachBinding binding = new ForeachBinding(PAIR_VAR, ForeachBindingType.PAIR_SLOT, slot, type); + for (String name : names) { + bindings.put(name.toUpperCase(Locale.ROOT), binding); + } + } + + private UnresolvedExpression firstArrayField(CalcitePlanContext context) { + return context.relBuilder.peek().getRowType().getFieldList().stream() + .filter(field -> field.getType() instanceof ArraySqlType) + .findFirst() + .map(field -> (UnresolvedExpression) new Field(new QualifiedName(field.getName()))) + .orElseThrow( + () -> + new SemanticCheckException( + "foreach auto_collections mode requires a multivalue field or JSON array")); + } + + /** + * Ensures the collection expression evaluates to a Calcite array. Non-array expressions (JSON + * array strings or {@code json_array()} calls) are wrapped in {@code foreach_json_array}, which + * parses the JSON and casts every element to one inferred type. + */ + private UnresolvedExpression asArrayExpression( + UnresolvedExpression collection, Foreach.Mode mode, CalcitePlanContext context) { + if (mode == Foreach.Mode.MULTIVALUE + || (mode == Foreach.Mode.AUTO_COLLECTIONS + && rexVisitor.analyze(collection, context).getType() instanceof ArraySqlType)) { + return collection; + } + return new Function( + BuiltinFunctionName.FOREACH_JSON_ARRAY.getName().getFunctionName(), + List.of( + collection, new Literal(jsonElementType(collection, context).name(), DataType.STRING))); + } + + /** + * Infers the element type a JSON array collection should be read as. JSON only distinguishes + * numbers and strings here, so the answer is DOUBLE (gson parses JSON numbers as doubles) or + * VARCHAR. Mixed content is rejected; expressions whose content is unknowable at plan time (e.g. + * a field holding a JSON string) default to VARCHAR. + */ + private SqlTypeName jsonElementType(UnresolvedExpression collection, CalcitePlanContext context) { + if (collection instanceof Function function + && BuiltinFunctionName.JSON_ARRAY + .getName() + .getFunctionName() + .equalsIgnoreCase(function.getFuncName())) { + return elementTypeOf( + function.getFuncArgs().stream() + .map(arg -> rexVisitor.analyze(arg, context).getType()) + .map(RelDataType::getFamily) + .collect(Collectors.toSet())); + } + if (collection instanceof Literal literal && literal.getType() == DataType.STRING) { + try { + List values = gson.fromJson(String.valueOf(literal.getValue()), List.class); + if (values != null) { + return elementTypeOf( + values.stream() + .filter(Objects::nonNull) + .map(v -> v instanceof Number ? SqlTypeFamily.NUMERIC : SqlTypeFamily.CHARACTER) + .collect(Collectors.toSet())); + } + } catch (JsonSyntaxException ignored) { + // Malformed JSON literals return null at runtime; type choice is irrelevant. + } + } + return SqlTypeName.VARCHAR; + } + + private SqlTypeName elementTypeOf(Set families) { + boolean numeric = families.contains(SqlTypeFamily.NUMERIC); + boolean character = families.contains(SqlTypeFamily.CHARACTER); + if (numeric && character) { + throw new SemanticCheckException( + "foreach json_array elements must be consistently strings or numbers"); + } + return numeric ? SqlTypeName.DOUBLE : SqlTypeName.VARCHAR; + } + + /** Row fields referenced by the eval clauses that must ride along inside the pairs. */ + private List capturedFields( + List evalClauses, + CalcitePlanContext context, + String itemName, + String iterName, + List targetAliases) { + Set excluded = + Stream.concat( + Stream.of(itemName, iterName, PLACEHOLDER_ITEM, PLACEHOLDER_ITER, PAIR_VAR), + targetAliases.stream()) + .map(name -> name.toUpperCase(Locale.ROOT)) + .collect(Collectors.toSet()); + Map rowFields = + context.relBuilder.peek().getRowType().getFieldList().stream() + .collect( + Collectors.toMap( + field -> field.getName().toUpperCase(Locale.ROOT), + field -> field, + (left, right) -> left, + LinkedHashMap::new)); + Map captured = new LinkedHashMap<>(); + evalClauses.forEach( + clause -> collectFieldReferences(clause.getExpression(), excluded, rowFields, captured)); + return new ArrayList<>(captured.values()); + } + + private void collectFieldReferences( + Node node, + Set excluded, + Map rowFields, + Map captured) { + if (node instanceof QualifiedName qualifiedName) { + String key = qualifiedName.toString().toUpperCase(Locale.ROOT); + RelDataTypeField field = rowFields.get(key); + if (field != null && !excluded.contains(key)) { + captured.putIfAbsent(key, field); + } + } + node.getChild().forEach(child -> collectFieldReferences(child, excluded, rowFields, captured)); + } + + // ---------------------------------------------------------------------- shared + + private void bind( + Map bindings, + ForeachBinding binding, + String defaultName, + String customName) { + bindings.put(defaultName.toUpperCase(Locale.ROOT), binding); + bindings.put(customName.toUpperCase(Locale.ROOT), binding); + } + + private String substituteTemplate(String template, Map bindings) { + String result = template; + for (Map.Entry entry : bindings.entrySet()) { + result = + result.replaceAll( + "(?i)<<" + Pattern.quote(entry.getKey()) + ">>", + Matcher.quoteReplacement(entry.getValue().value())); + } + return result; + } +} diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java index e1a4abbf89a..1c464043dbf 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java @@ -71,7 +71,6 @@ public enum BuiltinFunctionName { ARRAY(FunctionName.of("array")), FOREACH_JSON_ARRAY(FunctionName.of("foreach_json_array"), true), FOREACH_PAIR_COLLECTION(FunctionName.of("foreach_pair_collection"), true), - FOREACH_PAIR_ITEM(FunctionName.of("foreach_pair_item"), true), ARRAY_LENGTH(FunctionName.of("array_length")), ARRAY_SLICE(FunctionName.of("array_slice"), true), ARRAY_COMPACT(FunctionName.of("array_compact")), diff --git a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java index b93eefda881..7d2380b646e 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java @@ -5,6 +5,7 @@ package org.opensearch.sql.expression.function.CollectionUDF; +import java.util.Arrays; import java.util.List; import org.apache.calcite.adapter.enumerable.NotNullImplementor; import org.apache.calcite.adapter.enumerable.NullPolicy; @@ -12,15 +13,16 @@ import org.apache.calcite.linq4j.tree.Expression; import org.apache.calcite.linq4j.tree.Expressions; import org.apache.calcite.linq4j.tree.Types; -import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rex.RexCall; import org.apache.calcite.sql.type.SqlReturnTypeInference; import org.apache.calcite.sql.type.SqlTypeName; -import org.apache.calcite.sql.type.SqlTypeUtil; import org.opensearch.sql.expression.function.ImplementorUDF; import org.opensearch.sql.expression.function.UDFOperandMetadata; -/** Extracts an internal foreach pair slot while preserving the slot's inferred Calcite type. */ +/** + * Extracts one slot of an internal foreach pair: {@code foreach_pair_item(pair, index)}. Returns + * ANY; the foreach planner wraps every call in a CAST to the slot's plan-time type. + */ public class ForeachPairItemFunctionImpl extends ImplementorUDF { public ForeachPairItemFunctionImpl() { super(new ForeachPairItemImplementor(), NullPolicy.NONE); @@ -28,25 +30,11 @@ public ForeachPairItemFunctionImpl() { @Override public SqlReturnTypeInference getReturnTypeInference() { - return opBinding -> { - SqlTypeName typeName = SqlTypeName.valueOf(opBinding.getOperandLiteralValue(2, String.class)); - RelDataType type; - if (typeName == SqlTypeName.ARRAY) { - SqlTypeName componentTypeName = - SqlTypeName.valueOf(opBinding.getOperandLiteralValue(3, String.class)); - type = - SqlTypeUtil.createArrayType( - opBinding.getTypeFactory(), - opBinding - .getTypeFactory() - .createTypeWithNullability( - opBinding.getTypeFactory().createSqlType(componentTypeName), true), - true); - } else { - type = opBinding.getTypeFactory().createSqlType(typeName); - } - return opBinding.getTypeFactory().createTypeWithNullability(type, true); - }; + return opBinding -> + opBinding + .getTypeFactory() + .createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(SqlTypeName.ANY), true); } @Override @@ -69,13 +57,7 @@ public static Object eval(Object... args) { return null; } int index = ((Number) args[1]).intValue(); - if (args[0] instanceof Object[] pair) { - return index < 0 || index >= pair.length ? null : pair[index]; - } - List pair = (List) args[0]; - if (index < 0 || index >= pair.size()) { - return null; - } - return pair.get(index); + List pair = args[0] instanceof Object[] array ? Arrays.asList(array) : (List) args[0]; + return index < 0 || index >= pair.size() ? null : pair.get(index); } } diff --git a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java index 57bace94116..a2684cedcf8 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java @@ -69,8 +69,9 @@ public static RelDataType inferReturnTypeFromLambda( public static RexCall reInferReturnTypeForRexCallInsideLambda( RexCall rexCall, Map argTypes, RelDataTypeFactory typeFactory) { - if (rexCall.getKind() == SqlKind.CAST - || "CAST".equalsIgnoreCase(rexCall.getOperator().getName())) { + // A CAST's target type lives on the call itself and cannot be re-derived from its operands; + // re-inferring it trips SqlCastFunction's assertion. Keep the call as-is. + if (rexCall.getKind() == SqlKind.CAST) { return rexCall; } List filledOperands = new ArrayList<>(); @@ -94,6 +95,12 @@ public static RexCall reInferReturnTypeForRexCallInsideLambda( .getOperator() .inferReturnType( new RexCallBinding(typeFactory, rexCall.getOperator(), filledOperands, List.of())); + // A call whose operator can only infer ANY may carry a more precise type assigned when the + // call was built (e.g. foreach pair slot extraction); re-inference must not erase it. + if (returnType.getSqlTypeName() == SqlTypeName.ANY + && rexCall.getType().getSqlTypeName() != SqlTypeName.ANY) { + returnType = rexCall.getType(); + } return rexCall.clone(returnType, filledOperands); } } diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java index 76080b5abeb..b087b972570 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java @@ -77,7 +77,6 @@ import static org.opensearch.sql.expression.function.BuiltinFunctionName.FORALL; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_JSON_ARRAY; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_PAIR_COLLECTION; -import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_PAIR_ITEM; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FROM_DAYS; import static org.opensearch.sql.expression.function.BuiltinFunctionName.FROM_UNIXTIME; import static org.opensearch.sql.expression.function.BuiltinFunctionName.GET_FORMAT; @@ -1092,7 +1091,6 @@ void populate() { registerOperator(REDUCE, PPLBuiltinOperators.REDUCE); registerOperator(FOREACH_JSON_ARRAY, PPLBuiltinOperators.FOREACH_JSON_ARRAY); registerOperator(FOREACH_PAIR_COLLECTION, PPLBuiltinOperators.FOREACH_PAIR_COLLECTION); - registerOperator(FOREACH_PAIR_ITEM, PPLBuiltinOperators.FOREACH_PAIR_ITEM); // Register Json function register( diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java index 8d326623fea..4870bb88c96 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java @@ -837,28 +837,34 @@ public UnresolvedPlan visitEvalCommand(EvalCommandContext ctx) { @Override public UnresolvedPlan visitForeachCommand(OpenSearchPPLParser.ForeachCommandContext ctx) { java.util.Map options = new LinkedHashMap<>(); - List optionContexts = - ctx.foreachArgument().stream() - .map(OpenSearchPPLParser.ForeachArgumentContext::foreachOption) - .filter(java.util.Objects::nonNull) - .collect(Collectors.toList()); - List targetContexts = - ctx.foreachArgument().stream() - .map(OpenSearchPPLParser.ForeachArgumentContext::foreachTarget) - .filter(java.util.Objects::nonNull) - .collect(Collectors.toList()); - optionContexts.forEach( - option -> - options.put( - option.ident(0).getText().toLowerCase(Locale.ROOT), - normalizeForeachOptionValue(option.getChild(2).getText()))); - String mode = options.getOrDefault("mode", inferForeachMode(targetContexts)); - List patterns = - targetContexts.stream().map(this::getTextInQuery).collect(Collectors.toList()); - UnresolvedExpression collectionExpression = - "multifield".equalsIgnoreCase(mode) - ? null - : buildForeachCollectionExpression(targetContexts, mode); + List targets = new ArrayList<>(); + List patterns = new ArrayList<>(); + for (OpenSearchPPLParser.ForeachArgumentContext argument : ctx.foreachArgument()) { + OpenSearchPPLParser.ForeachOptionContext option = argument.foreachOption(); + if (option != null) { + options.put( + option.ident(0).getText().toLowerCase(Locale.ROOT), + stripForeachPlaceholderMarkers(option.getChild(2).getText())); + } else { + OpenSearchPPLParser.ForeachTargetContext target = argument.foreachTarget(); + patterns.add(getTextInQuery(target)); + targets.add( + target.logicalExpression() != null + ? expressionBuilder.visit(target.logicalExpression()) + : new Field(new QualifiedName(getTextInQuery(target)))); + } + } + Foreach.Mode mode = + options.containsKey("mode") + ? Foreach.Mode.of(options.get("mode")) + : inferForeachMode(targets); + UnresolvedExpression collectionExpression = null; + if (mode != Foreach.Mode.MULTIFIELD) { + if (targets.size() > 1 || (targets.isEmpty() && mode != Foreach.Mode.AUTO_COLLECTIONS)) { + throw new IllegalArgumentException("foreach collection modes accept exactly one field"); + } + collectionExpression = targets.isEmpty() ? null : targets.get(0); + } List evalClauses = ctx.foreachEvalCommand().foreachEvalClause().stream() .map( @@ -870,37 +876,19 @@ public UnresolvedPlan visitForeachCommand(OpenSearchPPLParser.ForeachCommandCont return new Foreach(mode, options, patterns, collectionExpression, evalClauses); } - private String normalizeForeachOptionValue(String value) { - if (value.startsWith("<<") && value.endsWith(">>")) { - return value.substring(2, value.length() - 2); - } - return value; - } - - private String inferForeachMode(List targets) { - if (targets.size() == 1 - && targets.get(0).logicalExpression() != null - && expressionBuilder.visit(targets.get(0).logicalExpression()) instanceof Function function - && "json_array".equalsIgnoreCase(function.getFuncName())) { - return "json_array"; - } - return "multifield"; + private String stripForeachPlaceholderMarkers(String value) { + return value.startsWith("<<") && value.endsWith(">>") + ? value.substring(2, value.length() - 2) + : value; } - private UnresolvedExpression buildForeachCollectionExpression( - List targets, String mode) { - if (targets.isEmpty() && "auto_collections".equalsIgnoreCase(mode)) { - return null; - } - if (targets.size() != 1) { - throw new IllegalArgumentException("foreach collection modes accept exactly one field"); - } - OpenSearchPPLParser.ForeachTargetContext target = targets.get(0); - if (target.logicalExpression() != null) { - return expressionBuilder.visit(target.logicalExpression()); - } - String fieldName = getTextInQuery(target); - return new Field(new QualifiedName(fieldName)); + /** A bare json_array(...) target implies json_array mode; anything else is multifield. */ + private Foreach.Mode inferForeachMode(List targets) { + return targets.size() == 1 + && targets.get(0) instanceof Function function + && "json_array".equalsIgnoreCase(function.getFuncName()) + ? Foreach.Mode.JSON_ARRAY + : Foreach.Mode.MULTIFIELD; } @Override From 0fd2ad6fea74902b52313324a6e381f07c30dbab Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 9 Jul 2026 06:17:27 +0000 Subject: [PATCH 06/12] Narrow foreachTarget grammar to fix ANTLR error-recovery regression foreachTarget's logicalExpression alternative let ANTLR consider the foreach rule during error recovery of unrelated malformed queries, changing expected-token sets in syntax error messages (caught by UnifiedRelevanceSearchTest#testMatchMissingArguments: 'match(' with a missing field started reporting 26 candidate tokens instead of the expected ','). foreach only ever needs a function call (json_array), a string literal (JSON array text), or a field/wildcard as its target, so accept exactly those instead of the whole expression grammar. Signed-off-by: Songkan Tang --- ppl/src/main/antlr/OpenSearchPPLParser.g4 | 5 +++-- .../org/opensearch/sql/ppl/parser/AstBuilder.java | 11 +++++++---- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4 index 0309167f344..65abc1a99d7 100644 --- a/ppl/src/main/antlr/OpenSearchPPLParser.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLParser.g4 @@ -387,9 +387,10 @@ foreachArgument ; foreachTarget - : wcFieldExpression + : functionCall + | stringLiteral + | wcFieldExpression | STAR - | logicalExpression ; foreachOption diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java index 4870bb88c96..c7f5233ea61 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java @@ -848,10 +848,13 @@ public UnresolvedPlan visitForeachCommand(OpenSearchPPLParser.ForeachCommandCont } else { OpenSearchPPLParser.ForeachTargetContext target = argument.foreachTarget(); patterns.add(getTextInQuery(target)); - targets.add( - target.logicalExpression() != null - ? expressionBuilder.visit(target.logicalExpression()) - : new Field(new QualifiedName(getTextInQuery(target)))); + if (target.functionCall() != null) { + targets.add(expressionBuilder.visit(target.functionCall())); + } else if (target.stringLiteral() != null) { + targets.add(expressionBuilder.visit(target.stringLiteral())); + } else { + targets.add(new Field(new QualifiedName(getTextInQuery(target)))); + } } } Foreach.Mode mode = From 55f1c9e46237c12bf705fa2835351c06759be9c5 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 9 Jul 2026 06:52:37 +0000 Subject: [PATCH 07/12] Assert full logical plans in foreach collection-mode unit tests Replace assertNotNull with verifyLogical so the tests pin the generated reduce/foreach_pair_collection/foreach_pair_item structure, placeholder slot indices, and json_array element-type inference. Signed-off-by: Songkan Tang --- .../ppl/calcite/CalcitePPLForeachTest.java | 71 ++++++++++++++----- 1 file changed, 54 insertions(+), 17 deletions(-) diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java index ae22f62e4a7..265040f76cf 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java @@ -5,7 +5,6 @@ package org.opensearch.sql.ppl.calcite; -import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; import org.apache.calcite.rel.RelNode; @@ -69,7 +68,25 @@ public void testForeachMultivaluePlansReduce() { + " itemstr=NUMBER nums [ eval total = total + NUMBER ] | fields total"; RelNode root = getRelNode(ppl); - assertNotNull(root); + String expectedLogical = + "LogicalProject(total=[reduce(foreach_pair_collection(array(1, 2, 3)), 0, (total," + + " __foreach_pair) -> +(total, foreach_pair_item(__foreach_pair, 0)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + } + + @Test + public void testForeachIterPlaceholderPlansReduce() { + String ppl = + "source=EMP | eval nums = array(10, 20, 30), total = 0 | foreach mode=multivalue" + + " iterstr=IDX nums [ eval total = total + IDX ] | fields total"; + RelNode root = getRelNode(ppl); + + String expectedLogical = + "LogicalProject(total=[reduce(foreach_pair_collection(array(10, 20, 30)), 0, (total," + + " __foreach_pair) -> +(total, foreach_pair_item(__foreach_pair, 1)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); } @Test @@ -79,7 +96,12 @@ public void testForeachJsonArrayPlansReduce() { + " <> ] | fields total"; RelNode root = getRelNode(ppl); - assertNotNull(root); + String expectedLogical = + "LogicalProject(total=[reduce(foreach_pair_collection(foreach_json_array('[1,2,3]':VARCHAR," + + " 'DOUBLE':VARCHAR)), 0.0E0:DOUBLE, (total, __foreach_pair) -> +(total," + + " foreach_pair_item(__foreach_pair, 0)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); } @Test @@ -89,7 +111,12 @@ public void testForeachJsonArrayFunctionWithoutModePlansReduce() { + " <> ] | fields total"; RelNode root = getRelNode(ppl); - assertNotNull(root); + String expectedLogical = + "LogicalProject(total=[reduce(foreach_pair_collection(foreach_json_array(JSON_ARRAY(FLAG(NULL_ON_NULL)," + + " 1, 2, 3), 'DOUBLE':VARCHAR)), 0.0E0:DOUBLE, (total, __foreach_pair) -> +(total," + + " foreach_pair_item(__foreach_pair, 0)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); } @Test @@ -108,7 +135,11 @@ public void testForeachAutoCollectionsPlansReduce() { + " eval total = total + <> ] | fields total"; RelNode root = getRelNode(ppl); - assertNotNull(root); + String expectedLogical = + "LogicalProject(total=[reduce(foreach_pair_collection(array(1, 2, 3)), 0, (total," + + " __foreach_pair) -> +(total, foreach_pair_item(__foreach_pair, 0)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); } @Test @@ -118,7 +149,11 @@ public void testForeachAutoCollectionsWithoutTargetPlansReduce() { + " eval total = total + <> ] | fields total"; RelNode root = getRelNode(ppl); - assertNotNull(root); + String expectedLogical = + "LogicalProject(total=[reduce(foreach_pair_collection(array(1, 2, 3)), 0, (total," + + " __foreach_pair) -> +(total, foreach_pair_item(__foreach_pair, 0)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); } @Test @@ -130,16 +165,18 @@ public void testForeachStringItemWithNumericIterPlansReduce() { + " word_and_num"; RelNode root = getRelNode(ppl); - assertNotNull(root); - } - - @Test - public void testForeachIterPlaceholderPlansReduce() { - String ppl = - "source=EMP | eval nums = array(10, 20, 30), total = 0 | foreach mode=multivalue" - + " iterstr=IDX nums [ eval total = total + IDX ] | fields total"; - RelNode root = getRelNode(ppl); - - assertNotNull(root); + // The captured field `nums` rides in pair slot 2; ITEM is slot 0 and ITER slot 1. + String expectedLogical = + "LogicalProject(word_and_num=[reduce(foreach_pair_collection(CASE(=('':VARCHAR, '')," + + " REGEXP_EXTRACT_ALL('ABCDE':VARCHAR, '.'), SPLIT('ABCDE':VARCHAR, '':VARCHAR))," + + " array('1', '2', '3', '4', '5')), array(), (word_and_num, __foreach_pair) ->" + + " mvappend(word_and_num, CONCAT(foreach_pair_item(__foreach_pair, 0)," + + " ITEM(foreach_pair_item(__foreach_pair, 2)," + + " CASE(<(foreach_pair_item(__foreach_pair, 1), 0)," + + " +(+(ARRAY_LENGTH(foreach_pair_item(__foreach_pair, 2))," + + " foreach_pair_item(__foreach_pair, 1)), 1), +(foreach_pair_item(__foreach_pair, 1)," + + " 1))))))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); } } From 9bdcd2dfe2d5e89a4dfb2070ee9cc40dc6564bd2 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 9 Jul 2026 09:10:27 +0000 Subject: [PATCH 08/12] Restore usage-based element type inference for field-backed JSON arrays The refactor's jsonElementType defaulted opaque expressions (an index field holding JSON text - the primary Splunk use of json_array mode) to VARCHAR, which broke numeric aggregation over such fields: reduce resolved to [ARRAY, DOUBLE, DOUBLE] and failed. Splunk returns 60 for a field holding "[10,20,30]" summed via foreach; the original branch matched that by inferring element type from usage. Bring back the usage scan for the opaque case only: if the item placeholder is consumed by arithmetic the elements are DOUBLE, else VARCHAR. json_array() calls and string literals keep the plan-time content inspection. New coverage: - CalcitePPLForeachTest: plan assertions for field-backed json_array with numeric and string usage - ForeachFieldJsonIT: field holding "[10,20,30]" sums to 60 (Splunk parity), string-content field concats, and native OpenSearch array fields (long field holding [1,2,3]) documented as rejected - the mapping types them as scalar BIGINT at plan time Signed-off-by: Songkan Tang --- .../sql/calcite/ForeachPlanner.java | 67 ++++++++++++--- .../calcite/remote/ForeachFieldJsonIT.java | 82 +++++++++++++++++++ .../ppl/calcite/CalcitePPLForeachTest.java | 32 ++++++++ 3 files changed, 171 insertions(+), 10 deletions(-) create mode 100644 integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java diff --git a/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java b/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java index e2e7475caf2..4a947d2a850 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java +++ b/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java @@ -30,6 +30,7 @@ import org.opensearch.sql.ast.Node; import org.opensearch.sql.ast.expression.DataType; import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.ForeachPlaceholder; import org.opensearch.sql.ast.expression.Function; import org.opensearch.sql.ast.expression.LambdaFunction; import org.opensearch.sql.ast.expression.Literal; @@ -71,6 +72,8 @@ class ForeachPlanner { /** Name of the lambda variable holding the internal pair inside generated reduce calls. */ private static final String PAIR_VAR = "__foreach_pair"; + private static final Set ARITHMETIC_OPERATORS = Set.of("+", "-", "*", "/", "%"); + private final CalciteRelNodeVisitor relVisitor; private final CalciteRexNodeVisitor rexVisitor; @@ -177,16 +180,15 @@ private RelNode planCollection(Foreach node, CalcitePlanContext context) { if (collection == null) { throw new SemanticCheckException("foreach " + mode + " mode requires a field"); } - collection = asArrayExpression(collection, mode, context); + String itemName = node.getOptions().getOrDefault(OPTION_ITEMSTR, PLACEHOLDER_ITEM); + String iterName = node.getOptions().getOrDefault(OPTION_ITERSTR, PLACEHOLDER_ITER); + collection = asArrayExpression(collection, mode, context, node.getEvalClauses(), itemName); RexNode collectionRex = rexVisitor.analyze(collection, context); if (!(collectionRex.getType() instanceof ArraySqlType arrayType)) { throw new SemanticCheckException("foreach " + mode + " mode requires a multivalue field"); } RelDataType itemType = arrayType.getComponentType(); RelDataType iterType = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER); - - String itemName = node.getOptions().getOrDefault(OPTION_ITEMSTR, PLACEHOLDER_ITEM); - String iterName = node.getOptions().getOrDefault(OPTION_ITERSTR, PLACEHOLDER_ITER); List targetAliases = node.getEvalClauses().stream().map(ForeachEvalClause::getTargetTemplate).toList(); List capturedFields = @@ -260,7 +262,11 @@ private UnresolvedExpression firstArrayField(CalcitePlanContext context) { * parses the JSON and casts every element to one inferred type. */ private UnresolvedExpression asArrayExpression( - UnresolvedExpression collection, Foreach.Mode mode, CalcitePlanContext context) { + UnresolvedExpression collection, + Foreach.Mode mode, + CalcitePlanContext context, + List evalClauses, + String itemName) { if (mode == Foreach.Mode.MULTIVALUE || (mode == Foreach.Mode.AUTO_COLLECTIONS && rexVisitor.analyze(collection, context).getType() instanceof ArraySqlType)) { @@ -269,16 +275,27 @@ private UnresolvedExpression asArrayExpression( return new Function( BuiltinFunctionName.FOREACH_JSON_ARRAY.getName().getFunctionName(), List.of( - collection, new Literal(jsonElementType(collection, context).name(), DataType.STRING))); + collection, + new Literal( + jsonElementType(collection, context, evalClauses, itemName).name(), + DataType.STRING))); } /** * Infers the element type a JSON array collection should be read as. JSON only distinguishes * numbers and strings here, so the answer is DOUBLE (gson parses JSON numbers as doubles) or - * VARCHAR. Mixed content is rejected; expressions whose content is unknowable at plan time (e.g. - * a field holding a JSON string) default to VARCHAR. + * VARCHAR. + * + *

For {@code json_array()} calls and string literals the content is visible at plan time; + * mixed content is rejected. For opaque expressions — typically a field holding JSON text, the + * primary Splunk use of json_array mode — content is unknowable, so infer from usage: an item + * placeholder consumed by arithmetic means numeric elements, anything else means strings. */ - private SqlTypeName jsonElementType(UnresolvedExpression collection, CalcitePlanContext context) { + private SqlTypeName jsonElementType( + UnresolvedExpression collection, + CalcitePlanContext context, + List evalClauses, + String itemName) { if (collection instanceof Function function && BuiltinFunctionName.JSON_ARRAY .getName() @@ -303,8 +320,38 @@ private SqlTypeName jsonElementType(UnresolvedExpression collection, CalcitePlan } catch (JsonSyntaxException ignored) { // Malformed JSON literals return null at runtime; type choice is irrelevant. } + return SqlTypeName.VARCHAR; + } + return itemUsedInArithmetic(evalClauses, itemName) ? SqlTypeName.DOUBLE : SqlTypeName.VARCHAR; + } + + private boolean itemUsedInArithmetic(List evalClauses, String itemName) { + return evalClauses.stream() + .anyMatch(clause -> itemUsedInArithmetic(clause.getExpression(), itemName, false)); + } + + private boolean itemUsedInArithmetic(Node node, String itemName, boolean inArithmetic) { + if (isItemReference(node, itemName)) { + return inArithmetic; + } + boolean arithmetic = + inArithmetic + || (node instanceof Function function + && ARITHMETIC_OPERATORS.contains(function.getFuncName())); + return node.getChild().stream() + .anyMatch(child -> itemUsedInArithmetic(child, itemName, arithmetic)); + } + + private boolean isItemReference(Node node, String itemName) { + String name; + if (node instanceof ForeachPlaceholder placeholder) { + name = placeholder.getName(); + } else if (node instanceof QualifiedName qualifiedName) { + name = qualifiedName.toString(); + } else { + return false; } - return SqlTypeName.VARCHAR; + return PLACEHOLDER_ITEM.equalsIgnoreCase(name) || itemName.equalsIgnoreCase(name); } private SqlTypeName elementTypeOf(Set families) { diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java new file mode 100644 index 00000000000..dbb7abd4d76 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java @@ -0,0 +1,82 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.remote; + +import static org.junit.Assert.assertThrows; +import static org.opensearch.sql.util.MatcherUtils.rows; +import static org.opensearch.sql.util.MatcherUtils.schema; +import static org.opensearch.sql.util.MatcherUtils.verifyDataRows; +import static org.opensearch.sql.util.MatcherUtils.verifySchema; + +import java.io.IOException; +import org.json.JSONObject; +import org.junit.jupiter.api.Test; +import org.opensearch.client.Request; +import org.opensearch.client.ResponseException; +import org.opensearch.sql.legacy.TestUtils; +import org.opensearch.sql.ppl.PPLIntegTestCase; + +/** Foreach collection modes over index fields (Splunk-parity scenarios). */ +public class ForeachFieldJsonIT extends PPLIntegTestCase { + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); + + if (!TestUtils.isIndexExist(client(), "test_foreach_field")) { + String mapping = + "{\"mappings\":{\"properties\":{" + + "\"jsonfield\":{\"type\":\"keyword\"}," + + "\"jsonstrs\":{\"type\":\"keyword\"}," + + "\"nativenums\":{\"type\":\"long\"}}}}"; + TestUtils.createIndexByRestClient(client(), "test_foreach_field", mapping); + + Request r = new Request("PUT", "/test_foreach_field/_doc/1?refresh=true"); + r.setJsonEntity( + "{\"jsonfield\": \"[10,20,30]\", \"jsonstrs\": \"[\\\"a\\\",\\\"b\\\"]\"," + + " \"nativenums\": [1, 2, 3]}"); + client().performRequest(r); + } + } + + @Test + public void testJsonArrayModeOnFieldWithNumericContent() throws IOException { + // Splunk: field holding "[10,20,30]" with foreach mode=json_array sums to 60. + JSONObject result = + executeQuery( + "source=test_foreach_field | eval total = 0 | foreach mode=json_array jsonfield [" + + " eval total = total + <> ] | fields total"); + verifySchema(result, schema("total", "double")); + verifyDataRows(result, rows(60.0)); + } + + @Test + public void testJsonArrayModeOnFieldWithStringContent() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach_field | eval r = '' | foreach mode=json_array jsonstrs [" + + " eval r = concat(r, <>) ] | fields r"); + verifySchema(result, schema("r", "string")); + verifyDataRows(result, rows("ab")); + } + + /** + * Native OpenSearch array fields (a long field holding [1,2,3]) are typed as scalar BIGINT at + * plan time because OpenSearch mappings do not distinguish scalars from arrays. foreach + * multivalue mode therefore rejects them; documents the current known limitation rather than the + * desired behavior. + */ + @Test + public void testNativeArrayFieldIsRejected() { + assertThrows( + ResponseException.class, + () -> + executeQuery( + "source=test_foreach_field | eval total = 0 | foreach mode=multivalue nativenums" + + " [ eval total = total + <> ] | fields total")); + } +} diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java index 265040f76cf..1ae097fe445 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLForeachTest.java @@ -179,4 +179,36 @@ public void testForeachStringItemWithNumericIterPlansReduce() { + " LogicalTableScan(table=[[scott, EMP]])\n"; verifyLogical(root, expectedLogical); } + + @Test + public void testForeachJsonArrayFieldNumericUsage() { + // Field content is opaque at plan time; arithmetic on <> selects DOUBLE elements. + String ppl = + "source=EMP | eval total = 0 | foreach mode=json_array ENAME [ eval total = total +" + + " <> ] | fields total"; + RelNode root = getRelNode(ppl); + + String expectedLogical = + "LogicalProject(total=[reduce(foreach_pair_collection(foreach_json_array($1," + + " 'DOUBLE':VARCHAR)), 0.0E0:DOUBLE, (total, __foreach_pair) -> +(total," + + " foreach_pair_item(__foreach_pair, 0)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + } + + @Test + public void testForeachJsonArrayFieldStringUsage() { + // No arithmetic on <> keeps field-backed JSON array elements as VARCHAR. + String ppl = + "source=EMP | eval r = '' | foreach mode=json_array ENAME [ eval r = concat(r, <>)" + + " ] | fields r"; + RelNode root = getRelNode(ppl); + + String expectedLogical = + "LogicalProject(r=[reduce(foreach_pair_collection(foreach_json_array($1," + + " 'VARCHAR':VARCHAR)), '':VARCHAR, (r, __foreach_pair) -> CONCAT(r," + + " foreach_pair_item(__foreach_pair, 0)))])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + } } From 5169eaf2cbe4319c99b8b000d81b106ae602d675 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 9 Jul 2026 09:45:42 +0000 Subject: [PATCH 09/12] Cover nested-field multivalue iteration in ForeachFieldJsonIT Nested-typed fields map to ARRAY at plan time so multivalue mode accepts and iterates them (verified: counting elements of a 2-element nested array returns 2). Pins the third field-backed collection shape alongside JSON-text fields (supported) and native scalar-mapped arrays (rejected). Signed-off-by: Songkan Tang --- .../calcite/remote/ForeachFieldJsonIT.java | 35 +++++++++++++------ 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java index dbb7abd4d76..064441c2b09 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java @@ -27,18 +27,17 @@ public void init() throws Exception { super.init(); enableCalcite(); - if (!TestUtils.isIndexExist(client(), "test_foreach_field")) { + if (!TestUtils.isIndexExist(client(), "test_foreach_field2")) { String mapping = - "{\"mappings\":{\"properties\":{" - + "\"jsonfield\":{\"type\":\"keyword\"}," - + "\"jsonstrs\":{\"type\":\"keyword\"}," - + "\"nativenums\":{\"type\":\"long\"}}}}"; - TestUtils.createIndexByRestClient(client(), "test_foreach_field", mapping); + "{\"mappings\":{\"properties\":{\"jsonfield\":{\"type\":\"keyword\"}," + + "\"jsonstrs\":{\"type\":\"keyword\"},\"nativenums\":{\"type\":\"long\"}," + + "\"nested_objs\":{\"type\":\"nested\",\"properties\":{\"a\":{\"type\":\"long\"}}}}}}"; + TestUtils.createIndexByRestClient(client(), "test_foreach_field2", mapping); - Request r = new Request("PUT", "/test_foreach_field/_doc/1?refresh=true"); + Request r = new Request("PUT", "/test_foreach_field2/_doc/1?refresh=true"); r.setJsonEntity( "{\"jsonfield\": \"[10,20,30]\", \"jsonstrs\": \"[\\\"a\\\",\\\"b\\\"]\"," - + " \"nativenums\": [1, 2, 3]}"); + + " \"nativenums\": [1, 2, 3], \"nested_objs\": [{\"a\": 1}, {\"a\": 2}]}"); client().performRequest(r); } } @@ -48,7 +47,7 @@ public void testJsonArrayModeOnFieldWithNumericContent() throws IOException { // Splunk: field holding "[10,20,30]" with foreach mode=json_array sums to 60. JSONObject result = executeQuery( - "source=test_foreach_field | eval total = 0 | foreach mode=json_array jsonfield [" + "source=test_foreach_field2 | eval total = 0 | foreach mode=json_array jsonfield [" + " eval total = total + <> ] | fields total"); verifySchema(result, schema("total", "double")); verifyDataRows(result, rows(60.0)); @@ -58,7 +57,7 @@ public void testJsonArrayModeOnFieldWithNumericContent() throws IOException { public void testJsonArrayModeOnFieldWithStringContent() throws IOException { JSONObject result = executeQuery( - "source=test_foreach_field | eval r = '' | foreach mode=json_array jsonstrs [" + "source=test_foreach_field2 | eval r = '' | foreach mode=json_array jsonstrs [" + " eval r = concat(r, <>) ] | fields r"); verifySchema(result, schema("r", "string")); verifyDataRows(result, rows("ab")); @@ -76,7 +75,21 @@ public void testNativeArrayFieldIsRejected() { ResponseException.class, () -> executeQuery( - "source=test_foreach_field | eval total = 0 | foreach mode=multivalue nativenums" + "source=test_foreach_field2 | eval total = 0 | foreach mode=multivalue nativenums" + " [ eval total = total + <> ] | fields total")); } + + /** + * Nested-typed fields map to ARRAY<ANY> at plan time, so multivalue mode iterates them. The + * lambda here only counts elements; it does not dereference the object item. + */ + @Test + public void testNestedFieldMultivalueIterates() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach_field2 | eval total = 0 | foreach mode=multivalue nested_objs [" + + " eval total = total + 1 ] | fields total"); + verifySchema(result, schema("total", "int")); + verifyDataRows(result, rows(2)); + } } From b8bed0cc1963de3c6f5791f7150bf593931bfb91 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 9 Jul 2026 10:03:08 +0000 Subject: [PATCH 10/12] Pin cross-feed behavior of foreach collection modes in IT Splunk silently no-ops when a mode is fed the wrong collection shape (verified against Splunk 10.4.0). Our behavior intentionally differs and these tests document it: - json_array mode fed a real array iterates it (total=6) - more permissive than Splunk's no-op - multivalue mode fed a JSON-text field fails at plan time with SemanticCheckException - louder than Splunk's no-op Signed-off-by: Songkan Tang --- .../calcite/remote/ForeachFieldJsonIT.java | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java index 064441c2b09..3f17e4382c6 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/ForeachFieldJsonIT.java @@ -92,4 +92,30 @@ public void testNestedFieldMultivalueIterates() throws IOException { verifySchema(result, schema("total", "int")); verifyDataRows(result, rows(2)); } + + /** + * Cross-feed behavior differs from Splunk by design. Splunk silently no-ops when a mode is fed + * the wrong collection shape; we either still work (json_array mode on a real array iterates it - + * more permissive) or fail loudly at plan time (multivalue mode on a JSON-text field). + */ + @Test + public void testJsonArrayModeOnRealArrayIterates() throws IOException { + JSONObject result = + executeQuery( + "source=test_foreach_field2 | eval nums = array(1, 2, 3), total = 0 | foreach" + + " mode=json_array nums [ eval total = total + <> ] | fields total | head" + + " 1"); + verifySchema(result, schema("total", "double")); + verifyDataRows(result, rows(6.0)); + } + + @Test + public void testMultivalueModeOnJsonTextFieldIsRejected() { + assertThrows( + ResponseException.class, + () -> + executeQuery( + "source=test_foreach_field2 | eval total = 0 | foreach mode=multivalue jsonfield" + + " [ eval total = total + <> ] | fields total")); + } } From 75c38c96cdc99754faf213eba76d04d9594c2c62 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 9 Jul 2026 12:11:50 +0000 Subject: [PATCH 11/12] Add user documentation for PPL foreach command Follows the structure of other command docs (timewrap, eval): syntax, parameters, placeholder table, notes on collection-mode accumulator semantics and type inference, and six runnable examples registered in docs/category.json for doctest. All example outputs verified against a live docTestCluster loaded with the doctest accounts dataset. Signed-off-by: Songkan Tang --- docs/category.json | 1 + docs/user/ppl/cmd/foreach.md | 190 +++++++++++++++++++++++++++++++++++ 2 files changed, 191 insertions(+) create mode 100644 docs/user/ppl/cmd/foreach.md diff --git a/docs/category.json b/docs/category.json index ac1dda966d2..0bf2a8db1e5 100644 --- a/docs/category.json +++ b/docs/category.json @@ -20,6 +20,7 @@ "user/ppl/cmd/fieldformat.md", "user/ppl/cmd/fields.md", "user/ppl/cmd/fillnull.md", + "user/ppl/cmd/foreach.md", "user/ppl/cmd/grok.md", "user/ppl/cmd/head.md", "user/ppl/cmd/join.md", diff --git a/docs/user/ppl/cmd/foreach.md b/docs/user/ppl/cmd/foreach.md new file mode 100644 index 00000000000..b86cd67d765 --- /dev/null +++ b/docs/user/ppl/cmd/foreach.md @@ -0,0 +1,190 @@ +# foreach + +The `foreach` command runs a templated `eval` expression for each field in a field list, each element of a multivalue (array) field, or each element of a JSON array. It eliminates repetitive `eval` statements when the same computation applies to many fields or to every element of a collection. + +## Syntax + +The `foreach` command has the following syntax: + +```syntax +foreach [mode=] [