misc fixes

Author: rok

python/pyarrow/__init__.py

@@ -234,7 +234,8 @@ def print_entry(label, value):
                          FixedSizeBinaryScalar, DictionaryScalar,
                          MapScalar, StructScalar, UnionScalar,
                          RunEndEncodedScalar, Bool8Scalar, ExtensionScalar,
-                         FixedShapeTensorScalar, JsonScalar, OpaqueScalar, UuidScalar)
+                         FixedShapeTensorScalar, JsonScalar, OpaqueScalar, UuidScalar,
+                         VariableShapeTensorScalar)
 
 # Buffers, allocation
 from pyarrow.lib import (DeviceAllocationType, Device, MemoryManager,

python/pyarrow/array.pxi

@@ -4974,35 +4974,37 @@ cdef class VariableShapeTensorArray(ExtensionArray):
             ]
           ]
         """
-        assert isinstance(obj, list), 'obj must be a list of numpy arrays'
+        if not isinstance(obj, list) or len(obj) == 0:
+            raise TypeError('obj must be a non-empty list of numpy arrays')
         numpy_type = obj[0].dtype
         arrow_type = from_numpy_dtype(numpy_type)
         ndim = obj[0].ndim
         permutations = [(-np.array(o.strides)).argsort(kind="stable") for o in obj]
         permutation = permutations[0]
         shapes = [np.take(o.shape, permutation) for o in obj]
 
-        if not all([o.dtype == numpy_type for o in obj]):
+        if not all(o.dtype == numpy_type for o in obj):
             raise TypeError('All numpy arrays must have matching dtype.')
 
-        if not all([o.ndim == ndim for o in obj]):
+        if not all(o.ndim == ndim for o in obj):
             raise ValueError('All numpy arrays must have matching ndim.')
 
-        if not all([np.array_equal(p, permutation) for p in permutations]):
+        if not all(np.array_equal(p, permutation) for p in permutations):
             raise ValueError('All numpy arrays must have matching permutation.')
 
         for shape in shapes:
             if len(shape) < 2:
                 raise ValueError(
-                    "Cannot convert 1D array or scalar to fixed shape tensor array")
+                    "Cannot convert 1D array or scalar to variable shape tensor array")
             if np.prod(shape) == 0:
                 raise ValueError("Expected a non-empty ndarray")
 
         values = array([np.ravel(o, order="K") for o in obj], list_(arrow_type))
         shapes = array(shapes, list_(int32(), list_size=ndim))
         struct_arr = StructArray.from_arrays([values, shapes], names=["data", "shape"])
 
-        return ExtensionArray.from_storage(variable_shape_tensor(arrow_type, ndim, permutation=permutation), struct_arr)
+        ext_type = variable_shape_tensor(arrow_type, ndim, permutation=permutation)
+        return ExtensionArray.from_storage(ext_type, struct_arr)
 
 
 cdef dict _array_classes = {

python/pyarrow/tests/test_extension_type.py

@@ -1639,8 +1639,10 @@ def test_tensor_array_from_numpy(np_type_str):
 
 
 @pytest.mark.numpy
-@pytest.mark.parametrize("value_type", (np.int8, np.int32, np.int64, np.float64))
+@pytest.mark.parametrize("value_type", (
+    "int8", "int32", "int64", "float64"))
 def test_variable_shape_tensor_class_methods(value_type):
+    value_type = getattr(np, value_type)
     ndim = 2
     shape_type = pa.list_(pa.int32(), ndim)
     arrow_type = pa.from_numpy_dtype(value_type)
@@ -1707,8 +1709,10 @@ def test_variable_shape_tensor_class_methods(value_type):
         [], dtype=value_type).reshape(shapes[1].as_py()))
 
 
-@pytest.mark.parametrize("value_type", (np.int8(), np.int64(), np.float32()))
+@pytest.mark.numpy
+@pytest.mark.parametrize("value_type", ("int8", "int64", "float32"))
 def test_variable_shape_tensor_array_from_numpy(value_type):
+    value_type = np.dtype(value_type).type()
     arrow_type = pa.from_numpy_dtype(value_type)
 
     arr = np.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]],
@@ -1729,38 +1733,31 @@ def test_variable_shape_tensor_array_from_numpy(value_type):
         pa.VariableShapeTensorArray.from_numpy_ndarray([arr.astype(np.int32()), arr])
 
     flat_arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=value_type)
-    bw = value_type.itemsize
 
     arr = flat_arr.reshape(1, 3, 4)
     tensor_array_from_numpy = pa.VariableShapeTensorArray.from_numpy_ndarray([arr])
     assert tensor_array_from_numpy.type.ndim == 3
     assert tensor_array_from_numpy.type.permutation == [0, 1, 2]
     assert tensor_array_from_numpy[0].to_tensor() == pa.Tensor.from_numpy(arr)
 
-    arr = as_strided(flat_arr, shape=(1, 2, 3, 2),
-                     strides=(bw * 12, bw * 6, bw, bw * 3))
-    tensor_array_from_numpy = pa.VariableShapeTensorArray.from_numpy_ndarray([arr])
-    assert tensor_array_from_numpy.type.ndim == 4
-    assert tensor_array_from_numpy.type.permutation == [0, 1, 3, 2]
-    assert tensor_array_from_numpy[0].to_tensor() == pa.Tensor.from_numpy(arr)
-
     arr = flat_arr.reshape(1, 2, 3, 2)
     result = pa.VariableShapeTensorArray.from_numpy_ndarray([arr])
     expected = np.array(
         [[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]], dtype=value_type)
     np.testing.assert_array_equal(result[0].to_numpy_ndarray(), expected)
 
     arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=value_type)
-    with pytest.raises(ValueError, match="Cannot convert 1D array or scalar to fixed"):
+    msg = "Cannot convert 1D array or scalar to variable"
+    with pytest.raises(ValueError, match=msg):
         pa.VariableShapeTensorArray.from_numpy_ndarray([arr])
 
     arr = np.array(1, dtype=value_type)
-    with pytest.raises(ValueError, match="Cannot convert 1D array or scalar to fixed"):
+    with pytest.raises(ValueError, match=msg):
         pa.VariableShapeTensorArray.from_numpy_ndarray([arr])
 
     arr = np.array([], dtype=value_type)
 
-    with pytest.raises(ValueError, match="Cannot convert 1D array or scalar to fixed"):
+    with pytest.raises(ValueError, match=msg):
         pa.VariableShapeTensorArray.from_numpy_ndarray([arr.reshape((0))])
 
     with pytest.raises(ValueError, match="Expected a non-empty ndarray"):
@@ -1770,6 +1767,45 @@ def test_variable_shape_tensor_array_from_numpy(value_type):
         pa.VariableShapeTensorArray.from_numpy_ndarray([arr.reshape((3, 0, 2))])
 
 
+@pytest.mark.numpy
+@pytest.mark.parametrize("dtype", (
+    "int8", "int16", "int32", "int64", "float32", "float64"))
+@pytest.mark.parametrize("order", ("F", "C"))
+def test_variable_shape_tensor_permutation_2d(dtype, order):
+    """Roundtrip 2D arrays with C and Fortran order."""
+    dtype = np.dtype(dtype)
+    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
+                   dtype=dtype, order=order)
+    expected_perm = [1, 0] if order == "F" else [0, 1]
+
+    result = pa.VariableShapeTensorArray.from_numpy_ndarray([arr])
+    assert result.type.permutation == expected_perm
+    assert result.type.ndim == 2
+
+    tensor = result[0].to_tensor()
+    assert tensor == pa.Tensor.from_numpy(arr)
+    assert list(tensor.shape) == list(arr.shape)
+    assert list(tensor.strides) == list(arr.strides)
+    np.testing.assert_array_equal(result[0].to_numpy_ndarray(), arr)
+
+    # Verify stored shape is the physical shape (permuted), per spec
+    stored_shape = result.storage.field("shape")[0].as_py()
+    assert stored_shape == list(np.take(arr.shape, expected_perm))
+
+
+@pytest.mark.numpy
+def test_variable_shape_tensor_permutation_multi_array():
+    """Multiple variable-shape 2D arrays with the same non-trivial permutation."""
+    arr1 = np.array([[10, 20], [30, 40], [50, 60]], dtype=np.int64, order="F")
+    arr2 = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64, order="F")
+    result = pa.VariableShapeTensorArray.from_numpy_ndarray([arr1, arr2])
+    assert result.type.permutation == [1, 0]
+    assert len(result) == 2
+    for i, expected in enumerate([arr1, arr2]):
+        np.testing.assert_array_equal(result[i].to_numpy_ndarray(), expected)
+        assert result[i].to_tensor() == pa.Tensor.from_numpy(expected)
+
+
 @pytest.mark.parametrize("tensor_type", (
     pa.fixed_shape_tensor(pa.int8(), [2, 2, 3]),
     pa.fixed_shape_tensor(pa.int8(), [2, 2, 3], permutation=[0, 2, 1]),
@@ -1969,7 +2005,7 @@ def test_variable_shape_tensor_type_is_picklable(pickle_module):
         'fixed_shape_tensor[value_type=int64, shape=[2,2,3], dim_names=[C,H,W]]'
     )
 ])
-def test_tensor_type_str(tensor_type, text, pickle_module):
+def test_tensor_type_str(tensor_type, text):
     tensor_type_str = tensor_type.__str__()
     assert text in tensor_type_str
 

python/pyarrow/types.pxi

@@ -2075,7 +2075,7 @@ cdef class VariableShapeTensorType(BaseExtensionType):
 
     def __reduce__(self):
         return variable_shape_tensor, (self.value_type, self.ndim,
-                                       self.permutation, self.dim_names, self.uniform_shape)
+                                       self.dim_names, self.permutation, self.uniform_shape)
 
     def __arrow_ext_scalar_class__(self):
         return VariableShapeTensorScalar
@@ -5936,8 +5936,10 @@ def variable_shape_tensor(DataType value_type, ndim, dim_names=None, permutation
         vector[optional[int64_t]] c_uniform_shape
         shared_ptr[CDataType] c_tensor_ext_type
 
-    assert value_type is not None
-    assert ndim is not None
+    if value_type is None:
+        raise TypeError('value_type must not be None')
+    if ndim is None:
+        raise TypeError('ndim must not be None')
 
     c_ndim = ndim