From 542eb21c103ff090b37c2dc4c117811d9b6c4286 Mon Sep 17 00:00:00 2001 From: Pratyush Adhikari Date: Sun, 16 Aug 2026 02:05:10 +0530 Subject: [PATCH 1/2] GH-50879: [C++] Implement replace_with_mask for List and LargeList types This commit adds support for variable-width list types (ListType and LargeListType) to the replace_with_mask compute kernel. It introduces a specialization of ReplaceMaskImpl that handles variable-length children safely by directly iterating over values and appending array slices, avoiding invalid length mutations. --- .../arrow/compute/kernels/vector_replace.cc | 127 +++++++++- .../compute/kernels/vector_replace_test.cc | 224 ++++++++++++++++++ 2 files changed, 348 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_replace.cc b/cpp/src/arrow/compute/kernels/vector_replace.cc index 6a9abfc03960..37056c3637d5 100644 --- a/cpp/src/arrow/compute/kernels/vector_replace.cc +++ b/cpp/src/arrow/compute/kernels/vector_replace.cc @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +#include "arrow/array/builder_nested.h" #include "arrow/compute/api_scalar.h" #include "arrow/compute/kernels/codegen_internal.h" #include "arrow/compute/kernels/common_internal.h" @@ -119,7 +120,8 @@ struct ReplaceMaskImpl {}; template struct ReplaceMaskImpl< - Type, enable_if_t::value || is_null_type::value)>> { + Type, enable_if_t::value || is_null_type::value || + is_var_length_list_type::value)>> { static Result ExecScalarMask(KernelContext* ctx, const ArraySpan& array, const BooleanScalar& mask, ExecValue replacements, int64_t replacements_offset, ExecResult* out) { @@ -322,6 +324,111 @@ struct ReplaceMaskImpl> { } }; +// Specialization for variable-size list types (list and large_list). +// Each list element is copied individually using AppendArraySlice, mirroring +// the enable_if_base_binary specialization's per-element builder approach. +template +struct ReplaceMaskImpl> { + using offset_type = typename Type::offset_type; + using BuilderType = typename TypeTraits::BuilderType; + + static Result ExecScalarMask(KernelContext* ctx, const ArraySpan& array, + const BooleanScalar& mask, + ExecValue replacements, + int64_t replacements_offset, + ExecResult* out) { + if (!mask.is_valid) { + // mask = null: output is all-null array + ARROW_ASSIGN_OR_RAISE( + auto replacement_array, + MakeArrayOfNull(array.type->GetSharedPtr(), array.length, ctx->memory_pool())); + out->value = std::move(replacement_array->data()); + return replacements_offset; + } else if (mask.value) { + // mask = true: output = replacement + if (replacements.is_scalar()) { + ARROW_ASSIGN_OR_RAISE( + auto replacement_array, + MakeArrayFromScalar(*replacements.scalar, array.length, ctx->memory_pool())); + out->value = std::move(replacement_array->data()); + } else { + // Zero-copy slice into the replacements array — same approach as base binary. + std::shared_ptr result = replacements.array.ToArrayData(); + result->offset += replacements_offset; + result->length = array.length; + // Null count from original replacements applies to the whole array, not this + // slice; mark as unknown so it is recomputed on demand. + result->null_count = kUnknownNullCount; + out->value = result; + } + return replacements_offset + array.length; + } else { + // mask = false: output = input (zero-copy) + out->value = array.ToArrayData(); + return replacements_offset; + } + } + + static Result ExecArrayMask(KernelContext* ctx, const ArraySpan& array, + const ArraySpan& mask, int64_t mask_offset, + ExecValue replacements, + int64_t replacements_offset, + ExecResult* out) { + // Build the output list array element-by-element. We cannot pre-allocate a flat + // buffer (unlike fixed-width types) because the child-array size of each list slot + // is unknown in advance, so we use a list builder and copy slots individually. + std::unique_ptr raw_builder; + RETURN_NOT_OK(MakeBuilderExactIndex(ctx->memory_pool(), + array.type->GetSharedPtr(), &raw_builder)); + auto& builder = checked_cast(*raw_builder); + RETURN_NOT_OK(builder.Reserve(array.length)); + + // Source offset tracks our position in `array` (the values argument). + int64_t source_offset = 0; + + // Narrow the mask span to [mask_offset, mask_offset + array.length). + ArraySpan adjusted_mask = mask; + adjusted_mask.offset += mask_offset; + adjusted_mask.length = std::min(adjusted_mask.length - mask_offset, array.length); + + RETURN_NOT_OK(VisitArraySpanInline( + adjusted_mask, + [&](bool replace) -> Status { + if (replace && replacements.is_scalar()) { + // Scalar replacement: append the scalar value once. + RETURN_NOT_OK(builder.AppendScalar(*replacements.scalar)); + } else { + const ArraySpan& source = replace ? replacements.array : array; + const int64_t offset = replace ? replacements_offset++ : source_offset; + // Check validity of the source element at `offset`. + const bool is_valid = + !source.MayHaveNulls() || + bit_util::GetBit(source.buffers[0].data, source.offset + offset); + if (is_valid) { + // AppendArraySlice copies one list element (offset, 1) including its + // child values and validity, correctly handling source.offset. + RETURN_NOT_OK(builder.AppendArraySlice(source, offset, 1)); + } else { + RETURN_NOT_OK(builder.AppendNull()); + } + } + source_offset++; + return Status::OK(); + }, + [&]() -> Status { + // Null mask entry → null output element. + RETURN_NOT_OK(builder.AppendNull()); + source_offset++; + return Status::OK(); + })); + + std::shared_ptr temp_output; + RETURN_NOT_OK(builder.Finish(&temp_output)); + out->value = std::move(temp_output->data()); + return replacements_offset; + } +}; + Status CheckReplaceMaskInputs(const DataType& value_type, int64_t arr_length, const ExecValue& mask_box, const DataType& replacements_type, @@ -862,8 +969,10 @@ void RegisterVectorFunction(FunctionRegistry* registry, GenerateTypeAgnosticVarBinaryBase(*ty), registry, func.get()); } - // TODO: list types - DCHECK_OK(registry->AddFunction(std::move(func))); + // Note: list types (LIST, LARGE_LIST) are NOT added here. RegisterVectorFunction is + // also used for fill_null_forward/fill_null_backward which do not yet support list + // types. The caller is responsible for adding list kernels when appropriate and for + // calling registry->AddFunction. // TODO(ARROW-9431): "replace_with_indices" } @@ -897,16 +1006,28 @@ void RegisterVectorReplace(FunctionRegistry* registry) { auto func = std::make_shared("replace_with_mask", Arity::Ternary(), replace_with_mask_doc); RegisterVectorFunction(registry, func); + // Add LIST and LARGE_LIST support. These are registered separately from + // RegisterVectorFunction because fill_null_forward/backward do not support list + // types yet; adding them here avoids requiring implementations for those kernels. + AddKernel(Type::LIST, ReplaceMask::GetSignature(Type::LIST), + ReplaceMask::Exec, ReplaceMaskChunked::Exec, registry, + func.get()); + AddKernel(Type::LARGE_LIST, ReplaceMask::GetSignature(Type::LARGE_LIST), + ReplaceMask::Exec, ReplaceMaskChunked::Exec, + registry, func.get()); + DCHECK_OK(registry->AddFunction(std::move(func))); } { auto func = std::make_shared("fill_null_forward", Arity::Unary(), fill_null_forward_doc); RegisterVectorFunction(registry, func); + DCHECK_OK(registry->AddFunction(std::move(func))); } { auto func = std::make_shared("fill_null_backward", Arity::Unary(), fill_null_backward_doc); RegisterVectorFunction(registry, func); + DCHECK_OK(registry->AddFunction(std::move(func))); } } } // namespace internal diff --git a/cpp/src/arrow/compute/kernels/vector_replace_test.cc b/cpp/src/arrow/compute/kernels/vector_replace_test.cc index 9dc8e70ab6a2..1051a43e7489 100644 --- a/cpp/src/arrow/compute/kernels/vector_replace_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_replace_test.cc @@ -2106,5 +2106,229 @@ TEST_F(TestFillNullType, TestFillOnNullType) { this->AssertFillNullArray(FillNullBackward, this->array(R"([null, null])"), this->array(R"([null, null])")); } + +// ---------------------------------------------------------------------------- +// Tests for replace_with_mask with list and large_list +// ---------------------------------------------------------------------------- + +// Helper: assert ReplaceWithMask output for list types, calling ValidateFull. +static void AssertReplaceWithMaskList(const Datum& values, const Datum& mask, + const Datum& replacements, + const Datum& expected) { + ASSERT_OK_AND_ASSIGN(auto actual, ReplaceWithMask(values, mask, replacements)); + if (actual.is_array()) { + ASSERT_OK(actual.make_array()->ValidateFull()); + } else if (actual.is_arraylike()) { + ASSERT_OK(actual.chunked_array()->ValidateFull()); + } + AssertDatumsEqual(expected, actual, /*verbose=*/true); +} + +// Scalar mask: false → copy input unchanged +TEST(TestReplaceWithMaskList, ScalarMaskFalse) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1, 2], [3, 4], [5]])"); + auto mask = std::make_shared(false); + auto replacements = ArrayFromJSON(ty, R"([])"); + auto expected = ArrayFromJSON(ty, R"([[1, 2], [3, 4], [5]])"); + AssertReplaceWithMaskList(values, Datum(mask), replacements, expected); +} + +// Scalar mask: true → all elements replaced by successive replacements +TEST(TestReplaceWithMaskList, ScalarMaskTrue) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1, 2], [3, 4]])"); + auto mask = std::make_shared(true); + auto replacements = ArrayFromJSON(ty, R"([[10, 20], [30]])"); + auto expected = ArrayFromJSON(ty, R"([[10, 20], [30]])"); + AssertReplaceWithMaskList(values, Datum(mask), replacements, expected); +} + +// Scalar mask: null → all outputs become null +TEST(TestReplaceWithMaskList, ScalarMaskNull) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1, 2], [3]])"); + auto mask = std::make_shared(); + mask->is_valid = false; + auto replacements = ArrayFromJSON(ty, R"([])"); + auto expected = ArrayFromJSON(ty, R"([null, null])"); + AssertReplaceWithMaskList(values, Datum(mask), replacements, expected); +} + +// Array mask: mixed true/false, no nulls in values or replacements +TEST(TestReplaceWithMaskList, ArrayMaskMixed) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1, 2], [3, 4], [5, 6]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, false, true])"); + auto replacements = ArrayFromJSON(ty, R"([[10, 20], [30, 40]])"); + auto expected = ArrayFromJSON(ty, R"([[10, 20], [3, 4], [30, 40]])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Array mask: all false → input unchanged +TEST(TestReplaceWithMaskList, ArrayMaskAllFalse) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1], [2], [3]])"); + auto mask = ArrayFromJSON(boolean(), R"([false, false, false])"); + auto replacements = ArrayFromJSON(ty, R"([])"); + auto expected = ArrayFromJSON(ty, R"([[1], [2], [3]])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Array mask: all true → all replaced +TEST(TestReplaceWithMaskList, ArrayMaskAllTrue) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1], [2], [3]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, true, true])"); + auto replacements = ArrayFromJSON(ty, R"([[10], [20], [30]])"); + auto expected = ArrayFromJSON(ty, R"([[10], [20], [30]])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Array mask: null mask entries → null output +TEST(TestReplaceWithMaskList, NullMaskEntries) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1, 2], [3, 4], [5]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, null, false])"); + auto replacements = ArrayFromJSON(ty, R"([[10, 20]])"); + auto expected = ArrayFromJSON(ty, R"([[10, 20], null, [5]])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Null elements in the values array: preserve nulls when mask is false +TEST(TestReplaceWithMaskList, NullsInValues) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([null, [1, 2], null])"); + auto mask = ArrayFromJSON(boolean(), R"([false, true, false])"); + auto replacements = ArrayFromJSON(ty, R"([[10]])"); + auto expected = ArrayFromJSON(ty, R"([null, [10], null])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Null elements in the replacements array +TEST(TestReplaceWithMaskList, NullsInReplacements) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1], [2], [3]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, true, false])"); + auto replacements = ArrayFromJSON(ty, R"([null, [20]])"); + auto expected = ArrayFromJSON(ty, R"([null, [20], [3]])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Scalar replacement (null scalar) +TEST(TestReplaceWithMaskList, ScalarReplacementNull) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1], [2], [3]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, false, true])"); + auto replacement_scalar = MakeNullScalar(ty); + auto expected = ArrayFromJSON(ty, R"([null, [2], null])"); + AssertReplaceWithMaskList(values, mask, Datum(replacement_scalar), expected); +} + +// Scalar replacement (valid scalar) +TEST(TestReplaceWithMaskList, ScalarReplacementValid) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1, 2], [3], [4, 5, 6]])"); + auto mask = ArrayFromJSON(boolean(), R"([false, true, true])"); + auto replacement_scalar = ScalarFromJSON(ty, R"([99, 100])"); + auto expected = ArrayFromJSON(ty, R"([[1, 2], [99, 100], [99, 100]])"); + AssertReplaceWithMaskList(values, mask, Datum(replacement_scalar), expected); +} + +// Empty arrays → empty output +TEST(TestReplaceWithMaskList, EmptyArrays) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([])"); + auto mask = ArrayFromJSON(boolean(), R"([])"); + auto replacements = ArrayFromJSON(ty, R"([])"); + auto expected = ArrayFromJSON(ty, R"([])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Variable-length children: each list slot has different child count +TEST(TestReplaceWithMaskList, VariableLengthChildren) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[], [1], [1, 2], [1, 2, 3]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, false, true, false])"); + auto replacements = ArrayFromJSON(ty, R"([[9, 8, 7, 6], []])"); + auto expected = ArrayFromJSON(ty, R"([[9, 8, 7, 6], [1], [], [1, 2, 3]])"); + AssertReplaceWithMaskList(values, mask, replacements, expected); +} + +// Chunked array input +TEST(TestReplaceWithMaskList, ChunkedArray) { + auto ty = list(int32()); + auto chunk0 = ArrayFromJSON(ty, R"([[1, 2], [3]])"); + auto chunk1 = ArrayFromJSON(ty, R"([[4, 5]])"); + auto values = std::make_shared(ArrayVector{chunk0, chunk1}); + auto mask = ArrayFromJSON(boolean(), R"([true, false, true])"); + auto replacements = ArrayFromJSON(ty, R"([[10, 20], [30, 40, 50]])"); + // Expected: chunked output with same chunk boundaries + auto exp0 = ArrayFromJSON(ty, R"([[10, 20], [3]])"); + auto exp1 = ArrayFromJSON(ty, R"([[30, 40, 50]])"); + auto expected = std::make_shared(ArrayVector{exp0, exp1}); + AssertReplaceWithMaskList(Datum(values), mask, replacements, Datum(expected)); +} + +// Scalar mask (true) with chunked array input +TEST(TestReplaceWithMaskList, ChunkedArrayScalarMaskTrue) { + auto ty = list(int32()); + auto chunk0 = ArrayFromJSON(ty, R"([[1], [2]])"); + auto chunk1 = ArrayFromJSON(ty, R"([[3]])"); + auto values = std::make_shared(ArrayVector{chunk0, chunk1}); + auto mask = std::make_shared(true); + auto replacements = ArrayFromJSON(ty, R"([[10], [20], [30]])"); + auto exp0 = ArrayFromJSON(ty, R"([[10], [20]])"); + auto exp1 = ArrayFromJSON(ty, R"([[30]])"); + auto expected = std::make_shared(ArrayVector{exp0, exp1}); + AssertReplaceWithMaskList(Datum(values), Datum(mask), replacements, Datum(expected)); +} + +// large_list: smoke test to ensure LARGE_LIST is registered +TEST(TestReplaceWithMaskLargeList, ArrayMaskMixed) { + auto ty = large_list(int32()); + auto values = ArrayFromJSON(ty, R"([[1, 2], [3, 4], [5]])"); + auto mask = ArrayFromJSON(boolean(), R"([false, true, false])"); + auto replacements = ArrayFromJSON(ty, R"([[10, 20]])"); + auto expected = ArrayFromJSON(ty, R"([[1, 2], [10, 20], [5]])"); + ASSERT_OK_AND_ASSIGN(auto actual, ReplaceWithMask(values, mask, replacements)); + ASSERT_OK(actual.make_array()->ValidateFull()); + AssertArraysEqual(*ArrayFromJSON(ty, R"([[1, 2], [10, 20], [5]])"), + *actual.make_array(), /*verbose=*/true); +} + +// large_list: null mask entry → null output +TEST(TestReplaceWithMaskLargeList, NullMaskEntry) { + auto ty = large_list(int32()); + auto values = ArrayFromJSON(ty, R"([[1], [2]])"); + auto mask = ArrayFromJSON(boolean(), R"([null, true])"); + auto replacements = ArrayFromJSON(ty, R"([[99]])"); + auto expected = ArrayFromJSON(ty, R"([null, [99]])"); + ASSERT_OK_AND_ASSIGN(auto actual, ReplaceWithMask(values, mask, replacements)); + ASSERT_OK(actual.make_array()->ValidateFull()); + AssertArraysEqual(*expected, *actual.make_array(), /*verbose=*/true); +} + +// Replacement length mismatch → Status::Invalid +TEST(TestReplaceWithMaskList, ReplacementLengthMismatch) { + auto ty = list(int32()); + auto values = ArrayFromJSON(ty, R"([[1], [2], [3]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, true, false])"); + // mask has 2 true entries but replacements has only 1 element + auto replacements = ArrayFromJSON(ty, R"([[10]])"); + auto result = ReplaceWithMask(values, mask, replacements); + ASSERT_FALSE(result.ok()); +} + +// Type mismatch → Status::Invalid +TEST(TestReplaceWithMaskList, TypeMismatch) { + auto values = ArrayFromJSON(list(int32()), R"([[1], [2]])"); + auto mask = ArrayFromJSON(boolean(), R"([true, false])"); + auto replacements = ArrayFromJSON(list(int64()), R"([[10]])"); + auto result = ReplaceWithMask(values, mask, replacements); + ASSERT_FALSE(result.ok()); +} + } // namespace compute } // namespace arrow + From 8f26029a855994857c6ec718a4d72b9c224c9633 Mon Sep 17 00:00:00 2001 From: Pratyush Adhikari Date: Sun, 16 Aug 2026 09:21:42 +0530 Subject: [PATCH 2/2] GH-50879: [C++] Apply clang-format fixes --- .../arrow/compute/kernels/vector_replace.cc | 20 +++++++++---------- .../compute/kernels/vector_replace_test.cc | 4 +--- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_replace.cc b/cpp/src/arrow/compute/kernels/vector_replace.cc index 37056c3637d5..ec832074549a 100644 --- a/cpp/src/arrow/compute/kernels/vector_replace.cc +++ b/cpp/src/arrow/compute/kernels/vector_replace.cc @@ -333,10 +333,8 @@ struct ReplaceMaskImpl> { using BuilderType = typename TypeTraits::BuilderType; static Result ExecScalarMask(KernelContext* ctx, const ArraySpan& array, - const BooleanScalar& mask, - ExecValue replacements, - int64_t replacements_offset, - ExecResult* out) { + const BooleanScalar& mask, ExecValue replacements, + int64_t replacements_offset, ExecResult* out) { if (!mask.is_valid) { // mask = null: output is all-null array ARROW_ASSIGN_OR_RAISE( @@ -370,16 +368,15 @@ struct ReplaceMaskImpl> { } static Result ExecArrayMask(KernelContext* ctx, const ArraySpan& array, - const ArraySpan& mask, int64_t mask_offset, - ExecValue replacements, - int64_t replacements_offset, - ExecResult* out) { + const ArraySpan& mask, int64_t mask_offset, + ExecValue replacements, + int64_t replacements_offset, ExecResult* out) { // Build the output list array element-by-element. We cannot pre-allocate a flat // buffer (unlike fixed-width types) because the child-array size of each list slot // is unknown in advance, so we use a list builder and copy slots individually. std::unique_ptr raw_builder; - RETURN_NOT_OK(MakeBuilderExactIndex(ctx->memory_pool(), - array.type->GetSharedPtr(), &raw_builder)); + RETURN_NOT_OK(MakeBuilderExactIndex(ctx->memory_pool(), array.type->GetSharedPtr(), + &raw_builder)); auto& builder = checked_cast(*raw_builder); RETURN_NOT_OK(builder.Reserve(array.length)); @@ -1012,7 +1009,8 @@ void RegisterVectorReplace(FunctionRegistry* registry) { AddKernel(Type::LIST, ReplaceMask::GetSignature(Type::LIST), ReplaceMask::Exec, ReplaceMaskChunked::Exec, registry, func.get()); - AddKernel(Type::LARGE_LIST, ReplaceMask::GetSignature(Type::LARGE_LIST), + AddKernel(Type::LARGE_LIST, + ReplaceMask::GetSignature(Type::LARGE_LIST), ReplaceMask::Exec, ReplaceMaskChunked::Exec, registry, func.get()); DCHECK_OK(registry->AddFunction(std::move(func))); diff --git a/cpp/src/arrow/compute/kernels/vector_replace_test.cc b/cpp/src/arrow/compute/kernels/vector_replace_test.cc index 1051a43e7489..b810e338cea3 100644 --- a/cpp/src/arrow/compute/kernels/vector_replace_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_replace_test.cc @@ -2113,8 +2113,7 @@ TEST_F(TestFillNullType, TestFillOnNullType) { // Helper: assert ReplaceWithMask output for list types, calling ValidateFull. static void AssertReplaceWithMaskList(const Datum& values, const Datum& mask, - const Datum& replacements, - const Datum& expected) { + const Datum& replacements, const Datum& expected) { ASSERT_OK_AND_ASSIGN(auto actual, ReplaceWithMask(values, mask, replacements)); if (actual.is_array()) { ASSERT_OK(actual.make_array()->ValidateFull()); @@ -2331,4 +2330,3 @@ TEST(TestReplaceWithMaskList, TypeMismatch) { } // namespace compute } // namespace arrow -