pyspark: StopWordsRemover with user defined functions (UDF)

Question

I have a dataframe consisting of text and languages

sf = spark.createDataFrame([
    ('eng', "I saw the red balloon"),
    ('eng', 'She was drinking tea from a black mug'),
    ('ger','Er ging heute sehr weit'),
    ('ger','Ich habe dich seit hundert Jahren nicht mehr gesehen')
], ["lang", "text"])
display(sf)

Output:

+----+--------------------+
|lang|                text|
+----+--------------------+
| eng|I saw the red bal...|
| eng|She was drinking ...|
| ger|Er ging heute seh...|
| ger|Ich habe dich sei...|
+----+--------------------+

I want to remove the stop word for each text, for this I create a dictionary:

from pyspark.ml.feature import StopWordsRemover

ger_stopwords = StopWordsRemover.loadDefaultStopWords("german")
eng_stopwords = StopWordsRemover.loadDefaultStopWords("english")
stopwords = {'eng':eng_stopwords,
            'ger':ger_stopwords}

And now I don't understand how can I apply stop words to a col('text') using udf ? Because transform() will not suit me in this case

score 3 · Accepted Answer · answered Nov 02 '21 at 16:32

I do not know exactly how to use StopWordsRemover, but based on what you did and on the documentation, I can offer this solution (without UDFs):

from functools import reduce


df = reduce(
    lambda a, b: a.unionAll(b),
    (
        StopWordsRemover(
            inputCol="splitted_text", outputCol="words", stopWords=value
        ).transform(
            sf.where(F.col("lang") == key).withColumn(
                "splitted_text", F.split("text", " ")
            )
        )
        for key, value in stopwords.items()
    ),
)

df.show()
+----+----------------------------------------------------+--------------------------------------------------------------+--------------------------------------+
|lang|text                                                |splitted_text                                                 |words                                 |
+----+----------------------------------------------------+--------------------------------------------------------------+--------------------------------------+
|eng |I saw the red balloon                               |[I, saw, the, red, balloon]                                   |[saw, red, balloon]                   |
|eng |She was drinking tea from a black mug               |[She, was, drinking, tea, from, a, black, mug]                |[drinking, tea, black, mug]           |
|ger |Er ging heute sehr weit                             |[Er, ging, heute, sehr, weit]                                 |[ging, heute, weit]                   |
|ger |Ich habe dich seit hundert Jahren nicht mehr gesehen|[Ich, habe, dich, seit, hundert, Jahren, nicht, mehr, gesehen]|[seit, hundert, Jahren, mehr, gesehen]|
+----+----------------------------------------------------+--------------------------------------------------------------+--------------------------------------+

pyspark: StopWordsRemover with user defined functions (UDF)

1 Answers1