2019年2月23日星期六

Customised Transformer in Spark mllib with Java

In retrospect, it may not be sensible for me to use Spark, or it could be easier to start building the application from scratch. If I did that, I might choose to write the application in Scala. With that said, I just want to note down the struggles I’ve been through in recent days, changing the underlying machine learning library of our existing application to Spark mllib, with little knowledge of Scala, Spark or machine learning. Despite the popularity of Spark, I found it very difficult to find examples or resources in Java.

First of all, version issue, our old application use Java 11, and Spark (mostly because of Scala) couldn’t work with it. All though Scala document says that 2.12 and 2.11 are compatible with Java 11, there’s also a disclaimer below says that the support of JDK 9+ is incomplete. As a result, the code compiles but failed at runtime with strange errors. In Java application, you should also take care that Scala version is correct, with Maven there is a maven-scala-plugin for that. Spark also brings a dependency on an old version of Guava, which may cause conflicts with other libraries.

Second, Spark requires that all the relevant code to be serializable. For some it is easy, others not so much. A lot of times I circumvent the problem by using static field. A good example is database connection, where in the official document it is recommended to either create connection for each partition, or use a connection pool. Just be careful that the shared code should be thread safe.

Third, and the biggest problem is implementing customised transformer. The source code is in Scala, and made massive use of trait, which becomes just an interface in Java. Sometimes Spark provides a Java compatible class, like in the case of Params, you should extend JavaParams instead. Mllib uses copy method to do hyper-parameter tuning (which I think is the main reason to implement Transformer rather than just write a udf in the first place), and does provide a defaultCopy method. But internally it use reflection a lot, and that makes it very difficult to use correctly. I only found at runtime that I need a constructor which requires a string argument, also I need to implement a getter method for each param. Similar error can also happen when using Params.set or Params.get method. The transformer should also implement MLWritable to be able save the model, MLWritable is also a trait in Scala, luckily the DefaultParamsWriter works out of the box. MLReadable is a bit different, the DefaultParamsReader somehow use reflection again and call the method on null. So we need to make Transformer.read method static to avoid NPE. The transformSchema method is important for a pipeline to know upfront that the schema of the dataset is valid.

Some minor things, although Param is declared to be generic, just write Param<MyOwnObject> will not work when saving the model. You need to extend Param class and override jsonEncode and jsonDecode method (as I can see, Mllib somehow uses Java serialization, json format, parquet format altogether). There is no single option to handle missing values, although Mllib provides a Inputer, you need to know the internal of the algorithm to apply the strategies (which seems to be a reasonable one). Dealing with unbalanced dataset (through a weight column) is only supported by a small number of algorithms. For now, I can only subsample my dataset to get a balanced set.
Finally, I do like the interface of PipelineStage, Transformer, the built-in cross validation and hyper-parameter tuning works out-of-the-box. A large part of the saved model is directly human readable, with some hazard of parquet files.

Code example (modified to a dummy example, not tested)