Function CheckAddDispatch

cpp/src/arrow/compute/function_test.cc:243–287

template <typename FunctionType, typename ExecType>
void CheckAddDispatch(FunctionType* func, ExecType exec) {
  using KernelType = typename FunctionType::KernelType;

  ASSERT_EQ(0, func->num_kernels());
  ASSERT_EQ(0, func->kernels().size());

  std::vector<InputType> in_types1 = {int32(), int32()};
  OutputType out_type1 = int32();

  ASSERT_OK(func->AddKernel(in_types1, out_type1, exec));
  ASSERT_OK(func->AddKernel({int32(), int8()}, int32(), exec));

  // Duplicate sig is okay
  ASSERT_OK(func->AddKernel(in_types1, out_type1, exec));

  // Add a kernel
  KernelType kernel({float64(), float64()}, float64(), exec);
  ASSERT_OK(func->AddKernel(kernel));

  ASSERT_EQ(4, func->num_kernels());
  ASSERT_EQ(4, func->kernels().size());

  // Try adding some invalid kernels
  ASSERT_RAISES(Invalid, func->AddKernel({}, int32(), exec));
  ASSERT_RAISES(Invalid, func->AddKernel({int32()}, int32(), exec));
  ASSERT_RAISES(Invalid, func->AddKernel({int8(), int8(), int8()}, int32(), exec));

  // Add valid and invalid kernel using kernel struct directly
  KernelType valid_kernel({boolean(), boolean()}, boolean(), exec);
  ASSERT_OK(func->AddKernel(valid_kernel));

  KernelType invalid_kernel({boolean()}, boolean(), exec);
  ASSERT_RAISES(Invalid, func->AddKernel(invalid_kernel));

  ASSERT_OK_AND_ASSIGN(const Kernel* dispatched, func->DispatchExact({int32(), int32()}));
  KernelSignature expected_sig(in_types1, out_type1);
  ASSERT_TRUE(dispatched->signature->Equals(expected_sig));

  // No kernel available
  ASSERT_RAISES(NotImplemented, func->DispatchExact({utf8(), utf8()}));

  // Wrong arity
  ASSERT_RAISES(Invalid, func->DispatchExact({}));
  ASSERT_RAISES(Invalid, func->DispatchExact({int32(), int32(), int32()}));
}

TEST(ScalarVectorFunction, DispatchExact) {
  ScalarFunction func1("scalar_test", Arity::Binary(), /*doc=*/FunctionDoc::Empty());