Get first example element from filtered aggregation pySpark

user_6696764 0
Viewed 121

Table:

+---------+---------+---------+---------+
|id       |path     |error    |message  |
+---------+---------+---------+---------+
|      1  |    a.a.a|   true  |   "aaa" |
+---------+---------+---------+---------+
|      2  |    a.a.a|   true  |   "bbb" |
+---------+---------+---------+---------+
|      2  |    a.a.a|   true  |   "bbc" |
+---------+---------+---------+---------+
|      2  |    a.a.b|   false |   "ccc" |
+---------+---------+---------+---------+

I have pySpark query:

data.groupBy('id', 'path')\
.agg(
    sum(when(col('error') == 'true', 1).otherwise(0)).alias('count'),
).show()

How can I add a new column with first example element that have col('error') == 'true'? I want a table with elements id, path, count, exampleItem.
Function first().alias('exampleItem') works, but returns elements that do not necessarily match the condition above.

pyspark apache-spark dataframe
2 Answers

What you can do is define a column that only contains the message in the case where the error is true, and null otherwise using the when function. Then, the first function used with ignorenulls set to true will give you what you expect.

d = [(1, "a.a.a", True, "aaa"), (2, "a.a.a", True, "bbb"),
     (2, "a.a.a", True, "bbc"), (2, "a.a.b", False, "ccc")]
data = spark.createDataFrame(d, ['id', 'path', 'error', 'message'])

data\
    .groupBy('id', 'path')\
    .agg(F.sum(F.when(F.col('error') == 'true', 1).otherwise(0)).alias('count'),
         F.first(F.when(F.col('error'), F.col('message')), ignorenulls=True).alias('exampleItem'))\
    .show()

+---+-----+-----+-----------+
| id| path|count|exampleItem|
+---+-----+-----+-----------+
|  2|a.a.a|    2|        bbb|
|  1|a.a.a|    1|        aaa|
|  2|a.a.b|    0|       null|
+---+-----+-----+-----------+

For the last row, the null value is due to the fact that no message meet the desired requirement.

user_88936860

You can use first with the ignorenulls=True option and a case when statement to get the first message with error = true.

from pyspark.sql import functions as F, Window

df2 = df.groupBy(
    'id', 'path'
).agg(
    F.sum(F.when(F.col('error') == 'true', 1).otherwise(0)).alias('count'),
    F.first(
        F.when(F.col('error') == 'true', F.col('message')),
        ignorenulls=True
    ).alias('exampleItem')
).orderBy('id', 'path')

df2.show()
+---+-----+-----+-----------+
| id| path|count|exampleItem|
+---+-----+-----+-----------+
|  1|a.a.a|    1|        aaa|
|  2|a.a.a|    2|        bbb|
|  2|a.a.b|    0|       null|
+---+-----+-----+-----------+

However, note that there is no ordering column defined in your sample dataframe, so first has no concrete meaning and will only return a random element with error = true. Presumably there should be a timestamp column which can be used to order the dataframe within each id/path partition.

user_141657300
Related