diff --git a/api/src/test/java/org/opensearch/sql/api/parser/UnifiedQueryParserTest.java b/api/src/test/java/org/opensearch/sql/api/parser/UnifiedQueryParserTest.java index 1b6b5181aef..40b3aa948fa 100644 --- a/api/src/test/java/org/opensearch/sql/api/parser/UnifiedQueryParserTest.java +++ b/api/src/test/java/org/opensearch/sql/api/parser/UnifiedQueryParserTest.java @@ -11,7 +11,7 @@ import static org.opensearch.sql.ast.dsl.AstDSL.agg; import static org.opensearch.sql.ast.dsl.AstDSL.aggregate; import static org.opensearch.sql.ast.dsl.AstDSL.alias; -import static org.opensearch.sql.ast.dsl.AstDSL.allFields; +import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta; import static org.opensearch.sql.ast.dsl.AstDSL.compare; import static org.opensearch.sql.ast.dsl.AstDSL.defaultStatsArgs; import static org.opensearch.sql.ast.dsl.AstDSL.eval; @@ -36,7 +36,7 @@ public class UnifiedQueryParserTest extends UnifiedQueryTestBase { public void testParseSource() { assertEqual( "source = catalog.employees", - project(relation(qualifiedName("catalog", "employees")), allFields())); + project(relation(qualifiedName("catalog", "employees")), AllFieldsExcludeMeta.of())); } @Test @@ -47,7 +47,7 @@ public void testParseFilter() { filter( relation(qualifiedName("catalog", "employees")), compare(">", field("age"), intLiteral(30))), - allFields())); + AllFieldsExcludeMeta.of())); } @Test @@ -58,7 +58,7 @@ public void testParseEval() { eval( relation(qualifiedName("catalog", "employees")), let(field("f"), function("abs", field("id")))), - allFields())); + AllFieldsExcludeMeta.of())); } @Test @@ -72,7 +72,7 @@ public void testParseStats() { emptyList(), exprList(alias("department", field("department"))), defaultStatsArgs()), - allFields())); + AllFieldsExcludeMeta.of())); } @Test diff --git a/core/src/main/java/org/opensearch/sql/ast/statement/Query.java b/core/src/main/java/org/opensearch/sql/ast/statement/Query.java index c6a78724b76..9f03e65d290 100644 --- a/core/src/main/java/org/opensearch/sql/ast/statement/Query.java +++ b/core/src/main/java/org/opensearch/sql/ast/statement/Query.java @@ -26,6 +26,7 @@ public class Query extends Statement { protected final UnresolvedPlan plan; protected final int fetchSize; private final QueryType queryType; + private final boolean includeMetadata; private HighlightConfig highlightConfig; @Override 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 5f58dea227e..f71532bc32d 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -59,6 +59,13 @@ public class CalcitePlanContext { */ @Getter @Setter private boolean isProjectVisited = false; + /** + * Whether to include metadata fields like _id, _index, _score in the result. When true, metadata + * fields are included in wildcard field selections. When false (default), metadata fields are + * excluded. + */ + @Getter @Setter private boolean includeMetadata = false; + private final Stack correlVar = new Stack<>(); private final Stack> windowPartitions = new Stack<>(); @@ -99,6 +106,7 @@ private CalcitePlanContext(CalcitePlanContext parent) { this.rexBuilder = parent.rexBuilder; // Share the same rexBuilder this.functionProperties = parent.functionProperties; this.highlightConfig = parent.highlightConfig; + this.includeMetadata = parent.includeMetadata; // Preserve parent's metadata setting 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 @@ -147,6 +155,13 @@ public static CalcitePlanContext create( return new CalcitePlanContext(config, sysLimit, queryType); } + public static CalcitePlanContext create( + FrameworkConfig config, SysLimit sysLimit, QueryType queryType, boolean includeMetadata) { + CalcitePlanContext context = new CalcitePlanContext(config, sysLimit, queryType); + context.setIncludeMetadata(includeMetadata); + return context; + } + /** * Executes {@code action} with the thread-local legacy flag set according to the supplied * settings. 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 15bfece5f46..5b33a05b57c 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -493,11 +493,18 @@ private RelNode handleAllFieldsProject(Project node, CalcitePlanContext context) "Invalid field exclusion: operation would exclude all fields from the result set"); } AllFields allFields = (AllFields) node.getProjectList().getFirst(); - if (!(allFields instanceof AllFieldsExcludeMeta)) { - // Should not remove nested fields for AllFieldsExcludeMeta. + + if (allFields instanceof AllFieldsExcludeMeta) { + // Do NOT remove nested fields for AllFieldsExcludeMeta + tryToRemoveMetaFields(context, true); // Force exclude metadata fields + } else { + // For AllFields (include_metadata=true), include metadata fields tryToRemoveNestedFields(context); + // Mark as project visited to prevent automatic metadata field removal + context.setProjectVisited(true); + // Don't force exclude metadata fields - let them remain + tryToRemoveMetaFields(context, false); } - tryToRemoveMetaFields(context, allFields instanceof AllFieldsExcludeMeta); return context.relBuilder.peek(); } @@ -645,6 +652,11 @@ private static void forceProjectExcept(RelBuilder relBuilder, Iterable * @param excludeByForce whether exclude metadata fields by force */ private static void tryToRemoveMetaFields(CalcitePlanContext context, boolean excludeByForce) { + // If include_metadata=true, never remove metadata fields + if (context.isIncludeMetadata() && !excludeByForce) { + return; + } + if (excludeByForce || !context.isProjectVisited()) { List originalFields = context.relBuilder.peek().getRowType().getFieldNames(); List metaFieldsRef = 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 b294a168cd8..9d4a75e8129 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -100,9 +100,22 @@ public void execute( QueryType queryType, HighlightConfig highlightConfig, ResponseListener listener) { + execute(plan, queryType, highlightConfig, false, listener); + } + + /** Execute with optional highlight config and include metadata flag. */ + public void execute( + UnresolvedPlan plan, + QueryType queryType, + HighlightConfig highlightConfig, + boolean includeMetadata, + ResponseListener listener) { if (shouldUseCalcite(queryType)) { - executeWithCalcite(plan, queryType, highlightConfig, listener); + executeWithCalcite(plan, queryType, highlightConfig, includeMetadata, listener); } else { + // Legacy engine always includes basic metadata (schema information) + // The includeMetadata flag doesn't affect legacy engine behavior since + // it already provides column names, types, and aliases in the schema executeWithLegacy(plan, queryType, listener, Optional.empty()); } } @@ -123,9 +136,22 @@ public void explain( HighlightConfig highlightConfig, ResponseListener listener, ExplainMode mode) { + explain(plan, queryType, highlightConfig, false, listener, mode); + } + + /** Explain with optional highlight config and include metadata flag. */ + public void explain( + UnresolvedPlan plan, + QueryType queryType, + HighlightConfig highlightConfig, + boolean includeMetadata, + ResponseListener listener, + ExplainMode mode) { if (shouldUseCalcite(queryType)) { - explainWithCalcite(plan, queryType, highlightConfig, listener, mode); + explainWithCalcite(plan, queryType, highlightConfig, includeMetadata, listener, mode); } else { + // Legacy engine provides basic explain information + // The includeMetadata flag doesn't significantly affect legacy explain behavior explainWithLegacy(plan, queryType, listener, mode, Optional.empty()); } } @@ -135,6 +161,15 @@ public void executeWithCalcite( QueryType queryType, HighlightConfig highlightConfig, ResponseListener listener) { + executeWithCalcite(plan, queryType, highlightConfig, false, listener); + } + + public void executeWithCalcite( + UnresolvedPlan plan, + QueryType queryType, + HighlightConfig highlightConfig, + boolean includeMetadata, + ResponseListener listener) { CalcitePlanContext.run( () -> { try { @@ -144,7 +179,10 @@ public void executeWithCalcite( long analyzeStart = System.nanoTime(); CalcitePlanContext context = CalcitePlanContext.create( - buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType); + buildFrameworkConfig(), + SysLimit.fromSettings(settings), + queryType, + includeMetadata); context.setHighlightConfig(highlightConfig); @@ -172,6 +210,7 @@ public void executeWithCalcite( } catch (Throwable t) { if (isCalciteFallbackAllowed(t) && !(t instanceof NonFallbackCalciteException)) { log.warn("Fallback to V2 query engine since got exception", t); + // Legacy engine provides basic metadata support, so fallback is acceptable executeWithLegacy(plan, queryType, listener, Optional.of(t)); } else { propagateCalciteError(t, listener); @@ -187,13 +226,26 @@ public void explainWithCalcite( HighlightConfig highlightConfig, ResponseListener listener, ExplainMode mode) { + explainWithCalcite(plan, queryType, highlightConfig, false, listener, mode); + } + + public void explainWithCalcite( + UnresolvedPlan plan, + QueryType queryType, + HighlightConfig highlightConfig, + boolean includeMetadata, + ResponseListener listener, + ExplainMode mode) { CalcitePlanContext.run( () -> { try { QueryProfiling.noop(); CalcitePlanContext context = CalcitePlanContext.create( - buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType); + buildFrameworkConfig(), + SysLimit.fromSettings(settings), + queryType, + includeMetadata); context.setHighlightConfig(highlightConfig); context.run( () -> { diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java index 3f6407e8873..f9be6635a7c 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java @@ -32,6 +32,8 @@ public class QueryPlan extends AbstractPlan { protected final HighlightConfig highlightConfig; + protected final boolean includeMetadata; + /** Constructor. */ public QueryPlan( QueryId queryId, @@ -39,7 +41,7 @@ public QueryPlan( UnresolvedPlan plan, QueryService queryService, ResponseListener listener) { - this(queryId, queryType, plan, queryService, listener, null); + this(queryId, queryType, plan, queryService, listener, null, false); } /** Constructor with highlight config. */ @@ -50,12 +52,25 @@ public QueryPlan( QueryService queryService, ResponseListener listener, HighlightConfig highlightConfig) { + this(queryId, queryType, plan, queryService, listener, highlightConfig, false); + } + + /** Constructor with highlight config and include metadata flag. */ + public QueryPlan( + QueryId queryId, + QueryType queryType, + UnresolvedPlan plan, + QueryService queryService, + ResponseListener listener, + HighlightConfig highlightConfig, + boolean includeMetadata) { super(queryId, queryType); this.plan = plan; this.queryService = queryService; this.listener = listener; this.pageSize = Optional.empty(); this.highlightConfig = highlightConfig; + this.includeMetadata = includeMetadata; } /** Constructor with page size. */ @@ -66,20 +81,33 @@ public QueryPlan( int pageSize, QueryService queryService, ResponseListener listener) { + this(queryId, queryType, plan, pageSize, queryService, listener, false); + } + + /** Constructor with page size and include metadata flag. */ + public QueryPlan( + QueryId queryId, + QueryType queryType, + UnresolvedPlan plan, + int pageSize, + QueryService queryService, + ResponseListener listener, + boolean includeMetadata) { super(queryId, queryType); this.plan = plan; this.queryService = queryService; this.listener = listener; this.pageSize = Optional.of(pageSize); this.highlightConfig = null; + this.includeMetadata = includeMetadata; } @Override public void execute() { if (pageSize.isPresent()) { - queryService.execute(new Paginate(pageSize.get(), plan), getQueryType(), listener); + queryService.execute(new Paginate(pageSize.get(), plan), getQueryType(), highlightConfig, includeMetadata, listener); } else { - queryService.execute(plan, getQueryType(), highlightConfig, listener); + queryService.execute(plan, getQueryType(), highlightConfig, includeMetadata, listener); } } @@ -91,7 +119,7 @@ public void explain( new NotImplementedException( "`explain` feature for paginated requests is not implemented yet.")); } else { - queryService.explain(plan, getQueryType(), highlightConfig, listener, mode); + queryService.explain(plan, getQueryType(), highlightConfig, includeMetadata, listener, mode); } } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 48e2b3ce5e0..c55f7dce39a 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -115,7 +115,8 @@ public AbstractPlan visitQuery( node.getPlan(), node.getFetchSize(), queryService, - context.getLeft()); + context.getLeft(), + node.isIncludeMetadata()); } else { // This should be picked up by the legacy engine. throw new UnsupportedCursorRequestException(); @@ -127,7 +128,8 @@ public AbstractPlan visitQuery( node.getPlan(), queryService, context.getLeft(), - node.getHighlightConfig()); + node.getHighlightConfig(), + node.isIncludeMetadata()); } } diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index dd73b26a8c3..ecbc03429d5 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -59,14 +59,14 @@ void init() { @Test public void create_from_query_should_success() { - Statement query = new Query(plan, 0, queryType); + Statement query = new Query(plan, 0, queryType, false); AbstractPlan queryExecution = factory.create(query, queryListener, explainListener); assertTrue(queryExecution instanceof QueryPlan); } @Test public void create_from_explain_should_success() { - Statement query = new Explain(new Query(plan, 0, queryType), queryType); + Statement query = new Explain(new Query(plan, 0, queryType, false), queryType); AbstractPlan queryExecution = factory.create(query, queryListener, explainListener); assertTrue(queryExecution instanceof ExplainPlan); } @@ -103,7 +103,7 @@ public void no_consumer_response_channel() { public void create_query_with_fetch_size_which_can_be_paged() { when(plan.accept(any(CanPaginateVisitor.class), any())).thenReturn(Boolean.TRUE); factory = new QueryPlanFactory(queryService); - Statement query = new Query(plan, 10, queryType); + Statement query = new Query(plan, 10, queryType, false); AbstractPlan queryExecution = factory.create(query, queryListener, explainListener); assertTrue(queryExecution instanceof QueryPlan); } @@ -112,7 +112,7 @@ public void create_query_with_fetch_size_which_can_be_paged() { public void create_query_with_fetch_size_which_cannot_be_paged() { when(plan.accept(any(CanPaginateVisitor.class), any())).thenReturn(Boolean.FALSE); factory = new QueryPlanFactory(queryService); - Statement query = new Query(plan, 10, queryType); + Statement query = new Query(plan, 10, queryType, false); assertThrows( UnsupportedCursorRequestException.class, () -> factory.create(query, queryListener, explainListener)); diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java index 128df14ff8e..b6e14a63fd4 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java @@ -9,6 +9,9 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -52,7 +55,7 @@ public void execute_no_page_size() { QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener); query.execute(); - verify(queryService, times(1)).execute(any(), any(), any(), any()); + verify(queryService, times(1)).execute(any(), any(), any(), anyBoolean(), any()); } @Test @@ -60,7 +63,8 @@ public void explain_no_page_size() { QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener); query.explain(explainListener, mode); - verify(queryService, times(1)).explain(plan, queryType, null, explainListener, mode); + verify(queryService, times(1)) + .explain(eq(plan), eq(queryType), isNull(), anyBoolean(), eq(explainListener), eq(mode)); } @Test diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/request/PPLQueryRequestFactory.java b/plugin/src/main/java/org/opensearch/sql/plugin/request/PPLQueryRequestFactory.java index bb87bf7fa91..e158ecb3937 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/request/PPLQueryRequestFactory.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/request/PPLQueryRequestFactory.java @@ -31,6 +31,7 @@ public class PPLQueryRequestFactory { private static final String QUERY_PARAMS_PRETTY = "pretty"; private static final String QUERY_PARAMS_PROFILE = "profile"; private static final String QUERY_PARAMS_FETCH_SIZE = "fetch_size"; + private static final String QUERY_PARAMS_INCLUDE_METADATA = "include_metadata"; /** * Build {@link PPLQueryRequest} from {@link RestRequest}. @@ -97,6 +98,13 @@ private static PPLQueryRequest parsePPLRequestFromPayload(RestRequest restReques "Invalid fetch_size parameter: must be a valid integer", e); } } + // Support include_metadata as a URL parameter if not already in the JSON body + if (!jsonContent.has(QUERY_PARAMS_INCLUDE_METADATA) + && restRequest.params().containsKey(QUERY_PARAMS_INCLUDE_METADATA)) { + jsonContent.put( + QUERY_PARAMS_INCLUDE_METADATA, + Boolean.parseBoolean(restRequest.params().get(QUERY_PARAMS_INCLUDE_METADATA))); + } PPLQueryRequest pplRequest = new PPLQueryRequest( jsonContent.getString(PPL_FIELD_NAME), diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java index 5c6266beee1..6c17b555f13 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java @@ -104,7 +104,8 @@ public String getName() { @Override protected Set responseParams() { Set responseParams = new HashSet<>(super.responseParams()); - responseParams.addAll(Arrays.asList("format", "mode", "sanitize", "fetch_size")); + responseParams.addAll( + Arrays.asList("format", "mode", "sanitize", "fetch_size", "include_metadata")); return responseParams; } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java index ea353066cb0..a7973085c3a 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -91,6 +91,9 @@ private AbstractPlan plan( ResponseListener explainListener) { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) ParseTree cst = parser.parse(request.getRequest()); + + boolean includeMetadata = request.getIncludeMetadata(); + Statement statement = cst.accept( new AstStatementBuilder( @@ -101,6 +104,7 @@ private AbstractPlan plan( .highlightConfig(request.getHighlightConfig()) .format(request.getFormat()) .explainMode(request.getExplainMode()) + .includeMetadata(includeMetadata) .build())); log.info( diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java index 06c7fe1c38e..11b31183088 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java @@ -26,6 +26,7 @@ public class PPLQueryRequest { private static final String DEFAULT_PPL_PATH = "/_plugins/_ppl"; private static final String FETCH_SIZE_FIELD = "fetch_size"; private static final String HIGHLIGHT_FIELD = "highlight"; + private static final String INCLUDE_METADATA_FIELD = "include_metadata"; private static final int MAX_HIGHLIGHT_FIELDS = 100; private static final int MAX_TAG_ENTRIES = 10; @@ -124,6 +125,20 @@ public int getFetchSize() { return jsonContent.optInt(FETCH_SIZE_FIELD, 0); } + /** + * Get whether to include metadata fields (_id, _index, _score, etc.) in the response. When + * enabled, metadata fields will be included alongside regular fields in wildcard field selections + * (e.g., fields *). + * + * @return true if metadata fields should be included, false otherwise (default: false) + */ + public boolean getIncludeMetadata() { + if (jsonContent == null) { + return false; + } + return jsonContent.optBoolean(INCLUDE_METADATA_FIELD, false); + } + /** * Get highlight config from the request. Supports both the simple array format ({@code ["*"]}) * and the rich OSD object format with {@code pre_tags}, {@code post_tags}, {@code fields}, and diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java index d2c1f610238..237e422b080 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java @@ -12,6 +12,7 @@ import lombok.Data; import lombok.RequiredArgsConstructor; import org.opensearch.sql.ast.expression.AllFields; +import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta; import org.opensearch.sql.ast.statement.Explain; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; @@ -37,7 +38,7 @@ public Statement visitPplStatement(OpenSearchPPLParser.PplStatementContext ctx) rawPlan = new Head(context.getFetchSize(), 0).attach(rawPlan); } UnresolvedPlan plan = addSelectAll(rawPlan); - Query query = new Query(plan, 0, PPL); + Query query = new Query(plan, 0, PPL, context.isIncludeMetadata()); if (context.getHighlightConfig() != null && context.getHighlightConfig().fields() != null && !context.getHighlightConfig().fields().isEmpty()) { @@ -76,13 +77,21 @@ public static class StatementBuilderContext { private final String format; private final String explainMode; + + /** Whether to include metadata fields like _id, _index, _score in the result. */ + private final boolean includeMetadata; } private UnresolvedPlan addSelectAll(UnresolvedPlan plan) { if ((plan instanceof Project) && !((Project) plan).isExcluded()) { return plan; } else { - return new Project(ImmutableList.of(AllFields.of())).attach(plan); + // Use AllFieldsExcludeMeta when include_metadata is false (default behavior) + // Use AllFields when include_metadata is true (include metadata fields) + boolean includeMetadata = context.isIncludeMetadata(); + AllFields allFields = includeMetadata ? AllFields.of() : AllFieldsExcludeMeta.of(); + + return new Project(ImmutableList.of(allFields)).attach(plan); } } } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelper.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelper.java index a67507be315..8072433b7c4 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelper.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelper.java @@ -7,7 +7,7 @@ import com.google.common.collect.ImmutableList; import lombok.experimental.UtilityClass; -import org.opensearch.sql.ast.expression.AllFields; +import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.setting.Settings; @@ -21,7 +21,7 @@ public UnresolvedPlan addSelectAll(UnresolvedPlan plan) { if ((plan instanceof Project) && !((Project) plan).isExcluded()) { return plan; } else { - return new Project(ImmutableList.of(AllFields.of())).attach(plan); + return new Project(ImmutableList.of(AllFieldsExcludeMeta.of())).attach(plan); } } diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index 0825f6d1def..2d9a099023a 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -100,15 +100,6 @@ public void onFailure(Exception e) { @Test public void testExecuteShouldPass() { - doAnswer( - invocation -> { - ResponseListener listener = invocation.getArgument(3); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); - return null; - }) - .when(queryService) - .execute(any(), any(), any(), any()); - pplService.execute( new PPLQueryRequest("search source=t a=1", null, QUERY), getQueryListener(false), @@ -117,15 +108,6 @@ public void testExecuteShouldPass() { @Test public void testExecuteCsvFormatShouldPass() { - doAnswer( - invocation -> { - ResponseListener listener = invocation.getArgument(3); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); - return null; - }) - .when(queryService) - .execute(any(), any(), any(), any()); - pplService.execute( new PPLQueryRequest("search source=t a=1", null, QUERY, "csv"), getQueryListener(false), @@ -134,15 +116,6 @@ public void testExecuteCsvFormatShouldPass() { @Test public void testExplainShouldPass() { - doAnswer( - invocation -> { - ResponseListener listener = invocation.getArgument(3); - listener.onResponse(new ExplainResponse(new ExplainResponseNode("test"))); - return null; - }) - .when(queryService) - .explain(any(), any(), any(), any(), any()); - pplService.explain( new PPLQueryRequest("search source=t a=1", null, EXPLAIN), new ResponseListener() { @@ -171,15 +144,6 @@ public void testExplainWithIllegalQueryShouldBeCaughtByHandler() { @Test public void testPrometheusQuery() { - doAnswer( - invocation -> { - ResponseListener listener = invocation.getArgument(3); - listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); - return null; - }) - .when(queryService) - .execute(any(), any(), any(), any()); - pplService.execute( new PPLQueryRequest("source = prometheus.http_requests_total", null, QUERY), getQueryListener(false), diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFlattenTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFlattenTest.java index 0cb3a90e557..7c8df8b7195 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFlattenTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFlattenTest.java @@ -64,16 +64,16 @@ public void testFlatten() { RelNode root = getRelNode(ppl); // Regarded as an identity scan. See RelBuilder#L2801 String expectedLogical = - "LogicalProject(DEPTNO=[$0], EMP=[$1], EMPNAME=[$2], EMPNO=[$3])\n" + "LogicalProject(DEPTNO=[$0], EMP=[$1], EMP.EMPNAME=[$2], EMP.EMPNO=[$3], EMPNAME=[$2], EMPNO=[$3])\n" + " LogicalTableScan(table=[[scott, DEPT]])\n"; verifyLogical(root, expectedLogical); String expectedSparkSql = - "SELECT `DEPTNO`, `EMP`, `EMP.EMPNAME` `EMPNAME`, `EMP.EMPNO` `EMPNO`\nFROM `scott`.`DEPT`"; + "SELECT `DEPTNO`, `EMP`, `EMP.EMPNAME`, `EMP.EMPNO`, `EMP.EMPNAME` `EMPNAME`, `EMP.EMPNO` `EMPNO`\nFROM `scott`.`DEPT`"; verifyPPLToSparkSQL(root, expectedSparkSql); String expectedResult = - "DEPTNO=10; EMP={7369, SMITH}; EMPNAME=SMITH; EMPNO=7369\n" - + "DEPTNO=20; EMP={7499, ALLEN}; EMPNAME=ALLEN; EMPNO=7499\n" - + "DEPTNO=30; EMP={7521, WARD}; EMPNAME=WARD; EMPNO=7521\n"; + "DEPTNO=10; EMP={7369, SMITH}; EMP.EMPNAME=SMITH; EMP.EMPNO=7369; EMPNAME=SMITH; EMPNO=7369\n" + + "DEPTNO=20; EMP={7499, ALLEN}; EMP.EMPNAME=ALLEN; EMP.EMPNO=7499; EMPNAME=ALLEN; EMPNO=7499\n" + + "DEPTNO=30; EMP={7521, WARD}; EMP.EMPNAME=WARD; EMP.EMPNO=7521; EMPNAME=WARD; EMPNO=7521\n"; verifyResult(root, expectedResult); } @@ -82,16 +82,16 @@ public void testFlattenWithAliases() { String ppl = "source=DEPT | flatten EMP as name, number"; RelNode root = getRelNode(ppl); String expectedLogical = - "LogicalProject(DEPTNO=[$0], EMP=[$1], name=[$2], number=[$3])\n" + "LogicalProject(DEPTNO=[$0], EMP=[$1], EMP.EMPNAME=[$2], EMP.EMPNO=[$3], name=[$2], number=[$3])\n" + " LogicalTableScan(table=[[scott, DEPT]])\n"; verifyLogical(root, expectedLogical); String expectedSparkSql = - "SELECT `DEPTNO`, `EMP`, `EMP.EMPNAME` `name`, `EMP.EMPNO` `number`\nFROM `scott`.`DEPT`"; + "SELECT `DEPTNO`, `EMP`, `EMP.EMPNAME`, `EMP.EMPNO`, `EMP.EMPNAME` `name`, `EMP.EMPNO` `number`\nFROM `scott`.`DEPT`"; verifyPPLToSparkSQL(root, expectedSparkSql); String expectedResult = - "DEPTNO=10; EMP={7369, SMITH}; name=SMITH; number=7369\n" - + "DEPTNO=20; EMP={7499, ALLEN}; name=ALLEN; number=7499\n" - + "DEPTNO=30; EMP={7521, WARD}; name=WARD; number=7521\n"; + "DEPTNO=10; EMP={7369, SMITH}; EMP.EMPNAME=SMITH; EMP.EMPNO=7369; name=SMITH; number=7369\n" + + "DEPTNO=20; EMP={7499, ALLEN}; EMP.EMPNAME=ALLEN; EMP.EMPNO=7499; name=ALLEN; number=7499\n" + + "DEPTNO=30; EMP={7521, WARD}; EMP.EMPNAME=WARD; EMP.EMPNO=7521; name=WARD; number=7521\n"; verifyResult(root, expectedResult); } @@ -104,12 +104,12 @@ public void testProject() { String ppl = "source=DEPT"; RelNode root = getRelNode(ppl); String expectedLogical = - "LogicalProject(DEPTNO=[$0], EMP=[$1])\n LogicalTableScan(table=[[scott, DEPT]])\n"; + "LogicalTableScan(table=[[scott, DEPT]])\n"; verifyLogical(root, expectedLogical); String expectedResult = - "DEPTNO=10; EMP={7369, SMITH}\n" - + "DEPTNO=20; EMP={7499, ALLEN}\n" - + "DEPTNO=30; EMP={7521, WARD}\n"; + "DEPTNO=10; EMP={7369, SMITH}; EMP.EMPNAME=SMITH; EMP.EMPNO=7369\n" + + "DEPTNO=20; EMP={7499, ALLEN}; EMP.EMPNAME=ALLEN; EMP.EMPNO=7499\n" + + "DEPTNO=30; EMP={7521, WARD}; EMP.EMPNAME=WARD; EMP.EMPNO=7521\n"; verifyResult(root, expectedResult); } diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java index ba31b75b38b..bcf39383b33 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstStatementBuilderTest.java @@ -16,6 +16,7 @@ import org.mockito.Mock; import org.opensearch.sql.ast.Node; import org.opensearch.sql.ast.expression.AllFields; +import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta; import org.opensearch.sql.ast.statement.Explain; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; @@ -36,12 +37,13 @@ public void buildQueryStatement() { assertEqual( "search source=t | where a=1", new Query( - project(filter(relation("t"), compare("=", field("a"), intLiteral(1))), AllFields.of()), + project(filter(relation("t"), compare("=", field("a"), intLiteral(1))), AllFieldsExcludeMeta.of()), 0, - PPL)); + PPL, + false)); assertEqual( "search source=t a=1", - new Query(project(search(relation("t"), "a:1"), AllFields.of()), 0, PPL)); + new Query(project(search(relation("t"), "a:1"), AllFieldsExcludeMeta.of()), 0, PPL, false)); } @Test @@ -51,13 +53,15 @@ public void buildExplainStatement() { new Explain( new Query( project( - filter(relation("t"), compare("=", field("a"), intLiteral(1))), AllFields.of()), + filter(relation("t"), compare("=", field("a"), intLiteral(1))), AllFieldsExcludeMeta.of()), 0, - PPL), + PPL, + false), PPL)); assertExplainEqual( "search source=t a=1", - new Explain(new Query(project(search(relation("t"), "a:1"), AllFields.of()), 0, PPL), PPL)); + new Explain( + new Query(project(search(relation("t"), "a:1"), AllFieldsExcludeMeta.of()), 0, PPL, false), PPL)); } @Test @@ -66,7 +70,8 @@ public void buildQueryStatementWithFetchSize() { assertEqualWithFetchSize( "search source=t a=1", 100, - new Query(project(head(search(relation("t"), "a:1"), 100, 0), AllFields.of()), 0, PPL)); + new Query( + project(head(search(relation("t"), "a:1"), 100, 0), AllFieldsExcludeMeta.of()), 0, PPL, false)); } @Test @@ -75,7 +80,7 @@ public void buildQueryStatementWithFetchSizeZero() { assertEqualWithFetchSize( "search source=t a=1", 0, - new Query(project(search(relation("t"), "a:1"), AllFields.of()), 0, PPL)); + new Query(project(search(relation("t"), "a:1"), AllFieldsExcludeMeta.of()), 0, PPL, false)); } @Test @@ -83,7 +88,8 @@ public void buildQueryStatementWithLargeFetchSize() { assertEqualWithFetchSize( "search source=t a=1", 10000, - new Query(project(head(search(relation("t"), "a:1"), 10000, 0), AllFields.of()), 0, PPL)); + new Query( + project(head(search(relation("t"), "a:1"), 10000, 0), AllFieldsExcludeMeta.of()), 0, PPL, false)); } @Test @@ -94,8 +100,7 @@ public void buildQueryStatementWithFetchSizeAndSmallerHead() { assertEqualWithFetchSize( "source=t | head 3", 10, - new Query(project(head(head(relation("t"), 3, 0), 10, 0), AllFields.of()), 0, PPL)); - } + new Query(project(head(head(relation("t"), 3, 0), 10, 0), AllFieldsExcludeMeta.of()), 0, PPL, false)); @Test public void buildQueryStatementWithFetchSizeSmallerThanHead() { @@ -105,7 +110,7 @@ public void buildQueryStatementWithFetchSizeSmallerThanHead() { assertEqualWithFetchSize( "source=t | head 100", 5, - new Query(project(head(head(relation("t"), 100, 0), 5, 0), AllFields.of()), 0, PPL)); + new Query(project(head(head(relation("t"), 100, 0), 5, 0), AllFields.of()), 0, PPL, false)); } @Test @@ -115,14 +120,15 @@ public void buildQueryStatementWithFetchSizeAndHeadWithOffset() { assertEqualWithFetchSize( "source=t | head 3 from 1", 10, - new Query(project(head(head(relation("t"), 3, 1), 10, 0), AllFields.of()), 0, PPL)); + new Query(project(head(head(relation("t"), 3, 1), 10, 0), AllFieldsExcludeMeta.of()), 0, PPL, false)); } @Test public void buildQueryStatementWithHighlight() { // Highlight config is set on the Query statement, not as an AST wrapper HighlightConfig config = new HighlightConfig(List.of("*")); - Query expected = new Query(project(search(relation("t"), "a:1"), AllFields.of()), 0, PPL); + Query expected = + new Query(project(search(relation("t"), "a:1"), AllFieldsExcludeMeta.of()), 0, PPL, false); expected.setHighlightConfig(config); assertEqualWithHighlight("search source=t a=1", config, expected); } @@ -130,7 +136,8 @@ public void buildQueryStatementWithHighlight() { @Test public void buildQueryStatementWithHighlightMultipleTerms() { HighlightConfig config = new HighlightConfig(List.of("error", "login")); - Query expected = new Query(project(search(relation("t"), "a:1"), AllFields.of()), 0, PPL); + Query expected = + new Query(project(search(relation("t"), "a:1"), AllFieldsExcludeMeta.of()), 0, PPL, false); expected.setHighlightConfig(config); assertEqualWithHighlight("search source=t a=1", config, expected); } @@ -141,7 +148,7 @@ public void buildQueryStatementWithHighlightNull() { assertEqualWithHighlight( "search source=t a=1", null, - new Query(project(search(relation("t"), "a:1"), AllFields.of()), 0, PPL)); + new Query(project(search(relation("t"), "a:1"), AllFieldsExcludeMeta.of()), 0, PPL, false)); } @Test @@ -149,7 +156,8 @@ public void buildQueryStatementWithHighlightAndFetchSize() { // Both fetch_size and highlight: Head wraps the plan, config is on the Query HighlightConfig config = new HighlightConfig(List.of("*")); Query expected = - new Query(project(head(search(relation("t"), "a:1"), 100, 0), AllFields.of()), 0, PPL); + new Query( + project(head(search(relation("t"), "a:1"), 100, 0), AllFieldsExcludeMeta.of()), 0, PPL, false); expected.setHighlightConfig(config); assertEqualWithHighlightAndFetchSize("search source=t a=1", config, 100, expected); } diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelperTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelperTest.java index 7c1264e0b63..c9aca678e67 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelperTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/UnresolvedPlanHelperTest.java @@ -15,7 +15,7 @@ import org.junit.runner.RunWith; import org.mockito.Mockito; import org.mockito.junit.MockitoJUnitRunner; -import org.opensearch.sql.ast.expression.AllFields; +import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta; import org.opensearch.sql.ast.expression.UnresolvedExpression; import org.opensearch.sql.ast.tree.Project; import org.opensearch.sql.ast.tree.Rename; @@ -39,7 +39,7 @@ public void addProjectForProjectExcludeOperator() { UnresolvedPlan plan = UnresolvedPlanHelper.addSelectAll(project); assertTrue(plan instanceof Project); - assertThat(((Project) plan).getProjectList(), Matchers.contains(AllFields.of())); + assertThat(((Project) plan).getProjectList(), Matchers.contains(AllFieldsExcludeMeta.of())); } @Test diff --git a/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java b/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java index dfb045f345a..a76c4368fd2 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java +++ b/sql/src/main/java/org/opensearch/sql/sql/parser/AstStatementBuilder.java @@ -25,7 +25,7 @@ public class AstStatementBuilder extends OpenSearchSQLParserBaseVisitor