Skip to content

Commit bab66ac

Browse files
feat(python): Clarify interaction between the CDeviceArray, the CArrayView, and the CArray (#409)
When device support was first added, the `CArrayView` was device-aware but the `CArray` was not. This worked well until it was clear that `__arrow_c_array__` needed to error if it did not represent a CPU array (and the `CArray` had no way to check). Now, the `CArray` has a `device_type` and `device_id`. A nice side-effect of this is that we get back the `view()` method (whose removal @jorisvandenbossche had lamented!). This also implements the device array protocol to help test apache/arrow#40717 . This protocol isn't finalized yet and I could remove that part until it is (although it doesn't seem likely to change). The non-cpu case is still hard to test without real-world CUDA support...this PR is just trying to get the right information in the right place as early as possible. ```python import nanoarrow as na array = na.c_array([1, 2, 3], na.int32()) array.device_type, array.device_id #> (1, 0) ``` --------- Co-authored-by: Dane Pitkin <[email protected]>
1 parent 00aa9c3 commit bab66ac

File tree

11 files changed

+281
-68
lines changed

11 files changed

+281
-68
lines changed

extensions/nanoarrow_device/src/nanoarrow/nanoarrow_device.c

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ struct ArrowDevice* ArrowDeviceCpu(void) {
115115

116116
void ArrowDeviceInitCpu(struct ArrowDevice* device) {
117117
device->device_type = ARROW_DEVICE_CPU;
118-
device->device_id = 0;
118+
device->device_id = -1;
119119
device->array_init = NULL;
120120
device->array_move = NULL;
121121
device->buffer_init = &ArrowDeviceCpuBufferInit;
@@ -135,7 +135,7 @@ struct ArrowDevice* ArrowDeviceCuda(ArrowDeviceType device_type, int64_t device_
135135
#endif
136136

137137
struct ArrowDevice* ArrowDeviceResolve(ArrowDeviceType device_type, int64_t device_id) {
138-
if (device_type == ARROW_DEVICE_CPU && device_id == 0) {
138+
if (device_type == ARROW_DEVICE_CPU) {
139139
return ArrowDeviceCpu();
140140
}
141141

extensions/nanoarrow_device/src/nanoarrow/nanoarrow_device_test.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ TEST(NanoarrowDevice, CheckRuntime) {
2828
TEST(NanoarrowDevice, CpuDevice) {
2929
struct ArrowDevice* cpu = ArrowDeviceCpu();
3030
EXPECT_EQ(cpu->device_type, ARROW_DEVICE_CPU);
31-
EXPECT_EQ(cpu->device_id, 0);
31+
EXPECT_EQ(cpu->device_id, -1);
3232
EXPECT_EQ(cpu, ArrowDeviceCpu());
3333

3434
void* sync_event = nullptr;

0 commit comments

Comments
 (0)