liaozhaorun
181994f063
- 移除 Python 多进程/多线程池，消除 DataFrame 序列化开销 - 采用 BFS 分层执行策略，每层表达式通过单次 with_columns 提交 - 利用 Polars Rust 引擎实现零拷贝并行计算 - 添加死锁检测机制处理依赖环
157 lines
4.9 KiB
Python
157 lines
4.9 KiB
Python
"""计算引擎。
|
||
|
||
执行并行运算，负责将执行计划应用到数据上。
|
||
|
||
利用 Polars 底层 Rust 引擎的原生并行能力，通过 BFS 分层执行策略
|
||
避免 Python 层面的多进程/多线程开销。
|
||
"""
|
||
|
||
from typing import Dict, List, Set
|
||
|
||
import polars as pl
|
||
|
||
from src.factors.engine.data_spec import ExecutionPlan
|
||
|
||
|
||
class ComputeEngine:
|
||
"""计算引擎 - 执行并行运算。
|
||
|
||
负责将执行计划应用到数据上，利用 Polars 底层 Rust 引擎的原生并行能力。
|
||
|
||
采用 BFS 分层执行策略：
|
||
1. 构建依赖图，识别各计划间的依赖关系
|
||
2. 按拓扑排序分层，每层包含互不依赖的计划
|
||
3. 将每层计划打包为表达式列表，通过单次 with_columns 提交
|
||
4. Polars 自动在所有 CPU 核心上并行计算，零拷贝内存
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
"""初始化计算引擎。"""
|
||
pass
|
||
|
||
def execute(
|
||
self,
|
||
plan: ExecutionPlan,
|
||
data: pl.DataFrame,
|
||
) -> pl.DataFrame:
|
||
"""执行单个计算计划。
|
||
|
||
Args:
|
||
plan: 执行计划
|
||
data: 输入数据（核心宽表）
|
||
|
||
Returns:
|
||
包含因子结果的 DataFrame
|
||
"""
|
||
# 检查依赖字段是否存在
|
||
missing_cols = plan.dependencies - set(data.columns)
|
||
if missing_cols:
|
||
raise ValueError(f"数据缺少必要的字段: {missing_cols}")
|
||
|
||
# 执行计算
|
||
return data.with_columns([plan.polars_expr.alias(plan.output_name)])
|
||
|
||
def execute_batch(
|
||
self,
|
||
plans: List[ExecutionPlan],
|
||
data: pl.DataFrame,
|
||
) -> pl.DataFrame:
|
||
"""顺序批量执行多个计算计划。
|
||
|
||
Args:
|
||
plans: 执行计划列表
|
||
data: 输入数据
|
||
|
||
Returns:
|
||
包含所有因子结果的 DataFrame
|
||
"""
|
||
result = data
|
||
for plan in plans:
|
||
result = self.execute(plan, result)
|
||
return result
|
||
|
||
def execute_parallel(
|
||
self,
|
||
plans: List[ExecutionPlan],
|
||
data: pl.DataFrame,
|
||
) -> pl.DataFrame:
|
||
"""分层并行执行计算计划（利用 Polars 原生并发优化）。
|
||
|
||
抛弃 Python 的多进程/多线程池，采用计算图拓扑分层（BFS DAG）。
|
||
将每一层互不依赖的表达式列表打包，通过单次 with_columns 交给 Polars，
|
||
由底层 Rust 引擎自动调度并行计算，实现零拷贝性能最大化。
|
||
|
||
Args:
|
||
plans: 执行计划列表
|
||
data: 输入数据
|
||
|
||
Returns:
|
||
包含所有因子结果的 DataFrame
|
||
|
||
Raises:
|
||
RuntimeError: 当存在依赖环或缺少基础依赖字段时
|
||
"""
|
||
if not plans:
|
||
return data
|
||
|
||
result = data
|
||
available_cols: Set[str] = set(result.columns)
|
||
|
||
# 复制一份计划列表用于迭代
|
||
remaining_plans = plans.copy()
|
||
|
||
while remaining_plans:
|
||
# 找出当前可以执行的所有独立计划（即依赖的所有列都已就绪）
|
||
current_layer: List[ExecutionPlan] = []
|
||
next_remaining: List[ExecutionPlan] = []
|
||
|
||
for plan in remaining_plans:
|
||
if plan.dependencies <= available_cols:
|
||
current_layer.append(plan)
|
||
else:
|
||
next_remaining.append(plan)
|
||
|
||
# 安全兜底：如果一轮遍历后没找到任何可执行计划，说明存在依赖环或数据缺失
|
||
if not current_layer:
|
||
missing = remaining_plans[0].dependencies - available_cols
|
||
raise RuntimeError(
|
||
f"计算发生死锁或缺少基础依赖字段！\n"
|
||
f"因子 '{remaining_plans[0].output_name}' 缺少: {missing}"
|
||
)
|
||
|
||
# 核心优化：利用 Polars 内部 Rust 级多线程引擎执行当前层
|
||
exprs = [plan.polars_expr.alias(plan.output_name) for plan in current_layer]
|
||
result = result.with_columns(exprs)
|
||
|
||
# 更新已就绪字段集合，为计算下一层做准备
|
||
for plan in current_layer:
|
||
available_cols.add(plan.output_name)
|
||
|
||
remaining_plans = next_remaining
|
||
|
||
return result
|
||
|
||
def compute(
|
||
self,
|
||
plans: List[ExecutionPlan],
|
||
data: pl.DataFrame,
|
||
parallel: bool = True,
|
||
) -> pl.DataFrame:
|
||
"""智能计算入口。
|
||
|
||
根据 parallel 参数自动选择执行模式：
|
||
- True: 使用分层并行执行（推荐）
|
||
- False: 使用顺序执行
|
||
|
||
Args:
|
||
plans: 执行计划列表
|
||
data: 输入数据
|
||
parallel: 是否使用并行执行
|
||
|
||
Returns:
|
||
包含所有因子结果的 DataFrame
|
||
"""
|
||
if parallel:
|
||
return self.execute_parallel(plans, data)
|
||
return self.execute_batch(plans, data)
|