diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 7c6313c326b5..da1a70b7b4e5 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -865,6 +865,7 @@ if(ARROW_COMPUTE) compute/kernels/scalar_string_ascii.cc compute/kernels/scalar_string_utf8.cc compute/kernels/scalar_temporal_binary.cc + compute/kernels/scalar_temporal_ternary.cc compute/kernels/scalar_temporal_unary.cc compute/kernels/scalar_validity.cc compute/kernels/util_internal.cc diff --git a/cpp/src/arrow/compute/api_scalar.cc b/cpp/src/arrow/compute/api_scalar.cc index f3b3dfbfbdfa..38ee736b49d2 100644 --- a/cpp/src/arrow/compute/api_scalar.cc +++ b/cpp/src/arrow/compute/api_scalar.cc @@ -953,6 +953,12 @@ SCALAR_EAGER_BINARY(MillisecondsBetween, "milliseconds_between") SCALAR_EAGER_BINARY(MicrosecondsBetween, "microseconds_between") SCALAR_EAGER_BINARY(NanosecondsBetween, "nanoseconds_between") +Result MonthDayNanoIntervalFromParts(const Datum& months, const Datum& days, + const Datum& nanoseconds, ExecContext* ctx) { + return CallFunction("month_day_nano_interval_from_parts", {months, days, nanoseconds}, + ctx); +} + // ---------------------------------------------------------------------- // Structural transforms Result MapLookup(const Datum& arg, MapLookupOptions options, ExecContext* ctx) { diff --git a/cpp/src/arrow/compute/api_scalar.h b/cpp/src/arrow/compute/api_scalar.h index 5acc8c7bb6b2..2838d6c5d2d5 100644 --- a/cpp/src/arrow/compute/api_scalar.h +++ b/cpp/src/arrow/compute/api_scalar.h @@ -1732,6 +1732,24 @@ ARROW_EXPORT Result WeeksBetween(const Datum& left, const Datum& right, ARROW_EXPORT Result MonthDayNanoBetween(const Datum& left, const Datum& right, ExecContext* ctx = NULLPTR); +/// \brief Construct month-day-nanosecond intervals from integer components +/// +/// Months and days are safely cast to int32, and nanoseconds to int64. +/// Scalar components are broadcast. If any component is null, the output is null. +/// Component values are preserved without normalization. +/// +/// \param[in] months integer array or scalar of months +/// \param[in] days integer array or scalar of days +/// \param[in] nanoseconds integer array or scalar of nanoseconds +/// \param[in] ctx the function execution context, optional +/// \return the resulting datum +/// +/// \since 27.0.0 +ARROW_EXPORT Result MonthDayNanoIntervalFromParts(const Datum& months, + const Datum& days, + const Datum& nanoseconds, + ExecContext* ctx = NULLPTR); + /// \brief DayTime Between finds the number of days and milliseconds between two values /// /// \param[in] left input treated as the start time diff --git a/cpp/src/arrow/compute/initialize.cc b/cpp/src/arrow/compute/initialize.cc index ec531e8d490f..b0a7c2609090 100644 --- a/cpp/src/arrow/compute/initialize.cc +++ b/cpp/src/arrow/compute/initialize.cc @@ -39,6 +39,7 @@ Status RegisterComputeKernels() { internal::RegisterScalarStringAscii(registry); internal::RegisterScalarStringUtf8(registry); internal::RegisterScalarTemporalBinary(registry); + internal::RegisterScalarTemporalTernary(registry); internal::RegisterScalarTemporalUnary(registry); internal::RegisterScalarValidity(registry); diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_benchmark.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_benchmark.cc index bd90bbb6a58d..5fb272392b8c 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_benchmark.cc @@ -19,6 +19,8 @@ #include "benchmark/benchmark.h" +#include "arrow/array.h" +#include "arrow/array/builder_time.h" #include "arrow/compute/api_scalar.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/random.h" @@ -166,6 +168,75 @@ static void BenchmarkAssumeTimezone(benchmark::State& state) { state.SetItemsProcessed(state.iterations() * array_size); } +void IntervalFromPartsSetArgs(benchmark::internal::Benchmark* bench) { + for (const auto size : {1024, 65536, 1048576}) { + for (const auto inverse_null_proportion : {0, 100, 2}) { + bench->Args({size, inverse_null_proportion}); + } + } +} + +struct IntervalPartsInputs { + Datum months; + Datum days; + std::shared_ptr nanoseconds; +}; + +template +static IntervalPartsInputs MakeIntervalPartsInputs(const RegressionArgs& args) { + auto rand = random::RandomArrayGenerator(kSeed); + IntervalPartsInputs inputs; + if constexpr (ScalarParts) { + inputs.months = Datum(int32_t{2}); + inputs.days = Datum(int32_t{-3}); + } else { + inputs.months = rand.Numeric(args.size, -12, 12, args.null_proportion); + inputs.days = rand.Numeric(args.size, -31, 31, args.null_proportion); + } + inputs.nanoseconds = std::static_pointer_cast( + rand.Numeric(args.size, kInt64Min, kInt64Max, args.null_proportion)); + return inputs; +} + +template +static void BenchmarkMonthDayNanoIntervalFromParts(benchmark::State& state) { + RegressionArgs args(state, /*size_is_bytes=*/false); + const IntervalPartsInputs inputs = MakeIntervalPartsInputs(args); + ExecContext* ctx = default_exec_context(); + + for (auto _ : state) { + ABORT_NOT_OK(MonthDayNanoIntervalFromParts(inputs.months, inputs.days, + inputs.nanoseconds, ctx)); + } +} + +template +static void BenchmarkMonthDayNanoIntervalBuilder(benchmark::State& state) { + RegressionArgs args(state, /*size_is_bytes=*/false); + const IntervalPartsInputs inputs = MakeIntervalPartsInputs(args); + const auto months = ScalarParts ? nullptr : inputs.months.array_as(); + const auto days = ScalarParts ? nullptr : inputs.days.array_as(); + + for (auto _ : state) { + MonthDayNanoIntervalBuilder builder; + ABORT_NOT_OK(builder.Reserve(args.size)); + for (int64_t i = 0; i < args.size; ++i) { + bool valid = inputs.nanoseconds->IsValid(i); + if constexpr (!ScalarParts) { + valid = valid && months->IsValid(i) && days->IsValid(i); + } + if (valid) { + builder.UnsafeAppend({ScalarParts ? int32_t{2} : months->Value(i), + ScalarParts ? int32_t{-3} : days->Value(i), + inputs.nanoseconds->Value(i)}); + } else { + builder.UnsafeAppendNull(); + } + } + ABORT_NOT_OK(builder.Finish()); + } +} + auto zoned = timestamp(TimeUnit::NANO, "Pacific/Marquesas"); auto non_zoned = timestamp(TimeUnit::NANO); auto time32_type = time32(TimeUnit::MILLI); @@ -247,6 +318,17 @@ BENCHMARK_TEMPLATE(BenchmarkStrptime, non_zoned)->Apply(SetArgs); BENCHMARK_TEMPLATE(BenchmarkStrptime, zoned)->Apply(SetArgs); BENCHMARK(BenchmarkAssumeTimezone)->Apply(SetArgs); +// Null proportions apply independently to each array input. Scalar months and +// days are valid constants, and nanoseconds remains an array in both cases. +BENCHMARK_TEMPLATE(BenchmarkMonthDayNanoIntervalFromParts, false) + ->Apply(IntervalFromPartsSetArgs); +BENCHMARK_TEMPLATE(BenchmarkMonthDayNanoIntervalFromParts, true) + ->Apply(IntervalFromPartsSetArgs); +BENCHMARK_TEMPLATE(BenchmarkMonthDayNanoIntervalBuilder, false) + ->Apply(IntervalFromPartsSetArgs); +BENCHMARK_TEMPLATE(BenchmarkMonthDayNanoIntervalBuilder, true) + ->Apply(IntervalFromPartsSetArgs); + // binary temporal benchmarks DECLARE_TEMPORAL_BINARY_BENCHMARKS_DATES_AND_TIMESTAMPS(YearsBetween); DECLARE_TEMPORAL_BINARY_BENCHMARKS_DATES_AND_TIMESTAMPS(QuartersBetween); diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_ternary.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_ternary.cc new file mode 100644 index 000000000000..a42f285f5191 --- /dev/null +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_ternary.cc @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/compute/kernels/common_internal.h" +#include "arrow/compute/registry_internal.h" +#include "arrow/util/bit_run_reader.h" +#include "arrow/util/logging_internal.h" + +namespace arrow { +namespace compute { +namespace internal { +namespace { + +struct MonthDayNanoIntervalFromPartsFunction : ScalarFunction { + using ScalarFunction::ScalarFunction; + + Result DispatchBest(std::vector* types) const override { + RETURN_NOT_OK(CheckArity(types->size())); + EnsureDictionaryDecoded(types); + for (const auto& type : *types) { + if (!is_integer(type.id()) && type.id() != Type::NA) { + return arrow::compute::detail::NoMatchingKernel(this, *types); + } + } + // The executor safely casts integer components to their physical widths. + (*types)[0] = int32(); + (*types)[1] = int32(); + (*types)[2] = int64(); + return ScalarFunction::DispatchExact(*types); + } +}; + +template +struct IntervalPartValues { + using CType = typename Type::c_type; + + explicit IntervalPartValues(const ExecValue& value) { + if (value.is_array() && value.array.buffers[1].data != nullptr) { + values = value.array.GetValues(1); + } else if (value.is_scalar() && value.scalar->is_valid) { + scalar = UnboxScalar::Unbox(*value.scalar); + } + } + + CType operator[](int64_t i) const { return values ? values[i] : scalar; } + + const CType* values = nullptr; + CType scalar = 0; +}; + +Status ExecMonthDayNanoIntervalFromParts(KernelContext*, const ExecSpan& batch, + ExecResult* out) { + if (batch.length == 0) return Status::OK(); + + const IntervalPartValues months(batch[0]); + const IntervalPartValues days(batch[1]); + const IntervalPartValues nanoseconds(batch[2]); + auto* output = out->array_span_mutable(); + auto* values = output->GetValues(1); + // The executor has already intersected the three input validity bitmaps. + // Visit only valid rows, including when an input is a null scalar. + arrow::internal::VisitSetBitRunsVoid( + output->buffers[0].data, output->offset, batch.length, + [&](int64_t start, int64_t length) { + for (int64_t i = start; i < start + length; ++i) { + values[i] = {months[i], days[i], nanoseconds[i]}; + } + }); + return Status::OK(); +} + +const FunctionDoc month_day_nano_interval_from_parts_doc{ + "Construct month-day-nanosecond intervals from separate components", + ("Months, days, and nanoseconds must be integer arrays or scalars.\n" + "Months and days are safely cast to int32, and nanoseconds to int64.\n" + "Scalar inputs are broadcast. If any component is null, the interval is null.\n" + "Components are preserved without normalization."), + {"months", "days", "nanoseconds"}}; + +} // namespace + +void RegisterScalarTemporalTernary(FunctionRegistry* registry) { + auto function = std::make_shared( + "month_day_nano_interval_from_parts", Arity::Ternary(), + month_day_nano_interval_from_parts_doc); + DCHECK_OK(function->AddKernel({int32(), int32(), int64()}, month_day_nano_interval(), + ExecMonthDayNanoIntervalFromParts)); + DCHECK_OK(registry->AddFunction(std::move(function))); +} + +} // namespace internal +} // namespace compute +} // namespace arrow diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index 86a81ffdd384..bec25847a2de 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -3886,5 +3886,174 @@ TEST_F(ScalarTemporalTest, DurationUnaryArithmetics) { ArrayFromJSON(duration(TimeUnit::SECOND), "[2, 1, null, 3, 0]")); CheckScalarUnary("sign", arr, ArrayFromJSON(int8(), "[1, -1, null, 1, 0]")); } + +TEST_F(ScalarTemporalTest, MonthDayNanoIntervalFromPartsBroadcasting) { + const char* function_name = "month_day_nano_interval_from_parts"; + const DatumVector arrays = {ArrayFromJSON(int32(), "[1, 2, 3]"), + ArrayFromJSON(int32(), "[-1, -2, -3]"), + ArrayFromJSON(int64(), "[10, 20, 30]")}; + const DatumVector scalars = {ScalarFromJSON(int32(), "4"), ScalarFromJSON(int32(), "5"), + ScalarFromJSON(int64(), "-6")}; + const char* expected_arrays[] = {"[[1, 5, -6], [2, 5, -6], [3, 5, -6]]", + "[[4, -1, -6], [4, -2, -6], [4, -3, -6]]", + "[[1, -1, -6], [2, -2, -6], [3, -3, -6]]", + "[[4, 5, 10], [4, 5, 20], [4, 5, 30]]", + "[[1, 5, 10], [2, 5, 20], [3, 5, 30]]", + "[[4, -1, 10], [4, -2, 20], [4, -3, 30]]", + "[[1, -1, 10], [2, -2, 20], [3, -3, 30]]"}; + + CheckScalar(function_name, scalars, + ScalarFromJSON(month_day_nano_interval(), "[4, 5, -6]")); + for (int array_mask = 1; array_mask < 8; ++array_mask) { + SCOPED_TRACE(array_mask); + DatumVector inputs; + for (int component = 0; component < 3; ++component) { + inputs.push_back(array_mask & (1 << component) ? arrays[component] + : scalars[component]); + } + CheckScalar( + function_name, inputs, + ArrayFromJSON(month_day_nano_interval(), expected_arrays[array_mask - 1])); + } +} + +TEST_F(ScalarTemporalTest, MonthDayNanoIntervalFromPartsNullsAndSlices) { + const char* function_name = "month_day_nano_interval_from_parts"; + // Give each input a different offset, including validity bits past the first byte. + auto months = + ArrayFromJSON(int32(), "[99, 1, 2, null, 4, 5, 6, 7, 8, null, 10, 11]")->Slice(1); + auto days = + ArrayFromJSON(int32(), "[99, 99, -1, null, -3, -4, -5, -6, -7, -8, -9, -10, -11]") + ->Slice(2); + auto nanos = + ArrayFromJSON(int64(), + "[99, 99, 99, 10, 20, 30, null, 50, 60, 70, 80, 90, 100, 110]") + ->Slice(3); + CheckScalar(function_name, {months, days, nanos}, + ArrayFromJSON(month_day_nano_interval(), + "[[1, -1, 10], null, null, null, [5, -5, 50], [6, -6, 60], " + "[7, -7, 70], [8, -8, 80], null, [10, -10, 100], " + "[11, -11, 110]]")); + + const DatumVector arrays = {ArrayFromJSON(int32(), "[1, 2, 3]"), + ArrayFromJSON(int32(), "[4, 5, 6]"), + ArrayFromJSON(int64(), "[7, 8, 9]")}; + const DatumVector scalars = {ScalarFromJSON(int32(), "1"), ScalarFromJSON(int32(), "4"), + ScalarFromJSON(int64(), "7")}; + for (int component = 0; component < 3; ++component) { + SCOPED_TRACE(component); + auto inputs = arrays; + inputs[component] = MakeNullScalar(arrays[component].type()); + CheckScalar(function_name, inputs, + ArrayFromJSON(month_day_nano_interval(), "[null, null, null]")); + inputs[component] = ArrayFromJSON(null(), "[null, null, null]"); + CheckScalar(function_name, inputs, + ArrayFromJSON(month_day_nano_interval(), "[null, null, null]")); + + inputs = scalars; + inputs[component] = MakeNullScalar(scalars[component].type()); + CheckScalar(function_name, inputs, MakeNullScalar(month_day_nano_interval())); + inputs[component] = MakeNullScalar(null()); + CheckScalar(function_name, inputs, MakeNullScalar(month_day_nano_interval())); + } +} + +TEST_F(ScalarTemporalTest, MonthDayNanoIntervalFromPartsIntegerInputs) { + const char* function_name = "month_day_nano_interval_from_parts"; + CheckScalar(function_name, + {ArrayFromJSON(int32(), "[-2147483648, 2147483647, 1, -1]"), + ArrayFromJSON(int32(), "[2147483647, -2147483648, -32, 32]"), + ArrayFromJSON(int64(), + "[-9223372036854775808, 9223372036854775807, " + "86400000000000, -86400000000000]")}, + ArrayFromJSON(month_day_nano_interval(), + "[[-2147483648, 2147483647, -9223372036854775808], " + "[2147483647, -2147483648, 9223372036854775807], " + "[1, -32, 86400000000000], [-1, 32, -86400000000000]]")); + + for (const auto& type : IntTypes()) { + SCOPED_TRACE(type->ToString()); + auto values = ArrayFromJSON(type, "[0, 1, 127]"); + CheckScalar(function_name, {values, values, values}, + ArrayFromJSON(month_day_nano_interval(), + "[[0, 0, 0], [1, 1, 1], [127, 127, 127]]")); + } + CheckScalar( + function_name, + {ArrayFromJSON(int8(), "[-1, 2]"), ArrayFromJSON(uint16(), "[32, 0]"), + ArrayFromJSON(int32(), "[-1000000000, 1]")}, + ArrayFromJSON(month_day_nano_interval(), "[[-1, 32, -1000000000], [2, 0, 1]]")); +} + +TEST_F(ScalarTemporalTest, MonthDayNanoIntervalFromPartsDifferentChunks) { + CheckScalarNonRecursive( + "month_day_nano_interval_from_parts", + {ChunkedArrayFromJSON(int32(), {"[1, 2]", "[]", "[3, null]"}), + ChunkedArrayFromJSON(int32(), {"[-1]", "[-2, -3, -4]"}), + ChunkedArrayFromJSON(int64(), {"[10, 20, 30]", "[40]"})}, + ChunkedArrayFromJSON(month_day_nano_interval(), + {"[[1, -1, 10], [2, -2, 20], [3, -3, 30], null]"})); +} + +TEST_F(ScalarTemporalTest, MonthDayNanoIntervalFromPartsDictionaryInputs) { + CheckScalar( + "month_day_nano_interval_from_parts", + {DictArrayFromJSON(dictionary(int8(), int16()), "[0, 1, null, 0]", "[1, -2]"), + DictArrayFromJSON(dictionary(int8(), uint16()), "[1, 0, 1, 0]", "[4, 5]"), + DictArrayFromJSON(dictionary(int8(), int64()), "[0, 0, 1, 1]", "[10, null]")}, + ArrayFromJSON(month_day_nano_interval(), "[[1, 5, 10], [-2, 4, 10], null, null]")); +} + +TEST_F(ScalarTemporalTest, MonthDayNanoIntervalFromPartsInvalidInputs) { + const char* function_name = "month_day_nano_interval_from_parts"; + const DatumVector inputs = {ArrayFromJSON(int32(), "[0]"), + ArrayFromJSON(int32(), "[0]"), + ArrayFromJSON(int64(), "[0]")}; + for (int component = 0; component < 3; ++component) { + SCOPED_TRACE(component); + for (const auto& bad_input : + {ArrayFromJSON(float64(), "[1.0]"), ArrayFromJSON(utf8(), "[\"1\"]"), + ArrayFromJSON(boolean(), "[true]")}) { + auto bad_inputs = inputs; + bad_inputs[component] = bad_input; + ASSERT_RAISES(NotImplemented, CallFunction(function_name, bad_inputs)); + } + auto bad_inputs = inputs; + bad_inputs[component] = ArrayFromJSON(inputs[component].type(), "[0, 1]"); + ASSERT_RAISES(Invalid, CallFunction(function_name, bad_inputs)); + } + + for (const char* value : {"[2147483648]", "[-2147483649]"}) { + for (int component : {0, 1}) { + auto bad_inputs = inputs; + bad_inputs[component] = ArrayFromJSON(int64(), value); + ASSERT_RAISES(Invalid, CallFunction(function_name, bad_inputs)); + } + } + ASSERT_RAISES( + Invalid, + CallFunction(function_name, {inputs[0], inputs[1], + ArrayFromJSON(uint64(), "[9223372036854775808]")})); + ASSERT_RAISES(Invalid, CallFunction(function_name, DatumVector{})); + ASSERT_RAISES(Invalid, CallFunction(function_name, {inputs[0], inputs[1]})); + ASSERT_RAISES( + Invalid, CallFunction(function_name, {inputs[0], inputs[1], inputs[2], inputs[2]})); +} + +TEST_F(ScalarTemporalTest, MonthDayNanoIntervalFromPartsEagerAPI) { + ASSERT_OK_AND_ASSIGN(auto result, + MonthDayNanoIntervalFromParts(ScalarFromJSON(int16(), "-2"), + ScalarFromJSON(int64(), "32"), + ScalarFromJSON(uint64(), "5"))); + AssertDatumsEqual(ScalarFromJSON(month_day_nano_interval(), "[-2, 32, 5]"), result); + + ASSERT_OK_AND_ASSIGN( + result, MonthDayNanoIntervalFromParts(ArrayFromJSON(int32(), "[1, null, -1]"), + ScalarFromJSON(int32(), "0"), + ArrayFromJSON(int64(), "[10, 20, -30]"))); + AssertDatumsEqual( + ArrayFromJSON(month_day_nano_interval(), "[[1, 0, 10], null, [-1, 0, -30]]"), + result); +} } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/registry_internal.h b/cpp/src/arrow/compute/registry_internal.h index d457f6fb95c0..7eeb55e2eff4 100644 --- a/cpp/src/arrow/compute/registry_internal.h +++ b/cpp/src/arrow/compute/registry_internal.h @@ -38,6 +38,7 @@ void RegisterScalarSetLookup(FunctionRegistry* registry); void RegisterScalarStringAscii(FunctionRegistry* registry); void RegisterScalarStringUtf8(FunctionRegistry* registry); void RegisterScalarTemporalBinary(FunctionRegistry* registry); +void RegisterScalarTemporalTernary(FunctionRegistry* registry); void RegisterScalarTemporalUnary(FunctionRegistry* registry); void RegisterScalarValidity(FunctionRegistry* registry); diff --git a/cpp/src/arrow/meson.build b/cpp/src/arrow/meson.build index ae77a3d99e21..f684537435b6 100644 --- a/cpp/src/arrow/meson.build +++ b/cpp/src/arrow/meson.build @@ -618,6 +618,7 @@ if needs_compute 'compute/kernels/scalar_string_ascii.cc', 'compute/kernels/scalar_string_utf8.cc', 'compute/kernels/scalar_temporal_binary.cc', + 'compute/kernels/scalar_temporal_ternary.cc', 'compute/kernels/scalar_temporal_unary.cc', 'compute/kernels/scalar_validity.cc', 'compute/kernels/util_internal.cc', diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index 15d426dca6c6..b242fee8cd92 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -1561,6 +1561,26 @@ null input value is converted into a null output value. * \(6) Any input type that can be cast to the resulting extension's storage type. This excludes extension types, unless being cast to the same extension type. +Temporal construction +~~~~~~~~~~~~~~~~~~~~~ + ++------------------------------------+------------+-------------------+-----------------------+ +| Function name | Arity | Input types | Output type | ++====================================+============+===================+=======================+ +| month_day_nano_interval_from_parts | Ternary | Integer | MonthDayNano interval | ++------------------------------------+------------+-------------------+-----------------------+ + +``month_day_nano_interval_from_parts`` constructs an interval from three +components: months, days, and nanoseconds. Integer inputs are safely cast to +Int32 for months and days, and Int64 for nanoseconds; values outside these +ranges raise an error. Floating-point and string inputs are not supported. + +The components are stored without normalization, including negative values. +For example, 32 days remain 32 days rather than being converted into months. +If any component is null, the resulting interval is null. Scalar components +are broadcast when combined with arrays, and three scalar inputs produce a +scalar interval. + Temporal component extraction ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/python/api/compute.rst b/docs/source/python/api/compute.rst index 2c4da3abd186..216aff2f79c3 100644 --- a/docs/source/python/api/compute.rst +++ b/docs/source/python/api/compute.rst @@ -443,6 +443,7 @@ Conversions cast ceil_temporal floor_temporal + month_day_nano_interval_from_parts round_temporal run_end_decode run_end_encode diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index d4d44e9347e0..fb0834e3e128 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -2628,6 +2628,106 @@ def test_strftime(): pc.strftime(tsa, options=pc.StrftimeOptions(fmt + "%z")) +def test_month_day_nano_interval_from_parts(): + # Python integers infer int64, including the month and day components. + result = pc.month_day_nano_interval_from_parts( + [1, -2, 0], [32, 3, -4], [90_000_000_000_000, -5, 6]) + expected = pa.array( + [(1, 32, 90_000_000_000_000), (-2, 3, -5), (0, -4, 6)], + type=pa.month_day_nano_interval()) + assert result.equals(expected) + + result = pc.month_day_nano_interval_from_parts( + pa.array([1, 2], type=pa.uint16()), + pa.array([3, -4], type=pa.int8()), + pa.array([2**32, 5], type=pa.uint64())) + assert result.equals(pa.array( + [(1, 3, 2**32), (2, -4, 5)], type=pa.month_day_nano_interval())) + + +def test_month_day_nano_interval_from_parts_scalars(): + result = pc.month_day_nano_interval_from_parts(1, -2, 3) + assert result.equals(pa.scalar( + (1, -2, 3), type=pa.month_day_nano_interval())) + + result = pc.month_day_nano_interval_from_parts(1, [2, None, -3], 4) + assert result.equals(pa.array( + [(1, 2, 4), None, (1, -3, 4)], type=pa.month_day_nano_interval())) + + result = pc.month_day_nano_interval_from_parts(None, 2, 3) + assert result.equals(pa.scalar(None, type=pa.month_day_nano_interval())) + + result = pc.month_day_nano_interval_from_parts( + [1, 2], pa.scalar(None, type=pa.int32()), 3) + assert result.equals(pa.array( + [None, None], type=pa.month_day_nano_interval())) + + +def test_month_day_nano_interval_from_parts_nulls(): + result = pc.month_day_nano_interval_from_parts( + [1, None, 3, 4], [2, 3, None, 5], [3, 4, 5, None]) + assert result.equals(pa.array( + [(1, 2, 3), None, None, None], type=pa.month_day_nano_interval())) + + +def test_month_day_nano_interval_from_parts_sliced_and_chunked(): + months = pa.array([99, 1, -2, 3, 99]).slice(1, 3) + days = pa.array([99, 4, 5, -6, 99]).slice(1, 3) + nanoseconds = pa.array([99, 7, -8, 9, 99]).slice(1, 3) + expected = pa.array( + [(1, 4, 7), (-2, 5, -8), (3, -6, 9)], + type=pa.month_day_nano_interval()) + result = pc.month_day_nano_interval_from_parts(months, days, nanoseconds) + assert result.equals(expected) + + # Chunk boundaries need not align between the component inputs. + months = pa.chunked_array([months[:1], months[1:]]) + days = pa.chunked_array([days[:2], days[2:]]) + result = pc.month_day_nano_interval_from_parts(months, days, nanoseconds) + assert result.combine_chunks().equals(expected) + + +@pytest.mark.parametrize("chunked", [False, True]) +def test_month_day_nano_interval_from_parts_empty(chunked): + months = (pa.chunked_array([], type=pa.int32()) if chunked + else pa.array([], type=pa.int32())) + result = pc.month_day_nano_interval_from_parts(months, 2, 3) + assert result.type == pa.month_day_nano_interval() + assert len(result) == 0 + + +def test_month_day_nano_interval_from_parts_boundaries(): + result = pc.month_day_nano_interval_from_parts( + [-2**31, 2**31 - 1], [2**31 - 1, -2**31], + [-2**63, 2**63 - 1]) + assert result.equals(pa.array( + [(-2**31, 2**31 - 1, -2**63), + (2**31 - 1, -2**31, 2**63 - 1)], + type=pa.month_day_nano_interval())) + + +@pytest.mark.parametrize("component", [0, 1]) +@pytest.mark.parametrize("value", [-2**31 - 1, 2**31]) +def test_month_day_nano_interval_from_parts_overflow(component, value): + parts = [1, 2, 3] + parts[component] = value + with pytest.raises(pa.ArrowInvalid, match="not in range"): + pc.month_day_nano_interval_from_parts(*parts) + + +def test_month_day_nano_interval_from_parts_unsigned_overflow(): + nanoseconds = pa.array([2**63], type=pa.uint64()) + with pytest.raises(pa.ArrowInvalid, match="not in range"): + pc.month_day_nano_interval_from_parts(1, 2, nanoseconds) + + +@pytest.mark.parametrize("months", [1.0, "1"]) +def test_month_day_nano_interval_from_parts_invalid_type(months): + with pytest.raises(pa.ArrowNotImplementedError, + match="no kernel matching input types"): + pc.month_day_nano_interval_from_parts(months, 2, 3) + + def _check_datetime_components(timestamps, timezone=None): ts = pd.to_datetime(timestamps).tz_localize( "UTC").tz_convert(timezone).to_series()