Using the output of a HttpOperator as input for a SQLExecuteQueryOperator when it is executed in chunks

[KennyPeeters]

Dear,
I had a requirement where I had to read data with an HttpOperator and where I had to write the data towards a database. The dataset was too large to be handled in one go, so I had to split it up with top and skip parameters towards the HttpOperator.

But, writing the data from the HttpOperator to the SQLExecuteQueryOperator was not that straightforward as first thought.

I tired the following approaches:

• task_group: The HttpOperator and SQLExecuteQueryOperator went into 9 mapped tasks without any relation with each other whlie "get_data" >> "write_data" had to be combined
• Normal python function: This does not have the partial and expand functionality
• python-function called by PythonOperator: If I am not mistaking you can't return a taskgroup with other operators within it

This has led me to the only option I could think of: Creating a custom HttpToDBOperator that takes the needed arguments to get the data from HttpOperator and write that dataset to the database with SQLExecuteQueryOperator.

Maybe not the most clean approach but it worked, but I am still wondering:

Is there an option I did not think of?

[dabla]

I understand the issue that Kenny has, as we also encountered that "problem" many times. The thing is that Airflow will always execute one mapped operator before another one, even if it's in a task group. At first when I saw a task group I expected it would execute the operators one after another for each mapped input, but that's not the case. What we would like to achieve is that a mapped task_group for example would execute all operators sequentially in the task_group for each input, and not execute all inputs to the first operator and then apply all outputs of that first operator to execute the second one. What would be even cooler is that those operators could be like chained and executed on the same worker, (thus avoiding the xcom communication between them so that the second operator would be able to use the input of the previous one) -> that last sentence is incorrect we could not avoid the xcoms I think.

As a temporary solution I would think of creating a custom SequentialOperator and pass the operators as parameters to that operator so that the sequential operator would then execute those operators, so the execute method of the SequentialOperator would then call the execute methods of the passed operators (I know there is more to it, it's just a simple explanation). "Problem" with this approach is that we would only see the SequentialOperator in the Graph, if it's a problem at all, but it would be nicer if we would have that option build-in in Airflow, like for example adding a boolean parameter sequential (which by default would be False of course) to the task_group and then the TaskGroup would know it has the execute all operators sequentially for each mapped task.

I'm willing to contribute for this but I'm not sure where I have to check this and if this would be accepted by the community as a possible feature? I've already looked at the TaskGroup class but don't know how I could achieve this logic, I suppose there is more to it than just the TaskGroup.

Also started a thread for this on slack.

[hussein-awala]

Could you provide some reproducible code? I believe that dynamic mapped task group is what your looking for.

Here are some of my comments on @dabla slack thread

No, the current behavior is exactly what you describe in the alternate behavior, the mapped flows will be independent, if you have any slowness in op1 in one of the mapped streams, the rest can be executed and terminated without any blocking

However, if you expand each task independently, you will have what you describe in first behavior, and this was the motivation to support expanding a task group

You can check https://cwiki.apache.org/confluence/plugins/servlet/mobile?contentId=191336946#content/view/191336946 where we distinguish between breadth-first execution in a sequence of mapped tasks and depth-first execution for mapped task group

I invite you to check the confluence page where we discussed the dynamic mapped task group feature.

[dabla]

So this example is not possible from the link above as the TaskGroupFactory only has the expand and expand_kwargs method I've checked the source code from 2.8:

lines = ["foo", "bar","baz"]
@task_group
def my_task_group(my_arg_1: str):
   task_1 = DummyOperator(task_id="task_1")
   task_2 = BashOperator(task_id="task_2", bash_command=f"echo ${section_1.my_arg_1}")
   task_3 = DummyOperator(task_id="task_3")
   task_1 >> [task_2, task_3]

mapped_task_group = my_task_group.map(my_arg_1=lines)

[dabla]

This is also the answer from ask astronomer:

To achieve depth-first execution with a task group in Apache Airflow, you can use the expand method on a @task_group-decorated function. This allows you to create a mapped task group where all tasks in the group are expanded together against the same inputs. Here's an example of how you can define and use a task group for depth-first execution:
from airflow.decorators import task, task_group

@task_group
def file_transforms(filename):
    converted = convert_to_yaml(filename)
    return replace_defaults(converted)

file_transforms.expand(filename=["data1.json", "data2.json"])

In this example, the file_transforms task group is expanded into two instances, one for each file. The tasks convert_to_yaml and replace_defaults within the group will each become two instances at runtime, and they will execute in a depth-first manner. This means that replace_defaults for "data1.json" can run before convert_to_yaml for "data2.json" is done, without having to wait for all instances of convert_to_yaml to complete.
Remember that task mapping within a mapped task group is not permitted, so you cannot have nested task mappings inside a mapped task group. If you need to perform branching or looping logic within a task group, you should use appropriate tasks like @task.branch for branching conditions and task mapping methods for loops.
If you have a more complex use case or need further assistance, please provide additional details about your specific workflow and requirements.

[dabla]

Here is an example of the DAG run which contains a task_group named "process_workspaces_info" on which we expand each workspace we want to process, as you can see it first executes the "get_workspace_info" task and starts scheduling taks for "check_workspaces_status" but those don't start yet, probably because it doesn't have enought workers which is logical as there are 97 mapped tasks, we don't have 97 workers available. Also we are using the CeleryKubernetesExecutor.

[dabla]

This is part of the code regarding the process_workspaces_info:

        @task_group
        def process_workspaces_info(data):
            workspaces_info_task = MSGraphAsyncOperator(
                task_id="get_workspace_info",
                conn_id="powerbi_api",
                url="myorg/admin/workspaces/getInfo",
                method="POST",
                query_parameters={
                    "lineage": True,
                    "datasourceDetails": True,
                    "datasetSchema": True,
                    "datasetExpressions": True,
                    "getArtifactUsers": True,
                },
                result_processor=lambda context, response: {"scanId": response["id"]},
                data=data,
            )

            check_workspaces_status_task = MSGraphSensor(
                task_id="check_workspaces_status",
                conn_id="powerbi_api",
                url="myorg/admin/workspaces/scanStatus/{scanId}",
                path_parameters=workspaces_info_task.output,
                result_processor=lambda context, result: {"actionId": result["id"]},
                timeout=350.0,
            )

            scan_workspaces_status_task = MSGraphAsyncOperator(
                task_id="scan_workspaces_status",
                conn_id="powerbi_api",
                url="myorg/admin/workspaces/scanResult/{actionId}",
                path_parameters=check_workspaces_status_task.output,
            )

            persist_workspaces_info_task = SQLRowsInsertOperator(
                task_id="persist_workspaces_info",
                conn_id=HANA_CONN_ID,
                table_name="table_name",
                commit_every=10,
                rows=scan_workspaces_status_task.output,
                rows_processor=to_workspaces,
                executemany=True,
            )

            persist_workspaces_datasource_instances_task = SQLRowsInsertOperator(
                task_id="persist_workspaces_datasource_instances",
                conn_id=HANA_CONN_ID,
                table_name="table_name",
                commit_every=10,
                rows=scan_workspaces_status_task.output,
                rows_processor=to_datasource_instances,
                executemany=True,
            )

            workspaces_info_task >> check_workspaces_status_task >> scan_workspaces_status_task >> persist_workspaces_info_task >> persist_workspaces_datasource_instances_task

        workspaces_task >> process_workspaces_info.expand(data=workspaces_task.output)

[potiuk]

I'm willing to contribute for this but I'm not sure where I have to check this and if this would be accepted by the community as a possible feature? I've already looked at the TaskGroup class but don't know how I could achieve this logic, I suppose there is more to it than just the TaskGroup.

Yes. If someone proposes it, that would be a good idea, but it's incredibly complex to implement taking into account that each task is executed separately and there is basically no execution priority no task affinity concepts in Airflow at all. So you would have to introduce all those concepts to Airflow to achieve it.

But in this case you can also implent your DAG using the knowledge of the parallelism you have - possibly - rather than implementing incredibly complex mechanism for scheduling to figure out how and when to prioritise tasks you can use the features that are there that make it possible for DAG author to "help the system quite a bit by adding limits:

https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/dynamic-task-mapping.html#placing-limits-on-mapped-tasks

Surely if you have 30 workers and put 100 mapped tasks to them, it's easy to imagine the "first" instances will take all the slots because they are eligible to execution way before the next tasks are. When you think about your "future" idea about the scheduled sequential execution, you would have to - basically automatically decide when not to schedule new tasks because they might be needed for the follow-up tasks, or implement affinity so - that when one task completes. the same task will automatically pick up some specific task. None of the concepts exist in Airfllow today.

Usually this comes down to trade-offs about utilisation.

But if you narrow it down to your case and if you smartly limit how many parallel instances of a task you can have, and have many more workers than that, you can achieve desired outcome likely:

Limiting parallel copies of a mapped task

If you wish to not have a large mapped task consume all available runner slots you can use the max_active_tis_per_dag setting on the task to restrict how many can be running at the same time.

Note, however, that this applies to all copies of that task against all active DagRuns, not just to this one specific DagRun.

Of course it's not solving all the possible complexities, but applying this to your specific case by DAG author is probably several orders of magnitude less complex than inventing a complex scheduling to handle the case, if you ask me.

But if you think you can propose and impement the more complex and geneic scheduling solution - surely writing AIP is a good start.

[KennyPeeters]

The scope of my question existed out of following steps:

• Retrieving a total number of records using an HttpOperator
• Calculate a number of chunks to be executed based on the total number and a (self defined) limit size
• Perform the actions: "load_data >> write_data" preferably executed for the calculated chunks (in my case 9) in that sequence

The issue I had:

• In a default setup the DAG would load the data for the 9 chunks and then write the data for the 9 chunks

While I kept on trying to get my use-case to work I also tested the pointers from this discussion, but I got it working on my own approach:

Above screenshot shows:

• get_data_1 is success, temporary_wait_..._1 is success and write_data_1 is scheduled
• get_data_3 is success and temporary_wait_..._3 is running

This exactly does what I was searching for and what I wanted to achieve with my "custom operator" where I reused HttpOperator and SQLExecuteQueryOperator. The differency is that, if I change it to those operators, it is now also shown in the UI.

Possible downsides/attention points that pops to mind:

• What will happen if I have loaded 10k task groups?
• Can the installation handle this?
• Is this better in performance or not (it's not)?

Anyway, wanted to share my findings and below code which reproduces this behavior:

    @task_group
    def create_dynamic_task_groups():
        for iteration in range(int(5)):
            @task(trigger_rule="one_success")
            def get_data():
                print("Getting data from ...")

            @task
            def temporary_wait_time_to_show_sequential_execution():
                # Generate a random sleep interval between 1 and 5 seconds
                sleep_interval = random.uniform(1, 5)

                # Sleep for the generated interval
                time.sleep(sleep_interval)

                # Continue with your code execution after the sleep
                print("Finished sleeping for", sleep_interval, "seconds")

            @task
            def write_data():
                print("Write data to ...")

            data_retrieval = get_data()
            wait = temporary_wait_time_to_show_sequential_execution()
            data_write = write_data()

            data_retrieval >> wait >> data_write

    dynamic_task_groups = create_dynamic_task_groups()

The code shows a "static" number of task-group creations instead of a dynamic task-group creation and this was actually the most difficult part as I was not able to access the xcom-value. I solved this by adding an airflow-variable that saved my number of iterations.

One problem I had with this, but I am sure it can be solved, is that the airflow-variable has to exist and that you cannot do a set in the DAG as this state is saved. This can also be because of our setup but, even though it was interesting to try it out, I did
not pursue it more as there are other priorities as well.

But, as David mentioned, I am also willing to collaborate on this topic.

[potiuk]

As mentioned above - the "sequential" execution is not the result of dynamic vs. static task groups. It's the matter of how many workers you hae and how many tasks are eligible to run. The solution with limiting number of running task instances should to allow you to control parallelism/sequentialism manually. And this is not really a problem of collaborating on it, it's more on comingup with the whole idea - basically from the scratch on how to do it (and whether to do it) - because that would mean some fundamental changes to how Airlfow queues, schedules and executes tasks.

[KennyPeeters]

@hussein-awala, I have found your example on the Airflow GitHub today and it is basically almost the same as what I had found yesterday. The only extra requirement I had, but already stepped away from it, was the "dynamic" number of amount of task groups