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

@@ -1751,16 +1751,16 @@ def test_variable_shape_tensor_array_from_numpy(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"):
+    with pytest.raises(ValueError, match="Cannot convert 1D array or scalar to variable"):
         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="Cannot convert 1D array or scalar to variable"):
         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="Cannot convert 1D array or scalar to variable"):
         pa.VariableShapeTensorArray.from_numpy_ndarray([arr.reshape((0))])
 
     with pytest.raises(ValueError, match="Expected a non-empty ndarray"):
@@ -1969,7 +1969,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