Skip to content

Commit

Permalink
Fix setting of global offset index
Browse files Browse the repository at this point in the history
  • Loading branch information
@sjperkins
sjperkins committed Feb 5, 2024
1 parent 84a77ee commit fe3aa4f
Show file tree
Hide file tree
Showing 7 changed files with 268 additions and 138 deletions.
225 changes: 120 additions & 105 deletions cpp/arcae/base_column_map.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,23 +96,35 @@ struct MapIterator {
MapIterator(const RangeIterator<ColumnMapping> & rit,
const ColumnMapping & map,
std::vector<std::size_t> chunk_index,
std::vector<std::size_t> global_index,
std::vector<std::size_t> strides,
bool done);

static MapIterator Make(const RangeIterator<ColumnMapping> & rit, bool done);
std::size_t nDim() const {
return chunk_index_.size();
};
std::size_t nDim() const { return chunk_index_.size(); };
std::size_t RowDim() const { return nDim() - 1; };

std::size_t RowDim() const {
return nDim() - 1;
std::size_t ChunkOffset() const;
std::size_t GlobalOffset() const {
return map_.get().FlatOffset(global_index_);
};

std::size_t ChunkOffset() const;
std::size_t GlobalOffset() const;
std::size_t RangeSize(std::size_t dim) const;
std::size_t MemStart(std::size_t dim) const;
std::size_t RangeSize(std::size_t dim) const {
return rit_.get().range_length_[dim];
}

const ColumnMap & DimMap(std::size_t dim) const {
return rit_.get().DimMap(dim);
}

std::size_t MemStart(std::size_t dim) const {
return rit_.get().mem_start_[dim];
}

std::size_t MemOffset(std::size_t dim, std::size_t offset=0) const;

const Range & DimRange(std::size_t dim) const {
return rit_.get().DimRange(dim);
}

MapIterator & operator++();
bool operator==(const MapIterator & other) const;
Expand All @@ -122,6 +134,98 @@ struct MapIterator {
};


template <typename ColumnMapping>
MapIterator<ColumnMapping>::MapIterator(const RangeIterator<ColumnMapping> & rit,
const ColumnMapping & map,
std::vector<std::size_t> chunk_index,
std::vector<std::size_t> strides,
bool done) :
rit_(std::cref(rit)),
map_(std::cref(map)),
chunk_index_(std::move(chunk_index)),
global_index_(chunk_index_.size(), 0),
strides_(std::move(strides)),
done_(done) {

for(std::size_t dim=0; dim < chunk_index_.size(); ++dim) {
global_index_[dim] = MemOffset(dim, chunk_index_[dim]);
}
}

template <typename ColumnMapping>
MapIterator<ColumnMapping>
MapIterator<ColumnMapping>::Make(const RangeIterator<ColumnMapping> & rit, bool done) {
auto chunk_index = decltype(MapIterator::chunk_index_)(rit.nDim(), 0);
auto strides = decltype(MapIterator::strides_)(rit.nDim(), 1);
using ItType = std::tuple<std::size_t, std::size_t>;

for(auto [dim, product]=ItType{1, 1}; dim < rit.nDim(); ++dim) {
product = strides[dim] = product * rit.range_length_[dim - 1];
}

return MapIterator{std::cref(rit), std::cref(rit.map_.get()),
std::move(chunk_index),
std::move(strides),
done};
}

template <typename ColumnMapping>
std::size_t
MapIterator<ColumnMapping>::ChunkOffset() const {
std::size_t offset = 0;
for(auto dim = std::size_t{0}; dim < nDim(); ++dim) {
offset += chunk_index_[dim] * strides_[dim];
}
return offset;
}

template <typename ColumnMapping>
std::size_t
MapIterator<ColumnMapping>::MemOffset(std::size_t dim, std::size_t offset) const {
auto base_offset = MemStart(dim);

if(DimRange(dim).type == Range::MAP) {
const auto & dim_map = DimMap(dim);
assert(offset < dim_map.size());
return base_offset + dim_map[offset].mem;
}

return base_offset + offset;
}

template <typename ColumnMapping>
MapIterator<ColumnMapping> & MapIterator<ColumnMapping>::operator++() {
assert(!done_);

// Iterate from fastest to slowest changing dimension
for(auto dim = std::size_t{0}; dim < nDim();) {
chunk_index_[dim]++;

// We've achieved a successful iteration in this dimension
if(chunk_index_[dim] < RangeSize(dim)) {
global_index_[dim] = MemOffset(dim, chunk_index_[dim]);
break;
// Reset to zero and retry in the next dimension
} else if(dim < RowDim()) {
chunk_index_[dim] = 0;
global_index_[dim] = MemOffset(dim, chunk_index_[dim]);
++dim;
// This was the slowest changing dimension so we're done
} else {
done_ = true;
break;
}
}

return *this;
}

template <typename ColumnMapping>
bool MapIterator<ColumnMapping>::operator==(const MapIterator<ColumnMapping> & other) const {
if(&rit_.get() != &other.rit_.get() || done_ != other.done_) return false;
return done_ ? true : chunk_index_ == other.chunk_index_;
}

// Iterates over the Disjoint Ranges defined by a ColumnMapping
template <typename ColumnMapping>
struct RangeIterator {
Expand Down Expand Up @@ -156,9 +260,9 @@ struct RangeIterator {
};

// Return the Maps for the given dimension
const ColumnMap & DimMaps(std::size_t dim) const {
const ColumnMap & DimMap(std::size_t dim) const {
assert(dim < nDim());
return map_.get().DimMaps(dim);
return map_.get().DimMap(dim);
};

// Return the currently selected Range of the given dimension
Expand Down Expand Up @@ -194,95 +298,6 @@ struct RangeIterator {
};


template <typename ColumnMapping>
MapIterator<ColumnMapping>::MapIterator(const RangeIterator<ColumnMapping> & rit,
const ColumnMapping & map,
std::vector<std::size_t> chunk_index,
std::vector<std::size_t> global_index,
std::vector<std::size_t> strides,
bool done) :
rit_(std::cref(rit)),
map_(std::cref(map)),
chunk_index_(std::move(chunk_index)),
global_index_(std::move(global_index)),
strides_(std::move(strides)),
done_(done) {}

template <typename ColumnMapping>
MapIterator<ColumnMapping>
MapIterator<ColumnMapping>::Make(const RangeIterator<ColumnMapping> & rit, bool done) {
auto chunk_index = decltype(MapIterator::chunk_index_)(rit.nDim(), 0);
auto global_index = decltype(MapIterator::global_index_)(rit.mem_start_);
auto strides = decltype(MapIterator::strides_)(rit.nDim(), 1);
using ItType = std::tuple<std::size_t, std::size_t>;

for(auto [dim, product]=ItType{1, 1}; dim < rit.nDim(); ++dim) {
product = strides[dim] = product * rit.range_length_[dim - 1];
}

return MapIterator{std::cref(rit), std::cref(rit.map_.get()),
std::move(chunk_index), std::move(global_index),
std::move(strides), done};
}

template <typename ColumnMapping>
std::size_t
MapIterator<ColumnMapping>::ChunkOffset() const {
std::size_t offset = 0;
for(auto dim = std::size_t{0}; dim < nDim(); ++dim) {
offset += chunk_index_[dim] * strides_[dim];
}
return offset;
}

template <typename ColumnMapping>
std::size_t
MapIterator<ColumnMapping>::GlobalOffset() const {
return map_.get().FlatOffset(global_index_);
}

template <typename ColumnMapping>
std::size_t
MapIterator<ColumnMapping>::RangeSize(std::size_t dim) const {
return rit_.get().range_length_[dim];
}

template <typename ColumnMapping>
std::size_t MapIterator<ColumnMapping>::MemStart(std::size_t dim) const {
return rit_.get().mem_start_[dim];
}


template <typename ColumnMapping>
MapIterator<ColumnMapping> & MapIterator<ColumnMapping>::operator++() {
assert(!done_);

// Iterate from fastest to slowest changing dimension
for(auto dim = std::size_t{0}; dim < nDim();) {
chunk_index_[dim]++;
global_index_[dim]++;
// We've achieved a successful iteration in this dimension
if(chunk_index_[dim] < RangeSize(dim)) { break; }
// Reset to zero and retry in the next dimension
else if(dim < RowDim()) {
chunk_index_[dim] = 0;
global_index_[dim] = MemStart(dim);
++dim;
}
// This was the slowest changing dimension so we're done
else { done_ = true; break; }
}

return *this;
}

template <typename ColumnMapping>
bool MapIterator<ColumnMapping>::operator==(const MapIterator<ColumnMapping> & other) const {
if(&rit_.get() != &other.rit_.get() || done_ != other.done_) return false;
return done_ ? true : chunk_index_ == other.chunk_index_;
}


template <typename ColumnMapping>
RangeIterator<ColumnMapping>::RangeIterator(ColumnMapping & column_map, bool done) :
map_(std::cref(column_map)),
Expand Down Expand Up @@ -334,7 +349,7 @@ void RangeIterator<ColumnMapping>::UpdateState() {
break;
}
case Range::MAP: {
const auto & dim_maps = DimMaps(dim);
const auto & dim_maps = DimMap(dim);
assert(range.start < dim_maps.size());
assert(range.end - 1 < dim_maps.size());
auto start = disk_start_[dim] = dim_maps[range.start].disk;
Expand Down Expand Up @@ -436,8 +451,8 @@ struct BaseColumnMap {
}


// Return the ColumnRange for the given dimension
const ColumnMap & DimMaps(std::size_t dim) const {
// Return the Dimension Map for the given dimension
const ColumnMap & DimMap(std::size_t dim) const {
assert(dim < nDim());
return maps_[dim];
}
Expand Down Expand Up @@ -482,7 +497,7 @@ struct BaseColumnMap {
template <typename T>
bool BaseColumnMap<T>::IsSimple() const {
for(std::size_t dim=0; dim < nDim(); ++dim) {
const auto & column_map = DimMaps(dim);
const auto & column_map = DimMap(dim);
const auto & column_range = DimRanges(dim);

// More than one range of row ids in a dimension
Expand Down
6 changes: 4 additions & 2 deletions cpp/tests/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
find_package(GTest CONFIG REQUIRED)

add_library(test_utils test_utils.cc)
add_compile_options("-fsanitize=address")
add_link_options("-fsanitize=address")
# add_compile_options("-fsanitize=address")
# add_link_options("-fsanitize=address")
add_compile_options("-g")
add_link_options("-g")

add_executable(basic_test basic_test.cc)
target_link_libraries(basic_test PRIVATE GTest::gtest_main arcae test_utils)
Expand Down
76 changes: 74 additions & 2 deletions cpp/tests/column_rw_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -708,7 +708,7 @@ TEST_F(ColumnWriteTest, WriteVisitorVariableNumeric) {
R"([[[0, 1],
[2, 3]],
[[4, 5],
[7, 8]]])"));
[6, 7]]])"));

auto var = GetArrayColumn<CT>(table, "VAR_2X3");
var.putColumnCells(casacore::RefRows(0, 0),
Expand All @@ -732,8 +732,23 @@ TEST_F(ColumnWriteTest, WriteVisitorVariableNumeric) {
CT{2}, CT{3}));
EXPECT_THAT(var.getColumnCells(casacore::RefRows(1, 1)),
::testing::ElementsAre(CT{4}, CT{5}, CT{0},
CT{7}, CT{8}, CT{0},
CT{6}, CT{7}, CT{0},
CT{0}, CT{0}, CT{0}));


// Variable data column, two channels, two correlations
ASSERT_OK_AND_ASSIGN(auto expected,
ArrayFromJSON(dtype,
R"([[[3, 2],
[1, 0]],
[[7, 6],
[5, 4]]])"));

ASSERT_OK_AND_ASSIGN(read_map, (ColumnReadMap::Make(var, {{}, {1, 0}, {1, 0}})));
read_visitor = ColumnReadVisitor(read_map);
visit_status = read_visitor.Visit(var.columnDesc().dataType());
ASSERT_OK(visit_status);
ASSERT_TRUE(expected->Equals(read_visitor.array_)) << *read_visitor.array_;
}


Expand Down Expand Up @@ -773,7 +788,64 @@ TEST_F(ColumnWriteTest, WriteVisitorVariableNumeric) {
CT{7}, CT{8}, CT{0},
CT{0}, CT{0}, CT{0}));
}
}

TEST_F(ColumnWriteTest, WriteVisitorVariableOutOfOrder) {
const auto & table = table_proxy_.table();
using CT = casacore::Int;

{
auto dtype = arrow::fixed_size_list(
arrow::fixed_size_list(
arrow::int32(), 2), 2);


// Variable data column, two channels, two correlations
ASSERT_OK_AND_ASSIGN(auto data,
ArrayFromJSON(dtype,
R"([[[0, 1],
[2, 3]],
[[4, 5],
[6, 7]]])"));

auto var = GetArrayColumn<CT>(table, "VAR_2X3");
var.putColumnCells(casacore::RefRows(0, 0), casacore::Array<CT>(IPos({2, 2, 1}), 0));
var.putColumnCells(casacore::RefRows(1, 1), casacore::Array<CT>(IPos({3, 3, 1}), 0));
ASSERT_OK_AND_ASSIGN(auto write_map, (ColumnWriteMap::Make(var, {{}, {0, 1}, {0, 1}}, data)));
auto write_visitor = ColumnWriteVisitor(write_map);
auto visit_status = write_visitor.Visit();
ASSERT_OK(visit_status);

ASSERT_OK_AND_ASSIGN(auto read_map, (ColumnReadMap::Make(var, {{}, {0, 1}, {0, 1}})));
auto read_visitor = ColumnReadVisitor(read_map);
visit_status = read_visitor.Visit(var.columnDesc().dataType());
ASSERT_OK(visit_status);
ASSERT_TRUE(data->Equals(read_visitor.array_));

// Sanity check values via casacore
EXPECT_THAT(var.getColumnCells(casacore::RefRows(0, 0)),
::testing::ElementsAre(CT{0}, CT{1},
CT{2}, CT{3}));
EXPECT_THAT(var.getColumnCells(casacore::RefRows(1, 1)),
::testing::ElementsAre(CT{4}, CT{5}, CT{0},
CT{6}, CT{7}, CT{0},
CT{0}, CT{0}, CT{0}));


// Variable data column, two channels, two correlations
ASSERT_OK_AND_ASSIGN(auto expected,
ArrayFromJSON(dtype,
R"([[[3, 2],
[1, 0]],
[[7, 6],
[5, 4]]])"));

ASSERT_OK_AND_ASSIGN(read_map, (ColumnReadMap::Make(var, {{}, {1, 0}, {1, 0}})));
read_visitor = ColumnReadVisitor(read_map);
visit_status = read_visitor.Visit(var.columnDesc().dataType());
ASSERT_OK(visit_status);
ASSERT_TRUE(expected->Equals(read_visitor.array_)) << *read_visitor.array_;
}
}

TEST_F(ColumnWriteTest, WriteVisitorVariableString) {
Expand Down
Loading

0 comments on commit fe3aa4f

Please sign in to comment.