Skip to content

GH-50879: [C++] Implement replace_with_mask for List and LargeList types #50880

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
pratyushadk wants to merge 2 commits into
base: main
Choose a base branch
from
+344 −3
Open

GH-50879: [C++] Implement replace_with_mask for List and LargeList types #50880

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
542eb21
GH-50879: [C++] Implement replace_with_mask for List and LargeList types
pratyushadk Aug 15, 2026
8f26029
GH-50879: [C++] Apply clang-format fixes
pratyushadk Aug 16, 2026
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
125 changes: 122 additions & 3 deletions cpp/src/arrow/compute/kernels/vector_replace.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -119,7 +120,8 @@ struct ReplaceMaskImpl {};

template <typename Type>
struct ReplaceMaskImpl<
Type, enable_if_t<!(is_base_binary_type<Type>::value || is_null_type<Type>::value)>> {
Type, enable_if_t<!(is_base_binary_type<Type>::value || is_null_type<Type>::value ||
is_var_length_list_type<Type>::value)>> {
static Result<int64_t> ExecScalarMask(KernelContext* ctx, const ArraySpan& array,
const BooleanScalar& mask, ExecValue replacements,
int64_t replacements_offset, ExecResult* out) {
Expand Down Expand Up @@ -322,6 +324,108 @@ struct ReplaceMaskImpl<Type, enable_if_base_binary<Type>> {
}
};

// Specialization for variable-size list types (list<T> and large_list<T>).
// Each list element is copied individually using AppendArraySlice, mirroring
// the enable_if_base_binary specialization's per-element builder approach.
template <typename Type>
struct ReplaceMaskImpl<Type, enable_if_var_size_list<Type>> {
using offset_type = typename Type::offset_type;
using BuilderType = typename TypeTraits<Type>::BuilderType;

static Result<int64_t> 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<ArrayData> 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<int64_t> 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<ArrayBuilder> raw_builder;
RETURN_NOT_OK(MakeBuilderExactIndex(ctx->memory_pool(), array.type->GetSharedPtr(),
&raw_builder));
auto& builder = checked_cast<BuilderType&>(*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<BooleanType>(
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<Array> 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,
Expand Down Expand Up @@ -862,8 +966,10 @@ void RegisterVectorFunction(FunctionRegistry* registry,
GenerateTypeAgnosticVarBinaryBase<ChunkedFunctor>(*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"
}
Expand Down Expand Up @@ -897,16 +1003,29 @@ void RegisterVectorReplace(FunctionRegistry* registry) {
auto func = std::make_shared<VectorFunction>("replace_with_mask", Arity::Ternary(),
replace_with_mask_doc);
RegisterVectorFunction<ReplaceMask, ReplaceMaskChunked>(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<ListType>::GetSignature(Type::LIST),
ReplaceMask<ListType>::Exec, ReplaceMaskChunked<ListType>::Exec, registry,
func.get());
AddKernel(Type::LARGE_LIST,
ReplaceMask<LargeListType>::GetSignature(Type::LARGE_LIST),
ReplaceMask<LargeListType>::Exec, ReplaceMaskChunked<LargeListType>::Exec,
registry, func.get());
DCHECK_OK(registry->AddFunction(std::move(func)));
}
{
auto func = std::make_shared<VectorFunction>("fill_null_forward", Arity::Unary(),
fill_null_forward_doc);
RegisterVectorFunction<FillNullForward, FillNullForwardChunked>(registry, func);
DCHECK_OK(registry->AddFunction(std::move(func)));
}
{
auto func = std::make_shared<VectorFunction>("fill_null_backward", Arity::Unary(),
fill_null_backward_doc);
RegisterVectorFunction<FillNullBackward, FillNullBackwardChunked>(registry, func);
DCHECK_OK(registry->AddFunction(std::move(func)));
}
}
} // namespace internal
Expand Down
222 changes: 222 additions & 0 deletions cpp/src/arrow/compute/kernels/vector_replace_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2106,5 +2106,227 @@ TEST_F(TestFillNullType, TestFillOnNullType) {
this->AssertFillNullArray(FillNullBackward, this->array(R"([null, null])"),
this->array(R"([null, null])"));
}

// ----------------------------------------------------------------------------
// Tests for replace_with_mask with list<int32> and large_list<int32>
// ----------------------------------------------------------------------------

// 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<BooleanScalar>(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<BooleanScalar>(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<BooleanScalar>();
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<ChunkedArray>(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<ChunkedArray>(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<ChunkedArray>(ArrayVector{chunk0, chunk1});
auto mask = std::make_shared<BooleanScalar>(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<ChunkedArray>(ArrayVector{exp0, exp1});
AssertReplaceWithMaskList(Datum(values), Datum(mask), replacements, Datum(expected));
}

// large_list<int32>: 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
Loading