如何在代码中覆盖Kedro数据集查询的WHERE子句？

问题场景 #

现有Kedro目录配置如下：

some_table.raw:
    type: pandas.GBQQueryDataset
    sql: SELECT * FROM database.table WHERE date >= {start_date}

需要在代码中动态传递start_date参数，期望调用方式类似：

catalog.load("some_table.raw", query={"start_date": "2024-01-01"})

但catalog.load不支持额外参数，且TemplatedConfigLoader只能绑定单一固定值，无法满足代码不同部分传递不同参数的需求。

三种可行解决方案 #

1. 自定义数据集类 #

继承GBQQueryDataset，扩展支持动态参数注入：

from kedro_datasets.pandas import GBQQueryDataset

class DynamicGBQQueryDataset(GBQQueryDataset):
    def __init__(self, sql: str, **kwargs):
        self.base_sql = sql
        super().__init__(sql=sql, **kwargs)
    
    def load(self, **query_params):
        # 渲染带参数的SQL语句
        self.sql = self.base_sql.format(**query_params)
        return super().load()

更新目录配置使用自定义数据集：

some_table.raw:
    type: your.module.path.DynamicGBQQueryDataset
    sql: SELECT * FROM database.table WHERE date >= {start_date}
    # 补充GBQ必填参数，如project_id、credentials等

代码中调用：

data = catalog.load("some_table.raw").load(start_date="2024-01-01")

2. 利用Kedro钩子动态修改配置 #

通过before_dataset_created钩子，在数据集实例化前替换SQL参数：

from kedro.framework.hooks import hook_impl
from kedro_datasets.pandas import GBQQueryDataset

class DynamicQueryHook:
    # 用类属性临时存储参数，避免全局变量污染
    _query_params = {}

    @classmethod
    def set_params(cls, **params):
        cls._query_params = params

    @hook_impl
    def before_dataset_created(self, dataset_name: str, dataset_config: dict):
        if dataset_name == "some_table.raw" and dataset_config["type"] == "pandas.GBQQueryDataset":
            if self._query_params:
                dataset_config["sql"] = dataset_config["sql"].format(**self._query_params)
                self._query_params = {}  # 重置参数，避免影响后续加载

在settings.py中注册钩子：

HOOKS = [DynamicQueryHook()]

代码中调用：

DynamicQueryHook.set_params(start_date="2024-01-01")
data = catalog.load("some_table.raw")

3. 编写数据集加载工厂函数 #

直接从配置中读取模板SQL，渲染后创建数据集实例：

from kedro.io import DataCatalog
from kedro_datasets.pandas import GBQQueryDataset

def load_gbq_with_params(catalog: DataCatalog, dataset_name: str, **query_params):
    # 获取原始数据集配置
    catalog_config = catalog._config_loader.get("catalog")[dataset_name]
    # 渲染SQL模板
    rendered_sql = catalog_config["sql"].format(**query_params)
    # 构建并返回加载后的数据
    dataset = GBQQueryDataset(
        sql=rendered_sql,
        **{k: v for k, v in catalog_config.items() if k not in ["type", "sql"]}
    )
    return dataset.load()

代码中调用：

data = load_gbq_with_params(catalog, "some_table.raw", start_date="2024-01-01")

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