-
Notifications
You must be signed in to change notification settings - Fork 4.2k
GH-50251: Add GetSpan() convenience method to ArrayData #50310
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.
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
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -101,6 +101,71 @@ void CheckDictionaryNullCount(const std::shared_ptr<DataType>& dict_type, | |
| ASSERT_EQ(arr->data()->MayHaveLogicalNulls(), expected_may_have_logical_nulls); | ||
| } | ||
|
|
||
| // Tests for ArrayData::GetSpan | ||
|
|
||
| TEST(TestArrayData, GetSpanWithOffset) { | ||
| using T = int32_t; | ||
|
|
||
| // Buffer with values 0..9 | ||
| std::vector<T> values = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; | ||
| auto value_buffer = Buffer::FromVector(values).ValueOrDie(); | ||
|
|
||
| // ArrayData with offset=3, length=4 → logically represents [3,4,5,6] | ||
| int64_t offset = 3; | ||
| int64_t length = 4; | ||
| auto data = ArrayData::Make(int32(), length, {nullptr /* null bitmap */, value_buffer}, | ||
| /*null_count=*/0, offset); | ||
|
|
||
| // Buffer index 1 is the values buffer | ||
| auto span = data->GetSpan<T>(/*i=*/1, /*length=*/length); | ||
|
|
||
| // Verify the offset is respected | ||
| ASSERT_EQ(span.size(), length); | ||
| ASSERT_EQ(span[0], 3); // values[offset] | ||
| ASSERT_EQ(span[length - 1], 6); // values[offset + length - 1] | ||
| ASSERT_EQ(span.data(), value_buffer->data_as<T>() + offset); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you also exercise calling |
||
| } | ||
|
|
||
| TEST(TestArrayData, GetSpanWithNullBufferCases) { | ||
| using T = int64_t; | ||
| std::vector<T> values = {10, 20, 30, 40}; | ||
| auto value_buffer = Buffer::FromVector(values).ValueOrDie(); | ||
| int64_t length = 4; | ||
| int64_t offset = 0; | ||
|
|
||
| // --- Case A: Values buffer is nullptr → expect empty span --- | ||
| { | ||
| auto data = ArrayData::Make(int64(), length, {nullptr, nullptr}, | ||
| /*null_count=*/0, offset); | ||
| auto span = data->GetSpan<T>(/*i=*/1, length); | ||
| ASSERT_TRUE(span.empty()); | ||
| ASSERT_EQ(span.data(), nullptr); | ||
|
Comment on lines
+138
to
+142
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is invalid. All null buffers in an integer array is only supported for length=0, can you fix the test accordingly? |
||
| } | ||
|
|
||
| // --- Case B: Null bitmap is nullptr, values buffer valid → span works --- | ||
| { | ||
| auto data = ArrayData::Make(int64(), length, {nullptr, value_buffer}, | ||
| /*null_count=*/0, offset); | ||
| auto span = data->GetSpan<T>(/*i=*/1, length); | ||
| ASSERT_EQ(span.size(), length); | ||
| ASSERT_EQ(span[0], 10); | ||
| ASSERT_EQ(span[3], 40); | ||
| } | ||
|
Comment on lines
+145
to
+153
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You're already testing this in |
||
|
|
||
| // --- Case C: Null bitmap exists and valid, values buffer valid → span works --- | ||
| { | ||
| // All bits set = no nulls | ||
| auto null_buffer = *Buffer::FromString(std::string("\xFF\xFF\xFF\xFF", 4)); | ||
| auto data = ArrayData::Make(int64(), length, {null_buffer, value_buffer}, | ||
| /*null_count=*/0, offset); | ||
| auto span = data->GetSpan<T>(/*i=*/1, length); | ||
| // The span ignores the null bitmap entirely – it just looks at buffer index 1 | ||
| ASSERT_EQ(span.size(), length); | ||
| ASSERT_EQ(span[0], 10); | ||
| ASSERT_EQ(span[3], 40); | ||
| } | ||
|
VedantRalekar marked this conversation as resolved.
|
||
| } | ||
|
|
||
| TEST_F(TestArray, TestNullCount) { | ||
| // These are placeholders | ||
| auto data = std::make_shared<Buffer>(nullptr, 0); | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -271,6 +271,28 @@ struct ARROW_EXPORT ArrayData { | |
| return GetValues<T>(i, offset); | ||
| } | ||
|
|
||
| /// \brief Access a buffer's data as a span | ||
| /// | ||
| /// \param i The buffer index | ||
| /// \param length The required length (in number of typed values) of the requested span | ||
| /// \pre i > 0 | ||
| /// \pre length <= the length of the buffer (in number of values) that's expected for | ||
| /// this array type | ||
| /// \return A span<const T> of the requested length, or empty if buffers[i] is null | ||
| template <typename T> | ||
| std::span<const T> GetSpan(int i, int64_t length) const { | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We have |
||
| if (!buffers[i]) { | ||
| return {}; | ||
| } | ||
|
|
||
| const int64_t buffer_length = buffers[i]->size() / static_cast<int64_t>(sizeof(T)); | ||
| assert(i > 0 && length + offset <= buffer_length); | ||
| ARROW_UNUSED(buffer_length); | ||
|
|
||
| return std::span<const T>(reinterpret_cast<const T*>(buffers[i]->data()) + offset, | ||
| length); | ||
| } | ||
|
|
||
| /// \brief Access a buffer's data as a typed C pointer | ||
| /// | ||
| /// \param i the buffer index | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can use
ASSERT_OK_AND_ASSIGNinstead of dying on error