Function CallRScalarUDF

r/src/compute.cpp:669–734

}

arrow::Status CallRScalarUDF(arrow::compute::KernelContext* context,
                             const arrow::compute::ExecSpan& span,
                             arrow::compute::ExecResult* result) {
  if (result->is_array_span()) {
    return arrow::Status::NotImplemented("ArraySpan result from R scalar UDF");
  }

  return SafeCallIntoRVoid(
      [&]() {
        auto kernel =
            reinterpret_cast<const arrow::compute::ScalarKernel*>(context->kernel());
        auto state = std::dynamic_pointer_cast<RScalarUDFKernelState>(kernel->data);

        cpp11::writable::list args_sexp(span.num_values());

        for (int i = 0; i < span.num_values(); i++) {
          const arrow::compute::ExecValue& exec_val = span[i];
          if (exec_val.is_array()) {
            args_sexp[i] = cpp11::to_r6<arrow::Array>(exec_val.array.ToArray());
          } else if (exec_val.is_scalar()) {
            args_sexp[i] = cpp11::to_r6<arrow::Scalar>(exec_val.scalar->GetSharedPtr());
          }
        }

        cpp11::sexp batch_length_sexp = cpp11::as_sexp(static_cast<double>(span.length));

        std::shared_ptr<arrow::DataType> output_type = result->type()->GetSharedPtr();
        cpp11::sexp output_type_sexp = cpp11::to_r6<arrow::DataType>(output_type);
        cpp11::writable::list udf_context = {batch_length_sexp, output_type_sexp};
        udf_context.names() = {"batch_length", "output_type"};

        cpp11::sexp func_result_sexp =
            cpp11::function(state->exec_func_)(udf_context, args_sexp);

        if (Rf_inherits(func_result_sexp, "Array")) {
          auto array = cpp11::as_cpp<std::shared_ptr<arrow::Array>>(func_result_sexp);

          // Error for an Array result of the wrong type
          if (!result->type()->Equals(array->type())) {
            return cpp11::stop(
                "Expected return Array or Scalar with type '%s' from user-defined "
                "function but got Array with type '%s'",
                result->type()->ToString().c_str(), array->type()->ToString().c_str());
          }

          result->value = std::move(array->data());
        } else if (Rf_inherits(func_result_sexp, "Scalar")) {
          auto scalar = cpp11::as_cpp<std::shared_ptr<arrow::Scalar>>(func_result_sexp);

          // handle a Scalar result of the wrong type
          if (!result->type()->Equals(scalar->type)) {
            return cpp11::stop(
                "Expected return Array or Scalar with type '%s' from user-defined "
                "function but got Scalar with type '%s'",
                result->type()->ToString().c_str(), scalar->type->ToString().c_str());
          }

          auto array = ValueOrStop(