GH-31769: [C++][Acero] Add spilling for hash join

cpp/src/arrow/CMakeLists.txt

@@ -406,6 +406,8 @@ if(ARROW_COMPUTE)
  compute/exec/query_context.cc
  compute/exec/sink_node.cc
  compute/exec/source_node.cc
+ compute/exec/spilling_util.cc
+ compute/exec/spilling_join.cc
  compute/exec/swiss_join.cc
  compute/exec/task_util.cc
  compute/exec/tpch_node.cc

cpp/src/arrow/compute/exec/CMakeLists.txt

@@ -37,6 +37,7 @@ add_arrow_compute_test(asof_join_node_test
  "arrow-compute"
  SOURCES
  asof_join_node_test.cc)
+add_arrow_compute_test(spilling_test PREFIX "arrow-compute")
 add_arrow_compute_test(tpch_node_test PREFIX "arrow-compute")
 add_arrow_compute_test(union_node_test PREFIX "arrow-compute")
 add_arrow_compute_test(util_test
@@ -47,6 +48,7 @@ add_arrow_compute_test(util_test
  task_util_test.cc)
 
 add_arrow_benchmark(expression_benchmark PREFIX "arrow-compute")
+add_arrow_benchmark(spilling_benchmark PREFIX "arrow-compute")
 
 add_arrow_benchmark(filter_benchmark
  PREFIX

cpp/src/arrow/compute/exec/accumulation_queue.cc

@@ -16,21 +16,19 @@
 // under the License.
 
 #include "arrow/compute/exec/accumulation_queue.h"
-
-#include <iterator>
+#include "arrow/compute/exec/key_hash.h"
+#include "arrow/util/atomic_util.h"
 
 namespace arrow {
-namespace util {
-using arrow::compute::ExecBatch;
+namespace compute {
+
 AccumulationQueue::AccumulationQueue(AccumulationQueue&& that) {
  this->batches_ = std::move(that.batches_);
- this->row_count_ = that.row_count_;
  that.Clear();
 }
 
 AccumulationQueue& AccumulationQueue::operator=(AccumulationQueue&& that) {
  this->batches_ = std::move(that.batches_);
- this->row_count_ = that.row_count_;
  that.Clear();
  return *this;
 }
@@ -39,20 +37,181 @@ void AccumulationQueue::Concatenate(AccumulationQueue&& that) {
  this->batches_.reserve(this->batches_.size() + that.batches_.size());
  std::move(that.batches_.begin(), that.batches_.end(),
  std::back_inserter(this->batches_));
- this->row_count_ += that.row_count_;
  that.Clear();
 }
 
 void AccumulationQueue::InsertBatch(ExecBatch batch) {
- row_count_ += batch.length;
  batches_.emplace_back(std::move(batch));
 }
 
-void AccumulationQueue::Clear() {
- row_count_ = 0;
- batches_.clear();
+void AccumulationQueue::SetBatch(size_t idx, ExecBatch batch) {
+ ARROW_DCHECK(idx < batches_.size());
+ batches_[idx] = std::move(batch);
+}
+
+size_t AccumulationQueue::CalculateRowCount() const {
+ size_t count = 0;
+ for (const ExecBatch& b : batches_) count += static_cast<size_t>(b.length);
+ return count;
+}
+
+void AccumulationQueue::Clear() { batches_.clear(); }
+
+Status SpillingAccumulationQueue::Init(QueryContext* ctx) {
+ ctx_ = ctx;
+ partition_locks_.Init(ctx_->max_concurrency(), kNumPartitions);
+ for (size_t ipart = 0; ipart < kNumPartitions; ipart++) {
+ Partition& part = partitions_[ipart];
+ part.task_group_read = ctx_->RegisterTaskGroup(
+ [this, ipart](size_t thread_index, int64_t batch_index) {
+ return partitions_[ipart].read_back_fn(
+ thread_index, static_cast<size_t>(batch_index),
+ std::move(partitions_[ipart].queue[batch_index]));
+ },
+ [this, ipart](size_t thread_index) {
+ return partitions_[ipart].on_finished(thread_index);
+ });
+ }
+ return Status::OK();
+}
+
+Status SpillingAccumulationQueue::InsertBatch(size_t thread_index, ExecBatch batch) {
+ Datum& hash_datum = batch.values.back();
+ const uint64_t* hashes =
+ reinterpret_cast<const uint64_t*>(hash_datum.array()->buffers[1]->data());
+ // `permutation` stores the indices of rows in the input batch sorted by partition.
+ std::vector<uint16_t> permutation(batch.length);
+ uint16_t part_starts[kNumPartitions + 1];
+ PartitionSort::Eval(
+ batch.length, kNumPartitions, part_starts,
+ /*partition_id=*/[&](int64_t i) { return partition_id(hashes[i]); },
+ /*output_fn=*/
+ [&permutation](int64_t input_pos, int64_t output_pos) {
+ permutation[output_pos] = static_cast<uint16_t>(input_pos);
+ });
+
+ int unprocessed_partition_ids[kNumPartitions];
+ RETURN_NOT_OK(partition_locks_.ForEachPartition(
+ thread_index, unprocessed_partition_ids,
+ /*is_prtn_empty_fn=*/
+ [&](int part_id) { return part_starts[part_id + 1] == part_starts[part_id]; },
+ /*process_prtn_fn=*/
+ [&](int locked_part_id_int) {
+ size_t locked_part_id = static_cast<size_t>(locked_part_id_int);
+ uint64_t num_total_rows_to_append =
+ part_starts[locked_part_id + 1] - part_starts[locked_part_id];
+
+ Partition& locked_part = partitions_[locked_part_id];
+
+ size_t offset = static_cast<size_t>(part_starts[locked_part_id]);
+ while (num_total_rows_to_append > 0) {
+ int num_rows_to_append =
+ std::min(static_cast<int>(num_total_rows_to_append),
+ static_cast<int>(ExecBatchBuilder::num_rows_max() -
+ locked_part.builder.num_rows()));
+
+ RETURN_NOT_OK(locked_part.builder.AppendSelected(
+ ctx_->memory_pool(), batch, num_rows_to_append, permutation.data() + offset,
+ batch.num_values()));
+
+ if (locked_part.builder.is_full()) {
+ ExecBatch batch = locked_part.builder.Flush();
+ Datum hash = std::move(batch.values.back());
+ batch.values.pop_back();
+ ExecBatch hash_batch({std::move(hash)}, batch.length);
+ if (locked_part_id < spilling_cursor_)
+ RETURN_NOT_OK(locked_part.file.SpillBatch(ctx_, std::move(batch)));
+ else
+ locked_part.queue.InsertBatch(std::move(batch));
+
+ if (locked_part_id >= hash_cursor_)
+ locked_part.hash_queue.InsertBatch(std::move(hash_batch));
+ }
+ offset += num_rows_to_append;
+ num_total_rows_to_append -= num_rows_to_append;
+ }
+ return Status::OK();
+ }));
+ return Status::OK();
+}
+
+const uint64_t* SpillingAccumulationQueue::GetHashes(size_t partition_idx,
+ size_t batch_idx) {
+ ARROW_DCHECK(partition_idx >= hash_cursor_.load());
+ Partition& partition = partitions_[partition_idx];
+ if (batch_idx > partition.hash_queue.batch_count()) {
+ const Datum& datum = partition.hash_queue[batch_idx].values[0];
+ return reinterpret_cast<const uint64_t*>(datum.array()->buffers[1]->data());
+ } else {
+ size_t hash_idx = partition.builder.num_cols();
+ KeyColumnArray kca = partition.builder.column(hash_idx - 1);
+ return reinterpret_cast<const uint64_t*>(kca.data(1));
+ }
+}
+
+Status SpillingAccumulationQueue::GetPartition(
+ size_t thread_index, size_t partition_idx,
+ std::function<Status(size_t, size_t, ExecBatch)> on_batch,
+ std::function<Status(size_t)> on_finished) {
+ bool is_in_memory = partition_idx >= spilling_cursor_.load();
+ Partition& partition = partitions_[partition_idx];
+ if (partition.builder.num_rows() > 0) {
+ ExecBatch batch = partition.builder.Flush();
+ batch.values.pop_back();
+ RETURN_NOT_OK(on_batch(thread_index,
+ /*batch_index=*/partition.queue.batch_count(),
+ std::move(batch)));
+ }
+
+ if (is_in_memory) {
+ ARROW_DCHECK(partition_idx >= hash_cursor_.load());
+ partition.read_back_fn = std::move(on_batch);
+ partition.on_finished = std::move(on_finished);
+ return ctx_->StartTaskGroup(partition.task_group_read, partition.queue.batch_count());
+ }
+
+ return partition.file.ReadBackBatches(
+ ctx_, on_batch,
+ [this, partition_idx, finished = std::move(on_finished)](size_t thread_index) {
+ RETURN_NOT_OK(partitions_[partition_idx].file.Cleanup());
+ return finished(thread_index);
+ });
+}
+
+size_t SpillingAccumulationQueue::CalculatePartitionRowCount(size_t partition) const {
+ return partitions_[partition].builder.num_rows() +
+ partitions_[partition].queue.CalculateRowCount();
 }
 
-ExecBatch& AccumulationQueue::operator[](size_t i) { return batches_[i]; }
-} // namespace util
+Result<bool> SpillingAccumulationQueue::AdvanceSpillCursor() {
+ size_t to_spill = spilling_cursor_.fetch_add(1);
+ if (to_spill >= kNumPartitions) {
+ ARROW_DCHECK(to_spill < 1000 * 1000 * 1000)
+ << "You've tried to advance the spill cursor over a billion times, you might "
+ "have a problem";
+ return false;
+ }
+
+ auto lock = partition_locks_.AcquirePartitionLock(static_cast<int>(to_spill));
+ Partition& partition = partitions_[to_spill];
+ size_t num_batches = partition.queue.batch_count();
+ for (size_t i = 0; i < num_batches; i++)
+ RETURN_NOT_OK(partition.file.SpillBatch(ctx_, std::move(partition.queue[i])));
+ return true;
+}
+
+Result<bool> SpillingAccumulationQueue::AdvanceHashCursor() {
+ size_t to_spill = hash_cursor_.fetch_add(1);
+ if (to_spill >= kNumPartitions) {
+ ARROW_DCHECK(to_spill < 1000 * 1000 * 1000)
+ << "You've tried to advance the spill cursor over a billion times, you might "
+ "have a problem";
+ return false;
+ }
+
+ auto lock = partition_locks_.AcquirePartitionLock(static_cast<int>(to_spill));
+ partitions_[to_spill].hash_queue.Clear();
+ return true;
+}
+} // namespace compute
 } // namespace arrow

cpp/src/arrow/compute/exec/accumulation_queue.h

@@ -21,16 +21,19 @@
 #include <vector>
 
 #include "arrow/compute/exec.h"
+#include "arrow/compute/exec/partition_util.h"
+#include "arrow/compute/exec/spilling_util.h"
+#include "arrow/compute/exec/task_util.h"
+#include "arrow/compute/light_array.h"
 
 namespace arrow {
-namespace util {
-using arrow::compute::ExecBatch;
+namespace compute {
 
 /// \brief A container that accumulates batches until they are ready to
 /// be processed.
-class AccumulationQueue {
+class ARROW_EXPORT AccumulationQueue {
  public:
- AccumulationQueue() : row_count_(0) {}
+ AccumulationQueue() = default;
  ~AccumulationQueue() = default;
 
  // We should never be copying ExecBatch around
@@ -42,16 +45,128 @@ class AccumulationQueue {
 
  void Concatenate(AccumulationQueue&& that);
  void InsertBatch(ExecBatch batch);
-int64_t row_count() { return row_count_; }
- size_t batch_count() { return batches_.size(); }
+void SetBatch(size_t idx, ExecBatch batch);
+ size_t batch_count() const { return batches_.size(); }
  bool empty() const { return batches_.empty(); }
+ size_t CalculateRowCount() const;
+
+ // Resizes the accumulation queue to contain size batches. The
+ // new batches will be empty and have length 0, but they will be
+ // usable (useful for concurrent modification of the AccumulationQueue
+ // of separate elements).
+ void Resize(size_t size) { batches_.resize(size); }
  void Clear();
- ExecBatch& operator[](size_t i);
+ ExecBatch& operator[](size_t i) { return batches_[i]; }
+ const ExecBatch& operator[](size_t i) const { return batches_[i]; }
 
  private:
- int64_t row_count_;
  std::vector<ExecBatch> batches_;
 };
 
-} // namespace util
+/// Accumulates batches in a queue that can be spilled to disk if needed
+///
+/// Each batch is partitioned by the lower bits of the hash column (which must be present)
+/// and rows are initially accumulated in batch builders (one per partition). As a batch
+/// builder fills up the completed batch is put into an in-memory accumulation queue (per
+/// partition).
+///
+/// When memory pressure is encountered the spilling queue's "spill cursor" can be
+/// advanced. This will cause a partition to be spilled to disk. Any future data
+/// arriving for that partition will go immediately to disk (after accumulating a full
+/// batch in the batch builder). Note that hashes are spilled separately from batches and
+/// have their own cursor. We assume that the Batch cursor is advanced faster than the
+/// spill cursor. Hashes are spilled separately to enable building a Bloom filter for
+/// spilled partitions.
+///
+/// Later, data is retrieved one partition at a time. Partitions that are in-memory will
+/// be delivered immediately in new thread tasks. Partitions that are on disk will be
+/// read from disk and delivered as they arrive.
+///
+/// This class assumes that data is fully accumulated before it is read-back. As such, do
+/// not call InsertBatch after calling GetPartition.
+class ARROW_EXPORT SpillingAccumulationQueue {
+ public:
+ // Number of partitions must be a power of two, since we assign partitions by
+ // looking at bottom few bits.
+ static constexpr int kLogNumPartitions = 6;
+ static constexpr int kNumPartitions = 1 << kLogNumPartitions;
+ Status Init(QueryContext* ctx);
+ // Assumes that the final column in batch contains 64-bit hashes of the columns.
+ Status InsertBatch(size_t thread_index, ExecBatch batch);
+ // Runs `on_batch` on each batch in the SpillingAccumulationQueue for the given
+ // partition. Each batch will have its own task. Once all batches have had their
+ // on_batch function run, `on_finished` will be called.
+ Status GetPartition(size_t thread_index, size_t partition_idx,
+ std::function<Status(size_t, size_t, ExecBatch)>
+ on_batch, // thread_index, batch_index, batch
+ std::function<Status(size_t)> on_finished);
+
+ // Returns hashes of the given partition and batch index.
+ // partition MUST be at least hash_cursor, as if partition < hash_cursor,
+ // these hashes will have been deleted.
+ const uint64_t* GetHashes(size_t partition_idx, size_t batch_idx);
+ inline size_t batch_count(size_t partition_idx) const {
+ const Partition& partition = partitions_[partition_idx];
+ size_t num_full_batches = partition_idx >= spilling_cursor_
+ ? partition.queue.batch_count()
+ : partition.file.num_batches();
+
+ return num_full_batches + (partition.builder.num_rows() > 0);
+ }
+
+ inline size_t row_count(size_t partition_idx, size_t batch_idx) const {
+ const Partition& partition = partitions_[partition_idx];
+ if (batch_idx < partition.hash_queue.batch_count())
+ return partition.hash_queue[batch_idx].length;
+ else
+ return partition.builder.num_rows();
+ }
+
+ static inline constexpr size_t partition_id(uint64_t hash) {
+ // Hash Table uses the top bits of the hash, so it is important
+ // to use the bottom bits of the hash for spilling to avoid
+ // a huge number of hash collisions per partition.
+ return static_cast<size_t>(hash & (kNumPartitions - 1));
+ }
+
+ // Returns the row count for the partition if it is still in-memory.
+ // Returns 0 if the partition has already been spilled.
+ size_t CalculatePartitionRowCount(size_t partition) const;
+
+ // Spills the next partition of batches to disk and returns true,
+ // or returns false if too many partitions have been spilled.
+ // The QueryContext's bytes_in_flight will be increased by the
+ // number of bytes spilled (unless the disk IO was very fast and
+ // the bytes_in_flight got reduced again).
+ //
+ // We expect that we always advance the SpillCursor faster than the
+ // HashCursor, and only advance the HashCursor when we've exhausted
+ // partitions for the SpillCursor.
+ Result<bool> AdvanceSpillCursor();
+ // Same as AdvanceSpillCursor but spills the hashes for the partition.
+ Result<bool> AdvanceHashCursor();
+ inline size_t spill_cursor() const { return spilling_cursor_.load(); }
+ inline size_t hash_cursor() const { return hash_cursor_.load(); }
+
+ private:
+ std::atomic<size_t> spilling_cursor_{0}; // denotes the first in-memory partition
+ std::atomic<size_t> hash_cursor_{0};
+
+ QueryContext* ctx_;
+ PartitionLocks partition_locks_;
+
+ struct Partition {
+ AccumulationQueue queue;
+ AccumulationQueue hash_queue;
+ ExecBatchBuilder builder;
+ SpillFile file;
+ int task_group_read;
+ std::function<Status(size_t, size_t, ExecBatch)> read_back_fn;
+ std::function<Status(size_t)> on_finished;
+ };
+
+ Partition partitions_[kNumPartitions];
+};
+
+} // namespace compute
 } // namespace arrow