Support filter params in graphlookup (opensearch-project#5134)

* Struct return array value instead of string

Signed-off-by: Heng Qian <qianheng@amazon.com>

* Support filter in GraphLookup

Signed-off-by: Heng Qian <qianheng@amazon.com>

* Fix IT

Signed-off-by: Heng Qian <qianheng@amazon.com>

* Add experimental tag in doc

Signed-off-by: Heng Qian <qianheng@amazon.com>

---------

Signed-off-by: Heng Qian <qianheng@amazon.com>
Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: qianheng-aws

core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java

@@ -18,6 +18,7 @@
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Literal;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
 
 /**
  * AST node for graphLookup command. Performs BFS graph traversal on a lookup table.
@@ -74,6 +75,11 @@ public enum Direction {
   /** Whether to use PIT (Point In Time) search for the lookup table to get complete results. */
   private final boolean usePIT;
 
+  /**
+   * Optional filter condition to restrict which lookup table documents participate in traversal.
+   */
+  private @Nullable final UnresolvedExpression filter;
+
   private UnresolvedPlan child;
 
   public String getDepthFieldName() {

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

@@ -2606,9 +2606,16 @@ public RelNode visitGraphLookup(GraphLookup node, CalcitePlanContext context) {
     // 3. Visit and materialize lookup table
     analyze(node.getFromTable(), context);
     tryToRemoveMetaFields(context, true);
+
+    // 4. Convert filter expression to RexNode against lookup table schema
+    RexNode filterRex = null;
+    if (node.getFilter() != null) {
+      filterRex = rexVisitor.analyze(node.getFilter(), context);
+    }
+
     RelNode lookupTable = builder.build();
 
-    // 4. Create LogicalGraphLookup RelNode
+    // 5. Create LogicalGraphLookup RelNode
     // The conversion rule will extract the OpenSearchIndex from the lookup table
     RelNode graphLookup =
         LogicalGraphLookup.create(
@@ -2623,7 +2630,8 @@ public RelNode visitGraphLookup(GraphLookup node, CalcitePlanContext context) {
             bidirectional,
             supportArray,
             batchMode,
-            usePIT);
+            usePIT,
+            filterRex);
 
     builder.push(graphLookup);
     return builder.peek();

core/src/main/java/org/opensearch/sql/calcite/plan/rel/GraphLookup.java

@@ -16,6 +16,7 @@
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.type.SqlTypeName;
 
 /**
@@ -51,6 +52,7 @@ public abstract class GraphLookup extends BiRel {
   protected final boolean supportArray;
   protected final boolean batchMode;
   protected final boolean usePIT;
+  @Nullable protected final RexNode filter;
 
   private RelDataType outputRowType;
 
@@ -72,6 +74,7 @@ public abstract class GraphLookup extends BiRel {
    *     pushdown)
    * @param batchMode Whether to batch all source start values into a single unified BFS
    * @param usePIT Whether to use PIT (Point In Time) search for complete results
+   * @param filter Optional filter condition for lookup table documents
    */
   protected GraphLookup(
       RelOptCluster cluster,
@@ -87,7 +90,8 @@ protected GraphLookup(
       boolean bidirectional,
       boolean supportArray,
       boolean batchMode,
-      boolean usePIT) {
+      boolean usePIT,
+      @Nullable RexNode filter) {
     super(cluster, traitSet, source, lookup);
     this.startField = startField;
     this.fromField = fromField;
@@ -99,6 +103,7 @@ protected GraphLookup(
     this.supportArray = supportArray;
     this.batchMode = batchMode;
     this.usePIT = usePIT;
+    this.filter = filter;
   }
 
   /** Returns the source table RelNode. */
@@ -181,6 +186,7 @@ public RelWriter explainTerms(RelWriter pw) {
         .item("bidirectional", bidirectional)
         .itemIf("supportArray", supportArray, supportArray)
         .itemIf("batchMode", batchMode, batchMode)
-        .itemIf("usePIT", usePIT, usePIT);
+        .itemIf("usePIT", usePIT, usePIT)
+        .itemIf("filter", filter, filter != null);
   }
 }

core/src/main/java/org/opensearch/sql/calcite/plan/rel/LogicalGraphLookup.java

@@ -12,6 +12,7 @@
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.plan.RelTraitSet;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rex.RexNode;
 
 /**
  * Logical RelNode for graphLookup command. TODO: need to support trim fields and several transpose
@@ -37,6 +38,7 @@ public class LogicalGraphLookup extends GraphLookup {
    * @param supportArray Whether to support array-typed fields
    * @param batchMode Whether to batch all source start values into a single unified BFS
    * @param usePIT Whether to use PIT (Point In Time) search for complete results
+   * @param filter Optional filter condition for lookup table documents
    */
   protected LogicalGraphLookup(
       RelOptCluster cluster,
@@ -52,7 +54,8 @@ protected LogicalGraphLookup(
       boolean bidirectional,
       boolean supportArray,
       boolean batchMode,
-      boolean usePIT) {
+      boolean usePIT,
+      @Nullable RexNode filter) {
     super(
         cluster,
         traitSet,
@@ -67,7 +70,8 @@ protected LogicalGraphLookup(
         bidirectional,
         supportArray,
         batchMode,
-        usePIT);
+        usePIT,
+        filter);
   }
 
   /**
@@ -85,6 +89,7 @@ protected LogicalGraphLookup(
    * @param supportArray Whether to support array-typed fields
    * @param batchMode Whether to batch all source start values into a single unified BFS
    * @param usePIT Whether to use PIT (Point In Time) search for complete results
+   * @param filter Optional filter condition for lookup table documents
    * @return A new LogicalGraphLookup instance
    */
   public static LogicalGraphLookup create(
@@ -99,7 +104,8 @@ public static LogicalGraphLookup create(
       boolean bidirectional,
       boolean supportArray,
       boolean batchMode,
-      boolean usePIT) {
+      boolean usePIT,
+      @Nullable RexNode filter) {
     RelOptCluster cluster = source.getCluster();
     RelTraitSet traitSet = cluster.traitSetOf(Convention.NONE);
     return new LogicalGraphLookup(
@@ -116,7 +122,8 @@ public static LogicalGraphLookup create(
         bidirectional,
         supportArray,
         batchMode,
-        usePIT);
+        usePIT,
+        filter);
   }
 
   @Override
@@ -135,6 +142,7 @@ public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
         bidirectional,
         supportArray,
         batchMode,
-        usePIT);
+        usePIT,
+        filter);
   }
 }

docs/user/ppl/cmd/graphlookup.md

@@ -1,5 +1,5 @@
 
-# graphLookup
+# graphLookup (Experimental)
 
 The `graphLookup` command performs recursive graph traversal on a collection using a breadth-first search (BFS) algorithm. It searches for documents matching a start value and recursively traverses connections between documents based on specified fields. This is useful for hierarchical data like organizational charts, social networks, or routing graphs.
 
@@ -8,7 +8,7 @@ The `graphLookup` command performs recursive graph traversal on a collection usi
 The `graphLookup` command has the following syntax:
 
 ```syntax
-graphLookup <lookupIndex> startField=<startField> fromField=<fromField> toField=<toField> [maxDepth=<maxDepth>] [depthField=<depthField>] [direction=(uni | bi)] [supportArray=(true | false)] [batchMode=(true | false)] [usePIT=(true | false)] as <outputField>
+graphLookup <lookupIndex> startField=<startField> fromField=<fromField> toField=<toField> [maxDepth=<maxDepth>] [depthField=<depthField>] [direction=(uni | bi)] [supportArray=(true | false)] [batchMode=(true | false)] [usePIT=(true | false)] [filter=(<condition>)] as <outputField>
 ```
 
 The following are examples of the `graphLookup` command syntax:
@@ -20,6 +20,7 @@ source = employees | graphLookup employees startField=reportsTo fromField=report
 source = employees | graphLookup employees startField=reportsTo fromField=reportsTo toField=name direction=bi as connections
 source = travelers | graphLookup airports startField=nearestAirport fromField=connects toField=airport supportArray=true as reachableAirports
 source = airports | graphLookup airports startField=airport fromField=connects toField=airport supportArray=true as reachableAirports
+source = employees | graphLookup employees startField=reportsTo fromField=reportsTo toField=name filter=(status = 'active' AND age > 18) as reportingHierarchy
 ```
 
 ## Parameters
@@ -38,6 +39,7 @@ The `graphLookup` command supports the following parameters.
 | `supportArray=(true \| false)` | Optional | When `true`, disables early visited-node filter pushdown to OpenSearch. Default is `false`. Set to `true` when `fromField` or `toField` contains array values to ensure correct traversal behavior. See [Array Field Handling](#array-field-handling) for details. |
 | `batchMode=(true \| false)` | Optional | When `true`, collects all start values from all source rows and performs a single unified BFS traversal. Default is `false`. The output changes to two arrays: `[Array<sourceRows>, Array<lookupResults>]`. See [Batch Mode](#batch-mode) for details. |
 | `usePIT=(true \| false)` | Optional | When `true`, enables PIT (Point In Time) search for the lookup table, allowing paginated retrieval of complete results without the `max_result_window` size limit. Default is `false`. See [PIT Search](#pit-search) for details. |
+| `filter=(<condition>)` | Optional | A filter condition to restrict which lookup table documents participate in the graph traversal. Only documents matching the condition are considered as candidates during BFS. Parentheses around the condition are required. Example: `filter=(status = 'active' AND age > 18)`. |
 | `as <outputField>` | Required | The name of the output array field that will contain all documents found during the graph traversal.                                                                                                                               |
 
 ## How It Works
@@ -329,6 +331,26 @@ source = employees
     as reportingHierarchy
 ```
 
+## Filtered Graph Traversal
+
+The `filter` parameter restricts which documents in the lookup table are considered during the BFS traversal. Only documents matching the filter condition participate as candidates at each traversal level.
+
+### Example
+
+The following query traverses only active employees in the reporting hierarchy:
+
+```ppl ignore
+source = employees
+  | graphLookup employees
+    startField=reportsTo
+    fromField=reportsTo
+    toField=name
+    filter=(status = 'active')
+    as reportingHierarchy
+```
+
+The filter is applied at the OpenSearch query level, so it combines efficiently with the BFS traversal queries. At each BFS level, the query sent to OpenSearch is effectively: `bool { filter: [user_filter, bfs_terms_query] }`.
+
 ## Limitations
 
 - The source input, which provides the starting point for the traversal, has a limitation of 100 documents to avoid performance issues.

docs/user/ppl/index.md

@@ -82,6 +82,8 @@ source=accounts
 | [addcoltotals command](cmd/addcoltotals.md) | 3.5 | stable (since 3.5) | Adds column values and appends a totals row. |
 | [transpose command](cmd/transpose.md) | 3.5 | stable (since 3.5) | Transpose rows to columns. |
 | [mvcombine command](cmd/mvcombine.md) | 3.5 | stable (since 3.4) | Combines values of a specified field across rows identical on all other fields. |
+| [graphlookup command](cmd/graphlookup.md) | 3.5 | experimental (since 3.5) | Performs recursive graph traversal on a collection using a BFS algorithm.|
+
 
   - [Syntax](cmd/syntax.md) - PPL query structure and command syntax formatting  
 * **Functions**  
@@ -101,4 +103,4 @@ source=accounts
 * **Optimization**  
   - [Optimization](../../user/optimization/optimization.rst)  
 * **Limitations**  
-  - [Limitations](limitations/limitations.md)  
+  - [Limitations](limitations/limitations.md)