Creating Multi-language Pipelines with Apache Spark or Avoid Having to Rewrite spaCy into Java

# This class is used to allow the Scala process to call into Python
# It may not run in the same Python process as your regular Python
# shell if you are running PySpark normally.
class PythonRegistrationProvider(object):
  """Provide an entry point for Scala to call to register functions."""
 
  def __init__(self, gateway):
    self.gateway = gateway
    self._sc = None
    self._session = None
    self._count = 0
 
  def registerFunction(self, ssc, jsession, function_name, params):
    jvm = self.gateway.jvm
    # If we don't have a reference to a running SparkContext
    # Get the SparkContext from the provided SparkSession.
    if not self._sc:
        master = ssc.master()
        jsc = jvm.org.apache.spark.api.java.JavaSparkContext(ssc)
        jsparkConf = ssc.conf()
        sparkConf = SparkConf(_jconf=jsparkConf)
        self._sc = SparkContext(
           master=master,
           conf=sparkConf,
           gateway=self.gateway,
           jsc=jsc)
        self._session = SparkSession.builder.getOrCreate()
    if function_name in functions_info:
        function_info = functions_info[function_name]
        if params:
            evaledParams = ast.literal_eval(params)
        else:
          evaledParams = []
        func = function_info.func(*evaledParams)
        ret_type = function_info.returnType()
        self._count = self._count + 1
        registration_name = function_name + str(self._count)
        udf = UserDefinedFunction(func, ret_type, registration_name)
        # Configure non-standard evaluation types (e.g. Arrow)
        udf.evalType = function_info.evalType()
        judf = udf._judf
        return judf
    else:
       return None
 
    class Java:
       package = "com.sparklingpandas.sparklingml.util.python"
       className = "PythonRegisterationProvider"
       implements = [package + "." + className]

/*** Abstract trait to implement in Python to allow Scala to call in to perform* registration.*/trait PythonRegisterationProvider {// Takes a SparkContext, SparkSession, String, and String// Returns UserDefinedPythonFunction but types + py4j :(def registerFunction(sc: SparkContext, session: Object,functionName: Object, params: Object): Object}

 // Call the registration provider from startup.py to get a Python UDF back.   val pythonUdf = Option(registrationProvider.registerFunction(     session.sparkContext,     session,     pythonFunctionName,     miniSerializeParams()))   val castUdf = pythonUdf.map(_.asInstanceOf[UserDefinedPythonFunction])

final val lang = new Param[String](this, "lang", "language for tokenization")/** @group getParam */final def getLang: String = $(lang)final def setLang(value: String): this.type = set(this.lang, value)def this() = this(Identifiable.randomUID("SpacyTokenizePython"))override val pythonFunctionName = "spaCytokenize"override protected def outputDataType = ArrayType(StringType)override protected def validateInputType(inputType: DataType): Unit = {if (inputType != StringType) {throw new IllegalArgumentException(s"Expected input type StringType instead found ${inputType}")}}override def copy(extra: ParamMap) = {defaultCopy(extra)}def miniSerializeParams() = {"[\"" + $(lang) + "\"]"}}

class SpacyTokenize(ScalarVectorizedTransformationFunction):
    @classmethod
    def setup(cls, sc, session, *args):
        pass
 
    @classmethod
    def func(cls, *args):
        lang = args[0]
 
        def inner(inputSeries):
            """Tokenize the inputString using spaCy for the provided language."""
            nlp = SpacyMagic.get(lang) # Optimization for spacy.load
 
            def tokenizeElem(elem):
                result_itr =  [token.text for token in nlp(elem)]
                return list(result_itr)
 
            return inputSeries.apply(tokenizeElem)
        return inner
 
    @classmethod
    def returnType(cls, *args):
        return ArrayType(StringType())

functions_info["spacytokenize"] = SpacyTokenize

# Spacy isn't serializable but loading it is semi-expensive
@ignore_unicode_prefix
class SpacyMagic(object):
    """Simple Spacy Magic to minimize loading time.
    >>> SpacyMagic.get("en")
    <spacy.lang.en.English ..."""
    _spacys = {}
    
    @classmethod
    def get(cls, lang):
        if lang not in cls._spacys:
            import spacy
            try:
                try:
                    cls._spacys[lang] = spacy.load(lang)
                except Exception:
                    spacy.cli.download(lang)
                    cls._spacys[lang] = spacy.load(lang)
            except Exception as e:
                raise Exception("Failed to find or download language {0}: {1}".format(lang, e))
    return cls._spacys[lang]

val data = session.load(...)val transformer = new SpacyTokenizePython()transformer.setLang("en")transformer.setInputCol("input")transformer.setOutputCol("tokens")val tokens = transformer.transform(data)val counts = tokens.groupBy("tokens").count()counts.write.format("json").save("...")