Repository: apache/airflow
Pull request: Make xcom_pull results order deterministic
closes https://github.com/apache/airflow/issues/11858

<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.

Feel free to ping committers for the review!

In case of existing issue, reference it using one of the following:

closes: #ISSUE
related: #ISSUE

How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->

---
**^ Add meaningful description above**

Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).


Problem statement:
Xcom_pull results order is not deterministic
For xcoms pushed from different tasks, having the same execution_date and in the same Dag, the result of the xcom_pull with key set as None is not deterministic in its order

**Apache Airflow version**: 2.0


**Kubernetes version (if you are using kubernetes)** (use `kubectl version`):

**Environment**:

- **Cloud provider or hardware configuration**:
- **OS** (e.g. from /etc/os-release):
- **Kernel** (e.g. `uname -a`):
- **Install tools**:
- **Others**:

**What happened**:
When tasks in the same dag push Xcoms, the result of the xcom_pull comes in any order. The order is not deterministic.

**What you expected to happen**:

I expected the order to be deterministic.

**How to reproduce it**:

Run the provided example_xcom dag in the repository multiple times

Hints from issue discussion:
Hmm, can you elaborate more on this @ephraimbuddy . For example if the XCom hasn't been pushed yet, you will get None. So the order of task that pushes XCom value does matter
This example dag below gives inconsistent results:

```
#
# 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 demonstrating the usage of XComs."""
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago

dag = DAG(
    'example_xcom',
    schedule_interval="@once",
    start_date=days_ago(2),
    default_args={'owner': 'airflow'},
    tags=['example']
)

value_1 = [1, 2, 3]
value_2 = {'a': 'b'}


def push(**kwargs):
    """Pushes an XCom without a specific target"""
    kwargs['ti'].xcom_push(key='value from pusher 1', value=value_1)


def push_by_returning(**kwargs):
    """Pushes an XCom without a specific target, just by returning it"""
    return value_2


def puller(**kwargs):
    """Pull all previously pushed XComs and check if the pushed values match the pulled values."""
    ti = kwargs['ti']

    # get value_1
    pulled_value_1 = ti.xcom_pull(key=None, task_ids='push')
    if pulled_value_1 != value_1:
        raise ValueError(f'The two values differ {pulled_value_1} and {value_1}')

    # get value_2
    pulled_value_2 = ti.xcom_pull(task_ids='push_by_returning')
    if pulled_value_2 != value_2:
        raise ValueError(f'The two values differ {pulled_value_2} and {value_2}')

    # get both value_1 and value_2
    pulled_value_1, pulled_value_2 = ti.xcom_pull(key=None, task_ids=['push', 'push_by_returning'])
    if pulled_value_1 != value_1:
        raise ValueError(f'The two values differ {pulled_value_1} and {value_1}')
    if pulled_value_2 != value_2:
        raise ValueError(f'The two values differ {pulled_value_2} and {value_2}')


push1 = PythonOperator(
    task_id='push',
    dag=dag,
    python_callable=push,
)

push2 = PythonOperator(
    task_id='push_by_returning',
    dag=dag,
    python_callable=push_by_returning,
)

pull = PythonOperator(
    task_id='puller',
    dag=dag,
    python_callable=puller,
)

pull << [push1, push2]
```
I checked the code, the xcom_pull is ordered both with execution_date and timestamp.
However, for the above DAG, the tasks have the same execution_date and may also have the same timestamp. In this case, the order of the xcom_pull result is not consistent.
I was able to reproduce this
I am able to reproduce and i checked the query and database entries and found following :

Query executed in order to fetch rows

SELECT xcom.key AS xcom_value 
FROM xcom 
WHERE xcom.task_id IN ('push', 'push_by_returning') AND xcom.dag_id = 'example_xcom' AND xcom.execution_date = '2020-11-01 00:00:00.000000' 
ORDER BY xcom.execution_date DESC, xcom.timestamp DESC;

![ScreenShoot](https://user-images.githubusercontent.com/11897651/97966276-a18ab000-1de1-11eb-9169-924174981b6f.png)

As shown in database entries in above table For execution Date 2020-11-01 and  2020-11-03  timestamps assigned to 'return_value' and 'value from pusher 1' are in different order.

Because of 'xcom.timestamp DESC;' in the query task with execution date '2020-11-03' is returning values in different order and gets failed.




To be honest I'm not sure why we order by `timestamp` - do we expect users to push same key (key, exec date, task_id, dag_id) few times in operator?

Imho the natural way to sort it would be to sort it by `key` column - the `execution_date` is already same for all entries if we do `xcom_pull`:
https://github.com/apache/airflow/blob/45ae145c25a19b4185c33ac0c4da920324b3987e/airflow/models/taskinstance.py#L1830-L1837

Hey @ephraimbuddy @turbaszek -> is it something we want to merge before 2.0.0rc1 or should we change the milestone?
> Hey @ephraimbuddy @turbaszek -> is it something we want to merge before 2.0.0rc1 or should we change the milestone?

I think it would be nice to have a deterministic result. However, I'm not sure if this broke something or not. Definitely, it's not something crucial

Commit message: Make Xcom_pull results order deterministic

closes https://github.com/apache/airflow/issues/11858

Patch to review:
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 82bfc6647734f..da646094082fa 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1883,16 +1883,24 @@ def xcom_pull(  # pylint: disable=inconsistent-return-statements
             task_ids=task_ids,
             include_prior_dates=include_prior_dates,
             session=session,
-        ).with_entities(XCom.value)
+        )
 
         # Since we're only fetching the values field, and not the
         # whole class, the @recreate annotation does not kick in.
         # Therefore we need to deserialize the fields by ourselves.
-
         if is_container(task_ids):
-            return [XCom.deserialize_value(xcom) for xcom in query]
+            vals_kv = {
+                result.task_id: XCom.deserialize_value(result)
+                for result in query.with_entities(XCom.task_id, XCom.value)
+            }
+
+            vals = []
+            for task_id in task_ids:
+                vals.append(vals_kv.get(task_id))
+
+            return vals
         else:
-            xcom = query.first()
+            xcom = query.with_entities(XCom.value).first()
             if xcom:
                 return XCom.deserialize_value(xcom)
 
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index ab9040d8adff6..0fd8a85cd4942 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1030,9 +1030,9 @@ def test_xcom_pull(self):
         # Pull the value pushed by the second task
         result = ti1.xcom_pull(task_ids='test_xcom_2', key='foo')
         self.assertEqual(result, 'baz')
-        # Pull the values pushed by both tasks
+        # Pull the values pushed by both tasks & Verify Order of task_ids pass & values returned
         result = ti1.xcom_pull(task_ids=['test_xcom_1', 'test_xcom_2'], key='foo')
-        self.assertEqual(result, ['baz', 'bar'])
+        self.assertEqual(result, ['bar', 'baz'])
 
     def test_xcom_pull_after_success(self):
         """


Code review comments:
diff_hunk: @@ -1883,16 +1883,24 @@ def xcom_pull(  # pylint: disable=inconsistent-return-statements
             task_ids=task_ids,
             include_prior_dates=include_prior_dates,
             session=session,
-        ).with_entities(XCom.value)
+        )
 
         # Since we're only fetching the values field, and not the
         # whole class, the @recreate annotation does not kick in.
         # Therefore we need to deserialize the fields by ourselves.
-
         if is_container(task_ids):
-            return [XCom.deserialize_value(xcom) for xcom in query]
+            vals_kv = {
+                result.task_id: XCom.deserialize_value(result)
+                for result in query.with_entities(XCom.task_id, XCom.value)
+            }
+
+            vals = []
+            for task_id in task_ids:
+                vals.append(vals_kv.get(task_id))
+
+            return vals
path: airflow/models/taskinstance.py
text: @user1:
```suggestion
            values_ordered_by_id = [vals_kv.get(task_id) for task_id in task_ids]
            return values_ordered_by_id
```