PostgreSQL递归查询转Python Pandas实现及DataFrame合并异常排查

一、原PL/pgSQL递归逻辑示例 #

假设原逻辑是查询树形部门结构，PL/pgSQL代码如下：

CREATE OR REPLACE FUNCTION get_dept_tree(root_dept_id INT)
RETURNS TABLE(dept_id INT, dept_name TEXT, parent_id INT, depth INT) AS $$
BEGIN
  RETURN QUERY
  WITH RECURSIVE dept_rec AS (
    -- 递归基础：根节点
    SELECT dept_id, dept_name, parent_id, 1 AS depth
    FROM departments
    WHERE dept_id = root_dept_id
    UNION ALL
    -- 递归部分：子节点
    SELECT d.dept_id, d.dept_name, d.parent_id, dr.depth + 1
    FROM departments d
    JOIN dept_rec dr ON d.parent_id = dr.dept_id
  )
  SELECT * FROM dept_rec;
END;
$$ LANGUAGE plpgsql;

二、转成plpython3u的递归实现（解决合并问题） #

方案1：列表累加结果（效率更高） #

优先用列表收集所有行数据，最后一次性转成DataFrame，从根源避免合并丢失：

CREATE OR REPLACE FUNCTION get_dept_tree_py(root_dept_id INT)
RETURNS TABLE(dept_id INT, dept_name TEXT, parent_id INT, depth INT) AS $$
import pandas as pd
import plpy

def recursive_scan(parent_id, current_depth, result_list):
    # 查询当前父节点的子节点
    query = """SELECT dept_id, dept_name, parent_id FROM departments WHERE parent_id = %s"""
    rows = plpy.execute(query, [parent_id])
    if not rows:
        return
    # 遍历子节点，加入结果列表并递归
    for row in rows:
        result_list.append({
            'dept_id': row['dept_id'],
            'dept_name': row['dept_name'],
            'parent_id': row['parent_id'],
            'depth': current_depth + 1
        })
        recursive_scan(row['dept_id'], current_depth + 1, result_list)

# 初始化结果列表，先加入根节点
root_query = """SELECT dept_id, dept_name, parent_id FROM departments WHERE dept_id = %s"""
root_rows = plpy.execute(root_query, [root_dept_id])
if not root_rows:
    return
root_row = root_rows[0]
result = [{
    'dept_id': root_row['dept_id'],
    'dept_name': root_row['dept_name'],
    'parent_id': root_row['parent_id'],
    'depth': 1
}]

# 启动递归扫描
recursive_scan(root_dept_id, 1, result)

# 转成DataFrame并返回结果
df = pd.DataFrame(result)
for _, row in df.iterrows():
    yield (row['dept_id'], row['dept_name'], row['parent_id'], row['depth'])

$$ LANGUAGE plpython3u;

方案2：直接用DataFrame递归合并（适合中间需要DataFrame处理的场景） #

确保每次递归返回的子DataFrame都合并到当前结果：

CREATE OR REPLACE FUNCTION get_dept_tree_py_df(root_dept_id INT)
RETURNS TABLE(dept_id INT, dept_name TEXT, parent_id INT, depth INT) AS $$
import pandas as pd
import plpy

def recursive_merge(parent_id, current_depth):
    query = """SELECT dept_id, dept_name, parent_id FROM departments WHERE parent_id = %s"""
    rows = plpy.execute(query, [parent_id])
    current_df = pd.DataFrame()
    for row in rows:
        # 构建当前子节点的临时DataFrame
        temp_df = pd.DataFrame([{
            'dept_id': row['dept_id'],
            'dept_name': row['dept_name'],
            'parent_id': row['parent_id'],
            'depth': current_depth + 1
        }])
        # 合并到当前结果
        current_df = pd.concat([current_df, temp_df], ignore_index=True)
        # 递归查询子节点并合并
        child_df = recursive_merge(row['dept_id'], current_depth + 1)
        current_df = pd.concat([current_df, child_df], ignore_index=True)
    return current_df

# 初始化根节点DataFrame
root_query = """SELECT dept_id, dept_name, parent_id FROM departments WHERE dept_id = %s"""
root_rows = plpy.execute(root_query, [root_dept_id])
if not root_rows:
    return
root_row = root_rows[0]
root_df = pd.DataFrame([{
    'dept_id': root_row['dept_id'],
    'dept_name': root_row['dept_name'],
    'parent_id': root_row['parent_id'],
    'depth': 1
}])

# 合并根节点与递归结果
final_df = pd.concat([root_df, recursive_merge(root_dept_id, 1)], ignore_index=True)
for _, row in final_df.iterrows():
    yield (row['dept_id'], row['dept_name'], row['parent_id'], row['depth'])

$$ LANGUAGE plpython3u;

三、DataFrame合并后仅存最后一个的原因排查 #

常见错误点 #

1. 递归内部重复初始化容器：每次递归都新建pd.DataFrame()，未传递之前的累加结果，导致历史数据丢失。
2. 忽略Pandas方法的返回值：pd.concat/append均返回新对象，原对象不会被修改，若未重新赋值，合并结果不会保留。
3. 未合并递归返回的子结果：调用递归后未将子节点返回的DataFrame合并到当前结果，导致子节点数据丢失。
4. 终止条件错误：未正确判断递归结束条件，导致仅执行最后一层递归。

错误代码示例（反面教材） #

# 错误：每次递归新建df且未合并子结果
def bad_recursive(parent_id, depth):
    df = pd.DataFrame()
    rows = plpy.execute("SELECT * FROM departments WHERE parent_id = %s", [parent_id])
    for row in rows:
        temp_df = pd.DataFrame([{'dept_id': row['dept_id'], 'depth': depth}])
        df = pd.concat([df, temp_df])
        # 错误：调用递归但未合并返回的child_df
        bad_recursive(row['dept_id'], depth +1)
    return df

内容的提问来源于stack exchange，提问作者Gerzzog