From f60ba85b6cc0064948d691fca7bbe996aca99afb Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 23 Oct 2026 23:26:11 -0700 Subject: [PATCH] [C++][Compute] Support unsigned indices in inverse permutation and scatter Signed-off-by: Rossi Sun --- cpp/src/arrow/compute/api_vector.h | 6 +- .../arrow/compute/kernels/vector_swizzle.cc | 171 +++++++++++++----- .../compute/kernels/vector_swizzle_test.cc | 150 +++++++++++---- docs/source/cpp/compute.rst | 2 +- python/pyarrow/_compute.pyx | 2 +- 5 files changed, 247 insertions(+), 84 deletions(-) diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index b55ed9e7be49..9896520c9f5a 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -408,7 +408,7 @@ class ARROW_EXPORT InversePermutationOptions : public FunctionOptions { /// length of the input indices. int64_t max_index = -1; /// \brief The data type for the output array of inverse permutation. Defaults to the - /// type of the input indices when `nullopt`. Must be a signed integer type. An + /// type of the input indices when `nullopt`. Must be an integer type. An /// invalid error will be reported if this type is not able to store the length of the /// input indices. std::optional> output_type; @@ -912,7 +912,7 @@ Result> PairwiseDiff(const Array& array, /// [1, 6, 3, null, 4, null, null] /// if max_index = 6. /// -/// \param[in] indices array-like indices +/// \param[in] indices array-like signed or unsigned integer indices /// \param[in] options configures the max index and the output type /// \param[in] ctx the function execution context, optional /// \return the resulting inverse permutation @@ -940,7 +940,7 @@ Result InversePermutation( /// if max_index = 6. /// /// \param[in] values datum to scatter -/// \param[in] indices array-like indices +/// \param[in] indices array-like signed or unsigned integer indices /// \param[in] options configures the max index of to scatter /// \param[in] ctx the function execution context, optional /// \return the resulting datum diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle.cc b/cpp/src/arrow/compute/kernels/vector_swizzle.cc index cf9f5379a694..1aca196794bd 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle.cc @@ -15,13 +15,18 @@ // specific language governing permissions and limitations // under the License.#include "arrow/compute/api_vector.h" +#include +#include + #include "arrow/compute/api_vector.h" #include "arrow/compute/function.h" #include "arrow/compute/kernels/codegen_internal.h" #include "arrow/compute/registry.h" #include "arrow/compute/registry_internal.h" +#include "arrow/util/bitmap_generate.h" #include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" +#include "arrow/util/macros.h" namespace arrow::compute::internal { @@ -44,8 +49,7 @@ const InversePermutationOptions* GetDefaultInversePermutationOptions() { using InversePermutationState = OptionsWrapper; /// Resolve the output type of inverse_permutation. The output type is specified in the -/// options, and if null, set it to the input type. The output type must be signed -/// integer. +/// options, and if null, set it to the input type. The output type must be integer. Result ResolveInversePermutationOutputType( KernelContext* ctx, const std::vector& input_types) { DCHECK_EQ(input_types.size(), 1); @@ -53,15 +57,56 @@ Result ResolveInversePermutationOutputType( std::shared_ptr output_type = InversePermutationState::Get(ctx).output_type.value_or(input_types[0].owned_type); - if (!is_signed_integer(output_type->id())) { - return Status::TypeError( - "Output type of inverse_permutation must be signed integer, got " + - output_type->ToString()); + if (!is_integer(output_type->id())) { + return Status::TypeError("Output type of inverse_permutation must be integer, got " + + output_type->ToString()); } return TypeHolder(std::move(output_type)); } +// Share the input implementation for signed and unsigned integers of the same width. +// The logical types still determine bounds checking and the output array type. +template +using UnsignedIntegerType = + typename CTypeTraits>::ArrowType; + +// This pass depends only on the output width, not the index type or input shape. +// Keep it separate from the indexed writes, and generate whole bitmap bytes +// instead of clearing one bit at a time for every missing destination. +template +ARROW_NOINLINE Status InversePermutationValidity(KernelContext* ctx, + const OutputCType* data, + int64_t input_length, + int64_t output_length, + std::shared_ptr* validity_buf) { + if (output_length == 0) { + return Status::OK(); + } + const auto invalid = static_cast(input_length); + const auto* first_invalid = std::find(data, data + output_length, invalid); + const int64_t offset = first_invalid - data; + if (offset == output_length) { + return Status::OK(); + } + ARROW_ASSIGN_OR_RAISE(*validity_buf, ctx->AllocateBitmap(output_length)); + auto* validity = (*validity_buf)->mutable_data(); + bit_util::SetBitmap(validity, 0, offset); + ::arrow::internal::GenerateBitsUnrolled(validity, offset, output_length - offset, + [&]() { return *first_invalid++ != invalid; }); + return Status::OK(); +} + +template +ARROW_NOINLINE Status InversePermutationIndexError(IndexCType index, + bool indices_are_signed) { + using SignedIndexCType = std::make_signed_t; + return Status::IndexError("Index out of bounds: ", + indices_are_signed + ? std::to_string(static_cast(index)) + : std::to_string(index)); +} + template struct InversePermutationImpl { using ThisType = InversePermutationImpl; @@ -78,7 +123,8 @@ struct InversePermutationImpl { int64_t output_length = options.max_index < 0 ? input_length : options.max_index + 1; std::shared_ptr output_type = options.output_type.value_or(input_type); - ThisType impl(ctx, indices, input_length, output_length); + ThisType impl(ctx, indices, input_length, output_length, + is_signed_integer(input_type->id())); RETURN_NOT_OK(VisitTypeInline(*output_type, &impl)); return ArrayData::Make(std::move(output_type), output_length, @@ -87,6 +133,8 @@ struct InversePermutationImpl { template enable_if_t::value, Status> Visit(const Type& output_type) { + // Keep output execution specialized so the loops remain inlineable. Sharing + // output signedness can cause aggressive outlining and regress throughput. using OutputCType = typename Type::c_type; RETURN_NOT_OK(CheckInput(output_type)); @@ -97,7 +145,7 @@ struct InversePermutationImpl { // contents - even if it is shadowed by the validity bit). The subsequent processing // will fill the valid values only. // - Otherwise (i.e. the output is "dense"), the validity buffer is lazily allocated - // and initialized all-true in the subsequent processing only when needed. The data + // and populated from the data in the subsequent processing only when needed. The data // buffer is preallocated and filled with "impossible" values (that is, input_length - // note that the range of inverse_permutation is [0, input_length)) for the subsequent // processing to detect validity. @@ -122,22 +170,26 @@ struct InversePermutationImpl { const ShapeType& indices_; const int64_t input_length_; const int64_t output_length_; + const bool indices_are_signed_; std::shared_ptr validity_buf_; std::shared_ptr data_buf_; InversePermutationImpl(KernelContext* ctx, const ShapeType& indices, - int64_t input_length, int64_t output_length) + int64_t input_length, int64_t output_length, + bool indices_are_signed) : ctx_(ctx), indices_(indices), input_length_(input_length), - output_length_(output_length) {} + output_length_(output_length), + indices_are_signed_(indices_are_signed) {} template Status CheckInput(const Type& output_type) { using OutputCType = typename Type::c_type; - if (static_cast(std::numeric_limits::max()) < input_length_) { + if (static_cast(std::numeric_limits::max()) < + static_cast(input_length_)) { return Status::Invalid( "Output type " + output_type.ToString() + " of inverse_permutation is insufficient to store indices of length " + @@ -203,12 +255,20 @@ struct InversePermutationImpl { DCHECK_NE(data_buf_, nullptr); OutputCType* data = data_buf_->mutable_data_as(); int64_t inverse = 0; + using SignedIndexCType = std::make_signed_t; + // Negative signed indices have their high bit set when read as unsigned. They + // must be rejected even when their unsigned representation fits output_length_. + const uint64_t index_limit = + indices_are_signed_ + ? std::min( + output_length_, + uint64_t{std::numeric_limits::max()} + 1) + : static_cast(output_length_); RETURN_NOT_OK(ExecType::VisitIndices( indices_, - [&](IndexCType index) { - if (ARROW_PREDICT_FALSE(index < 0 || - static_cast(index) >= output_length_)) { - return Status::IndexError("Index out of bounds: ", std::to_string(index)); + [=, &inverse, indices_are_signed = indices_are_signed_](IndexCType index) { + if (ARROW_PREDICT_FALSE(static_cast(index) >= index_limit)) { + return InversePermutationIndexError(index, indices_are_signed); } data[index] = static_cast(inverse); // If many nulls, set validity to true for valid values. @@ -223,19 +283,14 @@ struct InversePermutationImpl { return Status::OK(); })); - // If not many nulls, run another pass iterating over the data to set the validity - // to false if the value is "impossible". The validity buffer is on demand allocated - // and initialized all-true when the first "impossible" value is seen. + // If not many nulls, run another pass iterating over the data to determine validity + // from the "impossible" values. The validity buffer is allocated on demand + // only when an "impossible" value is seen. if constexpr (!likely_many_nulls) { - for (int64_t i = 0; i < output_length_; ++i) { - if (data[i] == static_cast(input_length_)) { - if (ARROW_PREDICT_FALSE(!validity_buf_)) { - RETURN_NOT_OK(AllocateValidityBufAndFill(true)); - validity = validity_buf_->mutable_data_as(); - } - bit_util::SetBitTo(validity, i, false); - } - } + using UnsignedOutputCType = std::make_unsigned_t; + RETURN_NOT_OK( + InversePermutationValidity(ctx_, data_buf_->data_as(), + input_length_, output_length_, &validity_buf_)); } return Status::OK(); @@ -251,6 +306,16 @@ struct InversePermutation { template static Status VisitIndices(const ArraySpan& span, ValidFunc&& valid_func, NullFunc&& null_func) { + // Avoid per-block validity dispatch when every input index is valid. + if (!span.MayHaveNulls()) { + if (span.length != 0) { + const auto* indices = span.GetValues(1); + for (int64_t i = 0; i < span.length; ++i) { + RETURN_NOT_OK(valid_func(indices[i])); + } + } + return Status::OK(); + } return VisitArraySpanInline(span, std::forward(valid_func), std::forward(null_func)); } @@ -275,10 +340,11 @@ struct InversePermutationChunked { template static Status VisitIndices(const std::shared_ptr& chunked_array, ValidFunc&& valid_func, NullFunc&& null_func) { + using ArrayExec = InversePermutation; for (const auto& chunk : chunked_array->chunks()) { ArraySpan span(*chunk->data()); - RETURN_NOT_OK(VisitArraySpanInline( - span, std::forward(valid_func), std::forward(null_func))); + RETURN_NOT_OK(ArrayExec::VisitIndices(span, std::forward(valid_func), + std::forward(null_func))); } return Status::OK(); } @@ -296,28 +362,37 @@ struct InversePermutationChunked { } }; +template +using UnsignedInversePermutation = InversePermutation>; + +template +using UnsignedInversePermutationChunked = + InversePermutationChunked>; + void RegisterVectorInversePermutation(FunctionRegistry* registry) { auto function = std::make_shared("inverse_permutation", Arity::Unary(), inverse_permutation_doc, GetDefaultInversePermutationOptions()); - auto add_kernel = [&function](Type::type type_id) { - VectorKernel kernel; - kernel.signature = - KernelSignature::Make({InputType(match::SameTypeId(type_id))}, - OutputType(ResolveInversePermutationOutputType)); - kernel.init = InversePermutationState::Init; - kernel.exec = GenerateInteger(type_id); - kernel.exec_chunked = - GenerateInteger( - type_id); - kernel.can_execute_chunkwise = false; - kernel.output_chunked = false; - DCHECK_OK(function->AddKernel(std::move(kernel))); + // Exact dispatch scans every kernel. Match integers once, then select the + // shared physical implementation without increasing per-call matching costs. + VectorKernel kernel; + kernel.signature = KernelSignature::Make( + {InputType(match::Integer())}, OutputType(ResolveInversePermutationOutputType)); + kernel.init = InversePermutationState::Init; + kernel.exec = [](KernelContext* ctx, const ExecSpan& span, ExecResult* result) { + const auto exec = GenerateInteger( + span[0].type()->id()); + return exec(ctx, span, result); }; - for (const auto& t : SignedIntTypes()) { - add_kernel(t->id()); - } + kernel.exec_chunked = [](KernelContext* ctx, const ExecBatch& batch, Datum* result) { + const auto exec = GenerateInteger(batch[0].type()->id()); + return exec(ctx, batch, result); + }; + kernel.can_execute_chunkwise = false; + kernel.output_chunked = false; + DCHECK_OK(function->AddKernel(std::move(kernel))); DCHECK_OK(registry->AddFunction(std::move(function))); } @@ -363,8 +438,8 @@ class ScatterMetaFunction : public MetaFunction { "Input and indices of scatter must have the same length, got " + std::to_string(values.length()) + " and " + std::to_string(indices.length())); } - if (!is_signed_integer(indices.type()->id())) { - return Status::TypeError("Indices of scatter must be of signed integer type, got ", + if (!is_integer(indices.type()->id())) { + return Status::TypeError("Indices of scatter must be of integer type, got ", indices.type()->ToString()); } // Internally invoke Take(values, InversePermutation(indices)) to implement scatter. @@ -377,7 +452,7 @@ class ScatterMetaFunction : public MetaFunction { // [b, g, d, c, e, null, null] if max_index = 6. InversePermutationOptions inverse_permutation_options{ scatter_options->max_index, - // Use the smallest possible uint type to store inverse permutation. + // Use the smallest possible signed integer type to store inverse permutation. InferSmallestInversePermutationType(values.length())}; ARROW_ASSIGN_OR_RAISE(auto inverse_permutation, CallFunction("inverse_permutation", {indices}, diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc index 22b78a016d93..35847dea8c8f 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc @@ -30,7 +30,7 @@ namespace arrow::compute { namespace { -using SmallSignedIntegerTypes = ::testing::Types; +using SmallIntegerTypes = ::testing::Types; } // namespace @@ -120,29 +120,54 @@ void TestInversePermutationForInputOutputTypes( void TestInversePermutation(const std::vector& indices_chunked_str, int64_t max_index, const std::string& expected_str, bool validity_must_be_null = false) { - TestInversePermutationForInputOutputTypes(SignedIntTypes(), SignedIntTypes(), - indices_chunked_str, max_index, expected_str, + TestInversePermutationForInputOutputTypes(IntTypes(), IntTypes(), indices_chunked_str, + max_index, expected_str, validity_must_be_null); } } // namespace -TEST(InversePermutation, InvalidOutputType) { - { - ARROW_SCOPED_TRACE("Output type unsigned"); - auto indices = ArrayFromJSON(int32(), "[]"); - ASSERT_RAISES_WITH_MESSAGE( - TypeError, - "Type error: Output type of inverse_permutation must be signed integer, got " - "uint32", - InversePermutation(indices, /*max_index=*/0, /*output_type=*/uint32())); +TEST(InversePermutation, InvalidInputType) { + // The integer matcher must not accept non-integers of the same width. + for (const auto& type : {boolean(), float32(), float64(), utf8()}) { + ARROW_SCOPED_TRACE(type->ToString()); + ASSERT_RAISES(NotImplemented, InversePermutation(ArrayFromJSON(type, "[null]"))); + } +} + +TEST(InversePermutation, DenseValidityBitmapBoundaries) { + for (int64_t length : {7, 8, 9, 63, 64, 65}) { + for (int64_t missing : {int64_t{0}, int64_t{1}, int64_t{2}, int64_t{3}, int64_t{4}, + int64_t{5}, int64_t{6}, int64_t{7}, length - 1}) { + if (missing >= length) continue; + ARROW_SCOPED_TRACE("Length: " + std::to_string(length) + + ", missing destination: " + std::to_string(missing)); + std::string indices = "["; + std::string expected = "["; + for (int64_t i = 0; i < length; ++i) { + if (i != 0) expected += ","; + if (i == missing) { + expected += "null"; + } else { + if (indices.size() > 1) indices += ","; + indices += std::to_string(i); + expected += std::to_string(i < missing ? i : i - 1); + } + } + indices += "]"; + expected += "]"; + TestInversePermutation({indices}, /*max_index=*/length - 1, expected); + } } +} + +TEST(InversePermutation, InvalidOutputType) { { ARROW_SCOPED_TRACE("Output type float"); auto indices = ArrayFromJSON(int32(), "[]"); ASSERT_RAISES_WITH_MESSAGE( TypeError, - "Type error: Output type of inverse_permutation must be signed integer, got " + "Type error: Output type of inverse_permutation must be integer, got " "float", InversePermutation(indices, /*max_index=*/0, /*output_type=*/float32())); } @@ -151,7 +176,7 @@ TEST(InversePermutation, InvalidOutputType) { auto indices = ArrayFromJSON(int32(), "[]"); ASSERT_RAISES_WITH_MESSAGE( TypeError, - "Type error: Output type of inverse_permutation must be signed integer, got " + "Type error: Output type of inverse_permutation must be integer, got " "string", InversePermutation(indices, /*max_index=*/0, /*output_type=*/utf8())); } @@ -167,7 +192,7 @@ TEST(InversePermutation, DefaultOptions) { } { ARROW_SCOPED_TRACE("Default options semantics"); - for (const auto& input_type : SignedIntTypes()) { + for (const auto& input_type : IntTypes()) { ARROW_SCOPED_TRACE("Input type: " + input_type->ToString()); auto indices = ArrayFromJSON(input_type, "[0]"); ASSERT_OK_AND_ASSIGN(Datum result, InversePermutation(indices)); @@ -192,6 +217,48 @@ TEST(InversePermutation, InvalidIndex) { } } +TEST(InversePermutation, IndexBounds) { + for (const auto& type : IntTypes()) { + ARROW_SCOPED_TRACE(type->ToString()); + const std::string invalid_index = is_signed_integer(type->id()) ? "-1" : "42"; + for (const Datum& indices : + {Datum(ArrayFromJSON(type, "[" + invalid_index + "]")), + Datum(ChunkedArrayFromJSON(type, {"[]", "[" + invalid_index + "]"}))}) { + ASSERT_RAISES_WITH_MESSAGE( + IndexError, "Index error: Index out of bounds: " + invalid_index, + InversePermutation(indices, /*max_index=*/1, /*output_type=*/int64())); + } + } + // Reading signed indices through their unsigned representation must still reject + // negative values, even if that representation would fit in the output. + for (const auto& indices : + {ArrayFromJSON(int8(), "[-1]"), ArrayFromJSON(int8(), "[-128]")}) { + ASSERT_RAISES(IndexError, InversePermutation(indices, /*max_index=*/300, + /*output_type=*/int64())); + } + // Do not narrow uint64 indices to int64 before checking bounds. + for (const std::string index : {"9223372036854775808", "18446744073709551615"}) { + for (const Datum& indices : + {Datum(ArrayFromJSON(uint64(), "[" + index + "]")), + Datum(ChunkedArrayFromJSON(uint64(), {"[]", "[" + index + "]"}))}) { + ASSERT_RAISES_WITH_MESSAGE( + IndexError, "Index error: Index out of bounds: " + index, + InversePermutation(indices, /*max_index=*/1, /*output_type=*/int64())); + } + } +} + +TEST(InversePermutation, SlicedIndices) { + for (const auto& type : IntTypes()) { + auto indices = ArrayFromJSON(type, "[42, 2, null, 0, 1, 42]")->Slice(1, 4); + auto expected = ArrayFromJSON(type, "[2, 3, 0, null]"); + AssertInversePermutation(indices, -1, type, expected, false); + auto chunked = std::make_shared( + ArrayVector{indices->Slice(0, 1), indices->Slice(1)}); + AssertInversePermutation(chunked, -1, type, expected, false); + } +} + TEST(InversePermutation, Basic) { { ARROW_SCOPED_TRACE("Basic"); @@ -251,7 +318,7 @@ class TestInversePermutationSmallOutputType : public ::testing::Test { } }; -TYPED_TEST_SUITE(TestInversePermutationSmallOutputType, SmallSignedIntegerTypes); +TYPED_TEST_SUITE(TestInversePermutationSmallOutputType, SmallIntegerTypes); TYPED_TEST(TestInversePermutationSmallOutputType, JustEnoughOutputType) { auto output_type = this->type_singleton(); @@ -260,7 +327,7 @@ TYPED_TEST(TestInversePermutationSmallOutputType, JustEnoughOutputType) { auto expected = ArrayFromJSON(output_type, "[" + std::to_string(input_length - 1) + "]"); DoTestInversePermutationForInputTypes( - SignedIntTypes(), + IntTypes(), [&](const std::shared_ptr& input_type) -> Result { return ConstantArrayGenerator::Zeroes(input_length, input_type); }, @@ -271,7 +338,7 @@ TYPED_TEST(TestInversePermutationSmallOutputType, InsufficientOutputType) { auto output_type = this->type_singleton(); int64_t input_length = static_cast(std::numeric_limits::max()) + 1; - for (const auto& input_type : SignedIntTypes()) { + for (const auto& input_type : IntTypes()) { ARROW_SCOPED_TRACE("Input type: " + input_type->ToString()); auto indices = ConstantArrayGenerator::Zeroes(input_length, input_type); ASSERT_RAISES_WITH_MESSAGE( @@ -440,7 +507,7 @@ void DoTestScatterForIndicesTypes( void DoTestScatter(const std::shared_ptr& values, const std::shared_ptr& indices, int64_t max_index, const std::shared_ptr& expected) { - DoTestScatterForIndicesTypes(SignedIntTypes(), values, indices, max_index, expected); + DoTestScatterForIndicesTypes(IntTypes(), values, indices, max_index, expected); } void TestScatter(const std::shared_ptr& value_type, @@ -468,19 +535,17 @@ TEST(Scatter, Invalid) { ARROW_SCOPED_TRACE("Invalid input type"); auto values = ArrayFromJSON(int32(), "[0]"); { - ARROW_SCOPED_TRACE("uint32"); - auto indices = ArrayFromJSON(uint32(), R"([0])"); + ARROW_SCOPED_TRACE("float"); + auto indices = ArrayFromJSON(float32(), R"([0])"); ASSERT_RAISES_WITH_MESSAGE( - TypeError, - "Type error: Indices of scatter must be of signed integer type, got uint32", + TypeError, "Type error: Indices of scatter must be of integer type, got float", Scatter(values, indices)); } { ARROW_SCOPED_TRACE("string"); auto indices = ArrayFromJSON(utf8(), R"(["a"])"); ASSERT_RAISES_WITH_MESSAGE( - TypeError, - "Type error: Indices of scatter must be of signed integer type, got string", + TypeError, "Type error: Indices of scatter must be of integer type, got string", Scatter(values, indices)); } } @@ -495,7 +560,7 @@ TEST(Scatter, DefaultOptions) { { ARROW_SCOPED_TRACE("Default options semantics"); auto values = ArrayFromJSON(utf8(), R"(["a"])"); - for (const auto& indices_type : SignedIntTypes()) { + for (const auto& indices_type : IntTypes()) { ARROW_SCOPED_TRACE("Indices type: " + indices_type->ToString()); auto indices = ArrayFromJSON(indices_type, "[0]"); ASSERT_OK_AND_ASSIGN(Datum result, Scatter(values, indices)); @@ -521,6 +586,30 @@ TEST(Scatter, InvalidIndex) { } } +TEST(Scatter, UnsignedIndexBounds) { + auto values = ArrayFromJSON(utf8(), R"(["a"])"); + for (const std::string index : {"9223372036854775808", "18446744073709551615"}) { + for (const Datum& indices : + {Datum(ArrayFromJSON(uint64(), "[" + index + "]")), + Datum(ChunkedArrayFromJSON(uint64(), {"[]", "[" + index + "]"}))}) { + ASSERT_RAISES_WITH_MESSAGE(IndexError, "Index error: Index out of bounds: " + index, + Scatter(values, indices, /*max_index=*/1)); + } + } +} + +TEST(Scatter, SlicedIndices) { + auto values = ArrayFromJSON(utf8(), R"(["a", "b", "c", "d"])"); + auto expected = ArrayFromJSON(utf8(), R"(["c", "d", "a", null])"); + for (const auto& type : IntTypes()) { + auto indices = ArrayFromJSON(type, "[42, 2, null, 0, 1, 42]")->Slice(1, 4); + AssertScatterAAA(values, indices, -1, expected); + auto chunked = std::make_shared( + ArrayVector{indices->Slice(0, 1), indices->Slice(1)}); + AssertScatterACC(values, chunked, -1, expected); + } +} + template class TestScatterSmallIndicesTypes : public ::testing::Test { protected: @@ -531,20 +620,19 @@ class TestScatterSmallIndicesTypes : public ::testing::Test { } }; -TYPED_TEST_SUITE(TestScatterSmallIndicesTypes, SmallSignedIntegerTypes); +TYPED_TEST_SUITE(TestScatterSmallIndicesTypes, SmallIntegerTypes); TYPED_TEST(TestScatterSmallIndicesTypes, MaxIntegerIndex) { auto values = ArrayFromJSON(utf8(), R"(["a"])"); auto indices_type = this->type_singleton(); int64_t max_integer = static_cast(std::numeric_limits::max()); - auto indices = ArrayFromJSON(indices_type, "[" + std::to_string(max_integer - 1) + "]"); - ASSERT_OK_AND_ASSIGN(auto expected_prefix_nulls, - MakeArrayOfNull(utf8(), max_integer - 1)); + auto indices = ArrayFromJSON(indices_type, "[" + std::to_string(max_integer) + "]"); + ASSERT_OK_AND_ASSIGN(auto expected_prefix_nulls, MakeArrayOfNull(utf8(), max_integer)); auto expected_suffix_value = values; ASSERT_OK_AND_ASSIGN(auto expected, Concatenate({expected_prefix_nulls, expected_suffix_value})); - DoTestScatterAAA(values, indices, /*max_index=*/max_integer - 1, expected); + DoTestScatterAAA(values, indices, /*max_index=*/max_integer, expected); } TEST(Scatter, Boolean) { diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index 15d426dca6c6..2aab3dcee3ec 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -1831,7 +1831,7 @@ These functions select and return a subset of their input. +---------------------+--------+----------------+--------------+---------------------+-------------------------------------+-------+ | filter | Binary | Any | Boolean | Input type 1 | :struct:`FilterOptions` | \(2) | +---------------------+--------+----------------+--------------+---------------------+-------------------------------------+-------+ -| inverse_permutation | Unary | Signed Integer | | Signed Integer \(4) | :struct:`InversePermutationOptions` | \(5) | +| inverse_permutation | Unary | Integer | | Integer \(4) | :struct:`InversePermutationOptions` | \(5) | +---------------------+--------+----------------+--------------+---------------------+-------------------------------------+-------+ | scatter | Binary | Any | Integer | Input type 1 | :struct:`ScatterOptions` | \(6) | +---------------------+--------+----------------+--------------+---------------------+-------------------------------------+-------+ diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx index 1c0c779f6037..fd188dfa0192 100644 --- a/python/pyarrow/_compute.pyx +++ b/python/pyarrow/_compute.pyx @@ -1474,7 +1474,7 @@ class InversePermutationOptions(_InversePermutationOptions): output_type : DataType, default None The data type for the output array of inverse permutation. If None, the output will be of the same type as the input indices, otherwise - must be a signed integer type. An invalid error will be reported if this type + must be an integer type. An invalid error will be reported if this type is not able to store the length of the input indices. """