From 04d89b913d247fc4f29b935eaaf74a09d5db9339 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 23 Jul 2026 11:08:54 +0800 Subject: [PATCH 01/10] Implement PPL format command Signed-off-by: Songkan Tang --- .../org/opensearch/sql/analysis/Analyzer.java | 6 + .../sql/ast/AbstractNodeVisitor.java | 5 + .../org/opensearch/sql/ast/tree/Format.java | 61 ++++++ .../sql/calcite/CalciteRelNodeVisitor.java | 8 + .../opensearch/sql/calcite/FormatPlanner.java | 184 ++++++++++++++++++ docs/user/ppl/cmd/format.md | 108 ++++++++++ docs/user/ppl/index.md | 1 + .../sql/calcite/CalciteNoPushdownIT.java | 1 + .../sql/calcite/remote/CalciteExplainIT.java | 14 ++ .../remote/CalciteFormatCommandIT.java | 64 ++++++ .../sql/ppl/NewAddedCommandsIT.java | 15 ++ ppl/src/main/antlr/OpenSearchPPLLexer.g4 | 4 + ppl/src/main/antlr/OpenSearchPPLParser.g4 | 19 ++ .../opensearch/sql/ppl/parser/AstBuilder.java | 46 +++++ .../sql/ppl/utils/PPLQueryDataAnonymizer.java | 19 ++ .../sql/ppl/calcite/CalcitePPLFormatTest.java | 153 +++++++++++++++ .../sql/ppl/parser/AstBuilderTest.java | 26 +++ .../ppl/utils/PPLQueryDataAnonymizerTest.java | 10 + 18 files changed, 744 insertions(+) create mode 100644 core/src/main/java/org/opensearch/sql/ast/tree/Format.java create mode 100644 core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java create mode 100644 docs/user/ppl/cmd/format.md create mode 100644 integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java create mode 100644 ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java 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 916cc00bc4a..f9373bd238a 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -75,6 +75,7 @@ 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.Format; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.Head; import org.opensearch.sql.ast.tree.Join; @@ -833,6 +834,11 @@ public LogicalPlan visitReverse(Reverse node, AnalysisContext context) { throw getOnlyForCalciteException("Reverse"); } + @Override + public LogicalPlan visitFormat(Format node, AnalysisContext context) { + throw getOnlyForCalciteException("Format"); + } + @Override public LogicalPlan visitSpath(SPath node, AnalysisContext context) { throw getOnlyForCalciteException("Spath"); 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 266f8f46f7d..36156d7c58c 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -64,6 +64,7 @@ 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.Format; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.Head; import org.opensearch.sql.ast.tree.Join; @@ -166,6 +167,10 @@ public T visitFlatten(Flatten node, C context) { return visitChildren(node, context); } + public T visitFormat(Format node, C context) { + return visitChildren(node, context); + } + public T visitTrendline(Trendline node, C context) { return visitChildren(node, context); } diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Format.java b/core/src/main/java/org/opensearch/sql/ast/tree/Format.java new file mode 100644 index 00000000000..7f2191b6b80 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Format.java @@ -0,0 +1,61 @@ +/* + * 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.Setter; +import lombok.ToString; +import org.opensearch.sql.ast.AbstractNodeVisitor; + +/** AST node that collapses input rows into a formatted search expression. */ +@Getter +@Setter +@ToString +@EqualsAndHashCode(callSuper = false) +@RequiredArgsConstructor +public class Format extends UnresolvedPlan { + + public static final String DEFAULT_MV_SEPARATOR = "OR"; + public static final int DEFAULT_MAX_RESULTS = 0; + public static final String DEFAULT_ROW_PREFIX = "("; + public static final String DEFAULT_COLUMN_PREFIX = "("; + public static final String DEFAULT_COLUMN_SEPARATOR = "AND"; + public static final String DEFAULT_COLUMN_END = ")"; + public static final String DEFAULT_ROW_SEPARATOR = "OR"; + public static final String DEFAULT_ROW_END = ")"; + public static final String DEFAULT_EMPTY_STRING = "NOT( )"; + + private final String mvSeparator; + private final int maxResults; + private final String rowPrefix; + private final String columnPrefix; + private final String columnSeparator; + private final String columnEnd; + private final String rowSeparator; + private final String rowEnd; + private final String emptyString; + private UnresolvedPlan child; + + @Override + public Format 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 visitor, C context) { + return visitor.visitFormat(this, context); + } +} 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 b7b52e16d82..77ff4b1df9e 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -133,6 +133,7 @@ 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.Format; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.GraphLookup.Direction; import org.opensearch.sql.ast.tree.Head; @@ -960,6 +961,13 @@ public RelNode visitReverse( return context.relBuilder.peek(); } + /** Collapses all user-facing input fields into one Splunk-compatible search expression. */ + @Override + public RelNode visitFormat(Format node, CalcitePlanContext context) { + visitChildren(node, context); + return new FormatPlanner().plan(node, context); + } + @Override public RelNode visitTranspose( org.opensearch.sql.ast.tree.Transpose node, CalcitePlanContext context) { diff --git a/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java new file mode 100644 index 00000000000..8ef67ec1a7d --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java @@ -0,0 +1,184 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite; + +import java.util.Comparator; +import java.util.List; +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.fun.SqlLibraryOperators; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.sql.type.SqlTypeUtil; +import org.apache.calcite.tools.RelBuilder; +import org.opensearch.sql.ast.tree.Format; + +/** Lowers the PPL {@code format} command to Calcite projects and a global aggregation. */ +public class FormatPlanner { + + /** Builds a single-row relation containing the formatted search expression. */ + public RelNode plan(Format node, CalcitePlanContext context) { + RelBuilder builder = context.relBuilder; + if (node.getMaxResults() > 0) { + builder.limit(0, node.getMaxResults()); + } + + List fields = + builder.peek().getRowType().getFieldList().stream() + .filter(field -> !field.getName().startsWith("_")) + .sorted(Comparator.comparing(RelDataTypeField::getName)) + .toList(); + + if (fields.isEmpty()) { + builder.values(new String[] {"search"}, node.getEmptyString()); + return builder.peek(); + } + + List formattedFields = + fields.stream().map(field -> formatField(field, node, context)).toList(); + RexNode joinedFields = + arrayJoin( + context, compactArray(context, formattedFields), " " + node.getColumnSeparator() + " "); + RexNode rowExpression = + concat( + context, + stringLiteral(node.getColumnPrefix() + " ", context), + joinedFields, + stringLiteral(" " + node.getColumnEnd(), context)); + RexNode nonEmptyRow = ifNotEmpty(joinedFields, rowExpression, context); + + builder.project(List.of(nonEmptyRow), List.of("__format_row")); + RexNode rowRef = builder.field("__format_row"); + builder.aggregate( + builder.groupKey(), + builder + .aggregateCall(SqlLibraryOperators.ARRAY_AGG, rowRef) + .filter(builder.isNotNull(rowRef)) + .as("__format_rows")); + + RexNode joinedRows = + arrayJoin(context, builder.field("__format_rows"), " " + node.getRowSeparator() + " "); + RexNode formatted = + concat( + context, + stringLiteral(node.getRowPrefix() + " ", context), + joinedRows, + stringLiteral(" " + node.getRowEnd(), context)); + RexNode result = + context.relBuilder.call( + SqlStdOperatorTable.CASE, + isNotEmpty(joinedRows, context), + formatted, + stringLiteral(node.getEmptyString(), context)); + builder.project(List.of(result), List.of("search"), true); + return builder.peek(); + } + + private RexNode formatField(RelDataTypeField field, Format node, CalcitePlanContext context) { + RelBuilder builder = context.relBuilder; + RexNode value = builder.field(field.getIndex()); + RexNode formattedValue; + RexNode hasValue = builder.isNotNull(value); + String fieldName = formatFieldName(field.getName()); + boolean rawValue = field.getName().equals("search") || field.getName().equals("query"); + + if (SqlTypeUtil.isArray(field.getType()) || SqlTypeUtil.isMultiset(field.getType())) { + RelDataType varchar = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR); + RelDataType varcharArray = context.rexBuilder.getTypeFactory().createArrayType(varchar, -1); + RexNode values = + builder.call( + SqlLibraryOperators.ARRAY_COMPACT, context.rexBuilder.makeCast(varcharArray, value)); + String repeatedValueSeparator = + rawValue + ? "\" " + node.getMvSeparator() + " \"" + : "\" " + node.getMvSeparator() + " " + fieldName + "=\""; + RexNode joinedValues = arrayJoin(context, values, repeatedValueSeparator); + formattedValue = + concat( + context, + stringLiteral(rawValue ? "( \"" : "( " + fieldName + "=\"", context), + joinedValues, + stringLiteral("\" )", context)); + hasValue = builder.call(SqlStdOperatorTable.AND, hasValue, isNotEmpty(joinedValues, context)); + } else { + RexNode stringValue = escapeValue(builder.cast(value, SqlTypeName.VARCHAR), context); + formattedValue = + concat( + context, + stringLiteral(rawValue ? "\"" : fieldName + "=\"", context), + stringValue, + stringLiteral("\"", context)); + } + + return builder.call( + SqlStdOperatorTable.CASE, + hasValue, + formattedValue, + context.rexBuilder.makeNullLiteral(formattedValue.getType())); + } + + private RexNode escapeValue(RexNode value, CalcitePlanContext context) { + RexNode escapedBackslashes = + context.relBuilder.call( + SqlStdOperatorTable.REPLACE, + value, + stringLiteral("\\", context), + stringLiteral("\\\\", context)); + return context.relBuilder.call( + SqlStdOperatorTable.REPLACE, + escapedBackslashes, + stringLiteral("\"", context), + stringLiteral("\\\"", context)); + } + + private String formatFieldName(String fieldName) { + if (fieldName.matches("[A-Za-z][A-Za-z0-9_]*")) { + return fieldName; + } + return "\"" + fieldName.replace("\\", "\\\\").replace("\"", "\\\"") + "\""; + } + + private RexNode compactArray(CalcitePlanContext context, List values) { + return context.relBuilder.call( + SqlLibraryOperators.ARRAY_COMPACT, + context.rexBuilder.makeCall(SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, values)); + } + + private RexNode arrayJoin(CalcitePlanContext context, RexNode values, String separator) { + return context.relBuilder.call( + SqlLibraryOperators.ARRAY_JOIN, values, stringLiteral(separator, context)); + } + + private RexNode ifNotEmpty(RexNode testedValue, RexNode result, CalcitePlanContext context) { + return context.relBuilder.call( + SqlStdOperatorTable.CASE, + isNotEmpty(testedValue, context), + result, + context.rexBuilder.makeNullLiteral(result.getType())); + } + + private RexNode isNotEmpty(RexNode value, CalcitePlanContext context) { + return context.relBuilder.call( + SqlStdOperatorTable.GREATER_THAN, + context.relBuilder.call(SqlStdOperatorTable.CHAR_LENGTH, value), + context.relBuilder.literal(0)); + } + + private RexNode concat(CalcitePlanContext context, RexNode... operands) { + RexNode result = operands[0]; + for (int i = 1; i < operands.length; i++) { + result = context.relBuilder.call(SqlStdOperatorTable.CONCAT, result, operands[i]); + } + return result; + } + + private RexNode stringLiteral(String value, CalcitePlanContext context) { + return context.rexBuilder.makeLiteral( + value, context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR), true); + } +} diff --git a/docs/user/ppl/cmd/format.md b/docs/user/ppl/cmd/format.md new file mode 100644 index 00000000000..e966282f4d2 --- /dev/null +++ b/docs/user/ppl/cmd/format.md @@ -0,0 +1,108 @@ +# format + +The `format` command collapses tabular input into one row containing a `search` string. Within each input row, it renders non-null fields as `field="value"` expressions joined by a column separator. It then joins the rendered rows with a row separator. + +`format` produces the expression as data. It does not parse or execute the generated `search` string. + +## Syntax + +```syntax +format [mvsep=""] [maxresults=] ["" "" "" "" "" ""] [emptystr=""] +``` + +If you specify positional delimiters, you must provide all six. + +## Parameters + +| Parameter | Default | Description | +| --- | --- | --- | +| `mvsep` | `OR` | Separator between values from a multivalue field. | +| `maxresults` | `0` | Maximum input rows to include. `0` means unlimited. | +| `row-prefix` | `(` | Prefix for the complete expression. | +| `column-prefix` | `(` | Prefix for each formatted input row. | +| `column-separator` | `AND` | Separator between fields in one row. | +| `column-end` | `)` | Suffix for each formatted input row. | +| `row-separator` | `OR` | Separator between formatted rows. | +| `row-end` | `)` | Suffix for the complete expression. | +| `emptystr` | `NOT( )` | Result when the input has no formattable fields or values. | + +## Behavior + +- Output contains one string field named `search`. +- Fields are rendered in lexicographic field-name order. +- Null and missing values are omitted. Empty strings are retained. +- Fields whose names begin with `_` are treated as internal fields and omitted. +- The field names `search` and `query` are omitted from their rendered predicates; their values are inserted directly into the expression. +- Field names containing special characters are double quoted. +- Scalar values are converted to strings. Double quotes and backslashes are escaped. +- Multivalue fields produce a parenthesized expression containing one predicate per non-null element. + +## Examples + +### Default formatting + +The following input: + +```text +status method +200 GET +500 POST +``` + +with this command: + +```ppl +source=logs +| fields status, method +| format +``` + +produces: + +```text +( ( method="GET" AND status="200" ) OR ( method="POST" AND status="500" ) ) +``` + +### Custom delimiters and row limit + +```ppl +source=logs +| fields status, method +| format maxresults=2 "[" "[" "&&" "]" "||" "]" +``` + +### Multivalue fields + +```ppl +source=alerts +| fields tags +| format mvsep="OR" "{" "[" "AND" "]" "AND" "}" +``` + +For a `tags` value of `[critical, network]`, the formatted value is: + +```text +{ [ ( tags="critical" OR tags="network" ) ] } +``` + +### Empty input + +```ppl +source=logs +| where status=999 +| fields status +| format emptystr="no matching data" +``` + +## Limitations + +- PPL currently supports `format` only as an explicit pipeline command. It is not implicitly appended to subsearches. +- The generated `search` value is not automatically injected into or executed by an outer query. +- Multivalue element escaping is limited by the backend array-expression support. Avoid quote and backslash characters in multivalue elements when the output will be parsed as a search expression. +- Row collection lowers to a global `ARRAY_AGG` without an aggregate order key. Distributed execution can change row order even when the input contains an upstream `sort`, so stable Splunk row-order parity is not guaranteed. + +## Related commands + +- [fields](fields.md) --- select the fields included in the formatted expression +- [eval](eval.md) --- build values before formatting +- [subquery](subquery.md) --- use the result of one query inside another query diff --git a/docs/user/ppl/index.md b/docs/user/ppl/index.md index 647c4568446..4525c6653a9 100644 --- a/docs/user/ppl/index.md +++ b/docs/user/ppl/index.md @@ -41,6 +41,7 @@ source=accounts | [where command](cmd/where.md) | 1.0 | stable (since 1.0) | Filter the search result using boolean expressions. | | [subquery command](cmd/subquery.md) | 3.0 | experimental (since 3.0) | Embed one PPL query inside another for complex filtering and data retrieval operations. | | [fields command](cmd/fields.md) | 1.0 | stable (since 1.0) | Keep or remove fields from the search result. | +| [format command](cmd/format.md) | 3.8 | experimental (since 3.8) | Collapse rows and fields into a single search-expression string. | | [rename command](cmd/rename.md) | 1.0 | stable (since 1.0) | Rename one or more fields in the search result. | | [eval command](cmd/eval.md) | 1.0 | stable (since 1.0) | Evaluate an expression and append the result to the search result. | | [foreach command](cmd/foreach.md) | 3.8 | experimental (since 3.8) | Run a templated evaluation for each selected field or collection element. | 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 6d0c88bf773..c3b3872228b 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 @@ -39,6 +39,7 @@ CalciteFieldFormatCommandIT.class, CalciteForeachCommandIT.class, ForeachFieldJsonIT.class, + CalciteFormatCommandIT.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 04dc2b0e74b..08386b9789d 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 @@ -2952,6 +2952,20 @@ public void testForeachExplain() throws IOException { && logical.contains(", 2)")); } + @Test + public void testFormatExplain() throws IOException { + String query = + StringUtils.format( + "source=%s | where account_number < 3 | fields firstname, account_number | format", + TEST_INDEX_BANK); + String logical = logicalPlan(explainQueryYaml(query)).toLowerCase(java.util.Locale.ROOT); + Assert.assertTrue( + "Expected format to lower to a global ARRAY_AGG and a search projection", + logical.contains("array_agg") + && logical.contains("array_join") + && logical.contains("search=")); + } + /** * 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/CalciteFormatCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java new file mode 100644 index 00000000000..59cea652bb0 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java @@ -0,0 +1,64 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.remote; + +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK; +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.sql.ppl.PPLIntegTestCase; + +public class CalciteFormatCommandIT extends PPLIntegTestCase { + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); + loadIndex(Index.BANK); + } + + @Test + public void testDefaultFormat() throws IOException { + JSONObject result = + executeQuery( + "source=" + + TEST_INDEX_BANK + + " | where account_number=1 | fields firstname, account_number | format"); + + verifySchema(result, schema("search", null, "string")); + verifyDataRows(result, rows("( ( account_number=\"1\" AND firstname=\"Amber JOHnny\" ) )")); + } + + @Test + public void testMultivalueAndCustomDelimiters() throws IOException { + JSONObject result = + executeQuery( + "source=" + + TEST_INDEX_BANK + + " | where account_number=1 | eval names=array(firstname, lastname) " + + "| fields names | format mvsep='OR' '{' '[' 'AND' ']' 'OR' '}'"); + + verifyDataRows( + result, rows("{ [ ( names=\"Amber JOHnny\" OR names=\"Duke Willmington\" ) ] }")); + } + + @Test + public void testEmptyResultFallback() throws IOException { + JSONObject result = + executeQuery( + "source=" + + TEST_INDEX_BANK + + " | where account_number < 0 | fields firstname " + + "| format emptystr='no matching data'"); + + verifyDataRows(result, rows("no matching data")); + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/NewAddedCommandsIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/NewAddedCommandsIT.java index a32dd9fb990..88a286317e4 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/NewAddedCommandsIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/NewAddedCommandsIT.java @@ -340,6 +340,21 @@ public void testMakeResults() throws IOException { } } + @Test + public void testFormatUnsupportedInV2() throws IOException { + JSONObject result; + try { + result = + executeQuery( + String.format( + "source=%s | where account_number=1 | fields firstname, account_number | format", + TEST_INDEX_BANK)); + } catch (ResponseException e) { + result = new JSONObject(TestUtils.getResponseBody(e.getResponse())); + } + verifyQuery(result); + } + @Test public void testMvExpandCommandBasicExpansion() throws IOException { JSONObject result; diff --git a/ppl/src/main/antlr/OpenSearchPPLLexer.g4 b/ppl/src/main/antlr/OpenSearchPPLLexer.g4 index fb072ae134f..9fc72adfacb 100644 --- a/ppl/src/main/antlr/OpenSearchPPLLexer.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLLexer.g4 @@ -32,6 +32,7 @@ SORT: 'SORT'; EVAL: 'EVAL'; FOREACH: 'FOREACH'; FIELDFORMAT: 'FIELDFORMAT'; +FORMAT: 'FORMAT'; HEAD: 'HEAD'; BIN: 'BIN'; TOP: 'TOP'; @@ -97,6 +98,9 @@ APPENDPIPE: 'APPENDPIPE'; COLUMN_NAME: 'COLUMN_NAME'; MVCOMBINE: 'MVCOMBINE'; NOMV: 'NOMV'; +MVSEP: 'MVSEP'; +MAXRESULTS: 'MAXRESULTS'; +EMPTYSTR: 'EMPTYSTR'; // EDGE_CLAUSE matches the entire edge clause pattern in graphLookup command. // This allows EDGE to be used as a field name elsewhere (e.g., eval edge=1). diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4 index a23a314537d..50dfd64b8d7 100644 --- a/ppl/src/main/antlr/OpenSearchPPLParser.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLParser.g4 @@ -97,6 +97,7 @@ commands | replaceCommand | mvcombineCommand | fieldformatCommand + | formatCommand | nomvCommand | graphLookupCommand | xyseriesCommand @@ -121,6 +122,7 @@ commandName | EVAL | FOREACH | FIELDFORMAT + | FORMAT | HEAD | BIN | TOP @@ -338,6 +340,20 @@ reverseCommand : REVERSE ; +formatCommand + : FORMAT formatOption* formatDelimiters? formatOption* + ; + +formatOption + : MVSEP EQUAL stringLiteral + | MAXRESULTS EQUAL integerLiteral + | EMPTYSTR EQUAL stringLiteral + ; + +formatDelimiters + : stringLiteral stringLiteral stringLiteral stringLiteral stringLiteral stringLiteral + ; + chartCommand : CHART chartOptions* statsAggTerm (OVER rowSplit)? (BY columnSplit)? chartOptions* | CHART chartOptions* statsAggTerm BY rowSplit (COMMA)? columnSplit chartOptions* @@ -1807,6 +1823,9 @@ searchableKeyWord | COUNTFIELD | SHOWCOUNT | MAXOUT + | MVSEP + | MAXRESULTS + | EMPTYSTR | PATH | INPUT | OUTPUT 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 efdbf26a205..3670918fdf6 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 @@ -90,6 +90,7 @@ 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.Format; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.GraphLookup.Direction; import org.opensearch.sql.ast.tree.Head; @@ -802,6 +803,51 @@ public UnresolvedPlan visitReverseCommand(OpenSearchPPLParser.ReverseCommandCont return new Reverse(); } + /** Format command. */ + @Override + public UnresolvedPlan visitFormatCommand(OpenSearchPPLParser.FormatCommandContext ctx) { + String mvSeparator = Format.DEFAULT_MV_SEPARATOR; + int maxResults = Format.DEFAULT_MAX_RESULTS; + String emptyString = Format.DEFAULT_EMPTY_STRING; + + for (OpenSearchPPLParser.FormatOptionContext option : ctx.formatOption()) { + if (option.MVSEP() != null) { + mvSeparator = StringUtils.unquoteText(getTextInQuery(option.stringLiteral())); + } else if (option.MAXRESULTS() != null) { + maxResults = Integer.parseInt(option.integerLiteral().getText()); + if (maxResults < 0) { + throw new SemanticCheckException("format maxresults must be non-negative"); + } + } else { + emptyString = StringUtils.unquoteText(getTextInQuery(option.stringLiteral())); + } + } + + List delimiters = + ctx.formatDelimiters() == null + ? List.of( + Format.DEFAULT_ROW_PREFIX, + Format.DEFAULT_COLUMN_PREFIX, + Format.DEFAULT_COLUMN_SEPARATOR, + Format.DEFAULT_COLUMN_END, + Format.DEFAULT_ROW_SEPARATOR, + Format.DEFAULT_ROW_END) + : ctx.formatDelimiters().stringLiteral().stream() + .map(literal -> StringUtils.unquoteText(getTextInQuery(literal))) + .toList(); + + return new Format( + mvSeparator, + maxResults, + delimiters.get(0), + delimiters.get(1), + delimiters.get(2), + delimiters.get(3), + delimiters.get(4), + delimiters.get(5), + emptyString); + } + /** Transpose command. */ @Override public UnresolvedPlan visitTransposeCommand(OpenSearchPPLParser.TransposeCommandContext ctx) { diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java index 398361fcea7..2d9cfc6dea5 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizer.java @@ -79,6 +79,7 @@ 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.Format; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.Head; import org.opensearch.sql.ast.tree.Join; @@ -612,6 +613,24 @@ public String visitForeach(Foreach node, String context) { return command.append(" [ eval ").append(evalClauses).append(" ]").toString(); } + @Override + public String visitFormat(Format node, String context) { + String child = node.getChild().getFirst().accept(this, context); + return StringUtils.format( + "%s | format mvsep=\"%s\" maxresults=%s \"%s\" \"%s\" \"%s\" \"%s\" \"%s\"" + + " \"%s\" emptystr=\"%s\"", + child, + MASK_LITERAL, + MASK_LITERAL, + MASK_LITERAL, + MASK_LITERAL, + MASK_LITERAL, + MASK_LITERAL, + MASK_LITERAL, + MASK_LITERAL, + MASK_LITERAL); + } + /** Build {@link LogicalSort}. */ @Override public String visitSort(Sort node, String context) { diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java new file mode 100644 index 00000000000..1cc96ff33ba --- /dev/null +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java @@ -0,0 +1,153 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.calcite; + +import static org.junit.Assert.assertThrows; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.test.CalciteAssert; +import org.junit.Test; + +public class CalcitePPLFormatTest extends CalcitePPLAbstractTest { + + public CalcitePPLFormatTest() { + super(CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL); + } + + @Test + public void testDefaultFormat() { + withPPLQuery("source=EMP | fields ENAME, JOB | head 2 | format") + .expectResult( + "search=( ( ENAME=\"SMITH\" AND JOB=\"CLERK\" ) OR ( ENAME=\"ALLEN\" AND" + + " JOB=\"SALESMAN\" ) )\n"); + } + + @Test + public void testLogicalAndSparkSqlPlan() { + RelNode root = getRelNode("source=EMP | fields ENAME | format maxresults=1"); + verifyLogical( + root, + "LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN($0, ' OR ':VARCHAR)), 0), ||(||('(" + + " ':VARCHAR, ARRAY_JOIN($0, ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT( )':VARCHAR)])\n" + + " LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0) FILTER $1])\n" + + " LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS" + + " NOT NULL($0), ||(||('ENAME=\"':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL," + + " '\\':VARCHAR, '\\\\':VARCHAR), '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR)," + + " null:VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR," + + " ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('ENAME=\"':VARCHAR," + + " REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR)," + + " '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR))," + + " ' )':VARCHAR), null:VARCHAR)], $f1=[IS NOT" + + " NULL(CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0)," + + " ||(||('ENAME=\"':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR," + + " '\\\\':VARCHAR), '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR), null:VARCHAR))), '" + + " AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT" + + " NULL($0), ||(||('ENAME=\"':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL," + + " '\\':VARCHAR, '\\\\':VARCHAR), '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR)," + + " null:VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR))])\n" + + " LogicalSort(fetch=[1])\n" + + " LogicalProject(ENAME=[$1])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"); + verifyPPLToSparkSQL( + root, + "SELECT CASE WHEN CHAR_LENGTH(ARRAY_JOIN(ARRAY_AGG(`__format_row`) FILTER (WHERE `$f1`)," + + " ' OR ')) > 0 THEN '( ' || ARRAY_JOIN(ARRAY_AGG(`__format_row`) FILTER (WHERE" + + " `$f1`), ' OR ') || ' )' ELSE 'NOT( )' END `search`\n" + + "FROM (SELECT CASE WHEN CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME`" + + " IS NOT NULL THEN 'ENAME=\"' || REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\'," + + " '\\\\'), '\"', '\\\"') || '\"' ELSE NULL END)), ' AND ')) > 0 THEN '( ' ||" + + " ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME` IS NOT NULL THEN 'ENAME=\"' ||" + + " REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\', '\\\\'), '\"', '\\\"') || '\"'" + + " ELSE NULL END)), ' AND ') || ' )' ELSE NULL END `__format_row`, CASE WHEN" + + " CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME` IS NOT NULL THEN" + + " 'ENAME=\"' || REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\', '\\\\'), '\"'," + + " '\\\"') || '\"' ELSE NULL END)), ' AND ')) > 0 THEN '( ' ||" + + " ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME` IS NOT NULL THEN 'ENAME=\"' ||" + + " REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\', '\\\\'), '\"', '\\\"') || '\"'" + + " ELSE NULL END)), ' AND ') || ' )' ELSE NULL END IS NOT NULL `$f1`\n" + + "FROM `scott`.`EMP`\n" + + "LIMIT 1) `t1`"); + } + + @Test + public void testCustomDelimitersAndMaxResults() { + withPPLQuery( + "source=EMP | fields ENAME, JOB | head 2 " + + "| format maxresults=1 \"[\" \"[\" \"&&\" \"]\" \"||\" \"]\"") + .expectResult("search=[ [ ENAME=\"SMITH\" && JOB=\"CLERK\" ] ]\n"); + } + + @Test + public void testMultivalueFormat() { + withPPLQuery( + "source=EMP | head 1 | eval tags=array(\"critical\", \"network\") " + + "| fields tags | format mvsep=\"mvseparator\" " + + "\"{\" \"[\" \"AND\" \"]\" \"AND\" \"}\"") + .expectResult("search={ [ ( tags=\"critical\" mvseparator tags=\"network\" ) ] }\n"); + } + + @Test + public void testEmptyMultivalueUsesFallback() { + withPPLQuery( + "source=EMP | head 1 | eval tags=array() | fields tags " + + "| format emptystr=\"empty\"") + .expectResult("search=empty\n"); + } + + @Test + public void testEmptyResultFallback() { + withPPLQuery( + "source=EMP | where EMPNO < 0 | fields ENAME | format emptystr=\"no matching data\"") + .expectResult("search=no matching data\n"); + } + + @Test + public void testNullOnlyRowUsesFallback() { + withPPLQuery("source=EMP | where EMPNO=7369 | fields COMM | format emptystr=\"empty\"") + .expectResult("search=empty\n"); + } + + @Test + public void testQuotesInValuesAreEscaped() { + withPPLQuery("source=EMP | head 1 | eval message='say \\\"hi\\\"' | fields message | format") + .expectResult("search=( ( message=\"say \\\"hi\\\"\" ) )\n"); + } + + @Test + public void testBackslashesInValuesAreEscaped() { + withPPLQuery("source=EMP | head 1 | eval path='a\\\\b' | fields path | format") + .expectResult("search=( ( path=\"a\\\\b\" ) )\n"); + } + + @Test + public void testInternalFieldsAreIgnored() { + withPPLQuery( + "source=EMP | head 1 | eval _private='hidden', visible='shown' " + + "| fields _private, visible | format") + .expectResult("search=( ( visible=\"shown\" ) )\n"); + } + + @Test + public void testSearchAndQueryFieldsDropFieldName() { + withPPLQuery( + "source=EMP | head 1 | eval search='status=200', query='method=GET', a='x' " + + "| fields search, query, a | format") + .expectResult("search=( ( a=\"x\" AND \"method=GET\" AND \"status=200\" ) )\n"); + } + + @Test + public void testSpecialCharactersInFieldNameAreQuoted() { + withPPLQuery("source=EMP | head 1 | eval a.b='x' | fields a.b | format") + .expectResult("search=( ( \"a.b\"=\"x\" ) )\n"); + } + + @Test + public void testAllSixDelimitersAreRequired() { + assertThrows( + RuntimeException.class, + () -> getRelNode("source=EMP | fields ENAME | format \"[\" \"[\" \"AND\"")); + } +} diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java index 0650b4d1e4c..95c6f97b5d6 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java @@ -77,6 +77,7 @@ import org.opensearch.sql.ast.expression.SpanUnit; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Chart; +import org.opensearch.sql.ast.tree.Format; import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.Join; import org.opensearch.sql.ast.tree.Kmeans; @@ -182,6 +183,31 @@ public void testSearchCommandWithoutSearch() { "source=t | where a=1", filter(relation("t"), compare("=", field("a"), intLiteral(1)))); } + @Test + public void testFormatCommandDefaults() { + assertEqual( + "source=t | format", + new Format( + Format.DEFAULT_MV_SEPARATOR, + Format.DEFAULT_MAX_RESULTS, + Format.DEFAULT_ROW_PREFIX, + Format.DEFAULT_COLUMN_PREFIX, + Format.DEFAULT_COLUMN_SEPARATOR, + Format.DEFAULT_COLUMN_END, + Format.DEFAULT_ROW_SEPARATOR, + Format.DEFAULT_ROW_END, + Format.DEFAULT_EMPTY_STRING) + .attach(relation("t"))); + } + + @Test + public void testFormatCommandOptionsAndDelimiters() { + assertEqual( + "source=t | format mvsep=\"MV\" maxresults=2 emptystr=\"empty\" " + + "\"[\" \"{\" \"&&\" \"}\" \"||\" \"]\"", + new Format("MV", 2, "[", "{", "&&", "}", "||", "]", "empty").attach(relation("t"))); + } + @Test public void testSearchCommandWithFilterBeforeSource() { assertEqual("search a=1 source=t", search(relation("t"), "a:1")); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java index 63f7dea81d7..3ccdc2726de 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java @@ -1268,6 +1268,16 @@ public void testForeachJsonArrayCommandMasksLiteralTarget() { "source=t | foreach mode=json_array '[1,2,3]' [ eval total = total + <> ]")); } + @Test + public void testFormatCommand() { + assertEquals( + "source=table | format mvsep=\"***\" maxresults=*** \"***\" \"***\" \"***\" \"***\"" + + " \"***\" \"***\" emptystr=\"***\"", + anonymize( + "source=sensitive | format mvsep=\"SECRET\" maxresults=5 \"[\" \"[\" \"AND\"" + + " \"]\" \"OR\" \"]\" emptystr=\"PRIVATE\"")); + } + @Test public void testUnion() { assertEquals( From c1fcc05036f30c854d2f16a72e415ffcc8c51ac5 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 23 Jul 2026 11:08:54 +0800 Subject: [PATCH 02/10] Support implicit format subsearches Signed-off-by: Songkan Tang --- .../org/opensearch/sql/analysis/Analyzer.java | 3 + .../sql/ast/expression/SearchSubquery.java | 39 ++++ .../subquery/RuntimeSearchScalarSubquery.java | 16 ++ .../org/opensearch/sql/ast/tree/Format.java | 6 +- .../org/opensearch/sql/ast/tree/Search.java | 24 ++- .../sql/calcite/CalciteRelNodeVisitor.java | 70 ++++++- .../sql/calcite/CalciteRexNodeVisitor.java | 16 ++ .../sql/calcite/DynamicSearchPlanBinder.java | 192 ++++++++++++++++++ .../opensearch/sql/calcite/FormatPlanner.java | 153 ++++++++++++-- .../sql/calcite/SearchPredicateCompiler.java | 13 ++ .../sql/executor/DynamicSearchExecutor.java | 101 +++++++++ .../opensearch/sql/executor/QueryService.java | 4 +- .../executor/analytics/TimewrapSignals.java | 7 + .../utils/TimewrapSignalsLeakTest.java | 17 ++ docs/user/ppl/cmd/format.md | 37 +++- .../analytics/AnalyticsEngineProfileIT.java | 44 ++++ .../remote/CalciteFormatCommandIT.java | 27 +++ .../CalciteImplicitFormatCommandIT.java | 99 +++++++++ .../plugin/config/OpenSearchPluginModule.java | 1 + .../plugin/rest/RestUnifiedQueryAction.java | 105 ++++++++-- ppl/src/main/antlr/OpenSearchPPLParser.g4 | 5 + .../sql/ppl/antlr/PPLSyntaxParser.java | 7 + .../opensearch/sql/ppl/parser/AstBuilder.java | 8 +- .../sql/ppl/parser/AstExpressionBuilder.java | 6 + .../parser/PPLSearchPredicateCompiler.java | 34 ++++ .../sql/ppl/antlr/PPLSyntaxParserTest.java | 8 + .../sql/ppl/calcite/CalcitePPLFormatTest.java | 62 ++++-- .../sql/ppl/calcite/CalcitePPLSearchTest.java | 161 +++++++++++++++ .../ppl/parser/AstSearchExpressionTest.java | 10 + .../PPLSearchPredicateCompilerTest.java | 44 ++++ .../ppl/utils/PPLQueryDataAnonymizerTest.java | 7 + 31 files changed, 1252 insertions(+), 74 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java create mode 100644 core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java create mode 100644 core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java create mode 100644 core/src/main/java/org/opensearch/sql/calcite/SearchPredicateCompiler.java create mode 100644 core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java create mode 100644 ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java create mode 100644 ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java 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 f9373bd238a..1c6c48a30af 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -300,6 +300,9 @@ public LogicalPlan visitLimit(Limit node, AnalysisContext context) { @Override public LogicalPlan visitSearch(Search node, AnalysisContext context) { + if (node.hasImplicitSubquery()) { + throw getOnlyForCalciteException("Implicit format subsearch"); + } LogicalPlan child = node.getChild().get(0).accept(this, context); Function queryStringFunc = AstDSL.function( diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java b/core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java new file mode 100644 index 00000000000..c7483281eda --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java @@ -0,0 +1,39 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ast.expression; + +import java.util.List; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import lombok.ToString; +import org.opensearch.sql.ast.tree.UnresolvedPlan; + +/** A subsearch used as a predicate term by the parent {@code search} command. */ +@Getter +@RequiredArgsConstructor +@EqualsAndHashCode(callSuper = false) +@ToString +public class SearchSubquery extends SearchExpression { + + private final UnresolvedPlan query; + + @Override + public String toQueryString() { + throw new IllegalStateException( + "An implicit search subquery must be bound before it can be converted to query_string"); + } + + @Override + public String toAnonymizedString() { + return "[ subsearch ]"; + } + + @Override + public List getChild() { + return List.of(); + } +} diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java b/core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java new file mode 100644 index 00000000000..816b8c11e13 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java @@ -0,0 +1,16 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ast.expression.subquery; + +import org.opensearch.sql.ast.tree.UnresolvedPlan; + +/** A runtime search scalar subquery that is planned independently from its parent query. */ +public class RuntimeSearchScalarSubquery extends ScalarSubquery { + + public RuntimeSearchScalarSubquery(UnresolvedPlan query) { + super(query); + } +} diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Format.java b/core/src/main/java/org/opensearch/sql/ast/tree/Format.java index 7f2191b6b80..318d57b6862 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Format.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Format.java @@ -30,7 +30,7 @@ public class Format extends UnresolvedPlan { public static final String DEFAULT_COLUMN_END = ")"; public static final String DEFAULT_ROW_SEPARATOR = "OR"; public static final String DEFAULT_ROW_END = ")"; - public static final String DEFAULT_EMPTY_STRING = "NOT( )"; + public static final String DEFAULT_EMPTY_STRING = "NOT ()"; private final String mvSeparator; private final int maxResults; @@ -41,6 +41,10 @@ public class Format extends UnresolvedPlan { private final String rowSeparator; private final String rowEnd; private final String emptyString; + + /** Whether the result is consumed by an enclosing search rather than returned to the user. */ + private boolean implicit; + private UnresolvedPlan child; @Override diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Search.java b/core/src/main/java/org/opensearch/sql/ast/tree/Search.java index 0a705f95c41..d574cf03741 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Search.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Search.java @@ -26,7 +26,7 @@ public class Search extends UnresolvedPlan { @EqualsAndHashCode.Include private final UnresolvedPlan child; - @EqualsAndHashCode.Include private final String queryString; + @EqualsAndHashCode.Include private final @Nullable String queryString; // Currently it's only for anonymizer private final @Nullable SearchExpression originalExpression; @@ -35,6 +35,28 @@ public Search(UnresolvedPlan child, String queryString) { this(child, queryString, null); } + /** Creates a search whose query string may need runtime subquery binding. */ + public static Search fromExpression(UnresolvedPlan child, SearchExpression originalExpression) { + return new Search( + child, + containsSubquery(originalExpression) ? null : originalExpression.toQueryString(), + originalExpression); + } + + public boolean hasImplicitSubquery() { + return queryString == null; + } + + private static boolean containsSubquery(SearchExpression expression) { + if (expression instanceof org.opensearch.sql.ast.expression.SearchSubquery) { + return true; + } + return expression.getChild().stream() + .filter(SearchExpression.class::isInstance) + .map(SearchExpression.class::cast) + .anyMatch(Search::containsSubquery); + } + @Override public List getChild() { return ImmutableList.of(child); 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 77ff4b1df9e..2f31d59eb65 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -107,12 +107,19 @@ import org.opensearch.sql.ast.expression.PatternMethod; import org.opensearch.sql.ast.expression.PatternMode; import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.SearchAnd; +import org.opensearch.sql.ast.expression.SearchExpression; +import org.opensearch.sql.ast.expression.SearchGroup; +import org.opensearch.sql.ast.expression.SearchNot; +import org.opensearch.sql.ast.expression.SearchOr; +import org.opensearch.sql.ast.expression.SearchSubquery; import org.opensearch.sql.ast.expression.Span; import org.opensearch.sql.ast.expression.SpanUnit; import org.opensearch.sql.ast.expression.UnresolvedExpression; import org.opensearch.sql.ast.expression.WindowFrame; import org.opensearch.sql.ast.expression.WindowFrame.FrameType; import org.opensearch.sql.ast.expression.WindowFunction; +import org.opensearch.sql.ast.expression.subquery.RuntimeSearchScalarSubquery; import org.opensearch.sql.ast.expression.subquery.SubqueryExpression; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.AddColTotals; @@ -330,17 +337,74 @@ private RelBuilder scan(RelOptTable tableSchema, CalcitePlanContext context) { public RelNode visitSearch(Search node, CalcitePlanContext context) { // Visit the Relation child to get the scan node.getChild().get(0).accept(this, context); + UnresolvedExpression query = + node.hasImplicitSubquery() + ? buildRuntimeSearchQuery(node.getOriginalExpression(), context) + : AstDSL.stringLiteral(node.getQueryString()); // Create query_string function Function queryStringFunc = - AstDSL.function( - "query_string", - AstDSL.unresolvedArg("query", AstDSL.stringLiteral(node.getQueryString()))); + AstDSL.function("query_string", AstDSL.unresolvedArg("query", query)); RexNode queryStringRex = rexVisitor.analyze(queryStringFunc, context); context.relBuilder.filter(queryStringRex); return context.relBuilder.peek(); } + /** Builds a query_string operand containing standard scalar subqueries for implicit format. */ + private UnresolvedExpression buildRuntimeSearchQuery( + SearchExpression expression, CalcitePlanContext context) { + if (expression instanceof SearchSubquery subquery) { + int maxResults = Math.max(0, context.sysLimit.subsearchLimit()); + Format format = + new Format( + Format.DEFAULT_MV_SEPARATOR, + maxResults, + Format.DEFAULT_ROW_PREFIX, + Format.DEFAULT_COLUMN_PREFIX, + Format.DEFAULT_COLUMN_SEPARATOR, + Format.DEFAULT_COLUMN_END, + Format.DEFAULT_ROW_SEPARATOR, + Format.DEFAULT_ROW_END, + Format.DEFAULT_EMPTY_STRING); + format.setImplicit(true); + format.attach(subquery.getQuery()); + return new RuntimeSearchScalarSubquery(format); + } + if (expression instanceof SearchGroup group) { + return concatSearch("(", buildRuntimeSearchQuery(group.getExpression(), context), ")"); + } + if (expression instanceof SearchNot not) { + return concatSearch("NOT(", buildRuntimeSearchQuery(not.getExpression(), context), ")"); + } + if (expression instanceof SearchAnd and) { + return concatSearch( + buildRuntimeSearchQuery(and.getLeft(), context), + " AND ", + buildRuntimeSearchQuery(and.getRight(), context)); + } + if (expression instanceof SearchOr or) { + return concatSearch( + buildRuntimeSearchQuery(or.getLeft(), context), + " OR ", + buildRuntimeSearchQuery(or.getRight(), context)); + } + return AstDSL.stringLiteral(expression.toQueryString()); + } + + private UnresolvedExpression concatSearch(Object... parts) { + UnresolvedExpression result = toSearchExpression(parts[0]); + for (int i = 1; i < parts.length; i++) { + result = AstDSL.function("concat", result, toSearchExpression(parts[i])); + } + return result; + } + + private UnresolvedExpression toSearchExpression(Object part) { + return part instanceof UnresolvedExpression expression + ? expression + : AstDSL.stringLiteral(part.toString()); + } + @Override public RelNode visitFilter(Filter 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 1bf1e217b51..7b5d761eb33 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -12,6 +12,7 @@ import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY; import java.math.BigDecimal; +import java.sql.SQLException; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -34,6 +35,7 @@ import org.apache.calcite.rex.RexLambdaRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexSubQuery; import org.apache.calcite.sql.SqlIntervalQualifier; import org.apache.calcite.sql.SqlOperator; import org.apache.calcite.sql.fun.SqlStdOperatorTable; @@ -77,6 +79,7 @@ import org.opensearch.sql.ast.expression.Xor; import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; import org.opensearch.sql.ast.expression.subquery.InSubquery; +import org.opensearch.sql.ast.expression.subquery.RuntimeSearchScalarSubquery; import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; import org.opensearch.sql.ast.expression.subquery.SubqueryExpression; import org.opensearch.sql.ast.tree.Sort.SortOption; @@ -865,6 +868,19 @@ public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) { @Override public RexNode visitScalarSubquery(ScalarSubquery node, CalcitePlanContext context) { + if (node instanceof RuntimeSearchScalarSubquery) { + CalcitePlanContext initContext = + CalcitePlanContext.create(context.config, context.sysLimit, context.queryType); + try { + return RexSubQuery.scalar(resolveSubqueryPlan(node.getQuery(), node, initContext)); + } finally { + try { + initContext.connection.close(); + } catch (SQLException e) { + throw new RuntimeException("Failed to close runtime search planning context", e); + } + } + } return context.relBuilder.scalarQuery( b -> { UnresolvedPlan subquery = node.getQuery(); diff --git a/core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java b/core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java new file mode 100644 index 00000000000..f011545a1f0 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java @@ -0,0 +1,192 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite; + +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelVisitor; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexShuttle; +import org.apache.calcite.rex.RexSubQuery; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.SqlTypeName; + +/** Finds and binds scalar subqueries used as runtime {@code search} predicates. */ +public final class DynamicSearchPlanBinder { + + private DynamicSearchPlanBinder() {} + + /** Returns the first unbound scalar subquery beneath a query_string query argument. */ + public static Optional find(RelNode plan) { + AtomicReference found = new AtomicReference<>(); + RexShuttle finder = + new RexShuttle() { + @Override + public RexNode visitCall(RexCall call) { + if (found.get() == null && isQueryString(call)) { + RexNode query = queryArgument(call); + if (query != null) { + query.accept( + new RexShuttle() { + @Override + public RexNode visitSubQuery(RexSubQuery subQuery) { + found.compareAndSet(null, subQuery); + return subQuery; + } + }); + } + } + return found.get() == null ? super.visitCall(call) : call; + } + }; + + new RelVisitor() { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) { + if (found.get() == null) { + node.accept(finder); + super.visit(node, ordinal, parent); + } + } + }.go(plan); + return Optional.ofNullable(found.get()); + } + + /** Replaces one runtime scalar query and folds the surrounding string concatenation. */ + public static RelNode bind(RelNode plan, RexSubQuery target, String compiledQueryString) { + RexBuilder rexBuilder = plan.getCluster().getRexBuilder(); + RexShuttle queryArgumentBinder = + new RexShuttle() { + @Override + public RexNode visitSubQuery(RexSubQuery subQuery) { + if (subQuery == target) { + return stringLiteral(rexBuilder, compiledQueryString); + } + return subQuery; + } + + @Override + public RexNode visitCall(RexCall call) { + RexNode visited = super.visitCall(call); + if (visited instanceof RexCall visitedCall) { + String constant = constantString(visitedCall); + if (constant != null) { + return stringLiteral(rexBuilder, constant); + } + } + return visited; + } + }; + RexShuttle binder = + new RexShuttle() { + @Override + public RexNode visitCall(RexCall call) { + if (isQueryString(call)) { + RexNode visited = call.accept(queryArgumentBinder); + return visited instanceof RexCall visitedCall + ? foldQueryStringArgument(visitedCall, rexBuilder) + : visited; + } + return super.visitCall(call); + } + }; + return rewrite(plan, binder); + } + + private static RelNode rewrite(RelNode node, RexShuttle binder) { + List inputs = node.getInputs(); + List rewrittenInputs = new ArrayList<>(inputs.size()); + boolean changed = false; + for (RelNode input : inputs) { + RelNode rewritten = rewrite(input, binder); + rewrittenInputs.add(rewritten); + changed |= rewritten != input; + } + RelNode withInputs = changed ? node.copy(node.getTraitSet(), rewrittenInputs) : node; + return withInputs.accept(binder); + } + + private static String constantString(RexCall call) { + if ((call.getKind() == SqlKind.CAST || call.getOperator().getName().equalsIgnoreCase("cast")) + && call.getOperands().size() == 1) { + return literalString(call.getOperands().getFirst()); + } + String name = call.getOperator().getName().toLowerCase(Locale.ROOT); + if (!(name.equals("concat") || name.equals("||"))) { + return null; + } + StringBuilder result = new StringBuilder(); + for (RexNode operand : call.getOperands()) { + String value = literalString(operand); + if (value == null) { + return null; + } + result.append(value); + } + return result.toString(); + } + + private static RexNode foldQueryStringArgument(RexCall queryString, RexBuilder rexBuilder) { + List operands = new ArrayList<>(queryString.getOperands()); + for (int i = 0; i < operands.size(); i++) { + RexNode operand = operands.get(i); + if (operand instanceof RexCall map && map.getOperands().size() >= 2) { + RexNode key = map.getOperands().getFirst(); + if (key instanceof RexLiteral literal + && "query".equalsIgnoreCase(literal.getValueAs(String.class))) { + RexNode value = map.getOperands().get(1); + String constant = literalString(value); + if (constant != null && !(value instanceof RexLiteral)) { + List mapOperands = new ArrayList<>(map.getOperands()); + mapOperands.set(1, stringLiteral(rexBuilder, constant)); + operands.set(i, map.clone(map.getType(), mapOperands)); + return queryString.clone(queryString.getType(), operands); + } + } + } + } + return queryString; + } + + private static RexNode stringLiteral(RexBuilder rexBuilder, String value) { + return rexBuilder.makeLiteral( + value, rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR), true); + } + + private static String literalString(RexNode node) { + if (node instanceof RexLiteral literal) { + return literal.getValueAs(String.class); + } + if (node instanceof RexCall call) { + return constantString(call); + } + return null; + } + + private static boolean isQueryString(RexCall call) { + return call.getOperator().getName().equalsIgnoreCase("query_string"); + } + + private static RexNode queryArgument(RexCall queryString) { + for (RexNode operand : queryString.getOperands()) { + if (operand instanceof RexCall map && map.getOperands().size() >= 2) { + RexNode key = map.getOperands().getFirst(); + if (key instanceof RexLiteral literal + && "query".equalsIgnoreCase(literal.getValueAs(String.class))) { + return map.getOperands().get(1); + } + } + } + return null; + } +} diff --git a/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java index 8ef67ec1a7d..7f793ba80b2 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java +++ b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java @@ -5,11 +5,16 @@ package org.opensearch.sql.calcite; +import java.util.ArrayList; import java.util.Comparator; import java.util.List; +import java.util.Optional; +import org.apache.calcite.rel.RelCollation; +import org.apache.calcite.rel.RelFieldCollation; 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.RexLambdaRef; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.fun.SqlLibraryOperators; import org.apache.calcite.sql.fun.SqlStdOperatorTable; @@ -17,6 +22,7 @@ import org.apache.calcite.sql.type.SqlTypeUtil; import org.apache.calcite.tools.RelBuilder; import org.opensearch.sql.ast.tree.Format; +import org.opensearch.sql.expression.function.PPLBuiltinOperators; /** Lowers the PPL {@code format} command to Calcite projects and a global aggregation. */ public class FormatPlanner { @@ -39,8 +45,79 @@ public RelNode plan(Format node, CalcitePlanContext context) { return builder.peek(); } + Optional scalarSearchField = + fields.stream() + .filter(field -> field.getName().equals("search")) + .filter( + field -> + !SqlTypeUtil.isArray(field.getType()) + && !SqlTypeUtil.isMultiset(field.getType())) + .findFirst(); + if (node.isImplicit() && scalarSearchField.isPresent()) { + return planImplicitSearchField(node, context, fields, scalarSearchField.get()); + } + List formattedFields = fields.stream().map(field -> formatField(field, node, context)).toList(); + RexNode nonEmptyRow = formatRow(formattedFields, node, context); + + aggregateRows(nonEmptyRow, context); + builder.project(List.of(formatAggregatedRows(node, context)), List.of("search"), true); + return builder.peek(); + } + + /** + * At an implicit subsearch boundary, a scalar {@code search} field has special behavior. Only the + * first row participates; a non-null {@code search} value is injected verbatim and suppresses the + * row's other fields. If that value is null, the first row is formatted normally without the + * {@code search} field. + */ + private RelNode planImplicitSearchField( + Format node, + CalcitePlanContext context, + List fields, + RelDataTypeField searchField) { + RelBuilder builder = context.relBuilder; + builder.limit(0, 1); + + RexNode rawSearch = builder.cast(builder.field(searchField.getIndex()), SqlTypeName.VARCHAR); + List fallbackFields = + fields.stream() + .filter(field -> field != searchField) + .map(field -> formatField(field, node, context)) + .toList(); + RexNode fallbackRow = formatRow(fallbackFields, node, context); + + builder.project(List.of(rawSearch, fallbackRow), List.of("__raw_search", "__format_row")); + RexNode rawSearchRef = builder.field("__raw_search"); + RexNode rowRef = builder.field("__format_row"); + builder.aggregate( + builder.groupKey(), + builder.aggregateCall(SqlStdOperatorTable.COUNT, rawSearchRef).as("__raw_search_count"), + builder.aggregateCall(SqlStdOperatorTable.MAX, rawSearchRef).as("__raw_search_value"), + builder.aggregateCall(SqlLibraryOperators.ARRAY_AGG, rowRef).as("__format_rows")); + + RexNode hasRawSearch = + builder.call( + SqlStdOperatorTable.GREATER_THAN, + builder.field("__raw_search_count"), + builder.literal(0)); + RexNode result = + builder.call( + SqlStdOperatorTable.CASE, + hasRawSearch, + builder.field("__raw_search_value"), + formatAggregatedRows(node, context)); + builder.project(List.of(result), List.of("search"), true); + return builder.peek(); + } + + private RexNode formatRow( + List formattedFields, Format node, CalcitePlanContext context) { + if (formattedFields.isEmpty()) { + RelDataType varchar = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR); + return context.rexBuilder.makeNullLiteral(varchar); + } RexNode joinedFields = arrayJoin( context, compactArray(context, formattedFields), " " + node.getColumnSeparator() + " "); @@ -50,19 +127,59 @@ public RelNode plan(Format node, CalcitePlanContext context) { stringLiteral(node.getColumnPrefix() + " ", context), joinedFields, stringLiteral(" " + node.getColumnEnd(), context)); - RexNode nonEmptyRow = ifNotEmpty(joinedFields, rowExpression, context); + return ifNotEmpty(joinedFields, rowExpression, context); + } - builder.project(List.of(nonEmptyRow), List.of("__format_row")); + private void aggregateRows(RexNode nonEmptyRow, CalcitePlanContext context) { + RelBuilder builder = context.relBuilder; + List ordering = inputOrdering(builder.peek()); + if (ordering.isEmpty()) { + builder.project(List.of(nonEmptyRow), List.of("__format_row")); + } else { + List projections = new ArrayList<>(); + projections.add(nonEmptyRow); + projections.addAll(builder.fields()); + List names = new ArrayList<>(); + names.add("__format_row"); + for (int i = 0; i < projections.size() - 1; i++) { + names.add("__format_order_" + i); + } + builder.project(projections, names, true); + } RexNode rowRef = builder.field("__format_row"); - builder.aggregate( - builder.groupKey(), - builder - .aggregateCall(SqlLibraryOperators.ARRAY_AGG, rowRef) - .filter(builder.isNotNull(rowRef)) - .as("__format_rows")); + RelBuilder.AggCall rows = builder.aggregateCall(SqlLibraryOperators.ARRAY_AGG, rowRef); + if (!ordering.isEmpty()) { + rows = rows.sort(ordering.stream().map(order -> orderExpression(order, builder)).toList()); + } + builder.aggregate(builder.groupKey(), rows.as("__format_rows")); + } + + private List inputOrdering(RelNode input) { + return input.getCluster().getMetadataQuery().collations(input).stream() + .map(RelCollation::getFieldCollations) + .filter(ordering -> !ordering.isEmpty()) + .findFirst() + .orElse(List.of()); + } - RexNode joinedRows = - arrayJoin(context, builder.field("__format_rows"), " " + node.getRowSeparator() + " "); + private RexNode orderExpression(RelFieldCollation order, RelBuilder builder) { + RexNode expression = builder.field(order.getFieldIndex() + 1); + if (order.getDirection().isDescending()) { + expression = builder.desc(expression); + } + if (order.nullDirection == RelFieldCollation.NullDirection.FIRST) { + expression = builder.nullsFirst(expression); + } else if (order.nullDirection == RelFieldCollation.NullDirection.LAST) { + expression = builder.nullsLast(expression); + } + return expression; + } + + private RexNode formatAggregatedRows(Format node, CalcitePlanContext context) { + RelBuilder builder = context.relBuilder; + RexNode nonNullRows = + builder.call(SqlLibraryOperators.ARRAY_COMPACT, builder.field("__format_rows")); + RexNode joinedRows = arrayJoin(context, nonNullRows, " " + node.getRowSeparator() + " "); RexNode formatted = concat( context, @@ -75,8 +192,7 @@ public RelNode plan(Format node, CalcitePlanContext context) { isNotEmpty(joinedRows, context), formatted, stringLiteral(node.getEmptyString(), context)); - builder.project(List.of(result), List.of("search"), true); - return builder.peek(); + return result; } private RexNode formatField(RelDataTypeField field, Format node, CalcitePlanContext context) { @@ -93,6 +209,7 @@ private RexNode formatField(RelDataTypeField field, Format node, CalcitePlanCont RexNode values = builder.call( SqlLibraryOperators.ARRAY_COMPACT, context.rexBuilder.makeCast(varcharArray, value)); + values = escapeArrayValues(values, varchar, context); String repeatedValueSeparator = rawValue ? "\" " + node.getMvSeparator() + " \"" @@ -106,12 +223,12 @@ private RexNode formatField(RelDataTypeField field, Format node, CalcitePlanCont stringLiteral("\" )", context)); hasValue = builder.call(SqlStdOperatorTable.AND, hasValue, isNotEmpty(joinedValues, context)); } else { - RexNode stringValue = escapeValue(builder.cast(value, SqlTypeName.VARCHAR), context); + RexNode stringValue = builder.cast(value, SqlTypeName.VARCHAR); formattedValue = concat( context, stringLiteral(rawValue ? "\"" : fieldName + "=\"", context), - stringValue, + escapeValue(stringValue, context), stringLiteral("\"", context)); } @@ -136,6 +253,14 @@ private RexNode escapeValue(RexNode value, CalcitePlanContext context) { stringLiteral("\\\"", context)); } + private RexNode escapeArrayValues( + RexNode values, RelDataType varchar, CalcitePlanContext context) { + RexLambdaRef element = new RexLambdaRef(0, "element", varchar); + RexNode escapeLambda = + context.rexBuilder.makeLambdaCall(escapeValue(element, context), List.of(element)); + return context.rexBuilder.makeCall(PPLBuiltinOperators.TRANSFORM, values, escapeLambda); + } + private String formatFieldName(String fieldName) { if (fieldName.matches("[A-Za-z][A-Za-z0-9_]*")) { return fieldName; diff --git a/core/src/main/java/org/opensearch/sql/calcite/SearchPredicateCompiler.java b/core/src/main/java/org/opensearch/sql/calcite/SearchPredicateCompiler.java new file mode 100644 index 00000000000..9249f28ffb4 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/calcite/SearchPredicateCompiler.java @@ -0,0 +1,13 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite; + +/** Compiles runtime PPL search predicate text to an OpenSearch query_string expression. */ +@FunctionalInterface +public interface SearchPredicateCompiler { + + String compile(String predicate); +} diff --git a/core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java b/core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java new file mode 100644 index 00000000000..030a2cad7e5 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java @@ -0,0 +1,101 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.executor; + +import java.util.List; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rex.RexSubQuery; +import org.opensearch.sql.calcite.DynamicSearchPlanBinder; +import org.opensearch.sql.calcite.SearchPredicateCompiler; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.exception.SemanticCheckException; + +/** Executes and binds implicit format subqueries before running the parent search. */ +public final class DynamicSearchExecutor { + + private DynamicSearchExecutor() {} + + @FunctionalInterface + public interface PlanExecutor { + void execute(RelNode plan, ResponseListener listener); + } + + public static void execute( + RelNode plan, + SearchPredicateCompiler compiler, + PlanExecutor subqueryExecutor, + PlanExecutor finalExecutor, + ResponseListener listener) { + bind( + plan, + compiler, + subqueryExecutor, + new ResponseListener<>() { + @Override + public void onResponse(RelNode bound) { + finalExecutor.execute(bound, listener); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); + } + + /** Executes all implicit format subqueries and returns a plan with literal search predicates. */ + public static void bind( + RelNode plan, + SearchPredicateCompiler compiler, + PlanExecutor subqueryExecutor, + ResponseListener listener) { + var dynamic = DynamicSearchPlanBinder.find(plan); + if (dynamic.isEmpty()) { + listener.onResponse(plan); + return; + } + + RexSubQuery subquery = dynamic.get(); + subqueryExecutor.execute( + subquery.rel, + new ResponseListener<>() { + @Override + public void onResponse(ExecutionEngine.QueryResponse response) { + try { + String formatted = extractScalarString(response); + String compiled = compiler.compile(formatted); + RelNode bound = DynamicSearchPlanBinder.bind(plan, subquery, compiled); + bind(bound, compiler, subqueryExecutor, listener); + } catch (Exception e) { + listener.onFailure(e); + } + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); + } + + private static String extractScalarString(ExecutionEngine.QueryResponse response) { + List rows = response.getResults(); + if (rows.size() != 1 || rows.getFirst().tupleValue().size() != 1) { + throw new SemanticCheckException( + "Implicit format subsearch must return exactly one row and one column"); + } + ExprValue value = rows.getFirst().tupleValue().values().iterator().next(); + if (value.isNull() || value.isMissing()) { + throw new SemanticCheckException("Implicit format subsearch returned a null search string"); + } + Object raw = value.value(); + if (!(raw instanceof String string)) { + throw new SemanticCheckException("Implicit format subsearch must return a string"); + } + return string; + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index 858ba0598e6..88b127a1272 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -17,7 +17,6 @@ import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; import lombok.Getter; -import lombok.RequiredArgsConstructor; import lombok.extern.log4j.Log4j2; import org.apache.calcite.jdbc.CalciteSchema; import org.apache.calcite.plan.RelOptUtil; @@ -48,7 +47,9 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.CalciteRelNodeVisitor; +import org.opensearch.sql.calcite.DynamicSearchPlanBinder; import org.opensearch.sql.calcite.OpenSearchSchema; +import org.opensearch.sql.calcite.SearchPredicateCompiler; import org.opensearch.sql.calcite.SysLimit; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType; @@ -64,6 +65,7 @@ import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.exception.CalciteUnsupportedException; import org.opensearch.sql.exception.NonFallbackCalciteException; +import org.opensearch.sql.executor.analytics.TimewrapSignals; import org.opensearch.sql.monitor.profile.MetricName; import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.ProfileMetric; diff --git a/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java b/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java index d7907e93d87..162d8dda52d 100644 --- a/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java +++ b/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java @@ -41,6 +41,13 @@ public static TimewrapSignals captureAndClear() { return new TimewrapSignals(active, unitName, series); } + /** Installs this snapshot on the current thread for an execution engine to capture. */ + public void install() { + CalcitePlanContext.stripNullColumns.set(active); + CalcitePlanContext.timewrapUnitName.set(unitName); + CalcitePlanContext.timewrapSeries.set(series); + } + /** * Applies the timewrap pivot to {@code response} if this snapshot is from a timewrap query; * otherwise returns it unchanged. The returned response carries over the input's profile/error. diff --git a/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java b/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java index eef98cdc4bf..94133c50d89 100644 --- a/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java +++ b/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java @@ -7,6 +7,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; @@ -21,6 +22,7 @@ import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.executor.ExecutionEngine.Schema.Column; +import org.opensearch.sql.executor.analytics.TimewrapSignals; /** * Leak guard for the timewrap pivot signals. {@code visitTimewrap} stashes its pivot state in the @@ -79,4 +81,19 @@ public void timewrapSignalsDoNotLeakOntoNextQueryOnSameThread() { "non-timewrap result gained artifact columns after a prior timewrap query"); assertEquals(rows, result.values()); } + + @Test + public void timewrapSignalsCanMoveAcrossInitPlanExecution() { + CalcitePlanContext.stripNullColumns.set(true); + CalcitePlanContext.timewrapUnitName.set("1|day|days|_before"); + CalcitePlanContext.timewrapSeries.set("relative"); + + TimewrapSignals snapshot = TimewrapSignals.captureAndClear(); + assertFalse(TimewrapPivot.isTimewrap()); + + snapshot.install(); + assertTrue(TimewrapPivot.isTimewrap()); + assertEquals("1|day|days|_before", CalcitePlanContext.timewrapUnitName.get()); + assertEquals("relative", CalcitePlanContext.timewrapSeries.get()); + } } diff --git a/docs/user/ppl/cmd/format.md b/docs/user/ppl/cmd/format.md index e966282f4d2..fed56da4274 100644 --- a/docs/user/ppl/cmd/format.md +++ b/docs/user/ppl/cmd/format.md @@ -2,7 +2,9 @@ The `format` command collapses tabular input into one row containing a `search` string. Within each input row, it renders non-null fields as `field="value"` expressions joined by a column separator. It then joins the rendered rows with a row separator. -`format` produces the expression as data. It does not parse or execute the generated `search` string. +As an explicit pipeline command, `format` produces the expression as data and does not execute the +generated `search` string. A subsearch used directly in a parent `search` expression has an implicit +final `format`; the generated string is parsed as a predicate and applied to the parent search. ## Syntax @@ -24,7 +26,7 @@ If you specify positional delimiters, you must provide all six. | `column-end` | `)` | Suffix for each formatted input row. | | `row-separator` | `OR` | Separator between formatted rows. | | `row-end` | `)` | Suffix for the complete expression. | -| `emptystr` | `NOT( )` | Result when the input has no formattable fields or values. | +| `emptystr` | `NOT ()` | Result when the input has no formattable fields or values. | ## Behavior @@ -32,10 +34,28 @@ If you specify positional delimiters, you must provide all six. - Fields are rendered in lexicographic field-name order. - Null and missing values are omitted. Empty strings are retained. - Fields whose names begin with `_` are treated as internal fields and omitted. -- The field names `search` and `query` are omitted from their rendered predicates; their values are inserted directly into the expression. +- The field names `search` and `query` are omitted from their rendered predicates. Their scalar + and multivalue elements are still quoted and escaped like other values. - Field names containing special characters are double quoted. - Scalar values are converted to strings. Double quotes and backslashes are escaped. - Multivalue fields produce a parenthesized expression containing one predicate per non-null element. +- At an implicit subsearch boundary, a scalar `search` field has special behavior: only the first + result row participates, and a non-null `search` value is inserted as predicate text instead of + formatting that row's other fields. A null `search` value falls back to formatting the first row's + other fields. The `query` field and multivalue `search` fields retain normal `format` behavior. + +## Implicit format in a search subquery + +When a bracketed subsearch appears in the parent `search` expression, PPL executes the subsearch, +implicitly formats its result, parses the resulting string as search predicate syntax, and then +executes the parent search. For example: + +```ppl +search source=logs [ search source=rules | where enabled=true | fields host, status ] +``` + +The subsearch is subject to `plugins.ppl.subsearch.maxout`. An explicit `format` command remains a +regular pipeline command and returns its generated string without executing it. ## Examples @@ -96,10 +116,13 @@ source=logs ## Limitations -- PPL currently supports `format` only as an explicit pipeline command. It is not implicitly appended to subsearches. -- The generated `search` value is not automatically injected into or executed by an outer query. -- Multivalue element escaping is limited by the backend array-expression support. Avoid quote and backslash characters in multivalue elements when the output will be parsed as a search expression. -- Row collection lowers to a global `ARRAY_AGG` without an aggregate order key. Distributed execution can change row order even when the input contains an upstream `sort`, so stable Splunk row-order parity is not guaranteed. +- Implicit format requires the Calcite query engine. +- Runtime binding is supported for bracketed subsearches in a parent `search` expression. Dynamic + subsearch interpolation in `eval` or other command positions is not supported. +- An implicit result is parsed only as a search predicate. Pipeline command syntax in a generated + `search` or `query` value is rejected. +- When upstream ordering metadata is available, row collection carries it into the aggregate order + key. Without an explicit upstream `sort`, distributed execution does not guarantee row order. ## Related commands diff --git a/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java b/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java index 1899213f960..7812d8d0366 100644 --- a/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java @@ -186,6 +186,31 @@ public void testPplProfileTasksHaveNodeAndTiming() throws IOException { assertTrue("at least one stage has tasks", foundTasks); } + @Test + public void testStaticSearchStillRunsThroughAnalyticsEngine() throws IOException { + ensureSetup(); + JSONObject result = + executeWithProfile( + "search source=" + INDEX + " name=alice | fields name", "/_plugins/_ppl"); + + assertEquals("alice", result.getJSONArray("datarows").getJSONArray(0).getString(0)); + } + + @Test + public void testImplicitFormatSubsearchRunsThroughAnalyticsEngine() throws IOException { + ensureSetup(); + JSONObject result = + executeWithProfile( + "search source=" + + INDEX + + " [ search source=" + + INDEX + + " name=alice | fields name | head 1 ] | fields name", + "/_plugins/_ppl"); + + assertEquals("alice", result.getJSONArray("datarows").getJSONArray(0).getString(0)); + } + @Test public void testPplExplainReturnsOnlyPlan() throws IOException { ensureSetup(); @@ -201,6 +226,25 @@ public void testPplExplainReturnsOnlyPlan() throws IOException { assertFalse("no profile in explain", calcite.has("profile")); } + @Test + public void testImplicitFormatExplainBindsRuntimePredicate() throws IOException { + ensureSetup(); + Request request = new Request("POST", "/_plugins/_ppl/_explain"); + request.setJsonEntity( + String.format( + Locale.ROOT, + "{\"query\": \"search source=%s [ search source=%s name=alice | fields name | head" + + " 1 ] | fields name\"}", + INDEX, + INDEX)); + Response response = client().performRequest(request); + JSONObject result = new JSONObject(entityAsString(response)); + + String logical = result.getJSONObject("calcite").get("logical").toString(); + assertTrue(logical, logical.contains("name:alice")); + assertFalse(logical, logical.contains("SCALAR_QUERY")); + } + @Test public void testSqlExplainReturnsOnlyPlan() throws IOException { ensureSetup(); diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java index 59cea652bb0..1bdef8bbcfd 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java @@ -61,4 +61,31 @@ public void testEmptyResultFallback() throws IOException { verifyDataRows(result, rows("no matching data")); } + + @Test + public void testDefaultEmptyResult() throws IOException { + JSONObject result = + executeQuery( + "source=" + + TEST_INDEX_BANK + + " | where account_number < 0 | fields firstname | format"); + + verifyDataRows(result, rows("NOT ()")); + } + + @Test + public void testUpstreamSortControlsRowOrder() throws IOException { + JSONObject result = + executeQuery( + "source=" + + TEST_INDEX_BANK + + " | where account_number in (1, 6, 13) | fields account_number " + + "| sort - account_number | format"); + + verifyDataRows( + result, + rows( + "( ( account_number=\"13\" ) OR ( account_number=\"6\" ) OR " + + "( account_number=\"1\" ) )")); + } } diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java new file mode 100644 index 00000000000..b53c5d3f8e6 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java @@ -0,0 +1,99 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.remote; + +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK; +import static org.opensearch.sql.util.MatcherUtils.rows; +import static org.opensearch.sql.util.MatcherUtils.verifyDataRows; + +import java.io.IOException; +import org.json.JSONObject; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.ppl.PPLIntegTestCase; + +/** End-to-end tests for implicit format subsearches on the query_string pushdown path. */ +public class CalciteImplicitFormatCommandIT extends PPLIntegTestCase { + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); + enabledOnlyWhenPushdownIsEnabled(); + loadIndex(Index.BANK); + } + + @Test + public void testImplicitFormatExecutesRawSearchField() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_BANK + + " [ search source=" + + TEST_INDEX_BANK + + " | head 1 | eval search='account_number=1' | fields search ]" + + " | fields account_number"); + + verifyDataRows(result, rows(1)); + } + + @Test + public void testImplicitFormatExecutesOrdinaryFields() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_BANK + + " [ search source=" + + TEST_INDEX_BANK + + " account_number=1 OR account_number=6 | fields account_number ]" + + " | fields account_number | sort account_number"); + + verifyDataRows(result, rows(1), rows(6)); + } + + @Test + public void testImplicitFormatUnderNotExpression() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_BANK + + " account_number=6 NOT [ search source=" + + TEST_INDEX_BANK + + " account_number=1 | fields account_number ]" + + " | fields account_number"); + + verifyDataRows(result, rows(6)); + } + + @Test + public void testImplicitFormatExplainUsesBoundPredicate() throws IOException { + String result = + explainQueryToString( + "search source=" + + TEST_INDEX_BANK + + " [ search source=" + + TEST_INDEX_BANK + + " | head 1 | eval search='account_number=1' | fields search ]" + + " | fields account_number"); + + assertTrue(result, result.contains("account_number:1")); + assertFalse(result, result.contains("SCALAR_QUERY")); + } + + @Test + public void testPostFormatEvalResultIsExecutedByParentSearch() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_BANK + + " [ search source=" + + TEST_INDEX_BANK + + " account_number=1 | fields account_number | format " + + "| eval search=replace(search, '1', '6') ]" + + " | fields account_number"); + + verifyDataRows(result, rows(6)); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 816f1071310..2bebc645d3b 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -37,6 +37,7 @@ import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer; import org.opensearch.sql.ppl.PPLService; import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; +import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler; import org.opensearch.sql.sql.SQLService; import org.opensearch.sql.sql.antlr.SQLSyntaxParser; import org.opensearch.sql.storage.StorageEngine; diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index 26c71060c2e..797b91a4ac6 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -40,13 +40,16 @@ import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.DynamicSearchExecutor; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.analytics.AnalyticsExecutionEngine; +import org.opensearch.sql.executor.analytics.TimewrapSignals; import org.opensearch.sql.lang.LangSpec; import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.QueryProfiling; import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; +import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler; import org.opensearch.sql.protocol.response.QueryResult; import org.opensearch.sql.protocol.response.format.ResponseFormatter; import org.opensearch.sql.protocol.response.format.SimpleJsonResponseFormatter; @@ -70,6 +73,8 @@ public class RestUnifiedQueryAction { private final org.opensearch.analytics.EngineContextProvider contextProvider; private final org.opensearch.sql.common.setting.Settings pluginSettings; private final org.opensearch.sql.executor.ExecutionDispatcher executionDispatcher; + private final PPLSearchPredicateCompiler searchPredicateCompiler = + new PPLSearchPredicateCompiler(); public RestUnifiedQueryAction( NodeClient client, @@ -232,25 +237,33 @@ private void doExecute( // string, so apply the equivalent top-level limit here before the system cap. plan = addFetchSizeLimit(plan, planContext, fetchSize); plan = addQuerySizeLimit(plan, planContext); - plan = - org.opensearch.sql.calcite.utils.CalciteToolsHelper.optimize( - plan, planContext); - RelNode finalPlan = plan; - Runnable executeTask = + // An init-plan must not consume or apply parent-result post-processing. Carry + // these signals across the async subquery callback and install them only for + // the final parent execution. + TimewrapSignals finalTimewrap = TimewrapSignals.captureAndClear(); + ResponseListener queryListener = + createQueryListener(queryType, profileCtx, closingListener); + DynamicSearchExecutor.PlanExecutor subqueryExecutor = + (subquery, subqueryListener) -> + analyticsEngine.execute( + subquery, planContext, queryCtx, subqueryListener); + DynamicSearchExecutor.PlanExecutor finalExecutor = profiling - ? () -> - analyticsEngine.executeWithProfile( - finalPlan, - planContext, - queryCtx, - createQueryListener(queryType, profileCtx, closingListener)) - : () -> - analyticsEngine.execute( - finalPlan, - planContext, - queryCtx, - createQueryListener(queryType, profileCtx, closingListener)); - executionDispatcher.dispatchTask(finalPlan, planContext, executeTask); + ? (bound, finalListener) -> { + finalTimewrap.install(); + analyticsEngine.executeWithProfile( + bound, planContext, queryCtx, finalListener); + } + : (bound, finalListener) -> { + finalTimewrap.install(); + analyticsEngine.execute(bound, planContext, queryCtx, finalListener); + }; + DynamicSearchExecutor.execute( + plan, + searchPredicateCompiler, + subqueryExecutor, + finalExecutor, + queryListener); } catch (Exception e) { closingListener.onFailure(e); } finally { @@ -298,14 +311,33 @@ private void doExplain( () -> { QueryRequestContext queryCtx = withParentTask(contextProvider.getContext(), parentTask); - try (UnifiedQueryContext context = buildContext(queryType, false, queryCtx)) { + UnifiedQueryContext context = buildContext(queryType, false, queryCtx); + ResponseListener closingListener = + wrapWithContextClose(context, listener); + try { UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context); RelNode plan = planner.plan(query); CalcitePlanContext planContext = context.getPlanContext(); plan = addQuerySizeLimit(plan, planContext); - analyticsEngine.explain(plan, mode, planContext, listener); + DynamicSearchExecutor.bind( + plan, + searchPredicateCompiler, + (subquery, subqueryListener) -> + analyticsEngine.execute( + subquery, planContext, queryCtx, subqueryListener), + new ResponseListener<>() { + @Override + public void onResponse(RelNode bound) { + analyticsEngine.explain(bound, mode, planContext, closingListener); + } + + @Override + public void onFailure(Exception e) { + closingListener.onFailure(e); + } + }); } catch (Exception e) { - listener.onFailure(e); + closingListener.onFailure(e); } finally { // explain plans a timewrap query (visitTimewrap sets thread-locals) but never // executes, so nothing captures-and-clears them — clear here to avoid leaking @@ -489,4 +521,35 @@ private static ActionListener wrapWithContextClose( } }); } + + private static ResponseListener wrapWithContextClose( + UnifiedQueryContext context, ResponseListener delegate) { + return new ResponseListener<>() { + @Override + public void onResponse(T response) { + try { + delegate.onResponse(response); + } finally { + closeContext(context); + } + } + + @Override + public void onFailure(Exception e) { + try { + delegate.onFailure(e); + } finally { + closeContext(context); + } + } + }; + } + + private static void closeContext(UnifiedQueryContext context) { + try { + context.close(); + } catch (Exception e) { + LOG.warn("Failed to close query context", e); + } + } } diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4 index 50dfd64b8d7..9fdc8ec9617 100644 --- a/ppl/src/main/antlr/OpenSearchPPLParser.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLParser.g4 @@ -167,6 +167,7 @@ searchCommand searchExpression : timeModifier # timeModifierExpression + | LT_SQR_PRTHS subSearch RT_SQR_PRTHS # implicitSubqueryExpression | LT_PRTHS searchExpression RT_PRTHS # groupedExpression | NOT searchExpression # notExpression | searchExpression OR searchExpression # orExpression @@ -174,6 +175,10 @@ searchExpression | searchTerm # termExpression ; +searchPredicate + : searchExpression EOF + ; + searchTerm : searchFieldComparison # searchComparisonTerm | searchFieldInList # searchInListTerm diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParser.java b/ppl/src/main/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParser.java index 1d4485e749e..1776a15da13 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParser.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParser.java @@ -24,6 +24,13 @@ public ParseTree parse(String query) { return parser.root(); } + /** Parses predicate text produced by an implicit subsearch format operation. */ + public OpenSearchPPLParser.SearchPredicateContext parseSearchPredicate(String predicate) { + OpenSearchPPLParser parser = createParser(createLexer(predicate)); + parser.addErrorListener(new SyntaxAnalysisErrorListener()); + return parser.searchPredicate(); + } + private OpenSearchPPLParser createParser(Lexer lexer) { return new OpenSearchPPLParser(new CommonTokenStream(lexer)); } 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 3670918fdf6..d1a36a72bed 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 @@ -229,12 +229,10 @@ public UnresolvedPlan visitSearchFrom(SearchFromContext ctx) { .get(); // Safe because we know size > 1 from the if condition } - // Convert to query string - String queryString = combined.toQueryString(); - - // Create Search node with relation and query string + // Static search expressions are converted immediately. An implicit subsearch remains in the + // expression tree until its formatted result is bound at execution time. Relation relation = (Relation) visitFromClause(ctx.fromClause()); - return new Search(relation, queryString, combined); + return Search.fromExpression(relation, combined); } } 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 b7031a1bd68..f74114d0a23 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 @@ -1045,6 +1045,12 @@ public UnresolvedExpression visitLogWithBaseSpan(OpenSearchPPLParser.LogWithBase } // Visitor methods for search expressions + @Override + public SearchExpression visitImplicitSubqueryExpression( + OpenSearchPPLParser.ImplicitSubqueryExpressionContext ctx) { + return new SearchSubquery(astBuilder.visitSubSearch(ctx.subSearch())); + } + @Override public SearchExpression visitGroupedExpression(OpenSearchPPLParser.GroupedExpressionContext ctx) { return new SearchGroup((SearchExpression) visit(ctx.searchExpression())); diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java new file mode 100644 index 00000000000..938e644988f --- /dev/null +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.parser; + +import org.opensearch.sql.ast.expression.SearchExpression; +import org.opensearch.sql.calcite.SearchPredicateCompiler; +import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; +import org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.SearchPredicateContext; + +/** Uses the PPL search-expression grammar to compile an implicit format result. */ +public class PPLSearchPredicateCompiler implements SearchPredicateCompiler { + + private final PPLSyntaxParser syntaxParser = new PPLSyntaxParser(); + + @Override + public String compile(String predicate) { + String trimmed = predicate == null ? "" : predicate.trim(); + if (trimmed.isEmpty()) { + return "*:*"; + } + if (trimmed.matches("(?i)NOT\\s*\\(\\s*\\)")) { + return "*:* AND NOT *:*"; + } + + SearchPredicateContext parsed = syntaxParser.parseSearchPredicate(trimmed); + AstBuilder astBuilder = new AstBuilder(trimmed); + SearchExpression expression = + (SearchExpression) new AstExpressionBuilder(astBuilder).visit(parsed.searchExpression()); + return expression.toQueryString(); + } +} 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 0c5ad9dd4fe..b569bdac36b 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 @@ -35,6 +35,14 @@ public void testSearchCommandIgnoreSearchKeywordShouldPass() { assertNotEquals(null, tree); } + @Test + public void testSearchCommandWithImplicitFormatSubsearchShouldPass() { + ParseTree tree = + new PPLSyntaxParser() + .parse("search source=outer [ search source=inner | fields status, host ]"); + assertNotNull(tree); + } + @Test public void testSearchCommandWithMultipleIndicesShouldPass() { ParseTree tree = new PPLSyntaxParser().parse("search source=t,u a=1 b=2"); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java index 1cc96ff33ba..80597edd428 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java @@ -25,14 +25,25 @@ public void testDefaultFormat() { + " JOB=\"SALESMAN\" ) )\n"); } + @Test + public void testUpstreamSortBecomesAggregateOrderKey() { + RelNode root = getRelNode("source=EMP | fields ENAME | sort ENAME | head 2 | format"); + String logical = org.apache.calcite.plan.RelOptUtil.toString(root); + + org.junit.Assert.assertTrue( + logical, logical.contains("ARRAY_AGG($0) WITHIN GROUP ([1 ASC-nulls-first])")); + verifyResult(root, "search=( ( ENAME=\"ADAMS\" ) OR ( ENAME=\"ALLEN\" ) )\n"); + } + @Test public void testLogicalAndSparkSqlPlan() { RelNode root = getRelNode("source=EMP | fields ENAME | format maxresults=1"); verifyLogical( root, - "LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN($0, ' OR ':VARCHAR)), 0), ||(||('(" - + " ':VARCHAR, ARRAY_JOIN($0, ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT( )':VARCHAR)])\n" - + " LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0) FILTER $1])\n" + "LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR))," + + " 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), '" + + " )':VARCHAR), 'NOT ()':VARCHAR)])\n" + + " LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)])\n" + " LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS" + " NOT NULL($0), ||(||('ENAME=\"':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL," + " '\\':VARCHAR, '\\\\':VARCHAR), '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR)," @@ -40,34 +51,21 @@ public void testLogicalAndSparkSqlPlan() { + " ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('ENAME=\"':VARCHAR," + " REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR)," + " '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR))," - + " ' )':VARCHAR), null:VARCHAR)], $f1=[IS NOT" - + " NULL(CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0)," - + " ||(||('ENAME=\"':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR," - + " '\\\\':VARCHAR), '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR), null:VARCHAR))), '" - + " AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT" - + " NULL($0), ||(||('ENAME=\"':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL," - + " '\\':VARCHAR, '\\\\':VARCHAR), '\"':VARCHAR, '\\\"':VARCHAR)), '\"':VARCHAR)," - + " null:VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR))])\n" + + " ' )':VARCHAR), null:VARCHAR)])\n" + " LogicalSort(fetch=[1])\n" + " LogicalProject(ENAME=[$1])\n" + " LogicalTableScan(table=[[scott, EMP]])\n"); verifyPPLToSparkSQL( root, - "SELECT CASE WHEN CHAR_LENGTH(ARRAY_JOIN(ARRAY_AGG(`__format_row`) FILTER (WHERE `$f1`)," - + " ' OR ')) > 0 THEN '( ' || ARRAY_JOIN(ARRAY_AGG(`__format_row`) FILTER (WHERE" - + " `$f1`), ' OR ') || ' )' ELSE 'NOT( )' END `search`\n" + "SELECT CASE WHEN CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY_AGG(`__format_row`)), ' OR '))" + + " > 0 THEN '( ' || ARRAY_JOIN(ARRAY_COMPACT(ARRAY_AGG(`__format_row`)), ' OR ') || '" + + " )' ELSE 'NOT ()' END `search`\n" + "FROM (SELECT CASE WHEN CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME`" + " IS NOT NULL THEN 'ENAME=\"' || REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\'," + " '\\\\'), '\"', '\\\"') || '\"' ELSE NULL END)), ' AND ')) > 0 THEN '( ' ||" + " ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME` IS NOT NULL THEN 'ENAME=\"' ||" + " REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\', '\\\\'), '\"', '\\\"') || '\"'" - + " ELSE NULL END)), ' AND ') || ' )' ELSE NULL END `__format_row`, CASE WHEN" - + " CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME` IS NOT NULL THEN" - + " 'ENAME=\"' || REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\', '\\\\'), '\"'," - + " '\\\"') || '\"' ELSE NULL END)), ' AND ')) > 0 THEN '( ' ||" - + " ARRAY_JOIN(ARRAY_COMPACT(ARRAY (CASE WHEN `ENAME` IS NOT NULL THEN 'ENAME=\"' ||" - + " REPLACE(REPLACE(CAST(`ENAME` AS STRING), '\\', '\\\\'), '\"', '\\\"') || '\"'" - + " ELSE NULL END)), ' AND ') || ' )' ELSE NULL END IS NOT NULL `$f1`\n" + + " ELSE NULL END)), ' AND ') || ' )' ELSE NULL END `__format_row`\n" + "FROM `scott`.`EMP`\n" + "LIMIT 1) `t1`"); } @@ -80,6 +78,14 @@ public void testCustomDelimitersAndMaxResults() { .expectResult("search=[ [ ENAME=\"SMITH\" && JOB=\"CLERK\" ] ]\n"); } + @Test + public void testAllEmptyDelimitersPreserveExpectedSpacing() { + withPPLQuery( + "source=EMP | where EMPNO=7369 | fields ENAME, JOB " + + "| format \"\" \"\" \"\" \"\" \"\" \"\"") + .expectResult("search= ENAME=\"SMITH\" JOB=\"CLERK\" \n"); + } + @Test public void testMultivalueFormat() { withPPLQuery( @@ -89,6 +95,14 @@ public void testMultivalueFormat() { .expectResult("search={ [ ( tags=\"critical\" mvseparator tags=\"network\" ) ] }\n"); } + @Test + public void testMultivalueElementsAreEscapedIndividually() { + withPPLQuery( + "source=EMP | head 1 | eval tags=array('say \\\"hi\\\"', 'a\\\\b') " + + "| fields tags | format") + .expectResult("search=( ( ( tags=\"say \\\"hi\\\"\" OR tags=\"a\\\\b\" ) ) )\n"); + } + @Test public void testEmptyMultivalueUsesFallback() { withPPLQuery( @@ -104,6 +118,12 @@ public void testEmptyResultFallback() { .expectResult("search=no matching data\n"); } + @Test + public void testDefaultEmptyResultUsesExpectedSpacing() { + withPPLQuery("source=EMP | where EMPNO < 0 | fields ENAME | format") + .expectResult("search=NOT ()\n"); + } + @Test public void testNullOnlyRowUsesFallback() { withPPLQuery("source=EMP | where EMPNO=7369 | fields COMM | format emptystr=\"empty\"") diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java index ce5bed94415..f66304878f1 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java @@ -7,7 +7,11 @@ import static org.junit.Assert.assertThrows; +import java.sql.PreparedStatement; +import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; import org.apache.calcite.plan.RelTraitDef; import org.apache.calcite.rel.RelNode; import org.apache.calcite.schema.SchemaPlus; @@ -17,6 +21,13 @@ import org.apache.calcite.tools.Programs; import org.junit.Ignore; import org.junit.Test; +import org.opensearch.sql.calcite.DynamicSearchPlanBinder; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.data.model.ExprValueUtils; +import org.opensearch.sql.executor.DynamicSearchExecutor; +import org.opensearch.sql.executor.ExecutionEngine; +import org.opensearch.sql.executor.pagination.Cursor; +import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler; public class CalcitePPLSearchTest extends CalcitePPLAbstractTest { public CalcitePPLSearchTest() { @@ -54,6 +65,156 @@ public void testSearchWithFilter() { verifyPPLToSparkSQL(root, expectedSparkSql); } + @Test + public void testSearchWithImplicitFormatSubsearchUsesScalarQuery() { + RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); + RelNode subquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + String logical = org.apache.calcite.plan.RelOptUtil.toString(root); + org.junit.Assert.assertTrue(logical, logical.contains("SCALAR_QUERY")); + org.junit.Assert.assertTrue(logical, logical.contains("ARRAY_AGG")); + org.junit.Assert.assertTrue(logical, logical.contains("query_string")); + org.junit.Assert.assertNotSame(root.getCluster(), subquery.getCluster()); + } + + @Test + public void testBoundImplicitFormatHasSameStaticQueryStringShape() { + RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); + var subquery = DynamicSearchPlanBinder.find(root).orElseThrow(); + + RelNode bound = DynamicSearchPlanBinder.bind(root, subquery, "DEPTNO:10"); + String logical = org.apache.calcite.plan.RelOptUtil.toString(bound); + + org.junit.Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); + org.junit.Assert.assertTrue( + logical, logical.contains("query_string(MAP('query', 'DEPTNO:10':VARCHAR))")); + } + + @Test + public void testBoundImplicitFormatCombinesWithStaticSearchPredicate() { + RelNode root = + getRelNode("search source=EMP JOB=CLERK [ search source=DEPT | fields DEPTNO | head 1 ]"); + var subquery = DynamicSearchPlanBinder.find(root).orElseThrow(); + + RelNode bound = DynamicSearchPlanBinder.bind(root, subquery, "DEPTNO:10"); + String logical = org.apache.calcite.plan.RelOptUtil.toString(bound); + + org.junit.Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); + org.junit.Assert.assertTrue( + logical, + logical.contains("query_string(MAP('query', '(JOB:CLERK) AND (DEPTNO:10)':VARCHAR))")); + } + + @Test + public void testMultipleImplicitFormatSubsearchesBindSequentially() { + RelNode root = + getRelNode( + "search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ] OR " + + "[ search source=EMP | fields JOB | head 1 ]"); + + RelNode firstBound = + DynamicSearchPlanBinder.bind( + root, DynamicSearchPlanBinder.find(root).orElseThrow(), "DEPTNO:10"); + org.junit.Assert.assertTrue(DynamicSearchPlanBinder.find(firstBound).isPresent()); + RelNode fullyBound = + DynamicSearchPlanBinder.bind( + firstBound, DynamicSearchPlanBinder.find(firstBound).orElseThrow(), "JOB:CLERK"); + String logical = org.apache.calcite.plan.RelOptUtil.toString(fullyBound); + + org.junit.Assert.assertTrue(DynamicSearchPlanBinder.find(fullyBound).isEmpty()); + org.junit.Assert.assertTrue( + logical, + logical.contains("query_string(MAP('query', '(DEPTNO:10 OR JOB:CLERK)':VARCHAR))")); + } + + @Test + public void testImplicitFormatSubsearchProducesSearchPredicateData() { + RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); + RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + + verifyResult(formatSubquery, "search=( ( DEPTNO=\"10\" ) )\n"); + } + + @Test + public void testImplicitFormatDoesNotExecuteMultivalueSearchFieldAsRawPredicate() { + RelNode root = + getRelNode( + "search source=EMP [ search source=DEPT | head 1 " + + "| eval search=array('DEPTNO=10', 'DEPTNO=20') | fields search ]"); + RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + + verifyResult(formatSubquery, "search=( ( ( \"DEPTNO=10\" OR \"DEPTNO=20\" ) ) )\n"); + } + + @Test + public void testImplicitFormatUsesFirstScalarSearchValueAndIgnoresOtherFields() { + RelNode root = + getRelNode( + "search source=EMP [ search source=DEPT " + + "| eval search=if(DEPTNO=10, 'DEPTNO=20', 'DEPTNO=10') " + + "| fields search, DEPTNO ]"); + RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + + verifyResult(formatSubquery, "search=DEPTNO=20\n"); + } + + @Test + public void testImplicitFormatUsesPostFormatEvalResult() { + RelNode root = + getRelNode( + "search source=EMP [ search source=DEPT | where DEPTNO=10 | fields DEPTNO " + + "| format | eval search=replace(search, '10', '20') ]"); + RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + + verifyResult(formatSubquery, "search=( ( DEPTNO=\"20\" ) )\n"); + } + + @Test + public void testImplicitFormatExecutesThenBindsParentSearch() { + RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); + AtomicReference executedParent = new AtomicReference<>(); + AtomicReference failure = new AtomicReference<>(); + + DynamicSearchExecutor.execute( + root, + new PPLSearchPredicateCompiler(), + (subquery, listener) -> executeScalarSubquery(subquery, listener), + (bound, listener) -> { + executedParent.set(bound); + listener.onResponse( + new ExecutionEngine.QueryResponse(null, Collections.emptyList(), Cursor.None)); + }, + new ResponseListener<>() { + @Override + public void onResponse(ExecutionEngine.QueryResponse response) {} + + @Override + public void onFailure(Exception e) { + failure.set(e); + } + }); + + org.junit.Assert.assertNull(failure.get()); + String logical = org.apache.calcite.plan.RelOptUtil.toString(executedParent.get()); + org.junit.Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); + org.junit.Assert.assertTrue( + logical, logical.contains("query_string(MAP('query', '((DEPTNO:10))':VARCHAR))")); + } + + private void executeScalarSubquery( + RelNode subquery, ResponseListener listener) { + try (PreparedStatement statement = org.apache.calcite.tools.RelRunners.run(subquery); + var resultSet = statement.executeQuery()) { + org.junit.Assert.assertTrue(resultSet.next()); + listener.onResponse( + new ExecutionEngine.QueryResponse( + null, + List.of(ExprValueUtils.tupleValue(Map.of("search", resultSet.getString(1)))), + Cursor.None)); + } catch (Exception e) { + listener.onFailure(e); + } + } + @Ignore("Fields used in search commands are not validated. Enable after fixing it.") @Test public void testSearchWithoutTimestampShouldThrow() { diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java index e89958c64dd..45ae05912e2 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java @@ -52,6 +52,16 @@ public void testFieldComparison() { assertEquals("logs", relation.getTableQualifiedName().toString()); } + @Test + public void testImplicitFormatSubsearchRemainsRuntimeBound() { + Node plan = buildPlan("search source=logs [ search source=rules | fields status, host ]"); + + assertTrue(plan instanceof Search); + Search search = (Search) plan; + assertTrue(search.hasImplicitSubquery()); + assertEquals(null, search.getQueryString()); + } + @Test public void testAndExpression() { String query = "search status=200 AND message=\"success\" source=logs"; diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java new file mode 100644 index 00000000000..042e1457edc --- /dev/null +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.parser; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import org.junit.Test; + +public class PPLSearchPredicateCompilerTest { + + private final PPLSearchPredicateCompiler compiler = new PPLSearchPredicateCompiler(); + + @Test + public void testCompileCanonicalFormatOutput() { + assertEquals( + "(((status:500 AND host:api\\-01)))", + compiler.compile("( ( status=\"500\" AND host=\"api-01\" ) )")); + } + + @Test + public void testCompileRuntimeSearchFieldValue() { + assertEquals( + "(status:>=500 OR host:api\\-01)", compiler.compile("status>=500 OR host=\"api-01\"")); + } + + @Test + public void testCompileEmptyFormatResultAsMatchNone() { + assertEquals("*:* AND NOT *:*", compiler.compile("NOT ()")); + } + + @Test + public void testCompileEmptyRawSearchFieldAsMatchAll() { + assertEquals("*:*", compiler.compile("")); + } + + @Test + public void testGeneratedPredicateCannotInjectPipelineCommands() { + assertThrows(RuntimeException.class, () -> compiler.compile("status=500 | head 1")); + } +} diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java index 3ccdc2726de..fc63992876c 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryDataAnonymizerTest.java @@ -41,6 +41,13 @@ public void testMakeResultsCommand() { assertEquals("makeresults", anonymize("makeresults count=5")); } + @Test + public void testImplicitFormatSubsearch() { + assertEquals( + "source=table [ subsearch ]", + anonymize("search source=outer [ search source=inner | fields host ]")); + } + @Test public void testTableFunctionCommand() { assertEquals( From de41da7c90f97fc889609acecc312512019777ae Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 23 Jul 2026 15:21:02 +0800 Subject: [PATCH 03/10] Use imports for format query classes Signed-off-by: Songkan Tang --- .../org/opensearch/sql/ast/tree/Search.java | 3 +- .../sql/calcite/remote/CalciteExplainIT.java | 18 +++++++- .../sql/ppl/calcite/CalcitePPLFormatTest.java | 6 ++- .../sql/ppl/calcite/CalcitePPLSearchTest.java | 46 ++++++++++--------- 4 files changed, 46 insertions(+), 27 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Search.java b/core/src/main/java/org/opensearch/sql/ast/tree/Search.java index d574cf03741..854b6b40d93 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Search.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Search.java @@ -14,6 +14,7 @@ import lombok.ToString; import org.opensearch.sql.ast.AbstractNodeVisitor; import org.opensearch.sql.ast.expression.SearchExpression; +import org.opensearch.sql.ast.expression.SearchSubquery; /** * Logical plan node for Search operation. Represents search expressions that get converted to @@ -48,7 +49,7 @@ public boolean hasImplicitSubquery() { } private static boolean containsSubquery(SearchExpression expression) { - if (expression instanceof org.opensearch.sql.ast.expression.SearchSubquery) { + if (expression instanceof SearchSubquery) { return true; } return expression.getChild().stream() 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 08386b9789d..b26059f5ac8 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 @@ -2920,7 +2920,7 @@ public void testNoMvBasic() throws IOException { String logical = logicalPlan(explainQueryYaml(query)); Assert.assertTrue( "Expected logical plan to contain ARRAY_JOIN function", - logical.toLowerCase(java.util.Locale.ROOT).contains("array_join")); + logical.toLowerCase(Locale.ROOT).contains("array_join")); } @Test @@ -2930,12 +2930,26 @@ public void testNoMvWithEval() throws IOException { "source=%s | eval full_name = concat(firstname, ' J.') | eval name_array =" + " array(full_name) | nomv name_array | fields name_array", TEST_INDEX_BANK); - String logical = logicalPlan(explainQueryYaml(query)).toLowerCase(java.util.Locale.ROOT); + String logical = logicalPlan(explainQueryYaml(query)).toLowerCase(Locale.ROOT); Assert.assertTrue( "Expected logical plan to contain both CONCAT and ARRAY_JOIN", logical.contains("concat") && logical.contains("array_join")); } + @Test + public void testFormatExplain() throws IOException { + String query = + StringUtils.format( + "source=%s | where account_number < 3 | fields firstname, account_number | format", + TEST_INDEX_BANK); + String logical = logicalPlan(explainQueryYaml(query)).toLowerCase(Locale.ROOT); + Assert.assertTrue( + "Expected format to lower to a global ARRAY_AGG and a search projection", + logical.contains("array_agg") + && logical.contains("array_join") + && logical.contains("search=")); + } + @Test public void testForeachExplain() throws IOException { String query = diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java index 80597edd428..d7eb977850a 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java @@ -7,8 +7,10 @@ import static org.junit.Assert.assertThrows; +import org.apache.calcite.plan.RelOptUtil; import org.apache.calcite.rel.RelNode; import org.apache.calcite.test.CalciteAssert; +import org.junit.Assert; import org.junit.Test; public class CalcitePPLFormatTest extends CalcitePPLAbstractTest { @@ -28,9 +30,9 @@ public void testDefaultFormat() { @Test public void testUpstreamSortBecomesAggregateOrderKey() { RelNode root = getRelNode("source=EMP | fields ENAME | sort ENAME | head 2 | format"); - String logical = org.apache.calcite.plan.RelOptUtil.toString(root); + String logical = RelOptUtil.toString(root); - org.junit.Assert.assertTrue( + Assert.assertTrue( logical, logical.contains("ARRAY_AGG($0) WITHIN GROUP ([1 ASC-nulls-first])")); verifyResult(root, "search=( ( ENAME=\"ADAMS\" ) OR ( ENAME=\"ALLEN\" ) )\n"); } diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java index f66304878f1..a8e69454d46 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java @@ -12,6 +12,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicReference; +import org.apache.calcite.plan.RelOptUtil; import org.apache.calcite.plan.RelTraitDef; import org.apache.calcite.rel.RelNode; import org.apache.calcite.schema.SchemaPlus; @@ -19,6 +20,8 @@ import org.apache.calcite.test.CalciteAssert; import org.apache.calcite.tools.Frameworks; import org.apache.calcite.tools.Programs; +import org.apache.calcite.tools.RelRunners; +import org.junit.Assert; import org.junit.Ignore; import org.junit.Test; import org.opensearch.sql.calcite.DynamicSearchPlanBinder; @@ -69,11 +72,11 @@ public void testSearchWithFilter() { public void testSearchWithImplicitFormatSubsearchUsesScalarQuery() { RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); RelNode subquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; - String logical = org.apache.calcite.plan.RelOptUtil.toString(root); - org.junit.Assert.assertTrue(logical, logical.contains("SCALAR_QUERY")); - org.junit.Assert.assertTrue(logical, logical.contains("ARRAY_AGG")); - org.junit.Assert.assertTrue(logical, logical.contains("query_string")); - org.junit.Assert.assertNotSame(root.getCluster(), subquery.getCluster()); + String logical = RelOptUtil.toString(root); + Assert.assertTrue(logical, logical.contains("SCALAR_QUERY")); + Assert.assertTrue(logical, logical.contains("ARRAY_AGG")); + Assert.assertTrue(logical, logical.contains("query_string")); + Assert.assertNotSame(root.getCluster(), subquery.getCluster()); } @Test @@ -82,11 +85,10 @@ public void testBoundImplicitFormatHasSameStaticQueryStringShape() { var subquery = DynamicSearchPlanBinder.find(root).orElseThrow(); RelNode bound = DynamicSearchPlanBinder.bind(root, subquery, "DEPTNO:10"); - String logical = org.apache.calcite.plan.RelOptUtil.toString(bound); + String logical = RelOptUtil.toString(bound); - org.junit.Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - org.junit.Assert.assertTrue( - logical, logical.contains("query_string(MAP('query', 'DEPTNO:10':VARCHAR))")); + Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); + Assert.assertTrue(logical, logical.contains("query_string(MAP('query', 'DEPTNO:10':VARCHAR))")); } @Test @@ -96,10 +98,10 @@ public void testBoundImplicitFormatCombinesWithStaticSearchPredicate() { var subquery = DynamicSearchPlanBinder.find(root).orElseThrow(); RelNode bound = DynamicSearchPlanBinder.bind(root, subquery, "DEPTNO:10"); - String logical = org.apache.calcite.plan.RelOptUtil.toString(bound); + String logical = RelOptUtil.toString(bound); - org.junit.Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - org.junit.Assert.assertTrue( + Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); + Assert.assertTrue( logical, logical.contains("query_string(MAP('query', '(JOB:CLERK) AND (DEPTNO:10)':VARCHAR))")); } @@ -114,14 +116,14 @@ public void testMultipleImplicitFormatSubsearchesBindSequentially() { RelNode firstBound = DynamicSearchPlanBinder.bind( root, DynamicSearchPlanBinder.find(root).orElseThrow(), "DEPTNO:10"); - org.junit.Assert.assertTrue(DynamicSearchPlanBinder.find(firstBound).isPresent()); + Assert.assertTrue(DynamicSearchPlanBinder.find(firstBound).isPresent()); RelNode fullyBound = DynamicSearchPlanBinder.bind( firstBound, DynamicSearchPlanBinder.find(firstBound).orElseThrow(), "JOB:CLERK"); - String logical = org.apache.calcite.plan.RelOptUtil.toString(fullyBound); + String logical = RelOptUtil.toString(fullyBound); - org.junit.Assert.assertTrue(DynamicSearchPlanBinder.find(fullyBound).isEmpty()); - org.junit.Assert.assertTrue( + Assert.assertTrue(DynamicSearchPlanBinder.find(fullyBound).isEmpty()); + Assert.assertTrue( logical, logical.contains("query_string(MAP('query', '(DEPTNO:10 OR JOB:CLERK)':VARCHAR))")); } @@ -193,18 +195,18 @@ public void onFailure(Exception e) { } }); - org.junit.Assert.assertNull(failure.get()); - String logical = org.apache.calcite.plan.RelOptUtil.toString(executedParent.get()); - org.junit.Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - org.junit.Assert.assertTrue( + Assert.assertNull(failure.get()); + String logical = RelOptUtil.toString(executedParent.get()); + Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); + Assert.assertTrue( logical, logical.contains("query_string(MAP('query', '((DEPTNO:10))':VARCHAR))")); } private void executeScalarSubquery( RelNode subquery, ResponseListener listener) { - try (PreparedStatement statement = org.apache.calcite.tools.RelRunners.run(subquery); + try (PreparedStatement statement = RelRunners.run(subquery); var resultSet = statement.executeQuery()) { - org.junit.Assert.assertTrue(resultSet.next()); + Assert.assertTrue(resultSet.next()); listener.onResponse( new ExecutionEngine.QueryResponse( null, From 699a3e07b06492a3503b614cb2a82cf4863cd983 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 23 Jul 2026 16:48:02 +0800 Subject: [PATCH 04/10] Use correlated scans for implicit format searches Signed-off-by: Songkan Tang --- .../subquery/RuntimeSearchScalarSubquery.java | 16 -- .../sql/calcite/CalcitePlanContext.java | 2 + .../sql/calcite/CalciteRelNodeVisitor.java | 9 +- .../sql/calcite/CalciteRexNodeVisitor.java | 16 -- .../sql/calcite/DynamicSearchPlanBinder.java | 192 ------------------ .../opensearch/sql/calcite/FormatPlanner.java | 18 +- .../sql/calcite/RuntimeSearchCorrelator.java | 157 ++++++++++++++ .../plan/DynamicQueryStringPushDown.java | 28 +++ .../sql/executor/DynamicSearchExecutor.java | 101 --------- .../opensearch/sql/executor/QueryService.java | 27 ++- .../analytics/AnalyticsEngineProfileIT.java | 20 +- .../remote/CalciteFormatCommandIT.java | 11 + .../CalciteImplicitFormatCommandIT.java | 35 +++- .../rules/EnumerableIndexScanRule.java | 3 +- .../scan/AbstractCalciteIndexScan.java | 13 ++ .../scan/CalciteEnumerableIndexScan.java | 41 ++++ .../storage/scan/CalciteLogicalIndexScan.java | 48 ++++- .../scan/context/DynamicQueryStringSpec.java | 73 +++++++ .../storage/scan/context/PushDownContext.java | 4 + .../plugin/config/OpenSearchPluginModule.java | 8 +- .../plugin/rest/RestUnifiedQueryAction.java | 57 +----- .../sql/ppl/calcite/CalcitePPLFormatTest.java | 21 +- .../sql/ppl/calcite/CalcitePPLSearchTest.java | 152 ++++++-------- .../PPLSearchPredicateCompilerTest.java | 5 + 24 files changed, 550 insertions(+), 507 deletions(-) delete mode 100644 core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java delete mode 100644 core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java create mode 100644 core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java create mode 100644 core/src/main/java/org/opensearch/sql/calcite/plan/DynamicQueryStringPushDown.java delete mode 100644 core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/DynamicQueryStringSpec.java diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java b/core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java deleted file mode 100644 index 816b8c11e13..00000000000 --- a/core/src/main/java/org/opensearch/sql/ast/expression/subquery/RuntimeSearchScalarSubquery.java +++ /dev/null @@ -1,16 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.ast.expression.subquery; - -import org.opensearch.sql.ast.tree.UnresolvedPlan; - -/** A runtime search scalar subquery that is planned independently from its parent query. */ -public class RuntimeSearchScalarSubquery extends ScalarSubquery { - - public RuntimeSearchScalarSubquery(UnresolvedPlan query) { - super(query); - } -} 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 c7f3bc373ac..24261ee1e17 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -68,6 +68,7 @@ public class CalcitePlanContext { ThreadLocal.withInitial(() -> true); @Getter @Setter private HighlightConfig highlightConfig; + @Getter @Setter private SearchPredicateCompiler searchPredicateCompiler; @Getter @Setter private boolean isResolvingJoinCondition = false; @Getter @Setter private boolean isResolvingSubquery = false; @Getter @Setter private boolean inCoalesceFunction = false; @@ -165,6 +166,7 @@ private CalcitePlanContext(CalcitePlanContext parent) { this.rexBuilder = parent.rexBuilder; // Share the same rexBuilder this.functionProperties = parent.functionProperties; this.highlightConfig = parent.highlightConfig; + this.searchPredicateCompiler = parent.searchPredicateCompiler; 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 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 2f31d59eb65..24e7c31b1cd 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -119,7 +119,7 @@ import org.opensearch.sql.ast.expression.WindowFrame; import org.opensearch.sql.ast.expression.WindowFrame.FrameType; import org.opensearch.sql.ast.expression.WindowFunction; -import org.opensearch.sql.ast.expression.subquery.RuntimeSearchScalarSubquery; +import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; import org.opensearch.sql.ast.expression.subquery.SubqueryExpression; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.AddColTotals; @@ -347,6 +347,11 @@ public RelNode visitSearch(Search node, CalcitePlanContext context) { RexNode queryStringRex = rexVisitor.analyze(queryStringFunc, context); context.relBuilder.filter(queryStringRex); + if (node.hasImplicitSubquery()) { + context.relBuilder.push( + RuntimeSearchCorrelator.correlate( + context.relBuilder.build(), context.getSearchPredicateCompiler())); + } return context.relBuilder.peek(); } @@ -368,7 +373,7 @@ private UnresolvedExpression buildRuntimeSearchQuery( Format.DEFAULT_EMPTY_STRING); format.setImplicit(true); format.attach(subquery.getQuery()); - return new RuntimeSearchScalarSubquery(format); + return new ScalarSubquery(format); } if (expression instanceof SearchGroup group) { return concatSearch("(", buildRuntimeSearchQuery(group.getExpression(), 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 7b5d761eb33..1bf1e217b51 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -12,7 +12,6 @@ import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY; import java.math.BigDecimal; -import java.sql.SQLException; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -35,7 +34,6 @@ import org.apache.calcite.rex.RexLambdaRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; -import org.apache.calcite.rex.RexSubQuery; import org.apache.calcite.sql.SqlIntervalQualifier; import org.apache.calcite.sql.SqlOperator; import org.apache.calcite.sql.fun.SqlStdOperatorTable; @@ -79,7 +77,6 @@ import org.opensearch.sql.ast.expression.Xor; import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; import org.opensearch.sql.ast.expression.subquery.InSubquery; -import org.opensearch.sql.ast.expression.subquery.RuntimeSearchScalarSubquery; import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; import org.opensearch.sql.ast.expression.subquery.SubqueryExpression; import org.opensearch.sql.ast.tree.Sort.SortOption; @@ -868,19 +865,6 @@ public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) { @Override public RexNode visitScalarSubquery(ScalarSubquery node, CalcitePlanContext context) { - if (node instanceof RuntimeSearchScalarSubquery) { - CalcitePlanContext initContext = - CalcitePlanContext.create(context.config, context.sysLimit, context.queryType); - try { - return RexSubQuery.scalar(resolveSubqueryPlan(node.getQuery(), node, initContext)); - } finally { - try { - initContext.connection.close(); - } catch (SQLException e) { - throw new RuntimeException("Failed to close runtime search planning context", e); - } - } - } return context.relBuilder.scalarQuery( b -> { UnresolvedPlan subquery = node.getQuery(); diff --git a/core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java b/core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java deleted file mode 100644 index f011545a1f0..00000000000 --- a/core/src/main/java/org/opensearch/sql/calcite/DynamicSearchPlanBinder.java +++ /dev/null @@ -1,192 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.calcite; - -import java.util.ArrayList; -import java.util.List; -import java.util.Locale; -import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; -import org.apache.calcite.rel.RelNode; -import org.apache.calcite.rel.RelVisitor; -import org.apache.calcite.rex.RexBuilder; -import org.apache.calcite.rex.RexCall; -import org.apache.calcite.rex.RexLiteral; -import org.apache.calcite.rex.RexNode; -import org.apache.calcite.rex.RexShuttle; -import org.apache.calcite.rex.RexSubQuery; -import org.apache.calcite.sql.SqlKind; -import org.apache.calcite.sql.type.SqlTypeName; - -/** Finds and binds scalar subqueries used as runtime {@code search} predicates. */ -public final class DynamicSearchPlanBinder { - - private DynamicSearchPlanBinder() {} - - /** Returns the first unbound scalar subquery beneath a query_string query argument. */ - public static Optional find(RelNode plan) { - AtomicReference found = new AtomicReference<>(); - RexShuttle finder = - new RexShuttle() { - @Override - public RexNode visitCall(RexCall call) { - if (found.get() == null && isQueryString(call)) { - RexNode query = queryArgument(call); - if (query != null) { - query.accept( - new RexShuttle() { - @Override - public RexNode visitSubQuery(RexSubQuery subQuery) { - found.compareAndSet(null, subQuery); - return subQuery; - } - }); - } - } - return found.get() == null ? super.visitCall(call) : call; - } - }; - - new RelVisitor() { - @Override - public void visit(RelNode node, int ordinal, RelNode parent) { - if (found.get() == null) { - node.accept(finder); - super.visit(node, ordinal, parent); - } - } - }.go(plan); - return Optional.ofNullable(found.get()); - } - - /** Replaces one runtime scalar query and folds the surrounding string concatenation. */ - public static RelNode bind(RelNode plan, RexSubQuery target, String compiledQueryString) { - RexBuilder rexBuilder = plan.getCluster().getRexBuilder(); - RexShuttle queryArgumentBinder = - new RexShuttle() { - @Override - public RexNode visitSubQuery(RexSubQuery subQuery) { - if (subQuery == target) { - return stringLiteral(rexBuilder, compiledQueryString); - } - return subQuery; - } - - @Override - public RexNode visitCall(RexCall call) { - RexNode visited = super.visitCall(call); - if (visited instanceof RexCall visitedCall) { - String constant = constantString(visitedCall); - if (constant != null) { - return stringLiteral(rexBuilder, constant); - } - } - return visited; - } - }; - RexShuttle binder = - new RexShuttle() { - @Override - public RexNode visitCall(RexCall call) { - if (isQueryString(call)) { - RexNode visited = call.accept(queryArgumentBinder); - return visited instanceof RexCall visitedCall - ? foldQueryStringArgument(visitedCall, rexBuilder) - : visited; - } - return super.visitCall(call); - } - }; - return rewrite(plan, binder); - } - - private static RelNode rewrite(RelNode node, RexShuttle binder) { - List inputs = node.getInputs(); - List rewrittenInputs = new ArrayList<>(inputs.size()); - boolean changed = false; - for (RelNode input : inputs) { - RelNode rewritten = rewrite(input, binder); - rewrittenInputs.add(rewritten); - changed |= rewritten != input; - } - RelNode withInputs = changed ? node.copy(node.getTraitSet(), rewrittenInputs) : node; - return withInputs.accept(binder); - } - - private static String constantString(RexCall call) { - if ((call.getKind() == SqlKind.CAST || call.getOperator().getName().equalsIgnoreCase("cast")) - && call.getOperands().size() == 1) { - return literalString(call.getOperands().getFirst()); - } - String name = call.getOperator().getName().toLowerCase(Locale.ROOT); - if (!(name.equals("concat") || name.equals("||"))) { - return null; - } - StringBuilder result = new StringBuilder(); - for (RexNode operand : call.getOperands()) { - String value = literalString(operand); - if (value == null) { - return null; - } - result.append(value); - } - return result.toString(); - } - - private static RexNode foldQueryStringArgument(RexCall queryString, RexBuilder rexBuilder) { - List operands = new ArrayList<>(queryString.getOperands()); - for (int i = 0; i < operands.size(); i++) { - RexNode operand = operands.get(i); - if (operand instanceof RexCall map && map.getOperands().size() >= 2) { - RexNode key = map.getOperands().getFirst(); - if (key instanceof RexLiteral literal - && "query".equalsIgnoreCase(literal.getValueAs(String.class))) { - RexNode value = map.getOperands().get(1); - String constant = literalString(value); - if (constant != null && !(value instanceof RexLiteral)) { - List mapOperands = new ArrayList<>(map.getOperands()); - mapOperands.set(1, stringLiteral(rexBuilder, constant)); - operands.set(i, map.clone(map.getType(), mapOperands)); - return queryString.clone(queryString.getType(), operands); - } - } - } - } - return queryString; - } - - private static RexNode stringLiteral(RexBuilder rexBuilder, String value) { - return rexBuilder.makeLiteral( - value, rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR), true); - } - - private static String literalString(RexNode node) { - if (node instanceof RexLiteral literal) { - return literal.getValueAs(String.class); - } - if (node instanceof RexCall call) { - return constantString(call); - } - return null; - } - - private static boolean isQueryString(RexCall call) { - return call.getOperator().getName().equalsIgnoreCase("query_string"); - } - - private static RexNode queryArgument(RexCall queryString) { - for (RexNode operand : queryString.getOperands()) { - if (operand instanceof RexCall map && map.getOperands().size() >= 2) { - RexNode key = map.getOperands().getFirst(); - if (key instanceof RexLiteral literal - && "query".equalsIgnoreCase(literal.getValueAs(String.class))) { - return map.getOperands().get(1); - } - } - } - return null; - } -} diff --git a/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java index 7f793ba80b2..af20ebde1d6 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java +++ b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java @@ -5,6 +5,8 @@ package org.opensearch.sql.calcite; +import static org.opensearch.sql.calcite.plan.OpenSearchConstants.METADATAFIELD_TYPE_MAP; + import java.util.ArrayList; import java.util.Comparator; import java.util.List; @@ -36,7 +38,7 @@ public RelNode plan(Format node, CalcitePlanContext context) { List fields = builder.peek().getRowType().getFieldList().stream() - .filter(field -> !field.getName().startsWith("_")) + .filter(field -> !METADATAFIELD_TYPE_MAP.containsKey(field.getName())) .sorted(Comparator.comparing(RelDataTypeField::getName)) .toList(); @@ -201,7 +203,6 @@ private RexNode formatField(RelDataTypeField field, Format node, CalcitePlanCont RexNode formattedValue; RexNode hasValue = builder.isNotNull(value); String fieldName = formatFieldName(field.getName()); - boolean rawValue = field.getName().equals("search") || field.getName().equals("query"); if (SqlTypeUtil.isArray(field.getType()) || SqlTypeUtil.isMultiset(field.getType())) { RelDataType varchar = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR); @@ -210,15 +211,12 @@ private RexNode formatField(RelDataTypeField field, Format node, CalcitePlanCont builder.call( SqlLibraryOperators.ARRAY_COMPACT, context.rexBuilder.makeCast(varcharArray, value)); values = escapeArrayValues(values, varchar, context); - String repeatedValueSeparator = - rawValue - ? "\" " + node.getMvSeparator() + " \"" - : "\" " + node.getMvSeparator() + " " + fieldName + "=\""; + String repeatedValueSeparator = "\" " + node.getMvSeparator() + " " + fieldName + "=\""; RexNode joinedValues = arrayJoin(context, values, repeatedValueSeparator); formattedValue = concat( context, - stringLiteral(rawValue ? "( \"" : "( " + fieldName + "=\"", context), + stringLiteral("( " + fieldName + "=\"", context), joinedValues, stringLiteral("\" )", context)); hasValue = builder.call(SqlStdOperatorTable.AND, hasValue, isNotEmpty(joinedValues, context)); @@ -227,7 +225,7 @@ private RexNode formatField(RelDataTypeField field, Format node, CalcitePlanCont formattedValue = concat( context, - stringLiteral(rawValue ? "\"" : fieldName + "=\"", context), + stringLiteral(fieldName + "=\"", context), escapeValue(stringValue, context), stringLiteral("\"", context)); } @@ -262,10 +260,10 @@ private RexNode escapeArrayValues( } private String formatFieldName(String fieldName) { - if (fieldName.matches("[A-Za-z][A-Za-z0-9_]*")) { + if (fieldName.matches("[A-Za-z_@][A-Za-z0-9_@-]*(\\.[A-Za-z_@][A-Za-z0-9_@-]*)*")) { return fieldName; } - return "\"" + fieldName.replace("\\", "\\\\").replace("\"", "\\\"") + "\""; + return "`" + fieldName.replace("`", "``") + "`"; } private RexNode compactArray(CalcitePlanContext context, List values) { diff --git a/core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java b/core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java new file mode 100644 index 00000000000..8b911e04c16 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java @@ -0,0 +1,157 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite; + +import com.google.common.collect.ImmutableSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.IntStream; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.CorrelationId; +import org.apache.calcite.rel.core.Filter; +import org.apache.calcite.rel.core.JoinRelType; +import org.apache.calcite.rel.logical.LogicalCorrelate; +import org.apache.calcite.rel.logical.LogicalFilter; +import org.apache.calcite.rel.logical.LogicalJoin; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCorrelVariable; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexShuttle; +import org.apache.calcite.rex.RexSubQuery; +import org.apache.calcite.rex.RexVisitorImpl; +import org.apache.calcite.util.ImmutableBitSet; +import org.opensearch.sql.calcite.plan.DynamicQueryStringPushDown; +import org.opensearch.sql.exception.SemanticCheckException; + +/** Rewrites runtime search scalar subqueries into a standard Calcite correlate. */ +public final class RuntimeSearchCorrelator { + + private RuntimeSearchCorrelator() {} + + /** + * Moves scalar subqueries out of a {@code query_string} filter and makes them the left input of a + * correlate. The right scan consumes their single-row output through a correlation variable, so + * its OpenSearch request is not built until every subquery result is available. + */ + public static RelNode correlate( + RelNode filterNode, SearchPredicateCompiler searchPredicateCompiler) { + if (!(filterNode instanceof Filter filter)) { + throw new IllegalStateException( + "Runtime search query must produce a filter, but got " + + filterNode.getClass().getSimpleName()); + } + List subqueries = findSubqueries(filter.getCondition()); + if (subqueries.isEmpty()) { + return filter; + } + + RexBuilder rexBuilder = filter.getCluster().getRexBuilder(); + RelNode queryInput = combineSubqueries(subqueries, rexBuilder); + CorrelationId correlationId = filter.getCluster().createCorrel(); + RexCorrelVariable correlation = + (RexCorrelVariable) rexBuilder.makeCorrel(queryInput.getRowType(), correlationId); + + Map correlatedInputs = new IdentityHashMap<>(); + int fieldOffset = 0; + for (RexSubQuery subquery : subqueries) { + correlatedInputs.put(subquery, rexBuilder.makeFieldAccess(correlation, fieldOffset)); + fieldOffset += subquery.rel.getRowType().getFieldCount(); + } + + RexNode correlatedCondition = + filter + .getCondition() + .accept( + new RexShuttle() { + @Override + public RexNode visitSubQuery(RexSubQuery subquery) { + RexNode correlatedInput = correlatedInputs.get(subquery); + return correlatedInput == null + ? super.visitSubQuery(subquery) + : correlatedInput; + } + }); + + RelNode right = + pushDownDynamicQueryString( + filter.getInput(), + correlatedCondition, + subqueries.stream().map(correlatedInputs::get).toList(), + correlationId, + searchPredicateCompiler); + LogicalCorrelate correlate = + LogicalCorrelate.create( + queryInput, + right, + correlationId, + ImmutableBitSet.range(queryInput.getRowType().getFieldCount()), + JoinRelType.INNER); + + int leftFieldCount = queryInput.getRowType().getFieldCount(); + List parentFields = + IntStream.range(0, filter.getRowType().getFieldCount()) + .mapToObj(index -> (RexNode) rexBuilder.makeInputRef(correlate, leftFieldCount + index)) + .toList(); + return LogicalProject.create( + correlate, List.of(), parentFields, filter.getRowType().getFieldNames()); + } + + private static List findSubqueries(RexNode condition) { + List subqueries = new ArrayList<>(); + Set seen = Collections.newSetFromMap(new IdentityHashMap<>()); + condition.accept( + new RexVisitorImpl(true) { + @Override + public Void visitSubQuery(RexSubQuery subquery) { + if (seen.add(subquery)) { + if (subquery.rel.getRowType().getFieldCount() != 1) { + throw new SemanticCheckException( + "Implicit format subsearch must return exactly one column"); + } + subqueries.add(subquery); + } + return null; + } + }); + return subqueries; + } + + private static RelNode combineSubqueries(List subqueries, RexBuilder rexBuilder) { + RelNode result = subqueries.getFirst().rel; + for (int i = 1; i < subqueries.size(); i++) { + result = + LogicalJoin.create( + result, + subqueries.get(i).rel, + List.of(), + rexBuilder.makeLiteral(true), + Set.of(), + JoinRelType.INNER); + } + return result; + } + + private static RelNode pushDownDynamicQueryString( + RelNode input, + RexNode condition, + List runtimePredicates, + CorrelationId correlationId, + SearchPredicateCompiler searchPredicateCompiler) { + if (input instanceof DynamicQueryStringPushDown pushDown) { + if (searchPredicateCompiler == null) { + throw new IllegalStateException("No PPL search predicate compiler is configured"); + } + return pushDown.pushDownDynamicQueryString( + condition, runtimePredicates, searchPredicateCompiler); + } + return LogicalFilter.create(input, condition, ImmutableSet.of(correlationId)); + } +} diff --git a/core/src/main/java/org/opensearch/sql/calcite/plan/DynamicQueryStringPushDown.java b/core/src/main/java/org/opensearch/sql/calcite/plan/DynamicQueryStringPushDown.java new file mode 100644 index 00000000000..43bb493e1b8 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/calcite/plan/DynamicQueryStringPushDown.java @@ -0,0 +1,28 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.plan; + +import java.util.List; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rex.RexNode; +import org.opensearch.sql.calcite.SearchPredicateCompiler; + +/** + * A scan that can defer a {@code query_string} predicate until its correlated input is available. + */ +public interface DynamicQueryStringPushDown { + + /** + * Pushes a dynamic {@code query_string} filter into this scan. + * + * @param condition filter condition containing a correlated query argument + * @param runtimePredicates correlated PPL predicate values that require runtime compilation + * @param compiler compiler from PPL search syntax to OpenSearch query-string syntax + * @return a scan that builds the OpenSearch request after the correlated value is available + */ + RelNode pushDownDynamicQueryString( + RexNode condition, List runtimePredicates, SearchPredicateCompiler compiler); +} diff --git a/core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java b/core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java deleted file mode 100644 index 030a2cad7e5..00000000000 --- a/core/src/main/java/org/opensearch/sql/executor/DynamicSearchExecutor.java +++ /dev/null @@ -1,101 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.executor; - -import java.util.List; -import org.apache.calcite.rel.RelNode; -import org.apache.calcite.rex.RexSubQuery; -import org.opensearch.sql.calcite.DynamicSearchPlanBinder; -import org.opensearch.sql.calcite.SearchPredicateCompiler; -import org.opensearch.sql.common.response.ResponseListener; -import org.opensearch.sql.data.model.ExprValue; -import org.opensearch.sql.exception.SemanticCheckException; - -/** Executes and binds implicit format subqueries before running the parent search. */ -public final class DynamicSearchExecutor { - - private DynamicSearchExecutor() {} - - @FunctionalInterface - public interface PlanExecutor { - void execute(RelNode plan, ResponseListener listener); - } - - public static void execute( - RelNode plan, - SearchPredicateCompiler compiler, - PlanExecutor subqueryExecutor, - PlanExecutor finalExecutor, - ResponseListener listener) { - bind( - plan, - compiler, - subqueryExecutor, - new ResponseListener<>() { - @Override - public void onResponse(RelNode bound) { - finalExecutor.execute(bound, listener); - } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); - } - - /** Executes all implicit format subqueries and returns a plan with literal search predicates. */ - public static void bind( - RelNode plan, - SearchPredicateCompiler compiler, - PlanExecutor subqueryExecutor, - ResponseListener listener) { - var dynamic = DynamicSearchPlanBinder.find(plan); - if (dynamic.isEmpty()) { - listener.onResponse(plan); - return; - } - - RexSubQuery subquery = dynamic.get(); - subqueryExecutor.execute( - subquery.rel, - new ResponseListener<>() { - @Override - public void onResponse(ExecutionEngine.QueryResponse response) { - try { - String formatted = extractScalarString(response); - String compiled = compiler.compile(formatted); - RelNode bound = DynamicSearchPlanBinder.bind(plan, subquery, compiled); - bind(bound, compiler, subqueryExecutor, listener); - } catch (Exception e) { - listener.onFailure(e); - } - } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); - } - - private static String extractScalarString(ExecutionEngine.QueryResponse response) { - List rows = response.getResults(); - if (rows.size() != 1 || rows.getFirst().tupleValue().size() != 1) { - throw new SemanticCheckException( - "Implicit format subsearch must return exactly one row and one column"); - } - ExprValue value = rows.getFirst().tupleValue().values().iterator().next(); - if (value.isNull() || value.isMissing()) { - throw new SemanticCheckException("Implicit format subsearch returned a null search string"); - } - Object raw = value.value(); - if (!(raw instanceof String string)) { - throw new SemanticCheckException("Implicit format subsearch must return a string"); - } - return string; - } -} diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index 88b127a1272..cff0552d2b1 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -47,7 +47,6 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.CalciteRelNodeVisitor; -import org.opensearch.sql.calcite.DynamicSearchPlanBinder; import org.opensearch.sql.calcite.OpenSearchSchema; import org.opensearch.sql.calcite.SearchPredicateCompiler; import org.opensearch.sql.calcite.SysLimit; @@ -65,7 +64,6 @@ import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.exception.CalciteUnsupportedException; import org.opensearch.sql.exception.NonFallbackCalciteException; -import org.opensearch.sql.executor.analytics.TimewrapSignals; import org.opensearch.sql.monitor.profile.MetricName; import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.ProfileMetric; @@ -88,6 +86,7 @@ public class QueryService { private DataSourceService dataSourceService; private Settings settings; private ExecutionDispatcher executionDispatcher = new DirectExecutionDispatcher(); + private SearchPredicateCompiler searchPredicateCompiler; public QueryService( Analyzer analyzer, @@ -101,7 +100,8 @@ public QueryService( planner, dataSourceService, settings, - new DirectExecutionDispatcher()); + new DirectExecutionDispatcher(), + null); } public QueryService( @@ -111,12 +111,31 @@ public QueryService( DataSourceService dataSourceService, Settings settings, ExecutionDispatcher executionDispatcher) { + this( + analyzer, + executionEngine, + planner, + dataSourceService, + settings, + executionDispatcher, + null); + } + + public QueryService( + Analyzer analyzer, + ExecutionEngine executionEngine, + Planner planner, + DataSourceService dataSourceService, + Settings settings, + ExecutionDispatcher executionDispatcher, + SearchPredicateCompiler searchPredicateCompiler) { this.analyzer = analyzer; this.executionEngine = executionEngine; this.planner = planner; this.dataSourceService = dataSourceService; this.settings = settings; this.executionDispatcher = executionDispatcher; + this.searchPredicateCompiler = searchPredicateCompiler; } @Getter(lazy = true) @@ -214,6 +233,7 @@ public void executeWithCalcite( buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType); context.setHighlightConfig(highlightConfig); + context.setSearchPredicateCompiler(searchPredicateCompiler); // Wrap analyze with ANALYZING stage tracking RelNode relNode = @@ -300,6 +320,7 @@ public void explainWithCalcite( CalcitePlanContext.create( buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType); context.setHighlightConfig(highlightConfig); + context.setSearchPredicateCompiler(searchPredicateCompiler); context.run( () -> { RelNode relNode = analyze(plan, context); diff --git a/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java b/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java index 7812d8d0366..84fa6730d52 100644 --- a/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/analytics/AnalyticsEngineProfileIT.java @@ -196,21 +196,6 @@ public void testStaticSearchStillRunsThroughAnalyticsEngine() throws IOException assertEquals("alice", result.getJSONArray("datarows").getJSONArray(0).getString(0)); } - @Test - public void testImplicitFormatSubsearchRunsThroughAnalyticsEngine() throws IOException { - ensureSetup(); - JSONObject result = - executeWithProfile( - "search source=" - + INDEX - + " [ search source=" - + INDEX - + " name=alice | fields name | head 1 ] | fields name", - "/_plugins/_ppl"); - - assertEquals("alice", result.getJSONArray("datarows").getJSONArray(0).getString(0)); - } - @Test public void testPplExplainReturnsOnlyPlan() throws IOException { ensureSetup(); @@ -227,7 +212,7 @@ public void testPplExplainReturnsOnlyPlan() throws IOException { } @Test - public void testImplicitFormatExplainBindsRuntimePredicate() throws IOException { + public void testImplicitFormatExplainShowsCorrelatedPlan() throws IOException { ensureSetup(); Request request = new Request("POST", "/_plugins/_ppl/_explain"); request.setJsonEntity( @@ -241,7 +226,8 @@ public void testImplicitFormatExplainBindsRuntimePredicate() throws IOException JSONObject result = new JSONObject(entityAsString(response)); String logical = result.getJSONObject("calcite").get("logical").toString(); - assertTrue(logical, logical.contains("name:alice")); + assertTrue(logical, logical.contains("LogicalCorrelate")); + assertTrue(logical, logical.contains("dynamicQueryString=$cor")); assertFalse(logical, logical.contains("SCALAR_QUERY")); } diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java index 1bdef8bbcfd..5134fd2e7ed 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteFormatCommandIT.java @@ -37,6 +37,17 @@ public void testDefaultFormat() throws IOException { verifyDataRows(result, rows("( ( account_number=\"1\" AND firstname=\"Amber JOHnny\" ) )")); } + @Test + public void testMetadataFieldIsIgnored() throws IOException { + JSONObject result = + executeQuery( + "source=" + + TEST_INDEX_BANK + + " | where account_number=1 | fields firstname, `_id` | format"); + + verifyDataRows(result, rows("( ( firstname=\"Amber JOHnny\" ) )")); + } + @Test public void testMultivalueAndCustomDelimiters() throws IOException { JSONObject result = diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java index b53c5d3f8e6..79b3cae8a83 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java @@ -53,6 +53,36 @@ public void testImplicitFormatExecutesOrdinaryFields() throws IOException { verifyDataRows(result, rows(1), rows(6)); } + @Test + public void testMultipleImplicitFormatSubsearchesFeedOneParentSearch() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_BANK + + " [ search source=" + + TEST_INDEX_BANK + + " account_number=1 | fields account_number ] OR [ search source=" + + TEST_INDEX_BANK + + " account_number=6 | fields account_number ]" + + " | fields account_number | sort account_number"); + + verifyDataRows(result, rows(1), rows(6)); + } + + @Test + public void testParentPipelineAggregatesDynamicSearchResults() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_BANK + + " [ search source=" + + TEST_INDEX_BANK + + " account_number=1 OR account_number=6 | fields account_number ]" + + " | stats count() as matched"); + + verifyDataRows(result, rows(2)); + } + @Test public void testImplicitFormatUnderNotExpression() throws IOException { JSONObject result = @@ -68,7 +98,7 @@ public void testImplicitFormatUnderNotExpression() throws IOException { } @Test - public void testImplicitFormatExplainUsesBoundPredicate() throws IOException { + public void testImplicitFormatExplainShowsCorrelatedDynamicScan() throws IOException { String result = explainQueryToString( "search source=" @@ -78,7 +108,8 @@ public void testImplicitFormatExplainUsesBoundPredicate() throws IOException { + " | head 1 | eval search='account_number=1' | fields search ]" + " | fields account_number"); - assertTrue(result, result.contains("account_number:1")); + assertTrue(result, result.contains("LogicalCorrelate")); + assertTrue(result, result.contains("dynamicQueryString=$cor")); assertFalse(result, result.contains("SCALAR_QUERY")); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/EnumerableIndexScanRule.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/EnumerableIndexScanRule.java index 4a6ac7fff7b..23b2863241d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/EnumerableIndexScanRule.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/EnumerableIndexScanRule.java @@ -35,7 +35,8 @@ protected EnumerableIndexScanRule(Config config) { @Override public boolean matches(RelOptRuleCall call) { CalciteLogicalIndexScan scan = call.rel(0); - return scan.getVariablesSet().isEmpty(); + return scan.getVariablesSet().isEmpty() + || scan.getPushDownContext().getDynamicQueryString() != null; } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java index f3773cfe246..4c319d6780b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java @@ -14,6 +14,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Supplier; import java.util.stream.Stream; import lombok.Getter; @@ -32,6 +33,7 @@ import org.apache.calcite.rel.RelWriter; import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.CorrelationId; import org.apache.calcite.rel.core.TableScan; import org.apache.calcite.rel.externalize.RelJsonWriter; import org.apache.calcite.rel.externalize.RelWriterImpl; @@ -104,10 +106,21 @@ public RelDataType deriveRowType() { return this.schema; } + @Override + public Set getVariablesSet() { + return pushDownContext.getDynamicQueryString() == null + ? super.getVariablesSet() + : pushDownContext.getDynamicQueryString().correlationIds(); + } + @Override public RelWriter explainTerms(RelWriter pw) { // Build explain string with context and request builder info String explainString = String.valueOf(pushDownContext); + if (pushDownContext.getDynamicQueryString() != null) { + explainString += + ", dynamicQueryString=" + pushDownContext.getDynamicQueryString().queryExpression(); + } if (pw instanceof RelJsonWriter) { // For JSON output, add structured items super.explainTerms(pw); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java index 7ba75e46ba0..1ffdd4b3e00 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java @@ -10,9 +10,11 @@ import org.apache.calcite.adapter.enumerable.EnumerableRelImplementor; import org.apache.calcite.adapter.enumerable.PhysType; import org.apache.calcite.adapter.enumerable.PhysTypeImpl; +import org.apache.calcite.adapter.enumerable.RexToLixTranslator; import org.apache.calcite.linq4j.AbstractEnumerable; import org.apache.calcite.linq4j.Enumerable; import org.apache.calcite.linq4j.Enumerator; +import org.apache.calcite.linq4j.tree.BlockBuilder; import org.apache.calcite.linq4j.tree.Blocks; import org.apache.calcite.linq4j.tree.Expression; import org.apache.calcite.linq4j.tree.Expressions; @@ -27,6 +29,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.checkerframework.checker.nullness.qual.Nullable; +import org.opensearch.index.query.QueryBuilders; import org.opensearch.sql.calcite.plan.Scannable; import org.opensearch.sql.calcite.plan.rule.OpenSearchRules; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; @@ -105,6 +108,23 @@ public Result implement(EnumerableRelImplementor implementor, Prefer pref) { pref.preferArray()); Expression scanOperator = implementor.stash(this, CalciteEnumerableIndexScan.class); + if (pushDownContext.getDynamicQueryString() != null) { + BlockBuilder builder = new BlockBuilder(); + RexToLixTranslator translator = + RexToLixTranslator.forAggregation( + implementor.getTypeFactory(), builder, null, implementor.getConformance()) + .setCorrelates(implementor::getCorrelVariableGetter); + List queryParts = + translator.translateList(pushDownContext.getDynamicQueryString().queryParts()).stream() + .map(expression -> (Expression) Expressions.convert_(expression, String.class)) + .toList(); + builder.add( + Expressions.return_( + null, + Expressions.call( + scanOperator, "scan", Expressions.newArrayInit(String.class, queryParts)))); + return implementor.result(physType, builder.toBlock()); + } return implementor.result(physType, Blocks.toBlock(Expressions.call(scanOperator, "scan"))); } @@ -115,10 +135,19 @@ public Result implement(EnumerableRelImplementor implementor, Prefer pref) { */ @Override public Enumerable<@Nullable Object> scan() { + return scan((String[]) null); + } + + /** Builds a scan whose query-string predicate comes from a correlated single-row input. */ + public Enumerable<@Nullable Object> scan(String @Nullable [] runtimeQueryParts) { return new AbstractEnumerable<>() { @Override public Enumerator enumerator() { OpenSearchRequestBuilder requestBuilder = pushDownContext.createRequestBuilder(); + if (runtimeQueryParts != null) { + String runtimeQuery = buildRuntimeQuery(runtimeQueryParts); + requestBuilder.pushDownFilterForCalcite(QueryBuilders.queryStringQuery(runtimeQuery)); + } return new OpenSearchIndexEnumerator( osIndex.getClient(), getRowType().getFieldNames(), @@ -130,4 +159,16 @@ public Enumerator enumerator() { } }; } + + private String buildRuntimeQuery(String[] queryParts) { + StringBuilder query = new StringBuilder(); + for (int i = 0; i < queryParts.length; i++) { + String part = queryParts[i]; + if (pushDownContext.getDynamicQueryString().runtimePredicateParts().contains(i)) { + part = pushDownContext.getDynamicQueryString().compiler().compile(part); + } + query.append(part); + } + return query.toString(); + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java index 2017437e7bd..4090cbbbbf9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java @@ -34,6 +34,8 @@ import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.rel.type.RelDataTypeField; import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.type.SqlTypeName; @@ -42,6 +44,8 @@ import org.apache.logging.log4j.Logger; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.sql.ast.tree.HighlightConfig; +import org.opensearch.sql.calcite.SearchPredicateCompiler; +import org.opensearch.sql.calcite.plan.DynamicQueryStringPushDown; import org.opensearch.sql.calcite.plan.HighlightPushDown; import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.calcite.utils.PPLHintUtils; @@ -59,6 +63,7 @@ import org.opensearch.sql.opensearch.storage.OpenSearchIndex; import org.opensearch.sql.opensearch.storage.scan.context.AbstractAction; import org.opensearch.sql.opensearch.storage.scan.context.AggSpec; +import org.opensearch.sql.opensearch.storage.scan.context.DynamicQueryStringSpec; import org.opensearch.sql.opensearch.storage.scan.context.FilterDigest; import org.opensearch.sql.opensearch.storage.scan.context.LimitDigest; import org.opensearch.sql.opensearch.storage.scan.context.OSRequestBuilderAction; @@ -70,7 +75,8 @@ /** The logical relational operator representing a scan of an OpenSearchIndex type. */ @Getter -public class CalciteLogicalIndexScan extends AbstractCalciteIndexScan implements HighlightPushDown { +public class CalciteLogicalIndexScan extends AbstractCalciteIndexScan + implements DynamicQueryStringPushDown, HighlightPushDown { private static final Logger LOG = LogManager.getLogger(CalciteLogicalIndexScan.class); public CalciteLogicalIndexScan( @@ -195,6 +201,46 @@ public AbstractRelNode pushDownFilter(Filter filter) { return null; } + @Override + public RelNode pushDownDynamicQueryString( + RexNode condition, List runtimePredicates, SearchPredicateCompiler compiler) { + RexNode queryExpression = queryArgument(condition); + if (queryExpression == null) { + throw new IllegalArgumentException( + "Dynamic query_string condition is missing its query argument"); + } + + CalciteLogicalIndexScan newScan = copy(); + newScan + .getPushDownContext() + .setDynamicQueryString( + DynamicQueryStringSpec.create(queryExpression, runtimePredicates, compiler)); + newScan + .getPushDownContext() + .add( + PushDownType.FILTER, + new FilterDigest(0, condition), + (OSRequestBuilderAction) requestBuilder -> {}); + return newScan; + } + + private RexNode queryArgument(RexNode condition) { + if (!(condition instanceof RexCall queryString) + || !queryString.getOperator().getName().equalsIgnoreCase("query_string")) { + return null; + } + for (RexNode operand : queryString.getOperands()) { + if (operand instanceof RexCall map && map.getOperands().size() >= 2) { + RexNode key = map.getOperands().getFirst(); + if (key instanceof RexLiteral literal + && "query".equalsIgnoreCase(literal.getValueAs(String.class))) { + return map.getOperands().get(1); + } + } + } + return null; + } + /** * Build schema for the current scan. Schema is the combination of all index fields and nested * fields in index fields. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/DynamicQueryStringSpec.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/DynamicQueryStringSpec.java new file mode 100644 index 00000000000..7247ce578a8 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/DynamicQueryStringSpec.java @@ -0,0 +1,73 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage.scan.context; + +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import org.apache.calcite.rel.core.CorrelationId; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexCorrelVariable; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexVisitorImpl; +import org.opensearch.sql.calcite.SearchPredicateCompiler; + +/** Runtime query-string input consumed by a correlated OpenSearch scan. */ +public record DynamicQueryStringSpec( + RexNode queryExpression, + List queryParts, + Set runtimePredicateParts, + SearchPredicateCompiler compiler) { + + /** Splits concatenation so only subsearch outputs are parsed as PPL predicates. */ + public static DynamicQueryStringSpec create( + RexNode queryExpression, List runtimePredicates, SearchPredicateCompiler compiler) { + List parts = new ArrayList<>(); + flattenConcatenation(queryExpression, parts); + Set predicateParts = new LinkedHashSet<>(); + for (int i = 0; i < parts.size(); i++) { + RexNode part = parts.get(i); + if (runtimePredicates.stream().anyMatch(predicate -> predicate == part)) { + predicateParts.add(i); + } + } + if (predicateParts.size() != runtimePredicates.size()) { + throw new IllegalArgumentException( + "Every runtime search predicate must be a query-string concatenation part"); + } + return new DynamicQueryStringSpec( + queryExpression, List.copyOf(parts), Set.copyOf(predicateParts), compiler); + } + + /** Correlation variables referenced by the runtime query expression. */ + public Set correlationIds() { + Set ids = new LinkedHashSet<>(); + for (RexNode queryPart : queryParts) { + queryPart.accept( + new RexVisitorImpl(true) { + @Override + public Void visitCorrelVariable(RexCorrelVariable correlVariable) { + ids.add(correlVariable.id); + return null; + } + }); + } + return Set.copyOf(ids); + } + + private static void flattenConcatenation(RexNode expression, List parts) { + if (expression instanceof RexCall call + && (call.getOperator().getName().equalsIgnoreCase("concat") + || call.getOperator().getName().equals("||"))) { + for (RexNode operand : call.getOperands()) { + flattenConcatenation(operand, parts); + } + return; + } + parts.add(expression); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownContext.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownContext.java index a622f948efb..8ee6ca300a8 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownContext.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownContext.java @@ -29,6 +29,7 @@ public class PushDownContext extends AbstractCollection { private boolean isAggregatePushed = false; @Setter private AggSpec aggSpec; + @Setter private DynamicQueryStringSpec dynamicQueryString; // Records the start pos of the query, which is updated by new added limit operations. private int startFrom = 0; @@ -53,6 +54,7 @@ public PushDownContext clone() { newContext.add(operation); } newContext.aggSpec = aggSpec; + newContext.dynamicQueryString = dynamicQueryString; return newContext; } @@ -69,6 +71,7 @@ public PushDownContext cloneWithoutSort() { } } newContext.aggSpec = aggSpec == null ? null : aggSpec.withoutBucketSort(); + newContext.dynamicQueryString = dynamicQueryString; return newContext; } @@ -114,6 +117,7 @@ public PushDownContext cloneForAggregate(Aggregate aggregate, @Nullable Project } newContext.add(operation); } + newContext.dynamicQueryString = dynamicQueryString; return newContext; } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 2bebc645d3b..40a2a60a229 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -129,7 +129,13 @@ public QueryPlanFactory queryPlanFactory( new ThreadPoolExecutionDispatcher(nodeClient.threadPool(), settings); QueryService queryService = new QueryService( - analyzer, executionEngine, planner, dataSourceService, settings, executionDispatcher); + analyzer, + executionEngine, + planner, + dataSourceService, + settings, + executionDispatcher, + new PPLSearchPredicateCompiler()); return new QueryPlanFactory(queryService); } } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index 797b91a4ac6..468895794e7 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -40,11 +40,9 @@ import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit; import org.opensearch.sql.common.response.ResponseListener; -import org.opensearch.sql.executor.DynamicSearchExecutor; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.analytics.AnalyticsExecutionEngine; -import org.opensearch.sql.executor.analytics.TimewrapSignals; import org.opensearch.sql.lang.LangSpec; import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.QueryProfiling; @@ -229,41 +227,23 @@ private void doExecute( ActionListener closingListener = wrapWithContextClose(context, listener); try { + CalcitePlanContext planContext = context.getPlanContext(); + planContext.setSearchPredicateCompiler(searchPredicateCompiler); UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context); RelNode plan = planner.plan(query); - CalcitePlanContext planContext = context.getPlanContext(); // PPL fetch_size caps the response to N rows (no cursor) — the V2 path attaches // a `head N` in AstStatementBuilder; the unified path parses only the query // string, so apply the equivalent top-level limit here before the system cap. plan = addFetchSizeLimit(plan, planContext, fetchSize); plan = addQuerySizeLimit(plan, planContext); - // An init-plan must not consume or apply parent-result post-processing. Carry - // these signals across the async subquery callback and install them only for - // the final parent execution. - TimewrapSignals finalTimewrap = TimewrapSignals.captureAndClear(); ResponseListener queryListener = createQueryListener(queryType, profileCtx, closingListener); - DynamicSearchExecutor.PlanExecutor subqueryExecutor = - (subquery, subqueryListener) -> - analyticsEngine.execute( - subquery, planContext, queryCtx, subqueryListener); - DynamicSearchExecutor.PlanExecutor finalExecutor = - profiling - ? (bound, finalListener) -> { - finalTimewrap.install(); - analyticsEngine.executeWithProfile( - bound, planContext, queryCtx, finalListener); - } - : (bound, finalListener) -> { - finalTimewrap.install(); - analyticsEngine.execute(bound, planContext, queryCtx, finalListener); - }; - DynamicSearchExecutor.execute( - plan, - searchPredicateCompiler, - subqueryExecutor, - finalExecutor, - queryListener); + if (profiling) { + analyticsEngine.executeWithProfile( + plan, planContext, queryCtx, queryListener); + } else { + analyticsEngine.execute(plan, planContext, queryCtx, queryListener); + } } catch (Exception e) { closingListener.onFailure(e); } finally { @@ -315,27 +295,12 @@ private void doExplain( ResponseListener closingListener = wrapWithContextClose(context, listener); try { + CalcitePlanContext planContext = context.getPlanContext(); + planContext.setSearchPredicateCompiler(searchPredicateCompiler); UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context); RelNode plan = planner.plan(query); - CalcitePlanContext planContext = context.getPlanContext(); plan = addQuerySizeLimit(plan, planContext); - DynamicSearchExecutor.bind( - plan, - searchPredicateCompiler, - (subquery, subqueryListener) -> - analyticsEngine.execute( - subquery, planContext, queryCtx, subqueryListener), - new ResponseListener<>() { - @Override - public void onResponse(RelNode bound) { - analyticsEngine.explain(bound, mode, planContext, closingListener); - } - - @Override - public void onFailure(Exception e) { - closingListener.onFailure(e); - } - }); + analyticsEngine.explain(plan, mode, planContext, closingListener); } catch (Exception e) { closingListener.onFailure(e); } finally { diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java index d7eb977850a..e4f7e755ebf 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFormatTest.java @@ -145,25 +145,34 @@ public void testBackslashesInValuesAreEscaped() { } @Test - public void testInternalFieldsAreIgnored() { + public void testUserFieldWithLeadingUnderscoreIsFormatted() { withPPLQuery( "source=EMP | head 1 | eval _private='hidden', visible='shown' " + "| fields _private, visible | format") - .expectResult("search=( ( visible=\"shown\" ) )\n"); + .expectResult("search=( ( _private=\"hidden\" AND visible=\"shown\" ) )\n"); } @Test - public void testSearchAndQueryFieldsDropFieldName() { + public void testSearchAndQueryAreOrdinaryFieldsForExplicitFormat() { withPPLQuery( "source=EMP | head 1 | eval search='status=200', query='method=GET', a='x' " + "| fields search, query, a | format") - .expectResult("search=( ( a=\"x\" AND \"method=GET\" AND \"status=200\" ) )\n"); + .expectResult( + "search=( ( a=\"x\" AND query=\"method=GET\" AND search=\"status=200\" ) )\n"); } @Test - public void testSpecialCharactersInFieldNameAreQuoted() { + public void testQualifiedFieldNameUsesPplSyntax() { withPPLQuery("source=EMP | head 1 | eval a.b='x' | fields a.b | format") - .expectResult("search=( ( \"a.b\"=\"x\" ) )\n"); + .expectResult("search=( ( a.b=\"x\" ) )\n"); + } + + @Test + public void testFieldNameWithSpacesUsesPplIdentifierQuoting() { + withPPLQuery( + "source=EMP | head 1 | rename ENAME as `display name` " + + "| fields `display name` | format") + .expectResult("search=( ( `display name`=\"SMITH\" ) )\n"); } @Test diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java index a8e69454d46..676ac5c2f36 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java @@ -7,30 +7,22 @@ import static org.junit.Assert.assertThrows; -import java.sql.PreparedStatement; -import java.util.Collections; import java.util.List; -import java.util.Map; import java.util.concurrent.atomic.AtomicReference; import org.apache.calcite.plan.RelOptUtil; import org.apache.calcite.plan.RelTraitDef; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelVisitor; +import org.apache.calcite.rel.logical.LogicalCorrelate; +import org.apache.calcite.rel.logical.LogicalJoin; import org.apache.calcite.schema.SchemaPlus; import org.apache.calcite.sql.parser.SqlParser; import org.apache.calcite.test.CalciteAssert; import org.apache.calcite.tools.Frameworks; import org.apache.calcite.tools.Programs; -import org.apache.calcite.tools.RelRunners; import org.junit.Assert; import org.junit.Ignore; import org.junit.Test; -import org.opensearch.sql.calcite.DynamicSearchPlanBinder; -import org.opensearch.sql.common.response.ResponseListener; -import org.opensearch.sql.data.model.ExprValueUtils; -import org.opensearch.sql.executor.DynamicSearchExecutor; -import org.opensearch.sql.executor.ExecutionEngine; -import org.opensearch.sql.executor.pagination.Cursor; -import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler; public class CalcitePPLSearchTest extends CalcitePPLAbstractTest { public CalcitePPLSearchTest() { @@ -69,69 +61,60 @@ public void testSearchWithFilter() { } @Test - public void testSearchWithImplicitFormatSubsearchUsesScalarQuery() { + public void testSearchWithImplicitFormatSubsearchUsesCorrelate() { RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); - RelNode subquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + LogicalCorrelate correlate = findCorrelate(root); String logical = RelOptUtil.toString(root); - Assert.assertTrue(logical, logical.contains("SCALAR_QUERY")); + + Assert.assertTrue(logical, logical.contains("LogicalCorrelate")); + Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); Assert.assertTrue(logical, logical.contains("ARRAY_AGG")); Assert.assertTrue(logical, logical.contains("query_string")); - Assert.assertNotSame(root.getCluster(), subquery.getCluster()); + Assert.assertSame(root.getCluster(), correlate.getLeft().getCluster()); + Assert.assertSame(root.getCluster(), correlate.getRight().getCluster()); } @Test - public void testBoundImplicitFormatHasSameStaticQueryStringShape() { + public void testImplicitFormatQueryStringReferencesCorrelatedValue() { RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); - var subquery = DynamicSearchPlanBinder.find(root).orElseThrow(); - - RelNode bound = DynamicSearchPlanBinder.bind(root, subquery, "DEPTNO:10"); - String logical = RelOptUtil.toString(bound); + String logical = RelOptUtil.toString(root); Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - Assert.assertTrue(logical, logical.contains("query_string(MAP('query', 'DEPTNO:10':VARCHAR))")); + Assert.assertTrue(logical, logical.contains("query_string")); + Assert.assertTrue(logical, logical.contains("$cor")); } @Test - public void testBoundImplicitFormatCombinesWithStaticSearchPredicate() { + public void testImplicitFormatCombinesWithStaticSearchPredicate() { RelNode root = getRelNode("search source=EMP JOB=CLERK [ search source=DEPT | fields DEPTNO | head 1 ]"); - var subquery = DynamicSearchPlanBinder.find(root).orElseThrow(); - - RelNode bound = DynamicSearchPlanBinder.bind(root, subquery, "DEPTNO:10"); - String logical = RelOptUtil.toString(bound); + String logical = RelOptUtil.toString(root); Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - Assert.assertTrue( - logical, - logical.contains("query_string(MAP('query', '(JOB:CLERK) AND (DEPTNO:10)':VARCHAR))")); + Assert.assertTrue(logical, logical.contains("JOB:CLERK")); + Assert.assertTrue(logical, logical.contains("$cor")); + Assert.assertTrue(logical, logical.contains("AND")); } @Test - public void testMultipleImplicitFormatSubsearchesBindSequentially() { + public void testMultipleImplicitFormatSubsearchesShareOneCorrelate() { RelNode root = getRelNode( "search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ] OR " + "[ search source=EMP | fields JOB | head 1 ]"); + LogicalCorrelate correlate = findCorrelate(root); + String logical = RelOptUtil.toString(root); - RelNode firstBound = - DynamicSearchPlanBinder.bind( - root, DynamicSearchPlanBinder.find(root).orElseThrow(), "DEPTNO:10"); - Assert.assertTrue(DynamicSearchPlanBinder.find(firstBound).isPresent()); - RelNode fullyBound = - DynamicSearchPlanBinder.bind( - firstBound, DynamicSearchPlanBinder.find(firstBound).orElseThrow(), "JOB:CLERK"); - String logical = RelOptUtil.toString(fullyBound); - - Assert.assertTrue(DynamicSearchPlanBinder.find(fullyBound).isEmpty()); - Assert.assertTrue( - logical, - logical.contains("query_string(MAP('query', '(DEPTNO:10 OR JOB:CLERK)':VARCHAR))")); + Assert.assertTrue(correlate.getLeft() instanceof LogicalJoin); + Assert.assertEquals(1, countCorrelates(root)); + Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); + Assert.assertTrue(logical, logical.contains("OR")); } @Test public void testImplicitFormatSubsearchProducesSearchPredicateData() { RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); - RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + RelNode formatSubquery = findCorrelate(root).getLeft(); verifyResult(formatSubquery, "search=( ( DEPTNO=\"10\" ) )\n"); } @@ -142,9 +125,10 @@ public void testImplicitFormatDoesNotExecuteMultivalueSearchFieldAsRawPredicate( getRelNode( "search source=EMP [ search source=DEPT | head 1 " + "| eval search=array('DEPTNO=10', 'DEPTNO=20') | fields search ]"); - RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + RelNode formatSubquery = findCorrelate(root).getLeft(); - verifyResult(formatSubquery, "search=( ( ( \"DEPTNO=10\" OR \"DEPTNO=20\" ) ) )\n"); + verifyResult( + formatSubquery, "search=( ( ( search=\"DEPTNO=10\" OR search=\"DEPTNO=20\" ) ) )\n"); } @Test @@ -154,7 +138,7 @@ public void testImplicitFormatUsesFirstScalarSearchValueAndIgnoresOtherFields() "search source=EMP [ search source=DEPT " + "| eval search=if(DEPTNO=10, 'DEPTNO=20', 'DEPTNO=10') " + "| fields search, DEPTNO ]"); - RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + RelNode formatSubquery = findCorrelate(root).getLeft(); verifyResult(formatSubquery, "search=DEPTNO=20\n"); } @@ -165,56 +149,38 @@ public void testImplicitFormatUsesPostFormatEvalResult() { getRelNode( "search source=EMP [ search source=DEPT | where DEPTNO=10 | fields DEPTNO " + "| format | eval search=replace(search, '10', '20') ]"); - RelNode formatSubquery = DynamicSearchPlanBinder.find(root).orElseThrow().rel; + RelNode formatSubquery = findCorrelate(root).getLeft(); verifyResult(formatSubquery, "search=( ( DEPTNO=\"20\" ) )\n"); } - @Test - public void testImplicitFormatExecutesThenBindsParentSearch() { - RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); - AtomicReference executedParent = new AtomicReference<>(); - AtomicReference failure = new AtomicReference<>(); - - DynamicSearchExecutor.execute( - root, - new PPLSearchPredicateCompiler(), - (subquery, listener) -> executeScalarSubquery(subquery, listener), - (bound, listener) -> { - executedParent.set(bound); - listener.onResponse( - new ExecutionEngine.QueryResponse(null, Collections.emptyList(), Cursor.None)); - }, - new ResponseListener<>() { - @Override - public void onResponse(ExecutionEngine.QueryResponse response) {} - - @Override - public void onFailure(Exception e) { - failure.set(e); - } - }); - - Assert.assertNull(failure.get()); - String logical = RelOptUtil.toString(executedParent.get()); - Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - Assert.assertTrue( - logical, logical.contains("query_string(MAP('query', '((DEPTNO:10))':VARCHAR))")); - } - - private void executeScalarSubquery( - RelNode subquery, ResponseListener listener) { - try (PreparedStatement statement = RelRunners.run(subquery); - var resultSet = statement.executeQuery()) { - Assert.assertTrue(resultSet.next()); - listener.onResponse( - new ExecutionEngine.QueryResponse( - null, - List.of(ExprValueUtils.tupleValue(Map.of("search", resultSet.getString(1)))), - Cursor.None)); - } catch (Exception e) { - listener.onFailure(e); - } + private LogicalCorrelate findCorrelate(RelNode root) { + AtomicReference result = new AtomicReference<>(); + new RelVisitor() { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) { + if (node instanceof LogicalCorrelate correlate) { + result.compareAndSet(null, correlate); + } + super.visit(node, ordinal, parent); + } + }.go(root); + Assert.assertNotNull(RelOptUtil.toString(root), result.get()); + return result.get(); + } + + private int countCorrelates(RelNode root) { + int[] count = {0}; + new RelVisitor() { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) { + if (node instanceof LogicalCorrelate) { + count[0]++; + } + super.visit(node, ordinal, parent); + } + }.go(root); + return count[0]; } @Ignore("Fields used in search commands are not validated. Enable after fixing it.") diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java index 042e1457edc..aa5e92aa996 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java @@ -27,6 +27,11 @@ public void testCompileRuntimeSearchFieldValue() { "(status:>=500 OR host:api\\-01)", compiler.compile("status>=500 OR host=\"api-01\"")); } + @Test + public void testCompileBacktickQuotedFieldName() { + assertEquals("display\\ name:value", compiler.compile("`display name`=\"value\"")); + } + @Test public void testCompileEmptyFormatResultAsMatchNone() { assertEquals("*:* AND NOT *:*", compiler.compile("NOT ()")); From 5ce9b9b31ddade3ee54e2aec0ae2f9c5d0bc2543 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 23 Jul 2026 18:24:42 +0800 Subject: [PATCH 05/10] Extract format planner helper field constants Signed-off-by: Songkan Tang --- .../opensearch/sql/calcite/FormatPlanner.java | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java index af20ebde1d6..66b5e546d3c 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java +++ b/core/src/main/java/org/opensearch/sql/calcite/FormatPlanner.java @@ -29,6 +29,14 @@ /** Lowers the PPL {@code format} command to Calcite projects and a global aggregation. */ public class FormatPlanner { + private static final String SEARCH_FIELD = "search"; + private static final String RAW_SEARCH_FIELD = "__raw_search"; + private static final String RAW_SEARCH_COUNT_FIELD = "__raw_search_count"; + private static final String RAW_SEARCH_VALUE_FIELD = "__raw_search_value"; + private static final String FORMAT_ROW_FIELD = "__format_row"; + private static final String FORMAT_ROWS_FIELD = "__format_rows"; + private static final String FORMAT_ORDER_FIELD_PREFIX = "__format_order_"; + /** Builds a single-row relation containing the formatted search expression. */ public RelNode plan(Format node, CalcitePlanContext context) { RelBuilder builder = context.relBuilder; @@ -43,13 +51,13 @@ public RelNode plan(Format node, CalcitePlanContext context) { .toList(); if (fields.isEmpty()) { - builder.values(new String[] {"search"}, node.getEmptyString()); + builder.values(new String[] {SEARCH_FIELD}, node.getEmptyString()); return builder.peek(); } Optional scalarSearchField = fields.stream() - .filter(field -> field.getName().equals("search")) + .filter(field -> field.getName().equals(SEARCH_FIELD)) .filter( field -> !SqlTypeUtil.isArray(field.getType()) @@ -64,7 +72,7 @@ public RelNode plan(Format node, CalcitePlanContext context) { RexNode nonEmptyRow = formatRow(formattedFields, node, context); aggregateRows(nonEmptyRow, context); - builder.project(List.of(formatAggregatedRows(node, context)), List.of("search"), true); + builder.project(List.of(formatAggregatedRows(node, context)), List.of(SEARCH_FIELD), true); return builder.peek(); } @@ -90,27 +98,27 @@ private RelNode planImplicitSearchField( .toList(); RexNode fallbackRow = formatRow(fallbackFields, node, context); - builder.project(List.of(rawSearch, fallbackRow), List.of("__raw_search", "__format_row")); - RexNode rawSearchRef = builder.field("__raw_search"); - RexNode rowRef = builder.field("__format_row"); + builder.project(List.of(rawSearch, fallbackRow), List.of(RAW_SEARCH_FIELD, FORMAT_ROW_FIELD)); + RexNode rawSearchRef = builder.field(RAW_SEARCH_FIELD); + RexNode rowRef = builder.field(FORMAT_ROW_FIELD); builder.aggregate( builder.groupKey(), - builder.aggregateCall(SqlStdOperatorTable.COUNT, rawSearchRef).as("__raw_search_count"), - builder.aggregateCall(SqlStdOperatorTable.MAX, rawSearchRef).as("__raw_search_value"), - builder.aggregateCall(SqlLibraryOperators.ARRAY_AGG, rowRef).as("__format_rows")); + builder.aggregateCall(SqlStdOperatorTable.COUNT, rawSearchRef).as(RAW_SEARCH_COUNT_FIELD), + builder.aggregateCall(SqlStdOperatorTable.MAX, rawSearchRef).as(RAW_SEARCH_VALUE_FIELD), + builder.aggregateCall(SqlLibraryOperators.ARRAY_AGG, rowRef).as(FORMAT_ROWS_FIELD)); RexNode hasRawSearch = builder.call( SqlStdOperatorTable.GREATER_THAN, - builder.field("__raw_search_count"), + builder.field(RAW_SEARCH_COUNT_FIELD), builder.literal(0)); RexNode result = builder.call( SqlStdOperatorTable.CASE, hasRawSearch, - builder.field("__raw_search_value"), + builder.field(RAW_SEARCH_VALUE_FIELD), formatAggregatedRows(node, context)); - builder.project(List.of(result), List.of("search"), true); + builder.project(List.of(result), List.of(SEARCH_FIELD), true); return builder.peek(); } @@ -136,24 +144,24 @@ private void aggregateRows(RexNode nonEmptyRow, CalcitePlanContext context) { RelBuilder builder = context.relBuilder; List ordering = inputOrdering(builder.peek()); if (ordering.isEmpty()) { - builder.project(List.of(nonEmptyRow), List.of("__format_row")); + builder.project(List.of(nonEmptyRow), List.of(FORMAT_ROW_FIELD)); } else { List projections = new ArrayList<>(); projections.add(nonEmptyRow); projections.addAll(builder.fields()); List names = new ArrayList<>(); - names.add("__format_row"); + names.add(FORMAT_ROW_FIELD); for (int i = 0; i < projections.size() - 1; i++) { - names.add("__format_order_" + i); + names.add(FORMAT_ORDER_FIELD_PREFIX + i); } builder.project(projections, names, true); } - RexNode rowRef = builder.field("__format_row"); + RexNode rowRef = builder.field(FORMAT_ROW_FIELD); RelBuilder.AggCall rows = builder.aggregateCall(SqlLibraryOperators.ARRAY_AGG, rowRef); if (!ordering.isEmpty()) { rows = rows.sort(ordering.stream().map(order -> orderExpression(order, builder)).toList()); } - builder.aggregate(builder.groupKey(), rows.as("__format_rows")); + builder.aggregate(builder.groupKey(), rows.as(FORMAT_ROWS_FIELD)); } private List inputOrdering(RelNode input) { @@ -180,7 +188,7 @@ private RexNode orderExpression(RelFieldCollation order, RelBuilder builder) { private RexNode formatAggregatedRows(Format node, CalcitePlanContext context) { RelBuilder builder = context.relBuilder; RexNode nonNullRows = - builder.call(SqlLibraryOperators.ARRAY_COMPACT, builder.field("__format_rows")); + builder.call(SqlLibraryOperators.ARRAY_COMPACT, builder.field(FORMAT_ROWS_FIELD)); RexNode joinedRows = arrayJoin(context, nonNullRows, " " + node.getRowSeparator() + " "); RexNode formatted = concat( From dc5db9af8c1c8f9fc45e1183866d58c021609322 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Thu, 23 Jul 2026 18:53:37 +0800 Subject: [PATCH 06/10] Simplify implicit format search planning Signed-off-by: Songkan Tang --- .../sql/api/UnifiedQueryPlanner.java | 7 +- .../sql/ast/expression/SearchSubquery.java | 2 +- .../org/opensearch/sql/ast/tree/Search.java | 2 +- .../sql/calcite/CalcitePlanContext.java | 2 - .../sql/calcite/CalciteRelNodeVisitor.java | 11 ++- .../opensearch/sql/executor/QueryService.java | 6 +- .../executor/analytics/TimewrapSignals.java | 7 -- .../utils/TimewrapSignalsLeakTest.java | 17 ----- docs/user/ppl/cmd/format.md | 5 +- .../plugin/config/OpenSearchPluginModule.java | 2 +- .../plugin/rest/RestUnifiedQueryAction.java | 76 ++++++------------- .../opensearch/sql/ppl/parser/AstBuilder.java | 2 +- .../parser/PPLSearchPredicateCompiler.java | 5 +- .../ppl/parser/AstSearchExpressionTest.java | 2 +- .../PPLSearchPredicateCompilerTest.java | 2 +- 15 files changed, 54 insertions(+), 94 deletions(-) diff --git a/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java b/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java index 9440833503f..55de9617b8b 100644 --- a/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java +++ b/api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java @@ -30,6 +30,7 @@ import org.opensearch.sql.exception.CalciteUnsupportedException; import org.opensearch.sql.exception.QueryEngineException; import org.opensearch.sql.exception.SemanticCheckException; +import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler; /** * {@code UnifiedQueryPlanner} provides a high-level API for parsing and analyzing queries using the @@ -137,13 +138,15 @@ public RelNode plan(String query) throws Exception { private static class CustomVisitorStrategy implements PlanningStrategy { private final UnifiedQueryContext context; private final UnifiedQueryParser parser; - private final CalciteRelNodeVisitor relNodeVisitor = - new CalciteRelNodeVisitor(new EmptyDataSourceService()); + private final CalciteRelNodeVisitor relNodeVisitor; @SuppressWarnings("unchecked") CustomVisitorStrategy(UnifiedQueryContext context) { this.context = context; this.parser = (UnifiedQueryParser) context.getParser(); + this.relNodeVisitor = + new CalciteRelNodeVisitor( + new EmptyDataSourceService(), PPLSearchPredicateCompiler.INSTANCE); } @Override diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java b/core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java index c7483281eda..ff0139cceeb 100644 --- a/core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java +++ b/core/src/main/java/org/opensearch/sql/ast/expression/SearchSubquery.java @@ -24,7 +24,7 @@ public class SearchSubquery extends SearchExpression { @Override public String toQueryString() { throw new IllegalStateException( - "An implicit search subquery must be bound before it can be converted to query_string"); + "An implicit search subquery must be planned as a correlated runtime input"); } @Override diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Search.java b/core/src/main/java/org/opensearch/sql/ast/tree/Search.java index 854b6b40d93..703674f6039 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Search.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Search.java @@ -36,7 +36,7 @@ public Search(UnresolvedPlan child, String queryString) { this(child, queryString, null); } - /** Creates a search whose query string may need runtime subquery binding. */ + /** Creates a search that may contain an implicit subquery. */ public static Search fromExpression(UnresolvedPlan child, SearchExpression originalExpression) { return new Search( 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 24261ee1e17..c7f3bc373ac 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -68,7 +68,6 @@ public class CalcitePlanContext { ThreadLocal.withInitial(() -> true); @Getter @Setter private HighlightConfig highlightConfig; - @Getter @Setter private SearchPredicateCompiler searchPredicateCompiler; @Getter @Setter private boolean isResolvingJoinCondition = false; @Getter @Setter private boolean isResolvingSubquery = false; @Getter @Setter private boolean inCoalesceFunction = false; @@ -166,7 +165,6 @@ private CalcitePlanContext(CalcitePlanContext parent) { this.rexBuilder = parent.rexBuilder; // Share the same rexBuilder this.functionProperties = parent.functionProperties; this.highlightConfig = parent.highlightConfig; - this.searchPredicateCompiler = parent.searchPredicateCompiler; 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 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 24e7c31b1cd..f64be0ed2a0 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -230,13 +230,21 @@ public class CalciteRelNodeVisitor extends AbstractNodeVisitor listener) @@ -233,7 +235,6 @@ public void executeWithCalcite( buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType); context.setHighlightConfig(highlightConfig); - context.setSearchPredicateCompiler(searchPredicateCompiler); // Wrap analyze with ANALYZING stage tracking RelNode relNode = @@ -320,7 +321,6 @@ public void explainWithCalcite( CalcitePlanContext.create( buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType); context.setHighlightConfig(highlightConfig); - context.setSearchPredicateCompiler(searchPredicateCompiler); context.run( () -> { RelNode relNode = analyze(plan, context); diff --git a/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java b/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java index 162d8dda52d..d7907e93d87 100644 --- a/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java +++ b/core/src/main/java/org/opensearch/sql/executor/analytics/TimewrapSignals.java @@ -41,13 +41,6 @@ public static TimewrapSignals captureAndClear() { return new TimewrapSignals(active, unitName, series); } - /** Installs this snapshot on the current thread for an execution engine to capture. */ - public void install() { - CalcitePlanContext.stripNullColumns.set(active); - CalcitePlanContext.timewrapUnitName.set(unitName); - CalcitePlanContext.timewrapSeries.set(series); - } - /** * Applies the timewrap pivot to {@code response} if this snapshot is from a timewrap query; * otherwise returns it unchanged. The returned response carries over the input's profile/error. diff --git a/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java b/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java index 94133c50d89..eef98cdc4bf 100644 --- a/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java +++ b/core/src/test/java/org/opensearch/sql/calcite/utils/TimewrapSignalsLeakTest.java @@ -7,7 +7,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; @@ -22,7 +21,6 @@ import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.executor.ExecutionEngine.Schema.Column; -import org.opensearch.sql.executor.analytics.TimewrapSignals; /** * Leak guard for the timewrap pivot signals. {@code visitTimewrap} stashes its pivot state in the @@ -81,19 +79,4 @@ public void timewrapSignalsDoNotLeakOntoNextQueryOnSameThread() { "non-timewrap result gained artifact columns after a prior timewrap query"); assertEquals(rows, result.values()); } - - @Test - public void timewrapSignalsCanMoveAcrossInitPlanExecution() { - CalcitePlanContext.stripNullColumns.set(true); - CalcitePlanContext.timewrapUnitName.set("1|day|days|_before"); - CalcitePlanContext.timewrapSeries.set("relative"); - - TimewrapSignals snapshot = TimewrapSignals.captureAndClear(); - assertFalse(TimewrapPivot.isTimewrap()); - - snapshot.install(); - assertTrue(TimewrapPivot.isTimewrap()); - assertEquals("1|day|days|_before", CalcitePlanContext.timewrapUnitName.get()); - assertEquals("relative", CalcitePlanContext.timewrapSeries.get()); - } } diff --git a/docs/user/ppl/cmd/format.md b/docs/user/ppl/cmd/format.md index fed56da4274..20d7fd3ff7d 100644 --- a/docs/user/ppl/cmd/format.md +++ b/docs/user/ppl/cmd/format.md @@ -117,8 +117,9 @@ source=logs ## Limitations - Implicit format requires the Calcite query engine. -- Runtime binding is supported for bracketed subsearches in a parent `search` expression. Dynamic - subsearch interpolation in `eval` or other command positions is not supported. +- Correlated runtime execution is supported for bracketed subsearches in a parent `search` + expression. Dynamic subsearch interpolation in `eval` or other command positions is not + supported. - An implicit result is parsed only as a search predicate. Pipeline command syntax in a generated `search` or `query` value is rejected. - When upstream ordering metadata is available, row collection carries it into the aggregate order diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 40a2a60a229..51191efcc08 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -135,7 +135,7 @@ public QueryPlanFactory queryPlanFactory( dataSourceService, settings, executionDispatcher, - new PPLSearchPredicateCompiler()); + PPLSearchPredicateCompiler.INSTANCE); return new QueryPlanFactory(queryService); } } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java index 468895794e7..26c71060c2e 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java @@ -47,7 +47,6 @@ import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.QueryProfiling; import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; -import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler; import org.opensearch.sql.protocol.response.QueryResult; import org.opensearch.sql.protocol.response.format.ResponseFormatter; import org.opensearch.sql.protocol.response.format.SimpleJsonResponseFormatter; @@ -71,8 +70,6 @@ public class RestUnifiedQueryAction { private final org.opensearch.analytics.EngineContextProvider contextProvider; private final org.opensearch.sql.common.setting.Settings pluginSettings; private final org.opensearch.sql.executor.ExecutionDispatcher executionDispatcher; - private final PPLSearchPredicateCompiler searchPredicateCompiler = - new PPLSearchPredicateCompiler(); public RestUnifiedQueryAction( NodeClient client, @@ -227,23 +224,33 @@ private void doExecute( ActionListener closingListener = wrapWithContextClose(context, listener); try { - CalcitePlanContext planContext = context.getPlanContext(); - planContext.setSearchPredicateCompiler(searchPredicateCompiler); UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context); RelNode plan = planner.plan(query); + CalcitePlanContext planContext = context.getPlanContext(); // PPL fetch_size caps the response to N rows (no cursor) — the V2 path attaches // a `head N` in AstStatementBuilder; the unified path parses only the query // string, so apply the equivalent top-level limit here before the system cap. plan = addFetchSizeLimit(plan, planContext, fetchSize); plan = addQuerySizeLimit(plan, planContext); - ResponseListener queryListener = - createQueryListener(queryType, profileCtx, closingListener); - if (profiling) { - analyticsEngine.executeWithProfile( - plan, planContext, queryCtx, queryListener); - } else { - analyticsEngine.execute(plan, planContext, queryCtx, queryListener); - } + plan = + org.opensearch.sql.calcite.utils.CalciteToolsHelper.optimize( + plan, planContext); + RelNode finalPlan = plan; + Runnable executeTask = + profiling + ? () -> + analyticsEngine.executeWithProfile( + finalPlan, + planContext, + queryCtx, + createQueryListener(queryType, profileCtx, closingListener)) + : () -> + analyticsEngine.execute( + finalPlan, + planContext, + queryCtx, + createQueryListener(queryType, profileCtx, closingListener)); + executionDispatcher.dispatchTask(finalPlan, planContext, executeTask); } catch (Exception e) { closingListener.onFailure(e); } finally { @@ -291,18 +298,14 @@ private void doExplain( () -> { QueryRequestContext queryCtx = withParentTask(contextProvider.getContext(), parentTask); - UnifiedQueryContext context = buildContext(queryType, false, queryCtx); - ResponseListener closingListener = - wrapWithContextClose(context, listener); - try { - CalcitePlanContext planContext = context.getPlanContext(); - planContext.setSearchPredicateCompiler(searchPredicateCompiler); + try (UnifiedQueryContext context = buildContext(queryType, false, queryCtx)) { UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context); RelNode plan = planner.plan(query); + CalcitePlanContext planContext = context.getPlanContext(); plan = addQuerySizeLimit(plan, planContext); - analyticsEngine.explain(plan, mode, planContext, closingListener); + analyticsEngine.explain(plan, mode, planContext, listener); } catch (Exception e) { - closingListener.onFailure(e); + listener.onFailure(e); } finally { // explain plans a timewrap query (visitTimewrap sets thread-locals) but never // executes, so nothing captures-and-clears them — clear here to avoid leaking @@ -486,35 +489,4 @@ private static ActionListener wrapWithContextClose( } }); } - - private static ResponseListener wrapWithContextClose( - UnifiedQueryContext context, ResponseListener delegate) { - return new ResponseListener<>() { - @Override - public void onResponse(T response) { - try { - delegate.onResponse(response); - } finally { - closeContext(context); - } - } - - @Override - public void onFailure(Exception e) { - try { - delegate.onFailure(e); - } finally { - closeContext(context); - } - } - }; - } - - private static void closeContext(UnifiedQueryContext context) { - try { - context.close(); - } catch (Exception e) { - LOG.warn("Failed to close query context", e); - } - } } 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 d1a36a72bed..9662bfb3d8c 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 @@ -230,7 +230,7 @@ public UnresolvedPlan visitSearchFrom(SearchFromContext ctx) { } // Static search expressions are converted immediately. An implicit subsearch remains in the - // expression tree until its formatted result is bound at execution time. + // expression tree so the planner can lower it to a correlated runtime input. Relation relation = (Relation) visitFromClause(ctx.fromClause()); return Search.fromExpression(relation, combined); } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java index 938e644988f..8f15ec24476 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompiler.java @@ -11,10 +11,13 @@ import org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.SearchPredicateContext; /** Uses the PPL search-expression grammar to compile an implicit format result. */ -public class PPLSearchPredicateCompiler implements SearchPredicateCompiler { +public final class PPLSearchPredicateCompiler implements SearchPredicateCompiler { + public static final PPLSearchPredicateCompiler INSTANCE = new PPLSearchPredicateCompiler(); private final PPLSyntaxParser syntaxParser = new PPLSyntaxParser(); + private PPLSearchPredicateCompiler() {} + @Override public String compile(String predicate) { String trimmed = predicate == null ? "" : predicate.trim(); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java index 45ae05912e2..1435bc55b52 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstSearchExpressionTest.java @@ -53,7 +53,7 @@ public void testFieldComparison() { } @Test - public void testImplicitFormatSubsearchRemainsRuntimeBound() { + public void testImplicitFormatSubsearchRemainsStructuredForPlanning() { Node plan = buildPlan("search source=logs [ search source=rules | fields status, host ]"); assertTrue(plan instanceof Search); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java index aa5e92aa996..7ca4d5c2af2 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/PPLSearchPredicateCompilerTest.java @@ -12,7 +12,7 @@ public class PPLSearchPredicateCompilerTest { - private final PPLSearchPredicateCompiler compiler = new PPLSearchPredicateCompiler(); + private final PPLSearchPredicateCompiler compiler = PPLSearchPredicateCompiler.INSTANCE; @Test public void testCompileCanonicalFormatOutput() { From 94b999ce6a38b85b47119c351e26fdbb7ff6b54e Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Mon, 27 Jul 2026 16:13:56 +0800 Subject: [PATCH 07/10] Use complete plans in implicit format tests Signed-off-by: Songkan Tang --- .../sql/ppl/calcite/CalcitePPLSearchTest.java | 93 ++++++++++--------- 1 file changed, 49 insertions(+), 44 deletions(-) diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java index 676ac5c2f36..0084c51f734 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java @@ -14,7 +14,6 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelVisitor; import org.apache.calcite.rel.logical.LogicalCorrelate; -import org.apache.calcite.rel.logical.LogicalJoin; import org.apache.calcite.schema.SchemaPlus; import org.apache.calcite.sql.parser.SqlParser; import org.apache.calcite.test.CalciteAssert; @@ -64,36 +63,42 @@ public void testSearchWithFilter() { public void testSearchWithImplicitFormatSubsearchUsesCorrelate() { RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); LogicalCorrelate correlate = findCorrelate(root); - String logical = RelOptUtil.toString(root); - - Assert.assertTrue(logical, logical.contains("LogicalCorrelate")); - Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - Assert.assertTrue(logical, logical.contains("ARRAY_AGG")); - Assert.assertTrue(logical, logical.contains("query_string")); + String expectedLogical = + """ + LogicalProject(EMPNO=[$1], ENAME=[$2], JOB=[$3], MGR=[$4], HIREDATE=[$5], SAL=[$6], COMM=[$7], DEPTNO=[$8]) + LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0}]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('DEPTNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('DEPTNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalSort(fetch=[1]) + LogicalProject(DEPTNO=[$0]) + LogicalTableScan(table=[[scott, DEPT]]) + LogicalFilter(condition=[query_string(MAP('query', $cor0.search))], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[scott, EMP]]) + """; + verifyLogical(root, expectedLogical); Assert.assertSame(root.getCluster(), correlate.getLeft().getCluster()); Assert.assertSame(root.getCluster(), correlate.getRight().getCluster()); } - @Test - public void testImplicitFormatQueryStringReferencesCorrelatedValue() { - RelNode root = getRelNode("search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ]"); - String logical = RelOptUtil.toString(root); - - Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - Assert.assertTrue(logical, logical.contains("query_string")); - Assert.assertTrue(logical, logical.contains("$cor")); - } - @Test public void testImplicitFormatCombinesWithStaticSearchPredicate() { RelNode root = getRelNode("search source=EMP JOB=CLERK [ search source=DEPT | fields DEPTNO | head 1 ]"); - String logical = RelOptUtil.toString(root); - - Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - Assert.assertTrue(logical, logical.contains("JOB:CLERK")); - Assert.assertTrue(logical, logical.contains("$cor")); - Assert.assertTrue(logical, logical.contains("AND")); + String expectedLogical = + """ + LogicalProject(EMPNO=[$1], ENAME=[$2], JOB=[$3], MGR=[$4], HIREDATE=[$5], SAL=[$6], COMM=[$7], DEPTNO=[$8]) + LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0}]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('DEPTNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('DEPTNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalSort(fetch=[1]) + LogicalProject(DEPTNO=[$0]) + LogicalTableScan(table=[[scott, DEPT]]) + LogicalFilter(condition=[query_string(MAP('query', CONCAT(CONCAT(CONCAT(CONCAT('(', 'JOB:CLERK':VARCHAR), ')'), ' AND ':VARCHAR), CONCAT(CONCAT('(', $cor0.search), ')'))))], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[scott, EMP]]) + """; + verifyLogical(root, expectedLogical); } @Test @@ -102,13 +107,27 @@ public void testMultipleImplicitFormatSubsearchesShareOneCorrelate() { getRelNode( "search source=EMP [ search source=DEPT | fields DEPTNO | head 1 ] OR " + "[ search source=EMP | fields JOB | head 1 ]"); - LogicalCorrelate correlate = findCorrelate(root); - String logical = RelOptUtil.toString(root); - - Assert.assertTrue(correlate.getLeft() instanceof LogicalJoin); - Assert.assertEquals(1, countCorrelates(root)); - Assert.assertFalse(logical, logical.contains("SCALAR_QUERY")); - Assert.assertTrue(logical, logical.contains("OR")); + String expectedLogical = + """ + LogicalProject(EMPNO=[$2], ENAME=[$3], JOB=[$4], MGR=[$5], HIREDATE=[$6], SAL=[$7], COMM=[$8], DEPTNO=[$9]) + LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0, 1}]) + LogicalJoin(condition=[true], joinType=[inner]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('DEPTNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('DEPTNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalSort(fetch=[1]) + LogicalProject(DEPTNO=[$0]) + LogicalTableScan(table=[[scott, DEPT]]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('JOB="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('JOB="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalSort(fetch=[1]) + LogicalProject(JOB=[$2]) + LogicalTableScan(table=[[scott, EMP]]) + LogicalFilter(condition=[query_string(MAP('query', CONCAT(CONCAT('(', CONCAT(CONCAT($cor0.search, ' OR ':VARCHAR), $cor0.search0)), ')')))], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[scott, EMP]]) + """; + verifyLogical(root, expectedLogical); } @Test @@ -169,20 +188,6 @@ public void visit(RelNode node, int ordinal, RelNode parent) { return result.get(); } - private int countCorrelates(RelNode root) { - int[] count = {0}; - new RelVisitor() { - @Override - public void visit(RelNode node, int ordinal, RelNode parent) { - if (node instanceof LogicalCorrelate) { - count[0]++; - } - super.visit(node, ordinal, parent); - } - }.go(root); - return count[0]; - } - @Ignore("Fields used in search commands are not validated. Enable after fixing it.") @Test public void testSearchWithoutTimestampShouldThrow() { From 64e526d17a975a92ca12208ddaa7732a850c1d39 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Tue, 28 Jul 2026 14:52:12 +0800 Subject: [PATCH 08/10] Harden implicit format runtime search planning Signed-off-by: Songkan Tang --- .../sql/calcite/CalcitePlanContext.java | 16 +++ .../sql/calcite/CalciteRelNodeVisitor.java | 5 +- .../sql/calcite/CalciteRexNodeVisitor.java | 17 ++- .../sql/calcite/RuntimeSearchCorrelator.java | 20 ++- .../calcite/CalciteRexNodeVisitorTest.java | 45 ++++++ .../calcite/RuntimeSearchCorrelatorTest.java | 128 ++++++++++++++++++ .../sql/calcite/remote/CalciteExplainIT.java | 53 ++++++-- .../CalciteImplicitFormatCommandIT.java | 33 +++++ .../explain_format_custom_multivalue.yaml | 14 ++ .../calcite/explain_format_default.yaml | 14 ++ .../calcite/explain_format_empty_result.yaml | 14 ++ .../explain_format_implicit_subsearch.yaml | 22 +++ .../scan/CalciteEnumerableIndexScanTest.java | 91 +++++++++++++ .../sql/ppl/calcite/CalcitePPLSearchTest.java | 32 +++++ 14 files changed, 481 insertions(+), 23 deletions(-) create mode 100644 core/src/test/java/org/opensearch/sql/calcite/RuntimeSearchCorrelatorTest.java create mode 100644 integ-test/src/test/resources/expectedOutput/calcite/explain_format_custom_multivalue.yaml create mode 100644 integ-test/src/test/resources/expectedOutput/calcite/explain_format_default.yaml create mode 100644 integ-test/src/test/resources/expectedOutput/calcite/explain_format_empty_result.yaml create mode 100644 integ-test/src/test/resources/expectedOutput/calcite/explain_format_implicit_subsearch.yaml create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScanTest.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 c7f3bc373ac..010d4bafa40 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -9,10 +9,13 @@ import java.sql.Connection; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.Stack; import java.util.function.BiFunction; import lombok.Getter; @@ -21,6 +24,7 @@ import org.apache.calcite.rex.RexCorrelVariable; import org.apache.calcite.rex.RexLambdaRef; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexSubQuery; import org.apache.calcite.tools.FrameworkConfig; import org.checkerframework.checker.nullness.qual.Nullable; import org.opensearch.sql.ast.expression.AggregateFunction; @@ -83,6 +87,10 @@ public class CalcitePlanContext { private final Stack correlVar = new Stack<>(); private final Stack> windowPartitions = new Stack<>(); + /** Identity marker for scalar subqueries produced by an implicit format inside search. */ + private final Set implicitFormatSubqueries = + Collections.newSetFromMap(new IdentityHashMap<>()); + @Getter public Map rexLambdaRefMap; /** @@ -205,6 +213,14 @@ public Optional peekCorrelVar() { } } + void registerImplicitFormatSubquery(RexSubQuery subquery) { + implicitFormatSubqueries.add(subquery); + } + + boolean isImplicitFormatSubquery(RexSubQuery subquery) { + return implicitFormatSubqueries.contains(subquery); + } + /** * Creates a clone of this context that shares the relBuilder with the parent. This allows lambda * expressions to reference fields from the current row while having their own lambda variable 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 f64be0ed2a0..f45e57f368d 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -357,7 +357,10 @@ public RelNode visitSearch(Search node, CalcitePlanContext context) { context.relBuilder.filter(queryStringRex); if (node.hasImplicitSubquery()) { context.relBuilder.push( - RuntimeSearchCorrelator.correlate(context.relBuilder.build(), searchPredicateCompiler)); + RuntimeSearchCorrelator.correlate( + context.relBuilder.build(), + searchPredicateCompiler, + context::isImplicitFormatSubquery)); } return context.relBuilder.peek(); } 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 1bf1e217b51..735773001ab 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -34,6 +34,7 @@ import org.apache.calcite.rex.RexLambdaRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexSubQuery; import org.apache.calcite.sql.SqlIntervalQualifier; import org.apache.calcite.sql.SqlOperator; import org.apache.calcite.sql.fun.SqlStdOperatorTable; @@ -79,6 +80,7 @@ import org.opensearch.sql.ast.expression.subquery.InSubquery; import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; import org.opensearch.sql.ast.expression.subquery.SubqueryExpression; +import org.opensearch.sql.ast.tree.Format; import org.opensearch.sql.ast.tree.Sort.SortOption; import org.opensearch.sql.ast.tree.Sort.SortOrder; import org.opensearch.sql.ast.tree.UnresolvedPlan; @@ -865,11 +867,16 @@ public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) { @Override public RexNode visitScalarSubquery(ScalarSubquery node, CalcitePlanContext context) { - return context.relBuilder.scalarQuery( - b -> { - UnresolvedPlan subquery = node.getQuery(); - return resolveSubqueryPlan(subquery, node, context); - }); + RexSubQuery subquery = + context.relBuilder.scalarQuery( + b -> { + UnresolvedPlan subqueryPlan = node.getQuery(); + return resolveSubqueryPlan(subqueryPlan, node, context); + }); + if (node.getQuery() instanceof Format format && format.isImplicit()) { + context.registerImplicitFormatSubquery(subquery); + } + return subquery; } @Override diff --git a/core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java b/core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java index 8b911e04c16..25169468dcb 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java +++ b/core/src/main/java/org/opensearch/sql/calcite/RuntimeSearchCorrelator.java @@ -12,6 +12,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Predicate; import java.util.stream.IntStream; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.CorrelationId; @@ -37,18 +38,22 @@ public final class RuntimeSearchCorrelator { private RuntimeSearchCorrelator() {} /** - * Moves scalar subqueries out of a {@code query_string} filter and makes them the left input of a - * correlate. The right scan consumes their single-row output through a correlation variable, so - * its OpenSearch request is not built until every subquery result is available. + * Moves the implicit-format scalar subqueries out of a {@code query_string} filter and makes them + * the left input of a correlate. The right scan consumes their single-row output through a + * correlation variable, so its OpenSearch request is not built until every runtime search + * subquery result is available. Other Calcite subquery kinds remain in their original expression. */ public static RelNode correlate( - RelNode filterNode, SearchPredicateCompiler searchPredicateCompiler) { + RelNode filterNode, + SearchPredicateCompiler searchPredicateCompiler, + Predicate isImplicitFormatSubquery) { if (!(filterNode instanceof Filter filter)) { throw new IllegalStateException( "Runtime search query must produce a filter, but got " + filterNode.getClass().getSimpleName()); } - List subqueries = findSubqueries(filter.getCondition()); + List subqueries = + findImplicitFormatSubqueries(filter.getCondition(), isImplicitFormatSubquery); if (subqueries.isEmpty()) { return filter; } @@ -104,14 +109,15 @@ public RexNode visitSubQuery(RexSubQuery subquery) { correlate, List.of(), parentFields, filter.getRowType().getFieldNames()); } - private static List findSubqueries(RexNode condition) { + static List findImplicitFormatSubqueries( + RexNode condition, Predicate isImplicitFormatSubquery) { List subqueries = new ArrayList<>(); Set seen = Collections.newSetFromMap(new IdentityHashMap<>()); condition.accept( new RexVisitorImpl(true) { @Override public Void visitSubQuery(RexSubQuery subquery) { - if (seen.add(subquery)) { + if (isImplicitFormatSubquery.test(subquery) && seen.add(subquery)) { if (subquery.rel.getRowType().getFieldCount() != 1) { throw new SemanticCheckException( "Implicit format subsearch must return exactly one column"); diff --git a/core/src/test/java/org/opensearch/sql/calcite/CalciteRexNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/calcite/CalciteRexNodeVisitorTest.java index 9be542f2086..bd02742a3ff 100644 --- a/core/src/test/java/org/opensearch/sql/calcite/CalciteRexNodeVisitorTest.java +++ b/core/src/test/java/org/opensearch/sql/calcite/CalciteRexNodeVisitorTest.java @@ -14,6 +14,7 @@ import java.util.List; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexSubQuery; import org.apache.calcite.sql.type.ArraySqlType; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.tools.FrameworkConfig; @@ -27,6 +28,9 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.sql.ast.expression.LambdaFunction; import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; +import org.opensearch.sql.ast.tree.Format; +import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.calcite.utils.CalciteToolsHelper; import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelBuilder; import org.opensearch.sql.datasource.DataSourceService; @@ -210,4 +214,45 @@ public void testPrepareLambdaForReduceFinalizerFunction() { lambdaContext.getRexLambdaRefMap().get("acc").getType().getSqlTypeName(), SqlTypeName.FLOAT); } + + @Test + public void testOnlyImplicitFormatScalarSubqueryIsRegisteredForRuntimeSearch() { + RexSubQuery implicitFormatRex = Mockito.mock(RexSubQuery.class); + RexSubQuery explicitFormatRex = Mockito.mock(RexSubQuery.class); + RexSubQuery ordinaryScalarRex = Mockito.mock(RexSubQuery.class); + when(relBuilder.scalarQuery(any())) + .thenReturn(implicitFormatRex, explicitFormatRex, ordinaryScalarRex); + + Format implicitFormat = + new Format( + Format.DEFAULT_MV_SEPARATOR, + Format.DEFAULT_MAX_RESULTS, + Format.DEFAULT_ROW_PREFIX, + Format.DEFAULT_COLUMN_PREFIX, + Format.DEFAULT_COLUMN_SEPARATOR, + Format.DEFAULT_COLUMN_END, + Format.DEFAULT_ROW_SEPARATOR, + Format.DEFAULT_ROW_END, + Format.DEFAULT_EMPTY_STRING); + implicitFormat.setImplicit(true); + Format explicitFormat = + new Format( + Format.DEFAULT_MV_SEPARATOR, + Format.DEFAULT_MAX_RESULTS, + Format.DEFAULT_ROW_PREFIX, + Format.DEFAULT_COLUMN_PREFIX, + Format.DEFAULT_COLUMN_SEPARATOR, + Format.DEFAULT_COLUMN_END, + Format.DEFAULT_ROW_SEPARATOR, + Format.DEFAULT_ROW_END, + Format.DEFAULT_EMPTY_STRING); + + visitor.visitScalarSubquery(new ScalarSubquery(implicitFormat), context); + visitor.visitScalarSubquery(new ScalarSubquery(explicitFormat), context); + visitor.visitScalarSubquery(new ScalarSubquery(Mockito.mock(UnresolvedPlan.class)), context); + + assertTrue(context.isImplicitFormatSubquery(implicitFormatRex)); + assertFalse(context.isImplicitFormatSubquery(explicitFormatRex)); + assertFalse(context.isImplicitFormatSubquery(ordinaryScalarRex)); + } } diff --git a/core/src/test/java/org/opensearch/sql/calcite/RuntimeSearchCorrelatorTest.java b/core/src/test/java/org/opensearch/sql/calcite/RuntimeSearchCorrelatorTest.java new file mode 100644 index 00000000000..4caf82a149f --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/calcite/RuntimeSearchCorrelatorTest.java @@ -0,0 +1,128 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import com.google.common.collect.ImmutableList; +import java.math.BigDecimal; +import java.util.List; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.volcano.VolcanoPlanner; +import org.apache.calcite.rel.logical.LogicalValues; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexSubQuery; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.junit.jupiter.api.Test; + +class RuntimeSearchCorrelatorTest { + + @Test + void findImplicitFormatSubqueriesIgnoresUnmarkedScalarSubqueries() { + JavaTypeFactoryImpl typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + RelOptCluster cluster = RelOptCluster.create(new VolcanoPlanner(), rexBuilder); + RexSubQuery implicitFormat = RexSubQuery.scalar(LogicalValues.createOneRow(cluster)); + RexSubQuery ordinaryScalar = RexSubQuery.scalar(LogicalValues.createOneRow(cluster)); + RelDataType integerType = typeFactory.createSqlType(SqlTypeName.INTEGER); + RexNode condition = + rexBuilder.makeCall( + integerType, SqlStdOperatorTable.PLUS, List.of(implicitFormat, ordinaryScalar)); + + assertEquals( + List.of(implicitFormat), + RuntimeSearchCorrelator.findImplicitFormatSubqueries( + condition, subquery -> subquery == implicitFormat)); + } + + @Test + void findImplicitFormatSubqueriesIgnoresInAndExistsSubqueries() { + JavaTypeFactoryImpl typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + RelOptCluster cluster = RelOptCluster.create(new VolcanoPlanner(), rexBuilder); + LogicalValues values = LogicalValues.createOneRow(cluster); + RexSubQuery exists = RexSubQuery.exists(values); + RexSubQuery in = + RexSubQuery.in(values, ImmutableList.of(rexBuilder.makeExactLiteral(BigDecimal.ZERO))); + + assertEquals( + List.of(), RuntimeSearchCorrelator.findImplicitFormatSubqueries(exists, subquery -> false)); + assertEquals( + List.of(), RuntimeSearchCorrelator.findImplicitFormatSubqueries(in, subquery -> false)); + } + + @Test + void findImplicitFormatSubqueryAlongsideExistsInConditionalString() { + JavaTypeFactoryImpl typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + RelOptCluster cluster = RelOptCluster.create(new VolcanoPlanner(), rexBuilder); + RelDataType varcharType = typeFactory.createSqlType(SqlTypeName.VARCHAR); + RelDataType searchRowType = typeFactory.builder().add("search", varcharType).build(); + LogicalValues searchValue = + LogicalValues.create( + cluster, + searchRowType, + ImmutableList.of(ImmutableList.of(rexBuilder.makeLiteral("( host=\"web-1\" )")))); + RexSubQuery implicitFormat = RexSubQuery.scalar(searchValue); + RexSubQuery exists = RexSubQuery.exists(LogicalValues.createOneRow(cluster)); + RexNode conditionalSearch = + rexBuilder.makeCall( + varcharType, + SqlStdOperatorTable.CASE, + List.of( + exists, + rexBuilder.makeLiteral("status:200"), + rexBuilder.makeLiteral("status:500"))); + RexNode condition = + rexBuilder.makeCall( + varcharType, + SqlStdOperatorTable.CONCAT, + List.of( + rexBuilder.makeCall( + varcharType, + SqlStdOperatorTable.CONCAT, + List.of(conditionalSearch, rexBuilder.makeLiteral(" AND "))), + implicitFormat)); + + assertEquals( + List.of(implicitFormat), + RuntimeSearchCorrelator.findImplicitFormatSubqueries( + condition, subquery -> subquery == implicitFormat)); + } + + @Test + void findImplicitFormatSubqueryAlongsideInSubqueryCastToString() { + JavaTypeFactoryImpl typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + RelOptCluster cluster = RelOptCluster.create(new VolcanoPlanner(), rexBuilder); + RelDataType varcharType = typeFactory.createSqlType(SqlTypeName.VARCHAR); + RelDataType searchRowType = typeFactory.builder().add("search", varcharType).build(); + LogicalValues searchValue = + LogicalValues.create( + cluster, + searchRowType, + ImmutableList.of(ImmutableList.of(rexBuilder.makeLiteral("( host=\"web-1\" )")))); + RexSubQuery implicitFormat = RexSubQuery.scalar(searchValue); + RexSubQuery in = + RexSubQuery.in( + LogicalValues.createOneRow(cluster), + ImmutableList.of(rexBuilder.makeExactLiteral(BigDecimal.ZERO))); + RexNode condition = + rexBuilder.makeCall( + varcharType, + SqlStdOperatorTable.CONCAT, + List.of(rexBuilder.makeCast(varcharType, in), implicitFormat)); + + assertEquals( + List.of(implicitFormat), + RuntimeSearchCorrelator.findImplicitFormatSubqueries( + condition, subquery -> subquery == implicitFormat)); + } +} 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 b26059f5ac8..35b3d9e52e1 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 @@ -2938,16 +2938,49 @@ public void testNoMvWithEval() throws IOException { @Test public void testFormatExplain() throws IOException { - String query = - StringUtils.format( - "source=%s | where account_number < 3 | fields firstname, account_number | format", - TEST_INDEX_BANK); - String logical = logicalPlan(explainQueryYaml(query)).toLowerCase(Locale.ROOT); - Assert.assertTrue( - "Expected format to lower to a global ARRAY_AGG and a search projection", - logical.contains("array_agg") - && logical.contains("array_join") - && logical.contains("search=")); + enabledOnlyWhenPushdownIsEnabled(); + String actual = + explainQueryYaml( + StringUtils.format( + "source=%s | where account_number < 3 | fields firstname, account_number | format", + TEST_INDEX_BANK)); + assertYamlEqualsIgnoreId(loadExpectedPlan("explain_format_default.yaml"), actual); + } + + @Test + public void testFormatCustomDelimitersAndMultivalueExplain() throws IOException { + enabledOnlyWhenPushdownIsEnabled(); + String actual = + explainQueryYaml( + StringUtils.format( + "source=%s | where account_number=1 | eval names=array(firstname, lastname) |" + + " fields names | format mvsep='OR' '{' '[' 'AND' ']' 'OR' '}'", + TEST_INDEX_BANK)); + assertYamlEqualsIgnoreId(loadExpectedPlan("explain_format_custom_multivalue.yaml"), actual); + } + + @Test + public void testFormatEmptyResultExplain() throws IOException { + enabledOnlyWhenPushdownIsEnabled(); + String actual = + explainQueryYaml( + StringUtils.format( + "source=%s | where account_number < 0 | fields firstname | format" + + " emptystr='no matching data'", + TEST_INDEX_BANK)); + assertYamlEqualsIgnoreId(loadExpectedPlan("explain_format_empty_result.yaml"), actual); + } + + @Test + public void testImplicitFormatSubsearchExplain() throws IOException { + enabledOnlyWhenPushdownIsEnabled(); + String actual = + explainQueryYaml( + StringUtils.format( + "search source=%s age=36 [ search source=%s account_number=6 OR" + + " account_number=13 | fields account_number ] | fields account_number", + TEST_INDEX_BANK, TEST_INDEX_BANK)); + assertYamlEqualsIgnoreId(loadExpectedPlan("explain_format_implicit_subsearch.yaml"), actual); } @Test diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java index 79b3cae8a83..a85bfbea849 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteImplicitFormatCommandIT.java @@ -5,6 +5,7 @@ package org.opensearch.sql.calcite.remote; +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT; import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK; import static org.opensearch.sql.util.MatcherUtils.rows; import static org.opensearch.sql.util.MatcherUtils.verifyDataRows; @@ -22,6 +23,7 @@ public void init() throws Exception { super.init(); enableCalcite(); enabledOnlyWhenPushdownIsEnabled(); + loadIndex(Index.ACCOUNT); loadIndex(Index.BANK); } @@ -69,6 +71,37 @@ public void testMultipleImplicitFormatSubsearchesFeedOneParentSearch() throws IO verifyDataRows(result, rows(1), rows(6)); } + @Test + public void testNestedImplicitFormatSubsearchesFeedParentSearch() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_ACCOUNT + + " [ search source=" + + TEST_INDEX_BANK + + " [ search source=" + + TEST_INDEX_ACCOUNT + + " firstname=Hattie | fields city, state ]" + + " | fields account_number ]" + + " | fields account_number, firstname, lastname, city, state"); + + verifyDataRows(result, rows(6, "Hattie", "Bond", "Dante", "TN")); + } + + @Test + public void testStaticSearchPredicateCombinesWithImplicitFormatSubsearch() throws IOException { + JSONObject result = + executeQuery( + "search source=" + + TEST_INDEX_BANK + + " age=36 [ search source=" + + TEST_INDEX_BANK + + " account_number=6 OR account_number=13 | fields account_number ]" + + " | fields account_number, firstname, age | sort account_number"); + + verifyDataRows(result, rows(6, "Hattie", 36)); + } + @Test public void testParentPipelineAggregatesDynamicSearchResults() throws IOException { JSONObject result = diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_format_custom_multivalue.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_custom_multivalue.yaml new file mode 100644 index 00000000000..9836adb6e06 --- /dev/null +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_custom_multivalue.yaml @@ -0,0 +1,14 @@ +calcite: + logical: | + LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('{ ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' }':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(AND(IS NOT NULL(array($1, $6)), >(CHAR_LENGTH(ARRAY_JOIN(transform(ARRAY_COMPACT(CAST(array($1, $6)):VARCHAR NOT NULL ARRAY NOT NULL), (element) -> REPLACE(REPLACE(element, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '" OR names="':VARCHAR)), 0)), ||(||('( names="':VARCHAR, ARRAY_JOIN(transform(ARRAY_COMPACT(CAST(array($1, $6)):VARCHAR NOT NULL ARRAY NOT NULL), (element) -> REPLACE(REPLACE(element, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '" OR names="':VARCHAR)), '" )':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('[ ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(AND(IS NOT NULL(array($1, $6)), >(CHAR_LENGTH(ARRAY_JOIN(transform(ARRAY_COMPACT(CAST(array($1, $6)):VARCHAR NOT NULL ARRAY NOT NULL), (element) -> REPLACE(REPLACE(element, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '" OR names="':VARCHAR)), 0)), ||(||('( names="':VARCHAR, ARRAY_JOIN(transform(ARRAY_COMPACT(CAST(array($1, $6)):VARCHAR NOT NULL ARRAY NOT NULL), (element) -> REPLACE(REPLACE(element, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '" OR names="':VARCHAR)), '" )':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), ' ]':VARCHAR), null:VARCHAR)]) + LogicalFilter(condition=[=($0, 1)]) + CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]]) + physical: | + EnumerableLimit(fetch=[10000]) + EnumerableCalc(expr#0=[{inputs}], expr#1=[CAST($t0):VARCHAR ARRAY], expr#2=[ARRAY_COMPACT($t1)], expr#3=[' OR ':VARCHAR], expr#4=[ARRAY_JOIN($t2, $t3)], expr#5=[CHAR_LENGTH($t4)], expr#6=[0], expr#7=[>($t5, $t6)], expr#8=['{ ':VARCHAR], expr#9=[||($t8, $t4)], expr#10=[' }':VARCHAR], expr#11=[||($t9, $t10)], expr#12=['NOT ()':VARCHAR], expr#13=[CASE($t7, $t11, $t12)], search=[$t13]) + EnumerableAggregate(group=[{}], __format_rows=[ARRAY_AGG($0) FILTER $1]) + EnumerableCalc(expr#0..1=[{inputs}], expr#2=['[ ':VARCHAR], expr#3=[array($t0, $t1)], expr#4=[IS NOT NULL($t3)], expr#5=[CAST($t3):VARCHAR NOT NULL ARRAY NOT NULL], expr#6=[ARRAY_COMPACT($t5)], expr#7=['\':VARCHAR], expr#8=['\\':VARCHAR], expr#9=[REPLACE(element, $t7, $t8)], expr#10=['"':VARCHAR], expr#11=['\"':VARCHAR], expr#12=[REPLACE($t9, $t10, $t11)], expr#13=[(element) -> REPLACE(REPLACE(element, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)], expr#14=[transform($t6, $t13)], expr#15=['" OR names="':VARCHAR], expr#16=[ARRAY_JOIN($t14, $t15)], expr#17=[CHAR_LENGTH($t16)], expr#18=[0], expr#19=[>($t17, $t18)], expr#20=[AND($t4, $t19)], expr#21=['( names="':VARCHAR], expr#22=[||($t21, $t16)], expr#23=['" )':VARCHAR], expr#24=[||($t22, $t23)], expr#25=[null:VARCHAR], expr#26=[CASE($t20, $t24, $t25)], expr#27=[ARRAY($t26)], expr#28=[ARRAY_COMPACT($t27)], expr#29=[' AND ':VARCHAR], expr#30=[ARRAY_JOIN($t28, $t29)], expr#31=[||($t2, $t30)], expr#32=[' ]':VARCHAR], expr#33=[||($t31, $t32)], expr#34=[CHAR_LENGTH($t30)], expr#35=[>($t34, $t18)], $f1=[$t33], $f2=[$t35]) + CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, lastname], FILTER->=($0, 1), PROJECT->[firstname, lastname]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","query":{"term":{"account_number":{"value":1,"boost":1.0}}},"_source":{"includes":["firstname","lastname"]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)]) diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_format_default.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_default.yaml new file mode 100644 index 00000000000..bba2e09b325 --- /dev/null +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_default.yaml @@ -0,0 +1,14 @@ +calcite: + logical: | + LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('account_number="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR), CASE(IS NOT NULL($1), ||(||('firstname="':VARCHAR, REPLACE(REPLACE(CAST($1):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('account_number="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR), CASE(IS NOT NULL($1), ||(||('firstname="':VARCHAR, REPLACE(REPLACE(CAST($1):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalFilter(condition=[<($0, 3)]) + CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]]) + physical: | + EnumerableLimit(fetch=[10000]) + EnumerableCalc(expr#0=[{inputs}], expr#1=[CAST($t0):VARCHAR ARRAY], expr#2=[ARRAY_COMPACT($t1)], expr#3=[' OR ':VARCHAR], expr#4=[ARRAY_JOIN($t2, $t3)], expr#5=[CHAR_LENGTH($t4)], expr#6=[0], expr#7=[>($t5, $t6)], expr#8=['( ':VARCHAR], expr#9=[||($t8, $t4)], expr#10=[' )':VARCHAR], expr#11=[||($t9, $t10)], expr#12=['NOT ()':VARCHAR], expr#13=[CASE($t7, $t11, $t12)], search=[$t13]) + EnumerableAggregate(group=[{}], __format_rows=[ARRAY_AGG($0) FILTER $1]) + EnumerableCalc(expr#0..1=[{inputs}], expr#2=['( ':VARCHAR], expr#3=[IS NOT NULL($t0)], expr#4=['account_number="':VARCHAR], expr#5=[CAST($t0):VARCHAR NOT NULL], expr#6=['\':VARCHAR], expr#7=['\\':VARCHAR], expr#8=[REPLACE($t5, $t6, $t7)], expr#9=['"':VARCHAR], expr#10=['\"':VARCHAR], expr#11=[REPLACE($t8, $t9, $t10)], expr#12=[||($t4, $t11)], expr#13=[||($t12, $t9)], expr#14=[null:VARCHAR], expr#15=[CASE($t3, $t13, $t14)], expr#16=[IS NOT NULL($t1)], expr#17=['firstname="':VARCHAR], expr#18=[CAST($t1):VARCHAR NOT NULL], expr#19=[REPLACE($t18, $t6, $t7)], expr#20=[REPLACE($t19, $t9, $t10)], expr#21=[||($t17, $t20)], expr#22=[||($t21, $t9)], expr#23=[CASE($t16, $t22, $t14)], expr#24=[ARRAY($t15, $t23)], expr#25=[ARRAY_COMPACT($t24)], expr#26=[' AND ':VARCHAR], expr#27=[ARRAY_JOIN($t25, $t26)], expr#28=[||($t2, $t27)], expr#29=[' )':VARCHAR], expr#30=[||($t28, $t29)], expr#31=[CHAR_LENGTH($t27)], expr#32=[0], expr#33=[>($t31, $t32)], $f1=[$t30], $f2=[$t33]) + CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname], FILTER-><($0, 3)], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","query":{"range":{"account_number":{"from":null,"to":3,"include_lower":true,"include_upper":false,"boost":1.0}}},"_source":{"includes":["account_number","firstname"]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)]) diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_format_empty_result.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_empty_result.yaml new file mode 100644 index 00000000000..8fb9787be17 --- /dev/null +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_empty_result.yaml @@ -0,0 +1,14 @@ +calcite: + logical: | + LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'no matching data':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($1), ||(||('firstname="':VARCHAR, REPLACE(REPLACE(CAST($1):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($1), ||(||('firstname="':VARCHAR, REPLACE(REPLACE(CAST($1):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalFilter(condition=[<($0, 0)]) + CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]]) + physical: | + EnumerableLimit(fetch=[10000]) + EnumerableCalc(expr#0=[{inputs}], expr#1=[CAST($t0):VARCHAR ARRAY], expr#2=[ARRAY_COMPACT($t1)], expr#3=[' OR ':VARCHAR], expr#4=[ARRAY_JOIN($t2, $t3)], expr#5=[CHAR_LENGTH($t4)], expr#6=[0], expr#7=[>($t5, $t6)], expr#8=['( ':VARCHAR], expr#9=[||($t8, $t4)], expr#10=[' )':VARCHAR], expr#11=[||($t9, $t10)], expr#12=['no matching data':VARCHAR], expr#13=[CASE($t7, $t11, $t12)], search=[$t13]) + EnumerableAggregate(group=[{}], __format_rows=[ARRAY_AGG($0) FILTER $1]) + EnumerableCalc(expr#0=[{inputs}], expr#1=['( ':VARCHAR], expr#2=[IS NOT NULL($t0)], expr#3=['firstname="':VARCHAR], expr#4=[CAST($t0):VARCHAR NOT NULL], expr#5=['\':VARCHAR], expr#6=['\\':VARCHAR], expr#7=[REPLACE($t4, $t5, $t6)], expr#8=['"':VARCHAR], expr#9=['\"':VARCHAR], expr#10=[REPLACE($t7, $t8, $t9)], expr#11=[||($t3, $t10)], expr#12=[||($t11, $t8)], expr#13=[null:VARCHAR], expr#14=[CASE($t2, $t12, $t13)], expr#15=[ARRAY($t14)], expr#16=[ARRAY_COMPACT($t15)], expr#17=[' AND ':VARCHAR], expr#18=[ARRAY_JOIN($t16, $t17)], expr#19=[||($t1, $t18)], expr#20=[' )':VARCHAR], expr#21=[||($t19, $t20)], expr#22=[CHAR_LENGTH($t18)], expr#23=[0], expr#24=[>($t22, $t23)], $f1=[$t21], $f2=[$t24]) + CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname], FILTER-><($0, 0), PROJECT->[firstname]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","query":{"range":{"account_number":{"from":null,"to":0,"include_lower":true,"include_upper":false,"boost":1.0}}},"_source":{"includes":["firstname"]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)]) diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_format_implicit_subsearch.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_implicit_subsearch.yaml new file mode 100644 index 00000000000..cccb4f668f9 --- /dev/null +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_format_implicit_subsearch.yaml @@ -0,0 +1,22 @@ +calcite: + logical: | + LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT]) + LogicalProject(account_number=[$1]) + LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0}]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('account_number="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(CASE(IS NOT NULL($0), ||(||('account_number="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\':VARCHAR, '\\':VARCHAR), '"':VARCHAR, '\"':VARCHAR)), '"':VARCHAR), null:VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalSort(fetch=[10000]) + LogicalProject(account_number=[$0]) + LogicalFilter(condition=[query_string(MAP('query', '(account_number:6 OR account_number:13)':VARCHAR))]) + CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]]) + CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[FILTER->query_string(MAP('query', CONCAT(CONCAT(CONCAT(CONCAT('(', 'age:36':VARCHAR), ')'), ' AND ':VARCHAR), CONCAT(CONCAT('(', $cor0.search), ')'))))], dynamicQueryString=CONCAT(CONCAT(CONCAT(CONCAT('(', 'age:36':VARCHAR), ')'), ' AND ':VARCHAR), CONCAT(CONCAT('(', $cor0.search), ')')), OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m"}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)]) + physical: | + EnumerableLimit(fetch=[10000]) + EnumerableCalc(expr#0..1=[{inputs}], account_number=[$t1]) + EnumerableCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0}]) + EnumerableCalc(expr#0=[{inputs}], expr#1=[CAST($t0):VARCHAR ARRAY], expr#2=[ARRAY_COMPACT($t1)], expr#3=[' OR ':VARCHAR], expr#4=[ARRAY_JOIN($t2, $t3)], expr#5=[CHAR_LENGTH($t4)], expr#6=[0], expr#7=[>($t5, $t6)], expr#8=['( ':VARCHAR], expr#9=[||($t8, $t4)], expr#10=[' )':VARCHAR], expr#11=[||($t9, $t10)], expr#12=['NOT ()':VARCHAR], expr#13=[CASE($t7, $t11, $t12)], search=[$t13]) + EnumerableAggregate(group=[{}], __format_rows=[ARRAY_AGG($0) FILTER $1]) + EnumerableCalc(expr#0=[{inputs}], expr#1=['( ':VARCHAR], expr#2=[IS NOT NULL($t0)], expr#3=['account_number="':VARCHAR], expr#4=[CAST($t0):VARCHAR NOT NULL], expr#5=['\':VARCHAR], expr#6=['\\':VARCHAR], expr#7=[REPLACE($t4, $t5, $t6)], expr#8=['"':VARCHAR], expr#9=['\"':VARCHAR], expr#10=[REPLACE($t7, $t8, $t9)], expr#11=[||($t3, $t10)], expr#12=[||($t11, $t8)], expr#13=[null:VARCHAR], expr#14=[CASE($t2, $t12, $t13)], expr#15=[ARRAY($t14)], expr#16=[ARRAY_COMPACT($t15)], expr#17=[' AND ':VARCHAR], expr#18=[ARRAY_JOIN($t16, $t17)], expr#19=[||($t1, $t18)], expr#20=[' )':VARCHAR], expr#21=[||($t19, $t20)], expr#22=[CHAR_LENGTH($t18)], expr#23=[0], expr#24=[>($t22, $t23)], $f1=[$t21], $f2=[$t24]) + CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number], FILTER->query_string(MAP('query', '(account_number:6 OR account_number:13)':VARCHAR)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","query":{"query_string":{"query":"(account_number:6 OR account_number:13)","fields":[],"type":"best_fields","default_operator":"or","max_determinized_states":10000,"enable_position_increments":true,"fuzziness":"AUTO","fuzzy_prefix_length":0,"fuzzy_max_expansions":50,"phrase_slop":0,"escape":false,"auto_generate_synonyms_phrase_query":true,"fuzzy_transpositions":true,"boost":1.0}},"_source":{"includes":["account_number"]}}, requestedTotalSize=10000, pageSize=null, startFrom=0)]) + CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[FILTER->query_string(MAP('query', CONCAT(CONCAT(CONCAT(CONCAT('(', 'age:36':VARCHAR), ')'), ' AND ':VARCHAR), CONCAT(CONCAT('(', $cor0.search), ')')))), PROJECT->[account_number]], dynamicQueryString=CONCAT(CONCAT(CONCAT(CONCAT('(', 'age:36':VARCHAR), ')'), ' AND ':VARCHAR), CONCAT(CONCAT('(', $cor0.search), ')')), OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","_source":{"includes":["account_number"]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)]) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScanTest.java new file mode 100644 index 00000000000..b28d18ee55f --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScanTest.java @@ -0,0 +1,91 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage.scan; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.opensearch.index.query.QueryBuilders.boolQuery; +import static org.opensearch.index.query.QueryBuilders.queryStringQuery; +import static org.opensearch.index.query.QueryBuilders.termQuery; + +import java.util.List; +import java.util.Optional; +import java.util.Set; +import org.apache.calcite.linq4j.Enumerator; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.plan.RelTraitSet; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexNode; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.calcite.SearchPredicateCompiler; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.monitor.ResourceStatus; +import org.opensearch.sql.monitor.ResourceStatus.ResourceType; +import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.monitor.OpenSearchResourceMonitor; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.opensearch.storage.OpenSearchIndex; +import org.opensearch.sql.opensearch.storage.scan.context.DynamicQueryStringSpec; +import org.opensearch.sql.opensearch.storage.scan.context.FilterDigest; +import org.opensearch.sql.opensearch.storage.scan.context.OSRequestBuilderAction; +import org.opensearch.sql.opensearch.storage.scan.context.PushDownContext; +import org.opensearch.sql.opensearch.storage.scan.context.PushDownType; + +class CalciteEnumerableIndexScanTest { + + @Test + void runtimeQueryStringIsConjoinedWithExistingPushedFilter() { + RelOptCluster cluster = mock(RelOptCluster.class); + RelTraitSet traitSet = mock(RelTraitSet.class); + RelOptTable table = mock(RelOptTable.class); + RelDataType schema = mock(RelDataType.class); + OpenSearchIndex osIndex = mock(OpenSearchIndex.class); + OpenSearchClient client = mock(OpenSearchClient.class); + OpenSearchResourceMonitor monitor = mock(OpenSearchResourceMonitor.class); + OpenSearchRequest request = mock(OpenSearchRequest.class); + OpenSearchRequestBuilder requestBuilder = + new OpenSearchRequestBuilder( + mock(OpenSearchExprValueFactory.class), 10000, mock(Settings.class)); + + when(schema.getFieldNames()).thenReturn(List.of("account_number")); + when(osIndex.createRequestBuilder()).thenReturn(requestBuilder); + when(osIndex.getClient()).thenReturn(client); + when(client.getNodeClient()).thenReturn(Optional.empty()); + when(osIndex.getMaxResultWindow()).thenReturn(10000); + when(osIndex.getQueryBucketSize()).thenReturn(1000); + when(osIndex.buildRequest(same(requestBuilder))).thenReturn(request); + when(osIndex.createOpenSearchResourceMonitor()).thenReturn(monitor); + when(monitor.getStatus()).thenReturn(ResourceStatus.healthy(ResourceType.MEMORY)); + + PushDownContext pushDownContext = new PushDownContext(osIndex); + pushDownContext.add( + PushDownType.FILTER, + new FilterDigest(0, mock(RexNode.class)), + (OSRequestBuilderAction) builder -> builder.pushDownFilterForCalcite(termQuery("age", 36))); + + RexNode runtimePart = mock(RexNode.class); + SearchPredicateCompiler compiler = mock(SearchPredicateCompiler.class); + when(compiler.compile("( account_number=\"6\" )")).thenReturn("account_number:6"); + pushDownContext.setDynamicQueryString( + new DynamicQueryStringSpec(runtimePart, List.of(runtimePart), Set.of(0), compiler)); + + CalciteEnumerableIndexScan scan = + new CalciteEnumerableIndexScan( + cluster, traitSet, List.of(), table, osIndex, schema, pushDownContext); + Enumerator enumerator = + scan.scan(new String[] {"( account_number=\"6\" )"}).enumerator(); + enumerator.close(); + + assertEquals( + boolQuery().filter(termQuery("age", 36)).filter(queryStringQuery("account_number:6")), + requestBuilder.getSourceBuilder().query()); + } +} diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java index 0084c51f734..1dfb070595d 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLSearchTest.java @@ -79,6 +79,7 @@ public void testSearchWithImplicitFormatSubsearchUsesCorrelate() { verifyLogical(root, expectedLogical); Assert.assertSame(root.getCluster(), correlate.getLeft().getCluster()); Assert.assertSame(root.getCluster(), correlate.getRight().getCluster()); + verifyResult(correlate.getLeft(), "search=( ( DEPTNO=\"10\" ) )\n"); } @Test @@ -99,6 +100,7 @@ public void testImplicitFormatCombinesWithStaticSearchPredicate() { LogicalTableScan(table=[[scott, EMP]]) """; verifyLogical(root, expectedLogical); + verifyResult(findCorrelate(root).getLeft(), "search=( ( DEPTNO=\"10\" ) )\n"); } @Test @@ -128,6 +130,36 @@ public void testMultipleImplicitFormatSubsearchesShareOneCorrelate() { LogicalTableScan(table=[[scott, EMP]]) """; verifyLogical(root, expectedLogical); + verifyResult( + findCorrelate(root).getLeft(), + "search=( ( DEPTNO=\"10\" ) ); search0=( ( JOB=\"CLERK\" ) )\n"); + } + + @Test + public void testNestedImplicitFormatSubsearchesUseNestedCorrelates() { + RelNode root = + getRelNode( + "search source=EMP [ search source=EMP [ search source=EMP EMPNO=7369 | fields EMPNO ]" + + " | eval EMPNO=7499 | fields EMPNO ]"); + String expectedLogical = + """ + LogicalProject(EMPNO=[$1], ENAME=[$2], JOB=[$3], MGR=[$4], HIREDATE=[$5], SAL=[$6], COMM=[$7], DEPTNO=[$8]) + LogicalCorrelate(correlation=[$cor1], joinType=[inner], requiredColumns=[{0}]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0)]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('EMPNO="':VARCHAR, REPLACE(REPLACE('7499':VARCHAR, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('EMPNO="':VARCHAR, REPLACE(REPLACE('7499':VARCHAR, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)]) + LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0}]) + LogicalProject(search=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT($0), ' OR ':VARCHAR)), ' )':VARCHAR), 'NOT ()':VARCHAR)]) + LogicalAggregate(group=[{}], __format_rows=[ARRAY_AGG($0) WITHIN GROUP ([1])]) + LogicalProject(__format_row=[CASE(>(CHAR_LENGTH(ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('EMPNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), 0), ||(||('( ':VARCHAR, ARRAY_JOIN(ARRAY_COMPACT(ARRAY(||(||('EMPNO="':VARCHAR, REPLACE(REPLACE(CAST($0):VARCHAR NOT NULL, '\\':VARCHAR, '\\\\':VARCHAR), '"':VARCHAR, '\\"':VARCHAR)), '"':VARCHAR))), ' AND ':VARCHAR)), ' )':VARCHAR), null:VARCHAR)], __format_order_0=[$0]) + LogicalFilter(condition=[query_string(MAP('query', 'EMPNO:7369':VARCHAR))]) + LogicalTableScan(table=[[scott, EMP]]) + LogicalFilter(condition=[query_string(MAP('query', $cor0.search))], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[scott, EMP]]) + LogicalFilter(condition=[query_string(MAP('query', $cor1.search))], variablesSet=[[$cor1]]) + LogicalTableScan(table=[[scott, EMP]]) + """; + verifyLogical(root, expectedLogical); } @Test From d9b2ec49d7ff580776e167ddc646b9092bc93b9f Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Tue, 28 Jul 2026 15:03:40 +0800 Subject: [PATCH 09/10] Update format command version to 3.9 Signed-off-by: Songkan Tang --- docs/user/ppl/index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user/ppl/index.md b/docs/user/ppl/index.md index 4525c6653a9..64c8e4e1db9 100644 --- a/docs/user/ppl/index.md +++ b/docs/user/ppl/index.md @@ -41,7 +41,7 @@ source=accounts | [where command](cmd/where.md) | 1.0 | stable (since 1.0) | Filter the search result using boolean expressions. | | [subquery command](cmd/subquery.md) | 3.0 | experimental (since 3.0) | Embed one PPL query inside another for complex filtering and data retrieval operations. | | [fields command](cmd/fields.md) | 1.0 | stable (since 1.0) | Keep or remove fields from the search result. | -| [format command](cmd/format.md) | 3.8 | experimental (since 3.8) | Collapse rows and fields into a single search-expression string. | +| [format command](cmd/format.md) | 3.9 | experimental (since 3.9) | Collapse rows and fields into a single search-expression string. | | [rename command](cmd/rename.md) | 1.0 | stable (since 1.0) | Rename one or more fields in the search result. | | [eval command](cmd/eval.md) | 1.0 | stable (since 1.0) | Evaluate an expression and append the result to the search result. | | [foreach command](cmd/foreach.md) | 3.8 | experimental (since 3.8) | Run a templated evaluation for each selected field or collection element. | From 669bee13720fc4e53759fe3ba018089be273c5c9 Mon Sep 17 00:00:00 2001 From: Songkan Tang Date: Wed, 29 Jul 2026 16:01:20 +0800 Subject: [PATCH 10/10] Resolve format command rebase integration Signed-off-by: Songkan Tang --- .../org/opensearch/sql/executor/QueryService.java | 8 +------- .../sql/calcite/remote/CalciteExplainIT.java | 14 -------------- ppl/src/main/antlr/OpenSearchPPLLexer.g4 | 1 - 3 files changed, 1 insertion(+), 22 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index 1dedf1f227f..e05d2fe3ef1 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -113,13 +113,7 @@ public QueryService( Settings settings, ExecutionDispatcher executionDispatcher) { this( - analyzer, - executionEngine, - planner, - dataSourceService, - settings, - executionDispatcher, - null); + analyzer, executionEngine, planner, dataSourceService, settings, executionDispatcher, null); } public QueryService( 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 35b3d9e52e1..b153a2d82d8 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 @@ -2999,20 +2999,6 @@ public void testForeachExplain() throws IOException { && logical.contains(", 2)")); } - @Test - public void testFormatExplain() throws IOException { - String query = - StringUtils.format( - "source=%s | where account_number < 3 | fields firstname, account_number | format", - TEST_INDEX_BANK); - String logical = logicalPlan(explainQueryYaml(query)).toLowerCase(java.util.Locale.ROOT); - Assert.assertTrue( - "Expected format to lower to a global ARRAY_AGG and a search projection", - logical.contains("array_agg") - && logical.contains("array_join") - && logical.contains("search=")); - } - /** * 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/ppl/src/main/antlr/OpenSearchPPLLexer.g4 b/ppl/src/main/antlr/OpenSearchPPLLexer.g4 index 9fc72adfacb..b17490578c2 100644 --- a/ppl/src/main/antlr/OpenSearchPPLLexer.g4 +++ b/ppl/src/main/antlr/OpenSearchPPLLexer.g4 @@ -32,7 +32,6 @@ SORT: 'SORT'; EVAL: 'EVAL'; FOREACH: 'FOREACH'; FIELDFORMAT: 'FIELDFORMAT'; -FORMAT: 'FORMAT'; HEAD: 'HEAD'; BIN: 'BIN'; TOP: 'TOP';