Skip to content

GH-14736: [C++] Propagate struct validity into union children #51017

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
Hei-MaoM wants to merge 1 commit into
base: main
Choose a base branch
from
+304 −0
Open

GH-14736: [C++] Propagate struct validity into union children #51017

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
222 changes: 222 additions & 0 deletions cpp/src/arrow/array/array_nested.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <string>
#include <utility>
Expand Down Expand Up @@ -1004,6 +1005,218 @@ Result<std::shared_ptr<Array>> FixedSizeListArray::Flatten(
// ----------------------------------------------------------------------
// Struct

namespace {

Result<std::shared_ptr<ArrayData>> ApplyValidityBitmap(
const std::shared_ptr<ArrayData>& data, const uint8_t* validity,
int64_t validity_offset, MemoryPool* pool);

// Sparse union children are aligned with the union's logical slots, so the
// parent validity can be applied to every child at the same positions.
Result<std::shared_ptr<ArrayData>> ApplyValidityBitmapToSparseUnion(
const std::shared_ptr<ArrayData>& data, const uint8_t* validity,
int64_t validity_offset, MemoryPool* pool) {
ARROW_ASSIGN_OR_RAISE(auto type_codes,
data->buffers[1]->CopySlice(data->offset * sizeof(int8_t),
data->length * sizeof(int8_t), pool));

std::vector<std::shared_ptr<ArrayData>> children;
children.reserve(data->child_data.size());
for (const auto& child : data->child_data) {
auto child_view = child->Copy();
if (data->offset != 0 || data->length != child_view->length) {
child_view = child_view->Slice(data->offset, data->length);
}
ARROW_ASSIGN_OR_RAISE(
child_view, ApplyValidityBitmap(child_view, validity, validity_offset, pool));
children.push_back(std::move(child_view));
}

return ArrayData::Make(data->type, data->length, {nullptr, std::move(type_codes)},
std::move(children), /*null_count=*/0, /*offset=*/0);
}

struct DenseUnionChildSegment {
bool is_null;
int32_t input_offset;
int64_t length;
};

// Dense union offsets may be shared by multiple slots. Rebuild each child's
// referenced values in logical order so a null parent never invalidates a value
// that is also referenced by a valid parent, while offsets remain non-decreasing.
Result<std::shared_ptr<ArrayData>> ApplyValidityBitmapToDenseUnion(
const std::shared_ptr<ArrayData>& data, const uint8_t* validity,
int64_t validity_offset, MemoryPool* pool) {
ARROW_ASSIGN_OR_RAISE(auto type_codes,
data->buffers[1]->CopySlice(data->offset * sizeof(int8_t),
data->length * sizeof(int8_t), pool));
ARROW_ASSIGN_OR_RAISE(auto value_offsets, data->buffers[2]->CopySlice(
data->offset * sizeof(int32_t),
data->length * sizeof(int32_t), pool));

const auto& union_type = checked_cast<const UnionType&>(*data->type);
const auto* input_type_codes = data->GetValuesSafe<int8_t>(1);
const auto* input_value_offsets = data->GetValuesSafe<int32_t>(2);
auto* output_value_offsets = reinterpret_cast<int32_t*>(value_offsets->mutable_data());

const int num_children = union_type.num_fields();
std::vector<std::vector<DenseUnionChildSegment>> child_segments(num_children);
std::vector<int64_t> child_lengths(num_children, 0);
std::vector<bool> has_previous(num_children, false);
std::vector<bool> previous_was_valid(num_children, false);
std::vector<int32_t> previous_input_offsets(num_children, 0);
std::vector<int32_t> previous_output_offsets(num_children, 0);

for (int64_t i = 0; i < data->length; ++i) {
const int child_id = union_type.child_ids()[input_type_codes[i]];
const bool is_valid = bit_util::GetBit(validity, validity_offset + i);
const int32_t input_offset = input_value_offsets[i];

bool reuse_previous =
has_previous[child_id] && previous_was_valid[child_id] == is_valid;
if (reuse_previous && is_valid) {
reuse_previous = previous_input_offsets[child_id] == input_offset;
}

if (reuse_previous) {
output_value_offsets[i] = previous_output_offsets[child_id];
} else {
const int64_t output_offset = child_lengths[child_id];
if (output_offset > std::numeric_limits<int32_t>::max()) {
return Status::CapacityError(
"Dense union child offset exceeds the maximum int32 value");
}
output_value_offsets[i] = static_cast<int32_t>(output_offset);
++child_lengths[child_id];

auto& segments = child_segments[child_id];
const bool extends_previous =
is_valid && has_previous[child_id] && previous_was_valid[child_id] &&
static_cast<int64_t>(previous_input_offsets[child_id]) + 1 == input_offset;
if (extends_previous) {
++segments.back().length;
} else {
segments.push_back({/*is_null=*/!is_valid, input_offset, /*length=*/1});
}
}

has_previous[child_id] = true;
previous_was_valid[child_id] = is_valid;
previous_input_offsets[child_id] = input_offset;
previous_output_offsets[child_id] = output_value_offsets[i];
}

std::vector<std::shared_ptr<ArrayData>> children(num_children);
for (int child_id = 0; child_id < num_children; ++child_id) {
auto input_child = MakeArray(data->child_data[child_id]);
ArrayVector fragments;
fragments.reserve(child_segments[child_id].size());
std::shared_ptr<Array> null_value;
for (const auto& segment : child_segments[child_id]) {
if (segment.is_null) {
if (!null_value) {
ARROW_ASSIGN_OR_RAISE(
null_value, MakeArrayOfNull(union_type.field(child_id)->type(), 1, pool));
}
fragments.push_back(null_value);
} else {
fragments.push_back(input_child->Slice(segment.input_offset, segment.length));
}
}

std::shared_ptr<Array> child;
if (fragments.empty()) {
ARROW_ASSIGN_OR_RAISE(child,
MakeArrayOfNull(union_type.field(child_id)->type(), 0, pool));
} else if (fragments.size() == 1) {
child = std::move(fragments[0]);
} else {
ARROW_ASSIGN_OR_RAISE(child, Concatenate(fragments, pool));
}
children[child_id] = child->data();
}

return ArrayData::Make(data->type, data->length,
{nullptr, std::move(type_codes), std::move(value_offsets)},
std::move(children),
/*null_count=*/0, /*offset=*/0);
}

Result<std::shared_ptr<ArrayData>> ApplyValidityBitmapToAlwaysNullType(
const std::shared_ptr<ArrayData>& data, const uint8_t* validity,
int64_t validity_offset, MemoryPool* pool) {
auto array = MakeArray(data);
ArrayVector fragments;
int64_t run_start = 0;
while (run_start < data->length) {
const bool is_valid = bit_util::GetBit(validity, validity_offset + run_start);
int64_t run_end = run_start + 1;
while (run_end < data->length &&
bit_util::GetBit(validity, validity_offset + run_end) == is_valid) {
++run_end;
}
const int64_t run_length = run_end - run_start;
if (is_valid) {
fragments.push_back(array->Slice(run_start, run_length));
} else {
ARROW_ASSIGN_OR_RAISE(auto nulls, MakeArrayOfNull(data->type, run_length, pool));
fragments.push_back(std::move(nulls));
}
run_start = run_end;
}

if (fragments.size() == 1) {
return fragments[0]->data();
}
ARROW_ASSIGN_OR_RAISE(auto result, Concatenate(fragments, pool));
return result->data();
}

Result<std::shared_ptr<ArrayData>> ApplyValidityBitmap(
const std::shared_ptr<ArrayData>& data, const uint8_t* validity,
int64_t validity_offset, MemoryPool* pool) {
if (internal::CountSetBits(validity, validity_offset, data->length) == data->length) {
return data;
}

if (data->type->id() == Type::SPARSE_UNION) {
return ApplyValidityBitmapToSparseUnion(data, validity, validity_offset, pool);
}
if (data->type->id() == Type::DENSE_UNION) {
return ApplyValidityBitmapToDenseUnion(data, validity, validity_offset, pool);
}

const auto layout = data->type->layout();
if (!layout.buffers.empty() && layout.buffers[0].kind == DataTypeLayout::BITMAP) {
std::shared_ptr<Buffer> null_bitmap;
if (data->buffers[0]) {
ARROW_ASSIGN_OR_RAISE(
null_bitmap, BitmapAnd(pool, data->buffers[0]->data(), data->offset, validity,
validity_offset, data->length, data->offset));
} else {
ARROW_ASSIGN_OR_RAISE(null_bitmap,
AllocateEmptyBitmap(data->offset + data->length, pool));
CopyBitmap(validity, validity_offset, data->length, null_bitmap->mutable_data(),
data->offset);
}

auto result = data->Copy();
result->buffers[0] = std::move(null_bitmap);
result->null_count = kUnknownNullCount;
return result;
}

if (data->type->id() == Type::NA) {
return data;
}
// Run-end encoded arrays also have no top-level validity bitmap. Rebuild
// them from valid and null runs instead of attaching an invalid buffer.
return ApplyValidityBitmapToAlwaysNullType(data, validity, validity_offset, pool);
}

} // namespace

struct StructArray::Impl {
mutable ArrayVector boxed_fields_;
};
Expand Down Expand Up @@ -1161,6 +1374,15 @@ Result<std::shared_ptr<Array>> StructArray::GetFlattenedField(int index,
if (data_->offset != 0 || data_->length != child_data->length) {
child_data = child_data->Slice(data_->offset, data_->length);
}

if (null_bitmap && is_union(child_data->type->id())) {
// Union arrays have no top-level validity bitmap. Propagate the struct's
// validity into their children while preserving the union type codes.
ARROW_ASSIGN_OR_RAISE(child_data, ApplyValidityBitmap(child_data, null_bitmap_data_,
data_->offset, pool));
return MakeArray(child_data);
}

std::shared_ptr<Buffer> child_null_bitmap = child_data->buffers[0];
const int64_t child_offset = child_data->offset;

Expand Down
82 changes: 82 additions & 0 deletions cpp/src/arrow/array/array_struct_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,88 @@ TEST(StructArray, Flatten) {
}
}

TEST(StructArray, FlattenSparseUnion) {
auto type_ids = ArrayFromJSON(int8(), "[2, 7, 2, 7, 2]");
auto ints = ArrayFromJSON(int64(), "[10, 20, 30, 40, 50]");
auto strs = ArrayFromJSON(utf8(), R"(["a", "b", "c", "d", "e"])");
ASSERT_OK_AND_ASSIGN(
auto union_array,
SparseUnionArray::Make(*type_ids, {ints, strs}, {"ints", "strs"}, {2, 7}));

std::shared_ptr<Buffer> all_valid;
BitmapFromVector<bool>({true, true, true, true, true}, &all_valid);
auto struct_type = struct_({field("union", union_array->type())});
auto struct_array = std::make_shared<StructArray>(struct_type, union_array->length(),
ArrayVector{union_array}, all_valid);

ASSERT_OK_AND_ASSIGN(auto flattened, struct_array->GetFlattenedField(0));
ASSERT_OK(flattened->ValidateFull());
ASSERT_EQ(flattened->data()->buffers[0], nullptr);
AssertArraysEqual(*union_array, *flattened, /*verbose=*/true);

std::shared_ptr<Buffer> validity;
BitmapFromVector<bool>({true, true, false, true, true}, &validity);
struct_array = std::make_shared<StructArray>(struct_type, union_array->length(),
ArrayVector{union_array}, validity);
auto sliced = std::static_pointer_cast<StructArray>(struct_array->Slice(1, 3));

ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(0));
ASSERT_OK(flattened->ValidateFull());
const auto& flattened_union = checked_cast<const SparseUnionArray&>(*flattened);
ASSERT_EQ(flattened_union.data()->buffers[0], nullptr);
EXPECT_EQ(flattened_union.type_code(0), 7);
EXPECT_EQ(flattened_union.type_code(1), 2);
EXPECT_EQ(flattened_union.type_code(2), 7);
AssertArraysEqual(*ArrayFromJSON(int64(), "[20, null, 40]"), *flattened_union.field(0),
/*verbose=*/true);
AssertArraysEqual(*ArrayFromJSON(utf8(), R"(["b", null, "d"])"),
*flattened_union.field(1), /*verbose=*/true);
}

TEST(StructArray, FlattenDenseUnionWithSharedOffsets) {
auto type_ids = ArrayFromJSON(int8(), "[7, 2, 2, 2, 7, 7]");
auto value_offsets = ArrayFromJSON(int32(), "[0, 0, 0, 0, 1, 2]");
auto ints = ArrayFromJSON(int64(), "[10]");
auto strs = ArrayFromJSON(utf8(), R"(["pad", "a", "b"])");
ASSERT_OK_AND_ASSIGN(auto union_array,
DenseUnionArray::Make(*type_ids, *value_offsets, {ints, strs},
{"ints", "strs"}, {2, 7}));

std::shared_ptr<Buffer> validity;
BitmapFromVector<bool>({true, true, false, true, true, false}, &validity);
auto struct_type = struct_({field("union", union_array->type())});
auto struct_array = std::make_shared<StructArray>(struct_type, union_array->length(),
ArrayVector{union_array}, validity);
auto sliced = std::static_pointer_cast<StructArray>(struct_array->Slice(1, 5));

ASSERT_OK_AND_ASSIGN(auto flattened, sliced->GetFlattenedField(0));
ASSERT_OK(flattened->ValidateFull());
const auto& flattened_union = checked_cast<const DenseUnionArray&>(*flattened);
ASSERT_EQ(flattened_union.data()->buffers[0], nullptr);
for (int64_t i = 0; i < flattened_union.length(); ++i) {
EXPECT_EQ(flattened_union.type_code(i),
checked_cast<const DenseUnionArray&>(*union_array).type_code(i + 1));
}

auto flattened_ints = flattened_union.field(0);
EXPECT_FALSE(flattened_ints->IsNull(flattened_union.value_offset(0)));
EXPECT_TRUE(flattened_ints->IsNull(flattened_union.value_offset(1)));
EXPECT_FALSE(flattened_ints->IsNull(flattened_union.value_offset(2)));
EXPECT_EQ(checked_cast<const Int64Array&>(*flattened_ints)
.Value(flattened_union.value_offset(0)),
10);
EXPECT_EQ(checked_cast<const Int64Array&>(*flattened_ints)
.Value(flattened_union.value_offset(2)),
10);

auto flattened_strs = flattened_union.field(1);
EXPECT_FALSE(flattened_strs->IsNull(flattened_union.value_offset(3)));
EXPECT_TRUE(flattened_strs->IsNull(flattened_union.value_offset(4)));
EXPECT_EQ(checked_cast<const StringArray&>(*flattened_strs)
.GetString(flattened_union.value_offset(3)),
"a");
}

/// ARROW-7740: Flattening a slice shouldn't affect the parent array.
TEST(StructArray, FlattenOfSlice) {
auto a = ArrayFromJSON(int32(), "[4, 5]");
Expand Down