This is an automated email from the ASF dual-hosted git repository.
pitrou 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 987e2312c3 GH-38868: [C++][Python] Add Array::ToTensor and fixed size
list support (#50929)
987e2312c3 is described below
commit 987e2312c3bab168bdfa1dd6570762b3c37f085b
Author: Antoine Prouvost <[email protected]>
AuthorDate: Thu Sep 3 12:09:07 2026 +0200
GH-38868: [C++][Python] Add Array::ToTensor and fixed size list support
(#50929)
### Rationale for this change
Enable multidimensional DLPack support for Array via `to_tensor`.
### What changes are included in this PR?
- Add virtual `Array::ToTensor`
- Add `NumericArray::ToTensor` for 1D arrays
- Add `FixedSizeListArray::ToTensor` for multidimensional arrays
- Add DLPack error suggestiong tensor convertion
- Add DLPack tests with `arr.to_tensor().__dlpack__()`
Note: Nulls are explicitly supported in `to_tensor` as unspecified data.
This was the current behaviour.
### Are these changes tested?
Yes
### Are there any user-facing changes?
New public Array function.
* GitHub Issue: #38868
Authored-by: AntoinePrv <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/array/array_base.cc | 14 ++++
cpp/src/arrow/array/array_base.h | 15 ++++
cpp/src/arrow/array/array_list_test.cc | 111 ++++++++++++++++++++++++++
cpp/src/arrow/array/array_nested.cc | 40 ++++++++++
cpp/src/arrow/array/array_nested.h | 8 ++
cpp/src/arrow/array/array_primitive.h | 18 +++++
cpp/src/arrow/array/array_test.cc | 62 ++++++++++++++
cpp/src/arrow/c/dlpack.cc | 95 ++++++++--------------
cpp/src/arrow/c/dlpack_test.cc | 39 ++++++---
cpp/src/arrow/extension/fixed_shape_tensor.cc | 2 +-
cpp/src/arrow/extension/fixed_shape_tensor.h | 6 +-
cpp/src/arrow/tensor.cc | 7 +-
cpp/src/arrow/tensor.h | 7 ++
cpp/src/arrow/tensor_test.cc | 2 +-
python/pyarrow/array.pxi | 60 +++++++-------
python/pyarrow/includes/libarrow.pxd | 6 +-
python/pyarrow/tests/test_dlpack.py | 62 ++++++++++++++
17 files changed, 441 insertions(+), 113 deletions(-)
diff --git a/cpp/src/arrow/array/array_base.cc
b/cpp/src/arrow/array/array_base.cc
index ce2e66655a..0add32a555 100644
--- a/cpp/src/arrow/array/array_base.cc
+++ b/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
diff --git a/cpp/src/arrow/array/array_base.h b/cpp/src/arrow/array/array_base.h
index 60df45357e..ed2e2356ea 100644
--- a/cpp/src/arrow/array/array_base.h
+++ b/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);
};
diff --git a/cpp/src/arrow/array/array_list_test.cc
b/cpp/src/arrow/array/array_list_test.cc
index 9908eb5886..4901a1e6fa 100644
--- a/cpp/src/arrow/array/array_list_test.cc
+++ b/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"
@@ -1858,4 +1859,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
diff --git a/cpp/src/arrow/array/array_nested.cc
b/cpp/src/arrow/array/array_nested.cc
index 55bce4d0de..6849aca089 100644
--- a/cpp/src/arrow/array/array_nested.cc
+++ b/cpp/src/arrow/array/array_nested.cc
@@ -33,6 +33,7 @@
#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"
@@ -40,6 +41,7 @@
#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
diff --git a/cpp/src/arrow/array/array_nested.h
b/cpp/src/arrow/array/array_nested.h
index bf84f802b1..57c87f900c 100644
--- a/cpp/src/arrow/array/array_nested.h
+++ b/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_;
};
diff --git a/cpp/src/arrow/array/array_primitive.h
b/cpp/src/arrow/array/array_primitive.h
index cebf47ad93..473ac16995 100644
--- a/cpp/src/arrow/array/array_primitive.h
+++ b/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_;
};
diff --git a/cpp/src/arrow/array/array_test.cc
b/cpp/src/arrow/array/array_test.cc
index dcfe1c76c3..668d9f00d8 100644
--- a/cpp/src/arrow/array/array_test.cc
+++ b/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();
diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc
index 6eb06c738c..4e25d50bb5 100644
--- a/cpp/src/arrow/c/dlpack.cc
+++ b/cpp/src/arrow/c/dlpack.cc
@@ -35,34 +35,24 @@ namespace arrow::dlpack {
namespace {
Result<DLDataType> GetDLDataType(const DataType& type) {
- DLDataType dtype;
+ auto dtype = DLDataType{};
dtype.lanes = 1;
dtype.bits = type.bit_width();
- switch (type.id()) {
- case Type::INT8:
- case Type::INT16:
- case Type::INT32:
- case Type::INT64:
- dtype.code = DLDataTypeCode::kDLInt;
- return dtype;
- case Type::UINT8:
- case Type::UINT16:
- case Type::UINT32:
- case Type::UINT64:
- dtype.code = DLDataTypeCode::kDLUInt;
- return dtype;
- case Type::HALF_FLOAT:
- case Type::FLOAT:
- case Type::DOUBLE:
- dtype.code = DLDataTypeCode::kDLFloat;
- return dtype;
- case Type::BOOL:
- // DLPack supports byte-packed boolean values
- return Status::TypeError("Bit-packed boolean data type not supported by
DLPack.");
- default:
- return Status::TypeError("DataType is not compatible with DLPack spec: ",
- type.ToString());
+ if (is_signed_integer(type.id())) {
+ dtype.code = DLDataTypeCode::kDLInt;
+ } else if (is_unsigned_integer(type.id())) {
+ dtype.code = DLDataTypeCode::kDLUInt;
+ } else if (is_floating(type.id())) {
+ dtype.code = DLDataTypeCode::kDLFloat;
+ } else if (type.id() == Type::BOOL) {
+ // DLPack supports byte-packed boolean values
+ return Status::TypeError("Bit-packed boolean data type not supported by
DLPack.");
+ } else {
+ return Status::TypeError(
+ "DataType is not compatible with DLPack spec: ", type.ToString(),
+ ", try converting to a Tensor for multi dimensional data support");
}
+ return {dtype};
}
template <typename DT, typename Vec>
@@ -78,6 +68,7 @@ struct ManagerCtx {
template <typename Vec>
struct ExportBufferParams {
std::shared_ptr<Buffer> buffer = nullptr;
+ /// Data offset in the buffer in bytes.
int64_t buffer_offset = 0;
/// Total number of values, i.e. the product of the shape.
int64_t size;
@@ -130,12 +121,12 @@ DT* ExportBuffer(ExportBufferParams<Vec>&& p) {
template <typename DT>
Result<DT*> ExportArrayImpl(const std::shared_ptr<Array>& arr, bool copy) {
- // Define DLDevice struct and check if array type is supported
- // by the DLPack protocol at the same time. Raise TypeError if not.
- // Supported data types: int, uint, float with no validity buffer.
+ if (arr->null_count() > 0) {
+ return Status::TypeError("Can only use DLPack on arrays with no nulls.");
+ }
ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(arr));
- // Define the DLDataType struct
+ // Define the DLDataType struct, or fail if the type is not supported.
const auto& type = *arr->type();
ARROW_ASSIGN_OR_RAISE(auto dtype, GetDLDataType(type));
@@ -167,6 +158,17 @@ Result<DT*> ExportArrayImpl(const std::shared_ptr<Array>&
arr, bool copy) {
return ExportBuffer<DT>(std::move(params));
}
+template <typename T>
+Result<DLDevice> ExportDeviceImpl(const std::shared_ptr<T>& a) {
+ // ArrayData reports the device of its buffers and children
+ if (a->data()->device_type() == DeviceAllocationType::kCPU) {
+ return {{.device_type = DLDeviceType::kDLCPU, .device_id = 0}};
+ } else {
+ return Status::NotImplemented(
+ "DLPack support is implemented only for buffers on CPU device.");
+ }
+}
+
} // namespace
Result<DLManagedTensor*> ExportArray(const std::shared_ptr<Array>& arr) {
@@ -179,29 +181,7 @@ Result<DLManagedTensorVersioned*>
ExportArrayVersioned(const std::shared_ptr<Arr
}
Result<DLDevice> ExportDevice(const std::shared_ptr<Array>& arr) {
- // Check if array is supported by the DLPack protocol.
- if (arr->null_count() > 0) {
- return Status::TypeError("Can only use DLPack on arrays with no nulls.");
- }
- const DataType& type = *arr->type();
- if (type.id() == Type::BOOL) {
- return Status::TypeError("Bit-packed boolean data type not supported by
DLPack.");
- }
- if (!is_integer(type.id()) && !is_floating(type.id())) {
- return Status::TypeError("DataType is not compatible with DLPack spec: ",
- type.ToString());
- }
-
- // Define DLDevice struct
- DLDevice device;
- if (arr->data()->buffers[1]->device_type() == DeviceAllocationType::kCPU) {
- device.device_id = 0;
- device.device_type = DLDeviceType::kDLCPU;
- return device;
- } else {
- return Status::NotImplemented(
- "DLPack support is implemented only for buffers on CPU device.");
- }
+ return ExportDeviceImpl(arr);
}
namespace {
@@ -265,16 +245,7 @@ Result<DLManagedTensorVersioned*>
ExportTensorVersioned(const std::shared_ptr<Te
}
Result<DLDevice> ExportDevice(const std::shared_ptr<Tensor>& t) {
- // Define DLDevice struct
- DLDevice device;
- if (t->data()->device_type() == DeviceAllocationType::kCPU) {
- device.device_id = 0;
- device.device_type = DLDeviceType::kDLCPU;
- return device;
- } else {
- return Status::NotImplemented(
- "DLPack support is implemented only for buffers on CPU device.");
- }
+ return ExportDeviceImpl(t);
}
} // namespace arrow::dlpack
diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc
index b46060fd8b..05de22237a 100644
--- a/cpp/src/arrow/c/dlpack_test.cc
+++ b/cpp/src/arrow/c/dlpack_test.cc
@@ -15,11 +15,13 @@
// specific language governing permissions and limitations
// under the License.
+#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <cstring>
#include <string>
#include <type_traits>
+#include <vector>
#include "arrow/array/array_base.h"
#include "arrow/buffer.h"
@@ -76,25 +78,28 @@ TYPED_TEST_SUITE(TestExportArray, ProducerTypes,
ProducerNames);
template <typename Producer>
void CheckDLTensor(const std::shared_ptr<Array>& arr,
const std::shared_ptr<DataType>& arrow_type,
- DLDataTypeCode dlpack_type, int64_t length) {
+ DLDataTypeCode dlpack_type, const std::vector<int64_t>&
shape,
+ const std::vector<int64_t>& strides) {
ASSERT_OK_AND_ASSIGN(auto* dlmtensor, Producer::Export(arr));
auto dltensor = dlmtensor->dl_tensor;
- const auto byte_width = arr->type()->byte_width();
+ ASSERT_EQ(arrow_type->id(), arr->type_id());
+ const auto byte_width = arrow_type->byte_width();
const auto start = arr->offset() * byte_width;
ASSERT_OK_AND_ASSIGN(auto sliced_buffer,
SliceBufferSafe(arr->data()->buffers[1], start));
if constexpr (Producer::copy) {
ASSERT_NE(sliced_buffer->data(), dltensor.data);
- ASSERT_EQ(0, std::memcmp(sliced_buffer->data(), dltensor.data, length *
byte_width));
+ ASSERT_EQ(
+ 0, std::memcmp(sliced_buffer->data(), dltensor.data, arr->length() *
byte_width));
} else {
ASSERT_EQ(sliced_buffer->data(), dltensor.data);
}
ASSERT_EQ(0, dltensor.byte_offset);
- ASSERT_EQ(length, dltensor.shape[0]);
- ASSERT_EQ(1, dltensor.ndim);
- ASSERT_EQ(1, *dltensor.strides); // Must be non-null with ndim>0 since 1.2
+ ASSERT_EQ(shape.size(), static_cast<size_t>(dltensor.ndim));
+ ASSERT_THAT(shape, ::testing::ElementsAreArray(dltensor.shape,
dltensor.ndim));
+ ASSERT_THAT(strides, ::testing::ElementsAreArray(dltensor.strides,
dltensor.ndim));
ASSERT_EQ(dlpack_type, dltensor.dtype.code);
ASSERT_EQ(arrow_type->bit_width(), dltensor.dtype.bits);
@@ -148,13 +153,13 @@ TYPED_TEST(TestExportArray, TestSupportedArray) {
for (auto [arrow_type, dlpack_type] : cases) {
const std::shared_ptr<Array> array =
ArrayFromJSON(arrow_type, "[1, 0, 10, 0, 2, 1, 3, 5, 1, 0]");
- CheckDLTensor<TypeParam>(array, arrow_type, dlpack_type, 10);
+ CheckDLTensor<TypeParam>(array, arrow_type, dlpack_type, {10}, {1});
ASSERT_OK_AND_ASSIGN(auto sliced_1, array->SliceSafe(1, 5));
- CheckDLTensor<TypeParam>(sliced_1, arrow_type, dlpack_type, 5);
+ CheckDLTensor<TypeParam>(sliced_1, arrow_type, dlpack_type, {5}, {1});
ASSERT_OK_AND_ASSIGN(auto sliced_2, array->SliceSafe(0, 5));
- CheckDLTensor<TypeParam>(sliced_2, arrow_type, dlpack_type, 5);
+ CheckDLTensor<TypeParam>(sliced_2, arrow_type, dlpack_type, {5}, {1});
ASSERT_OK_AND_ASSIGN(auto sliced_3, array->SliceSafe(3));
- CheckDLTensor<TypeParam>(sliced_3, arrow_type, dlpack_type, 7);
+ CheckDLTensor<TypeParam>(sliced_3, arrow_type, dlpack_type, {7}, {1});
}
ASSERT_EQ(allocated_bytes, arrow::default_memory_pool()->bytes_allocated());
@@ -164,7 +169,9 @@ TYPED_TEST(TestExportArray, TestErrors) {
const std::shared_ptr<Array> array_null = ArrayFromJSON(null(), "[]");
ASSERT_RAISES_WITH_MESSAGE(TypeError,
"Type error: DataType is not compatible with
DLPack spec: " +
- array_null->type()->ToString(),
+ array_null->type()->ToString() +
+ ", try converting to a Tensor for multi"
+ " dimensional data support",
TypeParam::Export(array_null));
const std::shared_ptr<Array> array_with_null = ArrayFromJSON(int8(), "[1,
100, null]");
@@ -176,13 +183,19 @@ TYPED_TEST(TestExportArray, TestErrors) {
ArrayFromJSON(utf8(), R"(["itsy", "bitsy", "spider"])");
ASSERT_RAISES_WITH_MESSAGE(TypeError,
"Type error: DataType is not compatible with
DLPack spec: " +
- array_string->type()->ToString(),
+ array_string->type()->ToString() +
+ ", try converting to a Tensor for multi"
+ " dimensional data support",
TypeParam::Export(array_string));
const std::shared_ptr<Array> array_boolean = ArrayFromJSON(boolean(),
"[true, false]");
ASSERT_RAISES_WITH_MESSAGE(
TypeError, "Type error: Bit-packed boolean data type not supported by
DLPack.",
- arrow::dlpack::ExportDevice(array_boolean));
+ TypeParam::Export(array_boolean));
+
+ // ExportDevice only reports the device, it does not validate the type
+ ASSERT_OK(arrow::dlpack::ExportDevice(array_boolean));
+ ASSERT_OK(arrow::dlpack::ExportDevice(array_null));
}
template <typename Producer>
diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc
b/cpp/src/arrow/extension/fixed_shape_tensor.cc
index 69eed2eed0..7f55375b82 100644
--- a/cpp/src/arrow/extension/fixed_shape_tensor.cc
+++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc
@@ -303,7 +303,7 @@ Result<std::shared_ptr<FixedShapeTensorArray>>
FixedShapeTensorArray::FromTensor
return std::static_pointer_cast<FixedShapeTensorArray>(ext_arr);
}
-const Result<std::shared_ptr<Tensor>> FixedShapeTensorArray::ToTensor() const {
+Result<std::shared_ptr<Tensor>> FixedShapeTensorArray::ToTensorWithNulls()
const {
// To convert an array of n dimensional tensors to a n+1 dimensional tensor
we
// interpret the array's length as the first dimension the new tensor.
diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.h
b/cpp/src/arrow/extension/fixed_shape_tensor.h
index eee44e1c81..0b3aafd26c 100644
--- a/cpp/src/arrow/extension/fixed_shape_tensor.h
+++ b/cpp/src/arrow/extension/fixed_shape_tensor.h
@@ -37,13 +37,17 @@ class ARROW_EXPORT FixedShapeTensorArray : public
ExtensionArray {
static Result<std::shared_ptr<FixedShapeTensorArray>> FromTensor(
const std::shared_ptr<Tensor>& tensor);
+ protected:
/// \brief Create a Tensor from FixedShapeTensorArray
///
/// This method will create a Tensor from a FixedShapeTensorArray, setting
its first
/// dimension as length equal to the FixedShapeTensorArray's length and the
remaining
/// dimensions as the FixedShapeTensorType's shape. Shape and dim_names will
be
/// permuted according to permutation stored in the FixedShapeTensorType
metadata.
- const Result<std::shared_ptr<Tensor>> ToTensor() const;
+ ///
+ /// Nulls are ignored, leaving the output tensor with unspecified values
where this
+ /// array has null entries.
+ Result<std::shared_ptr<Tensor>> ToTensorWithNulls() const override;
};
/// \brief Concrete type class for constant-size Tensor data.
diff --git a/cpp/src/arrow/tensor.cc b/cpp/src/arrow/tensor.cc
index b5988d7810..f2ff11a4f6 100644
--- a/cpp/src/arrow/tensor.cc
+++ b/cpp/src/arrow/tensor.cc
@@ -144,7 +144,7 @@ inline Status CheckTensorValidity(const
std::shared_ptr<DataType>& type,
return Status::Invalid("Null type is supplied");
}
if (!is_tensor_supported(type->id())) {
- return Status::Invalid(type->ToString(), " is not valid data type for a
tensor");
+ return Status::TypeError(type->ToString(), " is not valid data type for a
tensor");
}
if (!data) {
return Status::Invalid("Null data is supplied");
@@ -488,6 +488,11 @@ Status RecordBatchToTensor(const RecordBatch& batch, bool
null_to_nan, bool row_
} // namespace internal
+Result<std::shared_ptr<Tensor>> Tensor::FromArray(const
std::shared_ptr<Array>& array,
+ bool allow_nulls) {
+ return array->ToTensor(allow_nulls);
+}
+
/// Constructor with strides and dimension names
Tensor::Tensor(const std::shared_ptr<DataType>& type, const
std::shared_ptr<Buffer>& data,
const std::vector<int64_t>& shape, const std::vector<int64_t>&
strides,
diff --git a/cpp/src/arrow/tensor.h b/cpp/src/arrow/tensor.h
index 1300003c29..f327031343 100644
--- a/cpp/src/arrow/tensor.h
+++ b/cpp/src/arrow/tensor.h
@@ -27,6 +27,7 @@
#include "arrow/result.h"
#include "arrow/status.h"
#include "arrow/type.h"
+#include "arrow/type_fwd.h"
#include "arrow/type_traits.h"
#include "arrow/util/macros.h"
#include "arrow/util/visibility.h"
@@ -109,6 +110,12 @@ class ARROW_EXPORT Tensor {
return std::make_shared<Tensor>(type, data, shape, strides, dim_names);
}
+ /// \brief Attempt to create a Tensor from an Array.
+ ///
+ /// \see Array::ToTensor
+ static Result<std::shared_ptr<Tensor>> FromArray(const
std::shared_ptr<Array>& array,
+ bool allow_nulls = false);
+
virtual ~Tensor() = default;
/// Constructor with no dimension names or strides, data assumed to be
row-major
diff --git a/cpp/src/arrow/tensor_test.cc b/cpp/src/arrow/tensor_test.cc
index 2a2f564e79..8b0b4ff355 100644
--- a/cpp/src/arrow/tensor_test.cc
+++ b/cpp/src/arrow/tensor_test.cc
@@ -235,7 +235,7 @@ TEST(TestTensor, MakeFailureCases) {
ASSERT_RAISES(Invalid, Tensor::Make(nullptr, data, shape));
// invalid type
- ASSERT_RAISES(Invalid, Tensor::Make(binary(), data, shape));
+ ASSERT_RAISES(TypeError, Tensor::Make(binary(), data, shape));
// null data
ASSERT_RAISES(Invalid, Tensor::Make(float64(), nullptr, shape));
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
index e31feb1cb0..2b2130e992 100644
--- a/python/pyarrow/array.pxi
+++ b/python/pyarrow/array.pxi
@@ -1836,6 +1836,34 @@ cdef class Array(_PandasConvertible):
array = array.copy()
return array
+ def to_tensor(self, *, allow_nulls=False):
+ """
+ Convert this array to a pyarrow.Tensor.
+
+ This is supported when the data can reasonably be understood as a
+ multi-dimensional numeric tensor, such as numeric arrays (1D), nested
+ fixed size list arrays, and fixed shape tensor arrays.
+ The resulting tensor has a row major layout with the array elements
+ as the first dimension. The conversion is zero-copy.
+
+ Parameters
+ ----------
+ allow_nulls : bool, default `False`
+ When true, nulls are ignored, leaving the output tensor with
+ unspecified values where this array has null entries.
+ When false, nulls are rejected.
+
+ Returns
+ -------
+ pyarrow.Tensor
+ """
+ cdef:
+ shared_ptr[CTensor] ctensor
+ c_bool c_allow_nulls = allow_nulls
+ with nogil:
+ ctensor = GetResultValue(self.ap.ToTensor(c_allow_nulls))
+ return pyarrow_wrap_tensor(ctensor)
+
def to_pylist(self, *, maps_as_pydicts=None):
"""
Convert to a list of native Python objects.
@@ -4815,15 +4843,12 @@ cdef class ExtensionArray(Array):
-------
ext_array : ExtensionArray
"""
- cdef:
- shared_ptr[CExtensionArray] ext_array
-
if storage.type != typ.storage_type:
raise TypeError(f"Incompatible storage type {storage.type} "
f"for extension type {typ}")
- ext_array = make_shared[CExtensionArray](typ.sp_type, storage.sp_array)
- cdef Array result = pyarrow_wrap_array(<shared_ptr[CArray]> ext_array)
+ cdef Array result = pyarrow_wrap_array(
+ typ.ext_type.WrapArray(typ.sp_type, storage.sp_array))
result.validate()
return result
@@ -4944,31 +4969,6 @@ cdef class FixedShapeTensorArray(ExtensionArray):
return self.to_tensor().to_numpy()
- def to_tensor(self):
- """
- Convert fixed shape tensor extension array to a pyarrow.Tensor.
-
- The resulting Tensor will have (ndim + 1) dimensions.
- The size of the first dimension will be the length of the fixed shape
tensor array
- and the rest of the dimensions will match the permuted shape of the
fixed
- shape tensor.
-
- The conversion is zero-copy.
-
- Returns
- -------
- pyarrow.Tensor
- Tensor representing tensors in the fixed shape tensor array
concatenated
- along the first dimension.
- """
-
- cdef:
- CFixedShapeTensorArray* ext_array =
<CFixedShapeTensorArray*>(self.ap)
- CResult[shared_ptr[CTensor]] ctensor
- with nogil:
- ctensor = ext_array.ToTensor()
- return pyarrow_wrap_tensor(GetResultValue(ctensor))
-
@staticmethod
def from_numpy_ndarray(obj, dim_names=None):
"""
diff --git a/python/pyarrow/includes/libarrow.pxd
b/python/pyarrow/includes/libarrow.pxd
index 3e6a19ffe9..ffc02ffd79 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -280,6 +280,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
const shared_ptr[CArrayStatistics]& statistics() const
+ CResult[shared_ptr[CTensor]] ToTensor(c_bool allow_nulls) const
+
shared_ptr[CArray] MakeArray(const shared_ptr[CArrayData]& data)
CResult[shared_ptr[CArray]] MakeArrayOfNull(
const shared_ptr[CDataType]& type, int64_t length, CMemoryPool* pool)
@@ -3088,10 +3090,6 @@ cdef extern from "arrow/extension/fixed_shape_tensor.h"
namespace "arrow::extens
const vector[int64_t] permutation()
const vector[c_string] dim_names()
- cdef cppclass CFixedShapeTensorArray \
- " arrow::extension::FixedShapeTensorArray"(CExtensionArray):
- const CResult[shared_ptr[CTensor]] ToTensor() const
-
cdef extern from "arrow/extension/opaque.h" namespace "arrow::extension" nogil:
cdef cppclass COpaqueType \
diff --git a/python/pyarrow/tests/test_dlpack.py
b/python/pyarrow/tests/test_dlpack.py
index 7f5f98d866..f9aac892ce 100644
--- a/python/pyarrow/tests/test_dlpack.py
+++ b/python/pyarrow/tests/test_dlpack.py
@@ -135,6 +135,68 @@ def test_tensor_dlpack(np_type):
check_dlpack_export(t, expected)
+def multidim_arrays():
+ np_arr = np.arange(12, dtype=np.int32).reshape(3, 2, 2)
+ values = pa.array(np_arr.ravel(), type=pa.int32())
+ nested_list = pa.FixedSizeListArray.from_arrays(
+ pa.FixedSizeListArray.from_arrays(values, 2), 2)
+ return [
+ pytest.param(nested_list, np_arr, id="nested_fixed_size_list"),
+ pytest.param(
+ pa.FixedShapeTensorArray.from_numpy_ndarray(np_arr),
+ np_arr,
+ id="fixed_shape_tensor",
+ ),
+ ]
+
+
+@check_bytes_allocated
[email protected](('arr', 'expected'), multidim_arrays())
+def test_array_to_tensor_dlpack(arr, expected):
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ tensor = arr.to_tensor()
+ # A Tensor sharing an Array buffer is immutable, so it can only be exported
+ # through the versioned DLPack protocol.
+ assert not tensor.is_mutable
+ result = np.from_dlpack(DLPackForwarder(tensor, max_version=(1, 0)))
+ np.testing.assert_array_equal(result, expected, strict=True)
+ assert tensor.__dlpack_device__() == (1, 0)
+
+
+def multidim_arrays_with_nulls():
+ np_arr = np.arange(6, dtype=np.int32).reshape(3, 2)
+ # Masked entries keep defined values in the child array, so the tensor
+ # contents stay fully predictable.
+ nested_list = pa.FixedSizeListArray.from_arrays(
+ pa.array(np_arr.ravel(), type=pa.int32()), 2,
+ mask=pa.array([False, True, False]))
+ return [
+ pytest.param(nested_list, np_arr, id="fixed_size_list"),
+ pytest.param(
+ pa.ExtensionArray.from_storage(
+ pa.fixed_shape_tensor(pa.int32(), [2]), nested_list),
+ np_arr,
+ id="fixed_shape_tensor",
+ ),
+ ]
+
+
+@check_bytes_allocated
[email protected](('arr', 'expected'), multidim_arrays_with_nulls())
+def test_array_to_tensor_dlpack_nulls(arr, expected):
+ if Version(np.__version__) < Version("2.1.0"):
+ pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later")
+
+ with pytest.raises(pa.ArrowInvalid, match="Array contains nulls"):
+ arr.to_tensor()
+
+ tensor = arr.to_tensor(allow_nulls=True)
+ result = np.from_dlpack(DLPackForwarder(tensor, max_version=(1, 0)))
+ np.testing.assert_array_equal(result, expected, strict=True)
+
+
def dlpack_objects():
arr = pa.array([1, 2, 3], type=pa.int32())
return [
