Method Exec

cpp/src/arrow/compute/kernels/vector_replace.cc:390–450

template <typename Type>
struct ReplaceMaskChunked {
  static Status Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
    const Datum& mask = batch[1];
    const Datum& replacements = batch[2];

    // TODO(wesm): these assertions that the arguments cannot be ChunkedArray
    // should happen someplace more generic, not here
    if (!mask.is_array() && !mask.is_scalar()) {
      return Status::Invalid("Mask must be array or scalar, not ", batch[1].ToString());
    }

    if (!replacements.is_array() && !replacements.is_scalar()) {
      return Status::Invalid("Replacements must be array or scalar, not ",
                             replacements.ToString());
    }

    const ChunkedArray& arr = *batch[0].chunked_array();

    RETURN_NOT_OK(CheckReplaceMaskInputs(*arr.type(), arr.length(), GetExecValue(mask),
                                         *replacements.type(), replacements.length(),
                                         replacements.is_arraylike()));

    ExecValue replacements_val = GetExecValue(replacements);

    // Chunked array
    ArrayVector output_chunks;
    output_chunks.reserve(arr.num_chunks());

    int64_t mask_offset = 0;
    int64_t replacements_offset = 0;
    for (const std::shared_ptr<Array>& chunk : arr.chunks()) {
      if (chunk->length() == 0) continue;
      // Allocate a new array
      ExecResult chunk_result;
      if (is_fixed_width(out->type()->id())) {
        auto chunk_out = std::make_shared<ArrayData>(chunk->type(), chunk->length());
        chunk_out->buffers.resize(2);
        ARROW_ASSIGN_OR_RAISE(chunk_out->buffers[0],
                              ctx->AllocateBitmap(chunk->length()));
        const int64_t slot_width = out->type()->byte_width();
        ARROW_ASSIGN_OR_RAISE(chunk_out->buffers[1],
                              ctx->Allocate(slot_width * chunk->length()));
        chunk_result.value = chunk_out;
      }
      if (batch[1].is_scalar()) {
        ARROW_ASSIGN_OR_RAISE(
            replacements_offset,
            ReplaceMaskImpl<Type>::ExecScalarMask(
                ctx, *chunk->data(), batch[1].scalar_as<BooleanScalar>(),
                replacements_val, replacements_offset, &chunk_result));
      } else {
        ARROW_ASSIGN_OR_RAISE(replacements_offset,
                              ReplaceMaskImpl<Type>::ExecArrayMask(
                                  ctx, *chunk->data(), *batch[1].array(), mask_offset,
                                  replacements_val, replacements_offset, &chunk_result));
      }
      output_chunks.push_back(MakeArray(chunk_result.array_data()));
      mask_offset += chunk->length();
    }