From 2b3fbf95ba5fafac93ef79bb6a9d525eecfc8ebd Mon Sep 17 00:00:00 2001 From: Kanishk Dendukuri Date: Mon, 3 Aug 2026 21:52:06 -0700 Subject: [PATCH] GH-25505: [C++][Python] Add replace_with_indices compute kernel --- cpp/src/arrow/compute/api_vector.cc | 5 + cpp/src/arrow/compute/api_vector.h | 18 ++ .../arrow/compute/kernels/vector_replace.cc | 286 +++++++++++++++++- .../kernels/vector_replace_benchmark.cc | 46 +++ .../compute/kernels/vector_replace_test.cc | 218 +++++++++++++ docs/source/cpp/compute.rst | 10 +- docs/source/python/api/compute.rst | 1 + python/pyarrow/tests/test_compute.py | 62 ++++ 8 files changed, 642 insertions(+), 4 deletions(-) diff --git a/cpp/src/arrow/compute/api_vector.cc b/cpp/src/arrow/compute/api_vector.cc index 57b2eda69add..d6e92bb0102e 100644 --- a/cpp/src/arrow/compute/api_vector.cc +++ b/cpp/src/arrow/compute/api_vector.cc @@ -323,6 +323,11 @@ Result ReplaceWithMask(const Datum& values, const Datum& mask, return CallFunction("replace_with_mask", {values, mask, replacements}, ctx); } +Result ReplaceWithIndices(const Datum& values, const Datum& indices, + const Datum& replacements, ExecContext* ctx) { + return CallFunction("replace_with_indices", {values, indices, replacements}, ctx); +} + Result FillNullForward(const Datum& values, ExecContext* ctx) { return CallFunction("fill_null_forward", {values}, ctx); } diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index 2de7137cacdb..0e8d4a63c78f 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -477,6 +477,24 @@ ARROW_EXPORT Result ReplaceWithMask(const Datum& values, const Datum& mask, const Datum& replacements, ExecContext* ctx = NULLPTR); +/// \brief ReplaceWithIndices replaces the value at each position in `indices` +/// with the corresponding element from `replacements`. +/// +/// \param[in] values Array input to replace +/// \param[in] indices Integer array of positions to replace. Out-of-bounds and +/// null indices are ignored; if an index occurs more than once the last +/// replacement wins. +/// \param[in] replacements The replacement values (scalar or array). For an +/// array there must be as many replacement values as indices. +/// \param[in] ctx the function execution context, optional +/// +/// \return the resulting datum +/// +/// \note API not yet finalized +ARROW_EXPORT +Result ReplaceWithIndices(const Datum& values, const Datum& indices, + const Datum& replacements, ExecContext* ctx = NULLPTR); + /// \brief FillNullForward fill null values in forward direction /// /// The output array will be of the same type as the input values diff --git a/cpp/src/arrow/compute/kernels/vector_replace.cc b/cpp/src/arrow/compute/kernels/vector_replace.cc index ae39977aaea5..ab0a85a38286 100644 --- a/cpp/src/arrow/compute/kernels/vector_replace.cc +++ b/cpp/src/arrow/compute/kernels/vector_replace.cc @@ -22,6 +22,7 @@ #include "arrow/compute/kernels/util_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/util/bitmap_ops.h" +#include "arrow/util/int_util.h" #include "arrow/util/logging_internal.h" namespace arrow { @@ -450,6 +451,274 @@ struct ReplaceMaskChunked { } }; +Status CheckReplaceWithIndicesInputs(const DataType& value_type, + const DataType& indices_type, int64_t indices_length, + const DataType& replacements_type, + int64_t replacements_length, + bool replacements_is_array) { + // Needed for FixedSizeBinary/parameterized types + if (!value_type.Equals(replacements_type, /*check_metadata=*/false)) { + return Status::Invalid("Replacements must be of same type (expected ", + value_type.ToString(), " but got ", + replacements_type.ToString(), ")"); + } + if (!is_integer(indices_type.id())) { + return Status::Invalid("Indices must be of integer type (got ", + indices_type.ToString(), ")"); + } + if (replacements_is_array && replacements_length != indices_length) { + return Status::Invalid("Replacement array must be of appropriate length (expected ", + indices_length, " items but got ", replacements_length, + " items)"); + } + return Status::OK(); +} + +// Scatter the i-th replacement into output[indices[i] - values_base] for every index +template +void ReplaceWithIndicesScatter(const ArraySpan& indices, const Data& replacements, + bool replacements_bitmap, int64_t values_base, + int64_t values_length, const CopyBitmap& copy_bitmap, + uint8_t* out_bitmap, uint8_t* out_values, + int64_t out_offset, const DataType& value_type) { + int64_t i = 0; + VisitArraySpanInline( + indices, + [&](typename IndexType::c_type index) { + const int64_t idx = static_cast(index) - values_base; + if (idx >= 0 && idx < values_length) { + CopyDataUtils::CopyData(value_type, replacements, i, out_values, + out_offset + idx, /*length=*/1); + if (replacements_bitmap) { + copy_bitmap.SetBit(out_bitmap, out_offset + idx, i); + } else if (out_bitmap) { + bit_util::SetBitTo(out_bitmap, out_offset + idx, true); + } + } + ++i; + }, + [&]() { ++i; }); +} + +template +void DispatchIndexType(const ArraySpan& indices, Fn&& func) { + switch (indices.type->byte_width()) { + case 1: + return func(static_cast(nullptr)); + case 2: + return func(static_cast(nullptr)); + case 4: + return func(static_cast(nullptr)); + default: + DCHECK_EQ(indices.type->byte_width(), 8); + return func(static_cast(nullptr)); + } +} + +template +void ReplaceWithIndicesScatterDispatch(const ArraySpan& indices, const Data& replacements, + bool replacements_bitmap, int64_t values_base, + int64_t values_length, + const CopyBitmap& copy_bitmap, uint8_t* out_bitmap, + uint8_t* out_values, int64_t out_offset, + const DataType& value_type) { + DispatchIndexType(indices, [&](auto* index_type) { + using IndexType = std::remove_pointer_t; + ReplaceWithIndicesScatter( + indices, replacements, replacements_bitmap, values_base, values_length, + copy_bitmap, out_bitmap, out_values, out_offset, value_type); + }); +} + +template +struct ReplaceWithIndicesImpl {}; + +template +struct ReplaceWithIndicesImpl< + Type, enable_if_t::value || is_null_type::value)>> { + static Status ExecArrayIndices(KernelContext* ctx, const ArraySpan& array, + const ArraySpan& indices, int64_t values_base, + ExecValue replacements, ExecResult* out) { + ArrayData* out_arr = out->array_data().get(); + out_arr->length = array.length; + const int64_t out_offset = out_arr->offset; + uint8_t* out_values = out_arr->buffers[1]->mutable_data(); + // Start from a copy of the values, then overwrite the selected positions + CopyDataUtils::CopyData(*array.type, array, /*in_offset=*/0, out_values, + out_offset, array.length); + const bool replacements_bitmap = + replacements.is_array() ? replacements.array.MayHaveNulls() : true; + uint8_t* out_bitmap = out_arr->buffers[0]->mutable_data(); + out_arr->null_count = kUnknownNullCount; + if (array.MayHaveNulls()) { + arrow::internal::CopyBitmap(array.buffers[0].data, array.offset, array.length, + out_bitmap, out_offset); + } else { + bit_util::SetBitsTo(out_bitmap, out_offset, array.length, true); + } + + if (replacements.is_array()) { + const ArraySpan& source = replacements.array; + ReplaceWithIndicesScatterDispatch( + indices, source, replacements_bitmap, values_base, array.length, + CopyArrayBitmap{replacements_bitmap ? source.buffers[0].data : nullptr, + source.offset}, + out_bitmap, out_values, out_offset, *array.type); + } else { + const Scalar& source = *replacements.scalar; + ReplaceWithIndicesScatterDispatch( + indices, source, replacements_bitmap, values_base, array.length, + CopyScalarBitmap{source.is_valid}, out_bitmap, out_values, out_offset, + *array.type); + } + return Status::OK(); + } +}; + +template +struct ReplaceWithIndicesImpl> { + static Status ExecArrayIndices(KernelContext* ctx, const ArraySpan& array, + const ArraySpan& indices, int64_t values_base, + ExecValue replacements, ExecResult* out) { + out->value = array.ToArrayData(); + return Status::OK(); + } +}; + +template +struct ReplaceWithIndicesImpl> { + using BuilderType = typename TypeTraits::BuilderType; + + static Status ExecArrayIndices(KernelContext* ctx, const ArraySpan& array, + const ArraySpan& indices, int64_t values_base, + ExecValue replacements, ExecResult* out) { + // Var-length values can't be overwritten in place, so resolve for each output + // position whether a replacement targets it and then built output + std::vector replacement_for(array.length, -1); + DispatchIndexType(indices, [&](auto* index_type) { + using IndexType = std::remove_pointer_t; + int64_t i = 0; + VisitArraySpanInline( + indices, + [&](typename IndexType::c_type index) { + const int64_t idx = static_cast(index) - values_base; + if (idx >= 0 && idx < array.length) { + replacement_for[idx] = i; + } + ++i; + }, + [&]() { ++i; }); + }); + + BuilderType builder(array.type->GetSharedPtr(), ctx->memory_pool()); + RETURN_NOT_OK(builder.Reserve(array.length)); + RETURN_NOT_OK(builder.ReserveData(array.buffers[2].size)); + for (int64_t j = 0; j < array.length; ++j) { + const bool from_replacement = replacement_for[j] >= 0; + if (from_replacement && replacements.is_scalar()) { + const Scalar& scalar = *replacements.scalar; + if (scalar.is_valid) { + RETURN_NOT_OK(builder.Append(UnboxScalar::Unbox(scalar))); + } else { + RETURN_NOT_OK(builder.AppendNull()); + } + continue; + } + const ArraySpan& source = from_replacement ? replacements.array : array; + const int64_t offset = from_replacement ? replacement_for[j] : j; + if (!source.MayHaveNulls() || + bit_util::GetBit(source.buffers[0].data, source.offset + offset)) { + const uint8_t* data = source.buffers[2].data; + const auto* offsets = source.GetValues(1); + const auto offset0 = offsets[offset]; + const auto offset1 = offsets[offset + 1]; + RETURN_NOT_OK(builder.Append(data + offset0, offset1 - offset0)); + } else { + RETURN_NOT_OK(builder.AppendNull()); + } + } + std::shared_ptr temp_output; + RETURN_NOT_OK(builder.FinishInternal(&temp_output)); + // Builder type != logical type due to GenerateTypeAgnosticVarBinaryBase + temp_output->type = array.type->GetSharedPtr(); + out->value = std::move(temp_output); + return Status::OK(); + } +}; + +template +struct ReplaceWithIndices { + static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { + const ArraySpan& arr = batch[0].array; + const ExecValue& indices = batch[1]; + const ExecValue& replacements = batch[2]; + RETURN_NOT_OK(CheckReplaceWithIndicesInputs( + *arr.type, *indices.type(), indices.length(), *replacements.type(), + replacements.length(), replacements.is_array())); + RETURN_NOT_OK(arrow::internal::CheckIndexBounds(indices.array, arr.length)); + return ReplaceWithIndicesImpl::ExecArrayIndices( + ctx, arr, indices.array, /*values_base=*/0, replacements, out); + } + + static std::shared_ptr GetSignature(detail::GetTypeId get_id) { + return KernelSignature::Make( + {InputType(get_id.id), InputType(match::Integer()), InputType(get_id.id)}, + FirstType); + } +}; + +template +struct ReplaceWithIndicesChunked { + static Status Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) { + const Datum& indices = batch[1]; + const Datum& replacements = batch[2]; + + if (!indices.is_array()) { + return Status::Invalid("Indices must be array, not ", indices.ToString()); + } + if (!replacements.is_array() && !replacements.is_scalar()) { + return Status::Invalid("Replacements must be array or scalar, not ", + replacements.ToString()); + } + + const ChunkedArray& arr = *batch[0].chunked_array(); + RETURN_NOT_OK(CheckReplaceWithIndicesInputs( + *arr.type(), *indices.type(), indices.length(), *replacements.type(), + replacements.length(), replacements.is_arraylike())); + + ExecValue replacements_val = GetExecValue(replacements); + const ArraySpan& indices_span = *indices.array(); + RETURN_NOT_OK(arrow::internal::CheckIndexBounds(indices_span, arr.length())); + + // Each index refers to a position in the whole logical array, so track a base + // offset per chunk and apply the indices that fall within it + ArrayVector output_chunks; + output_chunks.reserve(arr.num_chunks()); + int64_t values_base = 0; + for (const std::shared_ptr& chunk : arr.chunks()) { + if (chunk->length() == 0) continue; + ExecResult chunk_result; + if (is_fixed_width(out->type()->id())) { + auto chunk_out = std::make_shared(chunk->type(), chunk->length()); + chunk_out->buffers.resize(2); + ARROW_ASSIGN_OR_RAISE(chunk_out->buffers[0], + ctx->AllocateBitmap(chunk->length())); + const int64_t slot_width = out->type()->byte_width(); + ARROW_ASSIGN_OR_RAISE(chunk_out->buffers[1], + ctx->Allocate(slot_width * chunk->length())); + chunk_result.value = chunk_out; + } + RETURN_NOT_OK(ReplaceWithIndicesImpl::ExecArrayIndices( + ctx, *chunk->data(), indices_span, values_base, replacements_val, + &chunk_result)); + output_chunks.push_back(MakeArray(chunk_result.array_data())); + values_base += chunk->length(); + } + + return ChunkedArray::Make(std::move(output_chunks), out->type()).Value(out); + } +}; + // This is for fixed-size types only template void FillNullInDirectionImpl(const ArraySpan& current_chunk, const uint8_t* null_bitmap, @@ -865,8 +1134,6 @@ void RegisterVectorFunction(FunctionRegistry* registry, } // TODO: list types DCHECK_OK(registry->AddFunction(std::move(func))); - - // TODO(ARROW-9431): "replace_with_indices" } } // namespace @@ -881,6 +1148,16 @@ const FunctionDoc replace_with_mask_doc( "Hence, for replacement arrays, len(replacements) == sum(mask == true)."), {"values", "mask", "replacements"}); +const FunctionDoc replace_with_indices_doc( + "Replace items at the given indices", + ("Given an array, an integer array of indices, and replacement values\n" + "(either scalar or array), the element of the array at position\n" + "indices[i] is replaced by the i-th replacement.\n" + "Null indices are ignored; if an index occurs more than once, the last\n" + "replacement wins. Out-of-bounds or negative indices raise an error.\n" + "Hence, for replacement arrays, len(replacements) == len(indices)."), + {"values", "indices", "replacements"}); + const FunctionDoc fill_null_forward_doc( "Carry non-null values forward to fill null slots", ("Given an array, propagate last valid observation forward to next valid\n" @@ -899,6 +1176,11 @@ void RegisterVectorReplace(FunctionRegistry* registry) { replace_with_mask_doc); RegisterVectorFunction(registry, func); } + { + auto func = std::make_shared("replace_with_indices", Arity::Ternary(), + replace_with_indices_doc); + RegisterVectorFunction(registry, func); + } { auto func = std::make_shared("fill_null_forward", Arity::Unary(), fill_null_forward_doc); diff --git a/cpp/src/arrow/compute/kernels/vector_replace_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_replace_benchmark.cc index 971a841de077..ff5810b6329b 100644 --- a/cpp/src/arrow/compute/kernels/vector_replace_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_replace_benchmark.cc @@ -81,10 +81,56 @@ static void ReplaceWithMaskHighSelectivityBench( state.SetBytesProcessed(state.iterations() * (len - offset) * 8); } +static void ReplaceWithIndicesLowSelectivityBench( + benchmark::State& state) { // NOLINT non-const reference + random::RandomArrayGenerator generator(kRandomSeed); + const int64_t len = state.range(0); + const int64_t offset = state.range(1); + + auto values = + generator.Int64(len, /*min=*/-65536, /*max=*/65536, /*null_probability=*/0.1) + ->Slice(offset); + const int64_t count = static_cast((len - offset) * 0.1); + auto indices = generator.Int32(count, /*min=*/0, /*max=*/(len - offset) - 1, + /*null_probability=*/0.1); + auto replacements = + generator.Int64(count, /*min=*/-65536, /*max=*/65536, /*null_probability=*/0.1); + + for (auto _ : state) { + ABORT_NOT_OK(ReplaceWithIndices(values, indices, replacements)); + } + state.SetBytesProcessed(state.iterations() * (len - offset) * 8); +} + +static void ReplaceWithIndicesHighSelectivityBench( + benchmark::State& state) { // NOLINT non-const reference + random::RandomArrayGenerator generator(kRandomSeed); + const int64_t len = state.range(0); + const int64_t offset = state.range(1); + + auto values = + generator.Int64(len, /*min=*/-65536, /*max=*/65536, /*null_probability=*/0.1) + ->Slice(offset); + const int64_t count = static_cast((len - offset) * 0.9); + auto indices = generator.Int32(count, /*min=*/0, /*max=*/(len - offset) - 1, + /*null_probability=*/0.1); + auto replacements = + generator.Int64(count, /*min=*/-65536, /*max=*/65536, /*null_probability=*/0.1); + + for (auto _ : state) { + ABORT_NOT_OK(ReplaceWithIndices(values, indices, replacements)); + } + state.SetBytesProcessed(state.iterations() * (len - offset) * 8); +} + BENCHMARK(ReplaceWithMaskLowSelectivityBench)->Args({kLongLength, 0}); BENCHMARK(ReplaceWithMaskLowSelectivityBench)->Args({kLongLength, 99}); BENCHMARK(ReplaceWithMaskHighSelectivityBench)->Args({kLongLength, 0}); BENCHMARK(ReplaceWithMaskHighSelectivityBench)->Args({kLongLength, 99}); +BENCHMARK(ReplaceWithIndicesLowSelectivityBench)->Args({kLongLength, 0}); +BENCHMARK(ReplaceWithIndicesLowSelectivityBench)->Args({kLongLength, 99}); +BENCHMARK(ReplaceWithIndicesHighSelectivityBench)->Args({kLongLength, 0}); +BENCHMARK(ReplaceWithIndicesHighSelectivityBench)->Args({kLongLength, 99}); } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/kernels/vector_replace_test.cc b/cpp/src/arrow/compute/kernels/vector_replace_test.cc index 9dc8e70ab6a2..156dcc8c1348 100644 --- a/cpp/src/arrow/compute/kernels/vector_replace_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_replace_test.cc @@ -66,6 +66,10 @@ class TestReplaceKernel : public ::testing::Test { return ArrayFromJSON(boolean(), value); } + std::shared_ptr indices(const std::string& value) { + return ArrayFromJSON(int32(), value); + } + Status AssertRaises(ReplaceFunction func, const Datum& array, const Datum& mask, const Datum& replacements) { auto result = func(array, mask, replacements, nullptr); @@ -184,6 +188,40 @@ class TestReplaceKernel : public ::testing::Test { EXPECT_OK_AND_ASSIGN(auto expected, builder->Finish()); return expected; } + + std::shared_ptr NaiveImplIndices( + const typename TypeTraits::ArrayType& array, const Int32Array& indices, + const typename TypeTraits::ArrayType& replacements) { + auto length = array.length(); + std::vector::CType> values(length); + std::vector is_valid(length); + for (int64_t i = 0; i < length; ++i) { + values[i] = array.IsValid(i) ? array.Value(i) : typename TypeTraits::CType{}; + is_valid[i] = array.IsValid(i); + } + for (int64_t i = 0; i < indices.length(); ++i) { + if (!indices.IsValid(i)) continue; + int64_t idx = indices.Value(i); + if (idx < 0 || idx >= length) continue; + if (replacements.IsValid(i)) { + values[idx] = replacements.Value(i); + is_valid[idx] = true; + } else { + is_valid[idx] = false; + } + } + auto builder = std::make_unique::BuilderType>( + default_type_instance(), default_memory_pool()); + for (int64_t i = 0; i < length; ++i) { + if (is_valid[i]) { + ARROW_EXPECT_OK(builder->Append(values[i])); + } else { + ARROW_EXPECT_OK(builder->AppendNull()); + } + } + EXPECT_OK_AND_ASSIGN(auto expected, builder->Finish()); + return expected; + } }; template @@ -442,6 +480,164 @@ TYPED_TEST(TestReplaceNumeric, ReplaceWithMaskRandom) { } } +TYPED_TEST(TestReplaceNumeric, ReplaceWithIndices) { + std::vector cases = { + {this->array("[]"), this->indices("[]"), this->array("[]"), this->array("[]")}, + {this->array("[0, 0, 0]"), this->indices("[]"), this->array("[]"), + this->array("[0, 0, 0]")}, + {this->array("[0, 0, 0]"), this->indices("[0, 2]"), this->array("[0, 0]"), + this->array("[0, 0, 0]")}, + {this->array("[0, 0, 0]"), this->indices("[1]"), this->scalar("0"), + this->array("[0, 0, 0]")}, + {this->array("[0, 0, 0]"), this->indices("[1]"), this->array("[null]"), + this->array("[0, null, 0]")}, + {this->array("[0, 0, 0]"), this->indices("[1]"), this->scalar("null"), + this->array("[0, null, 0]")}, + {this->array("[0, null, 0]"), this->indices("[0]"), this->array("[0]"), + this->array("[0, null, 0]")}, + }; + for (auto& c : cases) { + this->Assert(ReplaceWithIndices, c.input, c.mask, c.replacements, c.expected); + } +} + +TYPED_TEST(TestReplaceNumeric, ReplaceWithIndicesRandom) { + using ArrayType = typename TypeTraits::ArrayType; + using CType = typename TypeTraits::CType; + auto ty = this->type(); + + random::RandomArrayGenerator rand(/*seed=*/0); + const int64_t length = 1023; + std::vector values = {"0.01", "0"}; + // Clamp the range because date/time types don't print well with extreme values + values.push_back(std::to_string(static_cast(std::min( + 16384.0, static_cast(std::numeric_limits::max()))))); + auto options = key_value_metadata({"null_probability", "min", "max"}, values); + auto array = + checked_pointer_cast(rand.ArrayOf(*field("a", ty, options), length)); + const int64_t num_indices = 512; + auto indices = checked_pointer_cast( + rand.Int32(num_indices, /*min=*/0, /*max=*/static_cast(length) - 1, + /*null_probability=*/0.1)); + auto replacements = checked_pointer_cast( + rand.ArrayOf(*field("a", ty, options), num_indices)); + auto expected = this->NaiveImplIndices(*array, *indices, *replacements); + + this->Assert(ReplaceWithIndices, array, indices, replacements, expected); + for (int64_t slice = 1; slice <= 16; slice++) { + auto sliced_array = checked_pointer_cast(array->Slice(slice, 15)); + auto sliced_indices = checked_pointer_cast( + rand.Int32(num_indices, /*min=*/0, /*max=*/14, /*null_probability=*/0.1)); + auto new_expected = + this->NaiveImplIndices(*sliced_array, *sliced_indices, *replacements); + this->Assert(ReplaceWithIndices, sliced_array, sliced_indices, replacements, + new_expected); + } +} + +class TestReplaceWithIndices : public ::testing::Test { + protected: + std::shared_ptr array(const std::string& value) { + return ArrayFromJSON(int32(), value); + } + Datum scalar(const std::string& value) { return ScalarFromJSON(int32(), value); } + std::shared_ptr indices(const std::string& value) { + return ArrayFromJSON(int32(), value); + } + std::shared_ptr chunked(const std::vector& value) { + return ChunkedArrayFromJSON(int32(), value); + } + void Check(const Datum& values, const Datum& idx, const Datum& replacements, + const Datum& expected) { + ASSERT_OK_AND_ASSIGN(auto actual, + ReplaceWithIndices(values, idx, replacements, nullptr)); + ASSERT_OK(actual.make_array()->ValidateFull()); + AssertDatumsEqual(expected, actual, /*verbose=*/true); + } + void CheckChunked(const Datum& values, const Datum& idx, const Datum& replacements, + const Datum& expected) { + ASSERT_OK_AND_ASSIGN(auto actual, + ReplaceWithIndices(values, idx, replacements, nullptr)); + ASSERT_OK(actual.chunked_array()->ValidateFull()); + AssertDatumsEqual(expected, actual, /*verbose=*/true); + } +}; + +TEST_F(TestReplaceWithIndices, Basic) { + Check(array("[1, 2, 3, 4]"), indices("[0, 3]"), array("[10, 40]"), + array("[10, 2, 3, 40]")); + Check(array("[1, 2, 3, 4]"), indices("[3, 0]"), array("[40, 10]"), + array("[10, 2, 3, 40]")); + Check(array("[1, 2, 3]"), indices("[0, 2]"), scalar("9"), array("[9, 2, 9]")); + Check(array("[1, 2, 3]"), indices("[1]"), array("[null]"), array("[1, null, 3]")); + // Duplicate indices: last replacement wins + Check(array("[1, 2, 3, 4]"), indices("[0, 0]"), array("[10, 20]"), + array("[20, 2, 3, 4]")); + // Null indices are skipped + Check(array("[1, 2, 3, 4]"), indices("[null, 2]"), array("[10, 30]"), + array("[1, 2, 30, 4]")); +} + +TEST_F(TestReplaceWithIndices, IndexTypes) { + // Any integer width works for the indices + for (auto index_ty : + {int8(), uint8(), int16(), uint16(), int32(), uint32(), int64(), uint64()}) { + ASSERT_OK_AND_ASSIGN( + auto actual, + ReplaceWithIndices(array("[1, 2, 3, 4]"), ArrayFromJSON(index_ty, "[0, 3]"), + array("[10, 40]"), nullptr)); + ASSERT_OK(actual.make_array()->ValidateFull()); + AssertDatumsEqual(array("[10, 2, 3, 40]"), actual, /*verbose=*/true); + } +} + +TEST_F(TestReplaceWithIndices, Errors) { + // Replacement array length must match indices length + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr("Replacement array must be of appropriate length (expected 2 " + "items but got 1 items)"), + ReplaceWithIndices(array("[1, 2]"), indices("[0, 1]"), array("[0]"), nullptr) + .status()); + // Replacements must match the values type (checked for same-id parameterized types) + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("Replacements must be of same type"), + ReplaceWithIndices(ArrayFromJSON(fixed_size_binary(3), R"(["abc"])"), + indices("[0]"), ArrayFromJSON(fixed_size_binary(2), R"(["ab"])"), + nullptr) + .status()); + // Out-of-bounds and negative indices are errors + EXPECT_RAISES_WITH_MESSAGE_THAT( + IndexError, ::testing::HasSubstr("Index 9 out of bounds"), + ReplaceWithIndices(array("[1, 2, 3, 4]"), indices("[0, 9]"), array("[10, 40]"), + nullptr) + .status()); + EXPECT_RAISES_WITH_MESSAGE_THAT( + IndexError, ::testing::HasSubstr("Index -1 out of bounds"), + ReplaceWithIndices(array("[1, 2, 3, 4]"), indices("[-1, 2]"), array("[10, 30]"), + nullptr) + .status()); +} + +TEST_F(TestReplaceWithIndices, Chunked) { + CheckChunked(chunked({"[1, 2, 3]", "[4, 5, 6]"}), indices("[0, 4]"), array("[10, 50]"), + chunked({"[10, 2, 3]", "[4, 50, 6]"})); + // An empty chunk in the middle does not shift addressing + CheckChunked(chunked({"[1, 2]", "[]", "[3, 4]"}), indices("[1, 3]"), array("[20, 40]"), + chunked({"[1, 20]", "[3, 40]"})); + CheckChunked(chunked({"[1, 2]", "[3, 4]"}), indices("[0, 3]"), scalar("9"), + chunked({"[9, 2]", "[3, 9]"})); +} + +TEST_F(TestReplaceWithIndices, ChunkedErrors) { + // Bounds are checked against the total logical length + EXPECT_RAISES_WITH_MESSAGE_THAT( + IndexError, ::testing::HasSubstr("Index 6 out of bounds"), + ReplaceWithIndices(chunked({"[1, 2, 3]", "[4, 5, 6]"}), indices("[6]"), + array("[99]"), nullptr) + .status()); +} + TYPED_TEST(TestReplaceNumeric, ReplaceWithMaskErrors) { EXPECT_RAISES_WITH_MESSAGE_THAT( Invalid, @@ -1070,6 +1266,28 @@ TYPED_TEST(TestReplaceBinary, ReplaceWithMask) { } } +TYPED_TEST(TestReplaceBinary, ReplaceWithIndices) { + std::vector cases = { + {this->array("[]"), this->indices("[]"), this->array("[]"), this->array("[]")}, + // Replacement of a different length than the value it overwrites + {this->array(R"(["a", "bb", "ccc"])"), this->indices("[1]"), + this->array(R"(["dddd"])"), this->array(R"(["a", "dddd", "ccc"])")}, + {this->array(R"(["a", "bb", "ccc"])"), this->indices("[0, 2]"), + this->scalar(R"("z")"), this->array(R"(["z", "bb", "z"])")}, + {this->array(R"(["a", "bb"])"), this->indices("[0]"), this->array("[null]"), + this->array(R"([null, "bb"])")}, + // Duplicate indices: last wins + {this->array(R"(["a", "bb", "ccc"])"), this->indices("[0, 0]"), + this->array(R"(["x", "yy"])"), this->array(R"(["yy", "bb", "ccc"])")}, + // Null indices are skipped + {this->array(R"(["a", "bb", "ccc"])"), this->indices("[null, 2]"), + this->array(R"(["x", "yy"])"), this->array(R"(["a", "bb", "yy"])")}, + }; + for (auto& c : cases) { + this->Assert(ReplaceWithIndices, c.input, c.mask, c.replacements, c.expected); + } +} + TYPED_TEST(TestReplaceBinary, ReplaceWithMaskRandom) { using ArrayType = typename TypeTraits::ArrayType; auto ty = this->type(); diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index 1e067c52188d..cb60b3ce1613 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -1988,11 +1988,17 @@ replaced, based on the remaining inputs. +--------------------------+------------+-----------------------+--------------+--------------+--------------+-------+ | fill_null_forward | Unary | Fixed-width or binary | | | Input type 1 | \(1) | +--------------------------+------------+-----------------------+--------------+--------------+--------------+-------+ -| replace_with_mask | Ternary | Fixed-width or binary | Boolean | Input type 1 | Input type 1 | \(2) | +| replace_with_indices | Ternary | Fixed-width or binary | Integer | Input type 1 | Input type 1 | \(2) | ++--------------------------+------------+-----------------------+--------------+--------------+--------------+-------+ +| replace_with_mask | Ternary | Fixed-width or binary | Boolean | Input type 1 | Input type 1 | \(3) | +--------------------------+------------+-----------------------+--------------+--------------+--------------+-------+ * \(1) Valid values are carried forward/backward to fill null values. -* \(2) Each element in input 1 for which the corresponding Boolean in input 2 +* \(2) The element of input 1 at the position given by ``indices[i]`` (input 2) + is replaced with the i-th value from input 3. Null indices are ignored; if an + index occurs more than once, the last replacement wins. Out-of-bounds or + negative indices raise an error. +* \(3) Each element in input 1 for which the corresponding Boolean in input 2 is true is replaced with the next value from input 3. A null in input 2 results in a corresponding null in the output. diff --git a/docs/source/python/api/compute.rst b/docs/source/python/api/compute.rst index 6a4b04468d4c..2843f9c5cf9c 100644 --- a/docs/source/python/api/compute.rst +++ b/docs/source/python/api/compute.rst @@ -576,6 +576,7 @@ Structural Transforms list_value_length make_struct map_lookup + replace_with_indices replace_with_mask struct_field diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 1e08e73668e5..742455d959ed 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -1415,6 +1415,68 @@ def test_replace_with_mask_error_mask_length_mismatch(): pc.replace_with_mask(arr, mask, replacements) +def test_replace_with_indices_basic(): + """The i-th replacement is written at position indices[i].""" + arr = pa.array([1, 2, 3, 4]) + indices = pa.array([0, 3]) + replacements = pa.array([10, 40]) + expected = pa.array([10, 2, 3, 40]) + result = pc.replace_with_indices(arr, indices, replacements) + assert result.equals(expected) + + +def test_replace_with_indices_unordered(): + """Indices need not be sorted; replacement pairs with its index.""" + arr = pa.array([1, 2, 3, 4]) + result = pc.replace_with_indices(arr, pa.array([3, 0]), pa.array([40, 10])) + assert result.equals(pa.array([10, 2, 3, 40])) + + +def test_replace_with_indices_scalar_replacement(): + arr = pa.array([1, 2, 3]) + result = pc.replace_with_indices(arr, pa.array([0, 2]), pa.scalar(9)) + assert result.equals(pa.array([9, 2, 9])) + + +def test_replace_with_indices_string_type(): + arr = pa.array(['a', 'b', 'c', 'd']) + result = pc.replace_with_indices(arr, pa.array([0, 2]), pa.array(['x', 'y'])) + assert result.equals(pa.array(['x', 'b', 'y', 'd'])) + + +def test_replace_with_indices_null_replacement(): + arr = pa.array([1, 2, 3]) + result = pc.replace_with_indices(arr, pa.array([1]), pa.array([None])) + assert result.equals(pa.array([1, None, 3])) + + +def test_replace_with_indices_duplicate_last_wins(): + arr = pa.array([1, 2, 3, 4]) + result = pc.replace_with_indices(arr, pa.array([0, 0]), pa.array([10, 20])) + assert result.equals(pa.array([20, 2, 3, 4])) + + +def test_replace_with_indices_out_of_bounds_error(): + arr = pa.array([1, 2, 3, 4]) + with pytest.raises(pa.ArrowIndexError, match="out of bounds"): + pc.replace_with_indices(arr, pa.array([0, 9]), pa.array([10, 40])) + with pytest.raises(pa.ArrowIndexError, match="out of bounds"): + pc.replace_with_indices(arr, pa.array([-1, 2]), pa.array([10, 30])) + + +def test_replace_with_indices_null_index_skipped(): + arr = pa.array([1, 2, 3, 4]) + result = pc.replace_with_indices(arr, pa.array([None, 2]), pa.array([10, 30])) + assert result.equals(pa.array([1, 2, 30, 4])) + + +def test_replace_with_indices_error_length_mismatch(): + """Replacement array length must match the number of indices.""" + arr = pa.array([1, 2]) + with pytest.raises(pa.ArrowInvalid, match="expected 2.*but got 1"): + pc.replace_with_indices(arr, pa.array([0, 1]), pa.array([0])) + + def test_binary_join(): ar_list = pa.array([['foo', 'bar'], None, []]) expected = pa.array(['foo-bar', None, ''])