[PY]星火SQL：多列sessionization

Question

給定一個正龍i和數據幀

+-----+--+--+               
|group|n1|n2|                
+-----+--+--+                
| 1| 0| 0|                
| 1| 1| 1|                
| 1| 1| 5|                
| 1| 2| 2|                
| 1| 2| 6|                
| 1| 3| 3|                
| 1| 3| 7|                
| 1| 4| 4|                
| 1| 5| 1|                
| 1| 5| 5|                
+-----+--+--+ 

你會如何sessionize排在同group使得對於每對連續行r1，r2中的會議r2.n1>r1.n1，r2.n2>r1.n2，和max（r2.n1 - r1.n1，r2.n2 - r1.n2）< i？請注意，n1和n2值可能不是唯一的，這意味着構成會話的行在DataFrame中可能不連續。

作爲示例，對於給定的數據幀和i = 3的結果將是

+-----+--+--+-------+ 
|group|n1|n2|session| 
+-----+--+--+-------+ 
| 1| 0| 0|  1| 
| 1| 1| 1|  1| 
| 1| 1| 5|  2| 
| 1| 2| 2|  1| 
| 1| 2| 6|  2| 
| 1| 3| 3|  1| 
| 1| 3| 7|  2| 
| 1| 4| 4|  1| 
| 1| 5| 1|  3| 
| 1| 5| 5|  1| 
+-----+--+--+-------+ 

任何幫助或暗示將不勝感激。謝謝！

Answer 1

這看起來像你想用一個相同的數字來標記一個圖的所有連接部分。一個好的解決方案是使用graphframes：https://graphframes.github.io/quick-start.html

從你的數據框：

df = sc.parallelize([[1, 0, 0],[1, 1, 1],[1, 1, 5],[1, 2, 2],[1, 2, 6], 
        [1, 3, 3],[1, 3, 7],[1, 4, 4],[1, 5, 1],[1, 5, 5]]).toDF(["group","n1","n2"]) 

我們將創建一個包含獨特id是清單頂點數據幀：

import pyspark.sql.functions as psf 
v = df.select(psf.struct("n1", "n2").alias("id"), "group") 

    +-----+-----+ 
    | id|group| 
    +-----+-----+ 
    |[0,0]| 1| 
    |[1,1]| 1| 
    |[1,5]| 1| 
    |[2,2]| 1| 
    |[2,6]| 1| 
    |[3,3]| 1| 
    |[3,7]| 1| 
    |[4,4]| 1| 
    |[5,1]| 1| 
    |[5,5]| 1| 
    +-----+-----+ 

而且從您陳述的布爾條件定義的邊緣數據幀：

i = 3 
e = df.alias("r1").join(
    df.alias("r2"), 
    (psf.col("r1.group") == psf.col("r2.group")) 
    & (psf.col("r1.n1") < psf.col("r2.n1")) 
    & (psf.col("r1.n2") < psf.col("r2.n2")) 
    & (psf.greatest(
     psf.col("r2.n1") - psf.col("r1.n1"), 
     psf.col("r2.n2") - psf.col("r1.n2")) < i) 
).select(psf.struct("r1.n1", "r1.n2").alias("src"), psf.struct("r2.n1", "r2.n2").alias("dst")) 

    +-----+-----+ 
    | src| dst| 
    +-----+-----+ 
    |[0,0]|[1,1]| 
    |[0,0]|[2,2]| 
    |[1,1]|[2,2]| 
    |[1,1]|[3,3]| 
    |[1,5]|[2,6]| 
    |[1,5]|[3,7]| 
    |[2,2]|[3,3]| 
    |[2,2]|[4,4]| 
    |[2,6]|[3,7]| 
    |[3,3]|[4,4]| 
    |[3,3]|[5,5]| 
    |[4,4]|[5,5]| 
    +-----+-----+ 

現在找到所有連接的組件：

from graphframes import * 
g = GraphFrame(v, e) 
res = g.connectedComponents() 

    +-----+-----+------------+ 
    | id|group| component| 
    +-----+-----+------------+ 
    |[0,0]| 1|309237645312| 
    |[1,1]| 1|309237645312| 
    |[1,5]| 1| 85899345920| 
    |[2,2]| 1|309237645312| 
    |[2,6]| 1| 85899345920| 
    |[3,3]| 1|309237645312| 
    |[3,7]| 1| 85899345920| 
    |[4,4]| 1|309237645312| 
    |[5,1]| 1|292057776128| 
    |[5,5]| 1|309237645312| 
    +-----+-----+------------+