Skip to content

pandas.SQLQueryDataset + polars.PolarsDatabaseDataset: built-in Jinja2 templating via template_params #1385

Description

@nikanton

Description

Both pandas.SQLQueryDataset and the experimental polars.PolarsDatabaseDataset accept a SQL string or a filepath to a .sql file, but treat the file as a static blob. There is no way to parameterize it beyond SQLAlchemy bind parameters, which cannot change query structure — only scalar values inside a fixed shape. I would like first-class Jinja2 templating support baked into these datasets so that a single .sql.j2 file can serve multiple modes via load_args.template_params.

Context

Without templating, the path of least resistance for "the same query with slightly different filters" is to duplicate the file: one .sql per variant, each a copy-paste of the other with a few lines changed. The result is predictable — duplicated CTEs drift out of sync, a bug fix in one copy silently misses the others, and onboarding a new contributor means asking "which of these 90%-identical files is the source of truth?"

Two use cases would benefit enormously from a built-in templating option:

1. Custom training / testing windows driven from config. A data scientist wants to re-run the pipeline over a different date range, or compare results across several date ranges. With Jinja, the date range lives in globals.yml and the SQL file stays untouched — the dates simply flow through template_params:

raw_events_training:
  type: pandas.SQLQueryDataset
  filepath: sql/events.sql.j2
  credentials: warehouse
  load_args:
    template_params:
      start_date: "${globals:train_start_date}"
      end_date:   "${globals:train_end_date}"

2. No duplication between training and inference. Training and real-time scoring usually need the same feature-engineering SQL with a small filter difference. A Jinja-templated query lets both pipelines point at the same .sql.j2 file and differ only in a mode parameter that toggles the relevant branches inside the template:

events_train:
  type: pandas.SQLQueryDataset
  filepath: sql/events.sql.j2
  credentials: warehouse
  load_args:
    template_params:
      mode: train
      start_date: "${globals:train_start_date}"
      end_date:   "${globals:train_end_date}"

events_inference:
  type: pandas.SQLQueryDataset
  filepath: sql/events.sql.j2   # same file
  credentials: warehouse
  load_args:
    template_params:
      mode: inference

Without templating, the practical alternatives are (a) maintain two copies of each query and hope they stay in lockstep, or (b) redo feature engineering in Python after the fact, which opens the door to training/serving skew. Neither is great.

Possible Implementation

Add an opt-in template_params key inside load_args. When present, render the SQL (whether inline sql or file-loaded from filepath) through Jinja2 before execution. When absent, behaviour is unchanged — fully backwards compatible. Rendering is triggered only when template_params is present, regardless of the file extension, so the trigger is always explicit and opt-in.

Sketch for SQLQueryDataset.load (same shape would apply to PolarsDatabaseDataset):

def _render(self, raw_sql: str) -> str:
    params = self._load_args.pop("template_params", None)
    if params is None:
        return raw_sql

    from jinja2 import Environment, FileSystemLoader, StrictUndefined

    loader = (
        FileSystemLoader(str(self._filepath.parent))
        if self._filepath else None
    )
    env = Environment(loader=loader, undefined=StrictUndefined)
    template = (
        env.get_template(self._filepath.name) if self._filepath
        else env.from_string(raw_sql)
    )
    return template.render(**params)

Catalog usage on polars.PolarsDatabaseDataset looks the same:

events_inference:
  type: polars.PolarsDatabaseDataset
  filepath: sql/events.sql.j2
  table_name: events
  credentials: warehouse
  load_args:
    template_params:
      mode: inference

Macro support

A big part of the value of Jinja in SQL is reusing filter logic across queries via macros — e.g. a single _macros.sql.j2 that defines a window filter, imported by every query that needs it:

{# sql/_macros.sql.j2 #}
{%- macro date_window(mode, start_date=None, end_date=None) -%}
  {%- if mode == "train" %}
    AND event_time >= DATE '{{ start_date }}'
    AND event_time <  DATE '{{ end_date }}'
  {%- elif mode == "inference" %}
    AND event_time >= SYSDATE - INTERVAL '24' HOUR
  {%- endif -%}
{%- endmacro -%}
{# sql/events.sql.j2 #}
{%- from '_macros.sql.j2' import date_window -%}
SELECT ...
FROM events
WHERE 1=1
  {{ date_window(mode, start_date=start_date, end_date=end_date) }}

For this to work, the Jinja FileSystemLoader is pointed at the directory of the filepath, so macros that sit next to the query (e.g. sql/_macros.sql.j2 alongside sql/events.sql.j2) import with zero extra configuration. Users who don't use macros get this automatically at no cost.

Other design notes

  • Jinja as an optional dep: Jinja2 is already an indirect dep of Kedro (via cookiecutter), so no new required dependency, but the import can be done lazily inside the render helper to keep import time unaffected.
  • Composable with SQLAlchemy bind parameters: Jinja runs first to produce a SQL string; bind params apply after. Both mechanisms coexist and do different things (structure vs. values).

Happy to open a PR for pandas.SQLQueryDataset first (the stable path) and follow up with polars.PolarsDatabaseDataset afterwards.

Possible Alternatives

  • Duplicate SQL files per variant (the status quo when structural parameterization is needed): simple at first, expensive over time as duplicated CTEs drift out of sync.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions