[ARROW-10577] [Rust][DataFusion] Hash Aggregator stream finishes unexpectedly after going to Pending state

• Type: Bug
• Status: Resolved
• Priority: Blocker
• Resolution: Fixed
• Affects Version/s: 2.0.0
• Fix Version/s: 3.0.0
• Component/s: Rust, Rust - DataFusion
• Labels: pull-request-available
• External issue URL: https://github.com/apache/arrow/issues/26541
• Assignee: Rémi Dettai
• Reporter: Rémi Dettai
• Created: 13/Nov/20 08:19
• Updated: 11/Jan/23 08:14
• Resolved: 16/Nov/20 13:26

Description

This happens when executing a DataFusion query plan with hash aggregation where the data source is not ready on the first call by the Executor, and the async state machine is passed to a pending state

In the Stream implem of GroupedHashAggregateStream and HashAggregateStream, the state is set to self.finished = true on the first call to poll_next(). If the inner stream is Poll::Pending on the first call, this means that the next call resolves to Poll::Ready(None), thus finishing the stream instead of actually consuming the inner data.

I think that it does not happen with most current sources because they never trigger the Poll::Pending state. Parquet is implemented with a blocking call inside poll_next() (which is also problematic but an other issue), Memory yields directly, and CSV also always yields Poll::Ready

An analysis should be performed on all physical plans to check if the issue occurs in other places.

Issue Links

• links to: GitHub Pull Request #8658