2

Is it possible to retrieve multiple values from XCOM, pushed by a single task but with different keys?

I think I've seen examples to this:

# pulls one value
pulled_value = ti.xcom_pull(key=None, task_ids='my_task_id')

and to this:

# pulls multiple values but from multiple tasks
pulled_value_1, pulled_value_2 = ti.xcom_pull(key=None, task_ids=['my_task_id_1', 'my_task_id_2'])

What I need is would possibly look like this:

# pulls multiple values but from a single
pulled_value_1, pulled_value_2 = ti.xcom_pull(key=['my_key_1', 'my_key_2'], task_ids='my_task_id')

I can't find this in the documentation.
Is this even possible?
If yes, it makes a single database query in the background, or just repeats a single query multiple times? If not, how could I get similar behavior?

elaspog
  • 1,635
  • 3
  • 21
  • 51
  • I've checked the current master branch (https://github.com/apache/airflow/blob/main/airflow/models/xcom.py), and I did not find any possible function to handle this requirement, although in some use cases it could reduce the load on database. Would it be useful to add a feature like this? – elaspog Apr 10 '22 at 06:43

1 Answers1

3

Answering your questions:

  1. There is no such feature. xcom_pull accepts task_ids: Optional[Union[str, Iterable[str]]] but with the same key. You can open a PR to Airflow for adding the functionality you seek.

  2. As for number of queries: I assume that by "repeats a single query" you are asking if it execute a query per task_id. The answer is No. Airflow did this optimization in PR. In the source code you can see that xcom_pull() uses Xcom.get_many() and get_many creates a filter using IN which allows to compare against multiple values. You can see the relevant code lines here.

Elad Kalif
  • 14,110
  • 2
  • 17
  • 49
  • Thanks for your answer! In the example above I've mentioned `task_ids: Optional[Union[str, Iterable[str]]]` (see the example code `# pulls multiple values but from multiple tasks`) but this is not what I was looking for... I'm looking for a method to read multiple keys at once written to the XCOM by a single task. For any `Xcom.get_many()` method I could not find a solution where it expects a list of keys, like this: `key: Optional[Union[str, Iterable[str]]] = XCOM_RETURN_KEY`. – elaspog Apr 10 '22 at 07:27
  • 1
    I misread sorry.. edited. – Elad Kalif Apr 10 '22 at 07:28