Spark Scala API PoC

main.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StaticPointers #-}
{-# LANGUAGE TypeOperators #-}

module Main where

import Control.Monad (forM_)
import Control.Distributed.Closure
import Control.Distributed.Spark (getOrCreateSparkContext, newSparkConf, parallelize)
import Data.Int
import qualified Foreign.JNI as JNI
import Language.Java
import Language.Java.Inline
import Criterion.Main as Criterion
import System.IO.Unsafe (unsafePerformIO)

sincr :: J ('Iface "scala.Function1")
sincr = unsafePerformIO $
      [java| new scala.runtime.AbstractFunction1<Integer, Integer>() {
           public Integer apply(Integer x) {
               return x + 1;
           }
       } |]

zipAndCountScala :: RDD Int32 -> IO Int32
zipAndCountScala rdd = do
    klass <- JNI.findClass "java/lang/Integer"
    tag :: J ('Class "scala.reflect.ClassTag") <- [java| scala.reflect.ClassTag$.MODULE$.apply($klass) |]
    let rdd1 = rdd
    reify =<< [java| { return $rdd.zip($rdd1, $tag).count(); } |]

newtype RDD a = RDD (J ('Class "org.apache.spark.rdd.RDD"))
instance Coercible (RDD a) ('Class "org.apache.spark.rdd.RDD")

main :: IO ()
main = do
    conf <- newSparkConf "RDD benchmarks"
    sc   <- getOrCreateSparkContext conf
    forM_ [0,200..10000] $ \x -> do
      putStrLn $ "Size " ++ show x
      rdd <- parallelize sc (replicate x (1 :: Int32))
      rdd' <-
        callStatic
          (sing :: Sing "org.apache.spark.api.java.JavaRDD")
          "toRDD"
          [jvalue rdd]
      print =<< zipAndCountScala rdd'