ARROW-13561: [C++] Implement week kernel that accepts WeekOptions

cpp/src/arrow/compute/api_scalar.cc

@@ -245,8 +245,12 @@ static auto kMakeStructOptionsType = GetFunctionOptionsType<MakeStructOptions>(
     DataMember("field_nullability", &MakeStructOptions::field_nullability),
     DataMember("field_metadata", &MakeStructOptions::field_metadata));
 static auto kDayOfWeekOptionsType = GetFunctionOptionsType<DayOfWeekOptions>(
-    DataMember("one_based_numbering", &DayOfWeekOptions::one_based_numbering),
+    DataMember("count_from_zero", &DayOfWeekOptions::count_from_zero),
     DataMember("week_start", &DayOfWeekOptions::week_start));
+static auto kWeekOptionsType = GetFunctionOptionsType<WeekOptions>(
+    DataMember("week_starts_monday", &WeekOptions::week_starts_monday),
+    DataMember("count_from_zero", &WeekOptions::count_from_zero),
+    DataMember("first_week_is_fully_in_year", &WeekOptions::first_week_is_fully_in_year));
 static auto kNullOptionsType = GetFunctionOptionsType<NullOptions>(
     DataMember("nan_is_null", &NullOptions::nan_is_null));
 }  // namespace
@@ -406,12 +410,20 @@ MakeStructOptions::MakeStructOptions(std::vector<std::string> n)
 MakeStructOptions::MakeStructOptions() : MakeStructOptions(std::vector<std::string>()) {}
 constexpr char MakeStructOptions::kTypeName[];
 
-DayOfWeekOptions::DayOfWeekOptions(bool one_based_numbering, uint32_t week_start)
+DayOfWeekOptions::DayOfWeekOptions(bool count_from_zero, uint32_t week_start)
     : FunctionOptions(internal::kDayOfWeekOptionsType),
-      one_based_numbering(one_based_numbering),
+      count_from_zero(count_from_zero),
       week_start(week_start) {}
 constexpr char DayOfWeekOptions::kTypeName[];
 
+WeekOptions::WeekOptions(bool week_starts_monday, bool count_from_zero,
+                         bool first_week_is_fully_in_year)
+    : FunctionOptions(internal::kWeekOptionsType),
+      week_starts_monday(week_starts_monday),
+      count_from_zero(count_from_zero),
+      first_week_is_fully_in_year(first_week_is_fully_in_year) {}
+constexpr char WeekOptions::kTypeName[];
+
 NullOptions::NullOptions(bool nan_is_null)
     : FunctionOptions(internal::kNullOptionsType), nan_is_null(nan_is_null) {}
 constexpr char NullOptions::kTypeName[];
@@ -438,6 +450,7 @@ void RegisterScalarOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kSliceOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kMakeStructOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kDayOfWeekOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kWeekOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kNullOptionsType));
 }
 }  // namespace internal
@@ -629,6 +642,7 @@ SCALAR_EAGER_UNARY(Day, "day")
 SCALAR_EAGER_UNARY(DayOfYear, "day_of_year")
 SCALAR_EAGER_UNARY(ISOYear, "iso_year")
 SCALAR_EAGER_UNARY(ISOWeek, "iso_week")
+SCALAR_EAGER_UNARY(USWeek, "us_week")
 SCALAR_EAGER_UNARY(ISOCalendar, "iso_calendar")
 SCALAR_EAGER_UNARY(Quarter, "quarter")
 SCALAR_EAGER_UNARY(Hour, "hour")
@@ -648,6 +662,10 @@ Result<Datum> AssumeTimezone(const Datum& arg, AssumeTimezoneOptions options,
   return CallFunction("assume_timezone", {arg}, &options, ctx);
 }
 
+Result<Datum> Week(const Datum& arg, WeekOptions options, ExecContext* ctx) {
+  return CallFunction("week", {arg}, &options, ctx);
+}
+
 Result<Datum> Strftime(const Datum& arg, StrftimeOptions options, ExecContext* ctx) {
   return CallFunction("strftime", {arg}, &options, ctx);
 }

cpp/src/arrow/compute/api_scalar.h

@@ -317,12 +317,12 @@ class ARROW_EXPORT MakeStructOptions : public FunctionOptions {
 
 struct ARROW_EXPORT DayOfWeekOptions : public FunctionOptions {
  public:
-  explicit DayOfWeekOptions(bool one_based_numbering = false, uint32_t week_start = 1);
+  explicit DayOfWeekOptions(bool count_from_zero = true, uint32_t week_start = 1);
   constexpr static char const kTypeName[] = "DayOfWeekOptions";
   static DayOfWeekOptions Defaults() { return DayOfWeekOptions(); }
 
-  /// Number days from 1 if true and from 0 if false
-  bool one_based_numbering;
+  /// Number days from 0 if true and from 1 if false
+  bool count_from_zero;
   /// What day does the week start with (Monday=1, Sunday=7)
   uint32_t week_start;
 };
@@ -361,6 +361,33 @@ struct ARROW_EXPORT AssumeTimezoneOptions : public FunctionOptions {
   Nonexistent nonexistent;
 };
 
+struct ARROW_EXPORT WeekOptions : public FunctionOptions {
+ public:
+  explicit WeekOptions(bool week_starts_monday = true, bool count_from_zero = false,
+                       bool first_week_is_fully_in_year = false);
+  constexpr static char const kTypeName[] = "WeekOptions";
+  static WeekOptions Defaults() { return WeekOptions{}; }
+  static WeekOptions ISODefaults() {
+    return WeekOptions{/*week_starts_monday*/ true,
+                       /*count_from_zero=*/false,
+                       /*first_week_is_fully_in_year=*/false};
+  }
+  static WeekOptions USDefaults() {
+    return WeekOptions{/*week_starts_monday*/ false,
+                       /*count_from_zero=*/false,
+                       /*first_week_is_fully_in_year=*/false};
+  }
+
+  /// What day does the week start with (Monday=true, Sunday=false)
+  bool week_starts_monday;
+  /// Dates from current year that fall into last ISO week of the previous year return
+  /// 0 if true and 52 or 53 if false.
+  bool count_from_zero;
+  /// Must the first week be fully in January (true), or is a week that begins on
+  /// December 29, 30, or 31 considered to be the first week of the new year (false)?
+  bool first_week_is_fully_in_year;
+};
+
 /// @}
 
 /// \brief Get the absolute value of a value.
@@ -1008,7 +1035,8 @@ Result<Datum> ISOYear(const Datum& values, ExecContext* ctx = NULLPTR);
 
 /// \brief ISOWeek returns ISO week of year number for each element of `values`.
 /// First ISO week has the majority (4 or more) of its days in January.
-/// Week of the year starts with 1 and can run up to 53.
+/// ISO week starts on Monday. Year can have 52 or 53 weeks.
+/// Week numbering can start with 1.
 ///
 /// \param[in] values input to extract ISO week of year from
 /// \param[in] ctx the function execution context, optional
@@ -1018,6 +1046,34 @@ Result<Datum> ISOYear(const Datum& values, ExecContext* ctx = NULLPTR);
 /// \note API not yet finalized
 ARROW_EXPORT Result<Datum> ISOWeek(const Datum& values, ExecContext* ctx = NULLPTR);
 
+/// \brief USWeek returns US week of year number for each element of `values`.
+/// First US week has the majority (4 or more) of its days in January.
+/// US week starts on Sunday. Year can have 52 or 53 weeks.
+/// Week numbering starts with 1.
+///
+/// \param[in] values input to extract US week of year from
+/// \param[in] ctx the function execution context, optional
+/// \return the resulting datum
+///
+/// \since 6.0.0
+/// \note API not yet finalized
+ARROW_EXPORT Result<Datum> USWeek(const Datum& values, ExecContext* ctx = NULLPTR);
+
+/// \brief Week returns week of year number for each element of `values`.
+/// First ISO week has the majority (4 or more) of its days in January.
+/// Year can have 52 or 53 weeks. Week numbering can start with 0 or 1
+/// depending on DayOfWeekOptions.count_from_zero.
+///
+/// \param[in] values input to extract week of year from
+/// \param[in] options for setting numbering start
+/// \param[in] ctx the function execution context, optional
+/// \return the resulting datum
+///
+/// \since 6.0.0
+/// \note API not yet finalized
+ARROW_EXPORT Result<Datum> Week(const Datum& values, WeekOptions options = WeekOptions(),
+                                ExecContext* ctx = NULLPTR);
+
 /// \brief ISOCalendar returns a (ISO year, ISO week, ISO day of week) struct for
 /// each element of `values`.
 /// ISO week starts on Monday denoted by 1 and ends on Sunday denoted by 7.

cpp/src/arrow/compute/function_test.cc

@@ -104,6 +104,7 @@ TEST(FunctionOptions, Equality) {
   options.emplace_back(
       new MakeStructOptions({"col1"}, {false}, {key_value_metadata({{"key", "val"}})}));
   options.emplace_back(new DayOfWeekOptions(false, 1));
+  options.emplace_back(new WeekOptions(true, false, false));
   options.emplace_back(new CastOptions(CastOptions::Safe(boolean())));
   options.emplace_back(new CastOptions(CastOptions::Unsafe(int64())));
   options.emplace_back(new FilterOptions());

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

@@ -55,11 +55,14 @@ using arrow_vendored::date::literals::dec;
 using arrow_vendored::date::literals::jan;
 using arrow_vendored::date::literals::last;
 using arrow_vendored::date::literals::mon;
+using arrow_vendored::date::literals::sun;
 using arrow_vendored::date::literals::thu;
+using arrow_vendored::date::literals::wed;
 using internal::applicator::ScalarUnaryNotNull;
 using internal::applicator::SimpleUnary;
 
 using DayOfWeekState = OptionsWrapper<DayOfWeekOptions>;
+using WeekState = OptionsWrapper<WeekOptions>;
 using StrftimeState = OptionsWrapper<StrftimeOptions>;
 using AssumeTimezoneState = OptionsWrapper<AssumeTimezoneOptions>;
 
@@ -230,6 +233,18 @@ struct AssumeTimezoneExtractor
   }
 };
 
+template <template <typename...> class Op, typename Duration, typename InType,
+          typename OutType>
+struct TemporalComponentExtractWeek
+    : public TemporalComponentExtractBase<Op, Duration, InType, OutType> {
+  using Base = TemporalComponentExtractBase<Op, Duration, InType, OutType>;
+
+  static Status Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+    const WeekOptions& options = WeekState::Get(ctx);
+    return Base::ExecWithOptions(ctx, &options, batch, out);
+  }
+};
+
 // ----------------------------------------------------------------------
 // Extract year from temporal types
 //
@@ -301,7 +316,7 @@ struct DayOfWeek {
     for (int i = 0; i < 7; i++) {
       lookup_table[i] = i + 8 - options->week_start;
       lookup_table[i] = (lookup_table[i] > 6) ? lookup_table[i] - 7 : lookup_table[i];
-      lookup_table[i] += options->one_based_numbering;
+      lookup_table[i] += !options->count_from_zero;
     }
   }
 
@@ -362,31 +377,70 @@ struct ISOYear {
 };
 
 // ----------------------------------------------------------------------
-// Extract ISO week from temporal types
+// Extract week from temporal types
 //
-// First week of an ISO year has the majority (4 or more) of it's days in January.
+// First week of an ISO year has the majority (4 or more) of its days in January.
 // Last week of an ISO year has the year's last Thursday in it.
 // Based on
 // https://github.com/HowardHinnant/date/blob/6e921e1b1d21e84a5c82416ba7ecd98e33a436d0/include/date/iso_week.h#L1503
 
 template <typename Duration, typename Localizer>
-struct ISOWeek {
-  explicit ISOWeek(const FunctionOptions* options, Localizer&& localizer)
-      : localizer_(std::move(localizer)) {}
+struct Week {
+  explicit Week(const WeekOptions* options, Localizer&& localizer)
+      : localizer_(std::move(localizer)),
+        count_from_zero_(options->count_from_zero),
+        first_week_is_fully_in_year_(options->first_week_is_fully_in_year) {
+    if (options->week_starts_monday) {
+      if (first_week_is_fully_in_year_) {
+        wd_ = mon;
+      } else {
+        wd_ = thu;
+      }
+    } else {
+      if (first_week_is_fully_in_year_) {
+        wd_ = sun;
+      } else {
+        wd_ = wed;
+      }
+    }
+    if (count_from_zero_) {
+      days_offset_ = days{0};
+    } else {
+      days_offset_ = days{3};
+    }
+  }
 
   template <typename T, typename Arg0>
   T Call(KernelContext*, Arg0 arg, Status*) const {
     const auto t = floor<days>(localizer_.template ConvertTimePoint<Duration>(arg));
-    auto y = year_month_day{t + days{3}}.year();
-    auto start = localizer_.ConvertDays((y - years{1}) / dec / thu[last]) + (mon - thu);
-    if (t < start) {
-      --y;
-      start = localizer_.ConvertDays((y - years{1}) / dec / thu[last]) + (mon - thu);
+    auto y = year_month_day{t + days_offset_}.year();
+
+    if (first_week_is_fully_in_year_) {
+      auto start = localizer_.ConvertDays(y / jan / wd_[1]);
+      if (!count_from_zero_) {
+        if (t < start) {
+          --y;
+          start = localizer_.ConvertDays(y / jan / wd_[1]);
+        }
+      }
+      return static_cast<T>(floor<weeks>(t - start).count() + 1);
     }
-    return static_cast<T>(trunc<weeks>(t - start).count() + 1);
+
+    auto start = localizer_.ConvertDays((y - years{1}) / dec / wd_[last]) + (mon - thu);
+    if (!count_from_zero_) {
+      if (t < start) {
+        --y;
+        start = localizer_.ConvertDays((y - years{1}) / dec / wd_[last]) + (mon - thu);
+      }
+    }
+    return static_cast<T>(floor<weeks>(t - start).count() + 1);
   }
 
   Localizer localizer_;
+  arrow_vendored::date::weekday wd_;
+  arrow_vendored::date::days days_offset_;
+  const bool count_from_zero_;
+  const bool first_week_is_fully_in_year_;
 };
 
 // ----------------------------------------------------------------------
@@ -979,7 +1033,7 @@ const FunctionDoc day_of_week_doc{
      "represented by 6.\n"
      "`DayOfWeekOptions.week_start` can be used to set another starting day using\n"
      "the ISO numbering convention (1=start week on Monday, 7=start week on Sunday).\n"
-     "Day numbers can start at 0 or 1 based on `DayOfWeekOptions.one_based_numbering`.\n"
+     "Day numbers can start at 0 or 1 based on `DayOfWeekOptions.count_from_zero`.\n"
      "Null values emit null.\n"
      "An error is returned if the timestamps have a defined timezone but it\n"
      "cannot be found in the timezone database."),
@@ -1004,13 +1058,34 @@ const FunctionDoc iso_year_doc{
 
 const FunctionDoc iso_week_doc{
     "Extract ISO week of year number",
-    ("First ISO week has the majority (4 or more) of its days in January.\n"
+    ("First ISO week has the majority (4 or more) of its days in January."
+     "ISO week starts on Monday.\n"
      "Week of the year starts with 1 and can run up to 53.\n"
      "Null values emit null.\n"
      "An error is returned if the timestamps have a defined timezone but it\n"
      "cannot be found in the timezone database."),
     {"values"}};
 
+const FunctionDoc us_week_doc{
+    "Extract US week of year number",
+    ("First US week has the majority (4 or more) of its days in January."
+     "US week starts on Sunday.\n"
+     "Week of the year starts with 1 and can run up to 53.\n"
+     "Null values emit null.\n"
+     "An error is returned if the timestamps have a defined timezone but it\n"
+     "cannot be found in the timezone database."),
+    {"values"}};
+
+const FunctionDoc week_doc{
+    "Extract week of year number",
+    ("First week has the majority (4 or more) of its days in January.\n"
+     "Year can have 52 or 53 weeks. Week numbering can start with 0 or 1 using "
+     "DayOfWeekOptions.count_from_zero.\n"
+     "An error is returned if the timestamps have a defined timezone but it\n"
+     "cannot be found in the timezone database."),
+    {"values"},
+    "WeekOptions"};
+
 const FunctionDoc iso_calendar_doc{
     "Extract (ISO year, ISO week, ISO day of week) struct",
     ("ISO week starts on Monday denoted by 1 and ends on Sunday denoted by 7.\n"
@@ -1140,10 +1215,24 @@ void RegisterScalarTemporal(FunctionRegistry* registry) {
       "iso_year", {WithDates, WithTimestamps}, int64(), &iso_year_doc);
   DCHECK_OK(registry->AddFunction(std::move(iso_year)));
 
-  auto iso_week = MakeTemporal<ISOWeek, TemporalComponentExtract, Int64Type>(
-      "iso_week", {WithDates, WithTimestamps}, int64(), &iso_week_doc);
+  static const auto default_iso_week_options = WeekOptions::ISODefaults();
+  auto iso_week = MakeTemporal<Week, TemporalComponentExtractWeek, Int64Type>(
+      "iso_week", {WithDates, WithTimestamps}, int64(), &iso_week_doc,
+      &default_iso_week_options, WeekState::Init);
   DCHECK_OK(registry->AddFunction(std::move(iso_week)));
 
+  static const auto default_us_week_options = WeekOptions::USDefaults();
+  auto us_week = MakeTemporal<Week, TemporalComponentExtractWeek, Int64Type>(
+      "us_week", {WithDates, WithTimestamps}, int64(), &us_week_doc,
+      &default_us_week_options, WeekState::Init);
+  DCHECK_OK(registry->AddFunction(std::move(us_week)));
+
+  static const auto default_week_options = WeekOptions();
+  auto week = MakeTemporal<Week, TemporalComponentExtractWeek, Int64Type>(
+      "week", {WithDates, WithTimestamps}, int64(), &week_doc, &default_week_options,
+      WeekState::Init);
+  DCHECK_OK(registry->AddFunction(std::move(week)));
+
   auto iso_calendar = MakeSimpleUnaryTemporal<ISOCalendar>(
       "iso_calendar", {WithDates, WithTimestamps}, IsoCalendarType(), &iso_calendar_doc);
   DCHECK_OK(registry->AddFunction(std::move(iso_calendar)));