Add include_metadata request parameter for PPL queries opensearch-project#5235 (opensearch-project#5412)

Co-authored-by: Isha Gupta <igupta24@apple.com>
Signed-off-by: Isha Gupta <igupta24@apple.com>

Author: ishag4

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -359,7 +359,10 @@ public T visitAllFields(AllFields node, C context) {
   }
 
   public T visitAllFieldsExcludeMeta(AllFieldsExcludeMeta node, C context) {
-    return visitChildren(node, context);
+    // AllFieldsExcludeMeta is an AllFields, so fall back to visitAllFields by default. Visitors
+    // that need to tell the two apart override this method; visitors that don't would otherwise
+    // silently return null here and NPE in their caller.
+    return visitAllFields(node, context);
   }
 
   public T visitNestedAllTupleFields(NestedAllTupleFields node, C context) {

core/src/main/java/org/opensearch/sql/ast/statement/Query.java

@@ -26,8 +26,19 @@ public class Query extends Statement {
   protected final UnresolvedPlan plan;
   protected final int fetchSize;
   private final QueryType queryType;
+
+  /**
+   * Whether the request asked for metadata fields such as {@code _id}, {@code _index} and {@code
+   * _score} to be kept in the result.
+   */
+  private final boolean includeMetadata;
+
   private HighlightConfig highlightConfig;
 
+  public Query(UnresolvedPlan plan, int fetchSize, QueryType queryType) {
+    this(plan, fetchSize, queryType, false);
+  }
+
   @Override
   public <R, C> R accept(AbstractNodeVisitor<R, C> visitor, C context) {
     return visitor.visitQuery(this, context);

core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java

@@ -80,6 +80,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<RexCorrelVariable> correlVar = new Stack<>();
   private final Stack<List<RexNode>> windowPartitions = new Stack<>();
 
@@ -165,6 +172,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
@@ -219,6 +227,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.

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

@@ -726,10 +726,18 @@ private static void forceProjectExcept(RelBuilder relBuilder, Iterable<RexNode>
    *
    * <p>2. There is no other project ever visited in the main query
    *
+   * <p>Unless the request asked for metadata via {@code include_metadata}, in which case only the
+   * forced exclusion of case 1 still applies.
+   *
    * @param context CalcitePlanContext
    * @param excludeByForce whether exclude metadata fields by force
    */
   private static void tryToRemoveMetaFields(CalcitePlanContext context, boolean excludeByForce) {
+    // Join and subquery still strip metadata by force to keep their output schemas unambiguous.
+    if (context.isIncludeMetadata() && !excludeByForce) {
+      return;
+    }
+
     if (excludeByForce || !context.isProjectVisited()) {
       List<String> originalFields = context.relBuilder.peek().getRowType().getFieldNames();
       List<RexNode> metaFieldsRef =

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -152,9 +152,20 @@ public void execute(
       QueryType queryType,
       HighlightConfig highlightConfig,
       ResponseListener<ExecutionEngine.QueryResponse> 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<ExecutionEngine.QueryResponse> listener) {
     if (shouldUseCalcite(queryType)) {
-      executeWithCalcite(plan, queryType, highlightConfig, listener);
+      executeWithCalcite(plan, queryType, highlightConfig, includeMetadata, listener);
     } else {
+      // The V2 engine has no notion of metadata fields, so includeMetadata is ignored there.
       executeWithLegacy(plan, queryType, listener, Optional.empty());
     }
   }
@@ -175,20 +186,22 @@ public void explain(
       HighlightConfig highlightConfig,
       ResponseListener<ExecutionEngine.ExplainResponse> listener,
       ExplainMode mode) {
-    explain(plan, queryType, highlightConfig, listener, mode, null);
+    explain(plan, queryType, highlightConfig, false, listener, mode, null);
   }
 
-  /** Explain with optional highlight config and format. */
+  /** Explain with optional highlight config, include metadata flag, and format. */
   public void explain(
       UnresolvedPlan plan,
       QueryType queryType,
       HighlightConfig highlightConfig,
+      boolean includeMetadata,
       ResponseListener<ExecutionEngine.ExplainResponse> listener,
       ExplainMode mode,
       Format format) {
     if (shouldUseCalcite(queryType)) {
-      explainWithCalcite(plan, queryType, highlightConfig, listener, mode, format);
+      explainWithCalcite(plan, queryType, highlightConfig, includeMetadata, listener, mode, format);
     } else {
+      // The V2 engine has no notion of metadata fields, so includeMetadata is ignored there.
       explainWithLegacy(plan, queryType, listener, mode, Optional.empty());
     }
   }
@@ -198,6 +211,15 @@ public void executeWithCalcite(
       QueryType queryType,
       HighlightConfig highlightConfig,
       ResponseListener<ExecutionEngine.QueryResponse> listener) {
+    executeWithCalcite(plan, queryType, highlightConfig, false, listener);
+  }
+
+  public void executeWithCalcite(
+      UnresolvedPlan plan,
+      QueryType queryType,
+      HighlightConfig highlightConfig,
+      boolean includeMetadata,
+      ResponseListener<ExecutionEngine.QueryResponse> listener) {
     CalcitePlanContext.run(
         () -> {
           try {
@@ -209,7 +231,10 @@ public void executeWithCalcite(
                 () -> {
                   CalcitePlanContext context =
                       CalcitePlanContext.create(
-                          buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
+                          buildFrameworkConfig(),
+                          SysLimit.fromSettings(settings),
+                          queryType,
+                          includeMetadata);
 
                   context.setHighlightConfig(highlightConfig);
 
@@ -278,13 +303,14 @@ public void explainWithCalcite(
       HighlightConfig highlightConfig,
       ResponseListener<ExecutionEngine.ExplainResponse> listener,
       ExplainMode mode) {
-    explainWithCalcite(plan, queryType, highlightConfig, listener, mode, null);
+    explainWithCalcite(plan, queryType, highlightConfig, false, listener, mode, null);
   }
 
   public void explainWithCalcite(
       UnresolvedPlan plan,
       QueryType queryType,
       HighlightConfig highlightConfig,
+      boolean includeMetadata,
       ResponseListener<ExecutionEngine.ExplainResponse> listener,
       ExplainMode mode,
       Format format) {
@@ -296,7 +322,10 @@ public void explainWithCalcite(
                 () -> {
                   CalcitePlanContext context =
                       CalcitePlanContext.create(
-                          buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
+                          buildFrameworkConfig(),
+                          SysLimit.fromSettings(settings),
+                          queryType,
+                          includeMetadata);
                   context.setHighlightConfig(highlightConfig);
                   context.run(
                       () -> {

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java

@@ -33,14 +33,16 @@ public class QueryPlan extends AbstractPlan {
 
   protected final HighlightConfig highlightConfig;
 
+  protected final boolean includeMetadata;
+
   /** Constructor. */
   public QueryPlan(
       QueryId queryId,
       QueryType queryType,
       UnresolvedPlan plan,
       QueryService queryService,
       ResponseListener<ExecutionEngine.QueryResponse> listener) {
-    this(queryId, queryType, plan, queryService, listener, null);
+    this(queryId, queryType, plan, queryService, listener, null, false);
   }
 
   /** Constructor with highlight config. */
@@ -51,12 +53,25 @@ public QueryPlan(
       QueryService queryService,
       ResponseListener<ExecutionEngine.QueryResponse> 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<ExecutionEngine.QueryResponse> 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. */
@@ -67,20 +82,38 @@ public QueryPlan(
       int pageSize,
       QueryService queryService,
       ResponseListener<ExecutionEngine.QueryResponse> 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<ExecutionEngine.QueryResponse> 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);
     }
   }
 
@@ -98,7 +131,8 @@ public void explain(
           new NotImplementedException(
               "`explain` feature for paginated requests is not implemented yet."));
     } else {
-      queryService.explain(plan, getQueryType(), highlightConfig, listener, mode, format);
+      queryService.explain(
+          plan, getQueryType(), highlightConfig, includeMetadata, listener, mode, format);
     }
   }
 }

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java

@@ -117,7 +117,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();
@@ -129,7 +130,8 @@ public AbstractPlan visitQuery(
           node.getPlan(),
           queryService,
           context.getLeft(),
-          node.getHighlightConfig());
+          node.getHighlightConfig(),
+          node.isIncludeMetadata());
     }
   }
 

core/src/test/java/org/opensearch/sql/analysis/SelectExpressionAnalyzerTest.java

@@ -6,6 +6,7 @@
 package org.opensearch.sql.analysis;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
@@ -21,6 +22,8 @@
 import org.opensearch.sql.analysis.symbol.Namespace;
 import org.opensearch.sql.analysis.symbol.Symbol;
 import org.opensearch.sql.ast.dsl.AstDSL;
+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.expression.DSL;
 import org.opensearch.sql.expression.NamedExpression;
@@ -86,6 +89,25 @@ protected void assertAnalyzeEqual(
     assertEquals(Arrays.asList(expected), analyze(unresolvedExpression));
   }
 
+  /**
+   * {@link AllFieldsExcludeMeta} is an {@link org.opensearch.sql.ast.expression.AllFields}, so the
+   * V2 select-list analyzer must expand it the same way. It used to fall through to the default
+   * visitor method, which returns null and made {@link SelectExpressionAnalyzer#analyze} throw
+   * NullPointerException — surfacing as an HTTP 500 on any query wrapped in an implicit select-all.
+   */
+  @Test
+  public void all_fields_exclude_meta_expands_like_all_fields() {
+    SelectExpressionAnalyzer analyzer = new SelectExpressionAnalyzer(expressionAnalyzer);
+
+    List<NamedExpression> allFields =
+        analyzer.analyze(List.of(AllFields.of()), analysisContext, optimizer);
+    List<NamedExpression> excludeMeta =
+        analyzer.analyze(List.of(AllFieldsExcludeMeta.of()), analysisContext, optimizer);
+
+    assertFalse(allFields.isEmpty());
+    assertEquals(allFields, excludeMeta);
+  }
+
   @Test
   public void testContextWrapperIsolation() {
     // Test that context wrapper properly isolates optimizer instances, each wrapper should have its

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));