pitrou commented on code in PR #55169:
URL: https://github.com/apache/arrow/pull/55169#discussion_r4388834790
##########
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:
Can we make function parameters more explicit?
```suggestion
AssertInversePermutation(indices, /*xxx=*/ -1, type, expected, /*yyy=*/
false);
```
##########
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:
Return a `Result<std::shared_ptr<Buffer>>` here?
##########
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:
Perhaps max_index=41 to make sure that boundary cases are treated properly?
##########
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:
Instead of special-casing this, perhaps rewrite indices visiting using
`VisitBitRuns`?
##########
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:
Sounds almost like we would like a `GenerateInteger` that takes a
transformation function for the input type?
##########
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:
Can you explain what this function does? Perhaps by giving an example.
##########
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:
`AllocateBitmap` should already zero-initialize the bitmap, so I don't think
this is necessary.
--
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]
