GH-38868: [C++][Python] Add Array::ToTensor and fixed size list support

cpp/src/arrow/array/array_base.cc

@@ -323,6 +323,11 @@ Result<std::shared_ptr<Array>> Array::ViewOrCopyTo(
   return MakeArray(new_data);
 }
 
+Result<std::shared_ptr<Tensor>> Array::ToTensor() const {
+  return Status::NotImplemented("ToTensor is not implemented for Array type ",
+                                type()->name());
+}
+
 // ----------------------------------------------------------------------
 // NullArray
 

cpp/src/arrow/array/array_base.h

@@ -36,6 +36,8 @@
 
 namespace arrow {
 
+class Tensor;
+
 // ----------------------------------------------------------------------
 // User array accessor types
 
@@ -246,6 +248,15 @@ class ARROW_EXPORT Array {
   /// \return const std::shared_ptr<ArrayStatistics>&
   const std::shared_ptr<ArrayStatistics>& statistics() const { return data_->statistics; }
 
+  /// \brief Create a Tensor from this Array
+  ///
+  /// When the data can reasonably be understood as a multidimensional numeric Tensor,
+  /// return the data as such.
+  /// Examples include NumericArray, FixedShapeTensorArray, nested FixedSizeListArray.
+  /// Nulls are ignored, leaving the output tensor with unspecified values where this
+  /// array has null entries.
+  virtual Result<std::shared_ptr<Tensor>> ToTensor() const;
+
  protected:
   Array() = default;
   ARROW_DEFAULT_MOVE_AND_ASSIGN(Array);

cpp/src/arrow/array/array_list_test.cc

@@ -29,6 +29,7 @@
 #include "arrow/array/validate.h"
 #include "arrow/buffer.h"
 #include "arrow/status.h"
+#include "arrow/tensor.h"
 #include "arrow/testing/builder.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/type.h"
@@ -1821,4 +1822,105 @@ TEST_F(TestFixedSizeListArray, FlattenRecursively) {
                     *ArrayFromJSON(value_type_, "[0, 1, null, 3, 7, null, 2, 5]"));
 }
 
+namespace {
+
+/// The innermost values of the nested fixed size lists.
+std::shared_ptr<Array> LeafValues(std::shared_ptr<Array> array) {
+  while (array->type_id() == Type::FIXED_SIZE_LIST) {
+    const auto& fsl = checked_cast<const FixedSizeListArray&>(*array);
+    array =
+        fsl.values()->Slice(fsl.value_offset(0), array->length() * fsl.value_length());
+  }
+  return array;
+}
+
+template <typename T>
+void CheckToTensor(const std::shared_ptr<Array>& array, const std::vector<int64_t>& shape,
+                   std::initializer_list<T> values) {
+  const auto value_type = CTypeTraits<T>::type_singleton();
+  ASSERT_OK_AND_ASSIGN(
+      auto expected,
+      Tensor::Make(value_type, Buffer::Wrap(values.begin(), values.size()), shape));
+
+  ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+  ASSERT_OK(tensor->Validate());
+
+  AssertTypeEqual(*value_type, *tensor->type());
+  ASSERT_EQ(shape, tensor->shape());
+  ASSERT_TRUE(tensor->is_row_major());
+  ASSERT_TRUE(tensor->Equals(*expected));
+
+  // The tensor shares the values buffer, it does not copy
+  const auto leaf = LeafValues(array);
+  ASSERT_EQ(leaf->data()->buffers[1]->data() + leaf->offset() * sizeof(T),
+            tensor->data()->data());
+}
+
+}  // namespace
+
+TEST_F(TestFixedSizeListArray, ToTensor) {
+  auto array = ArrayFromJSON(fixed_size_list(int32(), 3),
+                             "[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]");
+  CheckToTensor<int32_t>(array, {4, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+
+  // Offset on the list array itself
+  CheckToTensor<int32_t>(array->Slice(2, 2), {2, 3}, {7, 8, 9, 10, 11, 12});
+
+  // Offset on the values array
+  auto values = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7, 8, 9]")->Slice(3);
+  ASSERT_OK_AND_ASSIGN(auto from_values, FixedSizeListArray::FromArrays(values, 3));
+  CheckToTensor<int32_t>(from_values, {2, 3}, {4, 5, 6, 7, 8, 9});
+
+  // Offsets on both the list array and its values
+  CheckToTensor<int32_t>(from_values->Slice(1), {1, 3}, {7, 8, 9});
+}
+
+TEST_F(TestFixedSizeListArray, ToTensorNested) {
+  auto array = ArrayFromJSON(fixed_size_list(fixed_size_list(float32(), 2), 3), R"([
+    [[1, 2], [3, 4], [5, 6]],
+    [[7, 8], [9, 10], [11, 12]],
+    [[13, 14], [15, 16], [17, 18]]
+  ])");
+  CheckToTensor<float>(array, {3, 3, 2},
+                       {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18});
+
+  CheckToTensor<float>(array->Slice(1, 2), {2, 3, 2},
+                       {7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18});
+
+  // Offsets accumulate at every level: the innermost values are offset by 2, the
+  // middle lists by 3 * 2 and the outer lists by 1 * 3 * 2.
+  auto values = ArrayFromJSON(float32(), "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]")
+                    ->Slice(2, 12);
+  ASSERT_OK_AND_ASSIGN(auto inner, FixedSizeListArray::FromArrays(values, 2));
+  ASSERT_OK_AND_ASSIGN(auto outer, FixedSizeListArray::FromArrays(inner, 3));
+  CheckToTensor<float>(outer, {2, 3, 2}, {2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13});
+  CheckToTensor<float>(outer->Slice(1), {1, 3, 2}, {8, 9, 10, 11, 12, 13});
+}
+
+TEST_F(TestFixedSizeListArray, ToTensorNulls) {
+  // Nulls are ignored, leaving unspecified values in the output tensor.
+  auto array = ArrayFromJSON(fixed_size_list(int32(), 2), "[[1, 2], null, [5, null]]");
+  ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+  ASSERT_OK(tensor->Validate());
+  ASSERT_EQ(std::vector<int64_t>({3, 2}), tensor->shape());
+}
+
+TEST_F(TestFixedSizeListArray, ToTensorZeroLength) {
+  auto array = ArrayFromJSON(fixed_size_list(int64(), 2), "[]");
+  ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+  ASSERT_OK(tensor->Validate());
+  ASSERT_EQ(std::vector<int64_t>({0, 2}), tensor->shape());
+}
+
+TEST_F(TestFixedSizeListArray, ToTensorUnsupportedType) {
+  ASSERT_RAISES(
+      NotImplemented,
+      ArrayFromJSON(fixed_size_list(utf8(), 1), R"([["a"], ["b"]])")->ToTensor());
+  ASSERT_RAISES(
+      Invalid,
+      ArrayFromJSON(fixed_size_list(boolean(), 2), "[[true, false]]")->ToTensor());
+  ASSERT_RAISES(Invalid,
+                ArrayFromJSON(fixed_size_list(date32(), 2), "[[1, 2]]")->ToTensor());
+}
+
 }  // namespace arrow

cpp/src/arrow/array/array_nested.cc

@@ -33,13 +33,15 @@
 #include "arrow/array/util.h"
 #include "arrow/buffer.h"
 #include "arrow/status.h"
+#include "arrow/tensor.h"
 #include "arrow/type.h"
 #include "arrow/type_fwd.h"
 #include "arrow/type_traits.h"
 #include "arrow/util/bit_util.h"
 #include "arrow/util/bitmap_generate.h"
 #include "arrow/util/bitmap_ops.h"
 #include "arrow/util/checked_cast.h"
+#include "arrow/util/int_util_overflow.h"
 #include "arrow/util/list_util.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/util/unreachable.h"
@@ -1001,6 +1003,48 @@ Result<std::shared_ptr<Array>> FixedSizeListArray::Flatten(
   return FlattenListArray(*this, memory_pool);
 }
 
+Result<std::shared_ptr<Tensor>> FixedSizeListArray::ToTensor() const {
+  const auto* data = this->data().get();
+  auto type = this->type();
+  int64_t offset = data->offset;
+  int64_t length = data->length;
+  std::vector<int64_t> shape{length};
+
+  // Iterate over nested fixed length container types.
+  // Each nested container increase the tensor dimension.
+  while (type->id() == Type::FIXED_SIZE_LIST) {
+    const auto* fsl = internal::checked_cast<const FixedSizeListType*>(type.get());
+    type = fsl->value_type();
+    data = data->child_data.front().get();
+
+    if (internal::MultiplyWithOverflow(offset, int64_t{fsl->list_size()}, &offset) ||
+        internal::AddWithOverflow(offset, data->offset, &offset) ||
+        internal::MultiplyWithOverflow(length, int64_t{fsl->list_size()}, &length)) {
+      return Status::Invalid("Flattened fixed size list does not fit in an int64");
+    }
+    shape.push_back(fsl->list_size());
+  }
+
+  // Only checking byte_width and leaving Tensor::Make error on unsupported types.
+  if (!is_fixed_width(*type)) {
+    return Status::NotImplemented("Expected a fixed width leaf type, got ", type->name());
+  }
+
+  std::shared_ptr<Buffer> buffer = nullptr;
+  if (const auto& buf = data->buffers[1]; buf != NULLPTR) {
+    const int64_t byte_width = type->byte_width();
+    int64_t boffset = 0;
+    int64_t blength = 0;
+    if (internal::MultiplyWithOverflow(offset, byte_width, &boffset) ||
+        internal::MultiplyWithOverflow(length, byte_width, &blength)) {
+      return Status::Invalid("Array byte size does not fit in an int64");
+    }
+    ARROW_ASSIGN_OR_RAISE(buffer, SliceBufferSafe(buf, boffset, blength));
+  }
+
+  return Tensor::Make(std::move(type), std::move(buffer), std::move(shape));
+}
+
 // ----------------------------------------------------------------------
 // Struct
 

cpp/src/arrow/array/array_nested.h

@@ -645,6 +645,14 @@ class ARROW_EXPORT FixedSizeListArray : public Array {
       std::shared_ptr<Buffer> null_bitmap = NULLPTR,
       int64_t null_count = kUnknownNullCount);
 
+  /// \brief Return a Tensor sharing the data.
+  ///
+  /// The output tensor has a row major layout with the number of elements as the first
+  /// dimension and the fixed size list as the remaining ones (possibly nested).
+  /// Nulls are ignored, leaving the output tensor with unspecified values where this
+  /// array has null entries.
+  Result<std::shared_ptr<Tensor>> ToTensor() const override;
+
  protected:
   void SetData(const std::shared_ptr<ArrayData>& data);
   int32_t list_size_;

cpp/src/arrow/array/array_primitive.h

@@ -25,11 +25,14 @@
 
 #include "arrow/array/array_base.h"
 #include "arrow/array/data.h"
+#include "arrow/buffer.h"
 #include "arrow/stl_iterator.h"
+#include "arrow/tensor.h"
 #include "arrow/type.h"
 #include "arrow/type_fwd.h"  // IWYU pragma: export
 #include "arrow/type_traits.h"
 #include "arrow/util/bit_util.h"
+#include "arrow/util/int_util_overflow.h"
 #include "arrow/util/macros.h"
 #include "arrow/util/visibility.h"
 
@@ -128,6 +131,23 @@ class NumericArray : public PrimitiveArray {
 
   IteratorType end() const { return IteratorType(*this, length()); }
 
+  /// \brief Return a one dimensional Tensor.
+  Result<std::shared_ptr<Tensor>> ToTensor() const override {
+    // Could be non-templated
+    const int64_t byte_width = type()->byte_width();
+    std::shared_ptr<Buffer> buffer;
+    if (data_->buffers[1] != NULLPTR) {
+      int64_t boffset = 0;
+      int64_t blength = 0;
+      if (internal::MultiplyWithOverflow(data_->offset, byte_width, &boffset) ||
+          internal::MultiplyWithOverflow(length(), byte_width, &blength)) {
+        return Status::Invalid("Array byte size does not fit in an int64");
+      }
+      ARROW_ASSIGN_OR_RAISE(buffer, SliceBufferSafe(data_->buffers[1], boffset, blength));
+    }
+    return Tensor::Make(type(), std::move(buffer), {length()});
+  }
+
  protected:
   NumericArray() : values_(NULLPTR) {}
 

cpp/src/arrow/array/array_test.cc

@@ -50,6 +50,7 @@
 #include "arrow/result.h"
 #include "arrow/scalar.h"
 #include "arrow/status.h"
+#include "arrow/tensor.h"
 #include "arrow/testing/builder.h"
 #include "arrow/testing/extension_type.h"
 #include "arrow/testing/gtest_compat.h"
@@ -1218,6 +1219,52 @@ TEST(TestPrimitiveArray, CtorNoValidityBitmap) {
   ASSERT_EQ(arr.data()->null_count, 0);
 }
 
+TEST(TestPrimitiveArray, ToTensor) {
+  const std::vector<int64_t> shape = {5};
+  const std::vector<int64_t> strides = {sizeof(int32_t)};
+
+  auto array = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5]");
+  ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+  ASSERT_OK(tensor->Validate());
+
+  EXPECT_EQ(int32(), tensor->type());
+  EXPECT_EQ(shape, tensor->shape());
+  EXPECT_EQ(strides, tensor->strides());
+  EXPECT_TRUE(tensor->is_contiguous());
+  EXPECT_TRUE(
+      TensorFromJSON(int32(), "[1, 2, 3, 4, 5]", shape, strides)->Equals(*tensor));
+}
+
+TEST(TestPrimitiveArray, ToTensorSliced) {
+  const std::vector<int64_t> shape = {3};
+  const std::vector<int64_t> strides = {sizeof(int64_t)};
+
+  auto array = ArrayFromJSON(int64(), "[1, 2, 3, 4, 5]")->Slice(2);
+  ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+  ASSERT_OK(tensor->Validate());
+
+  EXPECT_EQ(shape, tensor->shape());
+  EXPECT_TRUE(TensorFromJSON(int64(), "[3, 4, 5]", shape, strides)->Equals(*tensor));
+}
+
+TEST(TestPrimitiveArray, ToTensorNulls) {
+  // Nulls are ignored, leaving unspecified values in the output tensor.
+  const std::vector<int64_t> shape = {3};
+
+  auto array = ArrayFromJSON(int32(), "[1, null, 3]");
+  ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+  ASSERT_OK(tensor->Validate());
+
+  EXPECT_EQ(shape, tensor->shape());
+}
+
+TEST(TestPrimitiveArray, ToTensorUnsupportedType) {
+  auto array = ArrayFromJSON(date32(), "[1, 2, 3]");
+  ASSERT_RAISES(Invalid, array->ToTensor());
+
+  ASSERT_RAISES(NotImplemented, ArrayFromJSON(utf8(), R"(["a"])")->ToTensor());
+}
+
 class TestBuilder : public ::testing::Test {
  protected:
   MemoryPool* pool_ = default_memory_pool();