From ae2a70498964c2da02a39f482de349655da467f6 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 23 Oct 2026 23:26:11 -0700 Subject: [PATCH 1/4] [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 | 158 +++++++++++++----- .../compute/kernels/vector_swizzle_test.cc | 150 +++++++++++++---- docs/source/cpp/compute.rst | 2 +- python/pyarrow/_compute.pyx | 2 +- 5 files changed, 236 insertions(+), 82 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 4fae37c0fc54..5827da3dbdaa 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,50 @@ 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)); } +// 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 +117,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 +127,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 +139,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 +164,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 +249,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 +277,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 +300,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 +334,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(); } @@ -301,21 +361,27 @@ void RegisterVectorInversePermutation(FunctionRegistry* registry) { 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 = + GenerateKernelExec( + 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 = + GenerateKernelExec( + 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))); } @@ -361,8 +427,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. @@ -375,7 +441,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 c3c182ebdddd..a328d7d07655 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 8211c74ca39c..ee3ebd6d0e9a 100644 --- a/python/pyarrow/_compute.pyx +++ b/python/pyarrow/_compute.pyx @@ -1473,7 +1473,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. """ From 3a5c01e3eb1a343365390e926c020ab5365d9b73 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 30 Oct 2026 14:26:32 -0700 Subject: [PATCH 2/4] [C++][Compute] Use bitmap runs for inverse permutation traversal Signed-off-by: Rossi Sun --- .../arrow/compute/kernels/vector_swizzle.cc | 30 ++++++++++++------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle.cc b/cpp/src/arrow/compute/kernels/vector_swizzle.cc index 5827da3dbdaa..1277e0cddba8 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle.cc @@ -23,6 +23,7 @@ #include "arrow/compute/kernels/codegen_internal.h" #include "arrow/compute/registry.h" #include "arrow/compute/registry_internal.h" +#include "arrow/util/bit_run_reader.h" #include "arrow/util/bitmap_generate.h" #include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" @@ -85,7 +86,7 @@ ARROW_NOINLINE Status InversePermutationValidity(KernelContext* ctx, } ARROW_ASSIGN_OR_RAISE(*validity_buf, ctx->AllocateBitmap(output_length)); auto* validity = (*validity_buf)->mutable_data(); - bit_util::SetBitmap(validity, 0, offset); + bit_util::SetBitmap(validity, /*offset=*/0, /*length=*/offset); ::arrow::internal::GenerateBitsUnrolled(validity, offset, output_length - offset, [&]() { return *first_invalid++ != invalid; }); return Status::OK(); @@ -300,18 +301,25 @@ 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])); - } - } + if (span.length == 0) { return Status::OK(); } - return VisitArraySpanInline(span, std::forward(valid_func), - std::forward(null_func)); + const auto* indices = span.GetValues(1); + const uint8_t* validity = span.MayHaveNulls() ? span.buffers[0].data : nullptr; + return ::arrow::internal::VisitBitRuns( + validity, span.offset, span.length, + [&](int64_t start, int64_t length, bool valid) { + if (valid) { + for (int64_t i = start; i < start + length; ++i) { + RETURN_NOT_OK(valid_func(indices[i])); + } + } else { + for (int64_t i = 0; i < length; ++i) { + RETURN_NOT_OK(null_func()); + } + } + return Status::OK(); + }); } static Status Exec(KernelContext* ctx, const ExecSpan& span, ExecResult* result) { From 4580b180b1447697adc8ecf72a8f3fde07705f46 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 30 Oct 2026 15:15:56 -0700 Subject: [PATCH 3/4] [C++][Compute] Refine inverse permutation helpers and tests Signed-off-by: Rossi Sun --- .../arrow/compute/kernels/vector_swizzle.cc | 25 +++++++-------- .../compute/kernels/vector_swizzle_test.cc | 31 ++++++++++++------- 2 files changed, 31 insertions(+), 25 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle.cc b/cpp/src/arrow/compute/kernels/vector_swizzle.cc index 1277e0cddba8..15a952256b5c 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle.cc @@ -13,7 +13,7 @@ // "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/api_vector.h" +// under the License. #include #include @@ -70,26 +70,24 @@ Result ResolveInversePermutationOutputType( // 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) { +ARROW_NOINLINE Result> InversePermutationValidity( + KernelContext* ctx, const OutputCType* data, int64_t input_length, + int64_t output_length) { if (output_length == 0) { - return Status::OK(); + return nullptr; } 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(); + return nullptr; } - ARROW_ASSIGN_OR_RAISE(*validity_buf, ctx->AllocateBitmap(output_length)); - auto* validity = (*validity_buf)->mutable_data(); + ARROW_ASSIGN_OR_RAISE(auto validity_buf, ctx->AllocateBitmap(output_length)); + auto* validity = validity_buf->mutable_data(); bit_util::SetBitmap(validity, /*offset=*/0, /*length=*/offset); ::arrow::internal::GenerateBitsUnrolled(validity, offset, output_length - offset, [&]() { return *first_invalid++ != invalid; }); - return Status::OK(); + return validity_buf; } template @@ -283,9 +281,10 @@ struct InversePermutationImpl { // only when an "impossible" value is seen. if constexpr (!likely_many_nulls) { using UnsignedOutputCType = std::make_unsigned_t; - RETURN_NOT_OK( + ARROW_ASSIGN_OR_RAISE( + validity_buf_, InversePermutationValidity(ctx_, data_buf_->data_as(), - input_length_, output_length_, &validity_buf_)); + input_length_, output_length_)); } return Status::OK(); diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc index 35847dea8c8f..7be22d307055 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc @@ -136,6 +136,9 @@ TEST(InversePermutation, InvalidInputType) { } TEST(InversePermutation, DenseValidityBitmapBoundaries) { + // Leave one destination missing and check the inverse values and validity bitmap. + // For length = 8, missing = 3: [0, 1, 2, 4, 5, 6, 7] -> [0, 1, 2, null, 3, 4, 5, 6]. + // Vary output length and the first null bit around byte and word boundaries. 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}) { @@ -213,7 +216,7 @@ TEST(InversePermutation, InvalidIndex) { auto indices = ArrayFromJSON(int32(), "[42]"); ASSERT_RAISES_WITH_MESSAGE( IndexError, "Index error: Index out of bounds: 42", - InversePermutation(indices, /*max_index=*/1, /*output_type=*/int32())); + InversePermutation(indices, /*max_index=*/41, /*output_type=*/int32())); } } @@ -226,7 +229,7 @@ TEST(InversePermutation, IndexBounds) { 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())); + InversePermutation(indices, /*max_index=*/41, /*output_type=*/int64())); } } // Reading signed indices through their unsigned representation must still reject @@ -250,12 +253,15 @@ TEST(InversePermutation, IndexBounds) { TEST(InversePermutation, SlicedIndices) { for (const auto& type : IntTypes()) { - auto indices = ArrayFromJSON(type, "[42, 2, null, 0, 1, 42]")->Slice(1, 4); + auto indices = + ArrayFromJSON(type, "[42, 2, null, 0, 1, 42]")->Slice(/*offset=*/1, /*length=*/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); + AssertInversePermutation(indices, /*max_index=*/-1, type, expected, + /*validity_must_be_null=*/false); + auto chunked = std::make_shared(ArrayVector{ + indices->Slice(/*offset=*/0, /*length=*/1), indices->Slice(/*offset=*/1)}); + AssertInversePermutation(chunked, /*max_index=*/-1, type, expected, + /*validity_must_be_null=*/false); } } @@ -602,11 +608,12 @@ 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); + auto indices = + ArrayFromJSON(type, "[42, 2, null, 0, 1, 42]")->Slice(/*offset=*/1, /*length=*/4); + AssertScatterAAA(values, indices, /*max_index=*/-1, expected); + auto chunked = std::make_shared(ArrayVector{ + indices->Slice(/*offset=*/0, /*length=*/1), indices->Slice(/*offset=*/1)}); + AssertScatterACC(values, chunked, /*max_index=*/-1, expected); } } From 8b0a685da5dd6731269621b5a64daa7e1f6629cc Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 2 Nov 2026 19:51:03 -0800 Subject: [PATCH 4/4] [C++][Compute] Consolidate inverse permutation bounds tests --- .../arrow/compute/kernels/vector_swizzle.cc | 2 -- .../compute/kernels/vector_swizzle_test.cc | 34 ++++++++----------- 2 files changed, 15 insertions(+), 21 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle.cc b/cpp/src/arrow/compute/kernels/vector_swizzle.cc index 15a952256b5c..62b5b3a6c1a8 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle.cc @@ -126,8 +126,6 @@ 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)); diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc index 7be22d307055..d1d65cb471e8 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc @@ -204,32 +204,28 @@ TEST(InversePermutation, DefaultOptions) { } } -TEST(InversePermutation, InvalidIndex) { +TEST(InversePermutation, IndexBounds) { { - ARROW_SCOPED_TRACE("Negative index"); + ARROW_SCOPED_TRACE("Negative index with default options"); auto indices = ArrayFromJSON(int32(), "[-1]"); ASSERT_RAISES_WITH_MESSAGE(IndexError, "Index error: Index out of bounds: -1", InversePermutation(indices)); } - { - ARROW_SCOPED_TRACE("Exceeds max_index"); - auto indices = ArrayFromJSON(int32(), "[42]"); - ASSERT_RAISES_WITH_MESSAGE( - IndexError, "Index error: Index out of bounds: 42", - InversePermutation(indices, /*max_index=*/41, /*output_type=*/int32())); - } -} - -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=*/41, /*output_type=*/int64())); + std::vector invalid_indices{"42"}; + if (is_signed_integer(type->id())) { + invalid_indices.push_back("-1"); + } + for (const auto& invalid_index : invalid_indices) { + ARROW_SCOPED_TRACE(invalid_index); + 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=*/41, /*output_type=*/int64())); + } } } // Reading signed indices through their unsigned representation must still reject