diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index 4c625c45e489..c3ddf7b3de75 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -266,6 +266,23 @@ def array(object obj, type=None, mask=None, size=None, from_pandas=None, if type is not None and type.id == _Type_EXTENSION: extension_type = type type = type.storage_type + # GH-49644: when building a fixed_shape_tensor from a sequence of arrays, + # the converter only sees the flat storage type, so validate the + # tensor-specific constraints here where the type is still known. + if (isinstance(extension_type, FixedShapeTensorType) + and isinstance(obj, (list, tuple))): + if extension_type.permutation is not None: + raise NotImplementedError( + "Converting a sequence of arrays to a fixed_shape_tensor " + "with a permutation is not supported") + expected_shape = tuple(extension_type.shape) + for element in obj: + shape = getattr(element, "shape", None) + if (shape is not None and len(shape) >= 2 + and tuple(shape) != expected_shape): + raise ValueError( + f"Cannot convert array of shape {tuple(shape)} to a " + f"fixed_shape_tensor of shape {expected_shape}") if from_pandas is None: c_from_pandas = False @@ -4691,6 +4708,30 @@ cdef class FixedShapeTensorArray(ExtensionArray): 400 ] ] + + Create an extension array from a list of multi-dimensional NumPy arrays. + Each element is flattened in row-major (C) order, and its shape must match + the tensor shape. + + >>> import numpy as np + >>> pa.array([np.array([[1, 2], [3, 4]], dtype=np.int32), + ... np.array([[10, 20], [30, 40]], dtype=np.int32)], + ... type=tensor_type) + + [ + [ + 1, + 2, + 3, + 4 + ], + [ + 10, + 20, + 30, + 40 + ] + ] """ def to_numpy_ndarray(self): diff --git a/python/pyarrow/src/arrow/python/python_to_arrow.cc b/python/pyarrow/src/arrow/python/python_to_arrow.cc index e7ce54abcd8f..47290a7e7f10 100644 --- a/python/pyarrow/src/arrow/python/python_to_arrow.cc +++ b/python/pyarrow/src/arrow/python/python_to_arrow.cc @@ -908,13 +908,30 @@ class PyListConverter : public ListConverter { Status AppendNdarray(PyObject* value) { PyArrayObject* ndarray = reinterpret_cast(value); - if (PyArray_NDIM(ndarray) != 1) { - return Status::Invalid("Can only convert 1-dimensional array values"); - } if (PyArray_ISBYTESWAPPED(ndarray)) { // TODO return Status::NotImplemented("Byte-swapped arrays not supported"); } + OwnedRef flattened; + if (PyArray_NDIM(ndarray) != 1) { + // GH-49644: a fixed-size list (e.g. fixed-shape-tensor storage) is built + // from a multi- or 0-dimensional array by flattening it in C order. + if (this->list_type_->id() != Type::FIXED_SIZE_LIST) { + return Status::Invalid("Can only convert 1-dimensional array values of ", + this->list_type_->ToString(), " to a variable-sized list"); + } + // Get an aligned, C-contiguous array (copying only if needed). + PyObject* contiguous = + PyArray_CheckFromAny(value, nullptr, /*min_depth=*/0, /*max_depth=*/0, + NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_ALIGNED, nullptr); + RETURN_IF_PYERROR(); + flattened.reset( + PyArray_Ravel(reinterpret_cast(contiguous), NPY_CORDER)); + Py_DECREF(contiguous); + RETURN_IF_PYERROR(); + value = flattened.obj(); + ndarray = reinterpret_cast(value); + } const int64_t size = PyArray_SIZE(ndarray); RETURN_NOT_OK(AppendTo(this->list_type_, size)); RETURN_NOT_OK(this->list_builder_->ValidateOverflow(size)); diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 0205db2393ab..adc3e097b54a 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -2924,6 +2924,30 @@ def test_array_from_invalid_dim_raises(): pa.array(arr0d) +@pytest.mark.numpy +def test_fixed_size_list_from_multidim_ndarray(): + arr = pa.array([np.array([[1, 2, 3]], dtype=np.int64), + np.array([[4, 5, 6]], dtype=np.int64)], + type=pa.list_(pa.int64(), 3)) + assert arr.type == pa.list_(pa.int64(), 3) + assert arr.to_pylist() == [[1, 2, 3], [4, 5, 6]] + + arr = pa.array([np.array([[1, 2], [3, 4]], dtype=np.int64)], + type=pa.list_(pa.int64(), 4)) + assert arr.to_pylist() == [[1, 2, 3, 4]] + + with pytest.raises(pa.lib.ArrowInvalid): + pa.array([np.array([[1, 2], [3, 4]], dtype=np.int64)], + type=pa.list_(pa.int64(), 3)) + + with pytest.raises(pa.lib.ArrowInvalid, match=r"array values of .*int64"): + pa.array([np.array([[1, 2, 3]], dtype=np.int64)], + type=pa.list_(pa.int64())) + + arr = pa.array([np.array(1, dtype=np.int64)], type=pa.list_(pa.int64(), 1)) + assert arr.to_pylist() == [[1]] + + @pytest.mark.numpy def test_array_from_strided_bool(): # ARROW-6325 diff --git a/python/pyarrow/tests/test_extension_type.py b/python/pyarrow/tests/test_extension_type.py index 1adbd4e98070..35b801eca87e 100644 --- a/python/pyarrow/tests/test_extension_type.py +++ b/python/pyarrow/tests/test_extension_type.py @@ -1730,6 +1730,72 @@ def test_tensor_array_from_numpy(np_type_str): pa.FixedShapeTensorArray.from_numpy_ndarray(arr, dim_names=[0, 1]) +@pytest.mark.numpy +@pytest.mark.parametrize("np_type_str", ("int8", "int64", "float32")) +def test_tensor_array_from_list_of_ndarrays(np_type_str): + 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 + + expected = pa.FixedShapeTensorArray.from_numpy_ndarray(np.stack(elements)) + assert result.storage.equals(expected.storage) + + for scalar, original in zip(result, elements): + np.testing.assert_array_equal(scalar.to_numpy(), original) + + 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]) + + 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 + + with pytest.raises(ValueError, match="shape"): + pa.array([np.arange(6, dtype=np_dtype).reshape(3, 2)], type=tensor_type) + + 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) + + +@pytest.mark.numpy +def test_tensor_array_from_list_mixed_layout(): + # C- and F-ordered arrays with the same values must produce the same + # result, since the values are always flattened in C order. + 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]] + + +@pytest.mark.numpy +def test_tensor_array_from_list_of_0d_arrays(): + tensor_type = pa.fixed_shape_tensor(pa.int64(), ()) + result = pa.array([np.array(1, dtype=np.int64), np.array(2, dtype=np.int64)], + type=tensor_type) + assert result.type == tensor_type + assert result.storage.to_pylist() == [[1], [2]] + + @pytest.mark.numpy @pytest.mark.parametrize("tensor_type", ( pa.fixed_shape_tensor(pa.int8(), [2, 2, 3]),