在 SQL/Spark 中使用 Window 函数执行特定过滤器

Question

I currently have a large dataset, but for simplicity, it looks like this:

Person, Friend, Friendship_Score, Days_Known
Alice, Bob, 120, 56
Alice, Candy, 20, 23
Bob, Daniel, 24, 77
Bob, Alice, 120, 56
Candy, Alice, 20, 23
Daniel, Bob, 24, 77
Daniel, Ed, 56, 65
Daniel, Fin, 52, 54
Daniel, Gin, 22, 50
...

I want to use a window function on this dataset to make it look like this:

Alice, Bob, 120, 56
Bob, Daniel, 24, 77
Bob, Alice, 120, 56
Candy, Alice, 20, 23
Daniel, Bob, 24, 77
Daniel, Ed, 56, 65
Daniel, Fin, 52, 54

The logic behind the filter should be that for each person, we rank their friends in the order of how long they've known each other for (higher days_known value is at the top) and then only keep enough friends such that they have a friendship_score of 100.

For example, Alice would only need Bob because she has known him the longest and they have a friendship_score over 100. Bob would need both Daniel and Alice because Bob has known Daniel longer, but their friendship_score is only 24. However, after adding Alice, the next friend Bob has known the longest, the combined friendship_score is above 100.

I know we need to do some kind of window function and a rolling sum, but I am having trouble putting the ideas into code and was wondering if anyone could help with this. Thank you!

Answer 1

I don't have much experience with Spark but the docs indicate it supports both window functions and Selects From Selects which you will need to filter the result of the window function.

Notice that the running sum using window UNBOUNDED BELOW TO CURRENT ROW produces a sum bigger than 100 for the last record you want to keep. You really want a partial sum that does not include CURRENT ROW to filter and retain the right records. You can do this with the SUM window function and then subtract the score from current record. So your window function should read SUM(friendship_score) OVER (Partition By person Order By Days_Known desc ROWS UNBOUNDED BELOW TO CURRENT ROW) - friendship_score as prior_total_score

Select person, friend, friendship_score,days_known
From (
      Select *, SUM(friendship_score) OVER (Partition By person Order By Days_Known desc ROWS UNBOUNDED BELOW TO CURRENT ROW) - friendship_score as prior_total_score
      From MyTable
     )
Where prior_total_score < 100

You can add an Order By to the outer Select as desired.