PySpark 재귀·계층 데이터 — BOM·조직도를 재귀 CTE 없이 풀기

parent_child 테이블:
  parent | child
  A      | B
  A      | C
  B      | D
  D      | E      ← A 아래 4단계 깊이
 
질문: "A 의 모든 후손은?"  → {B, C, D, E}  (깊이를 미리 모름)

from pyspark.sql import functions as F
 
edges = spark.table("parent_child")            # parent, child
spark.sparkContext.setCheckpointDir("/tmp/hier-ckpt")
 
def descendants(root):
    # 시작: root 의 직접 자식
    frontier = edges.where(F.col("parent") == root).select("child")
    result = frontier
    iteration = 0
 
    while True:
        # 현재 frontier 의 자식들을 한 단계 확장
        next_level = (frontier.alias("f")
            .join(edges.alias("e"), F.col("f.child") == F.col("e.parent"))
            .select(F.col("e.child").alias("child"))
            .distinct())
 
        # 이미 찾은 것 제외 (새로 발견된 후손만)
        new_nodes = next_level.join(result, "child", "left_anti")
 
        if new_nodes.isEmpty():     # 더 이상 새 후손 없음 → 종료
            break
 
        result = result.unionByName(new_nodes).distinct()
        frontier = new_nodes
        iteration += 1
 
        # 반복 lineage 절단 (필수)
        if iteration % 3 == 0:
            result = result.checkpoint()
 
    return result

def expand_with_path(root):
    # 시작: 깊이 1, 경로 = [root, child]
    frontier = (edges.where(F.col("parent") == root)
        .select(
            F.col("child"),
            F.lit(1).alias("level"),
            F.array(F.lit(root), F.col("child")).alias("path")))
    result = frontier
 
    while True:
        nxt = (frontier.alias("f")
            .join(edges.alias("e"), F.col("f.child") == F.col("e.parent"))
            .select(
                F.col("e.child").alias("child"),
                (F.col("f.level") + 1).alias("level"),
                F.concat(F.col("f.path"), F.array(F.col("e.child"))).alias("path")))
 
        # 순환 방어: 경로 안에 이미 등장한 노드면 제외
        nxt = nxt.where(~F.array_contains(
            F.expr("slice(path, 1, size(path)-1)"), F.col("child")))
 
        new_nodes = nxt.join(result.select("child"), "child", "left_anti")
        if new_nodes.isEmpty():
            break
        result = result.unionByName(new_nodes)
        frontier = new_nodes
 
    return result   # child, level, path

MAX_DEPTH = 50      # 안전 상한 — 정상 계층이 이보다 깊을 리 없다면
if iteration > MAX_DEPTH:
    raise RuntimeError("순환 의심: 최대 깊이 초과")

# edges: parent, child, qty
frontier = (edges.where(F.col("parent") == product)
    .select("child", F.col("qty").alias("total_qty")))
 
# 확장 시 수량을 곱해서 누적
nxt = (frontier.alias("f")
    .join(edges.alias("e"), F.col("f.child") == F.col("e.parent"))
    .select(
        F.col("e.child").alias("child"),
        (F.col("f.total_qty") * F.col("e.qty")).alias("total_qty")))
 
# 같은 부품이 여러 경로로 쓰이면 최종 합산
result.groupBy("child").agg(F.sum("total_qty").alias("required_qty"))

from graphframes import GraphFrame
g = GraphFrame(vertices, edges)
 
# root 에서 도달 가능한 노드 탐색 (BFS)
paths = g.bfs(fromExpr=f"id = '{root}'", toExpr="id IS NOT NULL", maxPathLength=10)