Skip to content

GH-43146: [C++][Python] Validate list value nullability when casting #53948

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

Closed
1932819331hy-debug wants to merge 1 commit into from
+286 −8
Closed

GH-43146: [C++][Python] Validate list value nullability when casting #53948

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
137 changes: 129 additions & 8 deletions cpp/src/arrow/compute/kernels/scalar_cast_nested.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
#include <utility>
#include <vector>

#include "arrow/array/array_nested.h"
#include "arrow/array/builder_nested.h"
#include "arrow/array/builder_primitive.h"
#include "arrow/compute/api_scalar.h"
Expand All @@ -40,6 +41,96 @@ namespace compute::internal {

namespace {

bool CanMaskNestedValues(Type::type type) {
return is_list_like(type) || is_list_view(type) || type == Type::STRUCT ||
type == Type::DICTIONARY;
}

template <typename ListType, typename Visit>
Status VisitValidListValueRanges(const ArraySpan& in_array, Visit&& visit) {
for (int64_t i = 0; i < in_array.length; ++i) {
if (in_array.IsNull(i)) {
continue;
}
int64_t offset, length;
if constexpr (is_fixed_size_list_type<ListType>::value) {
length = checked_cast<const ListType&>(*in_array.type).list_size();
offset = (in_array.offset + i) * length;
} else {
const auto* offsets = in_array.GetValues<typename ListType::offset_type>(1);
offset = offsets[i];
if constexpr (is_list_view_type<ListType>::value) {
length = in_array.GetValues<typename ListType::offset_type>(2)[i];
} else {
length = offsets[i + 1] - offset;
}
}
RETURN_NOT_OK(visit(offset, length));
}
return Status::OK();
}

template <typename DestType>
Status CheckListValueNullability(const ArrayData& out_array) {
const auto& out_type = checked_cast<const DestType&>(*out_array.type);
if (out_type.value_field()->nullable() ||
out_array.child_data[0]->GetNullCount() == 0) {
return Status::OK();
}
const ArraySpan array(out_array);
auto values = array.child_data[0];
const int64_t values_offset = values.offset;
return VisitValidListValueRanges<DestType>(array, [&](int64_t offset, int64_t length) {
values.SetSlice(values_offset + offset, length);
if (values.GetNullCount() > 0) {
return Status::Invalid("List values have nulls. Can't cast to non-nullable ",
"value field '", out_type.value_field()->name(), "'");
}
return Status::OK();
});
}

template <typename SrcType>
Status MaskNestedListValues(KernelContext* ctx, const ArraySpan& in_array,
std::shared_ptr<ArrayData>* values) {
if (!CanMaskNestedValues((*values)->type->id())) {
return Status::OK();
}
if (in_array.GetNullCount() == 0) {
if constexpr (is_fixed_size_list_type<SrcType>::value) {
const int64_t size = checked_cast<const SrcType&>(*in_array.type).list_size();
if (in_array.offset == 0 && in_array.length * size == (*values)->length) {
return Status::OK();
}
} else if constexpr (!is_list_view_type<SrcType>::value) {
const auto* offsets = in_array.GetValues<typename SrcType::offset_type>(1);
if (offsets[0] == 0 && offsets[in_array.length] == (*values)->length) {
return Status::OK();
}
}
}

// Recursive casts must not validate descendants hidden by a null parent or slice.
ARROW_ASSIGN_OR_RAISE(
auto bitmap,
AllocateEmptyBitmap((*values)->offset + (*values)->length, ctx->memory_pool()));
RETURN_NOT_OK(
VisitValidListValueRanges<SrcType>(in_array, [&](int64_t offset, int64_t length) {
offset += (*values)->offset;
if ((*values)->buffers[0]) {
CopyBitmap((*values)->buffers[0]->data(), offset, length,
bitmap->mutable_data(), offset);
} else {
bit_util::SetBitsTo(bitmap->mutable_data(), offset, length, true);
}
return Status::OK();
}));
*values = (*values)->Copy();
(*values)->buffers[0] = std::move(bitmap);
(*values)->null_count = kUnknownNullCount;
return Status::OK();
}

// (Large)List<T> -> (Large)List<U>

// TODO(wesm): memory could be preallocated here and it would make
Expand Down Expand Up @@ -108,7 +199,8 @@ struct CastList {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
const CastOptions& options = CastState::Get(ctx);

auto child_type = checked_cast<const DestType&>(*out->type()).value_type();
const auto& out_type = checked_cast<const DestType&>(*out->type());
auto child_type = out_type.value_type();

const ArraySpan& in_array = batch[0].array;

Expand All @@ -118,6 +210,7 @@ struct CastList {
out_array->buffers[1] = in_array.GetBuffer(1);

std::shared_ptr<ArrayData> values = in_array.child_data[0].ToArrayData();
RETURN_NOT_OK(MaskNestedListValues<SrcType>(ctx, in_array, &values));

RETURN_NOT_OK(HandleOffsets(ctx, in_array, out_array, &values));

Expand All @@ -127,7 +220,7 @@ struct CastList {

DCHECK(cast_values.is_array());
out_array->child_data.push_back(cast_values.array());
return Status::OK();
return CheckListValueNullability<DestType>(*out_array);
}
};

Expand All @@ -148,7 +241,8 @@ struct CastFixedToVarList {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
const CastOptions& options = CastState::Get(ctx);

auto child_type = checked_cast<const DestType&>(*out->type()).value_type();
const auto& out_type = checked_cast<const DestType&>(*out->type());
auto child_type = out_type.value_type();

const ArraySpan& in_array = batch[0].array;

Expand All @@ -171,6 +265,7 @@ struct CastFixedToVarList {

// Handle values
std::shared_ptr<ArrayData> values = in_array.child_data[0].ToArrayData();
RETURN_NOT_OK(MaskNestedListValues<FixedSizeListType>(ctx, in_array, &values));
if (in_array.offset > 0) {
values = values->Slice(in_array.offset * list_size, in_array.length * list_size);
}
Expand All @@ -179,7 +274,7 @@ struct CastFixedToVarList {
DCHECK(cast_values.is_array());
out_array->child_data.push_back(cast_values.array());

return Status::OK();
return CheckListValueNullability<DestType>(*out_array);
}
};

Expand Down Expand Up @@ -269,6 +364,7 @@ struct CastVarToFixedList {

// Handle values
std::shared_ptr<ArrayData> values = in_array.child_data[0].ToArrayData();
RETURN_NOT_OK(MaskNestedListValues<SrcType>(ctx, in_array, &values));
ARROW_ASSIGN_OR_RAISE(Datum cast_values_datum,
Cast(values, child_type, options, ctx->exec_context()));

Expand All @@ -291,7 +387,7 @@ struct CastVarToFixedList {

out_array->child_data.emplace_back(std::move(cast_values));

return Status::OK();
return CheckListValueNullability<FixedSizeListType>(*out_array);
}
};

Expand All @@ -311,6 +407,7 @@ struct CastFixedList {

const ArraySpan& in_array = batch[0].array;
std::shared_ptr<ArrayData> values = in_array.child_data[0].ToArrayData();
RETURN_NOT_OK(MaskNestedListValues<FixedSizeListType>(ctx, in_array, &values));
ArrayData* out_array = out->array_data().get();
out_array->buffers[0] = in_array.GetBuffer(0);

Expand All @@ -322,7 +419,7 @@ struct CastFixedList {
Cast(values, child_type, options, ctx->exec_context()));
DCHECK(cast_values.is_array());
out_array->child_data.push_back(cast_values.array());
return Status::OK();
return CheckListValueNullability<FixedSizeListType>(*out_array);
}
};

Expand Down Expand Up @@ -382,7 +479,7 @@ struct CastStruct {
out_array->child_data.push_back(nulls->data());
} else {
const auto& in_field = in_type.field(in_field_index);
const auto& in_values = (in_array.child_data[in_field_index].ToArrayData()->Slice(
auto in_values = (in_array.child_data[in_field_index].ToArrayData()->Slice(
in_array.offset, in_array.length));
if (in_field->nullable() && !out_field->nullable() &&
in_values->GetNullCount() > 0) {
Expand All @@ -391,6 +488,12 @@ struct CastStruct {
" has nulls. Can't cast to non-nullable field '", out_field->name(),
"' of type ", out_field_type->ToString());
}
if (in_array.GetNullCount() > 0 && CanMaskNestedValues(in_values->type->id())) {
auto in_struct = checked_pointer_cast<StructArray>(in_array.ToArray());
ARROW_ASSIGN_OR_RAISE(auto flattened, in_struct->GetFlattenedField(
in_field_index, ctx->memory_pool()));
in_values = flattened->data();
}
ARROW_ASSIGN_OR_RAISE(Datum cast_values, Cast(in_values, out_field_type, options,
ctx->exec_context()));
DCHECK(cast_values.is_array());
Expand Down Expand Up @@ -443,7 +546,25 @@ struct CastMap {
in_array.offset, in_array.length));
}

RETURN_NOT_OK(CastListImpl::HandleOffsets(ctx, in_array, out_array, &entries));
const auto* offsets = in_array.GetValues<MapType::offset_type>(1);
if (in_array.GetNullCount() > 0 || offsets[0] != 0 ||
offsets[in_array.length] != entries->length) {
// Drop hidden entries before recursively casting their non-nullable keys.
auto in_map = checked_pointer_cast<MapArray>(in_array.ToArray());
ARROW_ASSIGN_OR_RAISE(auto flattened, in_map->Flatten(ctx->memory_pool()));
entries = flattened->data();
using offset_type = typename DestType::offset_type;
ARROW_ASSIGN_OR_RAISE(out_array->buffers[1],
ctx->Allocate(sizeof(offset_type) * (in_array.length + 1)));
auto* out_offsets = out_array->GetMutableValues<offset_type>(1);
out_offsets[0] = 0;
for (int64_t i = 0; i < in_array.length; ++i) {
out_offsets[i + 1] =
out_offsets[i] + (in_array.IsValid(i) ? offsets[i + 1] - offsets[i] : 0);
}
} else {
RETURN_NOT_OK(CastListImpl::HandleOffsets(ctx, in_array, out_array, &entries));
}

// Handle keys
const std::shared_ptr<ArrayData>& keys =
Expand Down
138 changes: 138 additions & 0 deletions cpp/src/arrow/compute/kernels/scalar_cast_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3641,6 +3641,144 @@ TEST(Cast, ListToListOptionsPassthru) {
}
}

TEST(Cast, ListValueNullability) {
auto list_types = [](const std::shared_ptr<Field>& value_field) {
return std::vector<std::shared_ptr<DataType>>{
list(value_field), large_list(value_field), fixed_size_list(value_field, 2)};
};
for (const auto& src_type : list_types(field("item", int32()))) {
for (const auto& value_type : {int32(), int64()}) {
for (const auto& dest_type :
list_types(field("item", value_type, /*nullable=*/false))) {
ARROW_SCOPED_TRACE("src_type = ", src_type->ToString(),
", dest_type = ", dest_type->ToString());
auto with_nulls = ArrayFromJSON(src_type, "[[1, null], [3, 4]]");
CheckCastFails(with_nulls, CastOptions::Safe(dest_type));
CheckCastFails(with_nulls, CastOptions::Unsafe(dest_type));
CheckCast(ArrayFromJSON(src_type, "[[1, 2], null, [3, 4]]"),
ArrayFromJSON(dest_type, "[[1, 2], null, [3, 4]]"));
CheckCast(MaskArrayWithNullsAt(with_nulls, {0}),
ArrayFromJSON(dest_type, "[null, [3, 4]]"));

// Nulls before and after a slice do not belong to the cast.
auto sliced = ArrayFromJSON(src_type, "[[null, 0], [1, 2], [3, 4], [null, null]]")
->Slice(1, 2);
CheckCast(sliced, ArrayFromJSON(dest_type, "[[1, 2], [3, 4]]"));
CheckCast(with_nulls->Slice(1, 0), ArrayFromJSON(dest_type, "[]"));
}
}
}
}

TEST(Cast, ListValueNullabilitySlicedChild) {
for (const auto& src_type :
{list(int32()), large_list(int32()), fixed_size_list(int32(), 2)}) {
auto dest_field = field("item", int64(), /*nullable=*/false);
auto dest_type = src_type->id() == Type::LIST ? list(dest_field)
: src_type->id() == Type::LARGE_LIST
? large_list(dest_field)
: fixed_size_list(dest_field, 2);
auto src = ArrayFromJSON(src_type, "[[1, 2], [3, 4]]");
auto values = ArrayFromJSON(int32(), "[null, null, 1, 2, 3, 4, null]");
auto data = src->data()->Copy();
data->child_data[0] = values->Slice(2, 4)->data();
CheckCast(MakeArray(data), ArrayFromJSON(dest_type, "[[1, 2], [3, 4]]"));

auto null_values = ArrayFromJSON(int32(), "[0, 0, 1, null, 3, 4]");
data->child_data[0] = null_values->Slice(2, 4)->data()->Copy();
data->child_data[0]->null_count = kUnknownNullCount;
CheckCastFails(MakeArray(data), CastOptions::Safe(dest_type));
}
}

TEST(Cast, ListValueNullabilityDictionaryValues) {
auto values =
DictArrayFromJSON(dictionary(int8(), int32()), "[0, 1, 1, 1]", "[null, 10]");
ASSERT_EQ(values->null_count(), 0);
for (bool nullable : {false, true}) {
auto src_field = field("item", values->type(), nullable);
auto dest_field = field("item", int32(), /*nullable=*/false);
for (const auto& types :
{std::make_pair(list(src_field), list(dest_field)),
std::make_pair(large_list(src_field), large_list(dest_field)),
std::make_pair(fixed_size_list(src_field, 2),
fixed_size_list(dest_field, 2))}) {
auto data = ArrayFromJSON(types.second, "[[1, 2], [3, 4]]")->data()->Copy();
data->type = types.first;
data->child_data[0] = values->data();
auto src = MakeArray(data);
CheckCastFails(src, CastOptions::Safe(types.second));
CheckCastFails(src, CastOptions::Unsafe(types.second));
CheckCast(MaskArrayWithNullsAt(src, {0}),
ArrayFromJSON(types.second, "[null, [10, 10]]"));
}
}
}

TEST(Cast, ListValueNullabilityNestedMasked) {
auto leaf_type = list(int32());
auto non_nullable_leaf = list(field("item", int32(), /*nullable=*/false));
for (const auto& types :
{std::make_pair(list(leaf_type), list(non_nullable_leaf)),
std::make_pair(large_list(leaf_type), large_list(non_nullable_leaf)),
std::make_pair(fixed_size_list(leaf_type, 2),
fixed_size_list(non_nullable_leaf, 2))}) {
auto src = ArrayFromJSON(types.first, "[[[null], [0]], [[1], [2]]]");
CheckCast(MaskArrayWithNullsAt(src, {0}),
ArrayFromJSON(types.second, "[null, [[1], [2]]]"));
CheckCast(src->Slice(1, 1), ArrayFromJSON(types.second, "[[[1], [2]]]"));

auto suffix = ArrayFromJSON(types.first, "[[[1], [2]], [[null], [0]]]");
CheckCast(suffix->Slice(0, 1), ArrayFromJSON(types.second, "[[[1], [2]]]"));
}

auto deep = ArrayFromJSON(list(list(leaf_type)), "[[[[null]]], [[[1]]]]");
CheckCast(MaskArrayWithNullsAt(deep, {0}),
ArrayFromJSON(list(list(non_nullable_leaf)), "[null, [[[1]]]]"));

auto struct_type = struct_({field("a", leaf_type)});
auto dest_struct_type = struct_({field("a", non_nullable_leaf)});
auto structs = ArrayFromJSON(struct_type, R"([{"a": [null]}, {"a": [1]}])");
CheckCast(MaskArrayWithNullsAt(structs, {0}),
ArrayFromJSON(dest_struct_type, R"([null, {"a": [1]}])"));

auto mixed = ArrayFromJSON(list(struct_type), R"([[{"a": [null]}], [{"a": [1]}]])");
CheckCast(MaskArrayWithNullsAt(mixed, {0}),
ArrayFromJSON(list(dest_struct_type), R"([null, [{"a": [1]}]])"));

auto map_type = map(utf8(), leaf_type);
auto dest_map_type = map(utf8(), non_nullable_leaf);
auto maps = ArrayFromJSON(map_type, R"([[["k", [null]]], [["k", [1]]]])");
CheckCast(MaskArrayWithNullsAt(maps, {0}),
ArrayFromJSON(dest_map_type, R"([null, [["k", [1]]]])"));
CheckCast(maps->Slice(1, 1), ArrayFromJSON(dest_map_type, R"([[["k", [1]]]])"));
auto map_suffix = ArrayFromJSON(map_type, R"([[["k", [1]]], [["k", [null]]]])");
CheckCast(map_suffix->Slice(0, 1), ArrayFromJSON(dest_map_type, R"([[["k", [1]]]])"));

auto key_map_type = map(leaf_type, int32());
auto dest_key_map_type = map(non_nullable_leaf, int32());
auto key_maps = ArrayFromJSON(key_map_type, "[[[[null], 1]], [[[1], 2]]]");
CheckCast(MaskArrayWithNullsAt(key_maps, {0}),
ArrayFromJSON(dest_key_map_type, "[null, [[[1], 2]]]"));
}

TEST(Cast, ListValueNullabilityNestedAndNullType) {
auto nested = ArrayFromJSON(list(list(int32())), "[[[1, null], []], []]");
auto non_nullable = list(list(field("item", int32(), /*nullable=*/false)));
CheckCastFails(nested, CastOptions::Safe(non_nullable));

auto null_values = ArrayFromJSON(list(null()), "[[null], []]");
CheckCastFails(null_values,
CastOptions::Safe(list(field("item", int32(), /*nullable=*/false))));
CheckCast(null_values->Slice(1, 1),
ArrayFromJSON(list(field("item", int32(), /*nullable=*/false)), "[[]]"));

for (const auto& src_type : {list_view(int32()), large_list_view(int32())}) {
CheckCastFails(ArrayFromJSON(src_type, "[[1, null], [3, 4]]"),
CastOptions::Safe(list(field("item", int32(), /*nullable=*/false))));
}
}

static void CheckFSLToFSL(const std::vector<std::shared_ptr<DataType>>& value_types,
const std::string& json_data,
const std::string& tweaked_val_bit_string,
Expand Down
Loading
Loading