Filtering items from a mapped task is broken

[GeryDeKocliko]

Apache Airflow version

3.0.0

If "Other Airflow 2 version" selected, which one?

No response

What happened?

Hi,

As outlined in the documentation, I’m attempting to filter out items in my dynamic task mapping by returning None within the filter task.

See : https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/dynamic-task-mapping.html#filtering-items-from-a-mapped-task

However, I’m observing that None values are still being passed to the subsequent expand, which is not the expected behavior.

What you think should happen instead?

The None value must not be passed to the next expand.

How to reproduce

from airflow.decorators import dag, task

@task
def generate_data():
    return [1, None, 3]


@task
def filter_item_from_mapped_task(item):
    if item is None:
        return None
    return item


@task
def print_item(item):
    assert item is not None
    print(item)


@dag(
    dag_id='test_bug_dag',
)
def test_bug_dag():
    data = generate_data()
    filtered_data = filter_item_from_mapped_task.expand(item=data)
    print_item.expand(item=filtered_data)

test_bug_dag()

Operating System

MacOS 14.6 (23G80)

Versions of Apache Airflow Providers

I haven’t added any additional providers

Deployment

Docker-Compose

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[tirkarthi]

Interesting, this works for [None] and [1, None] to run only one task but when None value is in the middle this error is present. I guess the map_index of 0 and 2 have xcom pushed but the underlying task seems 2 upstream tasks and tries to get map_index for 0 and 1 which causes 1 to have no XCom. I am checking on this.

[tirkarthi]

When the xcom values map 1:1 in length then the map_index of the downstream task can directly get the xcom value from upstream value using map_index of the downstream task. But when there is a difference like upstream has 2 mapped task where None and 1 is returned for map_index=0 and map_index=1 respectively then the downstream task only has to process value of upstream task at map_index=1 but the downstream task starts with map_index=0 and goes on to again fetch xcom value for map_index=0 which is None. I guess some change has to be made at below code that gets the index to search in LazyXComSequence. But I guess this is a high priority issue for the next release since filter doesn't seem to be working on basic use cases.

cc: @ashb for help since the dynamic task mapping was implemented by Ash and will have better idea.

airflow/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py

Lines 165 to 181 in 6a3bc01

	def _expand_mapped_field(self, key: str, value: Any, map_index: int, all_lengths: dict[str, int]) -> Any:
	def _find_index_for_this_field(index: int) -> int:
	# Need to use the original user input to retain argument order.
	for mapped_key in reversed(self.value):
	mapped_length = all_lengths[mapped_key]
	if mapped_length < 1:
	raise RuntimeError(f"cannot expand field mapped to length {mapped_length!r}")
	if mapped_key == key:
	return index % mapped_length
	index //= mapped_length
	return -1
	found_index = _find_index_for_this_field(map_index)
	if found_index < 0:
	return value
	if isinstance(value, Sequence):
	return value[found_index]

[tirkarthi]

I just tried it in Airflow 2 and 3 to see the difference. The index calculation is same in both Airflow 2 and 3 but the usage varies as below

In Airflow 2, Accessing LazySequence[0] goes on to get all the upstream task id xcom values and add offset for the calculated index and limit the result to 1

airflow/airflow-core/src/airflow/utils/db.py

Lines 1568 to 1576 in 36c8975

	def __getitem__(self, key: int | slice) -> T | Sequence[T]:
	if isinstance(key, int):
	if key >= 0:
	stmt = self._select_asc.offset(key)
	else:
	stmt = self._select_desc.offset(-1 - key)
	if (row := self._session.execute(stmt.limit(1)).one_or_none()) is None:
	raise IndexError(key)
	return self._process_row(row)

In Airflow 3, Accessing LazyXComSequence[0] goes on to get the upstream task id xcom values and uses the calculated index as map_index to the upstream which is different compared to Airflow 2

airflow/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py

Lines 151 to 157 in e2c9f9b

	def _get_item(self, index: int) -> T:
	# TODO: maybe we need to call SUPERVISOR_COMMS manually so we can handle not found here?
	return self._ti.xcom_pull(
	task_ids=self._xcom_arg.operator.task_id,
	key=self._xcom_arg.key,
	map_indexes=index,
	)