| 32 | }; |
| 33 | |
| 34 | void CheckDLTensor(const std::shared_ptr<Array>& arr, |
| 35 | const std::shared_ptr<DataType>& arrow_type, |
| 36 | DLDataTypeCode dlpack_type, int64_t length) { |
| 37 | ASSERT_OK_AND_ASSIGN(auto dlmtensor, arrow::dlpack::ExportArray(arr)); |
| 38 | auto dltensor = dlmtensor->dl_tensor; |
| 39 | |
| 40 | const auto byte_width = arr->type()->byte_width(); |
| 41 | const auto start = arr->offset() * byte_width; |
| 42 | ASSERT_OK_AND_ASSIGN(auto sliced_buffer, |
| 43 | SliceBufferSafe(arr->data()->buffers[1], start)); |
| 44 | ASSERT_EQ(sliced_buffer->data(), dltensor.data); |
| 45 | |
| 46 | ASSERT_EQ(0, dltensor.byte_offset); |
| 47 | ASSERT_EQ(NULL, dltensor.strides); |
| 48 | ASSERT_EQ(length, dltensor.shape[0]); |
| 49 | ASSERT_EQ(1, dltensor.ndim); |
| 50 | |
| 51 | ASSERT_EQ(dlpack_type, dltensor.dtype.code); |
| 52 | ASSERT_EQ(arrow_type->bit_width(), dltensor.dtype.bits); |
| 53 | ASSERT_EQ(1, dltensor.dtype.lanes); |
| 54 | ASSERT_EQ(DLDeviceType::kDLCPU, dltensor.device.device_type); |
| 55 | ASSERT_EQ(0, dltensor.device.device_id); |
| 56 | |
| 57 | ASSERT_OK_AND_ASSIGN(auto device, arrow::dlpack::ExportDevice(arr)); |
| 58 | ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type); |
| 59 | ASSERT_EQ(0, device.device_id); |
| 60 | |
| 61 | dlmtensor->deleter(dlmtensor); |
| 62 | } |
| 63 | |
| 64 | TEST_F(TestExportArray, TestSupportedArray) { |
| 65 | const std::vector<std::pair<std::shared_ptr<DataType>, DLDataTypeCode>> cases = { |