diff --git a/cpp/src/arrow/array/array_base.cc b/cpp/src/arrow/array/array_base.cc index ce2e66655af3..0add32a55525 100644 --- a/cpp/src/arrow/array/array_base.cc +++ b/cpp/src/arrow/array/array_base.cc @@ -323,6 +323,20 @@ Result> Array::ViewOrCopyTo( return MakeArray(new_data); } +Result> 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> 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 60df45357e5d..ed2e2356eac6 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& const std::shared_ptr& 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> 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> 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 8406bd1d8ed1..9c4202f32029 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" @@ -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 LeafValues(std::shared_ptr array) { + while (array->type_id() == Type::FIXED_SIZE_LIST) { + const auto& fsl = checked_cast(*array); + array = + fsl.values()->Slice(fsl.value_offset(0), array->length() * fsl.value_length()); + } + return array; +} + +template +void CheckToTensor(const std::shared_ptr& array, const std::vector& shape, + std::initializer_list values) { + const auto value_type = CTypeTraits::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(array, {4, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + + // Offset on the list array itself + CheckToTensor(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(from_values, {2, 3}, {4, 5, 6, 7, 8, 9}); + + // Offsets on both the list array and its values + CheckToTensor(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(array, {3, 3, 2}, + {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}); + + CheckToTensor(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(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(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({0, 0}), 1); + ASSERT_EQ(tensor->Value({0, 1}), 2); + ASSERT_EQ(tensor->Value({2, 0}), 5); + ASSERT_EQ(std::vector({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({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 c5a26a475c9e..fccd6c8ec97b 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> FixedSizeListArray::Flatten( return FlattenListArray(*this, memory_pool); } +Result> 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 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(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 = 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 bf84f802b1ab..57c87f900c6c 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& 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> ToTensorWithNulls() const override; + private: std::shared_ptr values_; }; diff --git a/cpp/src/arrow/array/array_primitive.h b/cpp/src/arrow/array/array_primitive.h index cebf47ad93d8..473ac16995d0 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> ToTensorWithNulls() const override { + // Could be non-templated + const int64_t byte_width = type()->byte_width(); + std::shared_ptr 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 dcfe1c76c301..668d9f00d883 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 shape = {5}; + const std::vector 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 shape = {3}; + const std::vector 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{0}, tensor->shape()); + EXPECT_EQ(std::vector{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({0}), 1); + ASSERT_EQ(tensor->Value({2}), 3); + EXPECT_EQ(std::vector{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 6eb06c738c53..4e25d50bb56f 100644 --- a/cpp/src/arrow/c/dlpack.cc +++ b/cpp/src/arrow/c/dlpack.cc @@ -35,34 +35,24 @@ namespace arrow::dlpack { namespace { Result 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 @@ -78,6 +68,7 @@ struct ManagerCtx { template struct ExportBufferParams { std::shared_ptr 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&& p) { template Result ExportArrayImpl(const std::shared_ptr& 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 ExportArrayImpl(const std::shared_ptr& arr, bool copy) { return ExportBuffer
(std::move(params)); } +template +Result ExportDeviceImpl(const std::shared_ptr& 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 ExportArray(const std::shared_ptr& arr) { @@ -179,29 +181,7 @@ Result ExportArrayVersioned(const std::shared_ptr ExportDevice(const std::shared_ptr& 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 ExportTensorVersioned(const std::shared_ptr ExportDevice(const std::shared_ptr& 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 b46060fd8be2..05de22237a9f 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 #include #include #include #include +#include #include "arrow/array/array_base.h" #include "arrow/buffer.h" @@ -76,25 +78,28 @@ TYPED_TEST_SUITE(TestExportArray, ProducerTypes, ProducerNames); template void CheckDLTensor(const std::shared_ptr& arr, const std::shared_ptr& arrow_type, - DLDataTypeCode dlpack_type, int64_t length) { + DLDataTypeCode dlpack_type, const std::vector& shape, + const std::vector& 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(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 = ArrayFromJSON(arrow_type, "[1, 0, 10, 0, 2, 1, 3, 5, 1, 0]"); - CheckDLTensor(array, arrow_type, dlpack_type, 10); + CheckDLTensor(array, arrow_type, dlpack_type, {10}, {1}); ASSERT_OK_AND_ASSIGN(auto sliced_1, array->SliceSafe(1, 5)); - CheckDLTensor(sliced_1, arrow_type, dlpack_type, 5); + CheckDLTensor(sliced_1, arrow_type, dlpack_type, {5}, {1}); ASSERT_OK_AND_ASSIGN(auto sliced_2, array->SliceSafe(0, 5)); - CheckDLTensor(sliced_2, arrow_type, dlpack_type, 5); + CheckDLTensor(sliced_2, arrow_type, dlpack_type, {5}, {1}); ASSERT_OK_AND_ASSIGN(auto sliced_3, array->SliceSafe(3)); - CheckDLTensor(sliced_3, arrow_type, dlpack_type, 7); + CheckDLTensor(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_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_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_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 diff --git a/cpp/src/arrow/extension/fixed_shape_tensor.cc b/cpp/src/arrow/extension/fixed_shape_tensor.cc index cd3d783479d6..e0162edbe9ac 100644 --- a/cpp/src/arrow/extension/fixed_shape_tensor.cc +++ b/cpp/src/arrow/extension/fixed_shape_tensor.cc @@ -319,7 +319,7 @@ Result> FixedShapeTensorArray::FromTensor return std::static_pointer_cast(ext_arr); } -const Result> FixedShapeTensorArray::ToTensor() const { +Result> 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 eee44e1c8164..0b3aafd26c4a 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> FromTensor( const std::shared_ptr& 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> ToTensor() const; + /// + /// Nulls are ignored, leaving the output tensor with unspecified values where this + /// array has null entries. + Result> 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 b5988d78106c..f2ff11a4f66e 100644 --- a/cpp/src/arrow/tensor.cc +++ b/cpp/src/arrow/tensor.cc @@ -144,7 +144,7 @@ inline Status CheckTensorValidity(const std::shared_ptr& 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> Tensor::FromArray(const std::shared_ptr& array, + bool allow_nulls) { + return array->ToTensor(allow_nulls); +} + /// Constructor with strides and dimension names Tensor::Tensor(const std::shared_ptr& type, const std::shared_ptr& data, const std::vector& shape, const std::vector& strides, diff --git a/cpp/src/arrow/tensor.h b/cpp/src/arrow/tensor.h index 1300003c2985..f3270313434e 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(type, data, shape, strides, dim_names); } + /// \brief Attempt to create a Tensor from an Array. + /// + /// \see Array::ToTensor + static Result> FromArray(const std::shared_ptr& 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 2a2f564e7910..8b0b4ff3551c 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 d8bbd001fddc..a578e3a5e6f9 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -1841,6 +1841,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. @@ -4820,15 +4848,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( ext_array) + cdef Array result = pyarrow_wrap_array( + typ.ext_type.WrapArray(typ.sp_type, storage.sp_array)) result.validate() return result @@ -4949,31 +4974,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 = (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 3e6a19ffe9e5..ffc02ffd79ae 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 09a510122fdf..3971b64f09ed 100644 --- a/python/pyarrow/tests/test_dlpack.py +++ b/python/pyarrow/tests/test_dlpack.py @@ -145,6 +145,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 +@pytest.mark.parametrize(('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 +@pytest.mark.parametrize(('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 [