How to count feature exists amount within a SparseVector

Viewed 452

I have a dataframe like this:

+---+-------+--------------+
| id|clicked|     features |
+---+-------+--------------+
|  7|  1.0  |(4,[3],[1.0]) |
|  8|  0.0  |(4,[2],[12.0])|
|  9|  0.0  |(4,[3],[0.1]) |
+---+-------+--------------+

and have converted it to a RDD:

Row(id=7, clicked=1.0, features=SparseVector(4, {3: 1.0}))
Row(id=8, clicked=0.0, features=SparseVector(4, {2: 12.0}))
Row(id=9, clicked=0.0, features=SparseVector(4, {3: 0.1}))

Now I want to find the number of exists for each feature. For example, in my dataframe/rdd, features column contains 5 features (0 to 4). Since feature in index 0, 1, and 4 do not contain any values, their amounts are 0. And feature index 2 is 1, feature index 3 is 2.

I want get that information within a dictionary data structure. How to do it?

{0:0, 1:0, 2:1, 3:2, 4:0}

I'm using PySpark, but an answer using Scala is also fine.

1 Answers

The best thing would be to do the counting before the SparseVector was created. If that is not an possible, you essentially have two options (until VectorUDTs are easily castable into arrays).

In both cases, the way to count the number of values that exist per feature is the same. Loop over the range of the size of the SparseVector and check if that index exists in the SparseVector.indices list. This will return counts for all features, including those where the count is 0.

A simpler way could be to create tuples of the form (index, 1) for each index in SparseVector.indices, but this would exclude from the final output any features without any values.

Option 1: Define a udf, explode, and aggregate:

import pyspark.sql.functions as f

featureCount_udf = f.udf(
    lambda r: [(x, int(x in r.indices)) for x in range(r.size)],
    ArrayType(
        StructType(
            [
                StructField("featureNumber", IntegerType()),
                StructField("count", IntegerType())
            ]
        )
    )
)

df.select(f.explode(featureCount_udf("features")).alias("features"))\
    .select("features.*")\
    .groupBy("featureNumber")\
    .agg(f.sum("count").alias("count"))\
    .show()
#+-------------+-----+
#|featureNumber|count|
#+-------------+-----+
#|            0|    0|
#|            2|    1|
#|            1|    0|
#|            3|    2|
#+-------------+-----+

Option 2: Convert to rdd and flatMap:

from operator import add

df.select("features")\
    .rdd\
    .flatMap(
        lambda r: [(x, int(x in r["features"].indices)) for x in range(r["features"].size)]
    )\
    .reduceByKey(add)\
    .toDF(["featureNumber", "count"])\
    .show()
#+-------------+-----+
#|featureNumber|count|
#+-------------+-----+
#|            0|    0|
#|            2|    1|
#|            1|    0|
#|            3|    2|
#+-------------+-----+

Here we flatMap each row to tuples of the form (featureNumber, containsValue). Then we can call reduceByKey to add the indicator variable for each feature.


Original Answer

If you want the output in a dictionary, you're going to have to call collect() at some point.

data = df.select("features").collect()

Now that you have the data as list of pyspark.sql.Rows, you can iterate through and use .indices and .size to identify which columns have values.

print([[int(x in r["features"].indices) for x in range(r["features"].size)] for r in data])
#[[0, 0, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]]

From this you can create a numpy array and sum the columns. Finally call enumerate on the result and pass that to the dict constructor.

Putting it all together:

mydict = dict(
    enumerate(
        np.array(
            [[int(x in r["features"].indices) for x in range(r["features"].size)]
             for r in data]
        ).sum(0)
    )
)
print(mydict)
#{0: 0, 1: 0, 2: 1, 3: 2}
Related