How do I get Pyspark to aggregate sets at two levels? - apache-spark

I need to aggregate rows in a DataFrame by collecting the values in a certain column in each group into a set. pyspark.sql.functions.collect_set does exactly what I need.
However, I need to do this for two columns in turn, because I need to group the input by one column, divide each group into subgroups by another column, and do some aggregation on each subgroup. I don't see how to get collect_set to create a set for each group.
Example:
df = spark.createDataFrame([('a', 'x', 11, 22), ('a', 'y', 33, 44), ('b', 'x', 55, 66), ('b', 'y', 77, 88),('a','x',12,23),('a','y',34,45),('b','x',56,67),('b','y',78,89)], ('col1', 'col2', 'col3', 'col4'))
df.show()
+----+----+----+----+
|col1|col2|col3|col4|
+----+----+----+----+
| a| x| 11| 22|
| a| y| 33| 44|
| b| x| 55| 66|
| b| y| 77| 88|
| a| x| 12| 23|
| a| y| 34| 45|
| b| x| 56| 67|
| b| y| 78| 89|
+----+----+----+----+
g1 = df.groupBy('col1', 'col2').agg(collect_set('col3'),collect_set('col4'))
g1.show()
+----+----+-----------------+-----------------+
|col1|col2|collect_set(col3)|collect_set(col4)|
+----+----+-----------------+-----------------+
| a| x| [12, 11]| [22, 23]|
| b| y| [78, 77]| [88, 89]|
| a| y| [33, 34]| [45, 44]|
| b| x| [56, 55]| [66, 67]|
+----+----+-----------------+-----------------+
g2 = g1.groupBy('col1').agg(collect_set('collect_set(col3)'),collect_set('collect_set(col4)'),count('col2'))
g2.show(truncate=False)
+----+--------------------------------------------+--------------------------------------------+-----------+
|col1|collect_set(collect_set(col3)) |collect_set(collect_set(col4)) |count(col2)|
+----+--------------------------------------------+--------------------------------------------+-----------+
|b |[WrappedArray(56, 55), WrappedArray(78, 77)]|[WrappedArray(66, 67), WrappedArray(88, 89)]|2 |
|a |[WrappedArray(33, 34), WrappedArray(12, 11)]|[WrappedArray(22, 23), WrappedArray(45, 44)]|2 |
+----+--------------------------+--------------------------------------------+-----------+
I'd like the result to look more like
+----+----------------+----------------+-----------+
|col1| ...col3... | ...col4... |count(col2)|
+----+----------------+----------------+-----------+
|b |[56, 55, 78, 77]|[66, 67, 88, 89]|2 |
|a |[33, 34, 12, 11]|[22, 23, 45, 44]|2 |
+----+----------------+----------------+-----------+
but I don't see an aggregate function to take the union of two or more sets, or a pyspark operation to flatten the "array of arrays" structure that shows up in g2.
Does pyspark provide a simple way to accomplish this? Or is there a totally different approach I should be taking?

In PySpark 2.4.5, you could use the now built-in flatten function.

You can flatten the columns with a UDF afterwards:
flatten = udf(lambda l: [x for i in l for x in i], ArrayType(IntegerType()))
I took the liberty of renaming the columns of g2 as col3 and and col4 to save typing. This gives:
g3 = g2.withColumn('col3flat', flatten('col3'))
>>> g3.show()
+----+--------------------+--------------------+-----+----------------+
|col1| col3| col4|count| col3flat|
+----+--------------------+--------------------+-----+----------------+
| b|[[78, 77], [56, 55]]|[[66, 67], [88, 89]]| 2|[78, 77, 56, 55]|
| a|[[12, 11], [33, 34]]|[[22, 23], [45, 44]]| 2|[12, 11, 33, 34]|
+----+--------------------+--------------------+-----+----------------+

You can accomplish the same with
from pyspark.sql.functions import collect_set, countDistinct
(
df.
groupby('col1').
agg(
collect_set('col3').alias('col3_vals'),
collect_set('col4').alias('col4_vals'),
countDistinct('col2').alias('num_grps')
).
show(truncate=False)
)
+----+----------------+----------------+--------+
|col1|col3_vals |col4_vals |num_grps|
+----+----------------+----------------+--------+
|b |[78, 56, 55, 77]|[66, 88, 67, 89]|2 |
|a |[33, 12, 34, 11]|[45, 22, 44, 23]|2 |
+----+----------------+----------------+--------+

Related

Spark dataframe filter with boolean list comprehension

I want to filter a spark dataframe sdf based on several columns being not null.
Imagine I have:
labels = ["A", "B, C"]
This would work:
sdf.where(sf.col(labels[0]).isNotNull() | sf.col(labels[1]).isNotNull() | sf.col(labels[2]).isNotNull())
But I would like to do something similar to a list comprehension if the list was much longer:
sdf.where(any([sf.col(l).isNotNull() for l in labels]))
(this does not work, {ValueError}Cannot convert column into bool: please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions.)
How can I achieve this?
You can use reduce from functools to iterate over your list of columns and apply your logic.
In your case, it looks like you want to grab all the rows where any column has a non-null value (so full null value rows should get filtered away).
from functools import reduce
import pyspark.sql.functions as F
labels = ["A", "B", "C"]
df = spark.createDataFrame(
[
(None, 1, "ABC"),
(1, None, "BCD"),
(None, None, None),
(2, 2, None),
(1, 3, "DEF"),
(2, 1, "EFG"),
(None, None, None),
(2, 2, None),
(None, 3, "HIJ"),
(None, None, None),
(2, 2, None),
(3, 1, "EFG"),
(3, 2, None),
(None, None, None),
(2, 2, None),
(3, 3, "HIJ"),
],
["A", "B", "C"]
)
df.filter(reduce(lambda x, y: x | y, (F.col(x).isNotNull() for x in labels))).show()
+----+----+----+
| A| B| C|
+----+----+----+
|null| 1| ABC|
| 1|null| BCD|
| 2| 2|null|
| 1| 3| DEF|
| 2| 1| EFG|
| 2| 2|null|
|null| 3| HIJ|
| 2| 2|null|
| 3| 1| EFG|
| 3| 2|null|
| 2| 2|null|
| 3| 3| HIJ|
+----+----+----+
As you can see, the rows with all null values are correctly filtered away. This is done by OR-ing the isNotNull() conditions.

How to fill up null values in Spark Dataframe based on other columns' value?

Given this dataframe:
+-----+-----+----+
|num_a|num_b| sum|
+-----+-----+----+
| 1| 1| 2|
| 12| 15| 27|
| 56| 11|null|
| 79| 3| 82|
| 111| 114| 225|
+-----+-----+----+
How would you fill up Null values in sum column if the value can be gathered from other columns? In this example 56+11 would be the value.
I've tried df.fillna with an udf, but that doesn't seems to work, as it was just getting the column name not the actual value. I would want to compute the value just for the rows with missing values, so creating a new column would not be a viable option.
If your requirement is UDF, then it can be done as:
import pyspark.sql.functions as F
from pyspark.sql.types import LongType
df = spark.createDataFrame(
[(1, 2, 3),
(12, 15, 27),
(56, 11, None),
(79, 3, 82)],
["num_a", "num_b", "sum"]
)
F.udf(returnType=LongType)
def fill_with_sum(num_a, num_b, sum):
return sum if sum is None else (num_a + num_b)
df = df.withColumn("sum", fill_with_sum(F.col("num_a"), F.col("num_b"), F.col("sum")))
[Out]:
+-----+-----+---+
|num_a|num_b|sum|
+-----+-----+---+
| 1| 2| 3|
| 12| 15| 27|
| 56| 11| 67|
| 79| 3| 82|
+-----+-----+---+
You can use coalesce function. Check this sample code
import pyspark.sql.functions as f
df = spark.createDataFrame(
[(1, 2, 3),
(12, 15, 27),
(56, 11, None),
(79, 3, 82)],
["num_a", "num_b", "sum"]
)
df.withColumn("sum", f.coalesce(f.col("sum"), f.col("num_a") + f.col("num_b"))).show()
Output is:
+-----+-----+---+
|num_a|num_b|sum|
+-----+-----+---+
| 1| 2| 3|
| 12| 15| 27|
| 56| 11| 67|
| 79| 3| 82|
+-----+-----+---+

PySpark - create multiple aggregative map columns without using UDF or join

I have a huge dataframe that looks similar to this:
+----+-------+-------+-----+
|name|level_A|level_B|hours|
+----+-------+-------+-----+
| Bob| 10| 3| 5|
| Bob| 10| 3| 15|
| Bob| 20| 3| 25|
| Sue| 30| 3| 35|
| Sue| 30| 7| 45|
+----+-------+-------+-----+
My desired output:
+----+--------------------+------------------+
|name| map_level_A| map_level_B|
+----+--------------------+------------------+
| Bob|{10 -> 20, 20 -> 25}| {3 -> 45}|
| Sue| {30 -> 80}|{7 -> 45, 3 -> 35}|
+----+--------------------+------------------+
Meaning, group by name, adding 2 MapType columns that map level_A and level_B to the sum of hours.
I know I can get that output using an UDF or a join operation.
However, in practice, the data is very big, and it's not 2 map columns, but rather tens of them, so join/UDF are just too costly.
Is there a more efficient way to do that?
You could consider using Window functions. You'll need a windowspec for each level_X partitioned by both name and level_X to calculate the sum of hours. Then group by name and create map from array of structs:
from pyspark.sql import Window
import pyspark.sql.functions as F
df = spark.createDataFrame([("Bob", 10, 3, 5), ("Bob", 10, 3, 15), ("Bob", 20, 3, 25),
("Sue", 30, 3, 35),("Sue", 30, 7, 45), ],
["name", "level_A", "level_B", "hours"])
wla = Window.partitionBy("name", "level_A")
wlb = Window.partitionBy("name", "level_B")
result = df.withColumn("hours_A", F.sum("hours").over(wla)) \
.withColumn("hours_B", F.sum("hours").over(wlb)) \
.groupBy("name") \
.agg(
F.map_from_entries(
F.collect_set(F.struct(F.col("level_A"), F.col("hours_A")))
).alias("map_level_A"),
F.map_from_entries(
F.collect_set(F.struct(F.col("level_B"), F.col("hours_B")))
).alias("map_level_B")
)
result.show()
#+----+--------------------+------------------+
#|name| map_level_A| map_level_B|
#+----+--------------------+------------------+
#| Sue| {30 -> 80}|{3 -> 35, 7 -> 45}|
#| Bob|{10 -> 20, 20 -> 25}| {3 -> 45}|
#+----+--------------------+------------------+

Descriptive statistics on values outside of group for each group

I have a Spark DataFrame like this:
edited: each name can appear multiple times, in any org.
df = sqlContext.createDataFrame(
[
('org_1', 'a', 1),
('org_1', 'a', 2),
('org_1', 'a', 3),
('org_1', 'b', 4),
('org_1', 'c', 5),
('org_2', 'a', 7),
('org_2', 'd', 4),
('org_2', 'e', 5),
('org_2', 'e', 10)
],
["org", "name", "value"]
)
I would like to calculate for each org and name: the mean, stddev and count of values from the rest of the names excluding that name within each org. E.g. For org_1, name b, mean = (1+2+3+5)/4
The DataFrame has ~450 million rows. I cannot use vectorized pandas_UDF because my Spark version is 2.2. There is also a constraint of spark.driver.maxResultSize of 4.0 GB.
I tried this on Pandas (filter rows within groups and take mean/std/count) on a DataFrame with only two columns (name and value). I haven't figured out how to do this with two levels of grouped columns (org and name).
def stats_fun(x):
return pd.Series({'data_mean': x['value'].mean(),
'data_std': x['value'].std(),
'data_n': x['value'].count(),
'anti_grp_mean': df[df['name'] != x.name]['value'].mean(),
'anti_grp_std': df[df['name'] != x.name]['value'].std(),
'anti_grp_n': df[df['name'] != x.name]['value'].count()
})
df.groupby('name').apply(stats_fun)
Is there a similar UDF function I can define on Spark? (This function would have to take in multiple columns). Otherwise, what is a more efficient way to do this?
A simple UDF can also work.
import pyspark.sql.functions as F
import numpy as np
from pyspark.sql.types import *
df = sql.createDataFrame(
[
('org_1', 'a', 1),
('org_1', 'a', 2),
('org_1', 'a', 3),
('org_1', 'b', 4),
('org_1', 'c', 5),
('org_2', 'a', 7),
('org_2', 'd', 4),
('org_2', 'e', 5),
('org_2', 'e', 10)
],
["org", "name", "value"]
)
+-----+----+-----+
| org|name|value|
+-----+----+-----+
|org_1| a| 1|
|org_1| a| 2|
|org_1| a| 3|
|org_1| b| 4|
|org_1| c| 5|
|org_2| a| 7|
|org_2| d| 4|
|org_2| e| 5|
|org_2| e| 10|
+-----+----+-----+
After applying groupby and collecting all elements in list, we apply udf to find statistics. After that, columns are exploded and split into multiple columns.
def _find_stats(a,b):
dict_ = zip(a,b)
stats = []
for name in a:
to_cal = [v for k,v in dict_ if k != name]
stats.append((name,float(np.mean(to_cal))\
,float(np.std(to_cal))\
,len(to_cal)))
print stats
return stats
find_stats = F.udf(_find_stats,ArrayType(ArrayType(StringType())))
cols = ['name', 'mean', 'stddev', 'count']
splits = [F.udf(lambda val:val[0],StringType()),\
F.udf(lambda val:val[1],StringType()),\
F.udf(lambda val:val[2],StringType()),\
F.udf(lambda val:val[3],StringType())]
df = df.groupby('org').agg(*[F.collect_list('name').alias('name'), F.collect_list('value').alias('value')])\
.withColumn('statistics', find_stats(F.col('name'), F.col('value')))\
.drop('name').drop('value')\
.select('org', F.explode('statistics').alias('statistics'))\
.select(['org']+[split_('statistics').alias(col_name) for split_,col_name in zip(splits,cols)])\
.dropDuplicates()
df.show()
+-----+----+-----------------+------------------+-----+
| org|name| mean| stddev|count|
+-----+----+-----------------+------------------+-----+
|org_1| c| 2.5| 1.118033988749895| 4|
|org_2| e| 5.5| 1.5| 2|
|org_2| a|6.333333333333333|2.6246692913372702| 3|
|org_2| d|7.333333333333333|2.0548046676563256| 3|
|org_1| a| 4.5| 0.5| 2|
|org_1| b| 2.75| 1.479019945774904| 4|
+-----+----+-----------------+------------------+-----+
If you also want 'value' column, you can insert that in the tuple in udf function and add one split udf.
Also, since there will be duplicates in the dataframe due to repetition of names, you can remove them using dropDuplicates.
Here is a way to do this using only DataFrame functions.
Just join your DataFrame to itself on the org column and use a where clause to specify that the name column should be different. Then we select the distinct rows of ('l.org', 'l.name', 'r.name', 'r.value') - essentially, we ignore the l.value column because we want to avoid double counting for the same (org, name) pair.
For example, this is how you could collect the other values for each ('org', 'name') pair:
import pyspark.sql.functions as f
df.alias('l').join(df.alias('r'), on='org')\
.where('l.name != r.name')\
.select('l.org', 'l.name', 'r.name', 'r.value')\
.distinct()\
.groupBy('l.org', 'l.name')\
.agg(f.collect_list('r.value').alias('other_values'))\
.show()
#+-----+----+------------+
#| org|name|other_values|
#+-----+----+------------+
#|org_1| a| [4, 5]|
#|org_1| b|[1, 2, 3, 5]|
#|org_1| c|[1, 2, 3, 4]|
#|org_2| a| [4, 5, 10]|
#|org_2| d| [7, 5, 10]|
#|org_2| e| [7, 4]|
#+-----+----+------------+
For the descriptive stats, you can use the mean, stddev, and count functions from pyspark.sql.functions:
df.alias('l').join(df.alias('r'), on='org')\
.where('l.name != r.name')\
.select('l.org', 'l.name', 'r.name', 'r.value')\
.distinct()\
.groupBy('l.org', 'l.name')\
.agg(
f.mean('r.value').alias('mean'),
f.stddev('r.value').alias('stddev'),
f.count('r.value').alias('count')
)\
.show()
#+-----+----+-----------------+------------------+-----+
#| org|name| mean| stddev|count|
#+-----+----+-----------------+------------------+-----+
#|org_1| a| 4.5|0.7071067811865476| 2|
#|org_1| b| 2.75| 1.707825127659933| 4|
#|org_1| c| 2.5|1.2909944487358056| 4|
#|org_2| a|6.333333333333333|3.2145502536643185| 3|
#|org_2| d|7.333333333333333|2.5166114784235836| 3|
#|org_2| e| 5.5|2.1213203435596424| 2|
#+-----+----+-----------------+------------------+-----+
Note that pyspark.sql.functions.stddev() returns the unbiased sample standard deviation. If you wanted the population standard deviation, use pyspark.sql.functions.stddev_pop():
df.alias('l').join(df.alias('r'), on='org')\
.where('l.name != r.name')\
.groupBy('l.org', 'l.name')\
.agg(
f.mean('r.value').alias('mean'),
f.stddev_pop('r.value').alias('stddev'),
f.count('r.value').alias('count')
)\
.show()
#+-----+----+-----------------+------------------+-----+
#| org|name| mean| stddev|count|
#+-----+----+-----------------+------------------+-----+
#|org_1| a| 4.5| 0.5| 2|
#|org_1| b| 2.75| 1.479019945774904| 4|
#|org_1| c| 2.5| 1.118033988749895| 4|
#|org_2| a|6.333333333333333|2.6246692913372702| 3|
#|org_2| d|7.333333333333333|2.0548046676563256| 3|
#|org_2| e| 5.5| 1.5| 2|
#+-----+----+-----------------+------------------+-----+
EDIT
As #NaomiHuang mentioned in the comments, you could also reduce l to the distinct org/name pairs before doing the join:
df.select('org', 'name')\
.distinct()\
.alias('l')\
.join(df.alias('r'), on='org')\
.where('l.name != r.name')\
.groupBy('l.org', 'l.name')\
.agg(f.collect_list('r.value').alias('other_values'))\
.show()
#+-----+----+------------+
#| org|name|other_values|
#+-----+----+------------+
#|org_1| a| [5, 4]|
#|org_1| b|[5, 1, 2, 3]|
#|org_1| c|[1, 2, 3, 4]|
#|org_2| a| [4, 5, 10]|
#|org_2| d| [7, 5, 10]|
#|org_2| e| [7, 4]|
#+-----+----+------------+

Spark : How to group by distinct values in DataFrame

I have a data in a file in the following format:
1,32
1,33
1,44
2,21
2,56
1,23
The code I am executing is following:
val sqlContext = new org.apache.spark.sql.SQLContext(sc)
import spark.implicits._
import sqlContext.implicits._
case class Person(a: Int, b: Int)
val ppl = sc.textFile("newfile.txt").map(_.split(","))
.map(p=> Person(p(0).trim.toInt, p(1).trim.toInt))
.toDF()
ppl.registerTempTable("people")
val result = ppl.select("a","b").groupBy('a).agg()
result.show
Expected Output is:
a 32, 33, 44, 23
b 21, 56
Instead of aggregation by sum, count, mean etc. I want every element in the row.
Try collect_set function inside agg()
val df = sc.parallelize(Seq(
(1,3), (1,6), (1,5), (2,1),(2,4)
(2,1))).toDF("a","b")
+---+---+
| a| b|
+---+---+
| 1| 3|
| 1| 6|
| 1| 5|
| 2| 1|
| 2| 4|
| 2| 1|
+---+---+
val df2 = df.groupBy("a").agg(collect_set("b")).show()
+---+--------------+
| a|collect_set(b)|
+---+--------------+
| 1| [3, 6, 5]|
| 2| [1, 4]|
+---+--------------+
And if you want duplicate entries , can use collect_list
val df3 = df.groupBy("a").agg(collect_list("b")).show()
+---+---------------+
| a|collect_list(b)|
+---+---------------+
| 1| [3, 6, 5]|
| 2| [1, 4, 1]|
+---+---------------+

Resources