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 index 842f46fb0b..a192966287 100644 --- 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 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 index 5993b2b4ec..496e6009e3 100644 --- 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 columns; + + /** + * Construct a logical write with given child node, table and column name list. + */ + public LogicalWrite(LogicalPlan child, Table table, List 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 index ccb951037a..341bcbc29e 100644 --- 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() { }, null)); - LogicalPlan write = LogicalPlanDSL.write(null, table); + LogicalPlan write = LogicalPlanDSL.write(null, table, Collections.emptyList()); assertNull(write.accept(new LogicalPlanNodeVisitor() { }, 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 index 137a830afd..7516aa1809 100644 --- 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())) ); }