apacheGH-50251: [C++] Add GetSpan to ArrayData (apache#50366)

dkp116 · pitrou · web-flow · commit 5f6ee02985d2 · 2026-07-07T17:26:55.000+02:00
### What changes are included in this PR? Add `GetSpan<T>(int i, int64_t length)` for read-only typed access to ArrayData. Add `GetMutableSpan<T>(int i, int64_t length)` for writable typed access to ArrayData. ### Are these changes tested? Yes, new unit tests were added in `array/data_test.cc` * GitHub Issue: apache#50251 Lead-authored-by: dkp116 <daneel668@gmail.com> Co-authored-by: Antoine Pitrou <antoine@python.org> Signed-off-by: Antoine Pitrou <antoine@python.org>
diff --git a/cpp/src/arrow/array/data.h b/cpp/src/arrow/array/data.h
@@ -256,6 +256,45 @@ 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<const T> of the requested length
+  template <typename T>
+  std::span<const T> GetSpan(int i, int64_t length) const {
+    if (!buffers[i]) {
+      assert(length == 0);
+      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>(buffers[i]->data_as<T>() + 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<T> of the requested length
+  template <typename T>
+  std::span<T> GetMutableSpan(int i, int64_t length) {
+    if (!buffers[i]) {
+      assert(length == 0);
+      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<T>(buffers[i]->mutable_data_as<T>() + this->offset, length);
+  }
 
   /// \brief Access a buffer's data as a typed C pointer
   ///
diff --git a/cpp/src/arrow/array/data_test.cc b/cpp/src/arrow/array/data_test.cc
@@ -15,9 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include <type_traits>
+
 #include <gtest/gtest.h>
 
-#include "arrow/array.h"
+#include "arrow/array/array_base.h"
 #include "arrow/array/data.h"
 #include "arrow/testing/gtest_util.h"
 
@@ -43,4 +45,65 @@ TEST(ArraySpan, SetSlice) {
   ASSERT_EQ(span.GetNullCount(), 1);
 }
 
+TEST(ArrayData, GetSpanRespectsOffset) {
+  std::vector<uint16_t> values = {1, 2, 3, 4, 5};
+
+  auto buffer = Buffer::FromVector(std::move(values));
+  auto data = ArrayData::Make(uint16(), /*length=*/3, {nullptr, buffer}, /*null_count=*/0,
+                              /*offset=*/1);
+  auto span = data->GetSpan<uint16_t>(1, 3);
+
+  const bool is_const_pointer = std::is_same_v<decltype(span)::pointer, const uint16_t*>;
+  ASSERT_TRUE(is_const_pointer);
+
+  EXPECT_EQ(span.size(), 3);
+  EXPECT_EQ(span[0], 2);
+  EXPECT_EQ(span[1], 3);
+  EXPECT_EQ(span[2], 4);
+}
+
+TEST(ArrayData, GetMutableSpanRespectsOffset) {
+  std::vector<uint16_t> values = {10, 20, 30, 40, 50};
+
+  auto buffer = std::make_shared<MutableBuffer>(reinterpret_cast<uint8_t*>(values.data()),
+                                                values.size() * sizeof(uint16_t));
+  std::vector<std::shared_ptr<Buffer>> buffers = {nullptr, buffer};
+
+  auto data =
+      ArrayData::Make(uint16(), /*length=*/3, buffers, /*null_count=*/0, /*offset=*/1);
+  auto span = data->GetMutableSpan<uint16_t>(1, 3);
+
+  const bool is_mut_pointer = std::is_same_v<decltype(span)::pointer, uint16_t*>;
+  ASSERT_TRUE(is_mut_pointer);
+
+  EXPECT_EQ(span.size(), 3);
+  EXPECT_EQ(span[0], 20);
+  EXPECT_EQ(span[1], 30);
+  EXPECT_EQ(span[2], 40);
+
+  span[0] = 200;
+  span[1] = 300;
+  span[2] = 400;
+
+  auto raw = reinterpret_cast<uint16_t*>(buffer->mutable_data());
+
+  EXPECT_EQ(raw[1], 200);
+  EXPECT_EQ(raw[2], 300);
+  EXPECT_EQ(raw[3], 400);
+}
+
+TEST(ArrayData, GetSpanNullBuffer) {
+  auto data = ArrayData::Make(uint16(), /*length=*/0, /*buffers=*/{nullptr, nullptr},
+                              /*null_count=*/0, /*offset=*/0);
+  auto span = data->GetSpan<uint16_t>(1, 0);
+  EXPECT_EQ(span.size(), 0);
+}
+
+TEST(ArrayData, GetMutableSpanNullBuffer) {
+  auto data = ArrayData::Make(uint16(), /*length=*/0, /*buffers=*/{nullptr, nullptr},
+                              /*null_count=*/0, /*offset=*/0);
+  auto span = data->GetMutableSpan<uint16_t>(1, 0);
+  EXPECT_EQ(span.size(), 0);
+}
+
 }  // namespace arrow
