zanmato1984 commented on code in PR #55169:
URL: https://github.com/apache/arrow/pull/55169#discussion_r4428954799
##########
cpp/src/arrow/compute/kernels/vector_swizzle_test.cc:
##########
@@ -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()));
Review Comment:
Changed to `max_index = 41`, so index `42` is exactly one past the allowed
maximum. The negative-index and large-`uint64` checks are preserved.
##########
cpp/src/arrow/compute/kernels/vector_swizzle_test.cc:
##########
@@ -120,29 +120,54 @@ void TestInversePermutationForInputOutputTypes(
void TestInversePermutation(const std::vector<std::string>&
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}) {
Review Comment:
Added an example in the test. For `length = 8` and `missing = 3`, indices
`[0, 1, 2, 4, 5, 6, 7]` produce `[0, 1, 2, null, 3, 4, 5, 6]`. The selected
lengths and missing positions exercise byte/word boundaries and unaligned
bitmap writes.
##########
cpp/src/arrow/compute/kernels/vector_swizzle.cc:
##########
@@ -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()) {
Review Comment:
Switched to `VisitBitRuns`, removing the separate no-null loop. In isolated
inverse-permutation benchmarks on Apple M4 with Clang/O2, the traversal change
gave about 1.5–1.9× speedup with 50% random nulls, but about 2.7× runtime with
strictly alternating nulls. The unified traversal therefore has a known
performance trade-off.
##########
cpp/src/arrow/compute/kernels/vector_swizzle.cc:
##########
@@ -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,
+ 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);
Review Comment:
This call sets the already-valid prefix `[0, offset)` to **ones**, rather
than zeroing it. The `0` is the starting bit offset. I kept the call and added
`/*offset=*/0` and `/*length=*/offset` to make that distinction explicit.
##########
cpp/src/arrow/compute/kernels/vector_swizzle_test.cc:
##########
@@ -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);
Review Comment:
Added the parameter-name comments to both calls, and checked the other new
tests for similar cases.
##########
cpp/src/arrow/compute/kernels/vector_swizzle.cc:
##########
@@ -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,
Review Comment:
Changed to return `Result<std::shared_ptr<Buffer>>`, removing the output
parameter.
##########
cpp/src/arrow/compute/kernels/vector_swizzle.cc:
##########
@@ -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>(
Review Comment:
That looks reasonable. I'll review #55291 and adapt this code to the shared
generator once it lands.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
