Fugue项目教程：深入理解engine_context的使用

2025-06-10 17:03:56作者：秋阔奎Evelyn

前言

在数据处理领域，我们经常需要在不同计算引擎（如Pandas、Spark、Dask、Ray等）之间切换。Fugue项目提供了一种优雅的解决方案，使得我们能够在不重写代码的情况下，将逻辑从一个引擎迁移到另一个引擎。本文将重点介绍Fugue中的engine_context功能，这是实现代码与执行引擎解耦的关键工具。

engine_context基础概念

engine_context是Fugue提供的一个上下文管理器，用于设置Fugue API函数的默认执行引擎。它的主要作用是简化代码，避免在多个Fugue函数调用中重复指定相同的执行引擎。

基本用法示例

import pandas as pd
import fugue.api as fa 

df = pd.DataFrame({"a": [1,2]})
df.to_parquet("/tmp/f.parquet")

def dummy(df:pd.DataFrame) -> pd.DataFrame:
    return df

with fa.engine_context("dask"):
    df = fa.load("/tmp/f.parquet")
    res = fa.transform(df, dummy, schema="*")
    fa.show(res)
    fa.save(res, "/tmp/f_out.parquet")

在这个例子中，所有Fugue操作（load、transform、show、save）都会默认使用Dask引擎执行，而不需要在每个函数调用中单独指定。

引擎覆盖机制

engine_context设置的默认引擎可以被单个函数调用覆盖，这提供了极大的灵活性：

with fa.engine_context(engine=None):  # 默认使用Pandas
    df = fa.load("/tmp/f.parquet")
    res = fa.transform(df, dummy, schema="*", engine="dask")  # 覆盖为Dask
    fa.show(res)

这种机制允许我们在保持大部分操作为轻量级Pandas处理的同时，将计算密集型的transform操作委托给分布式引擎。

与函数结合的最佳实践

将engine_context与函数结合使用，可以创建完全与引擎无关的业务逻辑：

def logic():
    df = fa.load("/tmp/f.parquet")
    res = fa.transform(df, dummy, schema="*")
    fa.show(res)

with fa.engine_context("dask"):
    logic()

更进一步，我们可以将引擎选择参数化，创建完全可移植的工作流：

def logic(engine):
    with fa.engine_context(engine):
        df = fa.load("/tmp/f.parquet")
        res = fa.transform(df, dummy, schema="*")
    return res

dask_df = logic("dask")  # 返回Dask DataFrame
pandas_df = logic(None)   # 返回Pandas DataFrame

复杂工作流构建

engine_context不仅限于简单的数据转换，还可以用于构建包含循环、条件判断等复杂逻辑的工作流：

from fugue.column import col, lit

with fa.engine_context():
    df = fa.load("/tmp/f.parquet")
    df = fa.assign(df, x=lit(1))
    for i in range(4):
        df = fa.assign(df, x=col("x")*lit(2))
    fa.show(df)