Create a Python transformer on sparsevector data type column in Pyspark ML

Viewed 1936

I have a dataframe with a column 'features'(each row in the dataframe represents a document). I used HashingTF to calculate the column 'tf' and I also created a custom transformer 'TermCount' (just as test) to calculate the 'total_terms' as follows:

from pyspark import SparkContext
from pyspark.sql import SQLContext,Row
from pyspark.ml.pipeline import Transformer
from pyspark.ml.param.shared import HasInputCol, HasOutputCol, Param
from pyspark.ml.feature import HashingTF
from pyspark.ml.util import keyword_only
from pyspark.mllib.linalg import SparseVector
from pyspark.sql.functions import udf

class TermCount(Transformer, HasInputCol, HasOutputCol):

    @keyword_only
    def __init__(self, inputCol=None, outputCol=None):
        super(TermCount, self).__init__()
        kwargs = self.__init__._input_kwargs
        self.setParams(**kwargs)

    @keyword_only
    def setParams(self, inputCol=None, outputCol=None):
        kwargs = self.setParams._input_kwargs
        return self._set(**kwargs)

    def _transform(self, dataset):

        def f(s):
            return len(s.values)

        out_col = self.getOutputCol()
        in_col = dataset[self.getInputCol()]
        return dataset.withColumn(out_col, udf(f)(in_col))

sc = SparkContext()
sqlContext = SQLContext(sc)
documents = sqlContext.createDataFrame([
    (0, "w1 w2 w3 w4 w1 w1 w1"),
    (1, "w2 w3 w4 w2"),
    (2, "w3 w4 w3"),
    (3, "w4")], ["doc_id", "doc_text"])

df = documents.map(lambda x : (x.doc_id,x.doc_text.split(" "))).toDF().withColumnRenamed("_1","doc_id").withColumnRenamed("_2","features")

htf = HashingTF(inputCol="features", outputCol="tf")
tf = htf.transform(df)

term_count_model=TermCount(inputCol="tf", outputCol="total_terms")
tc_df=term_count_model.transform(tf)
tc_df.show(truncate=False)
#+------+----------------------------+------------------------------------------------+-----------+
#|doc_id|features                    |tf                                              |total_terms|
#+------+----------------------------+------------------------------------------------+-----------+
#|0     |[w1, w2, w3, w4, w1, w1, w1]|(262144,[3738,3739,3740,3741],[4.0,1.0,1.0,1.0])|4          |
#|1     |[w2, w3, w4, w2]            |(262144,[3739,3740,3741],[2.0,1.0,1.0])         |3          |
#|2     |[w3, w4, w3]                |(262144,[3740,3741],[2.0,1.0])                  |2          |
#|3     |[w4]                        |(262144,[3741],[1.0])                           |1          |
#+------+----------------------------+------------------------------------------------+-----------+

Now, I need to add a similar transformer which receives 'tf' as an inputCol and compute the document frequency for each term (no_of_rows_contains_this_term / total_no_of_rows) to an outputCol of type Sparsevector and finally to get a result like this:

+------+----------------------------+------------------------------------------------+-----------+----------------------------------------------------+
|doc_id|features                    |tf                                              |total_terms| doc_freq                                           |
+------+----------------------------+------------------------------------------------+-----------+----------------------------------------------------+
|0     |[w1, w2, w3, w4, w1, w1, w1]|(262144,[3738,3739,3740,3741],[4.0,1.0,1.0,1.0])|4          |(262144,[3738,3739,3740,3741],[0.25,0.50,0.75,1.0]) |
|1     |[w2, w3, w4, w2]            |(262144,[3739,3740,3741],[2.0,1.0,1.0])         |3          |(262144,[3739,3740,3741],[0.50,0.75,1.0])           |
|2     |[w3, w4, w3]                |(262144,[3740,3741],[2.0,1.0])                  |2          |(262144,[3740,3741],[0.75,1.0])                     |
|3     |[w4]                        |(262144,[3741],[1.0])                           |1          |(262144,[3741],[1.0])                               |
+------+----------------------------+------------------------------------------------+-----------+----------------------------------------------------+
1 Answers
Related