From 6e8d07bf40f809d1572394c117349a76ca62fb65 Mon Sep 17 00:00:00 2001
From: Rok <rok@mihevc.org>
Date: Tue, 8 Feb 2022 21:54:39 +0100
Subject: [PATCH] Adding is_dst

---
 cpp/src/arrow/compute/api_scalar.cc           |  1 +
 cpp/src/arrow/compute/api_scalar.h            | 12 +++++
 .../compute/kernels/scalar_temporal_test.cc   | 11 ++++
 .../compute/kernels/scalar_temporal_unary.cc  | 51 ++++++++++++++++++-
 docs/source/cpp/compute.rst                   |  2 +
 5 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/cpp/src/arrow/compute/api_scalar.cc b/cpp/src/arrow/compute/api_scalar.cc
index 3a4d89e8e31..cead4ec4f61 100644
--- a/cpp/src/arrow/compute/api_scalar.cc
+++ b/cpp/src/arrow/compute/api_scalar.cc
@@ -773,6 +773,7 @@ SCALAR_EAGER_UNARY(Day, "day")
 SCALAR_EAGER_UNARY(DayOfYear, "day_of_year")
 SCALAR_EAGER_UNARY(Hour, "hour")
 SCALAR_EAGER_UNARY(YearMonthDay, "year_month_day")
+SCALAR_EAGER_UNARY(IsDaylightSavings, "is_dst")
 SCALAR_EAGER_UNARY(ISOCalendar, "iso_calendar")
 SCALAR_EAGER_UNARY(ISOWeek, "iso_week")
 SCALAR_EAGER_UNARY(ISOYear, "iso_year")
diff --git a/cpp/src/arrow/compute/api_scalar.h b/cpp/src/arrow/compute/api_scalar.h
index 96e3d491e45..c13b45b2798 100644
--- a/cpp/src/arrow/compute/api_scalar.h
+++ b/cpp/src/arrow/compute/api_scalar.h
@@ -1374,6 +1374,18 @@ ARROW_EXPORT Result<Datum> AssumeTimezone(const Datum& values,
                                           AssumeTimezoneOptions options,
                                           ExecContext* ctx = NULLPTR);
 
+/// \brief IsDaylightSavings extracts if currently observing daylight savings for each
+/// element of `values`
+///
+/// \param[in] values input to extract daylight savings indicator from
+/// \param[in] ctx the function execution context, optional
+/// \return the resulting datum
+///
+/// \since 8.0.0
+/// \note API not yet finalized
+ARROW_EXPORT Result<Datum> IsDaylightSavings(const Datum& values,
+                                             ExecContext* ctx = NULLPTR);
+
 /// \brief Finds either the FIRST, LAST, or ALL items with a key that matches the given
 /// query key in a map.
 ///
diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc
index f4fa2403c1f..120b7c0f60a 100644
--- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc
@@ -475,6 +475,8 @@ TEST_F(ScalarTemporalTest, TestTemporalComponentExtractionWithDifferentUnits) {
                      year_month_day);
     CheckScalarUnary("day_of_week", unit, times_seconds_precision, int64(), day_of_week);
     CheckScalarUnary("day_of_year", unit, times_seconds_precision, int64(), day_of_year);
+    ASSERT_RAISES(Invalid,
+                  IsDaylightSavings(ArrayFromJSON(unit, times_seconds_precision)));
     CheckScalarUnary("iso_year", unit, times_seconds_precision, int64(), iso_year);
     CheckScalarUnary("iso_week", unit, times_seconds_precision, int64(), iso_week);
     CheckScalarUnary("us_week", unit, times_seconds_precision, int64(), us_week);
@@ -570,6 +572,9 @@ TEST_F(ScalarTemporalTest, TestZoned1) {
   auto day_of_week = "[2, 1, 5, 1, 1, 0, 6, 2, 3, 5, 6, 5, 5, 5, 6, 5, null]";
   auto day_of_year =
       "[365, 60, 365, 137, 365, 364, 363, 364, 365, 2, 3, 365, 365, 362, 363, 365, null]";
+  std::string is_dst =
+      "[false, false, false, false, false, false, false, false, false, false, false, "
+      "false, false, false, false, false, null]";
   auto iso_year =
       "[1970, 2000, 1898, 2033, 2020, 2020, 2019, 2009, 2009, 2009, 2009, 2005, 2005, "
       "2008, 2008, 2011, null]";
@@ -604,6 +609,7 @@ TEST_F(ScalarTemporalTest, TestZoned1) {
   CheckScalarUnary("year_month_day", ArrayFromJSON(unit, times), year_month_day);
   CheckScalarUnary("day_of_week", unit, times, int64(), day_of_week);
   CheckScalarUnary("day_of_year", unit, times, int64(), day_of_year);
+  CheckScalarUnary("is_dst", unit, times, boolean(), is_dst);
   CheckScalarUnary("iso_year", unit, times, int64(), iso_year);
   CheckScalarUnary("iso_week", unit, times, int64(), iso_week);
   CheckScalarUnary("us_week", unit, times, int64(), us_week);
@@ -644,6 +650,9 @@ TEST_F(ScalarTemporalTest, TestZoned2) {
     auto day_of_week = "[3, 2, 6, 2, 2, 1, 0, 3, 4, 6, 0, 6, 5, 6, 0, 6, null]";
     auto day_of_year =
         "[1, 61, 1, 138, 1, 365, 364, 365, 1, 3, 4, 1, 365, 363, 364, 1, null]";
+    std::string is_dst =
+        "[false, true, false, false, true, true, true, true, true, true, true, true, "
+        "true, true, true, true, null]";
     auto iso_year =
         "[1970, 2000, 1898, 2033, 2020, 2020, 2020, 2009, 2009, 2009, 2010, 2005, 2005, "
         "2008, 2009, 2011, null]";
@@ -679,6 +688,7 @@ TEST_F(ScalarTemporalTest, TestZoned2) {
                      year_month_day);
     CheckScalarUnary("day_of_week", unit, times_seconds_precision, int64(), day_of_week);
     CheckScalarUnary("day_of_year", unit, times_seconds_precision, int64(), day_of_year);
+    CheckScalarUnary("is_dst", unit, times_seconds_precision, boolean(), is_dst);
     CheckScalarUnary("iso_year", unit, times_seconds_precision, int64(), iso_year);
     CheckScalarUnary("iso_week", unit, times_seconds_precision, int64(), iso_week);
     CheckScalarUnary("us_week", unit, times_seconds_precision, int64(), us_week);
@@ -710,6 +720,7 @@ TEST_F(ScalarTemporalTest, TestNonexistentTimezone) {
     ASSERT_RAISES(Invalid, YearMonthDay(timestamp_array));
     ASSERT_RAISES(Invalid, DayOfWeek(timestamp_array));
     ASSERT_RAISES(Invalid, DayOfYear(timestamp_array));
+    ASSERT_RAISES(Invalid, IsDaylightSavings(timestamp_array));
     ASSERT_RAISES(Invalid, ISOYear(timestamp_array));
     ASSERT_RAISES(Invalid, Week(timestamp_array));
     ASSERT_RAISES(Invalid, ISOCalendar(timestamp_array));
diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc
index 0bb170a790c..8bf82cc037a 100644
--- a/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc
@@ -48,7 +48,6 @@ using arrow_vendored::date::Monday;
 using arrow_vendored::date::months;
 using arrow_vendored::date::round;
 using arrow_vendored::date::Sunday;
-using arrow_vendored::date::sys_days;
 using arrow_vendored::date::sys_time;
 using arrow_vendored::date::trunc;
 using arrow_vendored::date::weekday;
@@ -133,6 +132,26 @@ struct AssumeTimezoneExtractor
   }
 };
 
+template <template <typename...> class Op, typename Duration, typename InType,
+          typename OutType>
+struct DaylightSavingsExtractor
+    : public TemporalComponentExtractBase<Op, Duration, InType, OutType> {
+  using Base = TemporalComponentExtractBase<Op, Duration, InType, OutType>;
+
+  static Status Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+    const auto& timezone = GetInputTimezone(batch.values[0]);
+    if (timezone.empty()) {
+      return Status::Invalid("Timestamps have no timezone. Cannot determine DST.");
+    }
+    ARROW_ASSIGN_OR_RAISE(auto tz, LocateZone(timezone));
+    using ExecTemplate = Op<Duration>;
+    auto op = ExecTemplate(nullptr, tz);
+    applicator::ScalarUnaryNotNullStateful<OutType, TimestampType, ExecTemplate> kernel{
+        op};
+    return kernel.Exec(ctx, batch, out);
+  }
+};
+
 template <template <typename...> class Op, typename Duration, typename InType,
           typename OutType>
 struct TemporalComponentExtractWeek
@@ -603,6 +622,22 @@ struct Nanosecond {
   }
 };
 
+// ----------------------------------------------------------------------
+// Extract if currently observing daylight savings
+
+template <typename Duration>
+struct IsDaylightSavings {
+  explicit IsDaylightSavings(const FunctionOptions* options, const time_zone* tz)
+      : tz_(tz) {}
+
+  template <typename T, typename Arg0>
+  T Call(KernelContext*, Arg0 arg, Status*) const {
+    return tz_->get_info(sys_time<Duration>{Duration{arg}}).save.count() != 0;
+  }
+
+  const time_zone* tz_;
+};
+
 // ----------------------------------------------------------------------
 // Round temporal values to given frequency
 
@@ -1444,6 +1479,14 @@ const FunctionDoc assume_timezone_doc{
     "AssumeTimezoneOptions",
     /*options_required=*/true};
 
+const FunctionDoc is_dst_doc{
+    "Extracts if currently observing daylight savings",
+    ("IsDaylightSavings returns true if a timestamp has a daylight saving\n"
+     "offset in the given timezone.\n"
+     "Null values emit null.\n"
+     "An error is returned if the values do not have a defined timezone."),
+    {"values"}};
+
 const FunctionDoc floor_temporal_doc{
     "Round temporal values down to nearest multiple of specified time unit",
     ("Null values emit null.\n"
@@ -1607,6 +1650,12 @@ void RegisterScalarTemporalUnary(FunctionRegistry* registry) {
                           &assume_timezone_doc, nullptr, AssumeTimezoneState::Init);
   DCHECK_OK(registry->AddFunction(std::move(assume_timezone)));
 
+  auto is_dst =
+      UnaryTemporalFactory<IsDaylightSavings, DaylightSavingsExtractor,
+                           BooleanType>::Make<WithTimestamps>("is_dst", boolean(),
+                                                              &is_dst_doc);
+  DCHECK_OK(registry->AddFunction(std::move(is_dst)));
+
   // Temporal rounding functions
   static const auto default_round_temporal_options = RoundTemporalOptions::Defaults();
   auto floor_temporal = UnaryTemporalFactory<FloorTemporal, TemporalComponentExtractRound,
diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst
index 13402718772..a706447db0d 100644
--- a/docs/source/cpp/compute.rst
+++ b/docs/source/cpp/compute.rst
@@ -1382,6 +1382,8 @@ For timestamps inputs with non-empty timezone, localized timestamp components wi
 +--------------------+------------+-------------------+---------------+----------------------------+-------+
 | hour               | Unary      | Timestamp, Time   | Int64         |                            |       |
 +--------------------+------------+-------------------+---------------+----------------------------+-------+
+| is_dst             | Unary      | Timestamp         | Boolean       |                            |       |
++--------------------+------------+-------------------+---------------+----------------------------+-------+
 | iso_week           | Unary      | Temporal          | Int64         |                            | \(2)  |
 +--------------------+------------+-------------------+---------------+----------------------------+-------+
 | iso_year           | Unary      | Temporal          | Int64         |                            | \(2)  |