Accelerate Spacy Tokenize with Spark 2.3's new Apache Arrow powered UDFS to compliment NLTK

Author: holdenk

sparklingml/feature/python_pipelines.py

@@ -125,7 +125,7 @@ def _transform(self, dataset):
         SpacyTokenize.setup(dataset._sc, dataset.sql_ctx, self.getLang())
         func = SpacyTokenize.func(self.getLang())
         ret_type = SpacyTokenize.returnType()
-        udf = UserDefinedFunction(func, ret_type)
+        udf = pandas_udf(func, ret_type)
         return dataset.withColumn(
             self.getOutputCol(), udf(self.getInputCol())
         )

sparklingml/transformation_functions.py

@@ -88,15 +88,16 @@ def get(cls, lang):
         return cls._spacys[lang]
 
 
-class SpacyTokenize(TransformationFunction):
+class SpacyTokenize(ScalarVectorizedTransformationFunction):
     """
     Tokenize input text using spacy.
     >>> spt = SpacyTokenize()
     >>> sp = spt.func("en")
-    >>> r = sp("hi boo")
+    >>> r = sp(pandas.Series(["hi boo"]))
     ...
     >>> r
-    [u'hi', u'boo']
+    0    [hi, boo]
+    dtype: object
     """
     @classmethod
     def setup(cls, sc, session, *args):
@@ -106,11 +107,15 @@ def setup(cls, sc, session, *args):
     def func(cls, *args):
         lang = args[0]
 
-        def inner(inputString):
+        def inner(inputSeries):
             """Tokenize the inputString using spacy for
             the provided language."""
             nlp = SpacyMagic.get(lang)
-            return list(map(lambda x: x.text, list(nlp(inputString))))
+
+            def tokenizeElem(elem):
+                return list(map(lambda token: token.text, list(nlp(unicode(elem)))))
+
+            return inputSeries.apply(tokenizeElem)
         return inner
 
     @classmethod