diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc index dcfe1c76c301..f4c657b91c4b 100644 --- a/cpp/src/arrow/array/array_test.cc +++ b/cpp/src/arrow/array/array_test.cc @@ -1137,7 +1137,39 @@ TEST_F(TestArray, TestBinaryViewAppendArraySlice) { AssertArraysEqual(*src, *dst); } - +TEST_F(TestArray, GetSpanRespectsOffset) { + auto data_buffer = Buffer::FromString("123456789abcdef0"); + auto data = ArrayData::Make(uint8(), 3, {nullptr, data_buffer}, 0, 1); + auto span = data->GetSpan(1, 3); + + EXPECT_EQ(span.size(), 3); + EXPECT_EQ(span[0], '2'); + EXPECT_EQ(span[1], '3'); + EXPECT_EQ(span[2], '4'); +} +TEST_F(TestArray, GetMutableSpanRespectsOffset) { + const int64_t nbytes = 16; + ASSERT_OK_AND_ASSIGN(auto uniq_buffer, AllocateResizableBuffer(nbytes, pool_)); + memset(uniq_buffer->mutable_data(), '0', nbytes); + std::shared_ptr buffer(std::move(uniq_buffer)); + std::vector> buffers = {nullptr, buffer}; + auto data = ArrayData::Make(uint8(), 3, buffers, 0, 1); + auto span = data->template GetMutableSpan(1, 3); + + EXPECT_EQ(span.size(), 3); + EXPECT_EQ(span[0], '0'); + EXPECT_EQ(span[1], '0'); + EXPECT_EQ(span[2], '0'); + + span[0] = 'X'; + span[1] = 'Y'; + span[2] = 'Z'; + + auto raw = buffer->mutable_data(); + EXPECT_EQ(raw[1], 'X'); + EXPECT_EQ(raw[2], 'Y'); + EXPECT_EQ(raw[3], 'Z'); +} TEST_F(TestArray, ValidateBuffersPrimitive) { auto empty_buffer = std::make_shared(""); auto null_buffer = Buffer::FromString("\xff"); diff --git a/cpp/src/arrow/array/data.h b/cpp/src/arrow/array/data.h index 92308e8a010f..b10a94fc869c 100644 --- a/cpp/src/arrow/array/data.h +++ b/cpp/src/arrow/array/data.h @@ -256,6 +256,37 @@ struct ARROW_EXPORT ArrayData { return NULLPTR; } } + /// \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 of the requested length + template + std::span GetSpan(int i, int64_t length) const { + const int64_t buffer_length = buffers[i]->size() / static_cast(sizeof(T)); + assert(i > 0 && length + offset <= buffer_length); + ARROW_UNUSED(buffer_length); + return std::span(buffers[i]->data_as() + this->offset, length); + } + + /// \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 of the requested length + template + std::span GetMutableSpan(int i, int64_t length) { + const int64_t buffer_length = buffers[i]->size() / static_cast(sizeof(T)); + assert(i > 0 && length + offset <= buffer_length); + ARROW_UNUSED(buffer_length); + return std::span(buffers[i]->mutable_data_as() + this->offset, length); + } /// \brief Access a buffer's data as a typed C pointer ///