Build a point-in-time correct training set and block feature leakage

id: point-in-time-training-set-leakage-gate
namespace: company.team

inputs:
  - id: join_strategy
    type: SELECT
    displayName: Feature join strategy
    description: >
      ASOF takes each feature as it was known before the label cutoff, which is
      correct. LATEST takes the newest feature value whatever its time, the
      common mistake that leaks the future into training. Pick LATEST to watch
      the gate block it.
    values:
      - ASOF
      - LATEST
    defaults: ASOF

  - id: embargo_days
    type: INT
    displayName: Embargo (days)
    description: Features must be at least this many days older than the label time.
      A gap covers the delay between an event and the moment it is really
      available.
    defaults: 7

  - id: min_coverage
    type: FLOAT
    displayName: Minimum feature coverage
    description: Share of labels that must get a feature value. Labels with no
      feature history before their cutoff count against it.
    defaults: 0.95

  - id: max_leakage_lift
    type: FLOAT
    displayName: Max leakage lift
    description: >
      Largest allowed gap between the feature/label correlation of the published
      set and of a known-correct point-in-time set. A large gap means the set
      predicts the label with information from after the cutoff.
    defaults: 0.05

variables:
  dataset_path: ml/churn/training_set.parquet
  manifest_kv: churn_training_set_manifest

tasks:
  - id: build_and_audit
    type: io.kestra.plugin.jdbc.duckdb.Queries
    description: >
      Generate a churn dataset (labels with a cutoff time plus weekly feature
      snapshots that keep changing after the cutoff), build the training set
      with the chosen join, then audit it: rows whose feature is newer than the
      cutoff, duplicate labels, coverage, and how much more the feature
      correlates with the label than it does in a correct point-in-time set.
      Writes the candidate as Parquet.
    outputFiles:
      - training
    fetchType: FETCH_ONE
    sql: |
      SELECT setseed(0.42);

      -- Labels: churn is decided at a per-customer cutoff between June and August.
      CREATE TABLE labels AS
      SELECT customer_id,
             TIMESTAMP '2026-06-01' + to_days(CAST(floor(random() * 60) AS INTEGER)) AS label_time,
             random() < 0.25 AS churned
      FROM range(1, 2001) t(customer_id);

      -- Weekly feature snapshots from April to the end of August. Churned customers file a few
      -- more tickets before the cutoff (weak, real signal) and many cancellation tickets after
      -- it (strong signal that does not exist yet at prediction time).
      CREATE TABLE feature_snapshots AS
      SELECT l.customer_id,
             TIMESTAMP '2026-04-01' + to_days(CAST(w * 7 AS INTEGER)) AS feature_time,
             CAST(floor(random() * 3) AS INTEGER)
               + CASE
                   WHEN l.churned AND TIMESTAMP '2026-04-01' + to_days(CAST(w * 7 AS INTEGER)) >= l.label_time
                     THEN 4 + CAST(floor(random() * 3) AS INTEGER)
                   WHEN l.churned THEN CAST(floor(random() * 2) AS INTEGER)
                   ELSE 0
                 END AS open_tickets_30d
      FROM labels l, range(0, 22) w(w);

      -- The reference point-in-time set, always built, so any candidate can be compared to it.
      CREATE TABLE pit_reference AS
      SELECT l.customer_id, l.label_time, l.churned, f.feature_time, f.open_tickets_30d
      FROM labels l
      ASOF LEFT JOIN feature_snapshots f
        ON l.customer_id = f.customer_id
       AND f.feature_time <= l.label_time - to_days({{ inputs.embargo_days }});

      CREATE TABLE training AS
      {% if inputs.join_strategy == 'ASOF' %}
      SELECT * FROM pit_reference;
      {% else %}
      SELECT l.customer_id, l.label_time, l.churned, f.feature_time, f.open_tickets_30d
      FROM labels l
      LEFT JOIN (
        SELECT * FROM feature_snapshots
        QUALIFY row_number() OVER (PARTITION BY customer_id ORDER BY feature_time DESC) = 1
      ) f ON l.customer_id = f.customer_id;
      {% endif %}

      COPY training TO '{{ outputFiles.training }}' (FORMAT PARQUET);

      SELECT
        (SELECT count(*) FROM training) AS rows,
        (SELECT count(*) FROM training WHERE feature_time > label_time - to_days({{ inputs.embargo_days }})) AS leaked_rows,
        (SELECT count(*) - count(DISTINCT customer_id) FROM training) AS duplicate_labels,
        round((SELECT avg(CASE WHEN feature_time IS NULL THEN 0 ELSE 1 END) FROM training), 4) AS coverage,
        round((SELECT corr(open_tickets_30d, CAST(churned AS INTEGER)) FROM training), 4) AS candidate_corr,
        round((SELECT corr(open_tickets_30d, CAST(churned AS INTEGER)) FROM pit_reference), 4) AS reference_corr,
        (SELECT strftime(max(feature_time), '%Y-%m-%d') FROM training) AS newest_feature,
        (SELECT strftime(min(label_time), '%Y-%m-%d') FROM training) AS earliest_label;

  - id: audit
    type: io.kestra.plugin.core.output.OutputValues
    description: The audit row, plus the lift of the candidate correlation over the
      point-in-time reference.
    values:
      rows: "{{ (outputs.build_and_audit.outputs | last).row.rows }}"
      leaked_rows: "{{ (outputs.build_and_audit.outputs | last).row.leaked_rows }}"
      duplicate_labels: "{{ (outputs.build_and_audit.outputs | last).row.duplicate_labels }}"
      coverage: "{{ (outputs.build_and_audit.outputs | last).row.coverage }}"
      candidate_corr: "{{ (outputs.build_and_audit.outputs | last).row.candidate_corr }}"
      reference_corr: "{{ (outputs.build_and_audit.outputs | last).row.reference_corr }}"
      leakage_lift: "{{ ((outputs.build_and_audit.outputs | last).row.candidate_corr -
        (outputs.build_and_audit.outputs | last).row.reference_corr) |
        numberFormat('0.0000') }}"

  - id: log_audit
    type: io.kestra.plugin.core.log.Log
    message: >-
      {{ inputs.join_strategy }} training set: {{ outputs.audit.values.rows }}
      labels, {{ outputs.audit.values.leaked_rows }} rows with a feature newer
      than cutoff minus {{ inputs.embargo_days }} days, {{
      outputs.audit.values.duplicate_labels }} duplicate labels, coverage {{
      outputs.audit.values.coverage }}. Feature/label correlation {{
      outputs.audit.values.candidate_corr }} vs {{
      outputs.audit.values.reference_corr }} in the point-in-time reference
      (lift {{ outputs.audit.values.leakage_lift }}).

  - id: leakage_gate
    type: io.kestra.plugin.core.flow.If
    description: Publish only a set with no future rows, no duplicate labels, enough
      coverage and no correlation lift over the point-in-time reference.
    condition: >-
      {{ (outputs.audit.values.leaked_rows | number) == 0 and
      (outputs.audit.values.duplicate_labels | number) == 0 and
      (outputs.audit.values.coverage | number) >= inputs.min_coverage and
      (outputs.audit.values.leakage_lift | number) <= inputs.max_leakage_lift }}
    then:
      - id: publish_training_set
        type: io.kestra.plugin.core.namespace.UploadFiles
        description: The audited set is the only one a training job should read.
        namespace: "{{ flow.namespace }}"
        filesMap:
          "{{ vars.dataset_path }}": "{{ outputs.build_and_audit.outputFiles.training }}"

      - id: record_manifest
        type: io.kestra.plugin.core.kv.Set
        description: What was published and the audit it passed, for the training job
          and for lineage.
        key: "{{ vars.manifest_kv }}"
        kvType: JSON
        value: |
          {"path": "{{ vars.dataset_path }}", "join": "{{ inputs.join_strategy }}", "embargo_days": {{ inputs.embargo_days }},
           "rows": {{ outputs.audit.values.rows }}, "coverage": {{ outputs.audit.values.coverage }},
           "feature_label_corr": {{ outputs.audit.values.candidate_corr }}, "execution_id": "{{ execution.id }}"}
    else:
      - id: block_publication
        type: io.kestra.plugin.core.execution.Fail
        description: The set would teach the model information it will not have in
          production. Nothing is published.
        errorMessage: >-
          Leakage gate failed for the {{ inputs.join_strategy }} set: {{
          outputs.audit.values.leaked_rows }} rows use a feature newer than the
          label cutoff, correlation {{ outputs.audit.values.candidate_corr }} vs
          {{ outputs.audit.values.reference_corr }} point-in-time (lift {{
          outputs.audit.values.leakage_lift }}, max {{ inputs.max_leakage_lift
          }}), coverage {{ outputs.audit.values.coverage }} (min {{
          inputs.min_coverage }}), {{ outputs.audit.values.duplicate_labels }}
          duplicate labels.

triggers:
  - id: nightly
    type: io.kestra.plugin.core.trigger.Schedule
    description: Rebuild the training set every night before the training job.
    cron: "0 2 * * *"