pitrou commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3958653503
##########
cpp/src/arrow/buffer.h:
##########
@@ -146,14 +147,66 @@ class ARROW_EXPORT Buffer {
}
}
- /// \brief Construct an immutable buffer that takes ownership of the contents
+ /// \brief Default data accessor used by TakeOwnership.
+ struct DefaultGetData {
Review Comment:
Can we just make this private?
Or perhaps just use
[`std::data`](https://en.cppreference.com/cpp/iterator/data)?
##########
python/pyarrow/tests/test_extension_type.py:
##########
@@ -1708,6 +1708,26 @@ def test_tensor_class_methods(np_type_str):
assert result.to_tensor().strides == (12 * bw, 1 * bw, 3 * bw, 6 * bw)
[email protected]
[email protected](
+ ("transpose", "permutation"),
+ [(False, [0, 1]), (True, [1, 0])]
+)
+def test_tensor_array_from_tensor(transpose, permutation):
+ arr = np.arange(24, dtype=np.int32).reshape(2, 3, 4)
+ arr = arr.transpose(0, 2, 1) if transpose else arr
+
+ result = pa.FixedShapeTensorArray.from_tensor(pa.Tensor.from_numpy(arr))
Review Comment:
Also validate the result here?
##########
python/pyarrow/tests/test_dlpack.py:
##########
@@ -374,3 +374,101 @@ def test_dlpack_cuda_not_supported():
with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
"only for buffers on CPU device."):
carr.__dlpack_device__()
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_tensor_from_dlpack(np_type):
+ def make_array():
+ base = np.arange(24, dtype=np_type).reshape((4, 6))
+ array = base[::2, 1::2]
+ assert not array.flags['C_CONTIGUOUS']
+ return array
+
+ # Non-contiguous, strided slice: DLPack carries explicit strides, so this
+ # should not need a copy on export.
+ tensor = pa.Tensor.from_dlpack(make_array())
+ assert isinstance(tensor, pa.Tensor)
+ gc.collect() # Attempts to free input array memory
+ np.testing.assert_array_equal(tensor.to_numpy(), make_array(), strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_array_from_dlpack(np_type):
+ expected = np.array([1, 2, 3, 4, 5], dtype=np_type)
+ arr = pa.Array.from_dlpack(expected)
+ arr.validate(full=True)
+ assert isinstance(arr, pa.Array)
+ np.testing.assert_array_equal(arr.to_numpy(), expected, strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_fixed_shape_tensor_array_from_dlpack(np_type):
+ source = np.arange(12, dtype=np_type).reshape((3, 2, 2))
Review Comment:
Can we also test with a transposed (non-row major) source as well?
##########
python/pyarrow/tests/test_extension_type.py:
##########
@@ -1708,6 +1708,26 @@ def test_tensor_class_methods(np_type_str):
assert result.to_tensor().strides == (12 * bw, 1 * bw, 3 * bw, 6 * bw)
[email protected]
[email protected](
+ ("transpose", "permutation"),
+ [(False, [0, 1]), (True, [1, 0])]
+)
+def test_tensor_array_from_tensor(transpose, permutation):
+ arr = np.arange(24, dtype=np.int32).reshape(2, 3, 4)
+ arr = arr.transpose(0, 2, 1) if transpose else arr
Review Comment:
What if the first dimension is transposed too? I suspect we'll get an error,
can we test that?
##########
python/pyarrow/tests/test_dlpack.py:
##########
@@ -374,3 +374,101 @@ def test_dlpack_cuda_not_supported():
with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
"only for buffers on CPU device."):
carr.__dlpack_device__()
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_tensor_from_dlpack(np_type):
+ def make_array():
+ base = np.arange(24, dtype=np_type).reshape((4, 6))
+ array = base[::2, 1::2]
+ assert not array.flags['C_CONTIGUOUS']
+ return array
+
+ # Non-contiguous, strided slice: DLPack carries explicit strides, so this
+ # should not need a copy on export.
+ tensor = pa.Tensor.from_dlpack(make_array())
+ assert isinstance(tensor, pa.Tensor)
+ gc.collect() # Attempts to free input array memory
+ np.testing.assert_array_equal(tensor.to_numpy(), make_array(), strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_array_from_dlpack(np_type):
+ expected = np.array([1, 2, 3, 4, 5], dtype=np_type)
+ arr = pa.Array.from_dlpack(expected)
+ arr.validate(full=True)
+ assert isinstance(arr, pa.Array)
+ np.testing.assert_array_equal(arr.to_numpy(), expected, strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_fixed_shape_tensor_array_from_dlpack(np_type):
+ source = np.arange(12, dtype=np_type).reshape((3, 2, 2))
+ arr = pa.FixedShapeTensorArray.from_dlpack(source)
+ arr.validate(full=True)
+ assert arr.type == pa.fixed_shape_tensor(pa.from_numpy_dtype(np_type), [2,
2])
+ assert arr.to_pylist() == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]
Review Comment:
Interesting. I'd have expected `[[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[8,
9], [10, 11]]]`, which is also what NumPy returns. It's not really important,
but slightly surprising IMHO.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
