Some fun with Apache Wayang and Spark / Tensorflow

Apache Wayang is a cross-platform data processing system. You write a pipeline once and Wayang's optimizer decides where each step runs, on Java streams, Spark, Flink, Postgres or other engines it has a plugin for. It can also split one job across several engines and move the data between them. Wayang started as a research project, went through the Apache Incubator and is now an Apache top-level project. I was on the first podling team and I'm on the PMC.

When I first wrote this post in 2022 the code was a quick sketch with calls like WayangClient.addWorkers. That API never existed, and Wayang doesn't add or remove Spark workers anyway, your cluster manager does that. I replaced the sketches with real code against the current Wayang API, checked against the Wayang repository. All three examples are the same word count, so you can see how little changes between languages. Register Java and Spark, and the optimizer decides which engine runs which step.

Wayang with Scala and Spark

import org.apache.wayang.api._
import org.apache.wayang.core.api.{Configuration, WayangContext}
import org.apache.wayang.java.Java
import org.apache.wayang.spark.Spark

object WordCount {
  def main(args: Array[String]): Unit = {
    // Register the engines you have, Wayang picks per step
    val wayang = new WayangContext(new Configuration)
      .withPlugin(Java.basicPlugin())
      .withPlugin(Spark.basicPlugin())

    val counts = new PlanBuilder(wayang)
      .withJobName("WordCount")
      .withUdfJarsOf(this.getClass)
      .readTextFile("file:///path/to/input.txt")
      .flatMap(_.split("\\W+"))
      .filter(_.nonEmpty)
      .map(word => (word.toLowerCase, 1))
      .reduceByKey(_._1, (a, b) => (a._1, a._2 + b._2))
      .collect()

    counts.foreach(println)
  }
}

PlanBuilder is the Scala entry point. Drop Spark.basicPlugin() and the same plan runs only on Java streams, add Flink.basicPlugin() and Flink becomes an option too. The pipeline code stays the same.

Wayang with Python and Spark

from pywy.dataquanta import WayangContext
from pywy.platforms.java import JavaPlugin
from pywy.platforms.spark import SparkPlugin

def word_count():
    WayangContext() \
        .register({JavaPlugin, SparkPlugin}) \
        .textfile("file:///path/to/input.txt") \
        .flatmap(lambda line: line.split(), str, str) \
        .filter(lambda word: word.strip() != "", str) \
        .map(lambda word: (word.lower(), 1), str, (str, int)) \
        .reduce_by_key(lambda t: t[0], lambda a, b: (a[0], int(a[1]) + int(b[1])), (str, int)) \
        .store_textfile("file:///tmp/wordcount-out.txt", (str, int))

if __name__ == "__main__":
    word_count()

The Python API is called pywy. It sends the plan as JSON to the Wayang REST API, so that service has to run first. The python folder in the repo shows how to build pywy and start the REST API.

Wayang with Java and Spark

import org.apache.wayang.api.JavaPlanBuilder;
import org.apache.wayang.basic.data.Tuple2;
import org.apache.wayang.core.api.Configuration;
import org.apache.wayang.core.api.WayangContext;
import org.apache.wayang.java.Java;
import org.apache.wayang.spark.Spark;

import java.util.Arrays;
import java.util.Collection;

public class WordCount {
    public static void main(String[] args) {
        WayangContext wayang = new WayangContext(new Configuration())
                .withPlugin(Java.basicPlugin())
                .withPlugin(Spark.basicPlugin());

        Collection<Tuple2<String, Integer>> counts = new JavaPlanBuilder(wayang)
                .withJobName("WordCount")
                .withUdfJarOf(WordCount.class)
                .readTextFile("file:///path/to/input.txt")
                .flatMap(line -> Arrays.asList(line.split("\\W+")))
                .filter(word -> !word.isEmpty())
                .map(word -> new Tuple2<>(word.toLowerCase(), 1))
                .reduceByKey(Tuple2::getField0,
                        (a, b) -> new Tuple2<>(a.getField0(), a.getField1() + b.getField1()))
                .collect();

        counts.forEach(System.out::println);
    }
}

JavaPlanBuilder gives Java the same fluent API as Scala. collect() runs the plan and brings the result back to the driver.

K-Means with Wayang and TensorFlow

The first version of this section passed K-Means output from a WayangClient.runKMeans call into a TensorFlow session. That call never existed, it was an experimental sketch of an idea, so I took it out. Here is what is real. K-Means in Wayang is a normal plan with a loop built from repeat. This is shortened from the K-Means app in the Wayang repo, where you also find Point and SelectNearestCentroid.

val points = planBuilder
  .readTextFile(inputFile)
  .map { line =>
    val fields = line.split(",")
    Point(fields(0).toDouble, fields(1).toDouble)
  }

val initialCentroids = planBuilder
  .loadCollection(Kmeans.createRandomCentroids(k))

val finalCentroids = initialCentroids.repeat(iterations, { currentCentroids =>
  points
    .mapJava(new SelectNearestCentroid)
    .withBroadcast(currentCentroids, "centroids")
    .reduceByKey(_.centroidId, _ + _)
    .map(_.average)
})

finalCentroids.map(_.toPoint).collect()

Each round broadcasts the current centroids, assigns every point to its nearest centroid and averages the points per centroid. Wayang decides if that runs on Java streams or Spark.

For TensorFlow, Wayang now ships a TensorFlow platform. Register it with Tensorflow.plugin() and the API gets dlTraining and predict operators, so model training can sit in the same plan as the data preparation. I don't have a short example for that I'd stand behind, so look at the tests in wayang-tensorflow to get started.

If you want to get involved, the project lives at wayang.apache.org. Contributions are welcome, code, docs and bug reports alike.