apache / apache/arrow

[C++][Acero] Scheduled tasks in SourceNode should be receptive to backpressure

Open
#47,641 0 comments 0 reactions 0 assignees View on GitHub
Component: C++ Type: bug
Dominant language
C++
Stars
17.1k
Forks
4.3k
Avg merge
3d 13h
Merged PRs (30d)
88

Description

### Describe the bug, including details regarding any error messages, version, and platform.

#### Current Behavior

The `SourceNode` is affected by backpressure. This means currently, that a `SourceNode` stops scheduling tasks once `PauseProducing` is invoked and awaits a future. Once that future is marked finished, execution continues. This means that any job that has been scheduled is not affected by backpressure.

#### Example

My `SourceNode` gets data, sends it to the `MapNode` which decompresses the data, and sends it to the `SinkNode` that streams it out.

The `generator` in my `SourceNode` creates an `ExecBatch` without awaiting any other futures:

```
arrow::Future> operator()() {
try {
return produceNextBatch();
} catch (const std::exception& exception) {
return arrow::Status::ExecutionError(exception.what());
}
};
```

Here, I noticed that this `SourceNode` generates and schedules all batches for propagation through the `ExecPlan`, *before* the first result is ever consumed.

Backpressure that eventually gets applied by the `SinkNode` after receiving the first few batches doesn't stop any of these scheduled tasks from starting. Once these start, they will be decompressed and will create a large memory spike.

I hoped that backpressure could prevent this memory spike.

#### Why does this happen

The limitation of the current backpressure implementation is that the `SourceNode` only checks backpressure when scheduling..

```
return generator_().Then(
[this](
const std::optional& morsel_or_end) -> Future> {
std::unique_lock lock(mutex_);
if (IsIterationEnd(morsel_or_end) || stop_requested_) {
return Break(batch_count_);
}
lock.unlock();
SliceAndDeliverMorsel(*morsel_or_end);
lock.lock();
if (!backpressure_future_.is_finished()) {
EVENT_ON_CURRENT_SPAN("SourceNode::BackpressureApplied");
return backpressure_future_.Then(
[]() -> ControlFlow { return Continue(); });
}
return Future>::MakeFinished(Continue());
}
```

.. not when the scheduled tasks are eventually executed in `SliceAndDeliverMorsel`:

```
plan_->query_context()->ScheduleTask(
[this, morsel_length, use_legacy_batching, initial_batch_index, morsel,
has_ordering = !ordering_.is_unordered()]() {
int64_t offset = 0;
int batch_index = initial_batch_index;
do {
int64_t batch_size =
std::min(morsel_length - offset, ExecPlan::kMaxBatchSize);
// In order for the legacy batching model to work we must
// not slice batches from the source
if (use_legacy_batching) {
batch_size = morsel_length;
}
ExecBatch batch = morsel.Slice(offset, batch_size);
UnalignedBufferHandling unaligned_buffer_handling =
plan_->query_context()->options().unaligned_buffer_handling.value_or(
GetDefaultUnalignedBufferHandling());
ARROW_RETURN_NOT_OK(
HandleUnalignedBuffers(&batch, unaligned_buffer_handling));
if (has_ordering) {
batch.index = batch_index;
}
offset += batch_size;
batch_index++;
ARROW_RETURN_NOT_OK(output_->InputReceived(this, std::move(batch)));
} while (offset < morsel.length);
return Status::OK();
},
"SourceNode::ProcessMorsel");
```

#### Possible mitigation

The solution could be to make the lambda scheduled by the `SourceNode` susceptible to backpressure by awaiting the corresponding future `backpressure_future_`.

I also saw that there exists a `ThrottledAsyncTaskScheduler` that performs similar tasks, but using this might be dangerous regarding the introduction of dead-locks.

I would also be interested whether there are other ways to circumvent the memory spikes in such cases.

### Component(s)

C++

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.