Pyspark higher-order SQL functions to create histograms from arrays? - pyspark

Question

I have a dataframe with a column of type array<sometype>.
More specifically, sometype = array<int> but I suspect this is irrelevant for the challenge at hand.
I'd like to assign to each unique element in the array the number of times it occurs there by using only higher-order functions and in linear time. The resulting column could be a map or an array of structs - doesn't matter.
For instance:
["val1", "val2", "val1", "val1", "val3", "val2", "val1"] --> { "val1": 4, "val3": 1, "val2": 2 }
I tried aggregate with map_concat (adding a single-element-map with an incremented counter) but the latter turned out to produce a multimap rather than overwrite an existing element with a new value, which foiled the plan.
Any other suggestion how this can be done?
Thanks

Answer 1

Try something like this.
df.show() #sample dataframe
#+---------------+
#| array|
#+---------------+
#| [1, 9, 1]|
#|[2, 2, 2, 1, 2]|
#|[3, 4, 4, 1, 4]|
#| [1, 4]|
#| [99, 99, 100]|
#| [92, 11, 92]|
#| [0, 0, 1]|
#+---------------+
Transform with filter:
from pyspark.sql import functions as F
df\
.withColumn("count",\
F.expr("""map_from_arrays(array_distinct(array),transform(array_distinct(array),\
x-> size(filter(array,y-> y=x))))"""))\
.show(truncate=False)
#+---------------+------------------------+
#|array |count |
#+---------------+------------------------+
#|[1, 9, 1] |[1 -> 2, 9 -> 1] |
#|[2, 2, 2, 1, 2]|[2 -> 4, 1 -> 1] |
#|[3, 4, 4, 1, 4]|[3 -> 1, 4 -> 3, 1 -> 1]|
#|[1, 4] |[1 -> 1, 4 -> 1] |
#|[99, 99, 100] |[99 -> 2, 100 -> 1] |
#|[92, 11, 92] |[92 -> 2, 11 -> 1] |
#|[0, 0, 1] |[0 -> 2, 1 -> 1] |
#+---------------+------------------------+
Or Transform with aggregate:
from pyspark.sql import functions as F
df\
.withColumn("count",\
F.expr("""map_from_arrays(array_distinct(array),transform(array_distinct(array),\
x-> aggregate(array, 0,(acc,t)->acc+IF(t=x,1,0))))"""))\
.show(truncate=False)
#+---------------+------------------------+
#|array |count |
#+---------------+------------------------+
#|[1, 9, 1] |[1 -> 2, 9 -> 1] |
#|[2, 2, 2, 1, 2]|[2 -> 4, 1 -> 1] |
#|[3, 4, 4, 1, 4]|[3 -> 1, 4 -> 3, 1 -> 1]|
#|[1, 4] |[1 -> 1, 4 -> 1] |
#|[99, 99, 100] |[99 -> 2, 100 -> 1] |
#|[92, 11, 92] |[92 -> 2, 11 -> 1] |
#|[0, 0, 1] |[0 -> 2, 1 -> 1] |
#+---------------+------------------------+
UPDATE:
A different way to do this could be if we know exactly all the elements we will be counting for. This way we can parallelize better as a struct is like a column and a struct of structs is basically a dataframe inside another.
elements=[1,9,2,3,4,99,100,92,11,0]
df.show() #sample dataframe
#+---------------+
#| array|
#+---------------+
#| [1, 9, 1]|
#|[2, 2, 2, 1, 2]|
#|[3, 4, 4, 1, 4]|
#| [1, 4]|
#| [99, 99, 100]|
#| [92, 11, 92]|
#| [0, 0, 1]|
#+---------------+
from pyspark.sql import functions as F
df.withColumn("struct", F.struct(*[(F.struct(F.expr("size(filter(array,x->x={}))"\
.format(y))).alias(str(y))) for y in elements]))\
.select("array",F.map_from_arrays(F.array(*[F.lit(x) for x in elements]),\
F.array(*[(F.col("struct.{}.col1".format(x)))\
for x in elements])).alias("count")).show(truncate=False)
#+---------------+-------------------------------------------------------------------------------------+
#|array |count |
#+---------------+-------------------------------------------------------------------------------------+
#|[1, 9, 1] |[1 -> 2, 9 -> 1, 2 -> 0, 3 -> 0, 4 -> 0, 99 -> 0, 100 -> 0, 92 -> 0, 11 -> 0, 0 -> 0]|
#|[2, 2, 2, 1, 2]|[1 -> 1, 9 -> 0, 2 -> 4, 3 -> 0, 4 -> 0, 99 -> 0, 100 -> 0, 92 -> 0, 11 -> 0, 0 -> 0]|
#|[3, 4, 4, 1, 4]|[1 -> 1, 9 -> 0, 2 -> 0, 3 -> 1, 4 -> 3, 99 -> 0, 100 -> 0, 92 -> 0, 11 -> 0, 0 -> 0]|
#|[1, 4] |[1 -> 1, 9 -> 0, 2 -> 0, 3 -> 0, 4 -> 1, 99 -> 0, 100 -> 0, 92 -> 0, 11 -> 0, 0 -> 0]|
#|[99, 99, 100] |[1 -> 0, 9 -> 0, 2 -> 0, 3 -> 0, 4 -> 0, 99 -> 2, 100 -> 1, 92 -> 0, 11 -> 0, 0 -> 0]|
#|[92, 11, 92] |[1 -> 0, 9 -> 0, 2 -> 0, 3 -> 0, 4 -> 0, 99 -> 0, 100 -> 0, 92 -> 2, 11 -> 1, 0 -> 0]|
#|[0, 0, 1] |[1 -> 1, 9 -> 0, 2 -> 0, 3 -> 0, 4 -> 0, 99 -> 0, 100 -> 0, 92 -> 0, 11 -> 0, 0 -> 2]|
#+---------------+-------------------------------------------------------------------------------------+
You could also try this(to get distinct count per row) using structs:
elements=[1,9,2,3,4,99,100,92,11,0]
from pyspark.sql import functions as F
collected=df.withColumn("struct", F.struct(*[(F.struct(F.expr("size(filter(array,x->x={}))"\
.format(y))).alias(str(y))) for y in elements]))\
.withColumn("vals", F.array(*[(F.col("struct.{}.col1".format(x))) for x in elements]))\
.select("array",F.arrays_zip(F.array(*[F.lit(x) for x in elements]),\
F.col("vals")).alias("count"))\
.withColumn("count", F.expr("""filter(count,x-> x.vals != 0)"""))\
.show(truncate=False)
#+---------------+------------------------+
#|array |count |
#+---------------+------------------------+
#|[1, 9, 1] |[[1, 2], [9, 1]] |
#|[2, 2, 2, 1, 2]|[[1, 1], [2, 4]] |
#|[3, 4, 4, 1, 4]|[[1, 1], [3, 1], [4, 3]]|
#|[1, 4] |[[1, 1], [4, 1]] |
#|[99, 99, 100] |[[99, 2], [100, 1]] |
#|[92, 11, 92] |[[92, 2], [11, 1]] |
#|[0, 0, 1] |[[1, 1], [0, 2]] |
#+---------------+------------------------+
Or, you can use map_from_entries with structs logic:
elements=[1,9,2,3,4,99,100,92,11,0]
from pyspark.sql import functions as F
collected=df.withColumn("struct", F.struct(*[(F.struct(F.expr("size(filter(array,x->x={}))"\
.format(y))).alias(str(y))) for y in elements]))\
.withColumn("vals", F.array(*[(F.col("struct.{}.col1".format(x))) for x in elements]))\
.withColumn("elems", F.array(*[F.lit(x) for x in elements]))\
.withColumn("count", F.map_from_entries(F.expr("""filter(arrays_zip(elems,vals),x-> x.vals != 0)""")))\
.select("array","count")\
.show(truncate=False)
#+---------------+------------------------+
#|array |count |
#+---------------+------------------------+
#|[1, 9, 1] |[1 -> 2, 9 -> 1] |
#|[2, 2, 2, 1, 2]|[1 -> 1, 2 -> 4] |
#|[3, 4, 4, 1, 4]|[1 -> 1, 3 -> 1, 4 -> 3]|
#|[1, 4] |[1 -> 1, 4 -> 1] |
#|[99, 99, 100] |[99 -> 2, 100 -> 1] |
#|[92, 11, 92] |[92 -> 2, 11 -> 1] |
#|[0, 0, 1] |[1 -> 1, 0 -> 2] |
#+---------------+------------------------+

Answer 2

While still hoping to see suggestions for linear time solutions, I'll post here a pretty ugly version that I'm currently using to address the problem in N*log(N):
decide whether there's a need for non-trivial aggregation
for the cases where aggregation needs to happen:
2.1. sort the list of elements (this is where the most time will be spent)
2.2. assign elements that are about to be repeated special null value, while also recording their positions in the sorted list
2.3. remove elements entries with nulls and shrink the remaining lists
2.4. from differences in the position values of the remaining elements derive how many repetitions for each n-gram there had been
h_df = in_df.withColumn('__isunique__', func.expr("size(mylist)==size(array_distinct(mylist))")) \
.withColumn('__sorted__', func.when(func.col('__isunique__'), func.col('mylist')) \
.otherwise(func.expr("array_sort(mylist)"))) \
.withColumn('__coded__', func.when(func.col('__isunique__'), func.lit(None)) \
.otherwise(func.expr("filter(transform(sequence(0, size(__sorted__)), \
n->named_struct('p', n, 'e', \
IF ((n>0 and __sorted__[n]!=__sorted__[n-1]) or n==size(__sorted__), \
__sorted__[n-1], \
null))), \
x -> isnotnull(x['e']))"))) \
.withColumn('__clen__', func.when(func.col('__isunique__'), 0).otherwise(func.size('__coded__'))) \
.withColumn('histo', func.when(func.col('__isunique__'), func.expr("map_from_arrays(__sorted__,array_repeat(1,size(__sorted__)))")) \
.otherwise(func.expr("map_from_entries(transform(sequence(0, __clen__-1), \
p -> struct(__coded__[p]['e'], \
IF (p==0, \
__coded__[p]['p'], \
__coded__[p]['p']-__coded__[p-1]['p']))))"))) \
.drop('__sorted__','__coded__','__clen__','__isunique__')
h_df.show(10,False)
==>
+------------------+
| mylist|
+------------------+
|[B, A, A, C, B, A]|
| [C, D, E]|
| [B]|
| [C, C, C]|
+------------------+
+------------------+------------------------+
|mylist |histo |
+------------------+------------------------+
|[B, A, A, C, B, A]|[A -> 3, B -> 2, C -> 1]|
|[C, D, E] |[C -> 1, D -> 1, E -> 1]|
|[B] |[B -> 1] |
|[C, C, C] |[C -> 3] |
+------------------+------------------------+