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

cpp/src/arrow/array/array_base.cc

@@ -323,6 +323,20 @@ Result<std::shared_ptr<Array>> Array::ViewOrCopyTo(
  return MakeArray(new_data);
 }
 
+Result<std::shared_ptr<Tensor>> Array::ToTensor(bool allow_nulls) const {
+ if (!allow_nulls && null_count() > 0) {
+ return Status::Invalid(
+ "Array contains nulls, explicitly pass `allow_nulls=true` to leave them "
+ "undefined.");
+ }
+
+ return ToTensorWithNulls();
+}
+
+Result<std::shared_ptr<Tensor>> Array::ToTensorWithNulls() const {
+ return Status::TypeError("ToTensor is not implemented for Array type ", type()->name());
+}
+
 // ----------------------------------------------------------------------
 // NullArray
 

cpp/src/arrow/array/array_base.h

@@ -29,6 +29,7 @@
 #include "arrow/result.h"
 #include "arrow/status.h"
 #include "arrow/type.h"
+#include "arrow/type_fwd.h"
 #include "arrow/util/bit_util.h"
 #include "arrow/util/macros.h"
 #include "arrow/util/visibility.h"
@@ -246,6 +247,17 @@ 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.
+ ///
+ /// \param[in] allow_nulls When true, nulls are ignored, leaving the output tensor with
+ /// unspecified values where this array has null entries. When false, nulls
+ /// are rejected.
+ Result<std::shared_ptr<Tensor>> ToTensor(bool allow_nulls = false) const;
+
  protected:
  Array() = default;
  ARROW_DEFAULT_MOVE_AND_ASSIGN(Array);
@@ -263,6 +275,9 @@ class ARROW_EXPORT Array {
  data_ = data;
  }
 
+ /// Implementation of ToTensor with nulls as undefined.
+ virtual Result<std::shared_ptr<Tensor>> ToTensorWithNulls() const;
+
  private:
  ARROW_DISALLOW_COPY_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,114 @@ 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});
+
+ // Slice offsets at each level contribute to the leaf offset, scaled by the list
+ // sizes above them.
+ 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));
+ // Offset 2 for initial value slice
+ CheckToTensor<float>(outer, {2, 3, 2}, {2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13});
+ // Offset of 2 (initial values) + 1 * 3 * 2 (outer slice times parent dimensions) = 8
+ CheckToTensor<float>(outer->Slice(1), {1, 3, 2}, {8, 9, 10, 11, 12, 13});
+}
+
+TEST_F(TestFixedSizeListArray, ToTensorNulls) {
+ auto array = ArrayFromJSON(fixed_size_list(int32(), 2), "[[1, 2], null, [5, null]]");
+
+ // Default behaviour is to not allow nulls
+ ASSERT_RAISES(Invalid, array->ToTensor());
+
+ // Nulls are ignored, leaving unspecified values in the output tensor.
+ ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor(/* allow_nulls= */ true));
+ ASSERT_OK(tensor->Validate());
+ ASSERT_EQ(tensor->Value<Int32Type>({0, 0}), 1);
+ ASSERT_EQ(tensor->Value<Int32Type>({0, 1}), 2);
+ ASSERT_EQ(tensor->Value<Int32Type>({2, 0}), 5);
+ 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(
+ TypeError,
+ ArrayFromJSON(fixed_size_list(utf8(), 1), R"([["a"], ["b"]])")->ToTensor());
+ ASSERT_RAISES(
+ TypeError,
+ ArrayFromJSON(fixed_size_list(boolean(), 2), "[[true, false]]")->ToTensor());
+ ASSERT_RAISES(TypeError,
+ 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,44 @@ Result<std::shared_ptr<Array>> FixedSizeListArray::Flatten(
  return FlattenListArray(*this, memory_pool);
 }
 
+Result<std::shared_ptr<Tensor>> FixedSizeListArray::ToTensorWithNulls() 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();
+ // Overflow cannot happen on a valid array (its data needs to fit in memory,
+ // therefore be smaller than INT64_MAX)
+ offset = offset * fsl->list_size() + data->offset;
+ length = length * fsl->list_size();
+ shape.push_back(fsl->list_size());
+ }
+
+ // Only checking byte_width which we need here and leaving Tensor::Make error on
+ // unsupported types.
+ if (!is_fixed_width(*type)) {
+ return Status::TypeError("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();
+ // Buffer guarantees this fits into an int64_t.
+ const int64_t byte_offset = offset * byte_width;
+ const int64_t byte_length = length * byte_width;
+ ARROW_ASSIGN_OR_RAISE(buffer, SliceBufferSafe(buf, byte_offset, byte_length));
+ }
+
+ return Tensor::Make(std::move(type), std::move(buffer), std::move(shape));
+}
+
 // ----------------------------------------------------------------------
 // Struct
 

cpp/src/arrow/array/array_nested.h

@@ -649,6 +649,14 @@ class ARROW_EXPORT FixedSizeListArray : public Array {
  void SetData(const std::shared_ptr<ArrayData>& data);
  int32_t list_size_;
 
+ /// \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>> ToTensorWithNulls() const override;
+
  private:
  std::shared_ptr<Array> values_;
 };

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"
 
@@ -138,6 +141,21 @@ class NumericArray : public PrimitiveArray {
  : NULLPTR;
  }
 
+ /// \brief Return a one dimensional Tensor.
+ Result<std::shared_ptr<Tensor>> ToTensorWithNulls() const override {
+ // Could be non-templated
+ const int64_t byte_width = type()->byte_width();
+ std::shared_ptr<Buffer> buffer;
+ if (data_->buffers[1] != NULLPTR) {
+ // Array guarantees this will not overflow.
+ const int64_t byte_offset = data_->offset * byte_width;
+ const int64_t byte_length = length() * byte_width;
+ ARROW_ASSIGN_OR_RAISE(buffer,
+ SliceBufferSafe(data_->buffers[1], byte_offset, byte_length));
+ }
+ return Tensor::Make(type(), std::move(buffer), {length()});
+ }
+
  const value_type* values_;
 };
 

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,67 @@ 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, ZeroLength) {
+ Int64Builder builder;
+ ASSERT_OK_AND_ASSIGN(auto array, builder.Finish());
+
+ ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor());
+ ASSERT_OK(tensor->Validate());
+
+ EXPECT_EQ(int64(), tensor->type());
+ EXPECT_EQ(std::vector<int64_t>{0}, tensor->shape());
+ EXPECT_EQ(std::vector<int64_t>{sizeof(int64_t)}, tensor->strides());
+}
+
+TEST(TestPrimitiveArray, ToTensorNulls) {
+ // Nulls are ignored, leaving unspecified values in the output tensor.
+ auto array = ArrayFromJSON(int32(), "[1, null, 3]");
+
+ // Default behaviour is to not allow nulls
+ ASSERT_RAISES(Invalid, array->ToTensor());
+
+ // Nulls are ignored, leaving unspecified values in the output tensor.
+ ASSERT_OK_AND_ASSIGN(auto tensor, array->ToTensor(/* allow_nulls= */ true));
+ ASSERT_OK(tensor->Validate());
+ ASSERT_EQ(tensor->Value<Int32Type>({0}), 1);
+ ASSERT_EQ(tensor->Value<Int32Type>({2}), 3);
+ EXPECT_EQ(std::vector<int64_t>{3}, tensor->shape());
+}
+
+TEST(TestPrimitiveArray, ToTensorUnsupportedType) {
+ auto array = ArrayFromJSON(date32(), "[1, 2, 3]");
+ ASSERT_RAISES(TypeError, array->ToTensor());
+ ASSERT_RAISES(TypeError, ArrayFromJSON(utf8(), R"(["a"])")->ToTensor());
+}
+
 class TestBuilder : public ::testing::Test {
  protected:
  MemoryPool* pool_ = default_memory_pool();