熊猫中的递归SQL CTE查询?


问题内容

如何将下面的递归SQL查询优雅地移植到Pandas python代码?不知如何在没有编写自己的递归函数的情况下看不到直接的方法?

Python示例代码:

import datetime
import numpy as np
import pandas as pd
import pandas.io.data
from pandas import Series, DataFrame

data = {
        'ID': [1,2,3,4,5,6,7,8],
        'Name': ['Keith','Josh','Robin','Raja','Tridip','Arijit','Amit','Dev'],
        'MgrID': [0,1,1,2,0,5,5,6]
    }

df = pd.DataFrame.from_dict(data)
df.set_index('ID', inplace=True, drop=False, append=False)
df.ix[df.query('MgrID >0')['MgrID']]

试图得到这个:

nLevel      ID          Name
================================
1           6            Arijit
2           8               Dev
1           1            Keith
2           2               Josh
2           3               Robin
3           4                 Raja
1           5            Tridip
2           7               Amit

递归SQL查询:

;WITH Employee (ID, Name, MgrID) AS 
(
    SELECT 1,      'Keith',      NULL   UNION ALL
    SELECT 2,      'Josh',       1      UNION ALL
    SELECT 3,      'Robin',      1      UNION ALL
    SELECT 4,      'Raja',       2      UNION ALL
    SELECT 5,      'Tridip',     NULL   UNION ALL
    SELECT 6,      'Arijit',     NULL      UNION ALL
    SELECT 7,      'Amit',       5      UNION ALL
    SELECT 8,      'Dev',        6   
)
,Hierarchy AS
(
    --  Anchor
    SELECT   ID
            ,Name
            ,MgrID
            ,nLevel = 1
            ,Family = ROW_NUMBER() OVER (ORDER BY Name)
    FROM Employee
    WHERE MgrID IS NULL

    UNION ALL
    --  Recursive query
    SELECT   E.ID
            ,E.Name
            ,E.MgrID
            ,H.nLevel+1
            ,Family
    FROM Employee   E
    JOIN Hierarchy  H ON E.MgrID = H.ID
)
SELECT nLevel ,ID,space(nLevel+(CASE WHEN nLevel > 1 THEN nLevel ELSE 0 END))+Name Name FROM Hierarchy ORDER BY Family, nLevel

问题答案:

首先,您需要更正python代码MgrID列表中的错字: [0,1,1,2,0,0,5,6]

其次,如果这项工作是在SQL中以递归方式完成的,为什么您期望Python / Pandas可以不用递归方法来完成呢?并不难:

def nlevel(id, mgr_dict=df.MgrID, _cache={0:0}):
    if id in _cache:
        return _cache[id]

    return 1+nlevel(mgr_dict[id],mgr_dict)

df['nLevel'] = df.ID.map(nlevel)

print df[['nLevel','ID','Name']]

然后,output(nLevel)是您所需要的(顺序除外,我从您的SQL中无法理解该顺序):

    nLevel  ID    Name
ID                    
1        1   1   Keith
2        2   2    Josh
3        2   3   Robin
4        3   4    Raja
5        1   5  Tridip
6        1   6  Arijit
7        2   7    Amit
8        2   8     Dev

[8 rows x 3 columns]