diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index d7b8ceb1bff..8c1dbee006f 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -74,6 +74,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; @@ -571,6 +572,11 @@ public LogicalPlan visitGraphLookup(GraphLookup node, AnalysisContext context) { throw getOnlyForCalciteException("graphlookup"); } + @Override + public LogicalPlan visitForeach(Foreach node, AnalysisContext context) { + throw getOnlyForCalciteException("foreach"); + } + /** Build {@link ParseExpression} to context and skip to child nodes. */ @Override public LogicalPlan visitParse(Parse node, AnalysisContext context) { 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 6d8415fd7ea..a32354883bf 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; @@ -174,6 +176,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); } @@ -254,6 +260,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..e11ecf541cd --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Foreach.java @@ -0,0 +1,77 @@ +/* + * 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 java.util.Locale; +import java.util.Map; +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 Mode mode; + private final Map options; + private final List fieldPatterns; + private final UnresolvedExpression collectionExpression; + 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); + } + + /** 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 + @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 3f81cdbae4a..162a4895805 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -17,10 +17,12 @@ 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; 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; @@ -78,6 +80,26 @@ public class CalcitePlanContext { @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<>(); + + /** Bare identifiers enabled by explicit options such as {@code itemstr=ITEM}. */ + @Getter private Map foreachIdentifierBindings = new HashMap<>(); + + /** Bindings that become active in lambda contexts cloned from this one. */ + private Map stagedForeachLambdaBindings = new HashMap<>(); + + /** Bare identifier bindings staged for a generated foreach lambda. */ + private Map stagedForeachIdentifierBindings = new HashMap<>(); + + /** Expressions computed by earlier assignments in the same foreach eval iteration. */ + @Getter private Map foreachComputedBindings = new HashMap<>(); + /** * Maps AggregateFunction AST nodes to their output field index for HAVING/post-aggregate * resolution. @@ -125,6 +147,12 @@ 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); + this.foreachIdentifierBindings = new HashMap<>(parent.foreachIdentifierBindings); + this.foreachIdentifierBindings.putAll(parent.stagedForeachIdentifierBindings); + this.foreachComputedBindings = new HashMap<>(parent.foreachComputedBindings); } public RexNode resolveJoinCondition( @@ -203,6 +231,48 @@ public static void clearTimewrapSignals() { timewrapSeries.set(null); } + public void pushForeachBindings( + Map bindings, Map identifierBindings) { + foreachBindings = new HashMap<>(bindings); + foreachIdentifierBindings = new HashMap<>(identifierBindings); + } + + public void stageForeachLambdaBindings( + Map bindings, Map identifierBindings) { + stagedForeachLambdaBindings = new HashMap<>(bindings); + stagedForeachIdentifierBindings = new HashMap<>(identifierBindings); + } + + public void putForeachComputedBinding(String name, RexNode expression) { + foreachComputedBindings.put(name.toUpperCase(java.util.Locale.ROOT), expression); + } + + public void clearForeachBindings() { + foreachBindings.clear(); + foreachIdentifierBindings.clear(); + stagedForeachLambdaBindings.clear(); + stagedForeachIdentifierBindings.clear(); + foreachComputedBindings.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, @Nullable RelDataType pairType) { + public ForeachBinding(String value, ForeachBindingType type) { + 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 13ccdf15197..0df3c571c6b 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -130,6 +130,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; @@ -214,12 +215,14 @@ public class CalciteRelNodeVisitor extends AbstractNodeVisitor 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 849b615f970..1bf1e217b51 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -19,6 +19,8 @@ import java.util.Locale; import java.util.Map; 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 javax.annotation.Nullable; @@ -53,6 +55,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; @@ -79,6 +82,8 @@ 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.CalcitePlanContext.ForeachBindingType; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType; import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; @@ -92,10 +97,13 @@ 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 public class CalciteRexNodeVisitor extends AbstractNodeVisitor { + private static final Pattern FOREACH_TEMPLATE = Pattern.compile("<<([A-Za-z0-9_]+)>>"); + private final CalciteRelNodeVisitor planVisitor; public RexNode analyze(UnresolvedExpression unresolved, CalcitePlanContext context) { @@ -134,6 +142,10 @@ public RexNode visitLiteral(Literal node, CalcitePlanContext context) { case NULL: return rexBuilder.makeNullLiteral(typeFactory.createSqlType(SqlTypeName.NULL)); case STRING: + RexNode foreachTemplate = foreachTemplateLiteral(value.toString(), context); + if (foreachTemplate != null) { + return foreachTemplate; + } if (value.toString().length() == 1) { // To align Spark/PostgreSQL, Char(1) is useful, such as cast('1' to boolean) should // return true @@ -399,9 +411,98 @@ private boolean isBooleanLiteral(RexNode node) { /** Resolve qualified name. Note, the name should be case-sensitive. */ @Override public RexNode visitQualifiedName(QualifiedName node, CalcitePlanContext context) { + String name = node.toString(); + RexNode computed = context.getForeachComputedBindings().get(name.toUpperCase(Locale.ROOT)); + if (computed != null) { + return computed; + } + ForeachBinding binding = + context.getForeachIdentifierBindings().get(name.toUpperCase(Locale.ROOT)); + if (binding != null) { + return foreachBindingToRexNode(node.toString(), binding, context); + } return QualifiedNameResolver.resolve(node, context); } + @Override + public RexNode visitForeachPlaceholder(ForeachPlaceholder node, CalcitePlanContext context) { + ForeachBinding binding = foreachBinding(node.getName(), context); + if (binding == null) { + throw new SemanticCheckException("Unresolved foreach placeholder <<" + node.getName() + ">>"); + } + return foreachBindingToRexNode(node.getName(), binding, context); + } + + private ForeachBinding foreachBinding(String name, CalcitePlanContext context) { + return context.getForeachBindings().get(name.toUpperCase(Locale.ROOT)); + } + + private RexNode foreachTemplateLiteral(String value, CalcitePlanContext context) { + Matcher matcher = FOREACH_TEMPLATE.matcher(value); + List parts = new ArrayList<>(); + int start = 0; + boolean replaced = false; + while (matcher.find()) { + ForeachBinding binding = foreachBinding(matcher.group(1), context); + if (binding == null) { + continue; + } + if (matcher.start() > start) { + parts.add(context.rexBuilder.makeLiteral(value.substring(start, matcher.start()))); + } + if (binding.type() == ForeachBindingType.PAIR_SLOT) { + RelDataType varchar = + context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR); + parts.add( + context.rexBuilder.makeCast( + varchar, foreachBindingToRexNode(matcher.group(1), binding, context), true, true)); + } else { + parts.add(context.rexBuilder.makeLiteral(binding.value())); + } + start = matcher.end(); + replaced = true; + } + if (!replaced) { + return null; + } + if (start < value.length()) { + parts.add(context.rexBuilder.makeLiteral(value.substring(start))); + } + if (parts.isEmpty()) { + return context.rexBuilder.makeLiteral(""); + } + RexNode result = parts.getFirst(); + for (int i = 1; i < parts.size(); i++) { + result = context.rexBuilder.makeCall(SqlStdOperatorTable.CONCAT, result, parts.get(i)); + } + return result; + } + + private RexNode foreachBindingToRexNode( + String name, ForeachBinding binding, CalcitePlanContext context) { + switch (binding.type()) { + case FIELD: + return context.relBuilder.field(binding.value()); + case PAIR_SLOT: + RexLambdaRef pair = context.getRexLambdaRefMap().get(binding.value()); + if (pair == null) { + throw new SemanticCheckException("Unresolved foreach lambda placeholder " + name); + } + // The slot's type is known at plan time, so assign it directly to the opaque extraction + // call. 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()); + } + } + @Override public RexNode visitAlias(Alias node, CalcitePlanContext context) { RexNode expr = analyze(node.getDelegated(), 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..d9ef71b2220 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/calcite/ForeachPlanner.java @@ -0,0 +1,630 @@ +/* + * 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.math.BigDecimal; +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.Compare; +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; +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; +import org.opensearch.sql.expression.function.PPLBuiltinOperators; +import org.opensearch.sql.expression.function.PPLFuncImpTable; + +/** + * 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 static final String STATE_VAR = "__foreach_state"; + + private static final Set ARITHMETIC_OPERATORS = Set.of("+", "-", "*", "/", "%"); + + 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) { + ForeachBindings bindings = + multifieldBindings(node.getFieldPatterns(), fieldName, node.getOptions()); + context.pushForeachBindings(bindings.values(), bindings.identifiers()); + try { + for (ForeachEvalClause clause : node.getEvalClauses()) { + RexNode expr = rexVisitor.analyze(clause.getExpression(), context); + String alias = substituteTemplate(clause.getTargetTemplate(), bindings.values()); + relVisitor.projectPlusOverriding( + List.of(context.relBuilder.alias(expr, alias)), List.of(alias), context); + } + } finally { + context.clearForeachBindings(); + } + } + return context.relBuilder.peek(); + } + + private ForeachBindings multifieldBindings( + List patterns, String fieldName, Map options) { + Map bindings = new LinkedHashMap<>(); + Map identifiers = new LinkedHashMap<>(); + bindOption( + bindings, + identifiers, + new ForeachBinding(fieldName, ForeachBindingType.FIELD), + PLACEHOLDER_FIELD, + OPTION_FIELDSTR, + options); + List orderedPatterns = + Stream.concat( + patterns.stream() + .filter(pattern -> !WildcardUtils.containsWildcard(pattern)) + .filter(pattern -> WildcardUtils.matchesWildcardPattern(pattern, fieldName)), + patterns.stream().filter(WildcardUtils::containsWildcard)) + .toList(); + for (String pattern : orderedPatterns) { + List segments = wildcardSegments(pattern, fieldName); + if (segments == null) { + continue; + } + bindOption( + bindings, + identifiers, + new ForeachBinding(String.join("", segments), ForeachBindingType.LITERAL), + PLACEHOLDER_MATCHSTR, + OPTION_MATCHSTR, + options); + for (int i = 0; i < segments.size(); i++) { + String defaultName = PLACEHOLDER_MATCHSEG + (i + 1); + bindOption( + bindings, + identifiers, + new ForeachBinding(segments.get(i), ForeachBindingType.LITERAL), + defaultName, + OPTION_MATCHSEG + (i + 1), + options); + } + break; + } + return new ForeachBindings(bindings, identifiers); + } + + /** + * 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"); + } + RexNode rawCollection = rexVisitor.analyze(collection, context); + boolean nativeArray = rawCollection.getType() instanceof ArraySqlType; + if ((mode == Foreach.Mode.MULTIVALUE && !nativeArray) + || (mode == Foreach.Mode.JSON_ARRAY && nativeArray)) { + return context.relBuilder.peek(); + } + collection = asArrayExpression(collection, nativeArray, context, node); + RexNode collectionRex = rexVisitor.analyze(collection, context); + if (!(collectionRex.getType() instanceof ArraySqlType arrayType)) { + return context.relBuilder.peek(); + } + RelDataType itemType = arrayType.getComponentType(); + RelDataType iterType = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER); + List targetAliases = + node.getEvalClauses().stream().map(ForeachEvalClause::getTargetTemplate).toList(); + List initialTargets = + targetAliases.stream() + .map(alias -> rexVisitor.analyze(new Field(new QualifiedName(alias)), context)) + .toList(); + List capturedFields = + capturedFields(node.getEvalClauses(), context, node.getOptions(), 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); + RexNode pairedCollectionRex = rexVisitor.analyze(pairedCollection, context); + + List targetTypes = initialTargets.stream().map(RexNode::getType).toList(); + List probedExpressions = + analyzeCollectionEval( + node, + context, + pairedCollectionRex, + state(initialTargets, context), + itemType, + iterType, + capturedFields, + targetAliases, + targetTypes); + targetTypes = probedExpressions.stream().map(RexNode::getType).toList(); + List castInitialTargets = new ArrayList<>(); + for (int i = 0; i < initialTargets.size(); i++) { + castInitialTargets.add( + context.rexBuilder.makeCast(targetTypes.get(i), initialTargets.get(i), true, true)); + } + RexNode initialState = state(castInitialTargets, context); + + ForeachBindings bindings = + collectionBindings(node, itemType, iterType, capturedFields, targetAliases, targetTypes); + context.stageForeachLambdaBindings(bindings.values(), bindings.identifiers()); + try { + CalcitePlanContext lambdaContext = + prepareStateLambdaContext(context, pairedCollectionRex, initialState); + List updatedTargets = analyzeClauses(node.getEvalClauses(), lambdaContext); + RexNode updatedState = state(updatedTargets, lambdaContext); + RexNode lambda = + context.rexBuilder.makeLambdaCall( + updatedState, + List.of( + lambdaContext.getRexLambdaRefMap().get(STATE_VAR), + lambdaContext.getRexLambdaRefMap().get(PAIR_VAR))); + RexNode reducedState = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, + BuiltinFunctionName.REDUCE, + pairedCollectionRex, + initialState, + lambda); + List outputs = new ArrayList<>(); + for (int i = 0; i < targetAliases.size(); i++) { + RexNode value = stateSlot(reducedState, i, targetTypes.get(i), context); + outputs.add(context.relBuilder.alias(value, targetAliases.get(i))); + } + relVisitor.projectPlusOverriding(outputs, targetAliases, context); + } finally { + context.clearForeachBindings(); + } + return context.relBuilder.peek(); + } + + private List analyzeCollectionEval( + Foreach node, + CalcitePlanContext context, + RexNode pairedCollection, + RexNode initialState, + RelDataType itemType, + RelDataType iterType, + List capturedFields, + List targetAliases, + List targetTypes) { + ForeachBindings bindings = + collectionBindings(node, itemType, iterType, capturedFields, targetAliases, targetTypes); + context.stageForeachLambdaBindings(bindings.values(), bindings.identifiers()); + try { + return analyzeClauses( + node.getEvalClauses(), + prepareStateLambdaContext(context, pairedCollection, initialState)); + } finally { + context.clearForeachBindings(); + } + } + + private CalcitePlanContext prepareStateLambdaContext( + CalcitePlanContext context, RexNode pairedCollection, RexNode initialState) { + LambdaFunction template = + new LambdaFunction( + new Literal(0, DataType.INTEGER), + List.of(new QualifiedName(STATE_VAR), new QualifiedName(PAIR_VAR))); + return rexVisitor.prepareLambdaContext( + context, + template, + List.of(pairedCollection, initialState), + BuiltinFunctionName.REDUCE.getName().getFunctionName(), + initialState.getType()); + } + + private List analyzeClauses( + List clauses, CalcitePlanContext lambdaContext) { + List expressions = new ArrayList<>(); + for (ForeachEvalClause clause : clauses) { + RexNode expression = rexVisitor.analyze(clause.getExpression(), lambdaContext); + expressions.add(expression); + lambdaContext.putForeachComputedBinding(clause.getTargetTemplate(), expression); + } + return expressions; + } + + private RexNode state(List values, CalcitePlanContext context) { + return PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.FOREACH_STATE, values.toArray(RexNode[]::new)); + } + + private RexNode stateSlot(RexNode state, int slot, RelDataType type, CalcitePlanContext context) { + return context.rexBuilder.makeCall( + type, + PPLBuiltinOperators.FOREACH_PAIR_ITEM, + List.of(state, context.rexBuilder.makeExactLiteral(BigDecimal.valueOf(slot)))); + } + + private ForeachBindings collectionBindings( + Foreach node, + RelDataType itemType, + RelDataType iterType, + List capturedFields, + List targetAliases, + List targetTypes) { + Map bindings = new LinkedHashMap<>(); + Map identifiers = new LinkedHashMap<>(); + bindPairPlaceholder( + bindings, identifiers, 0, itemType, PLACEHOLDER_ITEM, OPTION_ITEMSTR, node.getOptions()); + bindPairPlaceholder( + bindings, identifiers, 1, iterType, PLACEHOLDER_ITER, OPTION_ITERSTR, node.getOptions()); + for (int i = 0; i < capturedFields.size(); i++) { + RelDataTypeField field = capturedFields.get(i); + bindIdentifier(identifiers, field.getName(), PAIR_VAR, i + 2, field.getType()); + } + for (int i = 0; i < targetAliases.size(); i++) { + bindIdentifier(identifiers, targetAliases.get(i), STATE_VAR, i, targetTypes.get(i)); + } + return new ForeachBindings(bindings, identifiers); + } + + private void bindPairPlaceholder( + Map bindings, + Map identifiers, + int slot, + RelDataType type, + String placeholder, + String option, + Map options) { + ForeachBinding binding = new ForeachBinding(PAIR_VAR, ForeachBindingType.PAIR_SLOT, slot, type); + bindings.put(placeholder, binding); + if (options.containsKey(option)) { + String customName = options.get(option).toUpperCase(Locale.ROOT); + bindings.put(customName, binding); + identifiers.put(customName, binding); + } + } + + private void bindIdentifier( + Map identifiers, + String name, + String variable, + int slot, + RelDataType type) { + String key = name.toUpperCase(Locale.ROOT); + identifiers.put(key, new ForeachBinding(variable, ForeachBindingType.PAIR_SLOT, slot, type)); + } + + 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, + boolean nativeArray, + CalcitePlanContext context, + Foreach node) { + if (nativeArray) { + return collection; + } + return new Function( + BuiltinFunctionName.FOREACH_JSON_ARRAY.getName().getFunctionName(), + List.of( + collection, + new Literal(jsonElementType(collection, context, node).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. + * + *

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, Foreach node) { + 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; + } + return itemRequiresNumericType(node, context) ? SqlTypeName.DOUBLE : SqlTypeName.VARCHAR; + } + + private boolean itemRequiresNumericType(Foreach node, CalcitePlanContext context) { + String itemName = node.getOptions().getOrDefault(OPTION_ITEMSTR, PLACEHOLDER_ITEM); + boolean customIdentifier = node.getOptions().containsKey(OPTION_ITEMSTR); + return node.getEvalClauses().stream() + .anyMatch( + clause -> + itemRequiresNumericType( + clause.getExpression(), itemName, customIdentifier, context)); + } + + private boolean itemRequiresNumericType( + Node node, String itemName, boolean customIdentifier, CalcitePlanContext context) { + if (node instanceof Compare compare) { + if ((containsItem(compare.getLeft(), itemName, customIdentifier) + && isNumericExpression(compare.getRight(), context)) + || (containsItem(compare.getRight(), itemName, customIdentifier) + && isNumericExpression(compare.getLeft(), context))) { + return true; + } + } + if (node instanceof Function function) { + for (int i = 0; i < function.getFuncArgs().size(); i++) { + if (containsItem(function.getFuncArgs().get(i), itemName, customIdentifier) + && (ARITHMETIC_OPERATORS.contains(function.getFuncName()) + || PPLFuncImpTable.INSTANCE.requiresNumericArgument(function.getFuncName(), i))) { + return true; + } + } + } + return node.getChild().stream() + .anyMatch(child -> itemRequiresNumericType(child, itemName, customIdentifier, context)); + } + + private boolean containsItem(Node node, String itemName, boolean customIdentifier) { + if (isItemReference(node, itemName, customIdentifier)) { + return true; + } + return node.getChild().stream() + .anyMatch(child -> containsItem(child, itemName, customIdentifier)); + } + + private boolean isItemReference(Node node, String itemName, boolean customIdentifier) { + String name; + if (node instanceof ForeachPlaceholder placeholder) { + name = placeholder.getName(); + } else if (node instanceof QualifiedName qualifiedName) { + if (!customIdentifier) { + return false; + } + name = qualifiedName.toString(); + } else { + return false; + } + return PLACEHOLDER_ITEM.equalsIgnoreCase(name) || itemName.equalsIgnoreCase(name); + } + + private boolean isNumericExpression(Node node, CalcitePlanContext context) { + if (node instanceof Literal literal) { + return Set.of( + DataType.SHORT, + DataType.INTEGER, + DataType.LONG, + DataType.FLOAT, + DataType.DOUBLE, + DataType.DECIMAL) + .contains(literal.getType()); + } + try { + return rexVisitor.analyze((UnresolvedExpression) node, context).getType().getFamily() + == SqlTypeFamily.NUMERIC; + } catch (RuntimeException e) { + return false; + } + } + + 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, + Map options, + List targetAliases) { + Set excluded = + Stream.concat(Stream.of(PAIR_VAR), targetAliases.stream()) + .map(name -> name.toUpperCase(Locale.ROOT)) + .collect(Collectors.toSet()); + if (options.containsKey(OPTION_ITEMSTR)) { + excluded.add(options.get(OPTION_ITEMSTR).toUpperCase(Locale.ROOT)); + } + if (options.containsKey(OPTION_ITERSTR)) { + excluded.add(options.get(OPTION_ITERSTR).toUpperCase(Locale.ROOT)); + } + 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 bindOption( + Map bindings, + Map identifiers, + ForeachBinding binding, + String defaultName, + String option, + Map options) { + bindings.put(defaultName.toUpperCase(Locale.ROOT), binding); + if (options.containsKey(option)) { + String customName = options.get(option).toUpperCase(Locale.ROOT); + bindings.put(customName, binding); + identifiers.put(customName, 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; + } + + private record ForeachBindings( + Map values, Map identifiers) {} +} 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..e30d723ccfc 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,9 @@ 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_STATE(FunctionName.of("foreach_state"), 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..918a4715291 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairCollectionFunctionImpl.java @@ -0,0 +1,85 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.expression.function.CollectionUDF; + +import java.util.ArrayList; +import java.util.Arrays; +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 Arrays.asList(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..075453f38a6 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachPairItemFunctionImpl.java @@ -0,0 +1,63 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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; +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.sql.type.SqlReturnTypeInference; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.sql.expression.function.ImplementorUDF; +import org.opensearch.sql.expression.function.UDFOperandMetadata; + +/** + * Extracts one slot of an internal foreach pair: {@code foreach_pair_item(pair, index)}. Returns + * OTHER by default; the foreach planner assigns every call its inferred slot type explicitly. + */ +public class ForeachPairItemFunctionImpl extends ImplementorUDF { + public ForeachPairItemFunctionImpl() { + super(new ForeachPairItemImplementor(), NullPolicy.NONE); + } + + @Override + public SqlReturnTypeInference getReturnTypeInference() { + return opBinding -> + opBinding + .getTypeFactory() + .createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(SqlTypeName.OTHER), 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(); + 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/ForeachStateFunctionImpl.java b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachStateFunctionImpl.java new file mode 100644 index 00000000000..f71294f05b2 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachStateFunctionImpl.java @@ -0,0 +1,60 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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; +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; + +/** Packs the typed accumulator slots used by a collection-mode foreach eval. */ +public class ForeachStateFunctionImpl extends ImplementorUDF { + public ForeachStateFunctionImpl() { + super(new ForeachStateImplementor(), NullPolicy.NONE); + } + + @Override + public SqlReturnTypeInference getReturnTypeInference() { + return opBinding -> { + RelDataType slot = + opBinding + .getTypeFactory() + .createTypeWithNullability( + opBinding.getTypeFactory().createSqlType(SqlTypeName.OTHER), true); + return SqlTypeUtil.createArrayType(opBinding.getTypeFactory(), slot, true); + }; + } + + @Override + public UDFOperandMetadata getOperandMetadata() { + return null; + } + + public static class ForeachStateImplementor implements NotNullImplementor { + @Override + public Expression implement( + RexToLixTranslator translator, RexCall call, List translatedOperands) { + return Expressions.call( + Types.lookupMethod(ForeachStateFunctionImpl.class, "eval", Object[].class), + translatedOperands); + } + } + + public static Object eval(Object... args) { + return Arrays.asList(args); + } +} 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..e8f6d75d7ab 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,8 @@ 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.fun.SqlCastFunction; import org.apache.calcite.sql.type.SqlReturnTypeInference; import org.apache.calcite.sql.type.SqlTypeName; @@ -68,6 +70,11 @@ public static RelDataType inferReturnTypeFromLambda( public static RexCall reInferReturnTypeForRexCallInsideLambda( RexCall rexCall, Map argTypes, RelDataTypeFactory typeFactory) { + // 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 || rexCall.getOperator() instanceof SqlCastFunction) { + return rexCall; + } List filledOperands = new ArrayList<>(); List rexCallOperands = rexCall.getOperands(); for (RexNode rexNode : rexCallOperands) { @@ -89,6 +96,14 @@ public static RexCall reInferReturnTypeForRexCallInsideLambda( .getOperator() .inferReturnType( new RexCallBinding(typeFactory, rexCall.getOperator(), filledOperands, List.of())); + // An opaque helper call 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 + || returnType.getSqlTypeName() == SqlTypeName.OTHER) + && rexCall.getType().getSqlTypeName() != SqlTypeName.ANY + && rexCall.getType().getSqlTypeName() != SqlTypeName.OTHER) { + returnType = rexCall.getType(); + } return rexCall.clone(returnType, filledOperands); } } 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..d64f04bb9ad 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,9 @@ 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.ForeachStateFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.MVAppendFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.MVFindFunctionImpl; import org.opensearch.sql.expression.function.CollectionUDF.MVZipFunctionImpl; @@ -54,6 +57,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 +409,14 @@ 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 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 FOREACH_STATE = + new ForeachStateFunctionImpl().toUDF("foreach_state"); 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 b9350d18d84..151c4a96655 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,9 @@ 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.FOREACH_PAIR_COLLECTION; +import static org.opensearch.sql.expression.function.BuiltinFunctionName.FOREACH_STATE; 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; @@ -273,6 +276,7 @@ import java.math.BigDecimal; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -315,6 +319,8 @@ import org.opensearch.sql.calcite.utils.PPLOperandTypes; import org.opensearch.sql.calcite.utils.PlanUtils; import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils; +import org.opensearch.sql.data.type.ExprCoreType; +import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.exception.ExpressionEvaluationException; import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.expression.function.CollectionUDF.MVIndexFunctionImp; @@ -428,6 +434,63 @@ private PPLFuncImpTable(Builder builder, AggBuilder aggBuilder) { this.aggExternalFunctionRegistry = new ConcurrentHashMap<>(); } + /** Returns whether every known signature requires a numeric value at the given argument. */ + public boolean requiresNumericArgument(String functionName, int argumentIndex) { + Optional builtin = BuiltinFunctionName.of(functionName); + if (builtin.isEmpty()) { + return false; + } + List> implementations = + new ArrayList<>(functionRegistry.getOrDefault(builtin.get(), List.of())); + implementations.addAll(externalFunctionRegistry.getOrDefault(builtin.get(), List.of())); + boolean foundArgument = false; + for (Pair implementation : implementations) { + PPLTypeChecker checker = implementation.getKey().typeChecker(); + if (checker == null) { + return false; + } + try { + List> signatures = + checker.getParameterTypes().stream() + .filter(parameters -> argumentIndex < parameters.size()) + .toList(); + if (signatures.isEmpty()) { + return false; + } + foundArgument = true; + List acceptedTypes = + signatures.stream().map(parameters -> parameters.get(argumentIndex)).toList(); + if (acceptedTypes.stream().allMatch(ExprCoreType.numberTypes()::contains)) { + continue; + } + if (acceptedTypes.stream().anyMatch(type -> type != ExprCoreType.UNKNOWN) + || !requiresNumericByValidation(checker, signatures, argumentIndex)) { + return false; + } + } catch (RuntimeException e) { + return false; + } + } + return foundArgument; + } + + private boolean requiresNumericByValidation( + PPLTypeChecker checker, List> signatures, int argumentIndex) { + RelDataType numericType = TYPE_FACTORY.createSqlType(SqlTypeName.DOUBLE); + RelDataType stringType = TYPE_FACTORY.createSqlType(SqlTypeName.VARCHAR); + return signatures.stream() + .anyMatch( + signature -> { + List numericArguments = + new ArrayList<>(Collections.nCopies(signature.size(), numericType)); + if (!checker.checkOperandTypes(numericArguments)) { + return false; + } + numericArguments.set(argumentIndex, stringType); + return !checker.checkOperandTypes(numericArguments); + }); + } + /** * Register an operator from external services dynamically. * @@ -956,11 +1019,17 @@ void populate() { wrapSqlOperandTypeChecker( SqlLibraryOperators.REGEXP_REPLACE_3.getOperandTypeChecker(), REPLACE.name(), false)); registerOperator(UPPER, SqlStdOperatorTable.UPPER); - registerOperator(ABS, SqlStdOperatorTable.ABS); - registerOperator(ACOS, SqlStdOperatorTable.ACOS); - registerOperator(ASIN, SqlStdOperatorTable.ASIN); - registerOperator(ATAN, SqlStdOperatorTable.ATAN); - registerOperator(ATAN2, SqlStdOperatorTable.ATAN2); + registerOperator(ABS, SqlStdOperatorTable.ABS, PPLTypeChecker.family(SqlTypeFamily.NUMERIC)); + registerOperator( + ACOS, SqlStdOperatorTable.ACOS, PPLTypeChecker.family(SqlTypeFamily.NUMERIC)); + registerOperator( + ASIN, SqlStdOperatorTable.ASIN, PPLTypeChecker.family(SqlTypeFamily.NUMERIC)); + registerOperator( + ATAN, SqlStdOperatorTable.ATAN, PPLTypeChecker.family(SqlTypeFamily.NUMERIC)); + registerOperator( + ATAN2, + SqlStdOperatorTable.ATAN2, + PPLTypeChecker.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC)); // TODO, workaround to support sequence CompositeOperandTypeChecker. registerOperator( CEIL, @@ -1235,6 +1304,9 @@ void populate() { registerOperator(FILTER, PPLBuiltinOperators.FILTER); 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_STATE, PPLBuiltinOperators.FOREACH_STATE); // 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..28dff66a8ae --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/ForeachJsonArrayFunctionImpl.java @@ -0,0 +1,110 @@ +/* + * 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.JsonArray; +import com.google.gson.JsonElement; +import com.google.gson.JsonSyntaxException; +import java.math.BigDecimal; +import java.util.List; +import java.util.stream.StreamSupport; +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.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; + +/** 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.NONE); + } + + @Override + public SqlReturnTypeInference getReturnTypeInference() { + 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 null; + } + + public static class ForeachJsonArrayImplementor implements NotNullImplementor { + @Override + public Expression implement( + RexToLixTranslator translator, RexCall call, List translatedOperands) { + return Expressions.call( + Types.lookupMethod(ForeachJsonArrayFunctionImpl.class, "eval", Object[].class), + translatedOperands); + } + } + + public static Object eval(Object... args) { + if (args.length != 2 || args[0] == null) { + return null; + } + SqlTypeName elementType = SqlTypeName.valueOf(String.valueOf(args[1])); + try { + if (args[0] instanceof List values) { + return values.stream().map(value -> cast(value, elementType)).toList(); + } + JsonArray values = gson.fromJson(String.valueOf(args[0]), JsonArray.class); + if (values == null) { + return List.of(); + } + return StreamSupport.stream(values.spliterator(), false) + .map(value -> cast(value, elementType)) + .toList(); + } catch (JsonSyntaxException e) { + return List.of(); + } + } + + private static Object cast(Object value, SqlTypeName elementType) { + if (value instanceof JsonElement element) { + if (element.isJsonNull()) { + return elementType == SqlTypeName.VARCHAR ? "null" : null; + } + if (element.isJsonArray() || element.isJsonObject()) { + return element.toString(); + } + value = + element.getAsJsonPrimitive().isNumber() ? element.getAsNumber() : element.getAsString(); + } + if (value == null) { + return null; + } + return switch (elementType) { + case DOUBLE -> ((Number) value).doubleValue(); + case VARCHAR -> + value instanceof List || value instanceof java.util.Map + ? gson.toJson(value) + : String.valueOf(value); + case DECIMAL -> BigDecimal.valueOf(((Number) value).doubleValue()); + default -> value; + }; + } +} diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java index fee5bfcaa66..d073a6802aa 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java @@ -1943,6 +1943,24 @@ public void regex_command_throws_unsupported_exception_with_legacy_engine() { "Regex is supported only when plugins.calcite.enabled=true", exception.getMessage()); } + @Test + public void foreach_command_throws_unsupported_exception_with_legacy_engine() { + UnsupportedOperationException exception = + assertThrows( + UnsupportedOperationException.class, + () -> + analyze( + new org.opensearch.sql.ast.tree.Foreach( + org.opensearch.sql.ast.tree.Foreach.Mode.MULTIFIELD, + ImmutableMap.of(), + ImmutableList.of("integer_value"), + null, + ImmutableList.of()) + .attach(relation("schema")))); + assertEquals( + "foreach is supported only when plugins.calcite.enabled=true", exception.getMessage()); + } + @Test public void rex_command_throws_unsupported_operation_exception_in_legacy_engine() { UnsupportedOperationException exception = diff --git a/core/src/test/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachFunctionImplTest.java b/core/src/test/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachFunctionImplTest.java new file mode 100644 index 00000000000..e3c3724e0aa --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/expression/function/CollectionUDF/ForeachFunctionImplTest.java @@ -0,0 +1,49 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.expression.function.CollectionUDF; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.Arrays; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.expression.function.jsonUDF.ForeachJsonArrayFunctionImpl; + +public class ForeachFunctionImplTest { + + @Test + public void testPairCollectionPreservesNullElements() { + List pairs = + (List) + ForeachPairCollectionFunctionImpl.eval( + new Object[] {new Object[] {1, null, 3}, "captured"}); + + assertEquals(3, pairs.size()); + assertArrayEquals(new Object[] {1, 0, "captured"}, (Object[]) pairs.get(0)); + assertArrayEquals(new Object[] {null, 1, "captured"}, (Object[]) pairs.get(1)); + assertArrayEquals(new Object[] {3, 2, "captured"}, (Object[]) pairs.get(2)); + } + + @Test + public void testJsonArrayPreservesTopLevelNestedValuesAndNull() { + Object values = ForeachJsonArrayFunctionImpl.eval("[[1,2],{\"a\":1},null]", "VARCHAR"); + + assertEquals(List.of("[1,2]", "{\"a\":1}", "null"), values); + } + + @Test + public void testMalformedJsonArrayIsEmpty() { + assertEquals(List.of(), ForeachJsonArrayFunctionImpl.eval("not-json", "VARCHAR")); + } + + @Test + public void testStatePreservesHeterogeneousAndNullSlots() { + assertEquals( + Arrays.asList(3, null, "seen"), + ForeachStateFunctionImpl.eval(new Object[] {3, null, "seen"})); + } +} diff --git a/docs/category.json b/docs/category.json index 1e45f3e65bb..7296b089650 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..45dcb939f8a --- /dev/null +++ b/docs/user/ppl/cmd/foreach.md @@ -0,0 +1,191 @@ +# 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=] [