Skip to content

Commit

Permalink
fix(cpp): Fix offset handling in ViewArrayAs Range Helpers (#702)
Browse files Browse the repository at this point in the history
  • Loading branch information
@WillAyd
WillAyd authored Jan 8, 2025
1 parent 7844799 commit 7a80870
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 15 deletions.
2 changes: 1 addition & 1 deletion src/nanoarrow/common/utils_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -739,7 +739,7 @@ TEST(RandomAccessRangeTest, ConstructionAndPrinting) {
auto square = [](int64_t i) { return i * i; };

// the range is usable as a constant
const nanoarrow::internal::RandomAccessRange<decltype(square)> squares{square, 4};
const nanoarrow::internal::RandomAccessRange<decltype(square)> squares{square, 0, 4};

// gtest recognizes the range as a container and can print it
EXPECT_THAT(squares, testing::ElementsAre(0, 1, 4, 9));
Expand Down
23 changes: 9 additions & 14 deletions src/nanoarrow/hpp/view.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ class Maybe {
template <typename Get>
struct RandomAccessRange {
Get get;
int64_t offset;
int64_t size;

using value_type = decltype(std::declval<Get>()(0));
Expand All @@ -78,8 +79,8 @@ struct RandomAccessRange {
value_type operator*() const { return range->get(i); }
};

const_iterator begin() const { return {0, this}; }
const_iterator end() const { return {size, this}; }
const_iterator begin() const { return {offset, this}; }
const_iterator end() const { return {offset + size, this}; }
};

template <typename Next>
Expand Down Expand Up @@ -132,10 +133,8 @@ class ViewArrayAs {
struct Get {
const uint8_t* validity;
const void* values;
int64_t offset;

internal::Maybe<T> operator()(int64_t i) const {
i += offset;
if (validity == nullptr || ArrowBitGet(validity, i)) {
if (std::is_same<T, bool>::value) {
return ArrowBitGet(static_cast<const uint8_t*>(values), i);
Expand All @@ -155,8 +154,8 @@ class ViewArrayAs {
Get{
array_view->buffer_views[0].data.as_uint8,
array_view->buffer_views[1].data.data,
array_view->offset,
},
array_view->offset,
array_view->length,
} {}

Expand All @@ -165,8 +164,8 @@ class ViewArrayAs {
Get{
static_cast<const uint8_t*>(array->buffers[0]),
array->buffers[1],
/*offset=*/0,
},
array->offset,
array->length,
} {}

Expand All @@ -193,10 +192,8 @@ class ViewArrayAsBytes {
const uint8_t* validity;
const void* offsets;
const char* data;
int64_t offset;

internal::Maybe<ArrowStringView> operator()(int64_t i) const {
i += offset;
auto* offsets = static_cast<const OffsetType*>(this->offsets);
if (validity == nullptr || ArrowBitGet(validity, i)) {
return ArrowStringView{data + offsets[i], offsets[i + 1] - offsets[i]};
Expand All @@ -214,8 +211,8 @@ class ViewArrayAsBytes {
array_view->buffer_views[0].data.as_uint8,
array_view->buffer_views[1].data.data,
array_view->buffer_views[2].data.as_char,
array_view->offset,
},
array_view->offset,
array_view->length,
} {}

Expand All @@ -225,8 +222,8 @@ class ViewArrayAsBytes {
static_cast<const uint8_t*>(array->buffers[0]),
array->buffers[1],
static_cast<const char*>(array->buffers[2]),
/*offset=*/0,
},
array->offset,
array->length,
} {}

Expand All @@ -246,11 +243,9 @@ class ViewArrayAsFixedSizeBytes {
struct Get {
const uint8_t* validity;
const char* data;
int64_t offset;
int fixed_size;

internal::Maybe<ArrowStringView> operator()(int64_t i) const {
i += offset;
if (validity == nullptr || ArrowBitGet(validity, i)) {
return ArrowStringView{data + i * fixed_size, fixed_size};
}
Expand All @@ -266,9 +261,9 @@ class ViewArrayAsFixedSizeBytes {
Get{
array_view->buffer_views[0].data.as_uint8,
array_view->buffer_views[1].data.as_char,
array_view->offset,
fixed_size,
},
array_view->offset,
array_view->length,
} {}

Expand All @@ -277,9 +272,9 @@ class ViewArrayAsFixedSizeBytes {
Get{
static_cast<const uint8_t*>(array->buffers[0]),
static_cast<const char*>(array->buffers[1]),
/*offset=*/0,
fixed_size,
},
array->offset,
array->length,
} {}

Expand Down
88 changes: 88 additions & 0 deletions src/nanoarrow/hpp/view_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,3 +135,91 @@ TEST(NanoarrowHppTest, NanoarrowHppViewArrayStreamTest) {
EXPECT_EQ(stream_view.code(), ENOMEM);
EXPECT_STREQ(stream_view.error()->message, "foo bar");
}

TEST(NanoarrowHppTest, NanoarrowHppViewArrayOffsetTest) {
nanoarrow::UniqueSchema schema{};
ArrowSchemaInit(schema.get());
ASSERT_EQ(ArrowSchemaSetType(schema.get(), NANOARROW_TYPE_INT32), NANOARROW_OK);

nanoarrow::UniqueArray array{};
ASSERT_EQ(ArrowArrayInitFromSchema(array.get(), schema.get(), nullptr), NANOARROW_OK);
ASSERT_EQ(ArrowArrayStartAppending(array.get()), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendInt(array.get(), 0), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendInt(array.get(), 1), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendInt(array.get(), 2), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendInt(array.get(), 3), NANOARROW_OK);
ASSERT_EQ(ArrowArrayFinishBuildingDefault(array.get(), nullptr), NANOARROW_OK);
array->offset = 2;
array->length = 2;

EXPECT_THAT(nanoarrow::ViewArrayAs<int32_t>(array.get()), testing::ElementsAre(2, 3));

nanoarrow::UniqueArrayView array_view{};
ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), nullptr),
NANOARROW_OK);
ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), nullptr), NANOARROW_OK);

EXPECT_THAT(nanoarrow::ViewArrayAs<int32_t>(array_view.get()),
testing::ElementsAre(2, 3));
}

TEST(NanoarrowHppTest, NanoarrowHppViewArrayAsBytesOffsetTest) {
using namespace nanoarrow::literals;

nanoarrow::UniqueSchema schema{};
ArrowSchemaInit(schema.get());
ASSERT_EQ(ArrowSchemaSetType(schema.get(), NANOARROW_TYPE_STRING), NANOARROW_OK);

nanoarrow::UniqueArray array{};
ASSERT_EQ(ArrowArrayInitFromSchema(array.get(), schema.get(), nullptr), NANOARROW_OK);
ASSERT_EQ(ArrowArrayStartAppending(array.get()), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(array.get(), "foo"_asv), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(array.get(), "bar"_asv), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(array.get(), "baz"_asv), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(array.get(), "qux"_asv), NANOARROW_OK);
ASSERT_EQ(ArrowArrayFinishBuildingDefault(array.get(), nullptr), NANOARROW_OK);
array->offset = 2;
array->length = 2;

EXPECT_THAT(nanoarrow::ViewArrayAsBytes<32>(array.get()),
testing::ElementsAre("baz"_asv, "qux"_asv));

nanoarrow::UniqueArrayView array_view{};
ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), nullptr),
NANOARROW_OK);
ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), nullptr), NANOARROW_OK);
EXPECT_THAT(nanoarrow::ViewArrayAsBytes<32>(array_view.get()),
testing::ElementsAre("baz"_asv, "qux"_asv));
}

TEST(NanoarrowHppTest, NanoarrowHppViewArrayAsFixedSizeBytesOffsetTest) {
using namespace nanoarrow::literals;

constexpr int32_t FixedSize = 3;
nanoarrow::UniqueSchema schema{};
ArrowSchemaInit(schema.get());
ASSERT_EQ(ArrowSchemaSetTypeFixedSize(schema.get(), NANOARROW_TYPE_FIXED_SIZE_BINARY,
FixedSize),
NANOARROW_OK);

nanoarrow::UniqueArray array{};
ASSERT_EQ(ArrowArrayInitFromSchema(array.get(), schema.get(), nullptr), NANOARROW_OK);
ASSERT_EQ(ArrowArrayStartAppending(array.get()), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendBytes(array.get(), {{"foo"}, FixedSize}), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendBytes(array.get(), {{"bar"}, FixedSize}), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendBytes(array.get(), {{"baz"}, FixedSize}), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendBytes(array.get(), {{"qux"}, FixedSize}), NANOARROW_OK);
ASSERT_EQ(ArrowArrayFinishBuildingDefault(array.get(), nullptr), NANOARROW_OK);
array->offset = 2;
array->length = 2;

EXPECT_THAT(nanoarrow::ViewArrayAsFixedSizeBytes(array.get(), FixedSize),
testing::ElementsAre("baz"_asv, "qux"_asv));

nanoarrow::UniqueArrayView array_view{};
ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), nullptr),
NANOARROW_OK);
ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), nullptr), NANOARROW_OK);
EXPECT_THAT(nanoarrow::ViewArrayAsFixedSizeBytes(array_view.get(), FixedSize),
testing::ElementsAre("baz"_asv, "qux"_asv));
}

0 comments on commit 7a80870

Please sign in to comment.