AIP-76 Implement PartitionAtRuntime

[uranusjr]

Body

AIP-76

Committer

• I acknowledge that I am a maintainer/committer of the Apache Airflow project.

[anishgirianish]

Hi, I'd like to pick this up. I've been working on AllowedKeyMapper (#61931) and ProductMapper (#61937), and want to check my approach before writing code.

I think #58474 and PartitionAtRuntime should be tackled together since per-event partition key overrides are the transport mechanism for runtime-discovered keys.

Proposed approach: Extend OutletEventAccessor with an optional partition_key field so tasks can set per-event keys. _serialize_outlet_events would carry it alongside extra, and register_asset_changes_in_db would use it instead of falling back to dag_run.partition_key. This solves both #58474 and PartitionAtRuntime without a new Execution API endpoint. Does that seem right?

Two questions:

1. Should add_partition() allow arbitrary keys or only from a predefined set? The AIP example suggests arbitrary, but AllowedKeyMapper could optionally validate them. Worth adding?

2. PartitionMapper is consumer-side (to_downstream maps keys). PartitionAtRuntime is producer-side (creates keys). Should it be a separate abstraction, or still a PartitionMapper subclass? Given user-defined mappers aren't public for 3.2 (Implement rollup / windowing #59294), should this be internal-only for now?

Thoughts?

cc @uranusjr @Lee-W @dstandish

[Lee-W]

1. Should add_partition() allow arbitrary keys or only from a predefined set? The AIP example suggests arbitrary, but AllowedKeyMapper could optionally validate them. Worth adding?

arbitrary keys seem to make more sense to me and defer to AllowedKeyMapper to do its job

PartitionMapper is consumer-side (to_downstream maps keys). PartitionAtRuntime is producer-side (creates keys). Should it be a separate abstraction, or still a PartitionMapper subclass? Given user-defined mappers aren't public for 3.2 (#59294), should this be internal-only for now?

I think it'll be part of 3.3. We're in a stage of keeping 3.2 as stable as possible. We've not yet looked too deeply into the runtime thing, but feel free to propose your idea. and we can see how we can get it better together :) probably some manipulation in side task execution, but the interface needs some design

[Lee-W]

  # 1. TaskFlow
  @task(outlets=[my_asset])
  def process(*, outlet_events):
      outlet_events[my_asset].partition_key = compute_partition_key()

  # 2. @asset
  @asset(schedule=CronPartitionTimetable("0 * * * *", timezone="UTC"))
  def my_asset(self):
      self.partition_key = compute_partition_key()

  # 3. @asset.multi
  @asset.multi(outlets=[asset_a, asset_b])
  def my_multi_asset(*, outlet_events):
      outlet_events[asset_a].partition_key = compute_partition_key()

  # 4. Traditional operator
  class MyOperator(BaseOperator):
      def execute(self, context):
          context["outlet_events"][my_asset].partition_key = compute_partition_key()

We probably could have an interface like this? And we can have

[anishgirianish]

@Lee-W this sketch is really helpful, the consistency across TaskFlow, @asset, @asset.multi, and traditional operators is nice, and reusing the existing outlet_events accessor feels much cleaner than a separate PartitionAtRuntime(name=...) object as the AIP suggests. +1 from me on that direction.

One thing I wanted to run by you: all four forms in your sketch set a single partition key per run, which handles the override case beautifully (and resolves #58474 on the way). The AIP's motivating example, though, is the runtime fan-out looping over a discovered set and emitting one event per partition in a single task run:

@task(outlets=[my_asset])
def process(*, outlet_events):
    for customer in get_customers():
        outlet_events[my_asset].add_partition(key=customer, extra={"rows": ...})

One direction that composes cleanly with your sketch: add an add_partition(key=, extra=) method on the accessor alongside .partition_key = .... Under the hood, OutletEventAccessor would hold a list of pending (partition_key, extra) events assigning .partition_key / .extra mutates a default single event, add_partition appends. Today's single-event behavior is preserved as the default, and the same extension lands naturally on all four of your forms (self.add_partition(...) inside a @asset body, etc.).

One small AIP constraint worth surfacing somewhere in the final design: "Runtime partitioning cannot be combined with statically defined partitioning with PartitionByProduct."

Curious whether add_partition on the accessor matches where you were headed, or whether you had a different shape in mind for the fan-out case? If this direction resonates, I'd love to take a crack at the PR and bring it back to you for review 🙏