diff --git a/cpp/src/arrow/compute/kernels/scalar_cast_temporal.cc b/cpp/src/arrow/compute/kernels/scalar_cast_temporal.cc index d076186e5635..0632e776f5ff 100644 --- a/cpp/src/arrow/compute/kernels/scalar_cast_temporal.cc +++ b/cpp/src/arrow/compute/kernels/scalar_cast_temporal.cc @@ -24,12 +24,14 @@ #include "arrow/compute/kernels/scalar_cast_internal.h" #include "arrow/compute/kernels/temporal_internal.h" #include "arrow/util/bitmap_reader.h" +#include "arrow/util/int_util.h" #include "arrow/util/logging_internal.h" #include "arrow/util/time.h" #include "arrow/util/value_parsing.h" namespace arrow { +using internal::IntegersCanFit; using internal::ParseTimestampISO8601; using internal::ParseYYYY_MM_DD; @@ -495,6 +497,22 @@ struct CastFunctor +struct CastFunctor::value>> { + static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { + const auto& options = checked_cast(*ctx->state()).options; + if (!options.allow_int_overflow) { + RETURN_NOT_OK(IntegersCanFit(batch[0].array, *int32())); + } + CastNumberToNumberUnsafe(batch[0].type()->id(), Type::INT32, batch[0].array, + out->array_span_mutable()); + return Status::OK(); + } +}; + namespace { template @@ -523,9 +541,18 @@ std::shared_ptr GetDate32Cast() { // date32 -> date32 AddZeroCopyCast(Type::DATE32, date32(), date32(), func.get()); - // int32 -> date32 + // int32 -> date32 (zero copy) AddZeroCopyCast(Type::INT32, int32(), date32(), func.get()); + // Other integer types -> date32 + for (const auto& in_ty : IntTypes()) { + if (in_ty->id() == Type::INT32) { + continue; + } + auto exec = GenerateInteger(in_ty->id()); + DCHECK_OK(func->AddKernel(in_ty->id(), {in_ty}, date32(), std::move(exec))); + } + // date64 -> date32 AddSimpleCast(date64(), date32(), func.get()); diff --git a/cpp/src/arrow/compute/kernels/scalar_cast_test.cc b/cpp/src/arrow/compute/kernels/scalar_cast_test.cc index 864ec4afce0c..8c5b2f9801db 100644 --- a/cpp/src/arrow/compute/kernels/scalar_cast_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_cast_test.cc @@ -219,6 +219,9 @@ TEST(Cast, CanCast) { ExpectCanCast(null(), {smallint()}); ExpectCanCast(tinyint(), {smallint()}); // cast between compatible storage types + for (const auto& integer_type : kIntegerTypes) { + ExpectCanCast(integer_type, {date32()}); + } ExpectCanCast(date32(), {utf8(), large_utf8()}); ExpectCanCast(date64(), {utf8(), large_utf8()}); ExpectCanCast(timestamp(TimeUnit::NANO), {utf8(), large_utf8()}); @@ -2831,6 +2834,18 @@ TEST(Cast, DateToDate) { CheckCast(day_64_will_be_truncated, day_32, options); } +TEST(Cast, IntegerToDate32) { + auto expected = ArrayFromJSON(date32(), "[0, null, 1, 2]"); + for (const auto& integer_type : kIntegerTypes) { + CheckCast(ArrayFromJSON(integer_type, "[0, null, 1, 2]"), expected); + } + + CheckCastFails(ArrayFromJSON(int64(), "[0, 2147483648]"), + CastOptions::Safe(date32())); + CheckCastFails(ArrayFromJSON(uint64(), "[0, 4294967296]"), + CastOptions::Safe(date32())); +} + TEST(Cast, DateZeroCopy) { for (auto zero_copy_to_type : { date32(), diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 63c0eb0bea4e..c0315825421c 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -2292,6 +2292,17 @@ def test_cast_timestamp_to_int(): assert result.equals(expected) +@pytest.mark.numpy +def test_cast_date32_from_numpy_integers(): + expected = pa.array([datetime.date(1970, 1, 1), + datetime.date(1970, 1, 2), + datetime.date(1970, 1, 3)], type=pa.date32()) + for dtype in ['int8', 'uint8', 'int16', 'uint16', 'int32', 'uint32', + 'int64', 'uint64']: + result = pa.array(np.array([0, 1, 2], dtype=dtype), type=pa.date32()) + assert result.equals(expected) + + def test_cast_date32_to_int(): arr = pa.array([0, 1, 2], type='i4')