lookup - initial commit

Signed-off-by: Kacper Trochimiak <[email protected]>
opensearch-project · Jul 2, 2024 · 49b964d · 49b964d
1 parent 9fad78e
commit 49b964d
Show file tree

Hide file tree

Showing 7 changed files with 145 additions and 0 deletions.
diff --git a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLLexer.g4 b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLLexer.g4
@@ -65,6 +65,7 @@ NUM:                                'NUM';
 // ARGUMENT KEYWORDS
 KEEPEMPTY:                          'KEEPEMPTY';
 CONSECUTIVE:                        'CONSECUTIVE';
+APPENDONLY:                         'APPENDONLY';
 DEDUP_SPLITVALUES:                  'DEDUP_SPLITVALUES';
 PARTITIONS:                         'PARTITIONS';
 ALLNUM:                             'ALLNUM';

diff --git a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4
@@ -38,6 +38,7 @@ commands
    | renameCommand
    | statsCommand
    | dedupCommand
+   | lookupCommand
    | sortCommand
    | evalCommand
    | headCommand
@@ -107,6 +108,18 @@ dedupCommand
    : DEDUP (number = integerLiteral)? fieldList (KEEPEMPTY EQUAL keepempty = booleanLiteral)? (CONSECUTIVE EQUAL consecutive = booleanLiteral)?
    ;
 
+matchFieldWithOptAs
+   : orignalMatchField = fieldExpression (AS asMatchField = fieldExpression)?
+   ;
+
+copyFieldWithOptAs
+   : orignalCopyField = fieldExpression (AS asCopyField = fieldExpression)?
+   ;
+
+lookupCommand
+   : LOOKUP tableSource matchFieldWithOptAs (COMMA matchFieldWithOptAs)* (APPENDONLY EQUAL appendonly = booleanLiteral)? (copyFieldWithOptAs (COMMA copyFieldWithOptAs)*)*
+   ;
+
 sortCommand
    : SORT sortbyClause
    ;
@@ -848,6 +861,7 @@ keywordsCanBeId
    | RENAME
    | STATS
    | DEDUP
+   | LOOKUP
    | SORT
    | EVAL
    | HEAD

diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java
@@ -43,6 +43,7 @@
 import org.opensearch.sql.ast.tree.Head;
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Limit;
+import org.opensearch.sql.ast.tree.Lookup;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
@@ -208,6 +209,10 @@ public T visitDedupe(Dedupe node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitLookup(Lookup node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitHead(Head node, C context) {
     return visitChildren(node, context);
   }

diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/tree/Lookup.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/tree/Lookup.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import com.google.common.collect.ImmutableList;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Argument;
+import org.opensearch.sql.ast.expression.Map;
+
+import java.util.List;
+
+/** AST node represent Lookup operation. */
+
+public class Lookup extends UnresolvedPlan {
+  private UnresolvedPlan child;
+  private final String indexName;
+  private final List<Map> matchFieldList;
+  private final List<Argument> options;
+  private final List<Map> copyFieldList;
+
+  public Lookup(UnresolvedPlan child, String indexName, List<Map> matchFieldList, List<Argument> options, List<Map> copyFieldList) {
+    this.child = child;
+    this.indexName = indexName;
+    this.matchFieldList = matchFieldList;
+    this.options = options;
+    this.copyFieldList = copyFieldList;
+  }
+
+  public Lookup(String indexName, List<Map> matchFieldList, List<Argument> options, List<Map> copyFieldList) {
+    this.indexName = indexName;
+    this.matchFieldList = matchFieldList;
+    this.options = options;
+    this.copyFieldList = copyFieldList;
+  }
+
+  @Override
+  public Lookup attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+
+  public String getIndexName() {
+    return indexName;
+  }
+
+  public List<Map> getMatchFieldList() {
+    return matchFieldList;
+  }
+
+  public List<Argument> getOptions() {
+    return options;
+  }
+
+  public List<Map> getCopyFieldList() {
+    return copyFieldList;
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return ImmutableList.of(this.child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitLookup(this, context);
+  }
+}
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java
@@ -49,6 +49,7 @@
 import org.opensearch.sql.ast.tree.Filter;
 import org.opensearch.sql.ast.tree.Head;
 import org.opensearch.sql.ast.tree.Kmeans;
+import org.opensearch.sql.ast.tree.Lookup;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
 import org.opensearch.sql.ast.tree.Relation;
@@ -256,6 +257,11 @@ public LogicalPlan visitDedupe(Dedupe node, CatalystPlanContext context) {
         throw new IllegalStateException("Not Supported operation : dedupe ");
     }
 
+    @Override
+    public LogicalPlan visitLookup(Lookup node, CatalystPlanContext context) {
+        throw new IllegalStateException("Not Supported operation : lookup ");
+    }
+
     /**
      * Expression Analyzer.
      */

diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java
@@ -32,6 +32,7 @@
 import org.opensearch.sql.ast.tree.Filter;
 import org.opensearch.sql.ast.tree.Head;
 import org.opensearch.sql.ast.tree.Kmeans;
+import org.opensearch.sql.ast.tree.Lookup;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
@@ -195,6 +196,47 @@ public UnresolvedPlan visitDedupCommand(OpenSearchPPLParser.DedupCommandContext
     return new Dedupe(ArgumentFactory.getArgumentList(ctx), getFieldList(ctx.fieldList()));
   }
 
+  /** Lookup command */
+  @Override
+  public UnresolvedPlan visitLookupCommand(OpenSearchPPLParser.LookupCommandContext ctx) {
+    ArgumentFactory.getArgumentList(ctx);
+    ctx.tableSource();
+    ctx.copyFieldWithOptAs();
+    ctx.matchFieldWithOptAs();
+    return new Lookup(
+            ctx.tableSource().tableQualifiedName().getText(),
+            ctx.matchFieldWithOptAs().stream()
+                    .map(
+                            ct ->
+                                    new Map(
+                                            evaluateFieldExpressionContext(ct.orignalMatchField),
+                                            evaluateFieldExpressionContext(ct.asMatchField, ct.orignalMatchField)))
+                    .collect(Collectors.toList()),
+            ArgumentFactory.getArgumentList(ctx),
+            ctx.copyFieldWithOptAs().stream()
+                    .map(
+                            ct ->
+                                    new Map(
+                                            evaluateFieldExpressionContext(ct.orignalCopyField),
+                                            evaluateFieldExpressionContext(ct.asCopyField, ct.orignalCopyField)))
+                    .collect(Collectors.toList()));
+  }
+
+  private UnresolvedExpression evaluateFieldExpressionContext(
+          OpenSearchPPLParser.FieldExpressionContext f) {
+    return internalVisitExpression(f);
+  }
+
+  private UnresolvedExpression evaluateFieldExpressionContext(
+          OpenSearchPPLParser.FieldExpressionContext f0,
+          OpenSearchPPLParser.FieldExpressionContext f1) {
+    if (f0 == null) {
+      return internalVisitExpression(f1);
+    } else {
+      return internalVisitExpression(f0);
+    }
+  }
+
   /** Head command visitor. */
   @Override
   public UnresolvedPlan visitHeadCommand(OpenSearchPPLParser.HeadCommandContext ctx) {

diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/utils/ArgumentFactory.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/utils/ArgumentFactory.java
@@ -72,6 +72,13 @@ public static List<Argument> getArgumentList(OpenSearchPPLParser.DedupCommandCon
             : new Argument("consecutive", new Literal(false, DataType.BOOLEAN)));
   }
 
+  public static List<Argument> getArgumentList(OpenSearchPPLParser.LookupCommandContext ctx) {
+    return Arrays.asList(
+            ctx.appendonly != null
+                    ? new Argument("appendonly", getArgumentValue(ctx.appendonly))
+                    : new Argument("appendonly", new Literal(false, DataType.BOOLEAN)));
+  }
+
   /**
    * Get list of {@link Argument}.
    *