PPL Integration - Add implementation for KMeans algorithm

.github/workflows/sql-cli-release-workflow.yml

@@ -20,6 +20,17 @@ jobs:
       - name: Checkout SQL CLI
         uses: actions/checkout@v2
 
+      # dependencies: ml-commons
+      - name: Checkout ml-commons
+        uses: actions/checkout@v2
+        with:
+          repository: 'opensearch-project/ml-commons'
+          path: ml-commons
+          ref: 'main'
+      - name: Build ml-commons
+        working-directory: ./ml-commons
+        run: ./gradlew publishToMavenLocal
+
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:

.github/workflows/sql-cli-test-and-build-workflow.yml

@@ -17,6 +17,17 @@ jobs:
       - name: Checkout SQL CLI
         uses: actions/checkout@v2
 
+      # dependencies: ml-commons
+      - name: Checkout ml-commons
+        uses: actions/checkout@v2
+        with:
+          repository: 'opensearch-project/ml-commons'
+          path: ml-commons
+          ref: 'main'
+      - name: Build ml-commons
+        working-directory: ./ml-commons
+        run: ./gradlew publishToMavenLocal
+
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:

.github/workflows/sql-release-workflow.yml

@@ -15,6 +15,17 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
+      # dependencies: ml-commons
+      - name: Checkout ml-commons
+        uses: actions/checkout@v2
+        with:
+          repository: 'opensearch-project/ml-commons'
+          path: ml-commons
+          ref: 'main'
+      - name: Build ml-commons
+        working-directory: ./ml-commons
+        run: ./gradlew publishToMavenLocal
+
       - name: Checkout SQL
         uses: actions/checkout@v1
 

.github/workflows/sql-test-and-build-workflow.yml

@@ -17,6 +17,17 @@ jobs:
       uses: actions/setup-java@v1
       with:
         java-version: 1.14
+
+    # dependencies: ml-commons
+    - name: Checkout ml-commons
+      uses: actions/checkout@v2
+      with:
+        repository: 'opensearch-project/ml-commons'
+        path: ml-commons
+        ref: 'main'
+    - name: Build ml-commons
+      working-directory: ./ml-commons
+      run: ./gradlew publishToMavenLocal
 
     - name: Build with Gradle
       run: ./gradlew build assemble -Dopensearch.version=${{ env.OPENSEARCH_VERSION }}

core/build.gradle

@@ -46,6 +46,8 @@ dependencies {
     compile group: 'org.apache.commons', name: 'commons-lang3', version: '3.10'
     compile group: 'com.facebook.presto', name: 'presto-matching', version: '0.240'
     compile group: 'org.apache.commons', name: 'commons-math3', version: '3.6.1'
+    compile group: 'org.opensearch.ml', name:'opensearch-ml-client', version: '1.3.0.0'
+    compile group: 'org.opensearch', name: 'opensearch', version: "1.3.0-SNAPSHOT"
     compile project(':common')
 
     testImplementation('org.junit.jupiter:junit-jupiter:5.6.2')

core/src/main/java/org/opensearch/sql/analysis/Analyzer.java

@@ -36,6 +36,7 @@
 import org.opensearch.sql.ast.tree.Eval;
 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.Limit;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
@@ -47,6 +48,7 @@
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
 import org.opensearch.sql.data.model.ExprMissingValue;
+import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.Expression;
@@ -60,6 +62,7 @@
 import org.opensearch.sql.planner.logical.LogicalEval;
 import org.opensearch.sql.planner.logical.LogicalFilter;
 import org.opensearch.sql.planner.logical.LogicalLimit;
+import org.opensearch.sql.planner.logical.LogicalMLCommons;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalProject;
 import org.opensearch.sql.planner.logical.LogicalRareTopN;
@@ -366,6 +369,20 @@ public LogicalPlan visitValues(Values node, AnalysisContext context) {
     return new LogicalValues(valueExprs);
   }
 
+  /**
+   * Build {@link LogicalMLCommons} for Kmeans command.
+   */
+  @Override
+  public LogicalPlan visitKmeans(Kmeans node, AnalysisContext context) {
+    LogicalPlan child = node.getChild().get(0).accept(this, context);
+    List<Argument> options = node.getOptions();
+
+    TypeEnvironment currentEnv = context.peek();
+    currentEnv.define(new Symbol(Namespace.FIELD_NAME, "ClusterID"), ExprCoreType.INTEGER);
+
+    return new LogicalMLCommons(child, "kmeans", options);
+  }
+
   /**
    * The first argument is always "asc", others are optional.
    * Given nullFirst argument, use its value. Otherwise just use DEFAULT_ASC/DESC.

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -37,6 +37,7 @@
 import org.opensearch.sql.ast.tree.Eval;
 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.Limit;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
@@ -234,4 +235,8 @@ public T visitLimit(Limit node, C context) {
   public T visitSpan(Span node, C context) {
     return visitChildren(node, context);
   }
+
+  public T visitKmeans(Kmeans node, C context) {
+    return visitChildren(node, context);
+  }
 }

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

@@ -0,0 +1,40 @@
+package org.opensearch.sql.ast.tree;
+
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Argument;
+
+@Getter
+@Setter
+@ToString
+@EqualsAndHashCode(callSuper = true)
+@RequiredArgsConstructor
+@AllArgsConstructor
+public class Kmeans extends UnresolvedPlan {
+  private UnresolvedPlan child;
+
+  private final List<Argument> options;
+
+  @Override
+  public UnresolvedPlan attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitKmeans(this, context);
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return ImmutableList.of(this.child);
+  }
+}

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

@@ -0,0 +1,38 @@
+package org.opensearch.sql.planner.logical;
+
+import java.util.Collections;
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.ast.expression.Argument;
+
+/**
+ * ml-commons logical plan.
+ */
+@Getter
+@ToString
+@EqualsAndHashCode(callSuper = true)
+public class LogicalMLCommons extends LogicalPlan {
+  private final String algorithm;
+
+  private final List<Argument> arguments;
+
+  /**
+   * Constructor of LogicalMLCommons.
+   * @param child child logical plan
+   * @param algorithm algorithm name
+   * @param arguments arguments of the algorithm
+   */
+  public LogicalMLCommons(LogicalPlan child, String algorithm,
+                          List<Argument> arguments) {
+    super(Collections.singletonList(child));
+    this.algorithm = algorithm;
+    this.arguments = arguments;
+  }
+
+  @Override
+  public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitMLCommons(this, context);
+  }
+}

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

@@ -69,4 +69,8 @@ public R visitRareTopN(LogicalRareTopN plan, C context) {
   public R visitLimit(LogicalLimit plan, C context) {
     return visitNode(plan, context);
   }
+
+  public R visitMLCommons(LogicalMLCommons plan, C context) {
+    return visitNode(plan, context);
+  }
 }

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

@@ -0,0 +1,155 @@
+package org.opensearch.sql.planner.physical;
+
+import com.google.common.collect.ImmutableMap;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.opensearch.ml.client.MachineLearningClient;
+import org.opensearch.ml.common.dataframe.ColumnMeta;
+import org.opensearch.ml.common.dataframe.ColumnValue;
+import org.opensearch.ml.common.dataframe.DataFrame;
+import org.opensearch.ml.common.dataframe.DataFrameBuilder;
+import org.opensearch.ml.common.dataframe.Row;
+import org.opensearch.ml.common.dataset.DataFrameInputDataset;
+import org.opensearch.ml.common.parameter.FunctionName;
+import org.opensearch.ml.common.parameter.KMeansParams;
+import org.opensearch.ml.common.parameter.MLAlgoParams;
+import org.opensearch.ml.common.parameter.MLInput;
+import org.opensearch.ml.common.parameter.MLPredictionOutput;
+import org.opensearch.sql.ast.expression.Argument;
+import org.opensearch.sql.data.model.ExprDoubleValue;
+import org.opensearch.sql.data.model.ExprIntegerValue;
+import org.opensearch.sql.data.model.ExprStringValue;
+import org.opensearch.sql.data.model.ExprTupleValue;
+import org.opensearch.sql.data.model.ExprValue;
+
+/**
+ * ml-commons Physical operator to call machine learning interface to get results for
+ * algorithm execution.
+ */
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+public class MLCommonsOperator extends PhysicalPlan {
+  @Getter
+  private final PhysicalPlan input;
+
+  @Getter
+  private final String algorithm;
+
+  @Getter
+  private final List<Argument> arguments;
+
+  @Getter
+  private final MachineLearningClient machineLearningClient;
+
+  @EqualsAndHashCode.Exclude
+  private Iterator<ExprValue> iterator;
+
+  @Override
+  public void open() {
+    super.open();
+    DataFrame inputDataFrame = generateInputDataset();
+    MLAlgoParams mlAlgoParams = convertArgumentToMLParameter(arguments.get(0), algorithm);
+    MLInput mlinput = MLInput.builder()
+            .algorithm(FunctionName.valueOf(algorithm.toUpperCase()))
+            .parameters(mlAlgoParams)
+            .inputDataset(new DataFrameInputDataset(inputDataFrame))
+            .build();
+    MLPredictionOutput predictionResult = (MLPredictionOutput) machineLearningClient
+            .trainAndPredict(mlinput)
+            .actionGet(30, TimeUnit.SECONDS);
+    Iterator<Row> inputRowIter = inputDataFrame.iterator();
+    Iterator<Row> resultRowIter = predictionResult.getPredictionResult().iterator();
+    iterator = new Iterator<ExprValue>() {
+      @Override
+      public boolean hasNext() {
+          return inputRowIter.hasNext();
+      }
+
+      @Override
+      public ExprValue next() {
+        ImmutableMap.Builder<String, ExprValue> resultBuilder = new ImmutableMap.Builder<>();
+        resultBuilder.putAll(convertRowIntoExprValue(inputDataFrame.columnMetas(),
+                inputRowIter.next()));
+        resultBuilder.putAll(convertRowIntoExprValue(
+                predictionResult.getPredictionResult().columnMetas(),
+                resultRowIter.next()));
+        return ExprTupleValue.fromExprValueMap(resultBuilder.build());
+      }
+    };
+  }
+
+  @Override
+  public <R, C> R accept(PhysicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitMLCommons(this, context);
+  }
+
+  @Override
+  public boolean hasNext() {
+    return iterator.hasNext();
+  }
+
+  @Override
+  public ExprValue next() {
+    return iterator.next();
+  }
+
+  @Override
+  public List<PhysicalPlan> getChild() {
+    return Collections.singletonList(input);
+  }
+
+  protected MLAlgoParams convertArgumentToMLParameter(Argument argument, String algorithm) {
+    switch (FunctionName.valueOf(algorithm.toUpperCase())) {
+      case KMEANS:
+        return KMeansParams.builder().centroids((Integer) argument.getValue().getValue()).build();
+
+      default:
+        throw new IllegalArgumentException("unsupported argument type:"
+                + argument.getValue().getType());
+    }
+  }
+
+  private Map<String, ExprValue> convertRowIntoExprValue(ColumnMeta[] columnMetas, Row row) {
+    ImmutableMap.Builder<String, ExprValue> resultBuilder = new ImmutableMap.Builder<>();
+    for (int i = 0; i < columnMetas.length; i++) {
+      ColumnValue columnValue = row.getValue(i);
+      String resultKeyName = columnMetas[i].getName();
+      switch (columnValue.columnType()) {
+        case INTEGER:
+          resultBuilder.put(resultKeyName, new ExprIntegerValue(columnValue.intValue()));
+          break;
+        case DOUBLE:
+          resultBuilder.put(resultKeyName, new ExprDoubleValue(columnValue.doubleValue()));
+          break;
+        case STRING:
+          resultBuilder.put(resultKeyName, new ExprStringValue(columnValue.stringValue()));
+          break;
+        default:
+          break;
+      }
+    }
+    return resultBuilder.build();
+  }
+
+  private DataFrame generateInputDataset() {
+    List<Map<String, Object>> inputData = new LinkedList<>();
+    while (input.hasNext()) {
+      Map<String, Object> items = new HashMap<>();
+      input.next().tupleValue().forEach((key, value) -> {
+        items.put(key, value.value());
+      });
+      inputData.add(items);
+    }
+
+    return DataFrameBuilder.load(inputData);
+  }
+}
+

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

@@ -72,4 +72,8 @@ public R visitLimit(LimitOperator node, C context) {
     return visitNode(node, context);
   }
 
+  public R visitMLCommons(PhysicalPlan node, C context) {
+    return visitNode(node, context);
+  }
+
 }