Add columns field

Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Nov 26, 2022 · 0baec59 · 0baec59
1 parent 9a7ca08
commit 0baec59
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 6 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java
@@ -32,8 +32,8 @@
 @UtilityClass
 public class LogicalPlanDSL {
 
-  public static LogicalPlan write(LogicalPlan input, Table table) {
-    return new LogicalWrite(input, table);
+  public static LogicalPlan write(LogicalPlan input, Table table, List<String> columns) {
+    return new LogicalWrite(input, table, columns);
   }
 
   public static LogicalPlan aggregation(

diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalWrite.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalWrite.java
@@ -6,6 +6,7 @@
 package org.opensearch.sql.planner.logical;
 
 import java.util.Collections;
+import java.util.List;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.ToString;
@@ -15,16 +16,23 @@
  * Logical operator for insert statement.
  */
 @EqualsAndHashCode(callSuper = true)
+@Getter
 @ToString
 public class LogicalWrite extends LogicalPlan {
 
   /** Table that handles the write operation. */
-  @Getter
   private final Table table;
 
-  public LogicalWrite(LogicalPlan child, Table table) {
+  /** Optional column name list specified in insert statement. */
+  private final List<String> columns;
+
+  /**
+   * Construct a logical write with given child node, table and column name list.
+   */
+  public LogicalWrite(LogicalPlan child, Table table, List<String> columns) {
     super(Collections.singletonList(child));
     this.table = table;
+    this.columns = columns;
   }
 
   @Override

diff --git a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java
@@ -12,6 +12,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -86,7 +87,7 @@ public TableScanOperator build() {
     assertNull(tableScanBuilder.accept(new LogicalPlanNodeVisitor<Integer, Object>() {
     }, null));
 
-    LogicalPlan write = LogicalPlanDSL.write(null, table);
+    LogicalPlan write = LogicalPlanDSL.write(null, table, Collections.emptyList());
     assertNull(write.accept(new LogicalPlanNodeVisitor<Integer, Object>() {
     }, null));
 

diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java
@@ -282,7 +282,7 @@ void table_support_write_builder_should_be_replaced() {
 
     assertEquals(
         writeBuilder,
-        optimize(write(values(), table))
+        optimize(write(values(), table, Collections.emptyList()))
     );
   }