[BUG] leaks possible in the rapids shuffle if batches are received after the task completes

[abellina]

See comment below.

Old explanation:

When there is a limit query, tasks complete early and blocks that are being requested in flight could be leaked.

Update:

In looking through my logs I only see cases where we detected this in failure conditions, where the task was closed forcefully. I now think Spark is doing the right thing, even in limit cases, and draining iterators.

---

The short term solution is to make the RapidsShuffleIterator reject batches as they show up, which would make the RapidsShuffleClient close them on receipt. This can easily be done, but it's wasteful because we are going to be waiting for IO and holding up resources on the peer also when there is no compute to be done for these batches.

A longer term solution will be tracked separately.

[abellina]

I was having a hard time reproducing the issue again. The reason for this looks to be that we have a coalesce batches after an exchange. Since we normally like to get larger batches, they may well hide the fact that the iterator is requesting more than may eventually be consumed.

If we set the batch size small, concat batches will produce more batches and the LocalLimit nodes will see their limits reached, causing the tasks to complete => and us to receive batches after completion.

When you have an ORDER BY followed by LIMIT, this is a TakeOrderedAndProject. This setup needs to look at all the data, so we drain the batches received by the shuffle, and this bug doesn't apply here.