feat: scala-like api for dl

joker-star-l · joker-star-l · commit 141bcace2459 · 2025-03-05T20:02:19.000+08:00
diff --git a/wayang-api/wayang-api-scala-java/src/main/scala/org/apache/wayang/api/DataQuanta.scala b/wayang-api/wayang-api-scala-java/src/main/scala/org/apache/wayang/api/DataQuanta.scala
@@ -588,6 +588,18 @@ class DataQuanta[Out: ClassTag](val operator: ElementaryOperator, outputIndex: I
     predictOperator
   }
 
+  def predictJava[ThatOut: ClassTag, Result: ClassTag](
+    that: DataQuanta[ThatOut]
+    ): DataQuanta[Result] = {
+    val predictOperator = new PredictOperator(
+      implicitly[ClassTag[ThatOut]].runtimeClass,
+      implicitly[ClassTag[Result]].runtimeClass
+    )
+    this.connectTo(predictOperator, 0)
+    that.connectTo(predictOperator, 1)
+    predictOperator
+  }
+
   def dlTraining[ThatOut: ClassTag](
     model: DLModel,
     option: DLTrainingOperator.Option,
@@ -616,6 +628,24 @@ class DataQuanta[Out: ClassTag](val operator: ElementaryOperator, outputIndex: I
     dlTrainingOperator
   }
 
+
+  def dlTrainingJava[ThatOut: ClassTag](
+    model: DLModel,
+    option: DLTrainingOperator.Option,
+    that: DataQuanta[ThatOut]
+  ): DataQuanta[DLModel] = {
+    val dlTrainingOperator = new DLTrainingOperator(
+      model,
+      option,
+      implicitly[ClassTag[Out]].runtimeClass,
+      implicitly[ClassTag[ThatOut]].runtimeClass
+    )
+
+    this.connectTo(dlTrainingOperator, 0)
+    that.connectTo(dlTrainingOperator, 1)
+    dlTrainingOperator
+  }
+
   /**
     * Feeds this and a further instance into a [[CoGroupOperator]].
     *
diff --git a/wayang-api/wayang-api-scala-java/src/main/scala/org/apache/wayang/api/DataQuantaBuilder.scala b/wayang-api/wayang-api-scala-java/src/main/scala/org/apache/wayang/api/DataQuantaBuilder.scala
@@ -27,7 +27,8 @@ import java.util.{Collection => JavaCollection}
 import org.apache.wayang.api.graph.{Edge, EdgeDataQuantaBuilder, EdgeDataQuantaBuilderDecorator}
 import org.apache.wayang.api.util.{DataQuantaBuilderCache, TypeTrap}
 import org.apache.wayang.basic.data.{Record, Tuple2 => RT2}
-import org.apache.wayang.basic.operators.{GlobalReduceOperator, LocalCallbackSink, MapOperator, SampleOperator}
+import org.apache.wayang.basic.model.{DLModel, Model}
+import org.apache.wayang.basic.operators.{DLTrainingOperator, GlobalReduceOperator, LocalCallbackSink, MapOperator, SampleOperator}
 import org.apache.wayang.commons.util.profiledb.model.Experiment
 import org.apache.wayang.core.function.FunctionDescriptor.{SerializableBiFunction, SerializableBinaryOperator, SerializableFunction, SerializableIntUnaryOperator, SerializablePredicate}
 import org.apache.wayang.core.optimizer.ProbabilisticDoubleInterval
@@ -273,6 +274,30 @@ trait DataQuantaBuilder[+This <: DataQuantaBuilder[_, Out], Out] extends Logging
                          thatKeyUdf: SerializableFunction[ThatOut, Key]) =
     new JoinDataQuantaBuilder(this, that, thisKeyUdf, thatKeyUdf)
 
+  /**
+   * Feed the built [[DataQuanta]] of this and the given instance into a
+   * [[org.apache.wayang.basic.operators.DLTrainingOperator]].
+   *
+   * @param that       the other [[DataQuantaBuilder]] to join with
+   * @param model      model for the [[org.apache.wayang.basic.operators.DLTrainingOperator]]
+   * @param option     option for the [[org.apache.wayang.basic.operators.DLTrainingOperator]]
+   * @return a [[DLTrainingDataQuantaBuilder]]
+   */
+  def dlTraining[ThatOut](that: DataQuantaBuilder[_, ThatOut],
+                          model: DLModel,
+                          option: DLTrainingOperator.Option) =
+    new DLTrainingDataQuantaBuilder(this, that, model, option)
+
+  /**
+   * Feed the built [[DataQuanta]] of this and the given instance into a
+   * [[org.apache.wayang.basic.operators.PredictOperator]].
+   *
+   * @param that       the other [[DataQuantaBuilder]] to join with
+   * @return a [[PredictDataQuantaBuilder]]
+   */
+  def predict[ThatOut, Result](that: DataQuantaBuilder[_, ThatOut], resultType: Class[Result]) =
+    new PredictDataQuantaBuilder(this.asInstanceOf[DataQuantaBuilder[_, Model]], that, resultType)
+
   /**
     * Feed the built [[DataQuanta]] of this and the given instance into a
     * [[org.apache.wayang.basic.operators.CoGroupOperator]].
@@ -1336,6 +1361,53 @@ class JoinDataQuantaBuilder[In0, In1, Key](inputDataQuanta0: DataQuantaBuilder[_
 
 }
 
+/**
+ * [[DataQuantaBuilder]] implementation for [[org.apache.wayang.basic.operators.DLTrainingOperator]]s.
+ *
+ * @param inputDataQuanta0 [[DataQuantaBuilder]] for the first input [[DataQuanta]]
+ * @param inputDataQuanta1 [[DataQuantaBuilder]] for the first input [[DataQuanta]]
+ * @param model            model for the [[org.apache.wayang.basic.operators.DLTrainingOperator]]
+ * @param option           option for the [[org.apache.wayang.basic.operators.DLTrainingOperator]]
+ */
+class DLTrainingDataQuantaBuilder[In0, In1](inputDataQuanta0: DataQuantaBuilder[_, In0],
+                                        inputDataQuanta1: DataQuantaBuilder[_, In1],
+                                        model: DLModel,
+                                        option: DLTrainingOperator.Option)
+                                          (implicit javaPlanBuilder: JavaPlanBuilder)
+  extends BasicDataQuantaBuilder[DLTrainingDataQuantaBuilder[In0, In1], DLModel] {
+
+  // Since we are currently not looking at type parameters, we can statically determine the output type.
+  locally {
+    this.outputTypeTrap.dataSetType = dataSetType[DLModel]
+  }
+
+  override protected def build =
+    inputDataQuanta0.dataQuanta()
+      .dlTrainingJava(model, option, inputDataQuanta1.dataQuanta())(inputDataQuanta1.classTag)
+}
+
+/**
+ * [[DataQuantaBuilder]] implementation for [[org.apache.wayang.basic.operators.PredictOperator]]s.
+ *
+ * @param inputDataQuanta0 [[DataQuantaBuilder]] for the first input [[DataQuanta]]
+ * @param inputDataQuanta1 [[DataQuantaBuilder]] for the first input [[DataQuanta]]
+ */
+class PredictDataQuantaBuilder[In1, Out](inputDataQuanta0: DataQuantaBuilder[_, Model],
+                                         inputDataQuanta1: DataQuantaBuilder[_, In1],
+                                         outClass: Class[Out])
+                                         (implicit javaPlanBuilder: JavaPlanBuilder)
+  extends BasicDataQuantaBuilder[PredictDataQuantaBuilder[In1, Out], Out] {
+
+  // Since we are currently not looking at type parameters, we can statically determine the output type.
+  locally {
+    this.outputTypeTrap.dataSetType = dataSetType[Out]
+  }
+
+  override protected def build =
+    inputDataQuanta0.dataQuanta().
+      predictJava(inputDataQuanta1.dataQuanta())(inputDataQuanta1.classTag, ClassTag.apply(outClass))
+}
+
 /**
   * [[DataQuantaBuilder]] implementation for [[org.apache.wayang.basic.operators.CoGroupOperator]]s.
   *
diff --git a/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisIT.java b/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisIT.java
@@ -57,7 +57,7 @@ public class TensorflowIrisIT {
             "Iris-virginica", 2
     );
 
-    @Ignore
+    @Test
     public void test() {
         final Tuple<Operator, Operator> trainSource = fileOperation(TRAIN_PATH, true);
         final Tuple<Operator, Operator> testSource = fileOperation(TEST_PATH, false);
diff --git a/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisScalaLikeApiIT.java b/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisScalaLikeApiIT.java
@@ -18,21 +18,20 @@
 
 package org.apache.wayang.tests;
 
+import org.apache.wayang.api.*;
 import org.apache.wayang.basic.model.DLModel;
 import org.apache.wayang.basic.model.op.*;
 import org.apache.wayang.basic.model.op.nn.CrossEntropyLoss;
 import org.apache.wayang.basic.model.op.nn.Linear;
 import org.apache.wayang.basic.model.op.nn.Sigmoid;
 import org.apache.wayang.basic.model.optimizer.Adam;
 import org.apache.wayang.basic.model.optimizer.Optimizer;
-import org.apache.wayang.basic.operators.*;
+import org.apache.wayang.basic.operators.DLTrainingOperator;
 import org.apache.wayang.core.api.WayangContext;
-import org.apache.wayang.core.plan.wayangplan.Operator;
-import org.apache.wayang.core.plan.wayangplan.WayangPlan;
 import org.apache.wayang.core.util.Tuple;
 import org.apache.wayang.java.Java;
 import org.apache.wayang.tensorflow.Tensorflow;
-import org.junit.Ignore;
+import org.junit.Test;
 
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -56,22 +55,30 @@ public class TensorflowIrisScalaLikeApiIT {
             "Iris-virginica", 2
     );
 
-    @Ignore
+    @Test
     public void test() {
-        final Tuple<Operator, Operator> trainSource = fileOperation(TRAIN_PATH, true);
-        final Tuple<Operator, Operator> testSource = fileOperation(TEST_PATH, false);
+        WayangContext wayangContext = new WayangContext()
+                .with(Java.basicPlugin())
+                .with(Tensorflow.plugin());
+
+        JavaPlanBuilder plan = new JavaPlanBuilder(wayangContext);
+
+        final Tuple<DataQuantaBuilder<?, float[]>, DataQuantaBuilder<?, Integer>> trainSource =
+                fileOperation(plan, TRAIN_PATH, true);
+        final Tuple<DataQuantaBuilder<?, float[]>, DataQuantaBuilder<?, Integer>> testSource =
+                fileOperation(plan, TEST_PATH, false);
 
         /* training features */
-        Operator trainXSource = trainSource.field0;
+        DataQuantaBuilder<?, float[]> trainXSource = trainSource.field0;
 
         /* training labels */
-        Operator trainYSource = trainSource.field1;
+        DataQuantaBuilder<?, Integer> trainYSource = trainSource.field1;
 
         /* test features */
-        Operator testXSource = testSource.field0;
+        DataQuantaBuilder<?, float[]> testXSource = testSource.field0;
 
         /* test labels */
-        Operator testYSource = testSource.field1;
+        DataQuantaBuilder<?, Integer> testYSource = testSource.field1;
 
         /* model */
         Op l1 = new Linear(4, 32, true);
@@ -110,17 +117,15 @@ public void test() {
         option.setAccuracyCalculation(acc);
 
         /* training operator */
-        DLTrainingOperator<float[], Integer> trainingOperator = new DLTrainingOperator<>(
-                model, option, float[].class, Integer.class
-        );
+        DLTrainingDataQuantaBuilder<float[], Integer> trainingOperator =
+                trainXSource.dlTraining(trainYSource, model, option);
 
         /* predict operator */
-        PredictOperator<float[], float[]> predictOperator = new PredictOperator<>(
-                float[].class, float[].class
-        );
+        PredictDataQuantaBuilder<float[], float[]> predictOperator =
+                trainingOperator.predict(testXSource, float[].class);
 
         /* map to label */
-        MapOperator<float[], Integer> mapOperator = new MapOperator<>(array -> {
+        MapDataQuantaBuilder<float[], Integer> mapOperator = predictOperator.map(array -> {
             int maxIdx = 0;
             float maxVal = array[0];
             for (int i = 1; i < array.length; i++) {
@@ -130,69 +135,47 @@ public void test() {
                 }
             }
             return maxIdx;
-        }, float[].class, Integer.class);
+        });
 
         /* sink */
-        List<Integer> predicted = new ArrayList<>();
-        LocalCallbackSink<Integer> predictedSink = LocalCallbackSink.createCollectingSink(predicted, Integer.class);
-
-        List<Integer> groundTruth = new ArrayList<>();
-        LocalCallbackSink<Integer> groundTruthSink = LocalCallbackSink.createCollectingSink(groundTruth, Integer.class);
-
-        trainXSource.connectTo(0, trainingOperator, 0);
-        trainYSource.connectTo(0, trainingOperator, 1);
-        trainingOperator.connectTo(0, predictOperator, 0);
-        testXSource.connectTo(0, predictOperator, 1);
-        predictOperator.connectTo(0, mapOperator, 0);
-        mapOperator.connectTo(0, predictedSink, 0);
-        testYSource.connectTo(0, groundTruthSink, 0);
-
-        WayangPlan wayangPlan = new WayangPlan(predictedSink, groundTruthSink);
-
-        WayangContext wayangContext = new WayangContext();
-        wayangContext.register(Java.basicPlugin());
-        wayangContext.register(Tensorflow.plugin());
-        wayangContext.execute(wayangPlan);
+        List<Integer> predicted = new ArrayList<>(mapOperator.collect());
+        // fixme: Currently, wayang's scala-like api only supports a single collect,
+        //  so it is not possible to collect multiple result lists in a single plan.
+//        List<Integer> groundTruth = new ArrayList<>(testYSource.collect());
 
         System.out.println("predicted:    " + predicted);
-        System.out.println("ground truth: " + groundTruth);
-
-        float success = 0;
-        for (int i = 0; i < predicted.size(); i++) {
-            if (predicted.get(i).equals(groundTruth.get(i))) {
-                success += 1;
-            }
-        }
-        System.out.println("test accuracy: " + success / predicted.size());
+//        System.out.println("ground truth: " + groundTruth);
+
+//        float success = 0;
+//        for (int i = 0; i < predicted.size(); i++) {
+//            if (predicted.get(i).equals(groundTruth.get(i))) {
+//                success += 1;
+//            }
+//        }
+//        System.out.println("test accuracy: " + success / predicted.size());
     }
 
-    public static Tuple<Operator, Operator> fileOperation(URI uri, boolean random) {
-        TextFileSource textFileSource = new TextFileSource(uri.toString());
-        MapOperator<String, Tuple> mapOperator = new MapOperator<>(line -> {
-            String[] parts = line.split(",");
-            float[] x = new float[parts.length - 1];
-            for (int i = 0; i < x.length; i++) {
-                x[i] = Float.parseFloat(parts[i]);
-            }
-            int y = LABEL_MAP.get(parts[parts.length - 1]);
-            return new Tuple<>(x, y);
-        }, String.class, Tuple.class);
-
-        MapOperator<Tuple, float[]> mapX = new MapOperator<>(tuple -> (float[]) tuple.field0, Tuple.class, float[].class);
-        MapOperator<Tuple, Integer> mapY = new MapOperator<>(tuple -> (Integer) tuple.field1, Tuple.class, Integer.class);
+    public static Tuple<DataQuantaBuilder<?, float[]>, DataQuantaBuilder<?, Integer>>
+    fileOperation(JavaPlanBuilder plan, URI uri, boolean random) {
+        DataQuantaBuilder<?, String> textFileSource = plan.readTextFile(uri.toString());
 
         if (random) {
             Random r = new Random();
-            SortOperator<String, Integer> randomOperator = new SortOperator<>(e -> r.nextInt(), String.class, Integer.class);
-
-            textFileSource.connectTo(0, randomOperator, 0);
-            randomOperator.connectTo(0, mapOperator, 0);
-        } else {
-            textFileSource.connectTo(0, mapOperator, 0);
+            textFileSource = textFileSource.sort(e -> r.nextInt());
         }
 
-        mapOperator.connectTo(0, mapX, 0);
-        mapOperator.connectTo(0, mapY, 0);
+        MapDataQuantaBuilder<String, Tuple<float[], Integer>> mapXY = textFileSource.map(line -> {
+                    String[] parts = line.split(",");
+                    float[] x = new float[parts.length - 1];
+                    for (int i = 0; i < x.length; i++) {
+                        x[i] = Float.parseFloat(parts[i]);
+                    }
+                    int y = LABEL_MAP.get(parts[parts.length - 1]);
+                    return new Tuple<>(x, y);
+                });
+
+        MapDataQuantaBuilder<Tuple<float[], Integer>, float[]> mapX = mapXY.map(tuple -> tuple.field0);
+        MapDataQuantaBuilder<Tuple<float[], Integer>, Integer> mapY = mapXY.map(tuple -> tuple.field1);
 
         return new Tuple<>(mapX, mapY);
     }
