public:

static constexpr int kLogMiniBatchLength = 10;

static constexpr int kMiniBatchLength = 1 << kLogMiniBatchLength;

template <typename>

friend class TempVectorHolder;

public:

Status Init(MemoryPool* pool, int64_t size) {

num_vectors_ = 0;

top_ = 0;

buffer_size_ = PaddedAllocationSize(size) + kPadding + 2 * sizeof(uint64_t);

ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateResizableBuffer(size, pool));

// Ensure later operations don't accidentally read uninitialized memory.

std::memset(buffer->mutable_data(), 0xFF, size);

buffer_ = std::move(buffer);

return Status::OK();

}

private:

int64_t PaddedAllocationSize(int64_t num_bytes) {

// Round up allocation size to multiple of 8 bytes

// to avoid returning temp vectors with unaligned address.

//

// Also add padding at the end to facilitate loads and stores

// using SIMD when number of vector elements is not divisible

// by the number of SIMD lanes.

//

return ::arrow::bit_util::RoundUp(num_bytes, sizeof(int64_t)) + kPadding;

}

void alloc(uint32_t num_bytes, uint8_t** data, int* id) {

int64_t old_top = top_;

top_ += PaddedAllocationSize(num_bytes) + 2 * sizeof(uint64_t);

// Stack overflow check

ARROW_DCHECK(top_ <= buffer_size_);

*data = buffer_->mutable_data() + old_top + sizeof(uint64_t);

// We set 8 bytes before the beginning of the allocated range and

// 8 bytes after the end to check for stack overflow (which would

// result in those known bytes being corrupted).

reinterpret_cast<uint64_t*>(buffer_->mutable_data() + old_top)[0] = kGuard1;

reinterpret_cast<uint64_t*>(buffer_->mutable_data() + top_)[-1] = kGuard2;

*id = num_vectors_++;

}

void release(int id, uint32_t num_bytes) {

ARROW_DCHECK(num_vectors_ == id + 1);

int64_t size = PaddedAllocationSize(num_bytes) + 2 * sizeof(uint64_t);

ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[-1] ==

kGuard2);

ARROW_DCHECK(top_ >= size);

top_ -= size;

ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[0] ==

kGuard1);

--num_vectors_;

}

static constexpr uint64_t kGuard1 = 0x3141592653589793ULL;

static constexpr uint64_t kGuard2 = 0x0577215664901532ULL;

static constexpr int64_t kPadding = 64;

int num_vectors_;

int64_t top_;

std::unique_ptr<Buffer> buffer_;

int64_t buffer_size_;

friend class TempVectorStack;

public:

~TempVectorHolder() { stack_->release(id_, num_elements_ * sizeof(T)); }

T* mutable_data() { return reinterpret_cast<T*>(data_); }

TempVectorHolder(TempVectorStack* stack, uint32_t num_elements) {

stack_ = stack;

num_elements_ = num_elements;

stack_->alloc(num_elements * sizeof(T), &data_, &id_);

}

private:

TempVectorStack* stack_;

uint8_t* data_;

int id_;

uint32_t num_elements_;

public:

static void bits_to_indexes(int bit_to_search, int64_t hardware_flags,

const int num_bits, const uint8_t* bits, int* num_indexes,

uint16_t* indexes, int bit_offset = 0);

static void bits_filter_indexes(int bit_to_search, int64_t hardware_flags,

const int num_bits, const uint8_t* bits,

const uint16_t* input_indexes, int* num_indexes,

uint16_t* indexes, int bit_offset = 0);

// Input and output indexes may be pointing to the same data (in-place filtering).

static void bits_split_indexes(int64_t hardware_flags, const int num_bits,

const uint8_t* bits, int* num_indexes_bit0,

uint16_t* indexes_bit0, uint16_t* indexes_bit1,

int bit_offset = 0);

// Bit 1 is replaced with byte 0xFF.

static void bits_to_bytes(int64_t hardware_flags, const int num_bits,

const uint8_t* bits, uint8_t* bytes, int bit_offset = 0);

// Return highest bit of each byte.

static void bytes_to_bits(int64_t hardware_flags, const int num_bits,

const uint8_t* bytes, uint8_t* bits, int bit_offset = 0);

static bool are_all_bytes_zero(int64_t hardware_flags, const uint8_t* bytes,

uint32_t num_bytes);

private:

inline static uint64_t SafeLoadUpTo8Bytes(const uint8_t* bytes, int num_bytes);

inline static void SafeStoreUpTo8Bytes(uint8_t* bytes, int num_bytes, uint64_t value);

inline static void bits_to_indexes_helper(uint64_t word, uint16_t base_index,

int* num_indexes, uint16_t* indexes);

inline static void bits_filter_indexes_helper(uint64_t word,

const uint16_t* input_indexes,

int* num_indexes, uint16_t* indexes);

template <int bit_to_search, bool filter_input_indexes>

static void bits_to_indexes_internal(int64_t hardware_flags, const int num_bits,

const uint8_t* bits, const uint16_t* input_indexes,

int* num_indexes, uint16_t* indexes,

uint16_t base_index = 0);

#if defined(ARROW_HAVE_AVX2)

static void bits_to_indexes_avx2(int bit_to_search, const int num_bits,

const uint8_t* bits, int* num_indexes,

uint16_t* indexes, uint16_t base_index = 0);

static void bits_filter_indexes_avx2(int bit_to_search, const int num_bits,

const uint8_t* bits, const uint16_t* input_indexes,

int* num_indexes, uint16_t* indexes);

template <int bit_to_search>

static void bits_to_indexes_imp_avx2(const int num_bits, const uint8_t* bits,

int* num_indexes, uint16_t* indexes,

uint16_t base_index = 0);

template <int bit_to_search>

static void bits_filter_indexes_imp_avx2(const int num_bits, const uint8_t* bits,

const uint16_t* input_indexes,

int* num_indexes, uint16_t* indexes);

static void bits_to_bytes_avx2(const int num_bits, const uint8_t* bits, uint8_t* bytes);

static void bytes_to_bits_avx2(const int num_bits, const uint8_t* bytes, uint8_t* bits);

static bool are_all_bytes_zero_avx2(const uint8_t* bytes, uint32_t num_bytes);

#endif

int total = total_.load();

if (total == -1) return {};

return total;

DCHECK_NE(count_.load(), total_.load());

int count = count_.fetch_add(1) + 1;

if (count != total_.load()) return false;

return DoneOnce();

total_.store(total);

if (count_.load() != total) return false;

return DoneOnce();

public:

static int64_t FixBitAccess(int num_bytes_accessed_together, int64_t num_rows,

int bit_offset) {

int64_t num_bytes = bit_util::BytesForBits(num_rows + bit_offset);

int64_t num_bytes_safe =

std::max(static_cast<int64_t>(0LL), num_bytes - num_bytes_accessed_together + 1);

int64_t num_rows_safe =

std::max(static_cast<int64_t>(0LL), 8 * num_bytes_safe - bit_offset);

return std::min(num_rows_safe, num_rows);

}

static int64_t FixBinaryAccess(int num_bytes_accessed_together, int64_t num_rows,

int64_t length) {

int64_t num_rows_to_skip = bit_util::CeilDiv(length, num_bytes_accessed_together);

int64_t num_rows_safe =

std::max(static_cast<int64_t>(0LL), num_rows - num_rows_to_skip);

return num_rows_safe;

}

static int64_t FixVarBinaryAccess(int num_bytes_accessed_together, int64_t num_rows,

const uint32_t* offsets) {

// Do not process rows that could read past the end of the buffer using N

// byte loads/stores.

//

int64_t num_rows_safe = num_rows;

while (num_rows_safe > 0 &&

offsets[num_rows_safe] + num_bytes_accessed_together > offsets[num_rows]) {

--num_rows_safe;

}

return num_rows_safe;

}

static int FixSelection(int64_t num_rows_safe, int num_selected,

const uint16_t* selection) {

int num_selected_safe = num_selected;

while (num_selected_safe > 0 && selection[num_selected_safe - 1] >= num_rows_safe) {

--num_selected_safe;

}

return num_selected_safe;

}

const PostVisitCall& post_call) {

ARROW_ASSIGN_OR_RAISE(expr, Result<Expression>(pre(std::move(expr))));

auto call = expr.call();

if (!call) return expr;

bool at_least_one_modified = false;

std::vector<Expression> modified_arguments;

for (size_t i = 0; i < call->arguments.size(); ++i) {

ARROW_ASSIGN_OR_RAISE(auto modified_argument,

ModifyExpression(call->arguments[i], pre, post_call));

if (Identical(modified_argument, call->arguments[i])) {

continue;

}

if (!at_least_one_modified) {

modified_arguments = call->arguments;

at_least_one_modified = true;

}

modified_arguments[i] = std::move(modified_argument);

}

if (at_least_one_modified) {

// reconstruct the call expression with the modified arguments

auto modified_call = *call;

modified_call.arguments = std::move(modified_arguments);

return post_call(Expression(std::move(modified_call)), &expr);

}

return post_call(std::move(expr), NULLPTR);

public:

// All nodes should call TraceStartProducing or NoteStartProducing exactly once

// Most nodes will be fine with a call to NoteStartProducing since the StartProducing

// call is usually fairly cheap and simply schedules tasks to fetch the actual data.

// Create a span to record the StartProducing work

[[nodiscard]] ::arrow::internal::tracing::Scope TraceStartProducing(

std::string extra_details) const {

const T& self = cast();

std::string node_kind(self.kind_name());

util::tracing::Span span;

return START_SCOPED_SPAN(

span, node_kind + "::StartProducing",

{{"node.details", extra_details}, {"node.label", self.label()}});

}

// Record a call to StartProducing without creating with a span

void NoteStartProducing(std::string extra_details) const {

const T& self = cast();

std::string node_kind(self.kind_name());

EVENT_ON_CURRENT_SPAN(

node_kind + "::StartProducing",

{{"node.details", extra_details}, {"node.label", self.label()}});

}

// All nodes should call TraceInputReceived for each batch they receive. This call

// should track the time spent processing the batch. NoteInputReceived is available

// but usually won't be used unless a node is simply adding batches to a trivial queue.

// Create a span to record the InputReceived work

[[nodiscard]] ::arrow::internal::tracing::Scope TraceInputReceived(

const ExecBatch& batch) const {

const T& self = cast();

std::string node_kind(self.kind_name());

util::tracing::Span span;

return START_SCOPED_SPAN(

span, node_kind + "::InputReceived",

{{"node.label", self.label()}, {"node.batch_length", batch.length}});

}

// Record a call to InputReceived without creating with a span

void NoteInputReceived(const ExecBatch& batch) const {

const T& self = cast();

std::string node_kind(self.kind_name());

EVENT_ON_CURRENT_SPAN(

node_kind + "::InputReceived",

{{"node.label", self.label()}, {"node.batch_length", batch.length}});

}

// Create a span to record any "finish" work. This should NOT be called as part of

// InputFinished and many nodes may not need to call this at all. This should be used

// when a node has some extra work that has to be done once it has received all of its

// data. For example, an aggregation node calculating aggregations. This will

// typically be called as a result of InputFinished OR InputReceived.

[[nodiscard]] ::arrow::internal::tracing::Scope TraceFinish() const {

const T& self = cast();

std::string node_kind(self.kind_name());

util::tracing::Span span;

return START_SCOPED_SPAN(span, node_kind + "::Finish",

{{"node.label", self.label()}});

}

private:

const T& cast() const { return static_cast<const T&>(*this); }