Function PairwiseExecImpl

cpp/src/arrow/compute/kernels/vector_pairwise.cc:55–90

A generic pairwise implementation that can be reused by different ops.

/// A generic pairwise implementation that can be reused by different ops.
Status PairwiseExecImpl(KernelContext* ctx, const ArraySpan& input,
                        const ArrayKernelExec& scalar_exec, int64_t periods,
                        ArrayData* result) {
  // We only compute values in the region where the input-with-offset overlaps
  // the original input. The margin where these do not overlap gets filled with null.
  const auto margin_length = std::min(abs(periods), input.length);
  const auto computed_length = input.length - margin_length;
  const auto computed_start = periods > 0 ? margin_length : 0;
  const auto left_start = computed_start;
  const auto right_start = margin_length - computed_start;
  // prepare null bitmap
  int64_t null_count = margin_length;
  for (int64_t i = computed_start; i < computed_start + computed_length; i++) {
    if (input.IsValid(i) && input.IsValid(i - periods)) {
      bit_util::SetBit(result->buffers[0]->mutable_data(), i);
    } else {
      ++null_count;
    }
  }
  result->null_count = null_count;
  // prepare input span
  ArraySpan left(input);
  left.SetSlice(left_start, computed_length);
  ArraySpan right(input);
  right.SetSlice(right_start, computed_length);
  // prepare output span
  ArraySpan output_span;
  output_span.SetMembers(*result);
  output_span.offset = computed_start;
  output_span.length = computed_length;
  ExecResult output{output_span};
  // execute scalar function
  RETURN_NOT_OK(scalar_exec(ctx, ExecSpan({left, right}, computed_length), &output));

  return Status::OK();
}

Status PairwiseExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
  const auto& state = checked_cast<const PairwiseState&>(*ctx->state());