AggregateExec not cancellable

[pepijnve]

Describe the bug

The streams created by AggregateExec consume their input in a loop in their poll implementations without ever explicitly yielding. If the input stream never returns Pending (which is often the case for file input), Tokio will never have the opportunity to abort the running task.

This is often hidden by the presence of CoalesceExec in a query plan which runs the aggregation in a separate task. CoalesceExec uses RecordBatchReceiverStream which does return Pending itself and as a consequence is cancellable, but dropping/aborting the aggregation task will still not immediately stop it since Tokio can only stop a task when it yields.

To Reproduce

• Start datafusion-cli
• Execute SET datafusion.execution.target_partitions = 1;
• Execute SELECT sum(column) from table; on some relatively large table so that the query takes long enough to run
• Press ctrl-c

Expected behavior

• The running query is cancelled

Additional context

No response

[zhuqi-lucas]

take

[zhuqi-lucas]

Submitted the PR #16196

The PR is limited to solve aggregate with no group streaming, we can extend to more cases if it's not affecting performance?

[alamb]

I believe this sounds similar to from this PR from @jeffreyssmith2nd

• Un-cancellable Query when hitting many large files. #14036

@carols10cents has made a benchmark here:

• Cancellation benchmark #14818

[pepijnve]

Yes this is more or less the same issue. PR #14028 proposed adding a yield point at the leaf of the plan when moving from one file to the next. This PR adds yield points closer to the top of the plan tree just below the AggregateExec's stream by wrapping its input and then yields every 64 input batches. I was wondering if that should be row count or time interval based rather than batch count based.

I found #15314 in the meantime as well. This issue provides one concrete and easily reproducible example of a query that cannot be cancelled.

The comments on PR #14028 regarding Tokio's yield_now are interesting and relevant for PR #16196 as well. Seems like the code pattern should be
I can run some tests to see what the actual behavior is in the ST and MT Tokio runtimes if that helps.

Edit: conclusion in PR #14028 discussion was that calling wake_by_ref was fine.

[pepijnve]

@alamb how are you guys handling general housekeeping in the issue tracker? Should I close this issue since it's a specific example of #15314?

[alamb]

@alamb how are you guys handling general housekeeping in the issue tracker? Should I close this issue since it's a specific example of #15314?

That would be fine.

Alternately, I think we can close this ticket once #16196 is merged too

[pepijnve]

I've updated the issue title and description to be less broad and focus on cancelling aggregations instead. I hope I got the tokio details right.

[zhuqi-lucas]

Yes this is more or less the same issue. PR #14028 proposed adding a yield point at the leaf of the plan when moving from one file to the next. This PR adds yield points closer to the top of the plan tree just below the AggregateExec's stream by wrapping its input and then yields every 64 input batches. I was wondering if that should be row count or time interval based rather than batch count based.

I found #15314 in the meantime as well. This issue provides one concrete and easily reproducible example of a query that cannot be cancelled.

The comments on PR #14028 regarding Tokio's yield_now are interesting and relevant for PR #16196 as well. Seems like the code pattern should be I can run some tests to see what the actual behavior is in the ST and MT Tokio runtimes if that helps.

Edit: conclusion in PR #14028 discussion was that calling wake_by_ref was fine.

Thank you @pepijnve for review, why i was not using row count because we need to calculate batch_size * batch count, we want to not affect the performance for core logic for datafusion, even batch count 64, i am wandering if it will affect the core logic performance when we have huge data.

I am wandering if we can wrapper the yield logic outside the core exec logic in datafusion, such as in the datafusion-cli side if we only want to do the ctril c in datafusion-cli side.

But it seems more cases besides datafusion-cli which want to terminate the streaming, for example the customers who use grpc to terminate:

#14036 (comment)

We can be confident to merge, if we compare the aggregate clickbench to see no performance regression.

[zhuqi-lucas]

Updated no performance regression for the PR with huge aggregate testing:

#16196 (comment)

[pepijnve]

Just for context, I ran into this while working on a Java based application that drives the DataFusion queries. I want to be able to interrupt query execution from the Java side. From the Java side I'm basically calling runtime.block_on(async { record_batch_stream.next().await }). Since I can't use Java's thread interruption mechanism to unblock the next() call I was trying to use tokio::select! with a cancellation token instead and that failed to cancel. This is the same pattern as what the CLI does, just triggered via a different channel than a signal handler.

I still need to write an experiment for this, but based on my reading of the documentation even with a RecordBatchReceiverStream the aggregate would not actually get cancelled. It will look like it has cancelled because the top level task will yield, the JoinSet will get dropped, and the spawned tasks will be aborted. But tokio will only actually be able to terminate the task once the aggregate stream's poll loop yields.

[zhuqi-lucas]

Just for context, I ran into this while working on a Java based application that drives the DataFusion queries. I want to be able to interrupt query execution from the Java side. From the Java side I'm basically calling runtime.block_on(async { record_batch_stream.next().await }). Since I can't use Java's thread interruption mechanism to unblock the next() call I was trying to use tokio::select! with a cancellation token instead and that failed to cancel. This is the same pattern as what the CLI does, just triggered via a different channel than a signal handler.

I still need to write an experiment for this, but based on my reading of the documentation even with a RecordBatchReceiverStream the aggregate would not actually get cancelled. It will look like it has cancelled because the top level task will yield, the JoinSet will get dropped, and the spawned tasks will be aborted. But tokio will only actually be able to terminate the task once the aggregate stream's poll loop yields.

Thank you @pepijnve for this info, i tried the aggregate with more than 1 partition, it seems can be canceled, and it uses RecordBatchReceiverStream inside the CoalescePartitionExec, do you mean it's not the RecordBatchReceiverStream which help the cancellation?