ARROW-8934: [C++] Enable compute::Subtract with timestamp inputs to return duration

cpp/src/arrow/compute/kernels/codegen_internal.cc

@@ -26,6 +26,7 @@
 
 namespace arrow {
 namespace compute {
+namespace internal {
 
 void ExecFail(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
   ctx->SetStatus(Status::NotImplemented("This kernel is malformed"));
@@ -131,6 +132,12 @@ const std::vector<std::shared_ptr<DataType>>& FloatingPointTypes() {
   return g_floating_types;
 }
 
+const std::vector<TimeUnit::type>& AllTimeUnits() {
+  static std::vector<TimeUnit::type> units = {TimeUnit::SECOND, TimeUnit::MILLI,
+                                              TimeUnit::MICRO, TimeUnit::NANO};
+  return units;
+}
+
 const std::vector<std::shared_ptr<DataType>>& NumericTypes() {
   std::call_once(codegen_static_initialized, InitStaticData);
   return g_numeric_types;
@@ -172,5 +179,6 @@ Result<ValueDescr> FirstType(KernelContext*, const std::vector<ValueDescr>& desc
   return descrs[0];
 }
 
+}  // namespace internal
 }  // namespace compute
 }  // namespace arrow

cpp/src/arrow/compute/kernels/codegen_internal.h

@@ -54,6 +54,7 @@ using internal::FirstTimeBitmapWriter;
 using internal::GenerateBitsUnrolled;
 
 namespace compute {
+namespace internal {
 
 #ifdef ARROW_EXTRA_ERROR_CONTEXT
 
@@ -257,6 +258,9 @@ const std::vector<std::shared_ptr<DataType>>& UnsignedIntTypes();
 const std::vector<std::shared_ptr<DataType>>& IntTypes();
 const std::vector<std::shared_ptr<DataType>>& FloatingPointTypes();
 
+ARROW_EXPORT
+const std::vector<TimeUnit::type>& AllTimeUnits();
+
 // Returns a vector of example instances of parametric types such as
 //
 // * Decimal
@@ -884,5 +888,6 @@ ArrayKernelExec GenerateTemporal(detail::GetTypeId get_id) {
 // END of kernel generator-dispatchers
 // ----------------------------------------------------------------------
 
+}  // namespace internal
 }  // namespace compute
 }  // namespace arrow

cpp/src/arrow/compute/kernels/scalar_arithmetic.cc

@@ -24,6 +24,7 @@
 
 namespace arrow {
 namespace compute {
+namespace internal {
 namespace {
 
 template <typename T>
@@ -241,6 +242,7 @@ ArrayKernelExec NumericEqualTypesBinary(detail::GetTypeId get_id) {
     case Type::UINT32:
       return ScalarBinaryEqualTypes<UInt32Type, UInt32Type, Op>::Exec;
     case Type::INT64:
+    case Type::TIMESTAMP:
       return ScalarBinaryEqualTypes<Int64Type, Int64Type, Op>::Exec;
     case Type::UINT64:
       return ScalarBinaryEqualTypes<UInt64Type, UInt64Type, Op>::Exec;
@@ -255,26 +257,50 @@ ArrayKernelExec NumericEqualTypesBinary(detail::GetTypeId get_id) {
 }
 
 template <typename Op>
-void AddBinaryFunction(std::string name, FunctionRegistry* registry) {
+std::shared_ptr<ScalarFunction> MakeArithmeticFunction(std::string name) {
   auto func = std::make_shared<ScalarFunction>(name, Arity::Binary());
   for (const auto& ty : NumericTypes()) {
     auto exec = NumericEqualTypesBinary<Op>(ty);
     DCHECK_OK(func->AddKernel({ty, ty}, ty, exec));
   }
-  DCHECK_OK(registry->AddFunction(std::move(func)));
+  return func;
 }
 
 }  // namespace
 
-namespace internal {
-
 void RegisterScalarArithmetic(FunctionRegistry* registry) {
-  AddBinaryFunction<Add>("add", registry);
-  AddBinaryFunction<AddChecked>("add_checked", registry);
-  AddBinaryFunction<Subtract>("subtract", registry);
-  AddBinaryFunction<SubtractChecked>("subtract_checked", registry);
-  AddBinaryFunction<Multiply>("multiply", registry);
-  AddBinaryFunction<MultiplyChecked>("multiply_checked", registry);
+  // ----------------------------------------------------------------------
+  auto add = MakeArithmeticFunction<Add>("add");
+  DCHECK_OK(registry->AddFunction(std::move(add)));
+
+  // ----------------------------------------------------------------------
+  auto add_checked = MakeArithmeticFunction<AddChecked>("add_checked");
+  DCHECK_OK(registry->AddFunction(std::move(add_checked)));
+
+  // ----------------------------------------------------------------------
+  // subtract
+  auto subtract = MakeArithmeticFunction<Subtract>("subtract");
+
+  // Add subtract(timestamp, timestamp) -> duration
+  for (auto unit : AllTimeUnits()) {
+    InputType in_type(match::TimestampTypeUnit(unit));
+    auto exec = NumericEqualTypesBinary<Subtract>(Type::TIMESTAMP);
+    DCHECK_OK(subtract->AddKernel({in_type, in_type}, duration(unit), std::move(exec)));
+  }
+
+  DCHECK_OK(registry->AddFunction(std::move(subtract)));
+
+  // ----------------------------------------------------------------------
+  auto subtract_checked = MakeArithmeticFunction<SubtractChecked>("subtract_checked");
+  DCHECK_OK(registry->AddFunction(std::move(subtract_checked)));
+
+  // ----------------------------------------------------------------------
+  auto multiply = MakeArithmeticFunction<Multiply>("multiply");
+  DCHECK_OK(registry->AddFunction(std::move(multiply)));
+
+  // ----------------------------------------------------------------------
+  auto multiply_checked = MakeArithmeticFunction<MultiplyChecked>("multiply_checked");
+  DCHECK_OK(registry->AddFunction(std::move(multiply_checked)));
 }
 
 }  // namespace internal