From ab8237a6ec274dfe34bb1d17edd160ea93fc2ec9 Mon Sep 17 00:00:00 2001
From: Weston Pace <weston.pace@gmail.com>
Date: Wed, 10 Apr 2024 19:06:18 -0700
Subject: [PATCH] Reduce code duplication per PR review

---
 .../arrow/engine/substrait/type_internal.cc   | 60 ++++++++-----------
 1 file changed, 25 insertions(+), 35 deletions(-)

diff --git a/cpp/src/arrow/engine/substrait/type_internal.cc b/cpp/src/arrow/engine/substrait/type_internal.cc
index 1587112518426..5e7e364fe00c5 100644
--- a/cpp/src/arrow/engine/substrait/type_internal.cc
+++ b/cpp/src/arrow/engine/substrait/type_internal.cc
@@ -300,49 +300,39 @@ struct DataTypeToProtoImpl {
   }
   Status Visit(const Date64Type& t) { return EncodeUserDefined(t); }
 
+  template <typename Sub>
+  Status VisitTimestamp(const TimestampType& t,
+                        void (substrait::Type::*set_allocated_sub)(Sub*)) {
+    auto ts = SetWithThen(set_allocated_sub);
+    switch (t.unit()) {
+      case TimeUnit::SECOND:
+        ts->set_precision(0);
+        break;
+      case TimeUnit::MILLI:
+        ts->set_precision(3);
+        break;
+      case TimeUnit::MICRO:
+        ts->set_precision(6);
+        break;
+      case TimeUnit::NANO:
+        ts->set_precision(9);
+        break;
+      default:
+        return NotImplemented(t);
+    }
+    return Status::OK();
+  }
+
   Status Visit(const TimestampType& t) {
     if (t.timezone() == "") {
-      auto ts = SetWithThen(&substrait::Type::set_allocated_precision_timestamp);
-      switch (t.unit()) {
-        case TimeUnit::SECOND:
-          ts->set_precision(0);
-          break;
-        case TimeUnit::MILLI:
-          ts->set_precision(3);
-          break;
-        case TimeUnit::MICRO:
-          ts->set_precision(6);
-          break;
-        case TimeUnit::NANO:
-          ts->set_precision(9);
-          break;
-        default:
-          return NotImplemented(t);
-      }
+      return VisitTimestamp(t, &substrait::Type::set_allocated_precision_timestamp);
     } else {
       // Note: The timezone information is discarded here.  In Substrait the time zone
       // information is part of the function and not part of the type.  For example, to
       // convert a timestamp to a string, the time zone is passed as an argument to the
       // function.
-      auto ts = SetWithThen(&substrait::Type::set_allocated_precision_timestamp_tz);
-      switch (t.unit()) {
-        case TimeUnit::SECOND:
-          ts->set_precision(0);
-          break;
-        case TimeUnit::MILLI:
-          ts->set_precision(3);
-          break;
-        case TimeUnit::MICRO:
-          ts->set_precision(6);
-          break;
-        case TimeUnit::NANO:
-          ts->set_precision(9);
-          break;
-        default:
-          return NotImplemented(t);
-      }
+      return VisitTimestamp(t, &substrait::Type::set_allocated_precision_timestamp_tz);
     }
-    return Status::OK();
   }
 
   Status Visit(const Time32Type& t) { return EncodeUserDefined(t); }