Function TakeUsingVector

cpp/src/arrow/acero/hash_join_node_test.cc:446–499

Index < 0 means appending null values to all columns.

// Index < 0 means appending null values to all columns.
//
void TakeUsingVector(ExecContext* ctx, const std::vector<std::shared_ptr<Array>>& input,
                     const std::vector<int32_t> indices,
                     std::vector<std::shared_ptr<Array>>* result) {
  ASSERT_OK_AND_ASSIGN(
      std::shared_ptr<Buffer> buf,
      AllocateBuffer(indices.size() * sizeof(int32_t), ctx->memory_pool()));
  int32_t* buf_indices = reinterpret_cast<int32_t*>(buf->mutable_data());
  bool has_null_rows = false;
  for (size_t i = 0; i < indices.size(); ++i) {
    if (indices[i] < 0) {
      buf_indices[i] = 0;
      has_null_rows = true;
    } else {
      buf_indices[i] = indices[i];
    }
  }
  std::shared_ptr<Array> indices_array = MakeArray(ArrayData::Make(
      int32(), indices.size(), {nullptr, std::move(buf)}, /*null_count=*/0));

  result->resize(input.size());
  for (size_t i = 0; i < result->size(); ++i) {
    ASSERT_OK_AND_ASSIGN(Datum new_array, Take(input[i], indices_array));
    (*result)[i] = new_array.make_array();
  }
  if (has_null_rows) {
    for (size_t i = 0; i < result->size(); ++i) {
      if ((*result)[i]->data()->buffers[0] == NULLPTR) {
        ASSERT_OK_AND_ASSIGN(std::shared_ptr<Buffer> null_buf,
                             AllocateBitmap(indices.size(), ctx->memory_pool()));
        uint8_t* non_nulls = null_buf->mutable_data();
        memset(non_nulls, 0xFF, bit_util::BytesForBits(indices.size()));
        if ((*result)[i]->data()->buffers.size() == 2) {
          (*result)[i] = MakeArray(
              ArrayData::Make((*result)[i]->type(), indices.size(),
                              {std::move(null_buf), (*result)[i]->data()->buffers[1]}));
        } else {
          (*result)[i] = MakeArray(
              ArrayData::Make((*result)[i]->type(), indices.size(),
                              {std::move(null_buf), (*result)[i]->data()->buffers[1],
                               (*result)[i]->data()->buffers[2]}));
        }
      }
      (*result)[i]->data()->SetNullCount(kUnknownNullCount);
    }
    for (size_t i = 0; i < indices.size(); ++i) {
      if (indices[i] < 0) {
        for (size_t col = 0; col < result->size(); ++col) {
          uint8_t* non_nulls = (*result)[col]->data()->buffers[0]->mutable_data();
          bit_util::ClearBit(non_nulls, i);
        }
      }
    }
  }
}

// Generate random arrays given list of data types and null probabilities.
// Make sure that all generated records are unique.