(arrow) branch main updated: GH-37199: [C++] Expose a span converter for Buffer and ArraySpan (#38027)

[felipecrv]

This is an automated email from the ASF dual-hosted git repository.

felipecrv pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new 0552217efa GH-37199: [C++] Expose a span converter for Buffer and 
ArraySpan (#38027)
0552217efa is described below

commit 0552217efa4ba0a1a1a7857a86c92278ecf129c8
Author: 谢天 <jsjtxiet...@outlook.com>
AuthorDate: Tue Dec 19 07:39:02 2023 +0800

    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 <jsjtxiet...@outlook.com>
    Signed-off-by: Felipe Oliveira Carvalho <felipe...@gmail.com>
---
 cpp/src/arrow/array/data.h | 31 +++++++++++++++++++++++++++++++
 cpp/src/arrow/buffer.h     | 13 +++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/cpp/src/arrow/array/data.h b/cpp/src/arrow/array/data.h
index 4c2df83814..f29f164d19 100644
--- 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
index ae76550be2..52fd94ec1f 100644
--- 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_); }