Examples in scala

Signed-off-by: Nelson Arapé <[email protected]>
flyteorg · Jan 24, 2023 · 1350013 · 1350013
1 parent be2e860
commit 1350013
Show file tree

Hide file tree

Showing 8 changed files with 46 additions and 55 deletions.
diff --git a/...ekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/AddQuestionTask.scala b/...ekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/AddQuestionTask.scala
@@ -46,18 +46,3 @@ class AddQuestionTask
   override def run(input: AddQuestionTaskInput): AddQuestionTaskOutput =
     AddQuestionTaskOutput(ofString(s"${input.greeting.get} How are you?"))
 }
-
-object AddQuestionTask {
-
-  /** Binds input data to this task
-    *
-    * @param greeting
-    *   the input greeting message
-    * @return
-    *   a transformed instance of this class with input data
-    */
-  def apply(
-      greeting: SdkBindingData[String]
-  ): SdkTransform[AddQuestionTaskOutput] =
-    new AddQuestionTask().withInput("greeting", greeting)
-}
diff --git a/...ples-scala/src/main/scala/org/flyte/examples/flytekitscala/DynamicFibonacciWorkflow.scala b/...ples-scala/src/main/scala/org/flyte/examples/flytekitscala/DynamicFibonacciWorkflow.scala
@@ -23,9 +23,14 @@ import org.flyte.flytekitscala.{
   SdkScalaWorkflowBuilder
 }
 
+case class DynamicFibonacciWorkflowInput(n: SdkBindingData[Long])
 case class DynamicFibonacciWorkflowOutput(output: SdkBindingData[Long])
 class DynamicFibonacciWorkflow
-    extends SdkScalaWorkflow[DynamicFibonacciWorkflowOutput](
+    extends SdkScalaWorkflow[
+      DynamicFibonacciWorkflowInput,
+      DynamicFibonacciWorkflowOutput
+    ](
+      SdkScalaType[DynamicFibonacciWorkflowInput],
       SdkScalaType[DynamicFibonacciWorkflowOutput]
     ) {
 
@@ -34,7 +39,8 @@ class DynamicFibonacciWorkflow
 
     val fibonacci = builder.apply(
       "fibonacci",
-      new DynamicFibonacciWorkflowTask().withInput("n", n)
+      new DynamicFibonacciWorkflowTask(),
+      DynamicFibonacciWorkflowTaskInput(n)
     )
 
     builder.output("output", fibonacci.getOutputs.output)

diff --git a/...-scala/src/main/scala/org/flyte/examples/flytekitscala/DynamicFibonacciWorkflowTask.scala b/...-scala/src/main/scala/org/flyte/examples/flytekitscala/DynamicFibonacciWorkflowTask.scala
@@ -53,7 +53,11 @@ class DynamicFibonacciWorkflowTask
       else
         fib(
           n + 1,
-          builder(s"fib-${n + 1}", SumTask(value, prev)).getOutputs.c,
+          builder(
+            s"fib-${n + 1}",
+            new SumTask(),
+            SumTaskInput(value, prev)
+          ).getOutputs.c,
           value
         )
     }

diff --git a/...it-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/FibonacciWorkflow.scala b/...it-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/FibonacciWorkflow.scala
@@ -16,28 +16,42 @@
  */
 package org.flyte.examples.flytekitscala
 
-import org.flyte.flytekit.{SdkBindingData, SdkWorkflow, SdkWorkflowBuilder}
+import org.flyte.flytekit.SdkBindingData
 import org.flyte.flytekitscala.{
   SdkScalaWorkflowBuilder,
   SdkScalaType,
   SdkScalaWorkflow
 }
 
+case class FibonacciWorkflowInput(fib0: SdkBindingData[Long], fib1: SdkBindingData[Long])
 case class FibonacciWorkflowOutput(fib5: SdkBindingData[Long])
 
 class FibonacciWorkflow
-    extends SdkScalaWorkflow[FibonacciWorkflowOutput](
+    extends SdkScalaWorkflow[FibonacciWorkflowInput, FibonacciWorkflowOutput](
+      SdkScalaType[FibonacciWorkflowInput],
       SdkScalaType[FibonacciWorkflowOutput]
     ) {
 
   override def expand(builder: SdkScalaWorkflowBuilder): Unit = {
     val fib0 = builder.inputOfInteger("fib0", "Value for Fib0")
     val fib1 = builder.inputOfInteger("fib1", "Value for Fib1")
 
-    val fib2 = builder.apply("fib-2", SumTask(fib0, fib1)).getOutputs.c
-    val fib3 = builder.apply("fib-3", SumTask(fib1, fib2)).getOutputs.c
-    val fib4 = builder.apply("fib-4", SumTask(fib2, fib3)).getOutputs.c
-    val fib5 = builder.apply("fib-5", SumTask(fib3, fib4)).getOutputs.c
+    val fib2 = builder
+      .apply("fib-2", new SumTask(), SumTaskInput(fib0, fib1))
+      .getOutputs
+      .c
+    val fib3 = builder
+      .apply("fib-3", new SumTask(), SumTaskInput(fib1, fib2))
+      .getOutputs
+      .c
+    val fib4 = builder
+      .apply("fib-4", new SumTask(), SumTaskInput(fib2, fib3))
+      .getOutputs
+      .c
+    val fib5 = builder
+      .apply("fib-5", new SumTask(), SumTaskInput(fib3, fib4))
+      .getOutputs
+      .c
 
     builder.output("fib5", fib5, "Value for Fib5")
   }

diff --git a/flytekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/GreetTask.scala b/flytekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/GreetTask.scala
@@ -46,16 +46,3 @@ class GreetTask
   override def run(input: GreetTaskInput): GreetTaskOutput =
     GreetTaskOutput(ofString(s"Welcome, ${input.name.get()}!"))
 }
-
-object GreetTask {
-
-  /** Binds input data to this task
-    *
-    * @param name
-    *   the input name
-    * @return
-    *   a transformed instance of this class with input data
-    */
-  def apply(name: SdkBindingData[String]): SdkTransform[GreetTaskOutput] =
-    new GreetTask().withInput("name", name)
-}
diff --git a/flytekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/SumTask.scala b/flytekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/SumTask.scala
@@ -43,11 +43,3 @@ class SumTask
 
   override def isCacheSerializable: Boolean = true
 }
-
-object SumTask {
-  def apply(
-      a: SdkBindingData[Long],
-      b: SdkBindingData[Long]
-  ): SdkTransform[SumTaskOutput] =
-    new SumTask().withInput("a", a).withInput("b", b)
-}
diff --git a/...ekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/WelcomeWorkflow.scala b/...ekit-examples-scala/src/main/scala/org/flyte/examples/flytekitscala/WelcomeWorkflow.scala
@@ -16,7 +16,7 @@
  */
 package org.flyte.examples.flytekitscala
 
-import org.flyte.flytekit.{SdkBindingData, SdkWorkflow, SdkWorkflowBuilder}
+import org.flyte.flytekit.SdkBindingData
 import org.flyte.flytekitscala.{
   SdkScalaType,
   SdkScalaWorkflow,
@@ -49,10 +49,12 @@ import org.flyte.flytekitscala.{
   * | output: greeting(string) |
   */
 
+case class WelcomeWorkflowInput(name: SdkBindingData[String])
 case class WelcomeWorkflowOutput(greeting: SdkBindingData[String])
 
 class WelcomeWorkflow
-    extends SdkScalaWorkflow[WelcomeWorkflowOutput](
+    extends SdkScalaWorkflow[WelcomeWorkflowInput, WelcomeWorkflowOutput](
+      SdkScalaType[WelcomeWorkflowInput],
       SdkScalaType[WelcomeWorkflowOutput]
     ) {
 
@@ -61,11 +63,18 @@ class WelcomeWorkflow
     val name = builder.inputOfString("name", "The name for the welcome message")
 
     // uses the workflow input as the task input of the GreetTask
-    val greeting = builder.apply("greet", GreetTask(name)).getOutputs.greeting
+    val greeting = builder
+      .apply("greet", new GreetTask(), GreetTaskInput(name))
+      .getOutputs
+      .greeting
 
     // uses the output of the GreetTask as the task input of the AddQuestionTask
     val greetingWithQuestion = builder
-      .apply("add-question", AddQuestionTask(greeting))
+      .apply(
+        "add-question",
+        new AddQuestionTask(),
+        AddQuestionTaskInput(greeting)
+      )
       .getOutputs
       .greeting
 

diff --git a/flytekit-examples/src/main/java/org/flyte/examples/SimpleStructTask.java b/flytekit-examples/src/main/java/org/flyte/examples/SimpleStructTask.java
@@ -18,9 +18,7 @@
 
 import com.google.auto.value.AutoValue;
 import org.flyte.api.v1.Struct;
-import org.flyte.flytekit.SdkBindingData;
 import org.flyte.flytekit.SdkRunnableTask;
-import org.flyte.flytekit.SdkTransform;
 import org.flyte.flytekit.jackson.JacksonSdkType;
 
 /** Example Flyte task that takes a name as the input and outputs a simple greeting message. */
@@ -31,10 +29,6 @@ public SimpleStructTask() {
     super(JacksonSdkType.of(Input.class), JacksonSdkType.of(Output.class));
   }
 
-  public static SdkTransform<?> of(SdkBindingData<?> struct) {
-    return new SimpleStructTask().withInput("struct", struct);
-  }
-
   @AutoValue
   public abstract static class Input {
     public abstract Struct in();