From 24a1e58310ae7b439ce4b574ed519c10b3d3ad22 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 19 Mar 2026 23:18:11 -0700 Subject: [PATCH] GH-49514: [C++][Python] add make_date compute function --- cpp/src/arrow/compute/api_scalar.cc | 5 ++ cpp/src/arrow/compute/api_scalar.h | 14 +++ .../compute/kernels/scalar_temporal_binary.cc | 89 +++++++++++++++++++ .../compute/kernels/scalar_temporal_test.cc | 30 +++++++ python/pyarrow/tests/test_compute.py | 20 +++++ 5 files changed, 158 insertions(+) diff --git a/cpp/src/arrow/compute/api_scalar.cc b/cpp/src/arrow/compute/api_scalar.cc index b43eca542f36..a209f0736edf 100644 --- a/cpp/src/arrow/compute/api_scalar.cc +++ b/cpp/src/arrow/compute/api_scalar.cc @@ -900,6 +900,11 @@ SCALAR_EAGER_UNARY(USWeek, "us_week") SCALAR_EAGER_UNARY(USYear, "us_year") SCALAR_EAGER_UNARY(Year, "year") +Result MakeDate(const Datum& year, const Datum& month, const Datum& day, + ExecContext* ctx) { + return CallFunction("make_date", {year, month, day}, ctx); +} + Result AssumeTimezone(const Datum& arg, AssumeTimezoneOptions options, ExecContext* ctx) { return CallFunction("assume_timezone", {arg}, &options, ctx); diff --git a/cpp/src/arrow/compute/api_scalar.h b/cpp/src/arrow/compute/api_scalar.h index 8b341e865a16..a2562766cf34 100644 --- a/cpp/src/arrow/compute/api_scalar.h +++ b/cpp/src/arrow/compute/api_scalar.h @@ -1345,6 +1345,20 @@ Result Month(const Datum& values, ExecContext* ctx = NULLPTR); ARROW_EXPORT Result Day(const Datum& values, ExecContext* ctx = NULLPTR); +/// \brief MakeDate composes date32 values from year, month and day components. +/// +/// \param[in] year input year component +/// \param[in] month input month component +/// \param[in] day input day component +/// \param[in] ctx the function execution context, optional +/// \return the resulting datum +/// +/// \since 22.0.0 +/// \note API not yet finalized +ARROW_EXPORT +Result MakeDate(const Datum& year, const Datum& month, const Datum& day, + ExecContext* ctx = NULLPTR); + /// \brief YearMonthDay returns a struct containing the Year, Month and Day value for /// each element of `values`. /// diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_binary.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_binary.cc index 6d975d74e219..50e2061de8c8 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_binary.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_binary.cc @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. +#include #include #include +#include #include #include "arrow/builder.h" @@ -24,6 +26,8 @@ #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/kernels/temporal_internal.h" #include "arrow/compute/registry_internal.h" +#include "arrow/scalar.h" +#include "arrow/util/bit_util.h" #include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" #include "arrow/util/time.h" @@ -48,8 +52,10 @@ using chrono::jan; using chrono::last; using chrono::local_days; using chrono::local_time; +using chrono::month; using chrono::mon; using chrono::sun; +using chrono::day; using chrono::sys_days; using chrono::sys_time; using chrono::thu; @@ -57,6 +63,7 @@ using chrono::trunc; using chrono::wed; using chrono::weekday; using chrono::weeks; +using chrono::year; using chrono::year_month_day; using chrono::year_month_weekday; using chrono::years; @@ -311,6 +318,73 @@ using MicrosecondsBetween = UnitsBetween using NanosecondsBetween = UnitsBetween; +bool IsValueValid(const ArraySpan& array, int64_t i) { + return !array.MayHaveNulls() || + bit_util::GetBit(array.buffers[0].data, array.offset + i); +} + +Result ComposeDate32(int64_t year_component, int64_t month_component, + int64_t day_component) { + const int64_t year_min = static_cast(year::min()); + const int64_t year_max = static_cast(year::max()); + if (year_component < year_min || year_component > year_max) { + return Status::Invalid("Year component out of range for date32: ", year_component); + } + + const auto ymd = year{static_cast(year_component)} / + month{static_cast(month_component)} / + day{static_cast(day_component)}; + if (!ymd.ok()) { + return Status::Invalid("Invalid date components: year=", year_component, + ", month=", month_component, ", day=", day_component); + } + + const auto days_since_epoch = sys_days{ymd}.time_since_epoch().count(); + if (days_since_epoch < std::numeric_limits::min() || + days_since_epoch > std::numeric_limits::max()) { + return Status::Invalid("Composed date is out of range for date32: year=", + year_component, ", month=", month_component, + ", day=", day_component); + } + return static_cast(days_since_epoch); +} + +struct MakeDate { + static Status Exec(KernelContext*, const ExecSpan& batch, ExecResult* out) { + std::array arrays = {nullptr, nullptr, nullptr}; + std::array data = {nullptr, nullptr, nullptr}; + std::array scalars = {0, 0, 0}; + + for (int i = 0; i < batch.num_values(); ++i) { + if (batch[i].is_array()) { + arrays[i] = &batch[i].array; + data[i] = batch[i].array.GetValues(1); + } else { + const auto& scalar = *batch[i].scalar; + if (!scalar.is_valid) { + return Status::OK(); + } + scalars[i] = checked_cast(scalar).value; + } + } + + auto* out_values = out->array_span_mutable()->GetValues(1); + for (int64_t i = 0; i < batch.length; ++i) { + if ((arrays[0] && !IsValueValid(*arrays[0], i)) || + (arrays[1] && !IsValueValid(*arrays[1], i)) || + (arrays[2] && !IsValueValid(*arrays[2], i))) { + continue; + } + + const int64_t years = arrays[0] ? data[0][i] : scalars[0]; + const int64_t months = arrays[1] ? data[1][i] : scalars[1]; + const int64_t days = arrays[2] ? data[2][i] : scalars[2]; + ARROW_ASSIGN_OR_RAISE(out_values[i], ComposeDate32(years, months, days)); + } + return Status::OK(); + } +}; + // ---------------------------------------------------------------------- // Registration helpers @@ -453,9 +527,24 @@ const FunctionDoc nanoseconds_between_doc{ "Null values emit null."), {"start", "end"}}; +const FunctionDoc make_date_doc{ + "Compose date32 values from year, month, and day components", + ("For each set of input components, compose a date32 value.\n" + "Null values emit null.\n" + "An error is returned if any component set does not form a valid date."), + {"year", "month", "day"}}; + } // namespace void RegisterScalarTemporalBinary(FunctionRegistry* registry) { + // Temporal composition functions + auto make_date = + std::make_shared("make_date", Arity::Ternary(), make_date_doc); + ScalarKernel make_date_kernel({int64(), int64(), int64()}, date32(), MakeDate::Exec); + make_date_kernel.null_handling = NullHandling::INTERSECTION; + DCHECK_OK(make_date->AddKernel(std::move(make_date_kernel))); + DCHECK_OK(registry->AddFunction(std::move(make_date))); + // Temporal difference functions auto years_between = BinaryTemporalFactory::Make< diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index 60b2d791146a..6533dbd142e2 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -631,6 +631,36 @@ TEST_F(ScalarTemporalTest, TestTemporalComponentExtractionWithDifferentUnits) { } } +TEST_F(ScalarTemporalTest, TestMakeDate) { + auto years = ArrayFromJSON(int64(), year); + auto months = ArrayFromJSON(int64(), month); + auto days = ArrayFromJSON(int64(), day); + auto expected = ArrayFromJSON(date32(), date32s); + + ASSERT_OK_AND_ASSIGN(auto actual, CallFunction("make_date", {years, months, days})); + AssertDatumsEqual(expected, actual); + + ASSERT_OK_AND_ASSIGN( + auto from_wrapper, + MakeDate(ScalarFromJSON(int64(), "2024"), ArrayFromJSON(int8(), "[1, 2, null]"), + ScalarFromJSON(int16(), "1"))); + AssertDatumsEqual(ArrayFromJSON(date32(), "[19723, 19754, null]"), from_wrapper); + + ASSERT_OK_AND_ASSIGN( + auto null_scalar, + CallFunction("make_date", + {ScalarFromJSON(int64(), "null"), + ArrayFromJSON(int64(), "[1, 2, 3]"), + ScalarFromJSON(int64(), "1")})); + AssertDatumsEqual(ArrayFromJSON(date32(), "[null, null, null]"), null_scalar); + + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("Invalid date components"), + CallFunction("make_date", + {ArrayFromJSON(int64(), "[2021]"), ArrayFromJSON(int64(), "[2]"), + ArrayFromJSON(int64(), "[29]")})); +} + TEST_F(ScalarTemporalTest, TestOutsideNanosecondRange) { const char* times = R"(["1677-09-20T00:00:59.123456", "2262-04-13T23:23:23.999999"])"; auto unit = timestamp(TimeUnit::MICRO); diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 8c3b09f612cc..13a8fed0e6ce 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -2548,6 +2548,26 @@ def _check_datetime_components(timestamps, timezone=None): assert pc.week(tsa, options=week_options).equals(pa.array(iso_week)) +def test_make_date(): + year = pa.array([1970, 2000, 2019, None], type=pa.int32()) + month = pa.array([1, 2, 12, 1], type=pa.int8()) + day = pa.array([1, 29, 31, 1], type=pa.uint8()) + expected = pa.array([datetime.date(1970, 1, 1), + datetime.date(2000, 2, 29), + datetime.date(2019, 12, 31), + None], type=pa.date32()) + assert pc.make_date(year, month, day).equals(expected) + + broadcasted = pc.make_date(2024, pa.array([1, 2, None], type=pa.int16()), 1) + expected_broadcasted = pa.array([datetime.date(2024, 1, 1), + datetime.date(2024, 2, 1), + None], type=pa.date32()) + assert broadcasted.equals(expected_broadcasted) + + with pytest.raises(pa.ArrowInvalid, match="Invalid date components"): + pc.make_date([2021], [2], [29]) + + @pytest.mark.pandas def test_extract_datetime_components(request): timestamps = ["1970-01-01T00:00:59.123456789",