Skip to content

Commit

Permalink
Add C.41 Subarray constructor
Browse files Browse the repository at this point in the history 
+ Unstatus subarray_from_capnp
  • Loading branch information
@shaunrd0
shaunrd0 committed Jan 31, 2024
1 parent 8a9b23d commit 7e6d359
Show file tree
Hide file tree
Showing 5 changed files with 197 additions and 109 deletions.
10 changes: 7 additions & 3 deletions test/support/src/serialization_wrappers.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,11 +204,15 @@ void tiledb_subarray_serialize(
.ok());
// Deserialize
tiledb_subarray_t* deserialized_subarray;
auto layout = (*subarray)->subarray_->layout();
auto stats = (*subarray)->subarray_->stats();
shared_ptr<Logger> dummy_logger = make_shared<Logger>(HERE(), "");

tiledb::test::require_tiledb_ok(
ctx, tiledb_subarray_alloc(ctx, array, &deserialized_subarray));
REQUIRE(tiledb::sm::serialization::subarray_from_capnp(
builder, deserialized_subarray->subarray_)
.ok());
*deserialized_subarray->subarray_ =
tiledb::sm::serialization::subarray_from_capnp(
builder, array->array_.get(), layout, stats, dummy_logger);
*subarray = deserialized_subarray;
#endif
}
151 changes: 87 additions & 64 deletions tiledb/sm/serialization/query.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,89 +265,112 @@ Status subarray_to_capnp(
return Status::Ok();
}

Status subarray_from_capnp(
const capnp::Subarray::Reader& reader, Subarray* subarray) {
RETURN_NOT_OK(subarray->set_coalesce_ranges(reader.getCoalesceRanges()));
Subarray subarray_from_capnp(
const capnp::Subarray::Reader& reader,
const Array* array,
Layout layout,
stats::Stats* parent_stats,
shared_ptr<Logger> logger) {
bool coalesce_ranges = reader.getCoalesceRanges();
auto ranges_reader = reader.getRanges();

uint32_t dim_num = ranges_reader.size();
std::vector<RangeSetAndSuperset> range_subset(dim_num);
std::vector<bool> is_default(dim_num, false);
for (uint32_t i = 0; i < dim_num; i++) {
auto range_reader = ranges_reader[i];
Datatype type = Datatype::UINT8;
RETURN_NOT_OK(datatype_enum(range_reader.getType(), &type));
throw_if_not_ok(datatype_enum(range_reader.getType(), &type));
auto dim = array->array_schema_latest().dimension_ptr(i);

auto data = range_reader.getBuffer();
auto data_ptr = data.asBytes();
bool implicitly_initialized = range_reader.getHasDefaultRange();
range_subset[i] =
RangeSetAndSuperset(dim->type(), dim->domain(), true, coalesce_ranges);
is_default[i] = implicitly_initialized;
if (range_reader.hasBufferSizes()) {
auto ranges = range_buffers_from_capnp(range_reader);
RETURN_NOT_OK(subarray->set_ranges_for_dim(i, ranges));

// Set default indicator
subarray->set_is_default(i, range_reader.getHasDefaultRange());
// If the range is implicitly initialized, the RangeSetAndSuperset
// constructor will initialize the ranges to the domain.
if (!implicitly_initialized) {
// Edge case for dimension labels where there are only label ranges set.
if (ranges.empty()) {
range_subset[i] = RangeSetAndSuperset(
dim->type(), dim->domain(), false, coalesce_ranges);
}
// Add custom ranges, clearing any implicit ranges previously set.
for (const auto& range : ranges) {
throw_if_not_ok(range_subset[i].add_range_unrestricted(range));
}
}
} else {
// Handle 1.7 style ranges where there is a single range with no sizes
auto data = range_reader.getBuffer();
auto data_ptr = data.asBytes();
Range range(data_ptr.begin(), data.size());
RETURN_NOT_OK(subarray->set_ranges_for_dim(i, {range}));
subarray->set_is_default(i, range_reader.getHasDefaultRange());
throw_if_not_ok(range_subset[i].add_range_unrestricted(range));
}
}

std::vector<optional<Subarray::LabelRangeSubset>> label_range_subset(
dim_num, nullopt);
if (reader.hasLabelRanges()) {
subarray->add_default_label_ranges(dim_num);
auto label_ranges_reader = reader.getLabelRanges();
uint32_t label_num = label_ranges_reader.size();
for (uint32_t i = 0; i < label_num; i++) {
auto label_range_reader = label_ranges_reader[i];
auto dim_id = label_range_reader.getDimensionId();
auto dim_index = label_range_reader.getDimensionId();
auto dim = array->array_schema_latest().dimension_ptr(dim_index);
auto label_name = label_range_reader.getName();

// Deserialize ranges for this dim label
auto range_reader = label_range_reader.getRanges();
auto ranges = range_buffers_from_capnp(range_reader);

// Set ranges for this dim label on the subarray
subarray->set_label_ranges_for_dim(dim_id, label_name, ranges);
label_range_subset[dim_index] = {
label_name, dim->type(), coalesce_ranges};
}
}

std::unordered_map<std::string, std::vector<Range>> attr_range_subset;
if (reader.hasAttributeRanges()) {
std::unordered_map<std::string, std::vector<Range>> attr_ranges;
auto attr_ranges_reader = reader.getAttributeRanges();
if (attr_ranges_reader.hasEntries()) {
for (auto attr_ranges_entry : attr_ranges_reader.getEntries()) {
auto range_reader = attr_ranges_entry.getValue();
auto key = std::string_view{
attr_ranges_entry.getKey().cStr(),
attr_ranges_entry.getKey().size()};
attr_ranges[std::string{key}] = range_buffers_from_capnp(range_reader);
for (auto entry : attr_ranges_reader.getEntries()) {
auto range_reader = entry.getValue();
std::string key{entry.getKey().cStr(), entry.getKey().size()};
attr_range_subset[key] = range_buffers_from_capnp(range_reader);
}
}

for (const auto& attr_range : attr_ranges)
subarray->set_attribute_ranges(attr_range.first, attr_range.second);
}

// If cap'n proto object has stats set it on c++ object
Subarray s(array, layout, parent_stats, logger, true);
if (reader.hasStats()) {
stats::Stats* stats = subarray->stats();
// We should always have a stats here
if (stats != nullptr) {
RETURN_NOT_OK(stats_from_capnp(reader.getStats(), stats));
}
throw_if_not_ok(stats_from_capnp(reader.getStats(), s.stats()));
}

std::vector<unsigned> relevant_fragments;
if (reader.hasRelevantFragments()) {
auto relevant_fragments = reader.getRelevantFragments();
size_t count = relevant_fragments.size();
std::vector<unsigned> rf;
rf.reserve(count);
auto reader_rf = reader.getRelevantFragments();
size_t count = reader_rf.size();
relevant_fragments.reserve(count);
for (size_t i = 0; i < count; i++) {
rf.emplace_back(relevant_fragments[i]);
relevant_fragments.emplace_back(reader_rf[i]);
}

subarray->relevant_fragments() = RelevantFragments(rf);
}

return Status::Ok();
return {
array,
layout,
reader.hasStats() ? s.stats() : parent_stats,
logger,
range_subset,
is_default,
label_range_subset,
attr_range_subset,
relevant_fragments,
coalesce_ranges};
}

Status subarray_partitioner_to_capnp(
Expand Down Expand Up @@ -461,8 +484,8 @@ Status subarray_partitioner_from_capnp(
RETURN_NOT_OK(layout_enum(subarray_reader.getLayout(), &layout));

// Subarray, which is used to initialize the partitioner.
Subarray subarray(array, layout, query_stats, dummy_logger, true);
RETURN_NOT_OK(subarray_from_capnp(reader.getSubarray(), &subarray));
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query_stats, dummy_logger);
*partitioner = SubarrayPartitioner(
&config,
subarray,
Expand Down Expand Up @@ -519,10 +542,12 @@ Status subarray_partitioner_from_capnp(
partition_info->end_ = partition_info_reader.getEnd();
partition_info->split_multi_range_ =
partition_info_reader.getSplitMultiRange();
partition_info->partition_ =
Subarray(array, layout, query_stats, dummy_logger, true);
RETURN_NOT_OK(subarray_from_capnp(
partition_info_reader.getSubarray(), &partition_info->partition_));
partition_info->partition_ = subarray_from_capnp(
partition_info_reader.getSubarray(),
array,
layout,
query_stats,
dummy_logger);

if (compute_current_tile_overlap) {
throw_if_not_ok(partition_info->partition_.precompute_tile_overlap(
Expand All @@ -542,20 +567,18 @@ Status subarray_partitioner_from_capnp(
auto sr_reader = state_reader.getSingleRange();
const unsigned num_sr = sr_reader.size();
for (unsigned i = 0; i < num_sr; i++) {
auto subarray_reader_ = sr_reader[i];
state->single_range_.emplace_back(
array, layout, query_stats, dummy_logger, true);
Subarray& subarray_ = state->single_range_.back();
RETURN_NOT_OK(subarray_from_capnp(subarray_reader_, &subarray_));
auto subarray_reader = sr_reader[i];
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query_stats, dummy_logger);
state->single_range_.push_back(subarray);
}
auto m_reader = state_reader.getMultiRange();
const unsigned num_m = m_reader.size();
for (unsigned i = 0; i < num_m; i++) {
auto subarray_reader_ = m_reader[i];
state->multi_range_.emplace_back(
array, layout, query_stats, dummy_logger, true);
Subarray& subarray_ = state->multi_range_.back();
RETURN_NOT_OK(subarray_from_capnp(subarray_reader_, &subarray_));
auto subarray_reader = m_reader[i];
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query_stats, dummy_logger);
state->multi_range_.push_back(subarray);
}

// Overall mem budget
Expand Down Expand Up @@ -1131,9 +1154,9 @@ Status reader_from_capnp(
RETURN_NOT_OK(layout_enum(reader_reader.getLayout(), &layout));

// Subarray
Subarray subarray(array, layout, query->stats(), dummy_logger, true);
auto subarray_reader = reader_reader.getSubarray();
RETURN_NOT_OK(subarray_from_capnp(subarray_reader, &subarray));
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query->stats(), dummy_logger);
RETURN_NOT_OK(query->set_subarray_unsafe(subarray));

// Read state
Expand Down Expand Up @@ -1172,9 +1195,9 @@ Status index_reader_from_capnp(
RETURN_NOT_OK(layout_enum(reader_reader.getLayout(), &layout));

// Subarray
Subarray subarray(array, layout, query->stats(), dummy_logger, true);
auto subarray_reader = reader_reader.getSubarray();
RETURN_NOT_OK(subarray_from_capnp(subarray_reader, &subarray));
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query->stats(), dummy_logger);
RETURN_NOT_OK(query->set_subarray_unsafe(subarray));

// Read state
Expand Down Expand Up @@ -1214,9 +1237,9 @@ Status dense_reader_from_capnp(
RETURN_NOT_OK(layout_enum(reader_reader.getLayout(), &layout));

// Subarray
Subarray subarray(array, layout, query->stats(), dummy_logger, true);
auto subarray_reader = reader_reader.getSubarray();
RETURN_NOT_OK(subarray_from_capnp(subarray_reader, &subarray));
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query->stats(), dummy_logger);
RETURN_NOT_OK(query->set_subarray_unsafe(subarray));

// Read state
Expand Down Expand Up @@ -2253,9 +2276,9 @@ Status query_from_capnp(

// Subarray
if (writer_reader.hasSubarrayRanges()) {
Subarray subarray(array, layout, query->stats(), dummy_logger, true);
auto subarray_reader = writer_reader.getSubarrayRanges();
RETURN_NOT_OK(subarray_from_capnp(subarray_reader, &subarray));
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query->stats(), dummy_logger);
RETURN_NOT_OK(query->set_subarray_unsafe(subarray));
}
}
Expand Down Expand Up @@ -3203,9 +3226,9 @@ void ordered_dim_label_reader_from_capnp(
throw_if_not_ok(layout_enum(reader_reader.getLayout(), &layout));

// Subarray
Subarray subarray(array, layout, query->stats(), dummy_logger, false);
auto subarray_reader = reader_reader.getSubarray();
throw_if_not_ok(subarray_from_capnp(subarray_reader, &subarray));
Subarray subarray = subarray_from_capnp(
subarray_reader, array, layout, query->stats(), dummy_logger);
throw_if_not_ok(query->set_subarray_unsafe(subarray));

// OrderedDimLabelReader requires an initialized subarray for construction.
Expand Down
8 changes: 6 additions & 2 deletions tiledb/sm/serialization/query.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -251,8 +251,12 @@ Status subarray_to_capnp(
const Subarray* subarray,
capnp::Subarray::Builder* builder);

Status subarray_from_capnp(
const capnp::Subarray::Reader& reader, Subarray* subarray);
Subarray subarray_from_capnp(
const capnp::Subarray::Reader& reader,
const Array* array,
Layout layout,
stats::Stats* parent_stats,
shared_ptr<Logger> logger);

void ordered_dim_label_reader_to_capnp(
const Query& query,
Expand Down
26 changes: 26 additions & 0 deletions tiledb/sm/subarray/subarray.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,32 @@ Subarray::Subarray(
add_default_ranges();
}

Subarray::Subarray(
const Array* array,
Layout layout,
stats::Stats* stats,
shared_ptr<Logger> logger,
std::vector<RangeSetAndSuperset> range_subset,
std::vector<bool> is_default,
std::vector<optional<Subarray::LabelRangeSubset>> label_range_subset,
std::unordered_map<std::string, std::vector<Range>> attr_range_subset,
std::vector<unsigned int> relevant_fragments,
bool coalesce_ranges)
: stats_(stats)
, logger_(std::move(logger))
, array_(array->opened_array())
, layout_(layout)
, cell_order_(array_->array_schema_latest().cell_order())
, range_subset_(std::move(range_subset))
, label_range_subset_(std::move(label_range_subset))
, attr_range_subset_(std::move(attr_range_subset))
, is_default_(std::move(is_default))
, est_result_size_computed_(false)
, relevant_fragments_(relevant_fragments)
, coalesce_ranges_(coalesce_ranges)
, ranges_sorted_(false) {
}

Subarray::Subarray(const Subarray& subarray)
: Subarray() {
// Make a deep-copy clone
Expand Down
Loading

0 comments on commit 7e6d359

Please sign in to comment.