Skip to content

GH-54920: [C++][Compute] Support unsigned indices in inverse_permutation and scatter #55169

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zanmato1984 wants to merge 1 commit into
base: main
Choose a base branch
from
+247 −84
Open

GH-54920: [C++][Compute] Support unsigned indices in inverse_permutation and scatter #55169

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions cpp/src/arrow/compute/api_vector.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::shared_ptr<DataType>> output_type;
Expand Down Expand Up @@ -912,7 +912,7 @@ Result<std::shared_ptr<Array>> 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
Expand Down Expand Up @@ -940,7 +940,7 @@ Result<Datum> 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
Expand Down
171 changes: 123 additions & 48 deletions cpp/src/arrow/compute/kernels/vector_swizzle.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,18 @@
// specific language governing permissions and limitations
// under the License.#include "arrow/compute/api_vector.h"

#include <algorithm>
#include <type_traits>

#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 {

Expand All @@ -44,24 +49,64 @@ const InversePermutationOptions* GetDefaultInversePermutationOptions() {
using InversePermutationState = OptionsWrapper<InversePermutationOptions>;

/// 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<TypeHolder> ResolveInversePermutationOutputType(
KernelContext* ctx, const std::vector<TypeHolder>& input_types) {
DCHECK_EQ(input_types.size(), 1);
DCHECK_NE(input_types[0], nullptr);

std::shared_ptr<DataType> 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 <typename Type>
using UnsignedIntegerType =
typename CTypeTraits<std::make_unsigned_t<typename Type::c_type>>::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 <typename OutputCType>
ARROW_NOINLINE Status InversePermutationValidity(KernelContext* ctx,

@pitrou pitrou Oct 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Return a Result<std::shared_ptr<Buffer>> here?

const OutputCType* data,
int64_t input_length,
int64_t output_length,
std::shared_ptr<Buffer>* validity_buf) {
if (output_length == 0) {
return Status::OK();
}
const auto invalid = static_cast<OutputCType>(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);

@pitrou pitrou Oct 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AllocateBitmap should already zero-initialize the bitmap, so I don't think this is necessary.

::arrow::internal::GenerateBitsUnrolled(validity, offset, output_length - offset,
[&]() { return *first_invalid++ != invalid; });
return Status::OK();
}

template <typename IndexCType>
ARROW_NOINLINE Status InversePermutationIndexError(IndexCType index,
bool indices_are_signed) {
using SignedIndexCType = std::make_signed_t<IndexCType>;
return Status::IndexError("Index out of bounds: ",
indices_are_signed
? std::to_string(static_cast<SignedIndexCType>(index))
: std::to_string(index));
}

template <typename ExecType>
struct InversePermutationImpl {
using ThisType = InversePermutationImpl<ExecType>;
Expand All @@ -78,7 +123,8 @@ struct InversePermutationImpl {
int64_t output_length = options.max_index < 0 ? input_length : options.max_index + 1;
std::shared_ptr<DataType> 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,
Expand All @@ -87,6 +133,8 @@ struct InversePermutationImpl {

template <typename Type>
enable_if_t<is_integer_type<Type>::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));
Expand All @@ -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.
Expand All @@ -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<Buffer> validity_buf_;
std::shared_ptr<Buffer> 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 <typename Type>
Status CheckInput(const Type& output_type) {
using OutputCType = typename Type::c_type;

if (static_cast<int64_t>(std::numeric_limits<OutputCType>::max()) < input_length_) {
if (static_cast<uint64_t>(std::numeric_limits<OutputCType>::max()) <
static_cast<uint64_t>(input_length_)) {
return Status::Invalid(
"Output type " + output_type.ToString() +
" of inverse_permutation is insufficient to store indices of length " +
Expand Down Expand Up @@ -203,12 +255,20 @@ struct InversePermutationImpl {
DCHECK_NE(data_buf_, nullptr);
OutputCType* data = data_buf_->mutable_data_as<OutputCType>();
int64_t inverse = 0;
using SignedIndexCType = std::make_signed_t<IndexCType>;
// 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<uint64_t>(
output_length_,
uint64_t{std::numeric_limits<SignedIndexCType>::max()} + 1)
: static_cast<uint64_t>(output_length_);
RETURN_NOT_OK(ExecType::VisitIndices(
indices_,
[&](IndexCType index) {
if (ARROW_PREDICT_FALSE(index < 0 ||
static_cast<int64_t>(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<uint64_t>(index) >= index_limit)) {
return InversePermutationIndexError(index, indices_are_signed);
}
data[index] = static_cast<OutputCType>(inverse);
// If many nulls, set validity to true for valid values.
Expand All @@ -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<OutputCType>(input_length_)) {
if (ARROW_PREDICT_FALSE(!validity_buf_)) {
RETURN_NOT_OK(AllocateValidityBufAndFill(true));
validity = validity_buf_->mutable_data_as<uint8_t>();
}
bit_util::SetBitTo(validity, i, false);
}
}
using UnsignedOutputCType = std::make_unsigned_t<OutputCType>;
RETURN_NOT_OK(
InversePermutationValidity(ctx_, data_buf_->data_as<UnsignedOutputCType>(),
input_length_, output_length_, &validity_buf_));
}

return Status::OK();
Expand All @@ -251,6 +306,16 @@ struct InversePermutation {
template <typename ValidFunc, typename NullFunc>
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()) {
Comment on lines +309 to +310

@pitrou pitrou Oct 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of special-casing this, perhaps rewrite indices visiting using VisitBitRuns?

if (span.length != 0) {
const auto* indices = span.GetValues<typename IndexType::c_type>(1);
for (int64_t i = 0; i < span.length; ++i) {
RETURN_NOT_OK(valid_func(indices[i]));
}
}
return Status::OK();
}
return VisitArraySpanInline<IndexType>(span, std::forward<ValidFunc>(valid_func),
std::forward<NullFunc>(null_func));
}
Expand All @@ -275,10 +340,11 @@ struct InversePermutationChunked {
template <typename ValidFunc, typename NullFunc>
static Status VisitIndices(const std::shared_ptr<ChunkedArray>& chunked_array,
ValidFunc&& valid_func, NullFunc&& null_func) {
using ArrayExec = InversePermutation<Ignored, Type>;
for (const auto& chunk : chunked_array->chunks()) {
ArraySpan span(*chunk->data());
RETURN_NOT_OK(VisitArraySpanInline<IndexType>(
span, std::forward<ValidFunc>(valid_func), std::forward<NullFunc>(null_func)));
RETURN_NOT_OK(ArrayExec::VisitIndices(span, std::forward<ValidFunc>(valid_func),
std::forward<NullFunc>(null_func)));
}
return Status::OK();
}
Expand All @@ -296,28 +362,37 @@ struct InversePermutationChunked {
}
};

template <typename Ignored, typename Type>
using UnsignedInversePermutation = InversePermutation<Ignored, UnsignedIntegerType<Type>>;

template <typename Ignored, typename Type>
using UnsignedInversePermutationChunked =
InversePermutationChunked<Ignored, UnsignedIntegerType<Type>>;

void RegisterVectorInversePermutation(FunctionRegistry* registry) {
auto function = std::make_shared<VectorFunction>("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<InversePermutation, void, ArrayKernelExec>(type_id);
kernel.exec_chunked =
GenerateInteger<InversePermutationChunked, void, VectorKernel::ChunkedExec>(
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<UnsignedInversePermutation, void, ArrayKernelExec>(

@pitrou pitrou Oct 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds almost like we would like a GenerateInteger that takes a transformation function for the input type?

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<UnsignedInversePermutationChunked, void,
VectorKernel::ChunkedExec>(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)));
}
Expand Down Expand Up @@ -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.
Expand All @@ -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},
Expand Down
Loading
Loading