Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
65 changes: 65 additions & 0 deletions cpp/src/arrow/array/array_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Copy link
Copy Markdown
Member

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_ASSIGN instead of dying on error


// 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);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you also exercise calling GetSpan with a smaller length?

}

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The 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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're already testing this in GetSpanWithOffset, aren't you?


// --- 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);
}
Comment thread
VedantRalekar marked this conversation as resolved.
}

TEST_F(TestArray, TestNullCount) {
// These are placeholders
auto data = std::make_shared<Buffer>(nullptr, 0);
Expand Down
22 changes: 22 additions & 0 deletions cpp/src/arrow/array/data.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have ArrayData::GetMutableValues in addition to ArrayData::GetValues, why not add GetMutableSpan as well?

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
Expand Down