rok commented on code in PR #50203:
URL: https://github.com/apache/arrow/pull/50203#discussion_r3469329534
##########
python/pyarrow/tests/test_extension_type.py:
##########
@@ -1730,6 +1730,60 @@ def test_tensor_array_from_numpy(np_type_str):
pa.FixedShapeTensorArray.from_numpy_ndarray(arr, dim_names=[0, 1])
[email protected]
[email protected]("np_type_str", ("int8", "int64", "float32"))
+def test_tensor_array_from_list_of_ndarrays(np_type_str):
+ # GH-49644
+ np_dtype = np.dtype(np_type_str)
+ tensor_type = pa.fixed_shape_tensor(pa.from_numpy_dtype(np_dtype), (2, 3))
+
+ elements = [
+ np.arange(6, dtype=np_dtype).reshape(2, 3),
+ np.arange(6, 12, dtype=np_dtype).reshape(2, 3),
+ ]
+ result = pa.array(elements, type=tensor_type)
+ assert isinstance(result, pa.FixedShapeTensorArray)
+ assert result.type == tensor_type
+ assert len(result) == 2
+
+ # Must match the existing from_numpy_ndarray path on the same data
+ expected = pa.FixedShapeTensorArray.from_numpy_ndarray(np.stack(elements))
+ assert result.storage.equals(expected.storage)
+
+ # Each element round-trips back to the original ndarray (with its shape)
+ for scalar, original in zip(result, elements):
+ np.testing.assert_array_equal(scalar.to_numpy(), original)
+
+ # Higher-dimensional tensors work too
+ tensor_3d = pa.fixed_shape_tensor(pa.from_numpy_dtype(np_dtype), (2, 2, 3))
+ elements_3d = [np.arange(12, dtype=np_dtype).reshape(2, 2, 3)]
+ result_3d = pa.array(elements_3d, type=tensor_3d)
+ assert result_3d.type == tensor_3d
+ np.testing.assert_array_equal(result_3d[0].to_numpy(), elements_3d[0])
+
+ # None elements are allowed
+ result_with_null = pa.array([elements[0], None], type=tensor_type)
+ assert result_with_null.null_count == 1
+ assert result_with_null[1].as_py() is None
+
Review Comment:
Let's add another test here with mixed layouts. I suppose logic here
https://github.com/apache/arrow/pull/50203/changes#r3467238608 should make
everything C-contiguous and this would pass and return correct result.
```python
tensor_type = pa.fixed_shape_tensor(pa.int64(), (2, 3))
raw = [[1, 2, 3], [4, 5, 6]]
c_arr = np.array(raw, order='C')
f_arr = np.array(raw, order='F')
assert np.array_equal(c_arr, f_arr)
assert c_arr.tobytes('A') != f_arr.tobytes('A')
same = pa.array([c_arr, c_arr], type=tensor_type)
mixed = pa.array([c_arr, f_arr], type=tensor_type)
assert mixed.equals(same)
assert mixed.storage.to_pylist() == [[1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6]]
```
--
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]
