Add UT for all new classes

Signed-off-by: Chen Dai <[email protected]>

core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java

@@ -35,6 +35,7 @@
 import org.opensearch.sql.planner.physical.ValuesOperator;
 import org.opensearch.sql.planner.physical.WindowOperator;
 import org.opensearch.sql.storage.read.TableScanBuilder;
+import org.opensearch.sql.storage.write.TableWriteBuilder;
 
 /**
  * Default implementor for implementing logical to physical translation. "Default" here means all
@@ -129,6 +130,11 @@ public PhysicalPlan visitTableScanBuilder(TableScanBuilder plan, C context) {
     return plan.build();
   }
 
+  @Override
+  public PhysicalPlan visitTableWriteBuilder(TableWriteBuilder plan, C context) {
+    return plan.build(visitChild(plan, context));
+  }
+
   @Override
   public PhysicalPlan visitRelation(LogicalRelation node, C context) {
     throw new UnsupportedOperationException("Storage engine is responsible for "

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java

@@ -32,6 +32,10 @@
 @UtilityClass
 public class LogicalPlanDSL {
 
+  public static LogicalPlan write(LogicalPlan input, Table table) {
+    return new LogicalWrite(input, table);
+  }
+
   public static LogicalPlan aggregation(
       LogicalPlan input, List<NamedAggregator> aggregatorList, List<NamedExpression> groupByList) {
     return new LogicalAggregation(input, aggregatorList, groupByList);

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java

@@ -7,7 +7,7 @@
 package org.opensearch.sql.planner.logical;
 
 import org.opensearch.sql.storage.read.TableScanBuilder;
-import org.opensearch.sql.storage.TableWriteBuilder;
+import org.opensearch.sql.storage.write.TableWriteBuilder;
 
 /**
  * The visitor of {@link LogicalPlan}.
@@ -29,6 +29,10 @@ public R visitTableScanBuilder(TableScanBuilder plan, C context) {
     return visitNode(plan, context);
   }
 
+  public R visitWrite(LogicalWrite plan, C context) {
+    return visitNode(plan, context);
+  }
+
   public R visitTableWriteBuilder(TableWriteBuilder plan, C context) {
     return visitNode(plan, context);
   }

core/src/main/java/org/opensearch/sql/planner/logical/LogicalWrite.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.planner.logical;
+
+import java.util.Collections;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.storage.Table;
+
+/**
+ * Logical operator for insert statement.
+ */
+@EqualsAndHashCode(callSuper = true)
+@ToString
+public class LogicalWrite extends LogicalPlan {
+
+  /** Table that handles the write operation. */
+  @Getter
+  private final Table table;
+
+  public LogicalWrite(LogicalPlan child, Table table) {
+    super(Collections.singletonList(child));
+    this.table = table;
+  }
+
+  @Override
+  public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitWrite(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java

@@ -17,6 +17,7 @@
 import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort;
 import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder;
 import org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown;
+import org.opensearch.sql.planner.optimizer.rule.write.CreateTableWriteBuilder;
 
 /**
  * {@link LogicalPlan} Optimizer.
@@ -55,7 +56,8 @@ public static LogicalPlanOptimizer create() {
         TableScanPushDown.PUSH_DOWN_SORT,
         TableScanPushDown.PUSH_DOWN_LIMIT,
         TableScanPushDown.PUSH_DOWN_HIGHLIGHT,
-        TableScanPushDown.PUSH_DOWN_PROJECT));
+        TableScanPushDown.PUSH_DOWN_PROJECT,
+        new CreateTableWriteBuilder()));
   }
 
   /**

core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java

@@ -20,6 +20,7 @@
 import org.opensearch.sql.planner.logical.LogicalProject;
 import org.opensearch.sql.planner.logical.LogicalRelation;
 import org.opensearch.sql.planner.logical.LogicalSort;
+import org.opensearch.sql.planner.logical.LogicalWrite;
 import org.opensearch.sql.storage.Table;
 import org.opensearch.sql.storage.read.TableScanBuilder;
 
@@ -110,4 +111,14 @@ public static Property<LogicalPlan, Table> table() {
             ? Optional.of(((LogicalRelation) plan).getTable())
             : Optional.empty());
   }
+
+  /**
+   * Logical relation with table field.
+   */
+  public static Property<LogicalPlan, Table> writeTable() {
+    return Property.optionalProperty("table",
+        plan -> plan instanceof LogicalWrite
+            ? Optional.of(((LogicalWrite) plan).getTable())
+            : Optional.empty());
+  }
 }

core/src/main/java/org/opensearch/sql/planner/optimizer/rule/write/CreateTableWriteBuilder.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.planner.optimizer.rule.write;
+
+import static org.opensearch.sql.planner.optimizer.pattern.Patterns.writeTable;
+
+import com.facebook.presto.matching.Capture;
+import com.facebook.presto.matching.Captures;
+import com.facebook.presto.matching.Pattern;
+import lombok.Getter;
+import lombok.experimental.Accessors;
+import org.opensearch.sql.planner.logical.LogicalPlan;
+import org.opensearch.sql.planner.logical.LogicalWrite;
+import org.opensearch.sql.planner.optimizer.Rule;
+import org.opensearch.sql.storage.Table;
+import org.opensearch.sql.storage.write.TableWriteBuilder;
+
+/**
+ * Rule that replaces logical write operator with {@link TableWriteBuilder} for later optimization
+ * and transforming to physical operator.
+ */
+public class CreateTableWriteBuilder implements Rule<LogicalWrite> {
+
+  /** Capture the table inside matched logical relation operator. */
+  private final Capture<Table> capture;
+
+  /** Pattern that matches logical relation operator. */
+  @Accessors(fluent = true)
+  @Getter
+  private final Pattern<LogicalWrite> pattern;
+
+  /**
+   * Construct create table write builder rule.
+   */
+  public CreateTableWriteBuilder() {
+    this.capture = Capture.newCapture();
+    this.pattern = Pattern.typeOf(LogicalWrite.class)
+        .with(writeTable().capturedAs(capture));
+  }
+
+  @Override
+  public LogicalPlan apply(LogicalWrite plan, Captures captures) {
+    return captures.get(capture).createWriteBuilder(plan);
+  }
+}

core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java

@@ -7,7 +7,7 @@
 package org.opensearch.sql.planner.physical;
 
 import org.opensearch.sql.storage.TableScanOperator;
-import org.opensearch.sql.storage.TableWriteOperator;
+import org.opensearch.sql.storage.write.TableWriteOperator;
 
 /**
  * The visitor of {@link PhysicalPlan}.

core/src/main/java/org/opensearch/sql/storage/Table.java

@@ -10,8 +10,10 @@
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.executor.streaming.StreamingSource;
 import org.opensearch.sql.planner.logical.LogicalPlan;
+import org.opensearch.sql.planner.logical.LogicalWrite;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 import org.opensearch.sql.storage.read.TableScanBuilder;
+import org.opensearch.sql.storage.write.TableWriteBuilder;
 
 /**
  * Table.
@@ -76,9 +78,10 @@ default TableScanBuilder createScanBuilder() {
   /*
    * Create table write builder for logical to physical transformation.
    *
+   * @param plan logical write plan
    * @return table write builder
    */
-  default TableWriteBuilder createWriteBuilder() {
+  default TableWriteBuilder createWriteBuilder(LogicalWrite plan) {
     throw new UnsupportedOperationException(
         "Write operation is not supported on current table");
   }

core/src/main/java/org/opensearch/sql/storage/TableWriteBuilder.java → core/src/main/java/org/opensearch/sql/storage/write/TableWriteBuilder.java

@@ -3,17 +3,18 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.storage;
+package org.opensearch.sql.storage.write;
 
 import java.util.Collections;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor;
+import org.opensearch.sql.planner.physical.PhysicalPlan;
 
 /**
  * A {@link TableWriteBuilder} represents transition state between logical planning and physical
  * planning for table write operator. The concrete implementation class gets involved in the logical
- * optimization through this abstraction and thus transform to specific {@link TableWriteOperator} in
- * a certain data source.
+ * optimization through this abstraction and thus transform to specific {@link TableWriteOperator}
+ * in a certain storage engine.
  */
 public abstract class TableWriteBuilder extends LogicalPlan {
 
@@ -25,11 +26,12 @@ public TableWriteBuilder(LogicalPlan child) {
   }
 
   /**
-   * Build table write operator.
+   * Build table write operator with given child node.
    *
+   * @param child  child operator node
    * @return table write operator
    */
-  public abstract TableWriteOperator build();
+  public abstract TableWriteOperator build(PhysicalPlan child);
 
   @Override
   public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {

core/src/main/java/org/opensearch/sql/storage/TableWriteOperator.java → core/src/main/java/org/opensearch/sql/storage/write/TableWriteOperator.java

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.storage;
+package org.opensearch.sql.storage.write;
 
 import java.util.Collections;
 import java.util.List;

core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java

@@ -58,6 +58,8 @@
 import org.opensearch.sql.storage.Table;
 import org.opensearch.sql.storage.TableScanOperator;
 import org.opensearch.sql.storage.read.TableScanBuilder;
+import org.opensearch.sql.storage.write.TableWriteBuilder;
+import org.opensearch.sql.storage.write.TableWriteOperator;
 
 @ExtendWith(MockitoExtension.class)
 class DefaultImplementorTest {
@@ -212,4 +214,17 @@ public TableScanOperator build() {
     };
     assertEquals(tableScanOperator, tableScanBuilder.accept(implementor, null));
   }
+
+  @Test
+  public void visitTableWriteBuilderShouldBuildTableWriteOperator() {
+    LogicalPlan child = values();
+    TableWriteOperator tableWriteOperator = Mockito.mock(TableWriteOperator.class);
+    TableWriteBuilder logicalPlan = new TableWriteBuilder(child) {
+      @Override
+      public TableWriteOperator build(PhysicalPlan child) {
+        return tableWriteOperator;
+      }
+    };
+    assertEquals(tableWriteOperator, logicalPlan.accept(implementor, null));
+  }
 }

core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java

@@ -31,9 +31,12 @@
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.aggregation.Aggregator;
 import org.opensearch.sql.expression.window.WindowDefinition;
+import org.opensearch.sql.planner.physical.PhysicalPlan;
 import org.opensearch.sql.storage.Table;
 import org.opensearch.sql.storage.TableScanOperator;
 import org.opensearch.sql.storage.read.TableScanBuilder;
+import org.opensearch.sql.storage.write.TableWriteBuilder;
+import org.opensearch.sql.storage.write.TableWriteOperator;
 
 /**
  * Todo. Temporary added for UT coverage, Will be removed.
@@ -83,6 +86,19 @@ public TableScanOperator build() {
     assertNull(tableScanBuilder.accept(new LogicalPlanNodeVisitor<Integer, Object>() {
     }, null));
 
+    LogicalPlan write = LogicalPlanDSL.write(null, table);
+    assertNull(write.accept(new LogicalPlanNodeVisitor<Integer, Object>() {
+    }, null));
+
+    TableWriteBuilder tableWriteBuilder = new TableWriteBuilder(null) {
+      @Override
+      public TableWriteOperator build(PhysicalPlan child) {
+        return null;
+      }
+    };
+    assertNull(tableWriteBuilder.accept(new LogicalPlanNodeVisitor<Integer, Object>() {
+    }, null));
+
     LogicalPlan filter = LogicalPlanDSL.filter(relation, expression);
     assertNull(filter.accept(new LogicalPlanNodeVisitor<Integer, Object>() {
     }, null));

core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java

@@ -7,6 +7,7 @@
 package org.opensearch.sql.planner.optimizer;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
 import static org.opensearch.sql.data.model.ExprValueUtils.integerValue;
@@ -20,6 +21,8 @@
 import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project;
 import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation;
 import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort;
+import static org.opensearch.sql.planner.logical.LogicalPlanDSL.values;
+import static org.opensearch.sql.planner.logical.LogicalPlanDSL.write;
 
 import com.google.common.collect.ImmutableList;
 import java.util.Collections;
@@ -39,6 +42,7 @@
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 import org.opensearch.sql.storage.Table;
 import org.opensearch.sql.storage.read.TableScanBuilder;
+import org.opensearch.sql.storage.write.TableWriteBuilder;
 
 @ExtendWith(MockitoExtension.class)
 class LogicalPlanOptimizerTest {
@@ -270,6 +274,37 @@ public PhysicalPlan implement(LogicalPlan plan) {
     );
   }
 
+  @Test
+  void table_support_write_builder_should_be_replaced() {
+    Mockito.reset(table, tableScanBuilder);
+    TableWriteBuilder writeBuilder = Mockito.mock(TableWriteBuilder.class);
+    when(table.createWriteBuilder(any())).thenReturn(writeBuilder);
+
+    assertEquals(
+        writeBuilder,
+        optimize(write(values(), table))
+    );
+  }
+
+  @Test
+  void table_not_support_write_builder_should_report_error() {
+    Mockito.reset(table, tableScanBuilder);
+    Table table = new Table() {
+      @Override
+      public Map<String, ExprType> getFieldTypes() {
+        return null;
+      }
+
+      @Override
+      public PhysicalPlan implement(LogicalPlan plan) {
+        return null;
+      }
+    };
+
+    assertThrows(UnsupportedOperationException.class,
+        () -> table.createWriteBuilder(null));
+  }
+
   private LogicalPlan optimize(LogicalPlan plan) {
     final LogicalPlanOptimizer optimizer = LogicalPlanOptimizer.create();
     final LogicalPlan optimize = optimizer.optimize(plan);

core/src/test/java/org/opensearch/sql/planner/optimizer/pattern/PatternsTest.java

@@ -35,5 +35,6 @@ void source_is_empty() {
   void table_is_empty() {
     plan = Mockito.mock(LogicalFilter.class);
     assertFalse(Patterns.table().getFunction().apply(plan).isPresent());
+    assertFalse(Patterns.writeTable().getFunction().apply(plan).isPresent());
   }
 }