From e3fe3f0d048bc0ed58b4079b4eed25633bd9bed9 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Wed, 17 Jun 2026 17:54:09 +0530 Subject: [PATCH 1/2] Adding example dag for task state store with mapped tasks (cherry picked from commit 0112455111834a52de6824be7fc8084def992bcd) --- .../example_task_state_store_mapped.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py diff --git a/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py b/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py new file mode 100644 index 0000000000000..743f24e849bab --- /dev/null +++ b/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py @@ -0,0 +1,62 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Example DAG with mapped tasks to demonstrate task state store isolation per map_index.""" + +from __future__ import annotations + +import random +from datetime import datetime, timezone + +from airflow.sdk import DAG, task + +TABLES = ["orders", "customers", "products"] + +with DAG( + dag_id="example_task_state_store_mapped", + schedule=None, + start_date=datetime(2026, 1, 1), + catchup=False, + tags=["example", "task-state-store"], + doc_md=__doc__, +) as dag: + + @task + def get_tables() -> list[str]: + """Return the list of tables to process.""" + return TABLES + + @task + def process_table(table: str, task_state_store=None, ti=None) -> dict: + """Process one table — each mapped instance gets its own task state.""" + row_count = random.randint(100, 10000) + result = { + "table": table, + "map_index": ti.map_index, + "row_count": row_count, + "processed_at": datetime.now(tz=timezone.utc).isoformat(timespec="seconds"), + } + + task_state_store.set("table", table) + task_state_store.set("status", "complete") + task_state_store.set("row_count", row_count) + task_state_store.set("result", result) + + print(f"[map_index={ti.map_index}] Processed {table}: {row_count} rows") + return result + + tables = get_tables() + process_table.expand(table=tables) From 69afb83b94a1ba59d795933c211b6dd4c5ee3e90 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Wed, 17 Jun 2026 19:07:04 +0530 Subject: [PATCH 2/2] removed uneccessary fields --- .../airflow/example_dags/example_task_state_store_mapped.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py b/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py index 743f24e849bab..3fd8778af87c5 100644 --- a/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py +++ b/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py @@ -49,10 +49,7 @@ def process_table(table: str, task_state_store=None, ti=None) -> dict: "row_count": row_count, "processed_at": datetime.now(tz=timezone.utc).isoformat(timespec="seconds"), } - - task_state_store.set("table", table) task_state_store.set("status", "complete") - task_state_store.set("row_count", row_count) task_state_store.set("result", result) print(f"[map_index={ti.map_index}] Processed {table}: {row_count} rows")