ML cost model integration for Wayang

wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/ComplexityUtils.java

@@ -0,0 +1,74 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.wayang.basic.util;
+
+import org.apache.wayang.basic.operators.CoGroupOperator;
+import org.apache.wayang.basic.operators.FilterOperator;
+import org.apache.wayang.basic.operators.FlatMapOperator;
+import org.apache.wayang.basic.operators.GlobalReduceOperator;
+import org.apache.wayang.basic.operators.GroupByOperator;
+import org.apache.wayang.basic.operators.JoinOperator;
+import org.apache.wayang.basic.operators.LoopOperator;
+import org.apache.wayang.basic.operators.MapOperator;
+import org.apache.wayang.basic.operators.MapPartitionsOperator;
+import org.apache.wayang.basic.operators.MaterializedGroupByOperator;
+import org.apache.wayang.basic.operators.ReduceByOperator;
+import org.apache.wayang.basic.operators.ReduceOperator;
+import org.apache.wayang.basic.operators.SortOperator;
+import org.apache.wayang.core.optimizer.ComplexityClass;
+import org.apache.wayang.core.plan.wayangplan.Operator;
+
+public class ComplexityUtils {
+    /**
+     * Infer complexity class from a given operator's descriptors.
+     * @param operator
+     * @return {@link ComplexityClass#LOGARITHMIC}, {@link ComplexityClass#LINEAR}, {@link ComplexityClass#QUADRATIC} or {@link ComplexityClass#SUPERQUADRATIC}. {@link ComplexityClass#LINEAR} on default
+     */
+    public static ComplexityClass inferFromOperator(final Operator operator) {
+        if (operator instanceof final ReduceByOperator reduceBy) {
+            return reduceBy.getReduceDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final ReduceOperator reduce) {
+            return reduce.getReduceDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final GlobalReduceOperator globalReduce) {
+            return globalReduce.getReduceDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final CoGroupOperator coGroup) {
+            return coGroup.getKeyDescriptor0().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final GroupByOperator groupBy) {
+            return groupBy.getKeyDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final MaterializedGroupByOperator matGroupBy) {
+            return matGroupBy.getKeyDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final SortOperator sort) {
+            return sort.getKeyDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final JoinOperator join) {
+            return join.getKeyDescriptor0().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final MapOperator map) {
+            return map.getFunctionDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final FlatMapOperator flatMap) {
+            return flatMap.getFunctionDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final MapPartitionsOperator mapPartitions) {
+            return mapPartitions.getFunctionDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final FilterOperator filter) {
+            return filter.getPredicateDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        } else if (operator instanceof final LoopOperator loop) {
+            return loop.getCriterionDescriptor().getComplexityClass().orElse(ComplexityClass.LINEAR);
+        }
+
+        return ComplexityClass.LINEAR;
+    }
+}

wayang-commons/wayang-core/src/main/java/org/apache/wayang/core/function/FunctionDescriptor.java

@@ -18,6 +18,7 @@
 
 package org.apache.wayang.core.function;
 
+import org.apache.wayang.core.optimizer.ComplexityClass;
 import org.apache.wayang.core.optimizer.ProbabilisticDoubleInterval;
 import org.apache.wayang.core.optimizer.costs.LoadEstimator;
 import org.apache.wayang.core.optimizer.costs.LoadProfileEstimator;
@@ -32,10 +33,13 @@
  */
 public abstract class FunctionDescriptor implements Serializable {
 
-    public FunctionDescriptor() {}
+    public FunctionDescriptor() {
+    }
 
     private LoadProfileEstimator loadProfileEstimator;
 
+    private ComplexityClass complexityClass = null;
+
     public FunctionDescriptor(LoadProfileEstimator loadProfileEstimator) {
         this.setLoadProfileEstimator(loadProfileEstimator);
     }
@@ -48,6 +52,14 @@ public Optional<LoadProfileEstimator> getLoadProfileEstimator() {
         return Optional.ofNullable(this.loadProfileEstimator);
     }
 
+    public Optional<ComplexityClass> getComplexityClass(){
+        return Optional.ofNullable(complexityClass);
+    }
+
+    public void setComplexityClass(final ComplexityClass complexityClass){
+        this.complexityClass = complexityClass;
+    }
+
     /**
      * Utility method to retrieve the selectivity of a {@link FunctionDescriptor}
      *

wayang-commons/wayang-core/src/main/java/org/apache/wayang/core/optimizer/ComplexityClass.java

@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.wayang.core.optimizer;
+
+public enum ComplexityClass {
+    LINEAR,
+    LOGARITHMIC,
+    QUADRATIC,
+    SUPERQUADRATIC
+}

wayang-plugins/pom.xml

@@ -39,6 +39,6 @@
     <modules>
         <module>wayang-iejoin</module>
         <module>wayang-spatial</module>
+        <module>wayang-ml</module>
     </modules>
-
 </project>

wayang-plugins/wayang-ml/pom.xml

@@ -0,0 +1,157 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Licensed to the Apache Software Foundation (ASF) under one
+  ~ or more contributor license agreements.  See the NOTICE file
+  ~ distributed with this work for additional information
+  ~ regarding copyright ownership.  The ASF licenses this file
+  ~ to you under the Apache License, Version 2.0 (the
+  ~ "License"); you may not use this file except in compliance
+  ~ with the License.  You may obtain a copy of the License at
+  ~
+  ~   http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing,
+  ~ software distributed under the License is distributed on an
+  ~ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  ~ KIND, either express or implied.  See the License for the
+  ~ specific language governing permissions and limitations
+  ~ under the License.
+  ~
+-->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>org.apache.wayang</groupId>
+        <artifactId>wayang-plugins</artifactId>
+        <version>1.1.2-SNAPSHOT</version>
+    </parent>
+
+    <artifactId>wayang-ml</artifactId>
+    <version>1.1.2-SNAPSHOT</version>
+
+    <properties>
+        <java-module-name>org.apache.wayang.extensions.ml</java-module-name>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-api-sql</artifactId>
+          <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>com.microsoft.onnxruntime</groupId>
+            <artifactId>onnxruntime</artifactId>
+            <version>1.21.1</version>
+        </dependency>
+        <!--dependency>
+            <groupId>com.microsoft.onnxruntime</groupId>
+            <artifactId>onnxruntime_gpu</artifactId>
+            <version>1.18.0</version>
+        </dependency-->
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-core</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-basic</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-java</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-spark</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-flink</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.flink</groupId>
+            <artifactId>flink-java</artifactId>
+            <version>${flink.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-giraph</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-generic-jdbc</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.reflections</groupId>
+            <artifactId>reflections</artifactId>
+            <version>0.10.2</version>
+        </dependency>
+        <dependency>
+          <groupId>org.apache.wayang</groupId>
+          <artifactId>wayang-benchmark</artifactId>
+          <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.wayang</groupId>
+            <artifactId>wayang-api-python</artifactId>
+            <version>1.1.2-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-dbcp2</artifactId>
+            <version>2.7.0</version>
+        </dependency>
+        <dependency>
+          <groupId>org.apache.spark</groupId>
+          <artifactId>spark-core_2.12</artifactId>
+          <version>${spark.version}</version>
+        </dependency>
+        <dependency>
+          <groupId>org.apache.spark</groupId>
+          <artifactId>spark-graphx_2.12</artifactId>
+          <version>${spark.version}</version>
+        </dependency>
+        <dependency>
+          <groupId>org.apache.spark</groupId>
+          <artifactId>spark-mllib_2.12</artifactId>
+          <version>${spark.version}</version>
+        </dependency>
+        <dependency>
+          <groupId>com.google.protobuf</groupId>
+          <artifactId>protobuf-java</artifactId>
+          <version>3.16.3</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.calcite</groupId>
+            <artifactId>calcite-core</artifactId>
+            <version>${calcite.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.calcite</groupId>
+            <artifactId>calcite-linq4j</artifactId>
+            <version>${calcite.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.calcite</groupId>
+            <artifactId>calcite-file</artifactId>
+            <version>${calcite.version}</version>
+        </dependency>
+      </dependencies>
+    <build>
+        <resources>
+            <resource>
+                <directory>src/main/resources</directory>
+            </resource>
+        </resources>
+    </build>
+</project>

wayang-plugins/wayang-ml/src/main/java/org/apache/wayang/ml/MLContext.java

@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.wayang.ml;
+
+import java.util.Optional;
+
+import org.apache.logging.log4j.Level;
+import org.apache.wayang.core.api.Configuration;
+import org.apache.wayang.core.api.Job;
+import org.apache.wayang.core.api.WayangContext;
+import org.apache.wayang.core.plan.wayangplan.WayangPlan;
+import org.apache.wayang.core.util.ReflectionUtils;
+import org.apache.wayang.ml.encoding.OneHotMappings;
+import org.apache.wayang.ml.encoding.TreeEncoder;
+import org.apache.wayang.ml.util.Logging;
+
+/**
+ * This is the entry point for users to work with Wayang ML.
+ */
+public class MLContext extends WayangContext {
+    public MLContext() {
+        super();
+    }
+
+    public MLContext(final Configuration configuration) {
+        super(configuration);
+    }
+
+    /**
+     * Execute a plan.
+     *
+     * @param wayangPlan the plan to execute
+     * @param udfJars    JARs that declare the code for the UDFs
+     * @see ReflectionUtils#getDeclaringJar(Class)
+     */
+    @Override
+    public void execute(final WayangPlan wayangPlan, final String... udfJars) {
+        this.setLogLevel(Level.ERROR);
+        final Job wayangJob = this.createJob("", wayangPlan, udfJars);
+
+        final Configuration config = this.getConfiguration();
+        final Configuration jobConfig = wayangJob.getConfiguration();
+
+        wayangJob.execute();
+
+        if (config.getBooleanProperty("wayang.ml.experience.enabled")) {
+            final Optional<String> originalOption = config.getOptionalStringProperty("wayang.ml.experience.original");
+
+            final OneHotMappings mappings = new OneHotMappings();
+            final TreeEncoder encoder = new TreeEncoder(mappings);
+            final String original = originalOption.orElse(encoder.encode(wayangPlan, wayangJob.getOptimizationContext(), false).toString());
+
+            final Optional<String> choicesOption = config
+                    .getOptionalStringProperty("wayang.ml.experience.with-platforms");
+            final String withChoices = choicesOption
+                    .orElse(jobConfig.getStringProperty("wayang.ml.experience.with-platforms"));
+
+            final long execTime = jobConfig.getLongProperty("wayang.ml.experience.exec-time");
+
+            this.logExperience(original, withChoices, execTime);
+        }
+    }
+
+    private void logExperience(final String original, final String withChoices, final long execTime) {
+        if (!this.getConfiguration().getBooleanProperty("wayang.ml.experience.enabled")) {
+            return;
+        }
+
+        final String content = String.format("%s:%s:%d", original, withChoices, execTime);
+        Logging.writeToFile(content, this.getConfiguration().getStringProperty("wayang.ml.experience.file"));
+    }
+}