aboderinsamuel commented on code in PR #50203:
URL: https://github.com/apache/arrow/pull/50203#discussion_r3481633667
##########
python/pyarrow/tests/test_extension_type.py:
##########
@@ -1730,6 +1730,82 @@ 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
+
+ # A multi-dimensional element whose shape doesn't match the tensor shape is
+ # rejected, even when the total number of elements is the same (GH-49644).
+ with pytest.raises(ValueError, match="shape"):
+ pa.array([np.arange(6, dtype=np_dtype).reshape(3, 2)],
type=tensor_type)
+
+ # Permuted tensor types can't be built from a sequence (the flatten would
+ # store the wrong layout), so they're rejected for now.
+ permuted_type = pa.fixed_shape_tensor(
+ pa.from_numpy_dtype(np_dtype), (2, 3), permutation=[1, 0])
+ with pytest.raises(NotImplementedError, match="permutation"):
+ pa.array(elements, type=permuted_type)
+
+
[email protected]
+def test_tensor_array_from_list_mixed_layout():
+ # GH-49644: C- and F-ordered arrays with the same values must produce the
+ # same result, since the values are always flattened in C order.
Review Comment:
Done, removed
--
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]
