GH-37199: [C++] Expose a span converter for Buffer and ArraySpan (#38027

) ### Rationale for this change Convenience. We can have such a helper at the buffer and array data level. ### What changes are included in this PR? Add `Buffer::span_as`, `Buffer::mutuable_span_as` and `ArraySpan::GetSpan`. ### Are these changes tested? No, but I'm happy to add some test if needed. ### Are there any user-facing changes? Yes, new public functions. * Closes: #37199 Authored-by: jsjtxietian <[email protected]> Signed-off-by: Felipe Oliveira Carvalho <[email protected]>
apache · Dec 18, 2023 · 0552217 · 0552217
1 parent 01c461f
commit 0552217
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/cpp/src/arrow/array/data.h b/cpp/src/arrow/array/data.h
@@ -18,6 +18,7 @@
 #pragma once
 
 #include <atomic>  // IWYU pragma: export
+#include <cassert>
 #include <cstdint>
 #include <memory>
 #include <utility>
@@ -438,6 +439,36 @@ struct ARROW_EXPORT ArraySpan {
     return GetValues<T>(i, this->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
+  template <typename T>
+  util::span<const T> GetSpan(int i, int64_t length) const {
+    const int64_t buffer_length = buffers[i].size / static_cast<int64_t>(sizeof(T));
+    assert(i > 0 && length + offset <= buffer_length);
+    return util::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>
+  util::span<T> GetSpan(int i, int64_t length) {
+    const int64_t buffer_length = buffers[i].size / static_cast<int64_t>(sizeof(T));
+    assert(i > 0 && length + offset <= buffer_length);
+    return util::span<T>(buffers[i].mutable_data_as<T>() + this->offset, length);
+  }
+
   inline bool IsNull(int64_t i) const { return !IsValid(i); }
 
   inline bool IsValid(int64_t i) const {

diff --git a/cpp/src/arrow/buffer.h b/cpp/src/arrow/buffer.h
@@ -30,6 +30,7 @@
 #include "arrow/status.h"
 #include "arrow/type_fwd.h"
 #include "arrow/util/macros.h"
+#include "arrow/util/span.h"
 #include "arrow/util/visibility.h"
 
 namespace arrow {
@@ -233,6 +234,12 @@ class ARROW_EXPORT Buffer {
     return reinterpret_cast<const T*>(data());
   }
 
+  /// \brief Return the buffer's data as a span
+  template <typename T>
+  util::span<const T> span_as() const {
+    return util::span(data_as<T>(), static_cast<size_t>(size() / sizeof(T)));
+  }
+
   /// \brief Return a writable pointer to the buffer's data
   ///
   /// The buffer has to be a mutable CPU buffer (`is_cpu()` and `is_mutable()`
@@ -260,6 +267,12 @@ class ARROW_EXPORT Buffer {
     return reinterpret_cast<T*>(mutable_data());
   }
 
+  /// \brief Return the buffer's mutable data as a span
+  template <typename T>
+  util::span<T> mutable_span_as() const {
+    return util::span(mutable_data_as<T>(), static_cast<size_t>(size() / sizeof(T)));
+  }
+
   /// \brief Return the device address of the buffer's data
   uintptr_t address() const { return reinterpret_cast<uintptr_t>(data_); }